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 7789b248fb [Unity][Transform] Handle symbolic variables in LambdaLift
(#16411)
7789b248fb is described below
commit 7789b248fbe0776fc29efdc0d772a0deeb95f8a3
Author: Eric Lunderberg <[email protected]>
AuthorDate: Tue Jan 23 10:09:34 2024 -0600
[Unity][Transform] Handle symbolic variables in LambdaLift (#16411)
* [Unity][Transform] Handle symbolic variables in LambdaLift
Prior to this commit, symbolic variables used by a lambda function
would be duplicated between the caller and the lifted-out function.
In addition, shape inference within the lifted-out function was
performed without access to the symbolic variables, resulting in
unnecessary fallback from `R.Tensor([m, n])` to `R.Tensor(ndim=2)`.
This commit updates the `LambdaLift` transform to handle symbolic
variables. All symbolic variables have unique definitions across the
resulting `IRModule`, and shape inference in the lifted-out function
is aware of symbolic variables that have been exposed to it.
* Cleanup based on review comments
---
src/relax/transform/lambda_lift.cc | 381 +++++++++++------------
tests/python/relax/test_transform_lambda_lift.py | 150 +++++++--
2 files changed, 307 insertions(+), 224 deletions(-)
diff --git a/src/relax/transform/lambda_lift.cc
b/src/relax/transform/lambda_lift.cc
index c7caeab055..16bd8bfc91 100644
--- a/src/relax/transform/lambda_lift.cc
+++ b/src/relax/transform/lambda_lift.cc
@@ -236,95 +236,24 @@ class LambdaLifter : public ExprMutator {
using ExprMutator::VisitExpr_;
- void VisitBinding_(const VarBindingNode* binding) final {
- bool is_lambda = binding->value->IsInstance<FunctionNode>();
- if (is_lambda) {
- recur_vars_.push_back(binding->var);
- }
-
- Expr new_value = this->VisitExpr(binding->value);
+ void VisitBinding_(const VarBindingNode* binding, const FunctionNode*
func_node) final {
+ auto cache = current_lambda_var_;
+ current_lambda_var_ = binding->var;
- if (new_value->struct_info_.defined() &&
- !new_value->struct_info_.same_as(binding->var->struct_info_)) {
- binding->var->struct_info_ = GetStructInfo(new_value);
- binding->var->checked_type_ = new_value->checked_type_;
- }
- if (new_value.same_as(binding->value)) {
- builder_->EmitNormalized(GetRef<VarBinding>(binding));
- } else {
- builder_->EmitNormalized(VarBinding(binding->var, new_value));
- }
- if (is_lambda) {
- recur_vars_.pop_back();
+ auto new_value = VisitExpr(binding->value);
+ if (!rebind_map_.count(binding->var)) {
+ ReEmitBinding(binding, new_value);
}
- }
-
- Expr VisitExpr_(const CallNode* call_node) final {
- auto call = Downcast<Call>(ExprMutator::VisitExpr_(call_node));
- if (const auto* var_node = call_node->op.as<VarNode>()) {
- auto var = GetRef<Var>(var_node);
- bool has_closure = HasClosure(var);
- auto val = builder_->LookupBinding(var);
- if (const auto* fsinfo_node =
GetStructInfo(var).as<FuncStructInfoNode>()) {
- auto fsinfo = GetRef<FuncStructInfo>(fsinfo_node);
- if (!GetStructInfo(call).same_as(fsinfo)) {
- call->struct_info_ = fsinfo->ret;
- call->checked_type_ = GetStaticType(fsinfo->ret);
- }
- }
- // Call "relax.invoke_closure" to invoke closure
- Var clo_arg = var;
- if (has_closure && val->IsInstance<CallNode>()) {
- if (this->var_remap_.find(var->vid) != this->var_remap_.end()) {
- clo_arg = this->var_remap_.at(var->vid);
- }
- // if the original op was pure, we should use invoke_pure_closure
- Call orig_call = Downcast<Call>(val);
- bool purity;
- if (orig_call->op.as<OpNode>()) {
- auto orig_op = Downcast<Op>(orig_call->op);
- static const auto& purity_map = Op::GetAttrMap<Bool>("FPurity");
- purity = purity_map.count(orig_op) && purity_map[orig_op]->value;
- } else {
- purity = GetStructInfoAs<FuncStructInfoNode>(orig_call->op)->purity;
- }
-
- return Call(purity ? invoke_pure_closure_op_ : invoke_closure_op_,
- {clo_arg, Tuple(call_node->args)}, {},
- {GetStructInfo(GetRef<Expr>(call_node))});
- }
- auto it = lambda_map_.find(var);
- if (it != lambda_map_.end()) {
- // flatten nested call, e.g. call(y)(x) -> call(x, y))
- Array<relay::Expr> new_args;
- Array<StructInfo> params;
- for (const auto arg : call->args) {
- new_args.push_back(arg);
- params.push_back(StructInfoFromType(arg->checked_type()));
- }
- if (const auto* nest_call = it->second.as<CallNode>()) {
- // Update the StructInfo accordingly
- for (const auto arg : nest_call->args) {
- new_args.push_back(arg);
- params.push_back(StructInfoFromType(arg->checked_type()));
- }
- StructInfo new_func_sinfo;
- if (const auto* fsinfo =
GetStructInfo(nest_call->op).as<FuncStructInfoNode>()) {
- auto func_sinfo = GetRef<FuncStructInfo>(fsinfo);
- new_func_sinfo = FuncStructInfo(params, func_sinfo->ret);
- }
- nest_call->op->struct_info_ = new_func_sinfo;
- nest_call->op->checked_type_ = GetStaticType(new_func_sinfo);
- return Call(nest_call->op, new_args, call_node->attrs,
call_node->sinfo_args);
- }
- return Call(it->second, call->args, call_node->attrs,
call_node->sinfo_args);
- }
- }
- return std::move(call);
+ current_lambda_var_ = cache;
}
Expr VisitExpr_(const FunctionNode* func_node) final {
+ if (!current_lambda_var_) {
+ // Early bail-out for top-level functions
+ return ExprMutator::VisitExpr_(func_node);
+ }
+
auto func = GetRef<Function>(func_node);
String lift_func_name = [&]() {
@@ -336,20 +265,19 @@ class LambdaLifter : public ExprMutator {
return it->second;
}();
- auto global = GlobalVar(lift_func_name);
- Array<Var> free_vars = FreeVars(func);
Array<Var> captured_vars;
-
- Array<Var> typed_captured_vars;
- bool recursive = false;
- for (const auto& var : free_vars) {
- if (!recur_vars_.empty() && var == recur_vars_.back()) {
- recursive = true;
+ bool is_recursive = false;
+ bool is_closure = false;
+ for (const auto& var : FreeVars(func)) {
+ if (var.same_as(current_lambda_var_)) {
+ is_recursive = true;
} else {
+ is_closure = true;
captured_vars.push_back(var);
}
}
+ Array<Var> typed_captured_vars;
Map<Var, Expr> rebinding_map;
for (auto free_var : captured_vars) {
Var var = Var(free_var->name_hint(), GetStructInfo(free_var),
free_var->span);
@@ -357,155 +285,207 @@ class LambdaLifter : public ExprMutator {
rebinding_map.Set(free_var, var);
}
- // recursive call
- if (recursive) {
- if (!captured_vars.empty()) {
- Array<Expr> fvs;
- for (auto fv : captured_vars) {
- fvs.push_back(fv);
- }
- // it is required by block_blocker, will be updated later
- UpdateStructInfo(global, GetStructInfo(recur_vars_.back()));
- lambda_map_.emplace(recur_vars_.back(), Call(global, fvs));
- } else {
- if (recur_vars_.size() > 0) {
- lambda_map_.emplace(recur_vars_.back(), global);
- }
- }
+ tvm::Array<Var> lifted_func_params =
+ func_node->params.Map([this](Var var) { return VisitVarDef(var); });
+ for (const auto& var : typed_captured_vars) {
+ lifted_func_params.push_back(var);
}
- tvm::Array<Var> params;
- bool all_params_unchanged = true;
- for (Var param : func_node->params) {
- Var new_param = this->VisitVarDef(param);
- params.push_back(new_param);
- all_params_unchanged &= param.same_as(new_param);
+ auto gvar_lifted_func = GlobalVar(lift_func_name);
+ {
+ auto func_sinfo = Downcast<FuncStructInfo>(func_node->struct_info_);
+ if (is_closure) {
+ func_sinfo = FuncStructInfo(lifted_func_params.Map(GetStructInfo),
func_sinfo->ret,
+ func_sinfo->purity);
+ }
+ UpdateStructInfo(gvar_lifted_func, func_sinfo);
}
- Expr body = this->VisitWithNewScope(func_node->body);
- Expr visited_func;
+ Expr body = func_node->body;
- if (all_params_unchanged && body.same_as(func_node->body)) {
- visited_func = GetRef<Expr>(func_node);
- } else if (const auto& body_sinfo =
MatchStructInfo<ObjectStructInfo>(body)) {
- visited_func =
- Function(params, body, body_sinfo.value(), func_node->is_pure,
func_node->attrs);
- } else {
- visited_func =
- Function(params, body, func_node->ret_struct_info,
func_node->is_pure, func_node->attrs);
+ // Defining the rewrite rule prior to visiting the body, so that
+ // recursive closures can be updated.
+ if (is_recursive && is_closure) {
+ nested_closure_map_.emplace(
+ current_lambda_var_.value(),
+ Call(gvar_lifted_func, captured_vars.Map([](Var var) -> Expr {
return var; })));
}
- auto new_func = Downcast<Function>(visited_func);
- Function lifted_func;
- bool is_closure = IsClosure(captured_vars);
if (!is_closure) {
- lifted_func = Function(
- /*params=*/new_func->params,
- /*body=*/new_func->body,
- /*ret_struct_info=*/new_func->ret_struct_info,
- /*is_pure=*/new_func->is_pure,
- /*attrs=*/new_func->attrs,
- /*span=*/new_func->span);
- } else {
- // Flatten the Closure
- std::vector<Var> closure_params;
- closure_params.reserve(func->params.size() + typed_captured_vars.size());
- for (size_t i = 0; i < func->params.size(); ++i) {
- closure_params.emplace_back(func->params[i]);
- }
- for (size_t i = 0; i < typed_captured_vars.size(); ++i) {
- closure_params.emplace_back(typed_captured_vars[i]);
- }
+ rebind_map_.emplace(current_lambda_var_.value(), gvar_lifted_func);
+ }
- lifted_func = Function(/*params=*/closure_params,
- /*body=*/Bind(new_func->body, rebinding_map),
- /*ret_struct_info=*/new_func->ret_struct_info,
- /*is_pure=*/new_func->is_pure,
- /*attrs=*/new_func->attrs,
- /*span=*/func->span);
+ body = this->VisitWithNewScope(body, lifted_func_params);
+ StructInfo ret_struct_info = GetStructInfo(body);
+ body = Bind(body, rebinding_map);
- for (Var param : closure_params) {
- CHECK(param->checked_type_.defined())
- << "relax.Function requires params to contain checked_type_";
- }
+ Function lifted_func;
+ if (lifted_func_params.same_as(func_node->params) &&
body.same_as(func_node->body) &&
+ ret_struct_info.same_as(func_node->ret_struct_info)) {
+ lifted_func = GetRef<Function>(func_node);
+ } else {
+ lifted_func =
+ Function(lifted_func_params, body, ret_struct_info,
func_node->is_pure, func_node->attrs);
+ }
+
+ for (Var param : lifted_func->params) {
+ CHECK(param->checked_type_.defined())
+ << "relax.Function requires all parameters to contain checked_type_.
"
+ << "However, parameter " << param << " with struct info " <<
param->struct_info_
+ << " has no checked type";
}
ICHECK(lifted_func.defined());
+ if (is_closure || IsClosure(lifted_func)) {
+ closures_.insert(gvar_lifted_func);
+ }
+
// Add the lifted function to the module.
- global->struct_info_ = GetStructInfo(lifted_func);
- global->checked_type_ = lifted_func->checked_type_;
- builder_->UpdateFunction(global, lifted_func);
+ lifted_func = CopyWithNewVars(lifted_func);
+ gvar_lifted_func->struct_info_ = GetStructInfo(lifted_func);
+ gvar_lifted_func->checked_type_ = lifted_func->checked_type_;
- if (!is_closure) {
- return std::move(global);
- } else {
+ builder_->UpdateFunction(gvar_lifted_func, lifted_func);
+
+ Expr callable_value = gvar_lifted_func;
+ if (is_closure) {
// If we need to allocate a closure,
// we pass the variables in its environment here.
- Array<Expr> fvs;
- for (auto fv : captured_vars) {
- fvs.push_back(fv);
- }
+ Tuple arg_tuple(captured_vars.Map([](Var var) -> Expr { return var; }));
// Call make_closure intrinsic
- return Call(make_closure_op_, {global, Tuple(fvs)}, {}, {});
+ callable_value = Call(make_closure_op_, {gvar_lifted_func, arg_tuple},
{}, {});
}
+
+ return callable_value;
}
- bool HasClosure(const Var& var) {
- auto val = builder_->LookupBinding(var);
- if (const auto* value = val.as<GlobalVarNode>()) {
- IRModule ctx_mod = builder_->GetContextIRModule();
- ICHECK(ctx_mod->functions.size() > 0);
- BaseFunc func = ctx_mod->Lookup(GetRef<GlobalVar>(value));
- if (const auto* func_node = func.as<FunctionNode>()) {
- if (const auto* call_node = func_node->body.as<CallNode>()) {
- if (call_node->op == make_closure_op_) {
- return true;
- }
- } else if (const auto* seq_expr_node =
func_node->body.as<SeqExprNode>()) {
- // the return var points to a make_closure intrinsic
- if (const auto* var = seq_expr_node->body.as<VarNode>()) {
- return HasClosure(GetRef<Var>(var));
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = GetRef<Call>(call_node);
+
+ auto orig_sinfo = Downcast<StructInfo>(call->struct_info_);
+
+ if (auto opt_var = call->op.as<Var>()) {
+ auto var = opt_var.value();
+
+ // Call "relax.invoke_closure" to invoke closure
+
+ if (IsClosure(var) && builder_->LookupBinding(var).as<CallNode>()) {
+ // if the original op was pure, we should use invoke_pure_closure
+ Call orig_call = Downcast<Call>(builder_->LookupBinding(var));
+ bool is_pure = [&]() -> bool {
+ if (auto op = orig_call->op.as<Op>()) {
+ static const auto& purity_map = Op::GetAttrMap<Bool>("FPurity");
+ return purity_map.get(op.value(), Bool(false))->value;
+ } else if (const auto* func_sinfo =
+ orig_call->op->struct_info_.as<FuncStructInfoNode>())
{
+ return func_sinfo->purity;
+ } else {
+ LOG(FATAL) << "Could not determine purity of call to " <<
orig_call->op
+ << ", as it is neither a tvm::Op (type = \"" <<
orig_call->op->GetTypeKey()
+ << "\"), "
+ << "nor is is annotated with FuncStructInfo (sinfo = "
+ << orig_call->op->struct_info_ << ")";
}
- }
+ }();
+
+ auto prev = call;
+ call = Call(is_pure ? invoke_pure_closure_op_ : invoke_closure_op_,
+ {var, Tuple(call->args)}, {}, {orig_sinfo});
}
- } else if (const auto* func_node = val.as<FunctionNode>()) {
- if (const auto* call_node = func_node->body.as<CallNode>()) {
- if (call_node->op == make_closure_op_) {
- return true;
+ }
+
+ if (auto opt_var = call->op.as<Var>()) {
+ auto var = opt_var.value();
+ if (auto it = nested_closure_map_.find(var); it !=
nested_closure_map_.end()) {
+ Call nested_call = it->second;
+
+ Array<relay::Expr> new_args = call->args;
+ for (const auto arg : nested_call->args) {
+ new_args.push_back(arg);
}
+
+ auto prev = call;
+ call = Call(nested_call->op, new_args, call->attrs, call->sinfo_args);
+ }
+ }
+
+ return ExprMutator::VisitExpr_(call.get());
+ }
+
+ Expr VisitExpr_(const VarNode* op) override {
+ auto var = GetRef<Var>(op);
+ if (auto it = rebind_map_.find(var); it != rebind_map_.end()) {
+ return it->second;
+ }
+ return ExprMutator::VisitExpr_(op);
+ }
+
+ bool IsClosure(Expr val) {
+ if (auto opt_var = val.as<Var>()) {
+ if (closures_.count(opt_var.value())) {
+ return true;
}
- } else if (const auto* call_node = val.as<relax::CallNode>()) {
+ if (auto bound_value = builder_->LookupBinding(opt_var.value())) {
+ val = bound_value.value();
+ }
+ }
+
+ if (const auto* call_node = val.as<relax::CallNode>()) {
// recursive call
auto op = call_node->op;
- if (make_closure_op_ == op) {
+ if (auto local_var = op.as<Var>()) {
+ return IsClosure(local_var.value());
+ } else if (auto global_var = op.as<GlobalVar>()) {
+ return IsClosure(global_var.value());
+ } else {
+ return make_closure_op_ == op;
+ }
+
+ } else if (const auto* global_var = val.as<GlobalVarNode>()) {
+ if (closures_.count(GetRef<GlobalVar>(global_var))) {
return true;
}
- if (const auto* lv = op.as<VarNode>()) {
- return HasClosure(GetRef<Var>(lv));
+ IRModule ctx_mod = builder_->GetContextIRModule();
+ ICHECK(ctx_mod->functions.size() > 0);
+ BaseFunc func = ctx_mod->Lookup(GetRef<GlobalVar>(global_var));
+ const auto* func_node = func.as<FunctionNode>();
+ if (func_node) {
+ return IsClosure(func_node->body);
+ } else {
+ return false;
}
+
+ } else if (const auto* func_node = val.as<FunctionNode>()) {
+ return IsClosure(func_node->body);
+
+ } else if (const auto* seq_node = val.as<SeqExprNode>()) {
+ return IsClosure(seq_node->body);
+
+ } else {
+ return false;
}
- return false;
}
- bool IsClosure(const Array<Var>& captured_vars) { return
captured_vars.size() > 0; }
-
IRModule Lift() {
auto glob_funcs = mod_->functions;
- for (auto pair : glob_funcs) {
- if (auto* n = pair.second.as<FunctionNode>()) {
- auto func = GetRef<Function>(n);
- func = Function(func->params, VisitExpr(func->body),
func->ret_struct_info, func->is_pure,
- func->attrs);
- builder_->UpdateFunction(pair.first, func);
+ for (auto [gvar, base_func] : glob_funcs) {
+ if (auto opt = base_func.as<Function>()) {
+ // Must visit the function itself, and not just the function
+ // body, to ensure that EraseToWellDefined recognized symbolic
+ // variables that are exposed by the function signature.
+ auto func = Downcast<Function>(VisitExpr(opt.value()));
+ builder_->UpdateFunction(gvar, func);
}
}
return builder_->GetContextIRModule();
}
private:
- std::unordered_map<Var, Expr, ObjectPtrHash, ObjectPtrEqual> lambda_map_;
- Array<Var> recur_vars_;
+ std::unordered_map<Var, Call, ObjectPtrHash, ObjectPtrEqual>
nested_closure_map_;
+ std::unordered_map<Var, Expr, ObjectPtrHash, ObjectPtrEqual> rebind_map_;
+ std::unordered_set<Variant<GlobalVar, Var>, ObjectPtrHash, ObjectPtrEqual>
closures_;
+ Optional<Var> current_lambda_var_ = NullOpt;
IRModule mod_;
std::unordered_map<const FunctionNode*, String> lifted_names_;
@@ -519,9 +499,8 @@ class LambdaLifter : public ExprMutator {
namespace transform {
Pass LambdaLift() {
- runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
- [=](IRModule m, PassContext pc) { return relax::LambdaLifter(m).Lift();
};
- return CreateModulePass(pass_func, 1, "LambdaLift", {});
+ auto pass_func = [=](IRModule mod, PassContext pc) { return
relax::LambdaLifter(mod).Lift(); };
+ return tvm::transform::CreateModulePass(pass_func, 1, "LambdaLift", {});
}
TVM_REGISTER_GLOBAL("relax.transform.LambdaLift").set_body_typed(LambdaLift);
diff --git a/tests/python/relax/test_transform_lambda_lift.py
b/tests/python/relax/test_transform_lambda_lift.py
index 8f3daa06e2..f30afdae84 100644
--- a/tests/python/relax/test_transform_lambda_lift.py
+++ b/tests/python/relax/test_transform_lambda_lift.py
@@ -42,7 +42,7 @@ def test_basic():
"""Functions can be listed from local bindings to the IRModule"""
# the target IRModule
- @tvm.script.ir_module
+ @I.ir_module
class Expected:
@R.function(private=True)
def main_inner(
@@ -55,11 +55,10 @@ def test_basic():
def main(
x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
) -> R.Tensor((10, 5), "float32"):
- inner = Expected.main_inner
- gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
+ gv1: R.Tensor((10, 5), "float32") = Expected.main_inner(x1, y1)
return gv1
- @tvm.script.ir_module
+ @I.ir_module
class Before:
@R.function
def main(
@@ -84,18 +83,54 @@ def test_basic():
_check_save_roundtrip(after)
+def test_input_module_is_unmodified():
+ """The input module may not be modified
+
+ If the output requires new StructInfo, it must create a new relax
+ variable. It must not update the struct info of an existing relax
+ variable, as that variable may be used by another IRModule.
+ """
+
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
+ ) -> R.Tensor((2, 3), "float32"):
+ @R.function
+ def outer_func(
+ c1: R.Tensor((2, 3), "float32")
+ ) -> R.Callable((R.Tensor((2, 3), "float32"),), R.Tensor((2, 3),
"float32")):
+ @R.function
+ def inner_func(x1: R.Tensor((2, 3), "float32")) ->
R.Tensor((2, 3), "float32"):
+ s: R.Tensor((2, 3), "float32") = R.add(x1, c1)
+ return s
+
+ return inner_func
+
+ in_call = outer_func(x)
+ res = in_call(y)
+ return res
+
+ before = Before
+ copy_of_before = tvm.ir.load_json(tvm.ir.save_json(before))
+
+ transform.LambdaLift()(before)
+
+ tvm.ir.assert_structural_equal(before, copy_of_before)
+
+
def test_closure():
"""Lifting functions may require producing closures"""
# the expected IRModule
- @tvm.script.ir_module
+ @I.ir_module
class Expected:
@R.function
def main(
x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
) -> R.Tensor((2, 3), "float32"):
- outer_func = Expected.main_outer_func
- in_call = outer_func(x)
+ in_call = Expected.main_outer_func(x)
res = R.invoke_pure_closure(
in_call, (y,), sinfo_args=(R.Tensor((2, 3), dtype="float32"))
)
@@ -112,7 +147,7 @@ def test_closure():
return inner_func
# IRModule to perform Lambda Lifting
- @tvm.script.ir_module
+ @I.ir_module
class Before:
@R.function
def main(
@@ -144,7 +179,7 @@ def test_recursive():
"""The lifted function may be recursively defined"""
# the expected IRModule
- @tvm.script.ir_module
+ @I.ir_module
class Expected:
@R.function(private=True)
def main_while_loop(
@@ -174,7 +209,7 @@ def test_recursive():
return gv
# the IRModule to apply lambda lifting
- @tvm.script.ir_module
+ @I.ir_module
class Before:
@R.function
def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor:
@@ -218,22 +253,20 @@ def test_multi_func():
"""
# expected IRModule
- @tvm.script.ir_module
+ @I.ir_module
class Expected:
@R.function
def glob_func_1(
x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
) -> R.Tensor(None, "float32", ndim=2):
- inner = Expected.glob_func_1_inner
- gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
+ gv1: R.Tensor((10, 5), "float32") = Expected.glob_func_1_inner(x1,
y1)
return gv1
@R.function
def glob_func_2(
x11: R.Tensor((10, 5), "float32"), y11: R.Tensor((10, 5),
"float32")
) -> R.Tensor(None, "float32", ndim=2):
- inner = Expected.glob_func_2_inner
- gv11: R.Tensor((10, 5), "float32") = inner(x11, y11)
+ gv11: R.Tensor((10, 5), "float32") =
Expected.glob_func_2_inner(x11, y11)
return gv11
@R.function(private=True)
@@ -251,7 +284,7 @@ def test_multi_func():
return s1
# the IRModule to apply lambda lifting
- @tvm.script.ir_module
+ @I.ir_module
class Before:
@R.function
def glob_func_1(
@@ -291,7 +324,7 @@ def test_multi_func():
def test_no_local_func():
- @tvm.script.ir_module
+ @I.ir_module
class Before:
@T.prim_func
def sub(
@@ -318,7 +351,7 @@ def test_no_local_func():
def test_impure_function():
- @tvm.script.ir_module
+ @I.ir_module
class Expected:
@R.function(pure=False, private=True)
def main_inner() -> R.Tuple:
@@ -327,11 +360,10 @@ def test_impure_function():
@R.function(pure=False)
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
- inner = Expected.main_inner
- gv1 = inner()
+ gv1 = Expected.main_inner()
return x
- @tvm.script.ir_module
+ @I.ir_module
class Before:
@R.function(pure=False)
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
@@ -385,8 +417,7 @@ def test_lambda_function_with_same_name_as_global():
def main(
x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
) -> R.Tensor((10, 5), "float32"):
- inner = Expected.main_inner_0
- gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
+ gv1: R.Tensor((10, 5), "float32") = Expected.main_inner_0(x1, y1)
return gv1
@R.function(private=True)
@@ -404,5 +435,78 @@ def test_lambda_function_with_same_name_as_global():
assert_structural_equal(Expected, after)
+def test_symbolic_variable_defined_by_inner_func():
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
+ ) -> R.Tensor((10, 5), "float32"):
+ @R.function
+ def inner(x2: R.Tensor(("n", "m"), "float32"), y2: R.Tensor(("n",
"m"), "float32")):
+ sum_inner = R.add(x2, y2)
+ return sum_inner
+
+ sum_main = inner(x1, y1)
+ return sum_main
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
+ ) -> R.Tensor((10, 5), "float32"):
+ sum_main = Expected.main_inner(x1, y1)
+ return sum_main
+
+ @R.function(private=True)
+ def main_inner(
+ x2: R.Tensor(("n", "m"), "float32"), y2: R.Tensor(("n", "m"),
"float32")
+ ) -> R.Tensor(("n", "m"), "float32"):
+ sum_inner = R.add(x2, y2)
+ return sum_inner
+
+ After = transform.LambdaLift()(Before)
+ assert_structural_equal(Expected, After)
+
+
+def test_symbolic_variable_defined_by_outer_func():
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x1: R.Tensor(("n", "m"), "float32"), y1: R.Tensor(("n", "m"),
"float32")
+ ) -> R.Tensor(("n", "m"), "float32"):
+ n = T.int64()
+ m = T.int64()
+
+ @R.function
+ def inner(x2: R.Tensor((n, m), "float32"), y2: R.Tensor((n, m),
"float32")):
+ sum_inner = R.add(x2, y2)
+ return sum_inner
+
+ sum_main = inner(x1, y1)
+ return sum_main
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x1: R.Tensor(("n", "m"), "float32"), y1: R.Tensor(("n", "m"),
"float32")
+ ) -> R.Tensor(("n", "m"), "float32"):
+ sum_main = Expected.main_inner(x1, y1)
+ return sum_main
+
+ @R.function(private=True)
+ def main_inner(
+ x2: R.Tensor(("n", "m"), "float32"), y2: R.Tensor(("n", "m"),
"float32")
+ ) -> R.Tensor(("n", "m"), "float32"):
+ sum_inner = R.add(x2, y2)
+ return sum_inner
+
+ After = transform.LambdaLift()(Before)
+ assert_structural_equal(Expected, After)
+
+
if __name__ == "__main__":
tvm.testing.main()
