From 2d2c0860d2530cb23cee423cfd92b4ccd2025c1b Mon Sep 17 00:00:00 2001 From: Yuhx141 <155034285+Yuhx141@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:05:53 +0800 Subject: [PATCH] [FIX][Relax] Preserve private functions referenced as values --- src/relax/transform/expand_tuple_arguments.cc | 4 ++ src/relax/transform/global_var_reference.h | 64 +++++++++++++++++++ src/relax/transform/inline_functions.cc | 4 +- .../transform/remove_unused_parameters.cc | 3 + .../relax/test_transform_expand_tuple_args.py | 18 ++++++ ...test_transform_inline_private_functions.py | 20 ++++++ ...test_transform_remove_unused_parameters.py | 26 ++++++++ 7 files changed, 138 insertions(+), 1 deletion(-) create mode 100644 src/relax/transform/global_var_reference.h diff --git a/src/relax/transform/expand_tuple_arguments.cc b/src/relax/transform/expand_tuple_arguments.cc index 772dbcb8b5c2..9d9e59811f91 100644 --- a/src/relax/transform/expand_tuple_arguments.cc +++ b/src/relax/transform/expand_tuple_arguments.cc @@ -25,6 +25,8 @@ #include #include +#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 gvar_replacements; { PMap new_callees; for (const auto& [gvar, base_func] : mod->functions) { + if (value_references.count(gvar)) continue; if (auto func = base_func.as()) { if (auto opt = ExpandParams(func.value())) { auto new_func = opt.value(); diff --git a/src/relax/transform/global_var_reference.h b/src/relax/transform/global_var_reference.h new file mode 100644 index 000000000000..6c1cd3cf2e10 --- /dev/null +++ b/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 +#include + +#include +#include + +namespace tvm { +namespace relax { + +// References used as values require the original function definition and signature. +class GlobalVarValueCollector : public ExprVisitor { + public: + static std::unordered_set Collect( + const IRModule& mod) { + GlobalVarValueCollector collector; + for (const auto& [gvar, base_func] : mod->functions) { + if (auto func = base_func.as()) { + collector.VisitExpr(func.value()); + } + } + return std::move(collector.references_); + } + + private: + void VisitExpr_(const GlobalVarNode* op) final { references_.insert(ffi::GetRef(op)); } + + void VisitExpr_(const CallNode* op) final { + if (!op->op.as()) { + VisitExpr(op->op); + } + for (const auto& arg : op->args) { + VisitExpr(arg); + } + } + + std::unordered_set references_; +}; + +} // namespace relax +} // namespace tvm + +#endif // TVM_RELAX_TRANSFORM_GLOBAL_VAR_REFERENCE_H_ diff --git a/src/relax/transform/inline_functions.cc b/src/relax/transform/inline_functions.cc index f2cbb10f87cf..4ef3c710d61e 100644 --- a/src/relax/transform/inline_functions.cc +++ b/src/relax/transform/inline_functions.cc @@ -28,6 +28,7 @@ #include #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, Function> replacements; for (const auto& [gvar, base_func] : mod->functions) { if (auto opt = base_func.as()) { 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); } } diff --git a/src/relax/transform/remove_unused_parameters.cc b/src/relax/transform/remove_unused_parameters.cc index 9e318919d85f..074289520f4f 100644 --- a/src/relax/transform/remove_unused_parameters.cc +++ b/src/relax/transform/remove_unused_parameters.cc @@ -28,6 +28,7 @@ #include #include +#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> 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()) { if (auto callee_res = AnalyzeCallee(func.value())) { auto new_func = callee_res->func; diff --git a/tests/python/relax/test_transform_expand_tuple_args.py b/tests/python/relax/test_transform_expand_tuple_args.py index 11be7b6f7e6f..157c588fe2fd 100644 --- a/tests/python/relax/test_transform_expand_tuple_args.py +++ b/tests/python/relax/test_transform_expand_tuple_args.py @@ -80,5 +80,23 @@ def func(A: R.Tensor, B: R.Tensor, C: R.Tensor, D: R.Tensor) -> R.Tensor: tvm.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.ir.assert_structural_equal(after, Before) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_transform_inline_private_functions.py b/tests/python/relax/test_transform_inline_private_functions.py index 349c776b4022..d4007793f53a 100644 --- a/tests/python/relax/test_transform_inline_private_functions.py +++ b/tests/python/relax/test_transform_inline_private_functions.py @@ -102,5 +102,25 @@ def subroutine() -> R.Tensor([], "int64"): tvm.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.ir.assert_structural_equal(after, Before) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_transform_remove_unused_parameters.py b/tests/python/relax/test_transform_remove_unused_parameters.py index 116fed4838a2..0df66e6fe54c 100644 --- a/tests/python/relax/test_transform_remove_unused_parameters.py +++ b/tests/python/relax/test_transform_remove_unused_parameters.py @@ -236,5 +236,31 @@ def func(A: R.Tensor([m_func, n_func], "float32")) -> R.Tensor([m_func, n_func], tvm.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.ir.assert_structural_equal(after, Before) + + if __name__ == "__main__": tvm.testing.main()