[Unity][Transform] Replace eligible operators with in-place versions in dataflow blocks

include/tvm/relax/transform.h

@@ -572,6 +572,16 @@ TVM_DLL Pass ConvertToDataflow(int min_size = 2);
  */
 TVM_DLL Pass DeadCodeElimination(Array<runtime::String> entry_functions);
 
+/*!
+ * \brief Pass that changes calls to operators that can be done in-place
+ * (generally, these are elementwise operations) in dataflow blocks into in-place implementations.
+ * Supported operators will be replaced by calls to `call_tir_inplace` that invoke in-place
+ * PrimFunc implementations of those operators (which are based on the legalizations of those
+ * operators).
+ * \return The pass.
+ */
+TVM_DLL Pass DataflowUseInplaceCalls();
+
 /*!
  * \brief Automatic mixed precision pass. Currently the pass assumes the input module to be fp32
  * only, and will automatically cast fp32 to fp16 for certain ops.

python/tvm/relax/__init__.py

@@ -63,7 +63,13 @@
 from .exec_builder import ExecBuilder
 
 # Operator
-from .op.base import call_tir, call_pure_packed, call_dps_packed, call_tir_with_grad
+from .op.base import (
+    call_tir,
+    call_tir_inplace,
+    call_pure_packed,
+    call_dps_packed,
+    call_tir_with_grad,
+)
 
 # BlockBuilder
 from .block_builder import BlockBuilder

python/tvm/relax/testing/transform.py

@@ -17,14 +17,18 @@
 # pylint: disable=unused-argument, invalid-name, no-else-return, abstract-method, arguments-differ
 """Relax transformation passes for testing"""
 
+import logging
+import os
+from typing import Dict, List, Set, Tuple
 import tvm
 from tvm import ir, relax
 from tvm.ir import transform
 from tvm.ir.module import IRModule
 from tvm.ir.transform import PassContext
 from tvm.relax import PyExprMutator
-from tvm.relax.expr import Call
+from tvm.relax.expr import Call, DataflowBlock, Var
 from tvm.relay.backend.te_compiler import select_implementation
+from tvm.runtime.object import Object
 from tvm.target import Target
 
 
@@ -128,3 +132,95 @@ def transform(self):
 def ApplyEmptyCppMutator() -> tvm.ir.transform.Pass:
     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

python/tvm/relax/transform/__init__.py

@@ -30,6 +30,7 @@
     ConvertLayout,
     ConvertToDataflow,
     DataflowBlockPass,
+    DataflowUseInplaceCalls,
     DeadCodeElimination,
     DecomposeOpsForInference,
     DecomposeOpsForTraining,

python/tvm/relax/transform/transform.py

@@ -252,6 +252,22 @@ def RemovePurityChecking() -> tvm.ir.transform.Pass:
     return _ffi_api.RemovePurityChecking()  # type: ignore
 
 
+def DataflowUseInplaceCalls() -> tvm.ir.transform.Pass:
+    """
+    Pass that changes calls to operators that can be done in-place
+    (generally, these are elementwise operations) into in-place implementations.
+    Supported operators will be replaced by calls to `call_tir_inplace` that invoke
+    in-place PrimFunc implementations of those operators (which are based on the legalizations of
+    those operators).
+
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+        The pass
+    """
+    return _ffi_api.DataflowUseInplaceCalls()
+
+
 def LambdaLift() -> tvm.ir.transform.Pass:
     """A pass that lifts local functions into global.