[FIX][Relax] Preserve closure-referenced private functions

src/relax/transform/expand_tuple_arguments.cc

@@ -25,6 +25,8 @@
 #include <algorithm>
 #include <tuple>
 
+#include "global_var_reference.h"
+
 namespace tvm {
 namespace relax {
 
@@ -123,12 +125,14 @@ namespace transform {
 
 Pass ExpandTupleArguments() {
   auto pass_func = [=](IRModule mod, PassContext pc) -> IRModule {
+    auto value_references = GlobalVarValueCollector::Collect(mod);
     PMap<GlobalVar, GlobalVar> gvar_replacements;
 
     {
       PMap<GlobalVar, Function> new_callees;
 
       for (const auto& [gvar, base_func] : mod->functions) {
+        if (value_references.count(gvar)) continue;
         if (auto func = base_func.as<Function>()) {
           if (auto opt = ExpandParams(func.value())) {
             auto new_func = opt.value();

src/relax/transform/global_var_reference.h

@@ -0,0 +1,64 @@
+/*
+ * 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.
+ */
+
+#ifndef TVM_RELAX_TRANSFORM_GLOBAL_VAR_REFERENCE_H_
+#define TVM_RELAX_TRANSFORM_GLOBAL_VAR_REFERENCE_H_
+
+#include <tvm/ir/module.h>
+#include <tvm/relax/expr_functor.h>
+
+#include <unordered_set>
+#include <utility>
+
+namespace tvm {
+namespace relax {
+
+// References used as values require the original function definition and signature.
+class GlobalVarValueCollector : public ExprVisitor {
+ public:
+  static std::unordered_set<GlobalVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> Collect(
+      const IRModule& mod) {
+    GlobalVarValueCollector collector;
+    for (const auto& [gvar, base_func] : mod->functions) {
+      if (auto func = base_func.as<Function>()) {
+        collector.VisitExpr(func.value());
+      }
+    }
+    return std::move(collector.references_);
+  }
+
+ private:
+  void VisitExpr_(const GlobalVarNode* op) final { references_.insert(ffi::GetRef<GlobalVar>(op)); }
+
+  void VisitExpr_(const CallNode* op) final {
+    if (!op->op.as<GlobalVarNode>()) {
+      VisitExpr(op->op);
+    }
+    for (const auto& arg : op->args) {
+      VisitExpr(arg);
+    }
+  }
+
+  std::unordered_set<GlobalVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> references_;
+};
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_TRANSFORM_GLOBAL_VAR_REFERENCE_H_

src/relax/transform/inline_functions.cc

@@ -28,6 +28,7 @@
 #include <vector>
 
 #include "../../support/ordered_set.h"
+#include "global_var_reference.h"
 #include "utils.h"
 
 namespace tvm {
@@ -176,12 +177,13 @@ namespace transform {
 
 Pass InlinePrivateFunctions() {
   auto pass_func = [=](IRModule mod, PassContext pc) {
+    auto value_references = GlobalVarValueCollector::Collect(mod);
     ffi::Map<ffi::Variant<ffi::String, GlobalVar>, Function> replacements;
     for (const auto& [gvar, base_func] : mod->functions) {
       if (auto opt = base_func.as<relax::Function>()) {
         auto func = opt.value();
         bool is_private = !func->GetAttr(tvm::attr::kGlobalSymbol).has_value();
-        if (is_private) {
+        if (is_private && !value_references.count(gvar)) {
           replacements.Set(gvar, func);
         }
       }

src/relax/transform/remove_unused_parameters.cc

@@ -28,6 +28,7 @@
 #include <optional>
 #include <tuple>
 
+#include "global_var_reference.h"
 #include "utils.h"
 
 namespace tvm {
@@ -191,12 +192,14 @@ namespace transform {
 
 Pass RemoveUnusedParameters() {
   auto pass_func = [=](IRModule mod, PassContext pc) -> IRModule {
+    auto value_references = GlobalVarValueCollector::Collect(mod);
     PMap<GlobalVar, std::function<Call(Call)>> callsite_updaters;
 
     {
       IRModule new_callees;
 
       for (const auto& [gvar, base_func] : mod->functions) {
+        if (value_references.count(gvar)) continue;
         if (auto func = base_func.as<Function>()) {
           if (auto callee_res = AnalyzeCallee(func.value())) {
             auto new_func = callee_res->func;

tests/python/relax/test_transform_expand_tuple_args.py

@@ -81,5 +81,23 @@ def func(A: R.Tensor, B: R.Tensor, C: R.Tensor, D: R.Tensor) -> R.Tensor:
     tvm_tirx.ir.assert_structural_equal(After, Expected)
 
 
+def test_preserve_closure_callee():
+    @I.ir_module
+    class Before:
+        @R.function(private=True)
+        def pair_add(pair: R.Tuple(R.Tensor((2,), "float32"), R.Tensor((2,), "float32"))):
+            return R.add(pair[0], pair[1])
+
+        @R.function
+        def main(x: R.Tensor((2,), "float32"), y: R.Tensor((2,), "float32")):
+            closure = R.make_closure(Before.pair_add, ())
+            return R.invoke_pure_closure(closure, ((x, y),), ty_args=R.Tensor((2,), "float32"))
+
+    assert tvm.relax.analysis.check_well_formed(Before, check_ty=True)
+    after = tvm.relax.transform.ExpandTupleArguments()(Before)
+    assert tvm.relax.analysis.check_well_formed(after, check_ty=True)
+    tvm_tirx.ir.assert_structural_equal(after, Before)
+
+
 if __name__ == "__main__":
     tvm_tirx.testing.main()

tests/python/relax/test_transform_inline_private_functions.py

@@ -103,5 +103,25 @@ def subroutine() -> R.Tensor([], "int64"):
     tvm_tirx.ir.assert_structural_equal(Expected, After)
 
 
+def test_preserve_closure_callee():
+    @I.ir_module
+    class Before:
+        @R.function(private=True)
+        def lifted(value: R.Tensor((2,), "float32"), env: R.Tensor((2,), "float32")):
+            return R.add(value, env)
+
+        @R.function
+        def main(x: R.Tensor((2,), "float32"), y: R.Tensor((2,), "float32")):
+            direct = Before.lifted(x, y)
+            closure = R.make_closure(Before.lifted, (x,))
+            invoked = R.invoke_pure_closure(closure, (y,), ty_args=R.Tensor((2,), "float32"))
+            return R.add(direct, invoked)
+
+    assert relax.analysis.check_well_formed(Before, check_ty=True)
+    after = relax.transform.InlinePrivateFunctions()(Before)
+    assert relax.analysis.check_well_formed(after, check_ty=True)
+    tvm_tirx.ir.assert_structural_equal(after, Before)
+
+
 if __name__ == "__main__":
     tvm_tirx.testing.main()

tests/python/relax/test_transform_remove_unused_parameters.py

@@ -237,5 +237,31 @@ def func(A: R.Tensor([m_func, n_func], "float32")) -> R.Tensor([m_func, n_func],
     tvm_tirx.ir.assert_structural_equal(After, Expected)
 
 
+def test_preserve_closure_callee():
+    @I.ir_module
+    class Before:
+        @R.function(private=True)
+        def lifted(
+            value: R.Tensor((2,), "float32"),
+            env: R.Tensor((2,), "float32"),
+            unused: R.Tensor((2,), "float32"),
+        ):
+            return R.add(value, env)
+
+        @R.function
+        def main(
+            x: R.Tensor((2,), "float32"),
+            y: R.Tensor((2,), "float32"),
+            z: R.Tensor((2,), "float32"),
+        ):
+            closure = R.make_closure(Before.lifted, (x, z))
+            return R.invoke_pure_closure(closure, (y,), ty_args=R.Tensor((2,), "float32"))
+
+    assert tvm.relax.analysis.check_well_formed(Before, check_ty=True)
+    after = tvm.relax.transform.RemoveUnusedParameters()(Before)
+    assert tvm.relax.analysis.check_well_formed(after, check_ty=True)
+    tvm_tirx.ir.assert_structural_equal(after, Before)
+
+
 if __name__ == "__main__":
     tvm_tirx.testing.main()