Conversation
The insert_write_back_for_buffers pass placed every write-back copy_ at the end of the graph, arbitrarily extending the lifetime of the value being written back and wasting space in the memory plan. Now each copy_(buffer, value) is inserted at the earliest point that preserves the end-of-graph semantics: after the value is computed, after every reader of the buffer or any alias of it (they must observe the old contents), and after any mutation of the value or any alias of it (so we snapshot the final value). Aliases are found with a forward walk using schema alias_info, treating getitem, submodule calls, and schema-less targets conservatively. If the value written back by one copy may alias the buffer mutated by another, all copies fall back to the old end-of-graph placement in their original order. Fixes pytorch#7345
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20744
Note: Links to docs will display an error until the docs builds have been completed.
|
|
Hi @JPL11! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
|
This PR needs a
|
There was a problem hiding this comment.
Pull request overview
This PR updates insert_write_back_for_buffers_pass so write-back aten.copy_ nodes for mutated buffers/inputs are inserted at the earliest safe point in the FX graph (instead of always at the end), reducing live ranges and improving downstream memory planning.
Changes:
- Add schema-driven alias/mutation analysis utilities to determine the earliest safe insertion point for each write-back
copy_. - Insert write-backs earlier when independent, while falling back to end-of-graph insertion when write-backs may interfere via aliasing.
- Extend
exir/tests/test_passes.pywith new regression tests covering early insertion and “old value read” alias hazards, and update the existingtest_mutable_buffersexpected graph comment.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| exir/passes/insert_write_back_for_buffers_pass.py | Implements earliest-safe write-back insertion with conservative alias/mutation analysis and an independence fallback. |
| exir/tests/test_passes.py | Adds/updates tests to validate earlier insertion and alias-hazard ordering constraints. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
nil-is-all
added
the
module: exir
|
Pushed the ufmt formatting fix for the |
|
Triaged the six failures in the pull workflow run; none are from this diff:
For what it is worth, the write back tests and the full |
|
Thanks @JPL11, some CI failures are irrelevant. I'll trigger it again just-in-case. Could you sign the MetaCLA agreement if you haven't yet? |
|
Thanks @nil-is-all! CLA is signed already. The re-run did surface one real lint hit in my file that I had missed (flake8 B007, an unused loop variable from an enumerate I no longer needed), fixed and pushed. |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
meta-cla
Bot
added
the
CLA Signed
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Read this because #7345 has been picked up and dropped three times and this is the
first branch that actually implements it with tests. The approach is right: doing
a conservative alias walk before choosing the insertion point is the part that
makes the optimisation safe, and _may_alias_input and _mutates_input both
default to True when the schema cannot tell you, which is the correct direction to
fail in. node_order being built from enumerate(gm.graph.nodes) means the graph
is already topologically sorted, so the single forward pass in _collect_aliases
is sound rather than lucky.
The two tests are the right two. Asserting node_order[copy] == node_order[copy.args[1]] + 1
pins the optimisation, and test_mutable_buffers_write_back_after_old_value_reads
pins the safety condition that makes it non-trivial. That is more care than this
issue has had from anyone before.
One regression, and it is a crash rather than a behaviour change.
_insert_copy computes
last_placeholder = [node for node in gm.graph.nodes if node.op == "placeholder"][-1]unconditionally at line 151, before anything checks whether there is a write-back
to insert. A graph with no placeholders makes that an IndexError. Modules with
no inputs are unusual but legal and torch.export handles them, so this is
reachable:
class NoInput(torch.nn.Module):
def forward(self):
return torch.ones(3) * 2
ep = export(NoInput().eval(), args=())
insert_write_back_for_buffers_pass(ep)Run against the installed 1.4.1 pass and then against this branch's version of the
same file, same input:
main placeholder count 0 pass completes cleanly
this branch placeholder count 0 IndexError: list index out of range
insert_write_back_for_buffers_pass calls _insert_copy unconditionally at line
307, so there is no earlier return to save it. The graph in question has nothing
to write back, so today the pass is a no-op on it and afterwards it would throw.
Neither of the new tests catches this because both modules take an x. The
cheapest fix is to make the lookup tolerant rather than to guard the call site,
something like
placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
last_placeholder = placeholders[-1] if placeholders else Noneand then treating a None last_placeholder as "no placeholder floor" in
_insertion_point, where latest would start from return_node instead. A third
test with an input-free module would pin it.
One thing I could not check, so treat it as untested rather than as approval:
I did not verify the memory-planning win the issue asks for. I could not construct
a BUFFER_MUTATION output spec quickly outside your test harness, since raw
export alone did not produce one for me and your tests go through to_edge. So
I have verified the safety of the insertion logic by reading and the regression by
running, and I have taken the optimisation itself on the strength of your first
test rather than measuring a memory plan.
Environment for the runs above: executorch 1.4.1, torch 2.14.0+cpu, Python
3.12.10, Windows 11.
There was a problem hiding this comment.
🟡 Changes recommended
A critical alias-analysis issue can place a write-back before a later mutation.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
exir/passes/insert_write_back_for_buffers_pass.py:47
_may_alias_inputdiscards which input carries the schema's alias set. For example,aten.expand_asreturns an alias ofself, not itsotherinput; if the mutable buffer isother, this marks the expandedselfresult and all of its users as buffer aliases. That can move an otherwise safe write-back past unrelated work or trigger the multi-copy fallback, defeating the earliest-safe placement. Match the return alias set to the specific input, using the broad result only when that relationship is unavailable.
schema = getattr(node.target, "_schema", None)
if schema is None:
return True
return any(ret.alias_info is not None for ret in schema.returns)
- Files reviewed: 3/3 changed files
- Comments generated: 2
- Review effort level: Lite
There was a problem hiding this comment.
🔵 Needs a closer look
Two moderate findings remain regarding alias-set precision and alias-analysis performance.
Review details
Suppressed comments (2)
exir/passes/insert_write_back_for_buffers_pass.py:110
- Because
_may_alias_inputis only a boolean, this branch adds every FX input of any schema-known aliasing op to the alias set. For an op such asview_asorexpand_as, the second tensor supplies shape and does not alias the first tensor; if it is used later,_insertion_pointwaits for those unrelated users, and in multi-buffer cases this can create a false intersection that forces the whole write-back group to the end. Match return/input alias sets from the schema and reserve the all-input fallback for targets whose aliasing is actually unknown.
if _may_alias_input(node):
for arg in node.all_input_nodes:
if arg not in aliases:
aliases.add(arg)
changed = True
exir/passes/insert_write_back_for_buffers_pass.py:104
- This fixed-point walk can become quadratic on a long alias chain: because
node_orderis traversed forward, a seed near the end discovers only one predecessor per scan, so the loop rescans the whole graph for each backward step._collect_aliasesis also recomputed multiple times per mutable output (during independence checking and insertion), which can make export time grow substantially for large graphs with many views or mutable outputs. Consider a worklist over producer/user edges and caching each closure per seed.
changed = True
while changed:
changed = False
for node in node_order:
- Files reviewed: 3/3 changed files
- Comments generated: 0 new
- Review effort level: Lite
|
Addressed the two findings from the latest Copilot review in 1303796. Precision: the alias walk is now schema-driven — |
…_inputs (C901); rename shadowing loop var (F402)
There was a problem hiding this comment.
🟡 Changes recommended
Address the CI typing issue and avoid unnecessary alias analysis before approval.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
exir/passes/insert_write_back_for_buffers_pass.py:253
- When a graph has no write-back candidates, this still builds
node_orderand walks every node to construct the alias index. Because this pass runs for every exported program, models with no buffer/user-input copies now pay the full schema/alias-analysis cost without changing the graph; return early whenmutated_outputsis allNonebefore constructing these analyses.
alias_index = _AliasIndex(list(gm.graph.nodes))
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
There was a problem hiding this comment.
🟡 Changes recommended
Alias analysis must conservatively handle incomplete non-ATen schemas.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
exir/passes/insert_write_back_for_buffers_pass.py:133
- Because this accepts any schema-backed target, a custom op with a schema but no
Tensor(a!)annotation falls through as non-mutating. That can let the write-back run before a custom op mutates an alias ofreturn_node, leaving the buffer with a stale value; the existing CSE policy treats non-aten::schemas as untrusted for this reason. Treat non-ATen/unknown schemas conservatively (or otherwise reject missing mutation metadata).
schema = getattr(node.target, "_schema", None)
if schema is None:
return True
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
There was a problem hiding this comment.
🟡 Changes recommended
Aliasing write-back destinations need a conservative fallback to preserve their ordering.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
… the end-of-graph fallback
Fixes #7345
Summary
insert_write_back_for_buffers_passplaced every write-backcopy_at the end of the graph, arbitrarily extending the lifetime of the value being written back and wasting space in the memory plan.Each
copy_(buffer, value)is now inserted at the earliest point that preserves the end-of-graph semantics. The copy must come after:Aliases are found with a forward walk using schema
alias_info;getitem, submodule calls, and schema-less targets are treated conservatively as aliasing, and mutation detection usesalias_info.is_write(conservative when unknown). This matters becausereinplace_passruns before this pass, so in-place ops can be present; views are stillview_copyat this point sinceReplaceViewCopyWithViewPassruns later.If the value written back by one copy may alias the buffer mutated by another, the copies' relative order matters, so in that case all copies fall back to the old end-of-graph placement in their original order.
Test plan
Two new tests in
exir/tests/test_passes.py:test_mutable_buffers_write_back_is_inserted_early: the copy_ lands immediately after the value it writes back, before the rest of the graph.test_mutable_buffers_write_back_after_old_value_reads: regression test for the alias/old-read hazard — a read of the buffer's old value traced after the new value is computed keeps the write-back late.Existing
test_mutable_bufferspasses (its "After" graph comment updated for the new placement). Fullexir/tests/test_passes.py,test_memory_planning.py, andemit/test/test_emit.pypass locally (one pre-existing failure on clean main,test_to_out_variant_none_output, unrelated).cc @JacobSzwejbka @angelayi @metascroy