From 28071fb6e4f0c81201bf389d19e69ff3f6b045a5 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 16 Jan 2024 15:44:30 +0000 Subject: [PATCH 1/2] [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. --- src/relax/transform/lambda_lift.cc | 390 +++++++++--------- .../relax/test_transform_lambda_lift.py | 150 +++++-- 2 files changed, 316 insertions(+), 224 deletions(-) diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc index c7caeab05596..37c8d3f3168d 100644 --- a/src/relax/transform/lambda_lift.cc +++ b/src/relax/transform/lambda_lift.cc @@ -236,95 +236,25 @@ class LambdaLifter : public ExprMutator { using ExprMutator::VisitExpr_; - void VisitBinding_(const VarBindingNode* binding) final { - bool is_lambda = binding->value->IsInstance(); - if (is_lambda) { - recur_vars_.push_back(binding->var); + void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func_node) final { + auto cache = current_lambda_var_; + current_lambda_var_ = binding->var; + + // ExprMutator::VisitBinding_(binding, func_node); + auto new_value = VisitExpr(binding->value); + if (!rebind_map_.count(binding->var)) { + ReEmitBinding(binding, new_value); } - Expr new_value = this->VisitExpr(binding->value); - - 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(binding)); - } else { - builder_->EmitNormalized(VarBinding(binding->var, new_value)); - } - if (is_lambda) { - recur_vars_.pop_back(); - } + current_lambda_var_ = cache; } - Expr VisitExpr_(const CallNode* call_node) final { - auto call = Downcast(ExprMutator::VisitExpr_(call_node)); - if (const auto* var_node = call_node->op.as()) { - auto var = GetRef(var_node); - bool has_closure = HasClosure(var); - auto val = builder_->LookupBinding(var); - if (const auto* fsinfo_node = GetStructInfo(var).as()) { - auto fsinfo = GetRef(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()) { - 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(val); - bool purity; - if (orig_call->op.as()) { - auto orig_op = Downcast(orig_call->op); - static const auto& purity_map = Op::GetAttrMap("FPurity"); - purity = purity_map.count(orig_op) && purity_map[orig_op]->value; - } else { - purity = GetStructInfoAs(orig_call->op)->purity; - } - - return Call(purity ? invoke_pure_closure_op_ : invoke_closure_op_, - {clo_arg, Tuple(call_node->args)}, {}, - {GetStructInfo(GetRef(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 new_args; - Array 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()) { - // 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()) { - auto func_sinfo = GetRef(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); - } + Expr VisitExpr_(const FunctionNode* func_node) final { + if (!current_lambda_var_) { + // Early bail-out for top-level functions + return ExprMutator::VisitExpr_(func_node); } - return std::move(call); - } - Expr VisitExpr_(const FunctionNode* func_node) final { auto func = GetRef(func_node); String lift_func_name = [&]() { @@ -336,20 +266,19 @@ class LambdaLifter : public ExprMutator { return it->second; }(); - auto global = GlobalVar(lift_func_name); - Array free_vars = FreeVars(func); Array captured_vars; - - Array 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 typed_captured_vars; Map rebinding_map; for (auto free_var : captured_vars) { Var var = Var(free_var->name_hint(), GetStructInfo(free_var), free_var->span); @@ -357,155 +286,215 @@ class LambdaLifter : public ExprMutator { rebinding_map.Set(free_var, var); } - // recursive call - if (recursive) { - if (!captured_vars.empty()) { - Array 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 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 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(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(func_node); - } else if (const auto& body_sinfo = MatchStructInfo(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); + // recursive call + if (is_recursive && is_closure) { + // it is required by block_blocker, will be updated later + nested_closure_map_.emplace( + current_lambda_var_.value(), + Call(gvar_lifted_func, captured_vars.Map([](Var var) -> Expr { return var; }))); } - auto new_func = Downcast(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 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(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 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()) { - IRModule ctx_mod = builder_->GetContextIRModule(); - ICHECK(ctx_mod->functions.size() > 0); - BaseFunc func = ctx_mod->Lookup(GetRef(value)); - if (const auto* func_node = func.as()) { - if (const auto* call_node = func_node->body.as()) { - if (call_node->op == make_closure_op_) { - return true; - } - } else if (const auto* seq_expr_node = func_node->body.as()) { - // the return var points to a make_closure intrinsic - if (const auto* var = seq_expr_node->body.as()) { - return HasClosure(GetRef(var)); + Expr VisitExpr_(const CallNode* call_node) final { + auto call = GetRef(call_node); + + auto orig_sinfo = Downcast(call->struct_info_); + + if (auto opt_var = call->op.as()) { + auto var = opt_var.value(); + + // Call "relax.invoke_closure" to invoke closure + + if (bool is_closure = IsClosure(var); + is_closure && builder_->LookupBinding(var).as()) { + // if the original op was pure, we should use invoke_pure_closure + Call orig_call = Downcast(builder_->LookupBinding(var)); + bool is_pure = [&]() -> bool { + if (auto op = orig_call->op.as()) { + static const auto& purity_map = Op::GetAttrMap("FPurity"); + return purity_map.get(op.value(), Bool(false))->value; + } else if (const auto* func_sinfo = + orig_call->op->struct_info_.as()) { + 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()) { - if (const auto* call_node = func_node->body.as()) { - if (call_node->op == make_closure_op_) { - return true; + } + + if (auto opt_var = call->op.as()) { + auto var = opt_var.value(); + if (auto it = nested_closure_map_.find(var); it != nested_closure_map_.end()) { + Call nested_call = it->second; + + Array 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); } - } else if (const auto* call_node = val.as()) { + } + + return ExprMutator::VisitExpr_(call.get()); + } + + Expr VisitExpr_(const VarNode* op) override { + auto var = GetRef(op); + if (auto it = rebind_map_.find(var); it != rebind_map_.end()) { + return it->second; + } + return ExprMutator::VisitExpr_(op); + } + + bool IsClosure(Expr val) { + static int depth = -1; + struct Context { + explicit Context(int* ptr) : ptr(ptr) { (*ptr)++; } + ~Context() { (*ptr)--; } + int* ptr; + } context(&depth); + + if (auto opt_var = val.as()) { + if (closures_.count(opt_var.value())) { + return true; + } + if (auto bound_value = builder_->LookupBinding(opt_var.value())) { + val = bound_value.value(); + } + } + + if (const auto* call_node = val.as()) { // recursive call auto op = call_node->op; - if (make_closure_op_ == op) { + if (auto local_var = op.as()) { + return IsClosure(local_var.value()); + } else if (auto global_var = op.as()) { + return IsClosure(global_var.value()); + } else { + return make_closure_op_ == op; + } + + } else if (const auto* global_var = val.as()) { + if (closures_.count(GetRef(global_var))) { return true; } - if (const auto* lv = op.as()) { - return HasClosure(GetRef(lv)); + IRModule ctx_mod = builder_->GetContextIRModule(); + ICHECK(ctx_mod->functions.size() > 0); + BaseFunc func = ctx_mod->Lookup(GetRef(global_var)); + const auto* func_node = func.as(); + if (func_node) { + return IsClosure(func_node->body); + } else { + return false; } + + } else if (const auto* func_node = val.as()) { + return IsClosure(func_node->body); + + } else if (const auto* seq_node = val.as()) { + return IsClosure(seq_node->body); + + } else { + return false; } - return false; } - bool IsClosure(const Array& 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()) { - auto func = GetRef(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()) { + // 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(VisitExpr(opt.value())); + builder_->UpdateFunction(gvar, func); } } return builder_->GetContextIRModule(); } private: - std::unordered_map lambda_map_; - Array recur_vars_; + std::unordered_map nested_closure_map_; + std::unordered_map rebind_map_; + std::unordered_set, ObjectPtrHash, ObjectPtrEqual> closures_; + Optional current_lambda_var_ = NullOpt; IRModule mod_; std::unordered_map lifted_names_; @@ -519,9 +508,8 @@ class LambdaLifter : public ExprMutator { namespace transform { Pass LambdaLift() { - runtime::TypedPackedFunc 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 8f3daa06e200..f30afdae849c 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 main_inner( 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 inner( _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 main_outer_func(y: R.Tensor((2, 3), "float32")) -> R.Object: 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 main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), dtype="float32"): 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 glob_func_2_inner( 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 inner( def test_no_local_func(): - @tvm.script.ir_module + @I.ir_module class Before: @T.prim_func def sub( @@ -318,7 +351,7 @@ def before(c0: R.Tensor((16, 16), "float32"), x: R.Tensor(dtype="float32", ndim= 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 main_inner() -> R.Tuple: @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 @@ class Expected: 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 main_inner(): 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() From 43b7f8cdf59eae6bd6be94ebaa8dcffafa70d3e4 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 22 Jan 2024 19:46:49 +0000 Subject: [PATCH 2/2] Cleanup based on review comments --- src/relax/transform/lambda_lift.cc | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc index 37c8d3f3168d..16bd8bfc9110 100644 --- a/src/relax/transform/lambda_lift.cc +++ b/src/relax/transform/lambda_lift.cc @@ -240,7 +240,6 @@ class LambdaLifter : public ExprMutator { auto cache = current_lambda_var_; current_lambda_var_ = binding->var; - // ExprMutator::VisitBinding_(binding, func_node); auto new_value = VisitExpr(binding->value); if (!rebind_map_.count(binding->var)) { ReEmitBinding(binding, new_value); @@ -304,9 +303,9 @@ class LambdaLifter : public ExprMutator { Expr body = func_node->body; - // recursive call + // Defining the rewrite rule prior to visiting the body, so that + // recursive closures can be updated. if (is_recursive && is_closure) { - // it is required by block_blocker, will be updated later nested_closure_map_.emplace( current_lambda_var_.value(), Call(gvar_lifted_func, captured_vars.Map([](Var var) -> Expr { return var; }))); @@ -371,8 +370,7 @@ class LambdaLifter : public ExprMutator { // Call "relax.invoke_closure" to invoke closure - if (bool is_closure = IsClosure(var); - is_closure && builder_->LookupBinding(var).as()) { + if (IsClosure(var) && builder_->LookupBinding(var).as()) { // if the original op was pure, we should use invoke_pure_closure Call orig_call = Downcast(builder_->LookupBinding(var)); bool is_pure = [&]() -> bool { @@ -424,13 +422,6 @@ class LambdaLifter : public ExprMutator { } bool IsClosure(Expr val) { - static int depth = -1; - struct Context { - explicit Context(int* ptr) : ptr(ptr) { (*ptr)++; } - ~Context() { (*ptr)--; } - int* ptr; - } context(&depth); - if (auto opt_var = val.as()) { if (closures_.count(opt_var.value())) { return true;