Insert write-back copy_ nodes at the earliest safe point#20744
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. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
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.
| if node.op != "call_function": | ||
| return True | ||
| if node.target is operator.getitem: | ||
| return True | ||
| 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) |
| def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool: | ||
| """ | ||
| Whether this node may mutate the value passed to it as input_node. When we | ||
| cannot tell we conservatively answer True. | ||
| """ | ||
| if node.op == "output": | ||
| return False | ||
| if node.op != "call_function": | ||
| return True | ||
| schema = getattr(node.target, "_schema", None) | ||
| if schema is None: | ||
| return True | ||
| for i, arg in enumerate(node.args): | ||
| if arg is input_node and i < len(schema.arguments): | ||
| alias_info = schema.arguments[i].alias_info | ||
| if alias_info is not None and alias_info.is_write: | ||
| return True | ||
| schema_kwargs = {a.name: a for a in schema.arguments} | ||
| for name, arg in node.kwargs.items(): | ||
| if arg is input_node and name in schema_kwargs: | ||
| alias_info = schema_kwargs[name].alias_info | ||
| if alias_info is not None and alias_info.is_write: | ||
| return True | ||
| return False |
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 @metascroy