From 76a26fb4c81dd57792c5352b081c5981681921b2 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 12:57:34 -0700 Subject: [PATCH 1/7] Make in()/clone_in() rewrite the Func DAG eagerly The old two-phase wrapper model recorded wrappers in a map and applied them during lowering (WrapCalls). Because resolution read the pre-rewrite graph while the rewrite was deferred, the literal and effective graphs diverged, which was the common cause of several bugs: - Issue 3661: cloning the same Func twice crashed, because deep-copying a Func that already carried wrappers couldn't remap them. - "Deletion via cloning": a clone_in that redirected the only path to an already-wrapped Func orphaned it, leaving it in the environment but dead in the effective graph, tripping an assert in RealizationOrder. Custom in(g)/clone_in now rewrite the named consumers eagerly, so the graph always reflects reality and both bugs become unreachable. The consumer is frozen afterwards, since a later definition wouldn't be wrapped. Global f.in() is expressed as a global_wrapper link on the Func plus a follow flag on call-node FunctionPtrs: get() follows the link, so a call resolves to the wrapper as if every caller had been rewritten, and the deep_copy that lowering already does materializes it (rebuilding each call with the wrapper's name). Self-references and wrapper bodies are marked not to follow (via WeakenFunctionPtrs) so they don't cycle. This is retroactive, future-capturing, and chains (f.in().in()) for free, and removes the need for the deferred WrapCalls pass, which is deleted. Adds a fuzz test combining in/clone_in into deep chains and indirect wraps, an error test for adding a definition after wrapping, and a func_clone regression test. Reorders two update-after-wrap tests to define-then-wrap. Co-Authored-By: Claude Opus 4.8 --- Makefile | 2 - src/CMakeLists.txt | 2 - src/FindCalls.cpp | 19 +-- src/Func.cpp | 91 ++++++++-- src/Func.h | 25 ++- src/Function.cpp | 61 ++++++- src/Function.h | 7 + src/FunctionPtr.h | 18 +- src/IR.cpp | 7 +- src/IR.h | 3 +- src/InferArguments.cpp | 5 - src/Lower.cpp | 4 - src/PrintLoopNest.cpp | 4 - src/Schedule.cpp | 13 +- src/WrapCalls.cpp | 180 -------------------- src/WrapCalls.h | 24 --- test/correctness/func_clone.cpp | 61 ++++++- test/correctness/func_wrapper.cpp | 15 +- test/correctness/image_wrapper.cpp | 8 +- test/error/CMakeLists.txt | 1 + test/error/update_after_wrap.cpp | 23 +++ test/fuzz/CMakeLists.txt | 1 + test/fuzz/in_clone.cpp | 259 +++++++++++++++++++++++++++++ 23 files changed, 539 insertions(+), 294 deletions(-) delete mode 100644 src/WrapCalls.cpp delete mode 100644 src/WrapCalls.h create mode 100644 test/error/update_after_wrap.cpp create mode 100644 test/fuzz/in_clone.cpp diff --git a/Makefile b/Makefile index 8bcfc65de765..8e89763c1eb5 100644 --- a/Makefile +++ b/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) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index a5042f36a518..4669f94df4ab 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -222,7 +222,6 @@ target_sources( Var.h VectorizeLoops.h WasmExecutor.h - WrapCalls.h # keep-sorted end ) @@ -418,7 +417,6 @@ target_sources( Var.cpp VectorizeLoops.cpp WasmExecutor.cpp - WrapCalls.cpp # keep-sorted end ) diff --git a/src/FindCalls.cpp b/src/FindCalls.cpp index 8389d2a5ff7e..bcf8af7ea228 100644 --- a/src/FindCalls.cpp +++ b/src/FindCalls.cpp @@ -44,8 +44,7 @@ class FindCalls : public IRVisitor { void populate_environment_helper(const Function &f, std::map *env, std::vector *order, - bool recursive = true, - bool include_wrappers = false) { + bool recursive = true) { std::map::const_iterator iter = env->find(f.name()); if (iter != env->end()) { user_assert(iter->second.same_as(f)) @@ -73,12 +72,6 @@ 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); @@ -86,7 +79,7 @@ void populate_environment_helper(const Function &f, } 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 build_environment(const std::vector &f std::map env; std::vector 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 @@ -164,7 +157,7 @@ std::vector called_funcs_in_order_found(const std::vector &f std::map env; std::vector order; for (const Function &f : funcs) { - populate_environment_helper(f, &env, &order, true, true); + populate_environment_helper(f, &env, &order, true); } return order; } @@ -172,14 +165,14 @@ std::vector called_funcs_in_order_found(const std::vector &f std::map find_transitive_calls(const Function &f) { std::map res; std::vector order; - populate_environment_helper(f, &res, &order, true, false); + populate_environment_helper(f, &res, &order, true); return res; } std::map find_direct_calls(const Function &f) { std::map res; std::vector order; - populate_environment_helper(f, &res, &order, false, false); + populate_environment_helper(f, &res, &order, false); return res; } diff --git a/src/Func.cpp b/src/Func.cpp index f4e4298fae07..b73875fc5d2d 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -2170,6 +2170,9 @@ Func create_in_wrapper(Function wrapped_fn, const string &wrapper_name) { Func wrapper(wrapped_fn.new_function_in_same_group(wrapper_name)); vector args = Func(wrapped_fn).args(); wrapper(args) = Func(wrapped_fn)(args); + // The body's calls to wrapped_fn must not follow its global wrapper (or the + // wrapper would call itself); add_wrapper -> WeakenFunctionPtrs clears the + // follow flag on them. return wrapper; } @@ -2187,36 +2190,58 @@ Func create_clone_wrapper(Function wrapped_fn, const string &wrapper_name) { return wrapper; } +// The set of Func names that count as "reaching" the wrapped Func during +// custom-wrapper resolution: the wrapped Func itself plus all of its existing +// custom wrappers/clones. Because in()/clone_in() rewrite consumers eagerly, a +// Func that has already been wrapped calls the wrapper rather than the original, +// so a call to any existing custom wrapper must be treated as equivalent to a +// call to the original. The global wrapper (the "" entry) is deliberately +// excluded: custom wraps are independent of it. +std::set wrapper_stop_names(const Function &target) { + std::set names; + names.insert(target.name()); + for (const auto &w : target.wrappers()) { + if (!w.first.empty()) { + names.insert(Function(w.second).name()); + } + } + return names; +} + // Walk down the call graph from 'start'. Whenever we find a Func that directly -// calls 'target', record it and stop descending that branch — we don't want to -// pick up unrelated direct callers that happen to live deeper in the subtree. -void collect_direct_callers_of(const Function &target, +// calls the wrapped Func (or an existing wrapper of it, per 'stop_names'), +// record it and stop descending that branch — we don't want to pick up +// unrelated direct callers that happen to live deeper in the subtree. +void collect_direct_callers_of(const std::set &stop_names, const Function &start, std::set &visited, std::map &result) { - if (start.name() == target.name()) { + if (stop_names.count(start.name())) { + // 'start' is the wrapped Func itself or one of its existing wrappers; + // don't record it or descend through it. return; } if (!visited.insert(start.name()).second) { return; } std::map direct = find_direct_calls(start); - if (direct.count(target.name())) { - result.emplace(start.name(), start); - return; + for (const std::string &name : stop_names) { + if (direct.count(name)) { + result.emplace(start.name(), start); + return; + } } for (const auto &kv : direct) { - collect_direct_callers_of(target, kv.second, visited, result); + collect_direct_callers_of(stop_names, kv.second, visited, result); } } // Expand a user-supplied list of caller Funcs to the set of *direct* callers of -// 'target' that lie on a path from any of those callers down to 'target'. -// Funcs that already directly call 'target' pass through unchanged. If a Func -// has no static path to 'target' at all, leave it alone: the IR may not yet -// reflect a wrapper rewrite from a previous in()/clone_in(), and the existing -// in()/clone_in() semantics permit registering a wrapper for such Funcs. +// 'target' (through any existing wrappers) that lie on a path from any of those +// callers down to 'target'. If a Func has no static path to 'target' at all, +// leave it alone. vector resolve_transitive_callers(const Function &target, const vector &fs) { + std::set stop_names = wrapper_stop_names(target); vector out; std::set emitted; auto emit = [&](const Function &g) { @@ -2227,8 +2252,16 @@ vector resolve_transitive_callers(const Function &target, const vector direct_callers; std::set visited; - collect_direct_callers_of(target, f.function(), visited, direct_callers); + collect_direct_callers_of(stop_names, f.function(), visited, direct_callers); if (direct_callers.empty()) { + // No transitive path was found. That's legitimate only if 'f' itself + // directly calls the wrapped Func (e.g. 'f' is an existing wrapper of + // it); then we wrap 'f' directly. Otherwise 'f' does not use the + // wrapped Func at all, which is a user error. + user_assert(find_direct_calls(f.function()).count(target.name())) + << "Cannot wrap Func \"" << target.name() << "\" in \"" << f.name() + << "\" because \"" << f.name() << "\" does not call \"" + << target.name() << "\".\n"; emit(f.function()); } else { for (const auto &kv : direct_callers) { @@ -2256,14 +2289,38 @@ Func get_wrapper(Function wrapped_fn, string wrapper_name, const vector &f } Func wrapper = clone ? create_clone_wrapper(wrapped_fn, wrapper_name) : create_in_wrapper(wrapped_fn, wrapper_name); Function wrapper_fn = wrapper.function(); + if (fs.empty()) { - // Add global wrapper + // Global wrapper. Record it as wrapped_fn's global wrapper; calls to + // wrapped_fn follow global-wrapper links (see FuncRef::operator Expr + // and FunctionPtr::get), so all consumers -- present and future -- + // resolve to it, and deep_copy materializes that at the start of + // lowering. wrapped_fn.add_wrapper("", wrapper_fn); + wrapped_fn.set_global_wrapper(wrapper_fn); } else { for (const Func &f : fs) { user_assert(wrapped_fn.name() != f.name()) << "Cannot create wrapper of itself (\"" << wrapped_fn.name() << "\")\n"; wrapped_fn.add_wrapper(f.name(), wrapper_fn); + + // Eagerly redirect this consumer's calls to wrapped_fn to the + // new wrapper. + Function consumer(f.function()); + FunctionPtr replacement = wrapper_fn.get_contents(); + replacement.follow_global_wrappers = true; + if (consumer.get_contents().group() == wrapper_fn.get_contents().group()) { + // References within a FunctionGroup must be weak. + replacement.weaken(); + } + std::map subs; + subs[wrapped_fn.get_contents()] = replacement; + consumer.substitute_calls(subs); + + // The rewrite only touched the definitions that existed at this + // point, so freeze the consumer: adding more definitions to it + // afterwards would silently fail to be wrapped. + consumer.freeze(); } } return wrapper; @@ -3395,7 +3452,7 @@ FuncRef::operator Expr() const { << "Can't convert a reference Func \"" << func.name() << "\" to an Expr, because " << func.name() << " returns a Tuple.\n"; - return Call::make(func, args); + return Call::make(func, args, 0, /*follow_global_wrappers=*/true); } FuncTupleElementRef FuncRef::operator[](int i) const { @@ -3503,7 +3560,7 @@ Stage FuncTupleElementRef::operator=(const FuncRef &e) { } FuncTupleElementRef::operator Expr() const { - return Internal::Call::make(func_ref.function(), args, idx); + return Internal::Call::make(func_ref.function(), args, idx, /*follow_global_wrappers=*/true); } Realization Func::realize(std::vector sizes, const Target &target) { diff --git a/src/Func.h b/src/Func.h index 0bfb591871c7..65f77eb147b8 100644 --- a/src/Func.h +++ b/src/Func.h @@ -1279,10 +1279,14 @@ class Func { } // @} - /** Creates and returns a new identity Func that wraps this Func. During - * compilation, Halide replaces all calls to this Func done by 'f' - * with calls to the wrapper. If this Func is already wrapped for - * use in 'f', will return the existing wrapper. + /** Creates and returns a new identity Func that wraps this Func, and + * immediately rewrites 'f' to call the wrapper instead of this Func. If + * this Func is already wrapped for use in 'f', returns the existing wrapper + * without rewriting anything again. + * + * The rewrite is eager, so it only affects the definitions of 'f' that + * exist at the time of the call. 'f' is frozen afterwards: adding further + * definitions to it is an error, since they would not be wrapped. * * For example, g.in(f) would rewrite a pipeline like this: \code @@ -1380,15 +1384,18 @@ class Func { Func in(const std::vector &fs); /** Create and return a global identity wrapper, which wraps all calls to - * this Func by any other Func. If a global wrapper already exists, - * returns it. The global identity wrapper is only used by callers for - * which no custom wrapper has been specified. - */ + * this Func by any other Func. If a global wrapper already exists, returns + * it. Unlike the custom wrappers above, this doesn't rewrite consumers: + * calls to this Func are routed through the wrapper (the redirection is + * baked out at the top of lowering), so it applies to every caller -- those + * defined before and after -- that doesn't have a custom wrapper. It's + * independent of the custom wrappers, so the two can be used together. */ Func in(); /** Similar to \ref Func::in; however, instead of replacing the call to * this Func with an identity Func that refers to it, this replaces the - * call with a clone of this Func. + * call with a clone of this Func. Like in(), the rewrite is eager and + * freezes the consumers it rewrites. * * For example, f.clone_in(g) would rewrite a pipeline like this: \code diff --git a/src/Function.cpp b/src/Function.cpp index d9484e5aca0d..39dbe2012b4c 100644 --- a/src/Function.cpp +++ b/src/Function.cpp @@ -38,10 +38,19 @@ class WeakenFunctionPtrs : public IRMutator { Expr expr = IRMutator::visit(c); c = expr.as(); internal_assert(c); + // Match by the named Func (group slot), independent of any + // global-wrapper following, so this still finds self-references after + // the Func has been given a global wrapper. + FunctionPtr unfollowed = c->func; + unfollowed.follow_global_wrappers = false; if (c->func.defined() && - c->func.get() == func) { + unfollowed.get() == func) { FunctionPtr ptr = c->func; ptr.weaken(); + // These are a Func's own self-references, or a wrapper's call to the + // Func it wraps. Either way they must keep calling that Func rather + // than following its global wrapper (which would form a cycle). + ptr.follow_global_wrappers = false; expr = Call::make(c->type, c->name, c->args, c->call_type, ptr, c->value_index, c->image, c->param); @@ -112,6 +121,11 @@ struct FunctionContents { bool frozen = false; + // A weak pointer to this Func's global wrapper, if it has one (created by + // Func::in()). Lives in the same group. A call that follows global-wrapper + // links (see FunctionPtr) resolves through this to the wrapper. + FunctionPtr global_wrapper; + void accept(IRVisitor *visitor) const { func_schedule.accept(visitor); @@ -180,7 +194,17 @@ struct FunctionGroup { }; FunctionContents *FunctionPtr::get() const { - return &(group()->members[idx]); + if (!defined()) { + return nullptr; + } + FunctionContents *c = &(group()->members[idx]); + // Follow the chain of global wrappers to its end, if requested, so that a + // call resolves to the callee's global wrapper as if it had been rewritten. + while (follow_global_wrappers && c->global_wrapper.defined()) { + const FunctionPtr &next = c->global_wrapper; + c = &(next.group()->members[next.idx]); + } + return c; } template<> @@ -518,6 +542,17 @@ void Function::deep_copy(const FunctionPtr ©, DeepCopyMap &copied_map) const copy->output_buffers = contents->output_buffers; copy->func_schedule = contents->func_schedule.deep_copy(copied_map); + // Remap the global-wrapper link, if the wrapper is part of this copy. + if (contents->global_wrapper.defined()) { + auto it = copied_map.find(contents->global_wrapper); + if (it != copied_map.end()) { + FunctionPtr gw = it->second; + gw.weaken(); + gw.follow_global_wrappers = true; + copy->global_wrapper = gw; + } + } + // Copy the pure definition if (contents->init_def.defined()) { copy->init_def = contents->init_def.get_copy(); @@ -547,7 +582,8 @@ void Function::deep_copy(string name, const FunctionPtr ©, DeepCopyMap &copi void Function::define(const vector &args, vector values) { user_assert(!frozen()) << "Func " << name() << " cannot be given a new pure definition, " - << "because it has already been realized or used in the definition of another Func.\n"; + << "because it has already been realized, used in the definition of " + << "another Func, or been the target of a wrapper via in()/clone_in().\n"; user_assert(!has_extern_definition()) << "In pure definition of Func \"" << name() << "\":\n" << "Func with extern definition cannot be given a pure definition.\n"; @@ -686,7 +722,8 @@ void Function::define_update(const vector &_args, vector values, con << "Can't add an update definition without a pure definition first.\n"; user_assert(!frozen()) << "Func " << name() << " cannot be given a new update definition, " - << "because it has already been realized or used in the definition of another Func.\n"; + << "because it has already been realized, used in the definition of " + << "another Func, or been the target of a wrapper via in()/clone_in().\n"; for (auto &value : values) { user_assert(value.defined()) @@ -1202,6 +1239,20 @@ const map &Function::wrappers() const { return contents->func_schedule.wrappers(); } +void Function::set_global_wrapper(const Function &wrapper) { + // Self-references (and wrapper bodies) are already marked not to follow + // global wrappers when they are weakened (see WeakenFunctionPtrs), so + // pointing our global-wrapper link at 'wrapper' won't make them cycle. + FunctionPtr ptr = wrapper.contents; + ptr.weaken(); + ptr.follow_global_wrappers = true; + contents->global_wrapper = ptr; +} + +Function Function::global_wrapper() const { + return contents->global_wrapper.defined() ? Function(contents->global_wrapper) : Function(); +} + Function Function::new_function_in_same_group(const std::string &f) { int group_size = (int)(contents.group()->members.size()); contents.group()->members.resize(group_size + 1); @@ -1268,6 +1319,8 @@ Function &Function::substitute_calls(const map &substi internal_assert(it != substitutions.end()) << "Function not in environment: " << c->func->name << "\n"; FunctionPtr subs = it->second; + // Preserve whether this call routes through global wrappers. + subs.follow_global_wrappers = c->func.follow_global_wrappers; debug(4) << "...Replace call to Func \"" << c->name << "\" with " << "\"" << subs->name << "\"\n"; expr = Call::make(c->type, subs->name, c->args, c->call_type, diff --git a/src/Function.h b/src/Function.h index 55800f3457e5..ca469391698e 100644 --- a/src/Function.h +++ b/src/Function.h @@ -323,6 +323,13 @@ class Function { const std::map &wrappers() const; // @} + /** Set / get this Func's global wrapper (created by Func::in()). Calls that + * follow global-wrapper links resolve through it. */ + // @{ + void set_global_wrapper(const Function &wrapper); + Function global_wrapper() const; + // @} + /** Check if a Function is a trivial wrapper around another * Function, Buffer, or Parameter. Returns the Call node if it * is. Otherwise returns null. diff --git a/src/FunctionPtr.h b/src/FunctionPtr.h index f79000761caa..2e9321a8474e 100644 --- a/src/FunctionPtr.h +++ b/src/FunctionPtr.h @@ -35,6 +35,13 @@ struct FunctionPtr { /** The index of the function within the group. */ int idx = 0; + /** Whether get() follows global-wrapper links (created by Func::in()). Set + * on Call nodes so that a call to a Func resolves to that Func's global + * wrapper, as if every caller had been rewritten; left false on Func handles + * (so they still refer to the Func itself) and on a wrapper's own call to + * the Func it wraps. See FunctionContents::global_wrapper. */ + bool follow_global_wrappers = false; + /** Get a pointer to the group this Function belongs to. */ FunctionGroup *group() const { return weak ? weak : strong.get(); @@ -75,13 +82,16 @@ struct FunctionPtr { return weak || strong.defined(); } - /** Check if two FunctionPtrs refer to the same Function. */ + /** Check if two FunctionPtrs refer to the same Function, resolving through + * global-wrapper links (so a following pointer to a Func and a direct + * pointer to its wrapper are "the same"). */ bool same_as(const FunctionPtr &other) const { - return idx == other.idx && group() == other.group(); + return get() == other.get(); } - /** Pointer comparison, for using FunctionPtrs as keys in maps and - * sets. */ + /** Pointer comparison, for using FunctionPtrs as keys in maps and sets. + * Orders by the resolved Func (following global-wrapper links), matching + * same_as. */ bool operator<(const FunctionPtr &other) const { return get() < other.get(); } diff --git a/src/IR.cpp b/src/IR.cpp index c5158728f367..8910335251dd 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -595,14 +595,17 @@ Stmt Evaluate::make(Expr v) { return node; } -Expr Call::make(const Function &func, const std::vector &args, int idx) { +Expr Call::make(const Function &func, const std::vector &args, int idx, + bool follow_global_wrappers) { internal_assert(idx >= 0 && idx < func.outputs()) << "Value index out of range in call to halide function\n"; internal_assert(func.has_pure_definition() || func.has_extern_definition()) << "Call to undefined halide function\n"; + FunctionPtr fp = func.get_contents(); + fp.follow_global_wrappers = follow_global_wrappers; return make(func.output_types()[(size_t)idx], func.name(), args, Halide, - func.get_contents(), idx, Buffer<>(), Parameter()); + fp, idx, Buffer<>(), Parameter()); } namespace { diff --git a/src/IR.h b/src/IR.h index 16016fca819a..59f1fe42733d 100644 --- a/src/IR.h +++ b/src/IR.h @@ -876,7 +876,8 @@ struct Call : public ExprNode { Buffer<> image = Buffer<>(), Parameter param = Parameter()); /** Convenience constructor for calls to other halide functions */ - static Expr make(const Function &func, const std::vector &args, int idx = 0); + static Expr make(const Function &func, const std::vector &args, int idx = 0, + bool follow_global_wrappers = false); /** Convenience constructor for loads from concrete images */ static Expr make(const Buffer<> &image, const std::vector &args) { diff --git a/src/InferArguments.cpp b/src/InferArguments.cpp index 020d4184642d..15c422df9ac6 100644 --- a/src/InferArguments.cpp +++ b/src/InferArguments.cpp @@ -183,11 +183,6 @@ class InferArguments : public IRGraphVisitor { } } } - - // It also misses wrappers - for (const auto &p : func.wrappers()) { - Function(p.second).accept(this); - } } void include_parameter(const Parameter &p) { diff --git a/src/Lower.cpp b/src/Lower.cpp index cd7ccc9a03f4..b398e750b4d0 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -80,7 +80,6 @@ #include "UnrollLoops.h" #include "UnsafePromises.h" #include "VectorizeLoops.h" -#include "WrapCalls.h" namespace Halide { namespace Internal { @@ -160,9 +159,6 @@ void lower_impl(const vector &output_funcs, iter.second.lock_loop_levels(); } - // Substitute in wrapper Funcs - env = wrap_func_calls(env); - // Compute a realization order and determine group of functions which loops // are to be fused together auto [order, fused_groups] = realization_order(outputs, env); diff --git a/src/PrintLoopNest.cpp b/src/PrintLoopNest.cpp index 61aeb9fe8ef4..95c9f742113c 100644 --- a/src/PrintLoopNest.cpp +++ b/src/PrintLoopNest.cpp @@ -16,7 +16,6 @@ #include "SlidingWindow.h" #include "Target.h" #include "UniquifyVariableNames.h" -#include "WrapCalls.h" #include @@ -180,9 +179,6 @@ string print_loop_nest(const vector &output_funcs) { iter.second.lock_loop_levels(); } - // Substitute in wrapper Funcs - env = wrap_func_calls(env); - // Compute a realization order and determine group of functions which loops // are to be fused together auto [order, fused_groups] = realization_order(outputs, env); diff --git a/src/Schedule.cpp b/src/Schedule.cpp index a2583d0fb732..4e4928fc0546 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -366,13 +366,16 @@ FuncSchedule FuncSchedule::deep_copy( copy.contents->async = contents->async; copy.contents->ring_buffer = contents->ring_buffer; - // Deep-copy wrapper functions. + // Deep-copy wrapper functions. In a partial deep-copy (e.g. cloning a + // single Func via clone_in), the wrapper Funcs may not be among the Funcs + // being copied. Those wrappers describe redirections for callers of the + // original Func and don't apply to the copy, so drop them. for (const auto &iter : contents->wrappers) { - FunctionPtr &copied_func = copied_map[iter.second]; - internal_assert(copied_func.defined()) << Function(iter.second).name() << "\n"; - copy.contents->wrappers[iter.first] = copied_func; + const auto &copied_func = copied_map.find(iter.second); + if (copied_func != copied_map.end()) { + copy.contents->wrappers[iter.first] = copied_func->second; + } } - internal_assert(copy.contents->wrappers.size() == contents->wrappers.size()); return copy; } diff --git a/src/WrapCalls.cpp b/src/WrapCalls.cpp deleted file mode 100644 index 504c4f5d5388..000000000000 --- a/src/WrapCalls.cpp +++ /dev/null @@ -1,180 +0,0 @@ -#include "WrapCalls.h" -#include "FindCalls.h" -#include "Function.h" -#include "FunctionPtr.h" - -#include - -namespace Halide { -namespace Internal { - -using std::map; -using std::set; -using std::string; - -typedef map SubstitutionMap; - -namespace { - -void insert_func_wrapper_helper(map &func_wrappers_map, - const FunctionPtr &in_func, - const FunctionPtr &wrapped_func, - const FunctionPtr &wrapper) { - internal_assert(in_func.defined() && - wrapped_func.defined() && - wrapper.defined()); - internal_assert(func_wrappers_map[in_func].count(wrapped_func) == 0) - << "Should only have one wrapper for each function call in a Func\n"; - - SubstitutionMap &wrappers_map = func_wrappers_map[in_func]; - for (auto iter = wrappers_map.begin(); iter != wrappers_map.end(); ++iter) { - if (iter->second.same_as(wrapped_func)) { - debug(4) << "Merging wrapper of " << Function(in_func).name() - << " [" << Function(iter->first).name() - << ", " << Function(iter->second).name() - << "] with [" << Function(wrapped_func).name() << ", " - << Function(wrapper).name() << "]\n"; - iter->second = wrapper; - return; - } else if (wrapper.same_as(iter->first)) { - debug(4) << "Merging wrapper of " << Function(in_func).name() - << " [" << Function(wrapped_func).name() - << ", " << Function(wrapper).name() - << "] with [" << Function(iter->first).name() - << ", " << Function(iter->second).name() << "]\n"; - wrappers_map.emplace(wrapped_func, iter->second); - wrappers_map.erase(iter); - return; - } - } - wrappers_map[wrapped_func] = wrapper; -} - -void validate_custom_wrapper(const Function &in_func, const Function &wrapped, const Function &wrapper) { - map callees = find_direct_calls(in_func); - if (!callees.count(wrapper.name())) { - std::ostringstream callees_text; - for (const auto &it : callees) { - callees_text << " " << it.second.name() << "\n"; - } - - user_error - << "Cannot wrap \"" << wrapped.name() << "\" in \"" << in_func.name() - << "\" because \"" << in_func.name() << "\" does not call \"" - << wrapped.name() << "\"\n" - << "Direct callees of \"" << in_func.name() << "\" are:\n" - << callees_text.str(); - } -} - -} // anonymous namespace - -map wrap_func_calls(const map &env) { - map wrapped_env; - - map func_wrappers_map; // In Func -> [wrapped Func -> wrapper] - set global_wrappers; - - for (const auto &iter : env) { - wrapped_env.emplace(iter.first, iter.second); - func_wrappers_map[iter.second.get_contents()]; - } - - for (const auto &it : env) { - string wrapped_fname = it.first; - FunctionPtr wrapped_func = it.second.get_contents(); - const auto &wrappers = it.second.schedule().wrappers(); - - // Put the names of all wrappers of this Function into the set for - // faster comparison during the substitution. - set all_func_wrappers; - for (const auto &iter : wrappers) { - all_func_wrappers.insert(Function(iter.second).name()); - } - - for (const auto &iter : wrappers) { - string in_func = iter.first; - FunctionPtr wrapper = iter.second; - - if (in_func.empty()) { // Global wrapper - global_wrappers.insert(Function(wrapper).name()); - for (const auto &wrapped_env_iter : wrapped_env) { - in_func = wrapped_env_iter.first; - if ((wrapped_fname == in_func) || - (all_func_wrappers.find(in_func) != all_func_wrappers.end())) { - // The wrapper should still call the original function, - // so we don't want to rewrite the calls done by the - // wrapper. We also shouldn't rewrite the original - // function itself. - debug(4) << "Skip over replacing \"" << in_func - << "\" with \"" << Function(wrapper).name() << "\"\n"; - continue; - } - if (wrappers.count(in_func)) { - // If the 'in_func' already has custom wrapper for - // 'wrapped_func', don't substitute in the global wrapper. - // Custom wrapper always takes precedence over global wrapper - continue; - } - debug(4) << "Global wrapper: replacing reference of \"" - << wrapped_fname << "\" in \"" << in_func - << "\" with \"" << Function(wrapper).name() << "\"\n"; - insert_func_wrapper_helper(func_wrappers_map, - wrapped_env_iter.second.get_contents(), - wrapped_func, wrapper); - } - } else { // Custom wrapper - debug(4) << "Custom wrapper: replacing reference of \"" - << wrapped_fname << "\" in \"" << in_func << "\" with \"" - << Function(wrapper).name() << "\"\n"; - - const auto &in_func_iter = wrapped_env.find(in_func); - if (in_func_iter == wrapped_env.end()) { - // We find a wrapper definition of 'wrapped_func 'for 'in_func' - // which is not in this pipeline. We don't need to perform - // the substitution since no function in this pipeline will ever - // refer to 'in_func'. - // - // This situation might arise in the following case below: - // f(x) = x; - // g(x) = f(x) + 1; - // f.in(g); - // f.realize(..); - debug(4) << " skip custom wrapper for " << in_func << " [" << wrapped_fname - << " -> " << Function(wrapper).name() << "] since it's not in the pipeline\n"; - continue; - } - insert_func_wrapper_helper(func_wrappers_map, - wrapped_env[in_func].get_contents(), - wrapped_func, - wrapper); - } - } - } - - // Perform the substitution - for (auto &iter : wrapped_env) { - const auto &substitutions = func_wrappers_map[iter.second.get_contents()]; - if (!substitutions.empty()) { - iter.second.substitute_calls(substitutions); - } - } - - // Assert that the custom wrappers are actually used, i.e. if f.in(g) is - // called, but 'f' is never called inside 'g', this will throw a user error. - // Perform the check after the wrapper substitution to handle multi-fold - // wrappers, e.g. f.in(g).in(g). - for (const auto &iter : wrapped_env) { - const auto &substitutions = func_wrappers_map[iter.second.get_contents()]; - for (const auto &pair : substitutions) { - if (global_wrappers.find(Function(pair.second).name()) == global_wrappers.end()) { - validate_custom_wrapper(iter.second, Function(pair.first), Function(pair.second)); - } - } - } - - return wrapped_env; -} - -} // namespace Internal -} // namespace Halide diff --git a/src/WrapCalls.h b/src/WrapCalls.h deleted file mode 100644 index e54244bcf9f8..000000000000 --- a/src/WrapCalls.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef HALIDE_WRAP_CALLS_H -#define HALIDE_WRAP_CALLS_H - -/** \file - * - * Defines pass to replace calls to wrapped Functions with their wrappers. - */ - -#include -#include - -namespace Halide { -namespace Internal { - -class Function; - -/** Replace every call to wrapped Functions in the Functions' definitions with - * call to their wrapper functions. */ -std::map wrap_func_calls(const std::map &env); - -} // namespace Internal -} // namespace Halide - -#endif diff --git a/test/correctness/func_clone.cpp b/test/correctness/func_clone.cpp index 8bf0b4d80e87..151270acad6f 100644 --- a/test/correctness/func_clone.cpp +++ b/test/correctness/func_clone.cpp @@ -124,21 +124,63 @@ int multiple_funcs_sharing_clone_test() { return 0; } -int update_defined_after_clone_test() { +int clone_same_func_into_different_funcs_test() { + Func f("f"), g0("g0"), g1("g1"), g2("g2"), h("h"); + Var x("x"), y("y"); + + f(x, y) = x + y; + g0(x, y) = f(x, y); + g1(x, y) = f(x, y); + g2(x, y) = f(x, y); + h(x, y) = g0(x, y) + g1(x, y) + g2(x, y); + + // Cloning the same Func into two different callers should register two + // independent clones without tripping over each other's wrappers. + Func f_clone_in_g0 = f.clone_in(g0).compute_root(); + Func f_clone_in_g1 = f.clone_in(g1).compute_root(); + f.compute_root(); + g0.compute_root(); + g1.compute_root(); + g2.compute_root(); + + // Expect g0 and g1 to call their own clones, and g2 to call f directly. + CallGraphs expected = { + {h.name(), {g0.name(), g1.name(), g2.name()}}, + {g0.name(), {f_clone_in_g0.name()}}, + {g1.name(), {f_clone_in_g1.name()}}, + {g2.name(), {f.name()}}, + {f_clone_in_g0.name(), {}}, + {f_clone_in_g1.name(), {}}, + {f.name(), {}}, + }; + if (check_call_graphs(h, expected) != 0) { + return 1; + } + + Buffer im = h.realize({200, 200}); + auto func = [](int x, int y) { return 3 * (x + y); }; + if (check_image(im, func)) { + return 1; + } + return 0; +} + +int clone_of_func_with_update_test() { Func f("f"), g("g"); Var x("x"), y("y"); f(x, y) = x + y; g(x, y) = f(x, y); - Func clone = f.clone_in(g); - - // Update of 'g' is defined after f.clone_in(g) is called. g's updates should - // still call f's clone. + // clone_in() rewrites the call graph eagerly, so it only affects the stages + // of 'g' that exist when it is called. Define g's update first, then clone; + // both of g's calls to f are then redirected to the clone. RDom r(0, 100, 0, 100); r.where(r.x < r.y); g(r.x, r.y) += 2 * f(r.x, r.y); + Func clone = f.clone_in(g); + Param param; Var xi("xi"); @@ -346,8 +388,13 @@ int main(int argc, char **argv) { return 1; } - printf("Running update is defined after clone test\n"); - if (update_defined_after_clone_test() != 0) { + printf("Running clone same func into different funcs test\n"); + if (clone_same_func_into_different_funcs_test() != 0) { + return 1; + } + + printf("Running clone of func with update test\n"); + if (clone_of_func_with_update_test() != 0) { return 1; } diff --git a/test/correctness/func_wrapper.cpp b/test/correctness/func_wrapper.cpp index 94df8340861a..97287b232fb0 100644 --- a/test/correctness/func_wrapper.cpp +++ b/test/correctness/func_wrapper.cpp @@ -175,21 +175,22 @@ int global_wrapper_test() { return 0; } -int update_defined_after_wrapper_test() { +int wrapper_of_func_with_update_test() { Func f("f"), g("g"); Var x("x"), y("y"); f(x, y) = x + y; g(x, y) = f(x, y); - Func wrapper = f.in(g); - - // Update of 'g' is defined after f.in(g) is called. g's updates should - // still call f's wrapper. + // in() rewrites the call graph eagerly, so it only affects the stages of 'g' + // that exist when it is called. Define g's update first, then wrap; both of + // g's calls to f are then redirected to the wrapper. RDom r(0, 100, 0, 100); r.where(r.x < r.y); g(r.x, r.y) += 2 * f(r.x, r.y); + Func wrapper = f.in(g); + Param param; Var xi("xi"); @@ -559,8 +560,8 @@ int main(int argc, char **argv) { return 1; } - printf("Running update is defined after wrap test\n"); - if (update_defined_after_wrapper_test() != 0) { + printf("Running wrapper of func with update test\n"); + if (wrapper_of_func_with_update_test() != 0) { return 1; } diff --git a/test/correctness/image_wrapper.cpp b/test/correctness/image_wrapper.cpp index 91d5319c98fb..12c700212e4a 100644 --- a/test/correctness/image_wrapper.cpp +++ b/test/correctness/image_wrapper.cpp @@ -200,14 +200,14 @@ int update_defined_after_wrapper_test() { g(x, y) = img(x, y); - Func wrapper = img.in(g); - - // Update of 'g' is defined after img.in(g) is called. g's updates should - // still call img's wrapper. + // in() rewrites the call graph eagerly, so it only affects the stages of 'g' + // that exist when it is called. Define g's update first, then wrap. RDom r(0, 100, 0, 100); r.where(r.x < r.y); g(r.x, r.y) += 2 * img(r.x, r.y); + Func wrapper = img.in(g); + Param param; Var xi("xi"); diff --git a/test/error/CMakeLists.txt b/test/error/CMakeLists.txt index 234366c8d1b3..21918edd6f00 100644 --- a/test/error/CMakeLists.txt +++ b/test/error/CMakeLists.txt @@ -133,6 +133,7 @@ tests(GROUPS error uninitialized_param.cpp uninitialized_param_2.cpp unknown_target.cpp + update_after_wrap.cpp vector_tile.cpp vectorize_dynamic.cpp vectorize_too_little.cpp diff --git a/test/error/update_after_wrap.cpp b/test/error/update_after_wrap.cpp new file mode 100644 index 000000000000..1a9a1fc27b9b --- /dev/null +++ b/test/error/update_after_wrap.cpp @@ -0,0 +1,23 @@ +#include "Halide.h" +#include + +using namespace Halide; + +int main(int argc, char **argv) { + Func f("f"), g("g"); + Var x("x"), y("y"); + + f(x, y) = x + y; + g(x, y) = f(x, y); + + // Wrapping f in g redirects g's calls to f eagerly, and freezes g. + f.in(g); + + // Adding an update to g now would silently fail to be wrapped, so it is an + // error. + RDom r(0, 10); + g(r, r) += 1; + + printf("Success!\n"); + return 0; +} diff --git a/test/fuzz/CMakeLists.txt b/test/fuzz/CMakeLists.txt index 95716c2c134c..b36197bd4517 100644 --- a/test/fuzz/CMakeLists.txt +++ b/test/fuzz/CMakeLists.txt @@ -9,6 +9,7 @@ tests(GROUPS fuzz SOURCES bounds.cpp cse.cpp + in_clone.cpp lossless_cast.cpp simplify.cpp solve.cpp diff --git a/test/fuzz/in_clone.cpp b/test/fuzz/in_clone.cpp new file mode 100644 index 000000000000..5fee0e28cd2d --- /dev/null +++ b/test/fuzz/in_clone.cpp @@ -0,0 +1,259 @@ +#include "Halide.h" + +#include "fuzz_helpers.h" +#include +#include + +// Fuzz test for Func::in and Func::clone_in. We start from a simple randomly +// generated pipeline and apply a random sequence of in()/clone_in() calls, +// including clones of clones of clones and "indirect" wraps where the consumer +// passed is not a direct caller of the wrapped Func. in()/clone_in() rewrite +// the call graph eagerly and are all semantics-preserving, so the result of +// realizing the pipeline must not change no matter what we do to it. This +// stresses the wrapper/clone machinery (in particular deep-copying Funcs that +// already carry wrappers) and checks both that nothing triggers an internal +// error and that the numerics are preserved. + +namespace { + +using namespace Halide; + +// A data-only description of a simple base pipeline, so we can build it twice +// (once as a reference, once as the pipeline we mutate) from the same random +// choices. +struct NodeDesc { + // op == LEAF: value = ca*x + cb*y + cc. + // otherwise: a binary/unary combination of earlier nodes a and b. + enum Op { LEAF, + ADD, + SUB, + ADD_CONST, + SCALE_ADD } op; + int a = 0, b = 0; // indices of input nodes (< this node's index) + int ca = 0, cb = 0, cc = 0, k = 0; +}; + +struct PipelineDesc { + std::vector nodes; + int output = 0; +}; + +PipelineDesc generate_pipeline(FuzzingContext &fuzz) { + PipelineDesc desc; + int n = fuzz.ConsumeIntegralInRange(4, 8); + desc.nodes.resize(n); + + // Node 0 is the only leaf; every other node combines two earlier nodes, + // which gives a DAG rooted at the last node with plenty of shared + // sub-Funcs (so a Func can have several direct callers). + NodeDesc &leaf = desc.nodes[0]; + leaf.op = NodeDesc::LEAF; + leaf.ca = fuzz.ConsumeIntegralInRange(0, 3); + leaf.cb = fuzz.ConsumeIntegralInRange(0, 3); + leaf.cc = fuzz.ConsumeIntegralInRange(-4, 4); + + for (int i = 1; i < n; i++) { + NodeDesc &node = desc.nodes[i]; + node.op = (NodeDesc::Op)fuzz.ConsumeIntegralInRange(NodeDesc::ADD, NodeDesc::SCALE_ADD); + node.a = fuzz.ConsumeIntegralInRange(0, i - 1); + node.b = fuzz.ConsumeIntegralInRange(0, i - 1); + node.k = fuzz.ConsumeIntegralInRange(-4, 4); + } + desc.output = n - 1; + return desc; +} + +// Build the Funcs described by 'desc' into 'funcs'. +void build_funcs(const PipelineDesc &desc, Var x, Var y, std::vector &funcs) { + funcs.resize(desc.nodes.size()); + for (size_t i = 0; i < desc.nodes.size(); i++) { + const NodeDesc &node = desc.nodes[i]; + Func &f = funcs[i]; + switch (node.op) { + case NodeDesc::LEAF: + f(x, y) = node.ca * x + node.cb * y + node.cc; + break; + case NodeDesc::ADD: + f(x, y) = funcs[node.a](x, y) + funcs[node.b](x, y); + break; + case NodeDesc::SUB: + f(x, y) = funcs[node.a](x, y) - funcs[node.b](x, y); + break; + case NodeDesc::ADD_CONST: + f(x, y) = funcs[node.a](x, y) + node.k; + break; + case NodeDesc::SCALE_ADD: + f(x, y) = 2 * funcs[node.a](x, y) + funcs[node.b](x, y); + break; + } + } +} + +// The direct callees (among tracked Funcs) of each node's definition. +std::set direct_callees(const NodeDesc &node) { + switch (node.op) { + case NodeDesc::LEAF: + return {}; + case NodeDesc::ADD_CONST: + return {node.a}; + default: + return {node.a, node.b}; + } +} + +// A model of the call graph that mirrors what in()/clone_in() do, so we can +// generate only valid operations. Because the rewrite is eager, this call graph +// stays exactly in sync with the real pipeline. +struct CallGraphModel { + std::vector> callees; // what each Func currently calls + // For each Func, the custom (non-global) wrappers of it that exist, and the + // consumers already wrapped for it. Reusing a consumer, or wrapping a set of + // consumers that mixes wrapped and unwrapped ones, is what raises user + // errors, so we track these to avoid it. + std::vector> custom_wrappers; + std::vector> wrapped_consumers; + + int size() const { + return (int)callees.size(); + } + + int add(std::set node_callees) { + int idx = size(); + callees.push_back(std::move(node_callees)); + custom_wrappers.emplace_back(); + wrapped_consumers.emplace_back(); + return idx; + } + + // Resolution stops at the target or any of its custom wrappers, mirroring + // Func.cpp's resolve_transitive_callers. + bool is_stop(int target, int node) const { + return node == target || custom_wrappers[target].count(node); + } + + void collect(int target, int start, std::set &visited, std::set &result) const { + if (is_stop(target, start) || !visited.insert(start).second) { + return; + } + for (int c : callees[start]) { + if (is_stop(target, c)) { + result.insert(start); + return; + } + } + for (int c : callees[start]) { + collect(target, c, visited, result); + } + } + + // The direct callers of 'target' a wrap with consumer 'consumer' would + // rewrite. Empty if 'consumer' has no path to 'target'. + std::set resolve(int target, int consumer) const { + std::set visited, result; + collect(target, consumer, visited, result); + return result; + } +}; + +int compare(const Buffer &ref, const Buffer &got) { + for (int y = 0; y < ref.height(); y++) { + for (int x = 0; x < ref.width(); x++) { + if (ref(x, y) != got(x, y)) { + std::cerr << "Mismatch at (" << x << ", " << y << "): " + << "expected " << ref(x, y) << ", got " << got(x, y) << "\n"; + return 1; + } + } + } + return 0; +} + +} // namespace + +FUZZ_TEST(in_clone, FuzzingContext &fuzz) { + Var x("x"), y("y"); + const int W = 8, H = 8; + + PipelineDesc desc = generate_pipeline(fuzz); + + // Reference: build and realize without any wrappers. + Buffer reference; + { + std::vector funcs; + build_funcs(desc, x, y, funcs); + for (int i = 0; i < (int)funcs.size(); i++) { + if (i != desc.output) { + funcs[i].compute_root(); + } + } + reference = funcs[desc.output].realize({W, H}); + } + + // Test: build the same pipeline, then apply a random sequence of + // in()/clone_in() operations to grow it into something complex. + std::vector funcs; + build_funcs(desc, x, y, funcs); + + CallGraphModel model; + for (const NodeDesc &node : desc.nodes) { + model.add(direct_callees(node)); + } + const int output = desc.output; + + const int num_ops = fuzz.ConsumeIntegralInRange(5, 40); + const int max_funcs = 64; + for (int step = 0; step < num_ops && model.size() < max_funcs; step++) { + // Pick a wrap target (never the output, which has no callers) and a + // distinct consumer. + int target = fuzz.ConsumeIntegralInRange(0, model.size() - 1); + int consumer = fuzz.ConsumeIntegralInRange(0, model.size() - 1); + if (target == output || consumer == target) { + continue; + } + + std::set resolved = model.resolve(target, consumer); + + // Skip consumers that don't reach the target, and any wrap that would + // reuse a consumer already wrapped for this target (which is the only + // way these calls raise a user error). + if (resolved.empty()) { + continue; + } + bool conflict = false; + for (int f : resolved) { + if (model.wrapped_consumers[target].count(f)) { + conflict = true; + break; + } + } + if (conflict) { + continue; + } + + bool clone = fuzz.ConsumeBool(); + Func wrapper = clone ? funcs[target].clone_in(funcs[consumer]) + : funcs[target].in(funcs[consumer]); + + // Mirror the rewrite in the model: a clone recomputes what the target + // computed (same callees); a plain wrapper just reads from the target. + std::set wrapper_callees = clone ? model.callees[target] : std::set{target}; + int w = model.add(std::move(wrapper_callees)); + model.custom_wrappers[target].insert(w); + for (int f : resolved) { + model.wrapped_consumers[target].insert(f); + // The rewrite substitutes target -> wrapper, so it only redirects a + // consumer that calls the target directly. A consumer that reaches + // the target only through an existing wrapper is left unchanged + // (and the new wrapper is dead), matching Halide. + if (model.callees[f].erase(target)) { + model.callees[f].insert(w); + } + } + + funcs.push_back(wrapper); + funcs.back().compute_root(); + } + + Buffer result = funcs[output].realize({W, H}); + return compare(reference, result); +} From f3b2ca8eb0c9bbd8f9d133210f669309a4306796 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 15:16:36 -0700 Subject: [PATCH 2/7] Drop wrappers[""] and serialize the global-wrapper link Global .in() wrappers no longer occupy a "" entry in the wrappers map. Idempotency is decided by Function::global_wrapper(), and the freeze + follow-flag clearing that used to happen in add_wrapper("", W) now lives in set_global_wrapper. Function::global_wrapper() returns a strong, non-following handle to the immediate wrapper. Following there would make f.in() resolve to the end of the wrapper chain, so f.in().in() would wrap the wrong Func and, via copy_to_host, hit "Extern Func has itself as an argument". Serialize FunctionContents::global_wrapper as a WrapperRef so a round-tripped pipeline keeps its global-wrapper links. Co-Authored-By: Claude Opus 4.8 --- src/Deserialization.cpp | 14 ++++- src/Func.cpp | 135 +++++++++++++++++++++------------------- src/Function.cpp | 28 +++++++-- src/Function.h | 8 ++- src/Serialization.cpp | 15 ++++- src/halide_ir.fbs | 1 + 6 files changed, 126 insertions(+), 75 deletions(-) diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index e31c27eea492..a84e737ac544 100644 --- a/src/Deserialization.cpp +++ b/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) { diff --git a/src/Func.cpp b/src/Func.cpp index b586e982e700..d0b8f4073c3d 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -2339,80 +2339,87 @@ vector resolve_transitive_callers(const Function &target, const vector &fs_in, bool clone) { vector fs = fs_in.empty() ? fs_in : resolve_transitive_callers(wrapped_fn, fs_in); - // Either all Funcs in 'fs' have the same wrapper or they don't already - // have any wrappers. Otherwise, throw an error. If 'fs' is empty, then - // it is a global wrapper. const map &wrappers = wrapped_fn.wrappers(); wrapper_name += ("$" + std::to_string(wrappers.size())); - const auto &iter = fs.empty() ? wrappers.find("") : wrappers.find(fs[0].name()); - if (iter == wrappers.end()) { - // Make sure the other Funcs also don't have any wrappers + + if (fs.empty()) { + // Global wrapper (Func::in()). Idempotent: return the existing one, if + // any. Unlike custom wrappers it lives in a dedicated global_wrapper + // link rather than the wrappers map. + Function existing = wrapped_fn.global_wrapper(); + if (existing.get_contents().defined()) { + return Func(existing); + } + } else { + // Either all Funcs in 'fs' already share the same wrapper, or none of + // them have one. Otherwise it's an error. + const auto &iter = wrappers.find(fs[0].name()); + if (iter != wrappers.end()) { + internal_assert(iter->second.defined()); + validate_wrapper(wrapped_fn.name(), wrappers, fs, iter->second); + Function wrapper(iter->second); + internal_assert(wrapper.frozen()); + return Func(wrapper); + } for (size_t i = 1; i < fs.size(); ++i) { user_assert(wrappers.count(fs[i].name()) == 0) << "Cannot define the wrapper since " << fs[i].name() << " already has a wrapper while " << fs[0].name() << " doesn't \n"; } - Func wrapper = clone ? create_clone_wrapper(wrapped_fn, wrapper_name) : create_in_wrapper(wrapped_fn, wrapper_name); - Function wrapper_fn = wrapper.function(); - - // Build a profiler display name like ".in()" or - // ".in(, )" using the wrapped Func's display - // name and the consumers' display names (falling back to the - // IR-level name in each case). For .clone_in() use "clone_in". - auto display = [](const Function &f) { - return f.profiler_display_name().empty() ? f.name() : f.profiler_display_name(); - }; - std::string profiler_name = display(wrapped_fn) + (clone ? ".clone_in(" : ".in("); - for (size_t i = 0; i < fs.size(); i++) { - if (i > 0) { - profiler_name += ", "; - } - profiler_name += display(fs[i].function()); - } - profiler_name += ")"; - wrapper_fn.set_profiler_display_name(profiler_name); - - if (fs.empty()) { - // Global wrapper. Record it as wrapped_fn's global wrapper; calls to - // wrapped_fn follow global-wrapper links (see FuncRef::operator Expr - // and FunctionPtr::get), so all consumers -- present and future -- - // resolve to it, and deep_copy materializes that at the start of - // lowering. - wrapped_fn.add_wrapper("", wrapper_fn); - wrapped_fn.set_global_wrapper(wrapper_fn); - } else { - for (const Func &f : fs) { - user_assert(wrapped_fn.name() != f.name()) - << "Cannot create wrapper of itself (\"" << wrapped_fn.name() << "\")\n"; - wrapped_fn.add_wrapper(f.name(), wrapper_fn); - - // Eagerly redirect this consumer's calls to wrapped_fn to the - // new wrapper. - Function consumer(f.function()); - FunctionPtr replacement = wrapper_fn.get_contents(); - replacement.follow_global_wrappers = true; - if (consumer.get_contents().group() == wrapper_fn.get_contents().group()) { - // References within a FunctionGroup must be weak. - replacement.weaken(); - } - std::map subs; - subs[wrapped_fn.get_contents()] = replacement; - consumer.substitute_calls(subs); - - // The rewrite only touched the definitions that existed at this - // point, so freeze the consumer: adding more definitions to it - // afterwards would silently fail to be wrapped. - consumer.freeze(); - } + } + + Func wrapper = clone ? create_clone_wrapper(wrapped_fn, wrapper_name) : create_in_wrapper(wrapped_fn, wrapper_name); + Function wrapper_fn = wrapper.function(); + + // Build a profiler display name like ".in()" or + // ".in(, )" using the wrapped Func's display + // name and the consumers' display names (falling back to the + // IR-level name in each case). For .clone_in() use "clone_in". + auto display = [](const Function &f) { + return f.profiler_display_name().empty() ? f.name() : f.profiler_display_name(); + }; + std::string profiler_name = display(wrapped_fn) + (clone ? ".clone_in(" : ".in("); + for (size_t i = 0; i < fs.size(); i++) { + if (i > 0) { + profiler_name += ", "; } - return wrapper; + profiler_name += display(fs[i].function()); } - internal_assert(iter->second.defined()); - validate_wrapper(wrapped_fn.name(), wrappers, fs, iter->second); + profiler_name += ")"; + wrapper_fn.set_profiler_display_name(profiler_name); + + if (fs.empty()) { + // Global wrapper. Calls to wrapped_fn follow global-wrapper links (see + // FuncRef::operator Expr and FunctionPtr::get), so all consumers -- + // present and future -- resolve to it, and deep_copy materializes that + // at the start of lowering. + wrapped_fn.set_global_wrapper(wrapper_fn); + } else { + for (const Func &f : fs) { + user_assert(wrapped_fn.name() != f.name()) + << "Cannot create wrapper of itself (\"" << wrapped_fn.name() << "\")\n"; + wrapped_fn.add_wrapper(f.name(), wrapper_fn); + + // Eagerly redirect this consumer's calls to wrapped_fn to the new + // wrapper. + Function consumer(f.function()); + FunctionPtr replacement = wrapper_fn.get_contents(); + replacement.follow_global_wrappers = true; + if (consumer.get_contents().group() == wrapper_fn.get_contents().group()) { + // References within a FunctionGroup must be weak. + replacement.weaken(); + } + std::map subs; + subs[wrapped_fn.get_contents()] = replacement; + consumer.substitute_calls(subs); - Function wrapper(iter->second); - internal_assert(wrapper.frozen()); - return Func(wrapper); + // The rewrite only touched the definitions that existed at this + // point, so freeze the consumer: adding more definitions to it + // afterwards would silently fail to be wrapped. + consumer.freeze(); + } + } + return wrapper; } } // anonymous namespace diff --git a/src/Function.cpp b/src/Function.cpp index e64588eb7abb..15674b18f353 100644 --- a/src/Function.cpp +++ b/src/Function.cpp @@ -382,7 +382,8 @@ void Function::update_with_deserialization(const std::string &name, const std::vector &trace_tags, bool no_profiling, const std::string &profiler_display_name, - bool frozen) { + bool frozen, + const FunctionPtr &global_wrapper) { contents->name = name; contents->origin_name = origin_name; contents->output_types = output_types; @@ -406,6 +407,7 @@ void Function::update_with_deserialization(const std::string &name, contents->no_profiling = no_profiling; contents->profiler_display_name = profiler_display_name; contents->frozen = frozen; + contents->global_wrapper = global_wrapper; } namespace { @@ -1253,10 +1255,15 @@ const map &Function::wrappers() const { return contents->func_schedule.wrappers(); } -void Function::set_global_wrapper(const Function &wrapper) { - // Self-references (and wrapper bodies) are already marked not to follow - // global wrappers when they are weakened (see WeakenFunctionPtrs), so - // pointing our global-wrapper link at 'wrapper' won't make them cycle. +void Function::set_global_wrapper(Function &wrapper) { + wrapper.freeze(); + + // Weaken the wrapper's back-references to us, and mark them not to follow + // global wrappers -- otherwise the wrapper's own call to us would resolve + // through the link we're about to set and call itself. + WeakenFunctionPtrs weakener(contents.get()); + wrapper.mutate(&weakener); + FunctionPtr ptr = wrapper.contents; ptr.weaken(); ptr.follow_global_wrappers = true; @@ -1264,7 +1271,16 @@ void Function::set_global_wrapper(const Function &wrapper) { } Function Function::global_wrapper() const { - return contents->global_wrapper.defined() ? Function(contents->global_wrapper) : Function(); + if (!contents->global_wrapper.defined()) { + return Function(); + } + // Return a strong handle pointing directly at the immediate wrapper. The + // stored link is weak and follows further global-wrapper links; strip both + // so callers get this Func's own wrapper rather than the end of the chain. + FunctionPtr ptr = contents->global_wrapper; + ptr.strengthen(); + ptr.follow_global_wrappers = false; + return Function(ptr); } Function Function::new_function_in_same_group(const std::string &f) { diff --git a/src/Function.h b/src/Function.h index 8c0592a1459c..a7344d6a61c1 100644 --- a/src/Function.h +++ b/src/Function.h @@ -91,7 +91,8 @@ class Function { const std::vector &trace_tags, bool no_profiling, const std::string &profiler_display_name, - bool frozen); + bool frozen, + const FunctionPtr &global_wrapper); /** Get a handle on the halide function contents that this Function * represents. */ @@ -336,9 +337,10 @@ class Function { // @} /** Set / get this Func's global wrapper (created by Func::in()). Calls that - * follow global-wrapper links resolve through it. */ + * follow global-wrapper links resolve through it. set_global_wrapper freezes + * the wrapper and weakens its back-reference to this Func. */ // @{ - void set_global_wrapper(const Function &wrapper); + void set_global_wrapper(Function &wrapper); Function global_wrapper() const; // @} diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 2ca84e441c77..b1fe24af2815 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -1034,6 +1034,18 @@ Offset Serializer::serialize_function(FlatBufferBuilder &builde const bool no_profiling = function.should_not_profile(); const auto profiler_display_name_serialized = serialize_string(builder, function.profiler_display_name()); const bool frozen = function.frozen(); + + Offset global_wrapper_serialized = 0; + const Function global_wrapper = function.global_wrapper(); + if (global_wrapper.get_contents().defined()) { + auto global_wrapper_name_serialized = serialize_string(builder, global_wrapper.name()); + int func_index = -1; + if (auto it = this->func_mappings.find(global_wrapper.name()); it != this->func_mappings.end()) { + func_index = it->second; + } + global_wrapper_serialized = Serialize::CreateWrapperRef(builder, global_wrapper_name_serialized, func_index); + } + auto func = Serialize::CreateFunc(builder, name_serialized, origin_name_serialized, @@ -1057,7 +1069,8 @@ Offset Serializer::serialize_function(FlatBufferBuilder &builde builder.CreateVector(trace_tags_serialized), no_profiling, profiler_display_name_serialized, - frozen); + frozen, + global_wrapper_serialized); return func; } diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 60fdafdf6b5d..0e1fb38d5606 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -721,6 +721,7 @@ table Func { no_profiling: bool = false; profiler_display_name: string; frozen: bool = false; + global_wrapper: WrapperRef; } table Pipeline { From 57de2741ff3fd5b32e348c192d555b852f5ececa Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 15:38:45 -0700 Subject: [PATCH 3/7] Clarify Func::in() doc comment Describe the observable effect -- all past and future consumers are rewritten to call the wrapper -- instead of the follow-link mechanism. Co-Authored-By: Claude Opus 4.8 --- src/Func.h | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/Func.h b/src/Func.h index ef754f4c607d..396d2d436a0c 100644 --- a/src/Func.h +++ b/src/Func.h @@ -1408,13 +1408,10 @@ class Func { * this will throw an error. */ Func in(const std::vector &fs); - /** Create and return a global identity wrapper, which wraps all calls to - * this Func by any other Func. If a global wrapper already exists, returns - * it. Unlike the custom wrappers above, this doesn't rewrite consumers: - * calls to this Func are routed through the wrapper (the redirection is - * baked out at the top of lowering), so it applies to every caller -- those - * defined before and after -- that doesn't have a custom wrapper. It's - * independent of the custom wrappers, so the two can be used together. */ + /** Create and return a global identity wrapper, and rewrite all consumers + * of this Func -- both those defined before this call and those defined + * after -- to call the wrapper instead. Consumers with a custom wrapper of + * this Func are unaffected. If a global wrapper already exists, returns it. */ Func in(); /** Similar to \ref Func::in; however, instead of replacing the call to From 8bcb77de63d7c9f9655bec98383d6ac073f78f33 Mon Sep 17 00:00:00 2001 From: "halide-ci[bot]" <266445882+halide-ci[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:53:19 +0000 Subject: [PATCH 4/7] Apply pre-commit auto-fixes --- test/fuzz/in_clone.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/test/fuzz/in_clone.cpp b/test/fuzz/in_clone.cpp index 5fee0e28cd2d..ab46e833bee4 100644 --- a/test/fuzz/in_clone.cpp +++ b/test/fuzz/in_clone.cpp @@ -29,7 +29,7 @@ struct NodeDesc { SUB, ADD_CONST, SCALE_ADD } op; - int a = 0, b = 0; // indices of input nodes (< this node's index) + int a = 0, b = 0; // indices of input nodes (< this node's index) int ca = 0, cb = 0, cc = 0, k = 0; }; @@ -231,8 +231,7 @@ FUZZ_TEST(in_clone, FuzzingContext &fuzz) { } bool clone = fuzz.ConsumeBool(); - Func wrapper = clone ? funcs[target].clone_in(funcs[consumer]) - : funcs[target].in(funcs[consumer]); + Func wrapper = clone ? funcs[target].clone_in(funcs[consumer]) : funcs[target].in(funcs[consumer]); // Mirror the rewrite in the model: a clone recomputes what the target // computed (same callees); a plain wrapper just reads from the target. From d08c98969d53d9861be2159c3f462e46bbde8e7e Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 5 Aug 2026 09:23:23 -0700 Subject: [PATCH 5/7] Make global wrappers redirect autodiff calls and get_func handles The follow_global_wrappers flag is what makes a call node resolve to a Func's global wrapper (Func::in()). It was set only at the user-facing FuncRef chokepoints, so two consumer edges escaped it: - Autodiff builds its adjoint calls directly via Call::make, so a global wrapper never redirected them. The wrapper ended up with no consumers and was pruned, making a compute_at onto it an invalid location (seen in the anderson2021 cost-model schedule). - Pipeline::get_func returned a handle built from a Call node, inheriting its follow flag. After a first Func::in() the handle shifted to the new wrapper, so a second Func::in() wrapped the wrapper. Drop the default on Call::make's follow_global_wrappers argument so every consumer-edge site states intent, set it true in autodiff, and have get_func return a non-following handle. The rfactor self-reference and the ScheduleFunctions blend self-call stay non-following. Add tests: an rfactor+in test whose call graph pins the external edge (follows) versus the intermediate's self-reference (does not), and a test that Func::in() on a get_func handle is idempotent. Co-Authored-By: Claude Opus 4.8 --- src/Derivative.cpp | 6 ++-- src/Func.cpp | 5 ++- src/IR.h | 9 +++-- src/Pipeline.cpp | 8 ++++- src/ScheduleFunctions.cpp | 5 ++- test/correctness/func_wrapper.cpp | 56 +++++++++++++++++++++++++++++++ test/correctness/rfactor.cpp | 49 +++++++++++++++++++++++++++ 7 files changed, 130 insertions(+), 8 deletions(-) diff --git a/src/Derivative.cpp b/src/Derivative.cpp index 041fb721513a..cef5b9e94e36 100644 --- a/src/Derivative.cpp +++ b/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)) { diff --git a/src/Func.cpp b/src/Func.cpp index d0b8f4073c3d..db18fd7a590b 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -593,7 +593,10 @@ class SubstituteSelfReference : public IRMutator { vector args; args.insert(args.end(), c->args.begin(), c->args.end()); args.insert(args.end(), new_args.begin(), new_args.end()); - expr = Call::make(substitute, args, c->value_index); + // This rewrites a Func's self-reference into a self-reference of + // the rfactor intermediate, so it must not follow global wrappers. + expr = Call::make(substitute, args, c->value_index, + /*follow_global_wrappers=*/false); } return expr; } diff --git a/src/IR.h b/src/IR.h index e91b3ff87f7d..96185897b566 100644 --- a/src/IR.h +++ b/src/IR.h @@ -929,9 +929,12 @@ struct Call : public ExprNode { * unchanged if the new args are the same as the existing ones. */ Expr with(const std::vector &args) const; - /** Convenience constructor for calls to other halide functions */ - static Expr make(const Function &func, const std::vector &args, int idx = 0, - bool follow_global_wrappers = false); + /** Convenience constructor for calls to other halide functions. Pass + * follow_global_wrappers = true when constructing a consumer's call to + * 'func', so that a global wrapper (Func::in()) redirects it; pass false + * only for a Func's reference to itself. */ + static Expr make(const Function &func, const std::vector &args, int idx, + bool follow_global_wrappers); /** Convenience constructor for loads from concrete images */ static Expr make(const Buffer<> &image, const std::vector &args) { diff --git a/src/Pipeline.cpp b/src/Pipeline.cpp index 8bb53d8d3298..5d4bef6242d1 100644 --- a/src/Pipeline.cpp +++ b/src/Pipeline.cpp @@ -284,7 +284,13 @@ Func Pipeline::get_func(size_t index) { user_assert(index < order.size()) << "Index value passed is " << index << "; however, there are only " << order.size() << " functions in the pipeline.\n"; - return Func(env.find(order[index])->second); + // Return a stable handle to the named Func. The environment is built by + // walking Call nodes, whose FunctionPtrs may be marked to follow global + // wrappers; a Func handle must not follow, or it would silently shift to a + // wrapper created later (e.g. by a subsequent Func::in()). + FunctionPtr ptr = env.find(order[index])->second.get_contents(); + ptr.follow_global_wrappers = false; + return Func(Function(ptr)); } void Pipeline::compile_to(const std::map &output_files, diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp index 53d94f3382a9..23b6e041106a 100644 --- a/src/ScheduleFunctions.cpp +++ b/src/ScheduleFunctions.cpp @@ -150,7 +150,10 @@ class AddPredicates : public IRGraphMutator { if (type == ApplySplitResult::BlendProvides) { int idx = 0; for (Expr &v : values) { - v = select(cond, v, Call::make(func, args, idx++)); + // A Func referring to its own prior value; must not resolve + // through a global wrapper. + v = select(cond, v, Call::make(func, args, idx++, + /*follow_global_wrappers=*/false)); } return p->with(values, args, predicate); } else if (type == ApplySplitResult::PredicateProvides) { diff --git a/test/correctness/func_wrapper.cpp b/test/correctness/func_wrapper.cpp index 97287b232fb0..8cba8874712a 100644 --- a/test/correctness/func_wrapper.cpp +++ b/test/correctness/func_wrapper.cpp @@ -175,6 +175,57 @@ int global_wrapper_test() { return 0; } +int global_wrapper_via_get_func_test() { + // A Func handle from Pipeline::get_func is built by walking Call nodes, + // whose FunctionPtrs may be marked to follow global wrappers. The handle + // itself must not follow, or a second Func::in() on it would wrap the + // wrapper created by the first (producing a wrapper-of-a-wrapper). + Func f("f"), g("g"); + Var x("x"), y("y"); + f(x, y) = x + y; + g(x, y) = f(x, y); + + Pipeline p({g}); + Func f_from_env; + for (int i = 0; i < 2; i++) { + Func fn = p.get_func(i); + if (fn.name() == f.name()) { + f_from_env = fn; + } + } + if (!f_from_env.defined()) { + printf("get_func did not return f\n"); + return 1; + } + + Func w1 = f_from_env.in(); + Func w2 = f_from_env.in(); + if (w1.name() != w2.name()) { + printf("Func::in() on a get_func handle was not idempotent: %s vs %s\n", + w1.name().c_str(), w2.name().c_str()); + return 1; + } + + f.compute_root(); + w1.compute_root(); + + CallGraphs expected = { + {g.name(), {w1.name()}}, + {w1.name(), {f.name()}}, + {f.name(), {}}, + }; + if (check_call_graphs(g, expected) != 0) { + return 1; + } + + Buffer im = g.realize({50, 50}); + auto func = [](int x, int y) { return x + y; }; + if (check_image(im, func)) { + return 1; + } + return 0; +} + int wrapper_of_func_with_update_test() { Func f("f"), g("g"); Var x("x"), y("y"); @@ -560,6 +611,11 @@ int main(int argc, char **argv) { return 1; } + printf("Running global wrap via get_func test\n"); + if (global_wrapper_via_get_func_test() != 0) { + return 1; + } + printf("Running wrapper of func with update test\n"); if (wrapper_of_func_with_update_test() != 0) { return 1; diff --git a/test/correctness/rfactor.cpp b/test/correctness/rfactor.cpp index b3d598117168..7ab8b2396c59 100644 --- a/test/correctness/rfactor.cpp +++ b/test/correctness/rfactor.cpp @@ -60,6 +60,53 @@ int simple_rfactor_test() { return 0; } +template +int rfactor_wrapper_test() { + // A global wrapper on an rfactor intermediate. The reducing Func's call to + // the intermediate must follow the wrapper (external edge), but the + // intermediate's self-reference in its own update must not (following would + // make intm -> wrapper -> intm a cycle). + Func f("f"), g("g"); + Var x("x"), y("y"); + + f(x, y) = x + y; + f.compute_root(); + + g(x, y) = 40; + RDom r(10, 20, 30, 40); + g(r.x, r.y) = max(g(r.x, r.y) + f(r.x, r.y), g(r.x, r.y)); + g.reorder_storage(y, x); + + Var u("u"); + Func intm = g.update(0).rfactor(r.y, u); + Func intm_w = intm.in(); + intm.compute_root(); + intm_w.compute_root(); + + if (compile_module) { + // g calls the wrapper (not intm directly); the wrapper calls intm; and + // intm still self-references intm rather than the wrapper. + CallGraphs expected = { + {g.name(), {intm_w.name(), g.name()}}, + {intm_w.name(), {intm.name()}}, + {intm.name(), {f.name(), intm.name()}}, + {f.name(), {}}, + }; + if (check_call_graphs(g, expected) != 0) { + return 1; + } + } else { + Buffer im = g.realize({80, 80}); + auto func = [](int x, int y, int z) { + return (10 <= x && x <= 29) && (30 <= y && y <= 69) ? std::max(40 + x + y, 40) : 40; + }; + if (check_image(im, func)) { + return 1; + } + } + return 0; +} + template int reorder_split_rfactor_test() { Func f("f"), g("g"); @@ -1310,6 +1357,8 @@ int main(int argc, char **argv) { {"self assignment rfactor test", self_assignment_rfactor_test}, {"simple rfactor test: checking call graphs...", simple_rfactor_test}, {"simple rfactor test: checking output img correctness...", simple_rfactor_test}, + {"rfactor wrapper test: checking call graphs...", rfactor_wrapper_test}, + {"rfactor wrapper test: checking output img correctness...", rfactor_wrapper_test}, {"reorder split rfactor test: checking call graphs...", reorder_split_rfactor_test}, {"reorder split rfactor test: checking output img correctness...", reorder_split_rfactor_test}, {"multiple split rfactor test: checking call graphs...", multi_split_rfactor_test}, From 8d8a3f45bf34a3e4e3e37e7658a2a6b88d13d0ef Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 5 Aug 2026 11:35:02 -0700 Subject: [PATCH 6/7] Pass follow_global_wrappers in the IRGraph C++ printer Call::make's Function overload no longer defaults follow_global_wrappers, so the printer must supply it. Emit op->func.follow_global_wrappers so the reconstructed call matches the signature and round-trips faithfully. Co-Authored-By: Claude Opus 4.8 --- test/fuzz/IRGraphCXXPrinter.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/fuzz/IRGraphCXXPrinter.cpp b/test/fuzz/IRGraphCXXPrinter.cpp index 83a7e3739ad2..0d6ebc62faf1 100644 --- a/test/fuzz/IRGraphCXXPrinter.cpp +++ b/test/fuzz/IRGraphCXXPrinter.cpp @@ -236,7 +236,8 @@ void IRGraphCXXPrinter::visit(const Call *op) { // Variant 3: Convenience constructor for calls to other halide functions. // We wrap the FunctionPtr into a Function object to perfectly match // the expected `const Function &func` signature. - emit_node("Call", op, Internal::Function(op->func), op->args, op->value_index); + emit_node("Call", op, Internal::Function(op->func), op->args, op->value_index, + op->func.follow_global_wrappers); } else if (op->is_intrinsic()) { emit_node("Call", op, op->type, op->name, op->args, op->call_type); From a882a9a3b090bb0776e18ea9a18dbd9fa131bf55 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 6 Aug 2026 10:42:39 -0700 Subject: [PATCH 7/7] Make update_after_wrap error test call the wrapped Func The update added after f.in(g) previously did not reference f, so it would have been harmless without the freeze. Have it call f, so the test shows the actual hazard: the eager rewrite already happened, so the new update would call f directly rather than the wrapper. Co-Authored-By: Claude Opus 4.8 --- test/error/update_after_wrap.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/test/error/update_after_wrap.cpp b/test/error/update_after_wrap.cpp index 1a9a1fc27b9b..6292cc8c4115 100644 --- a/test/error/update_after_wrap.cpp +++ b/test/error/update_after_wrap.cpp @@ -10,13 +10,15 @@ int main(int argc, char **argv) { f(x, y) = x + y; g(x, y) = f(x, y); - // Wrapping f in g redirects g's calls to f eagerly, and freezes g. + // Wrapping f in g redirects g's existing calls to f to the wrapper, and + // freezes g. f.in(g); - // Adding an update to g now would silently fail to be wrapped, so it is an - // error. + // This update calls f, but the eager rewrite already happened, so it would + // call f directly rather than the wrapper -- inconsistent with g's original + // definition. Adding updates to a wrapped consumer is therefore an error. RDom r(0, 10); - g(r, r) += 1; + g(r, r) += f(r, r); printf("Success!\n"); return 0;