chore: import upstream snapshot with attribution

python/tvm/relax/testing/transform.py

@@ -0,0 +1,125 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-argument, invalid-name, no-else-return, abstract-method, arguments-differ
+"""Relax transformation passes for testing"""
+
+import logging
+import os
+
+import tvm_ffi
+
+import tvm
+from tvm.ir import Call
+from tvm.ir.module import IRModule
+from tvm.relax.expr import DataflowBlock, Var
+from tvm.runtime import Object
+
+
+def ApplyEmptyCppMutator() -> tvm.ir.transform.Pass:
+    """Create empty cpp mutator"""
+    packed_func = tvm.get_global_func("relax.testing.transform.ApplyEmptyCppMutator")
+    return packed_func()
+
+
+def dataflow_liveness_analysis(block: DataflowBlock) -> dict[Var, tuple[int, int]]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_liveness_analysis is exposed for testing only.")
+
+    live_ranges = tvm.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")(block)  # type: ignore
+    ret = {}
+    for var, live_range in live_ranges.items():
+        ret[var] = tuple(live_range)
+    return ret  # type: ignore
+
+
+def dataflow_alias_analysis(
+    block: DataflowBlock, inputs: list[Var]
+) -> tuple[dict[Var, set[int]], dict[int, list[set[int]]]]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_alias_analysis is exposed for testing only.")
+
+    alias_sets, tuple_map = tvm.get_global_func("relax.testing.transform.DataflowAliasAnalysis")(
+        block,
+        inputs,
+    )  # type: ignore
+    res_alias_sets = {}
+    res_tuple_map = {}
+    for var, alias_set in alias_sets.items():
+        res_alias_sets[var] = set(alias_set)
+    for idx, elem_alias_sets in tuple_map.items():
+        res_tuple_map[idx] = [set(alias_set) for alias_set in elem_alias_sets]
+    return res_alias_sets, res_tuple_map  # type: ignore
+
+
+@tvm_ffi.register_object("relax.transform.InplaceOpportunity")
+class InplaceOpportunity(Object):
+    """
+    Represents an opportunity to make a binding in-place. Exposed only for testing;
+    the constructor is not exposed.
+
+    Parameters:
+    -----------
+    binding_idx: int
+        Index of the binding within its block
+
+    arg_idxs: List[int]
+        Indices of arguments that are eligible to be used as in-place targets.
+    """
+
+    def __init__(self, _binding_idx, _arg_idxs):
+        raise NotImplementedError("Constructor for InplaceOpportunity not exposed!")
+
+
+def dataflow_inplace_analysis(
+    block: DataflowBlock, inputs: list[Var], mod: IRModule
+) -> tuple[list[tuple[int, set[int]]], list[tuple[int, set[int]]]]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_inplace_analysis is exposed for testing only.")
+    index_lists = tvm.get_global_func("relax.testing.transform.DataflowInplaceAnalysis")(
+        block, inputs, mod
+    )  # type: ignore
+
+    def convert(opp_list):
+        return list(map(lambda opp: (int(opp.binding_idx), set(map(int, opp.arg_idxs))), opp_list))
+
+    return (convert(index_lists[0]), convert(index_lists[1]))  # type: ignore
+
+
+def dataflow_single_inplace_call(
+    mod: IRModule, call: Call, inplace_indices: list[int]
+) -> tuple[Call, IRModule]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_single_inplace_call is exposed for testing only.")
+
+    ret = tvm.get_global_func("relax.testing.transform.SingleInplaceCall")(
+        mod,
+        call,
+        inplace_indices,
+    )  # type: ignore
+    return (ret[0], ret[1])  # type: ignore