This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new fab6db20e7 [Unity][Transform] Use parameter name in BundleModelParams
(#16309)
fab6db20e7 is described below
commit fab6db20e730edaa15f9b91e66336f3a01e381ba
Author: Eric Lunderberg <[email protected]>
AuthorDate: Mon Jan 8 16:06:10 2024 -0600
[Unity][Transform] Use parameter name in BundleModelParams (#16309)
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 +++
.../relax/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 f5ee8a07ac..a9cb719d26 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<Var>(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, Expr> var_to_expr_;
};
+Function BundleModelParams(const Function& func) {
+ ModelParamBundler mutator;
+ return Downcast<Function>(mutator(func));
+}
+
namespace transform {
Pass BundleModelParams() {
runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
[=](IRModule mod,
diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h
index 8b3525c628..802099f0ab 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 0248c438c1..e3528cc357 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 test_no_model_params():
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()
