From 76a26fb4c81dd57792c5352b081c5981681921b2 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 12:57:34 -0700 Subject: [PATCH 1/6] 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/6] 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/6] 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/6] 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/6] 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/6] 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);