From 3514a56c9c93939ce311cdce65881c329bc6fde1 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 11 Dec 2023 17:11:14 +0000 Subject: [PATCH] [Unity][Transform] Use parameter name in BundleModelParams Prior to this commit, the `BundleModelParams` would replace model parameters with `param_tuple[index]` within expressions. These nested expressions would then be normalized, resulting in `gv = param_tuple[index]` or `lv = param_tuple[index]` variable definitions. These auto-generated `gv` and `lv` names make it quite difficult to determine which model parameter is being used. This commit updates the `BundleModelParams` transform to explicitly produce the bound variable, `orig_param_name = param_tuple[index]`, preserving human-readable names from the parameters. --- src/relax/transform/bundle_model_params.cc | 7 +- src/relax/transform/utils.h | 12 +++ .../test_transform_bundle_model_params.py | 91 +++++++++++++++++++ 3 files changed, 109 insertions(+), 1 deletion(-) diff --git a/src/relax/transform/bundle_model_params.cc b/src/relax/transform/bundle_model_params.cc index f5ee8a07ac44..a9cb719d26d9 100644 --- a/src/relax/transform/bundle_model_params.cc +++ b/src/relax/transform/bundle_model_params.cc @@ -74,7 +74,7 @@ class ModelParamBundler : public ExprMutator { Expr VisitExpr_(const VarNode* op) override { auto var = GetRef(op); if (auto it = var_to_expr_.find(var); it != var_to_expr_.end()) { - return (*it).second; + return builder_->Emit((*it).second, op->name_hint()); } else { return ExprMutator::VisitExpr_(op); } @@ -84,6 +84,11 @@ class ModelParamBundler : public ExprMutator { Map var_to_expr_; }; +Function BundleModelParams(const Function& func) { + ModelParamBundler mutator; + return Downcast(mutator(func)); +} + namespace transform { Pass BundleModelParams() { runtime::TypedPackedFunc pass_func = [=](IRModule mod, diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h index 8b3525c628fe..802099f0ab4b 100644 --- a/src/relax/transform/utils.h +++ b/src/relax/transform/utils.h @@ -421,6 +421,18 @@ Expr EliminateCommonSubexpr(const Expr& expr, bool call_only = false); */ Expr CanonicalizeBindings(const Expr& expr); +/* \brief Remove use of trivial bindings + * + * Utility for converting from individual model parameters to a single + * parameter with a tuple of parameters. If the `kNumInput` attribute + * is absent, no model parameters are present, so no updates are made. + * + * \param func The function to be updated. + * + * \ret The updated function. + */ +Function BundleModelParams(const Function& func); + } // namespace relax } // namespace tvm diff --git a/tests/python/relax/test_transform_bundle_model_params.py b/tests/python/relax/test_transform_bundle_model_params.py index 0248c438c1fa..e3528cc357e4 100644 --- a/tests/python/relax/test_transform_bundle_model_params.py +++ b/tests/python/relax/test_transform_bundle_model_params.py @@ -102,5 +102,96 @@ def main( tvm.ir.assert_structural_equal(after, Expected) +def test_dataflow(): + """Parameters can be substituted into a dataflow block""" + + @tvm.script.ir_module + class Before: + @R.function + def main( + a: R.Tensor([16], "float32"), + b: R.Tensor([16], "float32"), + c: R.Tensor([16], "float32"), + ) -> R.Tensor([16], "float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + expr = a + expr = R.add(expr, b) + expr = R.add(expr, c) + R.output(expr) + return expr + + @tvm.script.ir_module + class Expected: + @R.function + def main( + a: R.Tensor([16], "float32"), + params: R.Tuple(R.Tensor([16], "float32"), R.Tensor([16], "float32")), + ) -> R.Tensor([16], "float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + expr = a + b = params[0] + expr = R.add(expr, b) + c = params[1] + expr = R.add(expr, c) + R.output(expr) + return expr + + mod = Before + after = relax.transform.BundleModelParams()(mod) + tvm.ir.assert_structural_equal(after, Expected) + + +def test_variable_names(): + """Parameters retain their names within the updated function + + For readability, the parameter names should be used to generate + the new variable names. + + Like `test_basic`, but explicitly checks the names of bound + variables. + """ + + @tvm.script.ir_module + class Before: + @R.function + def main( + a: R.Tensor([16], "float32"), + b: R.Tensor([16], "float32"), + c: R.Tensor([16], "float32"), + ) -> R.Tensor([16], "float32"): + R.func_attr({"num_input": 1}) + expr = a + expr = R.add(expr, b) + expr = R.add(expr, c) + return expr + + @tvm.script.ir_module + class Expected: + @R.function + def main( + a: R.Tensor([16], "float32"), + params: R.Tuple(R.Tensor([16], "float32"), R.Tensor([16], "float32")), + ) -> R.Tensor([16], "float32"): + R.func_attr({"num_input": 1}) + expr = a + b = params[0] + expr = R.add(expr, b) + c = params[1] + expr = R.add(expr, c) + return expr + + mod = Before + after = relax.transform.BundleModelParams()(mod) + tvm.ir.assert_structural_equal(after, Expected) + + for binding, expected_binding in zip( + after["main"].body.blocks[0].bindings, + Expected["main"].body.blocks[0].bindings, + ): + assert binding.var.name_hint == expected_binding.var.name_hint + + if __name__ == "__main__": tvm.testing.main()