src/relax/transform/lambda_lift.cc - tvm - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 /*!
  * \file tvm/relax/transform/lambda_lift.cc
  * \brief Lift local functions into global functions.
  */

 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/error.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/expr.h>
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/transform.h>

 #include <iostream>
 #include <vector>

 namespace tvm {
 namespace relax {

 namespace {

 /* \brief Collect names of functions to be lifted out */
 class LambdaNameCollector : ExprVisitor {
  public:
   static std::unordered_map<const FunctionNode*, ffi::String> Collect(const IRModule& mod) {
     LambdaNameCollector visitor;

     for (const auto& [gvar, base_func] : mod->functions) {
       visitor.previous_global_vars_.insert(gvar->name_hint);
     }

     for (const auto& [gvar, base_func] : mod->functions) {
       if (auto func = base_func.as<Function>()) {
         visitor.name_stack_.push_back(gvar->name_hint);
         visitor(func.value());
         visitor.name_stack_.pop_back();
       }
     }

     return visitor.Finalize();
   }

  private:
   void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func) override {
     if (auto opt = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
       ffi::String public_name = opt.value();

       // If a kGlobalSymbol exists, we must use the name exactly as it
       // appears, with no modifications.  Because these errors would
       // be raised from deep within an optimization pipeline, but
       // depends on small annotation changes from a user's initial
       // model definition, they are intentionally verbose to
       // (hopefully) provide sufficient context to a user encountering
       // the error.
       TVM_FFI_ICHECK(!previous_global_vars_.count(public_name))
           << "Function " << name_stack_.front() << " contains a lambda with kGlobalSymbol (\""
           << tvm::attr::kGlobalSymbol << "\" attribute of \"" << public_name << "\".  "
           << "However, the module already contains a GlobalVar with this name.  "
           << "If present, the kGlobalSymbol attribute must match the name of the GlobalVar, "
           << "and GlobalVar names must be unique across an IRModule.  "
           << "Lifting the " << public_name << " function out of " << name_stack_.front()
           << " would require violating one of these two conditions.";

       auto it = new_public_names_.find(public_name);
       TVM_FFI_ICHECK(it == new_public_names_.end())
           << "Function " << name_stack_.front() << " contains a lambda with kGlobalSymbol (\""
           << tvm::attr::kGlobalSymbol << "\" attribute of \"" << public_name << "\".  "
           << "However, the function " << it->second.front()
           << " also contains a lambda with the same value for kGlobalSymbol.  "
           << "If present, the kGlobalSymbol attribute must match the name of the GlobalVar, "
           << "and GlobalVar names must be unique across an IRModule.  "
           << "Lifting the " << public_name << " function out of both " << name_stack_.front()
           << " and " << it->second.front()
           << " would require violating one of these two conditions.";

       new_public_names_.insert({public_name, name_stack_});
       lifted_with_global_symbol_.insert({func, public_name});
     }

     name_stack_.push_back(binding->var->name);
     lambda_location_.insert({func, name_stack_});
     ExprVisitor::VisitBinding_(binding, func);
     name_stack_.pop_back();
   }

   // De-duplication of collected names
   std::unordered_map<const FunctionNode*, ffi::String> Finalize() const {
     // The functions which still must be assigned a name
     std::unordered_map<const FunctionNode*, ffi::Array<ffi::String>> remaining_to_name =
         lambda_location_;

     // Collecting the functions that now have a name.
     std::unordered_map<const FunctionNode*, ffi::String> lifted_names;

     // A lookup for names that are unavailable for use.
     std::unordered_set<ffi::String> unavailable_names = previous_global_vars_;

     // A helper function to generate de-duplicated names.  The
     // `proposed_name_generation_func` should be a function with
     // signature:
     //
     //     ffi::Optional<ffi::String> func(const FunctionNode*, const ffi::Array<ffi::String>&)
     //
     // The first argument will be the lambda function being lifted.
     // The second argument will be the nested location where that
     // lambda function was found.  The function should return the
     // proposed name for the lifted lambda function.  The proposed
     // name will be accepted if it does not conflict with any previous
     // names, and is unique for all lambda functions being lifted.
     //
     // This helper function is used to apply several different schemes
     // to generate the name of the lifted lambda function.  The
     // overall goal is to provide names that are unique (required by
     // IRModule), deterministic (required for unit testing), and
     // human-readable.
     auto attempt_name_generation = [&](const auto& proposed_name_generation_func) {
       if (remaining_to_name.empty()) {
         return;
       }

       std::unordered_map<ffi::String, const FunctionNode*> new_names;
       for (const auto& [func, location] : remaining_to_name) {
         if (ffi::Optional<ffi::String> opt_proposed_name =
                 proposed_name_generation_func(func, location)) {
           auto proposed_name = opt_proposed_name.value();

           if (unavailable_names.count(proposed_name)) {
             // The name is already used, either from a GlobalVar, or
             // from a previous round of attempted names.
           } else if (auto it = new_names.find(proposed_name); it != new_names.end()) {
             // The name is not unique within the current attempt.  Mark
             // the function as nullptr to previous any use of this name
             it->second = nullptr;
           } else {
             // The name is unique so far.  Track it for use.
             new_names.insert({proposed_name, func});
           }
         }
       }

       for (const auto& [name, func] : new_names) {
         if (func) {
           lifted_names.insert({func, name});
           remaining_to_name.erase(func);
         }
       }
     };

     // 1. Start with any publicly explosed names from kGlobalSymbol
     attempt_name_generation([&](const FunctionNode* func,
                                 const auto&) -> ffi::Optional<ffi::String> {
       if (auto it = lifted_with_global_symbol_.find(func); it != lifted_with_global_symbol_.end()) {
         return it->second;
       } else {
         return std::nullopt;
       }
     });

     // 2. Try concatenating the name of the relax variable with the
     // name of the function that contains it.
     attempt_name_generation([&](const FunctionNode*, const auto& location) -> ffi::String {
       std::stringstream stream;
       stream << location.front() << "_" << location.back();
       return stream.str();
     });

     // 3. Try concatenating the entire path together.  Don't include
     // paths of length 2, as they would already be attempted earlier.
     attempt_name_generation(
         [&](const FunctionNode*, const auto& location) -> ffi::Optional<ffi::String> {
           if (location.size() == 2) return std::nullopt;

           std::stringstream stream;
           bool is_first = true;
           for (const auto& loc : location) {
             if (is_first) {
               is_first = false;
             } else {
               stream << "_";
             }
             stream << loc;
           }
           return ffi::String(stream.str());
         });

     // 4. Fallback.  Count the number of times a relax variable with
     // that name was used.
     std::unordered_map<ffi::String, int> usage_count;
     attempt_name_generation([&](const FunctionNode*, const auto& location) -> ffi::String {
       std::stringstream stream;
       stream << location.front() << "_" << location.back();
       int usage = usage_count[stream.str()]++;
       stream << "_" << usage;

       return stream.str();
     });

     TVM_FFI_ICHECK(remaining_to_name.empty())
         << "Fallback failed to make unique names for all lifted lambda functions";

     return lifted_names;
   }

   ffi::Array<ffi::String> name_stack_;
   std::unordered_set<ffi::String> previous_global_vars_;
   std::unordered_map<ffi::String, ffi::Array<ffi::String>> new_public_names_;
   std::unordered_map<const FunctionNode*, ffi::String> lifted_with_global_symbol_;
   std::unordered_map<const FunctionNode*, ffi::Array<ffi::String>> lambda_location_;
 };

 }  // namespace

 /* The goal of this class is to lift out any nested functions into top-level
  * functions.
  *
  * We will lift a function out into a global which takes the set of the free
  * vars and then return the new created function.
  */
 class LambdaLifter : public ExprMutator {
  public:
   explicit LambdaLifter(const IRModule& module)
       : ExprMutator(module), mod_(module), lifted_names_(LambdaNameCollector::Collect(module)) {}

   using ExprMutator::VisitExpr_;

   void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func_node) final {
     auto cache = current_lambda_var_;
     current_lambda_var_ = binding->var;

     auto new_value = VisitExpr(binding->value);
     if (!rebind_map_.count(binding->var)) {
       ReEmitBinding(binding, new_value);
     }

     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 = ffi::GetRef<Function>(func_node);

     ffi::String lift_func_name = [&]() {
       auto it = lifted_names_.find(func_node);
       TVM_FFI_CHECK(it != lifted_names_.end(), InternalError)
           << "Found lambda function during mutation step, "
           << "but it wasn't found during the earlier name-generation step.";
       return it->second;
     }();

     ffi::Array<Var> captured_vars;
     bool is_recursive = false;
     bool is_closure = false;
     for (const auto& var : FreeVars(func)) {
       if (current_lambda_var_.same_as(var)) {
         is_recursive = true;
       } else {
         is_closure = true;
         captured_vars.push_back(var);
       }
     }

     ffi::Array<Var> typed_captured_vars;
     ffi::Map<Var, Expr> rebinding_map;
     for (auto free_var : captured_vars) {
       Var var = Var(free_var->name, GetType(free_var), free_var->span);
       typed_captured_vars.push_back(var);
       rebinding_map.Set(free_var, var);
     }

     tvm::ffi::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);
     }

     auto gvar_lifted_func = GlobalVar(lift_func_name);
     {
       auto func_ty = func_node->ty.as_or_throw<FuncType>();
       if (is_closure) {
         func_ty = FuncType(lifted_func_params.Map(GetType), func_ty->ret, func_ty->purity);
       }
       UpdateType(gvar_lifted_func, func_ty);
     }

     Expr body = func_node->body;

     // 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(Type::Missing(), gvar_lifted_func,
                                        captured_vars.Map([](Var var) -> Expr { return var; })));
     }

     if (!is_closure) {
       rebind_map_.emplace(current_lambda_var_.value(), gvar_lifted_func);
     }

     body = this->VisitWithNewScope(body, lifted_func_params);
     Type ret_ty = GetType(body);
     body = Bind(body, rebinding_map);

     Function lifted_func;
     if (lifted_func_params.same_as(func_node->params) && body.same_as(func_node->body) &&
         ret_ty.same_as(func_node->ret_ty)) {
       lifted_func = ffi::GetRef<Function>(func_node);
     } else {
       lifted_func =
           Function(lifted_func_params, body, ret_ty, func_node->is_pure, func_node->attrs);
     }

     TVM_FFI_ICHECK(lifted_func.defined());

     if (is_closure || IsClosure(lifted_func)) {
       closures_.insert(gvar_lifted_func);
     }

     // Add the lifted function to the module.
     lifted_func = CopyWithNewVars(lifted_func);
     gvar_lifted_func->ty = GetType(lifted_func);

     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.
       Tuple arg_tuple(captured_vars.Map([](Var var) -> Expr { return var; }));
       // Call make_closure intrinsic
       callable_value =
           Call(Type::Missing(), make_closure_op_, {gvar_lifted_func, arg_tuple}, {}, {});
     }

     return callable_value;
   }

   Expr VisitExpr_(const CallNode* call_node) final {
     auto call = ffi::GetRef<Call>(call_node);

     auto orig_ty = call->ty.as_or_throw<Type>();

     if (auto opt_var = call->op.as<Var>()) {
       auto var = opt_var.value();

       // Call "relax.invoke_closure" to invoke closure

       auto bound_value = LookupBinding(var);
       if (IsClosure(var) && bound_value.as<CallNode>()) {
         // if the original op was pure, we should use invoke_pure_closure
         Call orig_call = bound_value.value().as_or_throw<Call>();
         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(), false);
           } else if (const auto* func_ty = orig_call->op->ty.as<FuncTypeNode>()) {
             return func_ty->purity;
           } else {
             TVM_FFI_THROW(InternalError)
                 << "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 FuncType (ty = " << orig_call->op->ty << ")";
           }
         }();

         auto prev = call;
         call = Call(Type::Missing(), is_pure ? invoke_pure_closure_op_ : invoke_closure_op_,
                     {var, Tuple(call->args)}, {}, {orig_ty});
       }
     }

     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;

         ffi::Array<Expr> new_args = call->args;
         for (const auto arg : nested_call->args) {
           new_args.push_back(arg);
         }

         auto prev = call;
         call = Call(Type::Missing(), nested_call->op, new_args, call->attrs, call->ty_args);
       }
     }

     return ExprMutator::VisitExpr_(call.get());
   }

   Expr VisitExpr_(const VarNode* op) override {
     auto var = ffi::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;
       }
       if (auto bound_value = LookupBinding(opt_var.value())) {
         val = bound_value.value();
       }
     }

     if (const auto* call_node = val.as<tvm::CallNode>()) {
       // recursive call
       auto op = call_node->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_.same_as(op);
       }

     } else if (const auto* global_var = val.as<GlobalVarNode>()) {
       if (closures_.count(ffi::GetRef<GlobalVar>(global_var))) {
         return true;
       }
       IRModule ctx_mod = builder_->GetContextIRModule();
       TVM_FFI_ICHECK(ctx_mod->functions.size() > 0);
       BaseFunc func = ctx_mod->Lookup(ffi::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;
     }
   }

   IRModule Lift() {
     auto glob_funcs = mod_->functions;
     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 = VisitExpr(opt.value()).as_or_throw<Function>();
         builder_->UpdateFunction(gvar, func);
       }
     }
     return builder_->GetContextIRModule();
   }

  private:
   std::unordered_map<Var, Call> nested_closure_map_;
   std::unordered_map<Var, Expr> rebind_map_;
   std::unordered_set<ffi::Variant<GlobalVar, Var>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
       closures_;
   ffi::Optional<Var> current_lambda_var_ = std::nullopt;
   IRModule mod_;

   std::unordered_map<const FunctionNode*, ffi::String> lifted_names_;

   /*! \brief Cache ops that would be used later to reduce lookup overhead. */
   const Op& make_closure_op_ = Op::Get("relax.make_closure");
   const Op& invoke_closure_op_ = Op::Get("relax.invoke_closure");
   const Op& invoke_pure_closure_op_ = Op::Get("relax.invoke_pure_closure");
 };

 namespace transform {

 Pass LambdaLift() {
   auto pass_func = [=](IRModule mod, PassContext pc) { return relax::LambdaLifter(mod).Lift(); };
   return tvm::transform::CreateModulePass(pass_func, 1, "LambdaLift", {});
 }

 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def("relax.transform.LambdaLift", LambdaLift);
 }

 }  // namespace transform
 }  // namespace relax
 }  // namespace tvm
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	/*!
	* \file tvm/relax/transform/lambda_lift.cc
	* \brief Lift local functions into global functions.
	*/

	#include <tvm/ffi/cast.h>
	#include <tvm/ffi/error.h>
	#include <tvm/ffi/reflection/registry.h>
	#include <tvm/relax/analysis.h>
	#include <tvm/relax/expr.h>
	#include <tvm/relax/expr_functor.h>
	#include <tvm/relax/transform.h>

	#include <iostream>
	#include <vector>

	namespace tvm {
	namespace relax {

	namespace {

	/* \brief Collect names of functions to be lifted out */
	class LambdaNameCollector : ExprVisitor {
	public:
	static std::unordered_map<const FunctionNode*, ffi::String> Collect(const IRModule& mod) {
	LambdaNameCollector visitor;

	for (const auto& [gvar, base_func] : mod->functions) {
	visitor.previous_global_vars_.insert(gvar->name_hint);
	}

	for (const auto& [gvar, base_func] : mod->functions) {
	if (auto func = base_func.as<Function>()) {
	visitor.name_stack_.push_back(gvar->name_hint);
	visitor(func.value());
	visitor.name_stack_.pop_back();
	}
	}

	return visitor.Finalize();
	}

	private:
	void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func) override {
	if (auto opt = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
	ffi::String public_name = opt.value();

	// If a kGlobalSymbol exists, we must use the name exactly as it
	// appears, with no modifications. Because these errors would
	// be raised from deep within an optimization pipeline, but
	// depends on small annotation changes from a user's initial
	// model definition, they are intentionally verbose to
	// (hopefully) provide sufficient context to a user encountering
	// the error.
	TVM_FFI_ICHECK(!previous_global_vars_.count(public_name))
	<< "Function " << name_stack_.front() << " contains a lambda with kGlobalSymbol (\""
	<< tvm::attr::kGlobalSymbol << "\" attribute of \"" << public_name << "\". "
	<< "However, the module already contains a GlobalVar with this name. "
	<< "If present, the kGlobalSymbol attribute must match the name of the GlobalVar, "
	<< "and GlobalVar names must be unique across an IRModule. "
	<< "Lifting the " << public_name << " function out of " << name_stack_.front()
	<< " would require violating one of these two conditions.";

	auto it = new_public_names_.find(public_name);
	TVM_FFI_ICHECK(it == new_public_names_.end())
	<< "Function " << name_stack_.front() << " contains a lambda with kGlobalSymbol (\""
	<< tvm::attr::kGlobalSymbol << "\" attribute of \"" << public_name << "\". "
	<< "However, the function " << it->second.front()
	<< " also contains a lambda with the same value for kGlobalSymbol. "
	<< "If present, the kGlobalSymbol attribute must match the name of the GlobalVar, "
	<< "and GlobalVar names must be unique across an IRModule. "
	<< "Lifting the " << public_name << " function out of both " << name_stack_.front()
	<< " and " << it->second.front()
	<< " would require violating one of these two conditions.";

	new_public_names_.insert({public_name, name_stack_});
	lifted_with_global_symbol_.insert({func, public_name});
	}

	name_stack_.push_back(binding->var->name);
	lambda_location_.insert({func, name_stack_});
	ExprVisitor::VisitBinding_(binding, func);
	name_stack_.pop_back();
	}

	// De-duplication of collected names
	std::unordered_map<const FunctionNode*, ffi::String> Finalize() const {
	// The functions which still must be assigned a name
	std::unordered_map<const FunctionNode*, ffi::Array<ffi::String>> remaining_to_name =
	lambda_location_;

	// Collecting the functions that now have a name.
	std::unordered_map<const FunctionNode*, ffi::String> lifted_names;

	// A lookup for names that are unavailable for use.
	std::unordered_set<ffi::String> unavailable_names = previous_global_vars_;

	// A helper function to generate de-duplicated names. The
	// `proposed_name_generation_func` should be a function with
	// signature:
	//
	// ffi::Optional<ffi::String> func(const FunctionNode*, const ffi::Array<ffi::String>&)
	//
	// The first argument will be the lambda function being lifted.
	// The second argument will be the nested location where that
	// lambda function was found. The function should return the
	// proposed name for the lifted lambda function. The proposed
	// name will be accepted if it does not conflict with any previous
	// names, and is unique for all lambda functions being lifted.
	//
	// This helper function is used to apply several different schemes
	// to generate the name of the lifted lambda function. The
	// overall goal is to provide names that are unique (required by
	// IRModule), deterministic (required for unit testing), and
	// human-readable.
	auto attempt_name_generation = [&](const auto& proposed_name_generation_func) {
	if (remaining_to_name.empty()) {
	return;
	}

	std::unordered_map<ffi::String, const FunctionNode*> new_names;
	for (const auto& [func, location] : remaining_to_name) {
	if (ffi::Optional<ffi::String> opt_proposed_name =
	proposed_name_generation_func(func, location)) {
	auto proposed_name = opt_proposed_name.value();

	if (unavailable_names.count(proposed_name)) {
	// The name is already used, either from a GlobalVar, or
	// from a previous round of attempted names.
	} else if (auto it = new_names.find(proposed_name); it != new_names.end()) {
	// The name is not unique within the current attempt. Mark
	// the function as nullptr to previous any use of this name
	it->second = nullptr;
	} else {
	// The name is unique so far. Track it for use.
	new_names.insert({proposed_name, func});
	}
	}
	}

	for (const auto& [name, func] : new_names) {
	if (func) {
	lifted_names.insert({func, name});
	remaining_to_name.erase(func);
	}
	}
	};

	// 1. Start with any publicly explosed names from kGlobalSymbol
	attempt_name_generation([&](const FunctionNode* func,
	const auto&) -> ffi::Optional<ffi::String> {
	if (auto it = lifted_with_global_symbol_.find(func); it != lifted_with_global_symbol_.end()) {
	return it->second;
	} else {
	return std::nullopt;
	}
	});

	// 2. Try concatenating the name of the relax variable with the
	// name of the function that contains it.
	attempt_name_generation([&](const FunctionNode*, const auto& location) -> ffi::String {
	std::stringstream stream;
	stream << location.front() << "_" << location.back();
	return stream.str();
	});

	// 3. Try concatenating the entire path together. Don't include
	// paths of length 2, as they would already be attempted earlier.
	attempt_name_generation(
	[&](const FunctionNode*, const auto& location) -> ffi::Optional<ffi::String> {
	if (location.size() == 2) return std::nullopt;

	std::stringstream stream;
	bool is_first = true;
	for (const auto& loc : location) {
	if (is_first) {
	is_first = false;
	} else {
	stream << "_";
	}
	stream << loc;
	}
	return ffi::String(stream.str());
	});

	// 4. Fallback. Count the number of times a relax variable with
	// that name was used.
	std::unordered_map<ffi::String, int> usage_count;
	attempt_name_generation([&](const FunctionNode*, const auto& location) -> ffi::String {
	std::stringstream stream;
	stream << location.front() << "_" << location.back();
	int usage = usage_count[stream.str()]++;
	stream << "_" << usage;

	return stream.str();
	});

	TVM_FFI_ICHECK(remaining_to_name.empty())
	<< "Fallback failed to make unique names for all lifted lambda functions";

	return lifted_names;
	}

	ffi::Array<ffi::String> name_stack_;
	std::unordered_set<ffi::String> previous_global_vars_;
	std::unordered_map<ffi::String, ffi::Array<ffi::String>> new_public_names_;
	std::unordered_map<const FunctionNode*, ffi::String> lifted_with_global_symbol_;
	std::unordered_map<const FunctionNode*, ffi::Array<ffi::String>> lambda_location_;
	};

	} // namespace

	/* The goal of this class is to lift out any nested functions into top-level
	* functions.
	*
	* We will lift a function out into a global which takes the set of the free
	* vars and then return the new created function.
	*/
	class LambdaLifter : public ExprMutator {
	public:
	explicit LambdaLifter(const IRModule& module)
	: ExprMutator(module), mod_(module), lifted_names_(LambdaNameCollector::Collect(module)) {}

	using ExprMutator::VisitExpr_;

	void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func_node) final {
	auto cache = current_lambda_var_;
	current_lambda_var_ = binding->var;

	auto new_value = VisitExpr(binding->value);
	if (!rebind_map_.count(binding->var)) {
	ReEmitBinding(binding, new_value);
	}

	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 = ffi::GetRef<Function>(func_node);

	ffi::String lift_func_name = [&]() {
	auto it = lifted_names_.find(func_node);
	TVM_FFI_CHECK(it != lifted_names_.end(), InternalError)
	<< "Found lambda function during mutation step, "
	<< "but it wasn't found during the earlier name-generation step.";
	return it->second;
	}();

	ffi::Array<Var> captured_vars;
	bool is_recursive = false;
	bool is_closure = false;
	for (const auto& var : FreeVars(func)) {
	if (current_lambda_var_.same_as(var)) {
	is_recursive = true;
	} else {
	is_closure = true;
	captured_vars.push_back(var);
	}
	}

	ffi::Array<Var> typed_captured_vars;
	ffi::Map<Var, Expr> rebinding_map;
	for (auto free_var : captured_vars) {
	Var var = Var(free_var->name, GetType(free_var), free_var->span);
	typed_captured_vars.push_back(var);
	rebinding_map.Set(free_var, var);
	}

	tvm::ffi::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);
	}

	auto gvar_lifted_func = GlobalVar(lift_func_name);
	{
	auto func_ty = func_node->ty.as_or_throw<FuncType>();
	if (is_closure) {
	func_ty = FuncType(lifted_func_params.Map(GetType), func_ty->ret, func_ty->purity);
	}
	UpdateType(gvar_lifted_func, func_ty);
	}

	Expr body = func_node->body;

	// 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(Type::Missing(), gvar_lifted_func,
	captured_vars.Map([](Var var) -> Expr { return var; })));
	}

	if (!is_closure) {
	rebind_map_.emplace(current_lambda_var_.value(), gvar_lifted_func);
	}

	body = this->VisitWithNewScope(body, lifted_func_params);
	Type ret_ty = GetType(body);
	body = Bind(body, rebinding_map);

	Function lifted_func;
	if (lifted_func_params.same_as(func_node->params) && body.same_as(func_node->body) &&
	ret_ty.same_as(func_node->ret_ty)) {
	lifted_func = ffi::GetRef<Function>(func_node);
	} else {
	lifted_func =
	Function(lifted_func_params, body, ret_ty, func_node->is_pure, func_node->attrs);
	}

	TVM_FFI_ICHECK(lifted_func.defined());

	if (is_closure \|\| IsClosure(lifted_func)) {
	closures_.insert(gvar_lifted_func);
	}

	// Add the lifted function to the module.
	lifted_func = CopyWithNewVars(lifted_func);
	gvar_lifted_func->ty = GetType(lifted_func);

	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.
	Tuple arg_tuple(captured_vars.Map([](Var var) -> Expr { return var; }));
	// Call make_closure intrinsic
	callable_value =
	Call(Type::Missing(), make_closure_op_, {gvar_lifted_func, arg_tuple}, {}, {});
	}

	return callable_value;
	}

	Expr VisitExpr_(const CallNode* call_node) final {
	auto call = ffi::GetRef<Call>(call_node);

	auto orig_ty = call->ty.as_or_throw<Type>();

	if (auto opt_var = call->op.as<Var>()) {
	auto var = opt_var.value();

	// Call "relax.invoke_closure" to invoke closure

	auto bound_value = LookupBinding(var);
	if (IsClosure(var) && bound_value.as<CallNode>()) {
	// if the original op was pure, we should use invoke_pure_closure
	Call orig_call = bound_value.value().as_or_throw<Call>();
	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(), false);
	} else if (const auto* func_ty = orig_call->op->ty.as<FuncTypeNode>()) {
	return func_ty->purity;
	} else {
	TVM_FFI_THROW(InternalError)
	<< "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 FuncType (ty = " << orig_call->op->ty << ")";
	}
	}();

	auto prev = call;
	call = Call(Type::Missing(), is_pure ? invoke_pure_closure_op_ : invoke_closure_op_,
	{var, Tuple(call->args)}, {}, {orig_ty});
	}
	}

	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;

	ffi::Array<Expr> new_args = call->args;
	for (const auto arg : nested_call->args) {
	new_args.push_back(arg);
	}

	auto prev = call;
	call = Call(Type::Missing(), nested_call->op, new_args, call->attrs, call->ty_args);
	}
	}

	return ExprMutator::VisitExpr_(call.get());
	}

	Expr VisitExpr_(const VarNode* op) override {
	auto var = ffi::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;
	}
	if (auto bound_value = LookupBinding(opt_var.value())) {
	val = bound_value.value();
	}
	}

	if (const auto* call_node = val.as<tvm::CallNode>()) {
	// recursive call
	auto op = call_node->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_.same_as(op);
	}

	} else if (const auto* global_var = val.as<GlobalVarNode>()) {
	if (closures_.count(ffi::GetRef<GlobalVar>(global_var))) {
	return true;
	}
	IRModule ctx_mod = builder_->GetContextIRModule();
	TVM_FFI_ICHECK(ctx_mod->functions.size() > 0);
	BaseFunc func = ctx_mod->Lookup(ffi::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;
	}
	}

	IRModule Lift() {
	auto glob_funcs = mod_->functions;
	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 = VisitExpr(opt.value()).as_or_throw<Function>();
	builder_->UpdateFunction(gvar, func);
	}
	}
	return builder_->GetContextIRModule();
	}

	private:
	std::unordered_map<Var, Call> nested_closure_map_;
	std::unordered_map<Var, Expr> rebind_map_;
	std::unordered_set<ffi::Variant<GlobalVar, Var>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
	closures_;
	ffi::Optional<Var> current_lambda_var_ = std::nullopt;
	IRModule mod_;

	std::unordered_map<const FunctionNode*, ffi::String> lifted_names_;

	/! \brief Cache ops that would be used later to reduce lookup overhead. /
	const Op& make_closure_op_ = Op::Get("relax.make_closure");
	const Op& invoke_closure_op_ = Op::Get("relax.invoke_closure");
	const Op& invoke_pure_closure_op_ = Op::Get("relax.invoke_pure_closure");
	};

	namespace transform {

	Pass LambdaLift() {
	auto pass_func = [=](IRModule mod, PassContext pc) { return relax::LambdaLifter(mod).Lift(); };
	return tvm::transform::CreateModulePass(pass_func, 1, "LambdaLift", {});
	}

	TVM_FFI_STATIC_INIT_BLOCK() {
	namespace refl = tvm::ffi::reflection;
	refl::GlobalDef().def("relax.transform.LambdaLift", LambdaLift);
	}

	} // namespace transform
	} // namespace relax
	} // namespace tvm