This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fcfc05bb29 [Transform] Allow explicit name of bundled model parameters
(#16597)
fcfc05bb29 is described below
commit fcfc05bb291894a0b7bfcaefd4affddf587f72ea
Author: Eric Lunderberg <[email protected]>
AuthorDate: Thu Feb 22 11:19:12 2024 -0600
[Transform] Allow explicit name of bundled model parameters (#16597)
In `BundleModelParams`, allow the user to specify a name for the tuple
parameters. If unspecified, defaults to the previous name
`"model_params"`.
---
python/tvm/relax/transform/transform.py | 11 ++++--
src/relax/transform/bundle_model_params.cc | 14 ++++----
src/relax/transform/utils.h | 5 ++-
.../relax/test_transform_bundle_model_params.py | 40 ++++++++++++++++++++++
4 files changed, 61 insertions(+), 9 deletions(-)
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index b2aaa3e331..c017f0cda7 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -852,7 +852,7 @@ def LiftTransformParams() -> tvm.ir.transform.Pass:
return _ffi_api.LiftTransformParams() # type: ignore
-def BundleModelParams() -> tvm.ir.transform.Pass:
+def BundleModelParams(param_tuple_name: Optional[str] = None) ->
tvm.ir.transform.Pass:
"""Bundle several model parameters into a single tuple paramters
For each function, if the function has the attribute "num_input",
@@ -860,13 +860,20 @@ def BundleModelParams() -> tvm.ir.transform.Pass:
Run-time parameters (e.g. activations) are the first `num_input`
parameters, and the remainder are compile-time weights.
+ Parameters
+ ----------
+ param_tuple_name: Optional[str]
+
+ The name of the tuple parameter. If unspecified, defaults to
+ "model_params".
+
Returns
-------
ret : tvm.transform.Pass
The registered pass for lifting transformation of parameters.
"""
- return _ffi_api.BundleModelParams() # type: ignore
+ return _ffi_api.BundleModelParams(param_tuple_name) # type: ignore
def LegalizeOps(
diff --git a/src/relax/transform/bundle_model_params.cc
b/src/relax/transform/bundle_model_params.cc
index a9cb719d26..f5798049ef 100644
--- a/src/relax/transform/bundle_model_params.cc
+++ b/src/relax/transform/bundle_model_params.cc
@@ -35,7 +35,8 @@ namespace relax {
class ModelParamBundler : public ExprMutator {
public:
- ModelParamBundler() {}
+ explicit ModelParamBundler(Optional<String> param_tuple_name)
+ : param_tuple_name_(param_tuple_name) {}
Expr VisitExpr_(const FunctionNode* op) override {
Function func = GetRef<Function>(op);
@@ -59,7 +60,7 @@ class ModelParamBundler : public ExprMutator {
param_tuple.push_back(GetStructInfo(func->params[i]));
}
- Var var_param_tuple("model_params", TupleStructInfo(param_tuple));
+ Var var_param_tuple(param_tuple_name_.value_or("model_params"),
TupleStructInfo(param_tuple));
params.push_back(var_param_tuple);
for (size_t i = num_input; i < func->params.size(); i++) {
@@ -81,21 +82,22 @@ class ModelParamBundler : public ExprMutator {
}
private:
+ Optional<String> param_tuple_name_;
Map<Var, Expr> var_to_expr_;
};
-Function BundleModelParams(const Function& func) {
- ModelParamBundler mutator;
+Function BundleModelParams(const Function& func, Optional<String>
param_tuple_name) {
+ ModelParamBundler mutator(param_tuple_name);
return Downcast<Function>(mutator(func));
}
namespace transform {
-Pass BundleModelParams() {
+Pass BundleModelParams(Optional<String> param_tuple_name) {
runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
[=](IRModule mod,
PassContext pc) {
IRModule updates;
- ModelParamBundler mutator;
+ ModelParamBundler mutator(param_tuple_name);
for (const auto& [gvar, func] : mod->functions) {
if (auto opt = func.as<relax::Function>()) {
diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h
index 802099f0ab..1ad714972c 100644
--- a/src/relax/transform/utils.h
+++ b/src/relax/transform/utils.h
@@ -429,9 +429,12 @@ Expr CanonicalizeBindings(const Expr& expr);
*
* \param func The function to be updated.
*
+ * \param param_tuple_name The name of the tuple parameter. If
+ * unspecified, defaults to "model_params"
+ *
* \ret The updated function.
*/
-Function BundleModelParams(const Function& func);
+Function BundleModelParams(const Function& func, Optional<String>
param_tuple_name = NullOpt);
} // 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 e3528cc357..415a883f16 100644
--- a/tests/python/relax/test_transform_bundle_model_params.py
+++ b/tests/python/relax/test_transform_bundle_model_params.py
@@ -193,5 +193,45 @@ def test_variable_names():
assert binding.var.name_hint == expected_binding.var.name_hint
+def test_bundled_param_name():
+ """The tuple parameter can have an explicit name"""
+
+ @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"),
+ custom_tuple_name: R.Tuple(R.Tensor([16], "float32"),
R.Tensor([16], "float32")),
+ ) -> R.Tensor([16], "float32"):
+ R.func_attr({"num_input": 1})
+ expr = a
+ b = custom_tuple_name[0]
+ expr = R.add(expr, b)
+ c = custom_tuple_name[1]
+ expr = R.add(expr, c)
+ return expr
+
+ mod = Before
+ after = relax.transform.BundleModelParams("custom_tuple_name")(mod)
+ tvm.ir.assert_structural_equal(after, Expected)
+
+ for param, expected_param in zip(after["main"].params,
Expected["main"].params):
+ assert param.name_hint == expected_param.name_hint
+
+
if __name__ == "__main__":
tvm.testing.main()
