Eager in/clone_in

Makefile

@@ -627,7 +627,6 @@ SOURCE_FILES = \
   Var.cpp \
   VectorizeLoops.cpp \
   WasmExecutor.cpp \
-  WrapCalls.cpp
 # keep-sorted end
 
  C_TEMPLATE_FILES = \
@@ -816,7 +815,6 @@ HEADER_FILES = \
   Var.h \
   VectorizeLoops.h \
   WasmExecutor.h \
-  WrapCalls.h
 # keep-sorted end
 
 OBJECTS = $(SOURCE_FILES:%.cpp=$(BUILD_DIR)/%.o)

src/CMakeLists.txt

@@ -224,7 +224,6 @@ target_sources(
     Var.h
     VectorizeLoops.h
     WasmExecutor.h
-    WrapCalls.h
     # keep-sorted end
 )
 
@@ -420,7 +419,6 @@ target_sources(
     Var.cpp
     VectorizeLoops.cpp
     WasmExecutor.cpp
-    WrapCalls.cpp
     # keep-sorted end
 )
 

src/Derivative.cpp

@@ -729,7 +729,8 @@ void ReverseAccumulationVisitor::propagate_adjoints(
                 calls.reserve(rhs_tuple.size());
                 for (int i = 0; i < (int)rhs_tuple.size(); i++) {
                     calls.push_back(Call::make(
-                        adjoint_funcs[func_key].function(), args, i));
+                        adjoint_funcs[func_key].function(), args, i,
+                        /*follow_global_wrappers=*/true));
                 }
                 prev_adjoint(args) = Tuple(calls);
                 adjoint_funcs[prev_func_key] = prev_adjoint;
@@ -741,7 +742,8 @@ void ReverseAccumulationVisitor::propagate_adjoints(
                 for (int i = 0; i < (int)output_exprs.size(); i++) {
                     expr_adjoints[output_exprs[i]] =
                         Call::make(adjoint_funcs[func_key].function(),
-                                   update_args, i);
+                                   update_args, i,
+                                   /*follow_global_wrappers=*/true);
                 }
 
                 for (Expr &e : reverse_view(expr_list)) {

src/Deserialization.cpp

@@ -507,12 +507,24 @@ void Deserializer::deserialize_function(const Serialize::Func *function, Functio
     const bool no_profiling = function->no_profiling();
     const std::string profiler_display_name = deserialize_string(function->profiler_display_name());
     const bool frozen = function->frozen();
+
+    FunctionPtr global_wrapper;
+    if (const auto *global_wrapper_ref = function->global_wrapper()) {
+        const int32_t func_index = global_wrapper_ref->func_index();
+        if (auto it = this->reverse_function_mappings.find(func_index); it != this->reverse_function_mappings.end() && func_index != -1) {
+            global_wrapper = it->second;
+            // Global-wrapper links are weak (same-group) and are followed
+            // during call resolution (see FunctionPtr::get).
+            global_wrapper.follow_global_wrappers = true;
+        }
+    }
+
     hl_function.update_with_deserialization(name, origin_name, output_types, required_types,
                                             required_dim, args, func_schedule, init_def, updates,
                                             debug_file, output_buffers, extern_arguments, extern_function_name,
                                             name_mangling, extern_function_device_api, extern_proxy_expr,
                                             trace_loads, trace_stores, trace_realizations, trace_tags,
-                                            no_profiling, profiler_display_name, frozen);
+                                            no_profiling, profiler_display_name, frozen, global_wrapper);
 }
 
 Stmt Deserializer::deserialize_stmt(Serialize::Stmt type_code, const void *stmt) {

src/FindCalls.cpp

@@ -44,8 +44,7 @@ class FindCalls : public IRVisitor {
 void populate_environment_helper(const Function &f,
                                  std::map<std::string, Function> *env,
                                  std::vector<Function> *order,
-                                 bool recursive = true,
-                                 bool include_wrappers = false) {
+                                 bool recursive = true) {
     std::map<std::string, Function>::const_iterator iter = env->find(f.name());
     if (iter != env->end()) {
         user_assert(iter->second.same_as(f))
@@ -73,20 +72,14 @@ void populate_environment_helper(const Function &f,
         }
     }
 
-    if (include_wrappers) {
-        for (const auto &it : f.schedule().wrappers()) {
-            insert_func(Function{it.second}, &calls.calls, &calls.order);
-        }
-    }
-
     if (!recursive) {
         for (const Function &g : calls.order) {
             insert_func(g, env, order);
         }
     } else {
         insert_func(f, env, order);
         for (const Function &g : calls.order) {
-            populate_environment_helper(g, env, order, recursive, include_wrappers);
+            populate_environment_helper(g, env, order, recursive);
         }
     }
 }
@@ -97,7 +90,7 @@ std::map<std::string, Function> build_environment(const std::vector<Function> &f
     std::map<std::string, Function> env;
     std::vector<Function> order;
     for (const Function &f : funcs) {
-        populate_environment_helper(f, &env, &order, true, true);
+        populate_environment_helper(f, &env, &order, true);
     }
 
     // Validate the environment: no Parameter (ImageParam, Generator
@@ -180,22 +173,22 @@ std::vector<Function> called_funcs_in_order_found(const std::vector<Function> &f
     std::map<std::string, Function> env;
     std::vector<Function> order;
     for (const Function &f : funcs) {
-        populate_environment_helper(f, &env, &order, true, true);
+        populate_environment_helper(f, &env, &order, true);
     }
     return order;
 }
 
 std::map<std::string, Function> find_transitive_calls(const Function &f) {
     std::map<std::string, Function> res;
     std::vector<Function> order;
-    populate_environment_helper(f, &res, &order, true, false);
+    populate_environment_helper(f, &res, &order, true);
     return res;
 }
 
 std::map<std::string, Function> find_direct_calls(const Function &f) {
     std::map<std::string, Function> res;
     std::vector<Function> order;
-    populate_environment_helper(f, &res, &order, false, false);
+    populate_environment_helper(f, &res, &order, false);
     return res;
 }