[Unity][Transform] Implement ExpandTupleArguments (#16115)

[Unity] Implement ExpandTupleArguments transform

Currently, the `FuseOps` and `FuseTIR` passes have a large amount of
added complexity to identify and handle partial use of tuple
arguments.  The handling partial use of tuples could be significantly
simpler if performed in multiple steps.

1. Perform `FuseOps`.  Any tuple variables that are used by the fused
   function are passed as-is.

2. Expand any parameters that are passed as a tuple.  Any unused
   tensors that were included in a partially-used tuple will be
   converted to unused parameters.

3. Remove any unused parameters.  Any unused tensors that were
   included in a partially-used tuple will be removed in this
   step.

4. Perform `FuseTIR`.  No checking for tuple arguments, either partial
   or full, is required at this step.

This PR implements `relax.transform.ExpandTupleArguments`, which is
step (2) in this process.

include/tvm/relax/transform.h

@@ -275,6 +275,12 @@ TVM_DLL Pass LiftTransformParams();
  */
 TVM_DLL Pass UpdateVDevice(VDevice new_vdevice, int64_t index);
 
+/*! \brief Expand tuple arguments to internal functions
+ *
+ * \return The Pass
+ */
+TVM_DLL Pass ExpandTupleArguments();
+
 /*! \brief Remove unused outputs from internal functions
  *
  * \return The Pass

python/tvm/relax/transform/__init__.py

@@ -33,6 +33,7 @@
     DecomposeOpsForInference,
     DecomposeOpsForTraining,
     EliminateCommonSubexpr,
+    ExpandTupleArguments,
     FewShotTuning,
     FoldConstant,
     FunctionPass,

python/tvm/relax/transform/transform.py

@@ -558,6 +558,16 @@ def FoldConstant() -> tvm.ir.transform.Pass:
     return _ffi_api.FoldConstant()  # type: ignore
 
 
+def ExpandTupleArguments() -> tvm.ir.transform.Pass:
+    """Expand tuple arguments to internal functions
+
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+    """
+    return _ffi_api.ExpandTupleArguments()  # type: ignore
+
+
 def RemoveUnusedOutputs() -> tvm.ir.transform.Pass:
     """Remove unused outputs from internal functions
 

src/relax/transform/expand_tuple_arguments.cc

@@ -0,0 +1,187 @@
+/*
+ * 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.
+ */
+
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+
+#include <algorithm>
+#include <tuple>
+
+namespace tvm {
+namespace relax {
+
+namespace {
+
+template <typename T, typename U>
+using PMap = std::unordered_map<T, U, ObjectPtrHash, ObjectPtrEqual>;
+
+Optional<Function> ExpandParams(Function func) {
+  bool is_exposed = func->attrs.GetAttr<String>(tvm::attr::kGlobalSymbol).defined();
+  if (is_exposed) return NullOpt;
+
+  bool has_tuple_param = std::any_of(
+      func->params.begin(), func->params.end(),
+      [](const Var& param) -> bool { return param->struct_info_.as<TupleStructInfoNode>(); });
+
+  if (!has_tuple_param) return NullOpt;
+
+  Array<Var> params;
+  Array<Binding> bindings;
+
+  std::function<void(const Var&)> expand_param = [&](const Var& param) {
+    if (auto sinfo = param->struct_info_.as<TupleStructInfoNode>()) {
+      Array<Expr> internal_tuple;
+      for (size_t i = 0; i < sinfo->fields.size(); i++) {
+        auto name = static_cast<const std::stringstream&>(std::stringstream()
+                                                          << param->name_hint() << "_" << i)
+                        .str();
+        Var new_param(name, sinfo->fields[i]);
+        internal_tuple.push_back(new_param);
+        expand_param(new_param);
+      }
+      bindings.push_back(VarBinding(param, Tuple(internal_tuple)));
+    } else {
+      params.push_back(param);
+    }
+  };
+
+  for (const auto& param : func->params) {
+    expand_param(param);
+  }
+
+  FuncStructInfo new_sinfo(params.Map([](const auto& var) { return GetStructInfo(var); }),
+                           func->ret_struct_info,
+                           Downcast<FuncStructInfo>(func->struct_info_)->purity);
+
+  auto write_ptr = func.CopyOnWrite();
+  write_ptr->params = params;
+  write_ptr->body = SeqExpr({BindingBlock(bindings)}, func->body);
+  write_ptr->struct_info_ = new_sinfo;
+
+  return func;
+}
+
+class TupleExpander : public ExprMutator {
+ public:
+  explicit TupleExpander(PMap<GlobalVar, GlobalVar> callees) : replacements_(callees) {}
+
+  using ExprMutator::VisitExpr_;
+
+  Expr VisitExpr_(const CallNode* op) override {
+    auto node = Downcast<Call>(ExprMutator::VisitExpr_(op));
+
+    if (auto gvar = node->op.as<GlobalVar>()) {
+      if (auto it = replacements_.find(gvar.value()); it != replacements_.end()) {
+        Array<Expr> new_args;
+
+        std::function<void(const Expr&)> expand_arg = [&](const Expr& arg) {
+          if (auto sinfo = arg->struct_info_.as<TupleStructInfoNode>()) {
+            for (size_t i = 0; i < sinfo->fields.size(); i++) {
+              expand_arg(TupleGetItem(arg, i));
+            }
+          } else {
+            new_args.push_back(arg);
+          }
+        };
+
+        for (const auto& arg : node->args) {
+          expand_arg(arg);
+        }
+
+        auto write_ptr = node.CopyOnWrite();
+        write_ptr->op = it->second;
+        write_ptr->args = new_args;
+      }
+    }
+
+    return node;
+  }
+
+  PMap<GlobalVar, GlobalVar> replacements_;
+};
+
+}  // namespace
+
+namespace transform {
+
+Pass ExpandTupleArguments() {
+  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
+      [=](IRModule mod, PassContext pc) -> IRModule {
+    PMap<GlobalVar, GlobalVar> gvar_replacements;
+
+    {
+      PMap<GlobalVar, Function> new_callees;
+
+      for (const auto& [gvar, base_func] : mod->functions) {
+        if (auto func = base_func.as<Function>()) {
+          if (auto opt = ExpandParams(func.value())) {
+            auto new_func = opt.value();
+            GlobalVar new_gvar(gvar->name_hint, new_func->checked_type_);
+            new_gvar->struct_info_ = new_func->struct_info_;
+            gvar_replacements[gvar] = new_gvar;
+            new_callees[new_gvar] = new_func;
+          }
+        }
+      }
+
+      if (gvar_replacements.empty()) {
+        return mod;
+      }
+      auto write_ptr = mod.CopyOnWrite();
+      for (auto [old_gvar, new_gvar] : gvar_replacements) {
+        write_ptr->Remove(old_gvar);
+        write_ptr->Add(new_gvar, new_callees.at(new_gvar));
+      }
+    }
+
+    TupleExpander mutator(std::move(gvar_replacements));
+
+    IRModule caller_updates;
+
+    for (const auto& [gvar, base_func] : mod->functions) {
+      if (auto func = base_func.as<Function>()) {
+        auto mutated = Downcast<Function>(mutator.VisitExpr(func.value()));
+        if (!mutated.same_as(base_func)) {
+          caller_updates->Add(gvar, mutated);
+        }
+      }
+    }
+
+    if (caller_updates->functions.size()) {
+      mod.CopyOnWrite()->Update(caller_updates);
+    }
+    return mod;
+  };
+  auto inner_pass = CreateModulePass(pass_func, 0, "ExpandTupleArgumentsInner", {});
+
+  return tvm::transform::Sequential(
+      {
+          inner_pass,
+          CanonicalizeBindings(),
+          DeadCodeElimination({}),
+      },
+      "ExpandTupleArguments");
+}
+
+TVM_REGISTER_GLOBAL("relax.transform.ExpandTupleArguments").set_body_typed(ExpandTupleArguments);
+
+}  // namespace transform
+
+}  // namespace relax
+}  // namespace tvm

tests/python/relax/test_transform_expand_tuple_args.py

@@ -0,0 +1,79 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I, relax as R, tir as T
+
+
+class BaseCompare(tvm.testing.CompareBeforeAfter):
+    transform = tvm.relax.transform.ExpandTupleArguments()
+
+
+class TestSimple(BaseCompare):
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(A: R.Tensor, B: R.Tensor):
+            return Before.func((A, B))
+
+        @R.function(private=True)
+        def func(args: R.Tuple([R.Tensor, R.Tensor])) -> R.Tensor:
+            return args[0]
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(A: R.Tensor, B: R.Tensor):
+            return Expected.func(A, B)
+
+        @R.function(private=True)
+        def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+            return A
+
+
+class TestNested(BaseCompare):
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(A: R.Tensor, B: R.Tensor, C: R.Tensor, D: R.Tensor) -> R.Tensor:
+            return Before.func(((A, B), (C, D)))
+
+        @R.function(private=True)
+        def func(
+            args: R.Tuple(
+                [
+                    R.Tuple([R.Tensor, R.Tensor]),
+                    R.Tuple([R.Tensor, R.Tensor]),
+                ]
+            )
+        ) -> R.Tensor:
+            return args[0][1]
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(A: R.Tensor, B: R.Tensor, C: R.Tensor, D: R.Tensor) -> R.Tensor:
+            return Expected.func(A, B, C, D)
+
+        @R.function(private=True)
+        def func(A: R.Tensor, B: R.Tensor, C: R.Tensor, D: R.Tensor) -> R.Tensor:
+            return B
+
+
+if __name__ == "__main__":
+    tvm.testing.main()