MergeOptimizer (and structural dedup in FrozenFunctionGraph) collapses equivalent subgraphs into a single node. This avoids redundant computation — an identical subgraph is computed once and its result shared. But that sharing is exactly wrong for buffer management: when downstream inplace/aliasing logic needs two structurally-identical nodes to remain distinct buffers, merging them silently breaks it.
Cases where merging is wrong
-
Fresh allocations feeding inplace updates. Consider:
zeros(5)[1:].set(y)
zeros(5)[2:].set(z)
The two zeros(5) allocations are structurally identical and get merged. But each set wants its own fresh buffer to write into — merging makes both updates alias the same zeros buffer, which is incorrect once the inplace rewrites are set up.
-
Deliberate input duplication to break aliasing. We sometimes duplicate an input (e.g. via IndexedElemwise, or an explicit x.identity()) specifically to allow an inplace operation or to break aliasing between an op's inputs. If two such duplicates have the same structure, Merge can fold them back into the same variable, undoing the duplication.
These two pull in opposite directions: dedup is good for avoiding recomputation, but actively wrong once inplace operations have been arranged around distinct buffers.
Proposal
Establish a standard procedure for marking nodes as "must not be merged with structurally-equal siblings" while still allowing two whole graphs to be considered equivalent if they share the same global structure.
One concrete idea: give Alloc/zeros-style ops a pre-graph counter (a per-construction unique tag) so two same-valued allocations are not merged with each other within a single graph — yet two separate graphs built the same way remain equivalent (the counters line up positionally / are normalized when comparing graphs, not treated as free constants).
This needs to interact cleanly with:
MergeOptimizer / MergeFeatureFrozenFunctionGraph structural equality- the inplace/destroy-handler and
IndexedElemwise rewrites that rely on distinct buffers
Goal
A documented, consistent way to express "these must stay distinct buffers" that both the merge machinery and graph-equivalence checks respect, instead of ad-hoc handling per rewrite.
MergeOptimizer(and structural dedup inFrozenFunctionGraph) collapses equivalent subgraphs into a single node. This avoids redundant computation — an identical subgraph is computed once and its result shared. But that sharing is exactly wrong for buffer management: when downstream inplace/aliasing logic needs two structurally-identical nodes to remain distinct buffers, merging them silently breaks it.Cases where merging is wrong
Fresh allocations feeding inplace updates. Consider:
The two
zeros(5)allocations are structurally identical and get merged. But eachsetwants its own fresh buffer to write into — merging makes both updates alias the same zeros buffer, which is incorrect once the inplace rewrites are set up.Deliberate input duplication to break aliasing. We sometimes duplicate an input (e.g. via
IndexedElemwise, or an explicitx.identity()) specifically to allow an inplace operation or to break aliasing between an op's inputs. If two such duplicates have the same structure, Merge can fold them back into the same variable, undoing the duplication.These two pull in opposite directions: dedup is good for avoiding recomputation, but actively wrong once inplace operations have been arranged around distinct buffers.
Proposal
Establish a standard procedure for marking nodes as "must not be merged with structurally-equal siblings" while still allowing two whole graphs to be considered equivalent if they share the same global structure.
One concrete idea: give
Alloc/zeros-style ops a pre-graph counter (a per-construction unique tag) so two same-valued allocations are not merged with each other within a single graph — yet two separate graphs built the same way remain equivalent (the counters line up positionally / are normalized when comparing graphs, not treated as free constants).This needs to interact cleanly with:
MergeOptimizer/MergeFeatureFrozenFunctionGraphstructural equalityIndexedElemwiserewrites that rely on distinct buffersGoal
A documented, consistent way to express "these must stay distinct buffers" that both the merge machinery and graph-equivalence checks respect, instead of ad-hoc handling per rewrite.