Eager in/clone_in - #9277
Open
abadams wants to merge 8 commits into
Open
Conversation
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 <noreply@anthropic.com>
# Conflicts: # src/Func.cpp # src/IR.cpp # test/error/CMakeLists.txt # test/fuzz/CMakeLists.txt
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
alexreinking
self-requested a review
August 4, 2026 23:18
alexreinking
added
the
release_notes
Member
|
Looks like this breaks Anderson2021. |
Member
Author
|
Yeah, something aint right. Investigating. |
Member
Author
|
Autodiff was generating call nodes directly instead of using Func's operator(), and it was doing it without the new flag that says to follow global wrappers, so it wasn't seeing a global wrapper. A Func was then compute_at its own global wrapper, but that global wrapper had no consumers, so it wasn't in the pipeline, so that Func's schedule was invalid (compute_at something not in the pipeline). Fixing. |
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 <noreply@anthropic.com>
alexreinking
approved these changes
Aug 5, 2026
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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This changes clone_in and in to just eagerly rewrite all consumers to call via a wrapper (or copy, for clone_in), instead of maintaining a map of wrappers that needs to be reasoned about on all future calls to DAG-rewriting scheduling directives (e.g. eager_inline).
The argless Func::in() rewrites all existing consumers via a sleight of hand. Existing consumers hold a reference to this Function object via a FunctionPtr. This PR makes it so that when FunctionPtrs dereference to a Function (FunctionPtr::get()), they follow global wrapper chains instead of returning the first Function found, but only if a flag on the FunctionPtr is set. This means you can change the meaning of all existing FunctionPtrs that are defined with the flag set without a Function needing to keep a list of all its callers (which turns into a memory management nightmare - what if there's a caller in another thread that's currently going out of scope?)
Func::in with args doesn't have this problem, because the consumers to rewrite are handed to it as args, so they can just be rewritten then and there.
Also adds a fuzzer to try out weird in/clone_in combos to look for oddities. I think we can probably make this a more general Func-DAG-rewriting fuzzer in future.
Fixes #3661
Breaking changes
Changes the behavior of in and clone_in to be eager, but I can't think of any existing code that this would break in practice. People treat it that way already. It's observible to the user if they introspect on the RHS of a calling Func - they'll see that the rewrite has occurred already, instead of being done later during lowering.