Make inner graph Ops immutable#2243
Open
ricardoV94 wants to merge 5 commits into
Open
Conversation
ricardoV94
force-pushed
the
real_frozen_graphs
branch
5 times, most recently
from
fe15202 to
726ee00
Compare
11 tasks
ricardoV94
force-pushed
the
real_frozen_graphs
branch
4 times, most recently
from
ea00cf7 to
67d4867
Compare
Introduce FrozenFunctionGraph and FrozenApply: an immutable, hash-consable representation of an inner graph, plus an AbstractApply base shared by the mutable Apply and the immutable FrozenApply. FunctionGraph.freeze(dedup_nodes=...) builds one, baking a destroy-aware toposort by forwarding the source graph's orderings (a DestroyHandler, if present, makes the order safe for a backend to funcify as-is; a regular graph stays the plain toposort at no cost). Apply.clone_with_new_inputs rebuilds immutable nodes instead of mutating in place. Composite and ScalarLoop adopt the frozen representation for their inner graphs. The clone_inner_graph(s) machinery is retained here and removed once every inner-graph op is immutable. Claude-Session: https://claude.ai/code/session_0134bV4HGae6Pk5xw1yUuf2z
Scan and OpFromGraph carry a single frozen inner graph (op.fgraph) and are optimized
by a graph rewrite instead of at link time. Each inner-graph op has its own rewrite
that dispatches the inner-graph inplace contract on the linker via
functools.singledispatch with a raising base, so every backend registers an
implementation explicitly (scan/rewriting/{c,numba,jit}.py, tensor/rewriting/{numba,
ofg}.py). numba owns its Scan memory model (+ potentially_overwritten_reads) and the
boundary deepcopies in scan/rewriting/numba.py; numba funcify just funcifies op.fgraph
and reads codegen metadata off the op (no graph mutation). The shared orchestration
lives in compile/aliasing.py::rewrite_inner_graph; the rewrites register at position
49.6 tagged minimum_compile + compile_inner_graph.
The rewrite owns unfreeze -> optimize -> freeze (the optimized graph still carries a
DestroyHandler, so its freeze is destroy-aware) and the ops store the frozen graph
verbatim: Scan.clone_with_inner_graph is copy(self) + inner.freeze() (no rebuild);
OpFromGraph rebuilds via construct_nominal_fgraph and re-attaches a DestroyHandler
gated on inplace.
Claude-Session: https://claude.ai/code/session_0134bV4HGae6Pk5xw1yUuf2z
ScipyWrapperOp (Minimize / Root) carries a frozen inner graph and registers a per-linker optimize_inner_graph rewrite (noinplace, via the shared compile/aliasing.py::rewrite_inner_graph helper), like Scan and OpFromGraph. Claude-Session: https://claude.ai/code/session_0134bV4HGae6Pk5xw1yUuf2z
Cloning an outer graph used to optionally deep-clone the inner graphs of HasInnerGraph ops (clone_inner_graph(s)=True). Now that every inner-graph op (Scan/OpFromGraph/Composite/ScipyWrapperOp) is immutable -- clone() returns self -- that deep-clone is always a no-op, so drop the kwarg and the branch throughout the clone machinery (Apply.clone, Apply.clone_with_new_inputs, FrozenApply.clone, clone, clone_node_and_cache, clone_get_equiv, FunctionGraph.clone, rebuild_collect_shared) and clone_node_and_cache's op-clone caching. Claude-Session: https://claude.ai/code/session_0134bV4HGae6Pk5xw1yUuf2z
ricardoV94
force-pushed
the
real_frozen_graphs
branch
from
67d4867 to
f77b580
Compare
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
1 participant
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.
Closes #2232
Closes #2028
Closes #2033
Closes #2194 by adding a method (should it be the default) where theres's no node deduplication, as the toposort index becomes part of the hash
No more inner graph "cloning" of inner graphs