Repository navigation
Conversation
🔗 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! |
|
|
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
Unresolved critical and moderate alias-analysis issues remain, and the dependent-copy fallback lacks coverage.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (2)
exir/passes/insert_write_back_for_buffers_pass.py:41
_may_alias_inputis independent of which input reaches this node, so any aliased return is treated as an alias of every input. For example,aten.expand_as(x, state)returns a view ofx, notstate; whenstateis the seed, later users of that view are incorrectly counted as state readers. This can push an otherwise safe write-back past unrelated work (and can also trigger the multi-copy end-of-graph fallback), defeating the pass's earliest-placement goal. Match the return alias set to the specific input alias set, while retaining the conservative fallback 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)
exir/passes/insert_write_back_for_buffers_pass.py:199
- The new
independent=Falsefallback is not covered by the tests: the existing two-copy case uses distinct write-back sources and destinations, so this safety branch is never exercised. Add a case where one write-back value may alias another destination (for example, one mutable buffer is copied from another) and assert both the preserved copy order and final outputs.
independent = True
if len(copies) > 1:
mutated_aliases: Set[torch.fx.Node] = set()
return_aliases: Set[torch.fx.Node] = set()
for mutated_node, return_node in copies:
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Lite
|
Thank you for the comprehensive review. As you recommended, I have fixed the problem in b6457f4. The process of lookup by placeholders has become tolerant, as the program finds placeholders[-1] or returns None in case placeholders is empty. In addition, I have added the last_placeholder parameter to the _insertion_point function, which allows us to treat None as the absence of the placeholder floor. Finally, I added the third test case that you requested, which involves an input-free module. First, I executed the code using the old implementation to generate IndexError and, after that, I confirmed that the new implementation works fine the same way in torch 2.14.0 / executorch 1.4.1 / Python 3.12, as seen in your reports. In addition to this, the commit also deals with two Copilot feedbacks issued in July, which were indeed relevant issues: _may_alias_input now returns False for the output node (since there is no value generated, the function might trigger end-of-graph fallback by mistake), and _mutates_input also inspects the list/tuple arguments which assures that the foreach-style operation mutating the tensor contained in the list is detected, thus covering the one missing element of the contract that states “answer conservatively True.” The portion that you were unable to review states that merely exporting raw data does not guarantee you will receive BUFFER_MUTATION output specification since the buffer has to be lifted and mutated. The simplest way of doing that is to call to_edge shortly after calling export with the parameters module and arguments, and the program will be exported. Please note that the memory-planning win itself can be proven by the assertion node_order[copy] == node_order[copy.args[1]] + 1 (which states that write-back occurs immediately after the value recorded), but I fully agree that measuring the full memory-usage plan would serve as more reliable proof and I would be happy to add it if you are interested. |
64aadf1 to
290e226
Compare
290e226 to
c834270
Compare
c834270 to
57c0388
Compare
57c0388 to
82a8901
Compare
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
…_inputs (C901); rename shadowing loop var (F402)
… the end-of-graph fallback
reinplace_pass marks an annotation-only rewrite with _share_alloc_with_arg_idx and the memory planner then places that output on the argument's storage. The alias index only consulted schemas, so such an output never joined its argument's alias closure and a reader of it could be ordered after the write-back copy_. _storage_sharing_inputs reads the annotation, the same edge ReplaceViewCopyWithViewPass follows in _alias_source. _aliasing_inputs adds it to the schema-derived set, and _mutates_input treats the shared argument as written, since the output lands in its storage. The schema walk in _mutates_input moves to _schema_writes_input to keep the function under the complexity limit. Regression test: a maximum of the buffer carrying the annotation, read after the buffer's new value is computed; the write-back must follow that read.
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 @nil-is-all @metascroy