Insert write-back copy_ nodes at the earliest safe point

[JPL11]

Fixes #7345

Summary

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.

Each copy_(buffer, value) is now inserted at the earliest point that preserves the end-of-graph semantics. The copy must come after:

• the value itself,
• every reader of the buffer or any alias of it (they must observe the old contents) — this is the alias hazard that blocked the earlier attempt on the issue,
• every mutation of the value or any alias of it (so we snapshot the final value),
• all placeholders.

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 uses alias_info.is_write (conservative when unknown). This matters because reinplace_pass runs before this pass, so in-place ops can be present; views are still view_copy at this point since ReplaceViewCopyWithViewPass runs 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_buffers passes (its "After" graph comment updated for the new placement). Full exir/tests/test_passes.py, test_memory_planning.py, and emit/test/test_emit.py pass locally (one pre-existing failure on clean main, test_to_out_variant_none_output, unrelated).

cc @JacobSzwejbka @angelayi @nil-is-all @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20744

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit b7e0183 with merge base 89809c4 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @JPL11!

Thank you for your pull request and welcome to our community.

Action Required

In 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.

Process

In 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 CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: JPL11 / name: Jacky Li (1af26c7)

[Copilot]

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.py with new regression tests covering early insertion and “old value read” alias hazards, and update the existing test_mutable_buffers expected 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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[JPL11]

Pushed the ufmt formatting fix for the lintrunner job. The lintrunner-mypy failures are in backends/cortex_m/quantizer/pattern_matcher.py and backends/arm/_passes/arm_pass_utils.py, which this PR does not touch, so those look like they need a fix on main.

[JPL11]

Triaged the six failures in the pull workflow run; none are from this diff:

• unittest / linux, unittest-editable / linux, test-arm-backend-no-driver (test_pytest_ops_tosa): all die in environment setup building the tosa-tools serialization wheel (ERROR: Use build.verbose instead of cmake.verbose for scikit-build-core >= 0.10). Looks like a new scikit-build-core release broke that third party package today; the exir test suites never ran.
• android / build-android and unittest-nxp-neutron (2h timeout): both also failed on main's last completed pull run (29037897153), so pre existing.
• test-binary-size-linux-gcc: Fail 52168 > 48500. That measures the compiled C++ size_test binary; this PR only touches exir/*.py, so it cannot move that number. The threshold was last calibrated 2026-03-06 and something on main appears to have outgrown it.

For what it is worth, the write back tests and the full exir/tests/test_passes.py pass locally on this branch.

[nil-is-all]

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?

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[JPL11]

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. lintrunner --skip MYPY is clean locally on both changed files now. The remaining red jobs should be the external ones from the triage above (tosa-tools wheel setup, android and nxp from main, the binary size threshold).

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[ErenAta16]

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 None

and 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.

[Copilot]

🟡 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_order and 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 when mutated_outputs is all None before constructing these analyses.

    alias_index = _AliasIndex(list(gm.graph.nodes))

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🟡 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 of return_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

[Copilot]

🟡 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

[Copilot]

🔵 Needs a closer look

The alias-aware ordering and fallback behavior require final human review.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Lite

[JPL11]

@pytorchbot label "release notes: exir"

[Copilot]

Copilot review overview

🟡 Changes recommended

Reinplace annotation-only storage-sharing aliases may be omitted, allowing unsafe early write-back placement.

Review effort: Lite
Findings: 1

Open (1)

• Memory planner ignores reinplace annotation alias edges · New

[Copilot]

Copilot review overview

🔵 Needs a closer look

Include _share_alloc_with_arg_idx aliases in the analysis and add regression coverage.

Review effort: Lite
Findings: None

Resolved since last review (1)

• Memory planner ignores reinplace annotation alias edges

[Copilot]

Copilot review overview

🔵 Needs a closer look

Alias-sensitive graph ordering warrants final human review.

Review effort: Lite
Findings: None

[JPL11]

@JacobSzwejbka, you opened #7345 and offered there to review a PR for it. Could you take a look at this one when you have a moment? @nil-is-all, if Jacob is not the right reviewer for exir/passes anymore, could you point me to whoever is?

Where the PR stands:

• Each write-back copy_ is inserted at the earliest point that keeps the end-of-graph semantics: after the value, after every reader of the buffer or any alias of it, after every mutation of the value or any alias of it, and after all placeholders. If two write-backs may alias each other, all of them fall back to the old end-of-graph placement in their original order.
• Alias analysis reads alias_info from aten:: schemas only, matching the cse_pass policy, and treats getitem, submodule calls and schema-less targets as aliasing. The alias adjacency is built once per graph and closures are cached, and the pass skips the analysis when there is nothing to write back.
• The no-placeholder crash that @ErenAta16 found is fixed, with a third test for input-free modules. Four tests now cover early insertion, the old-read hazard, the aliased fallback and the input-free case.
• Copilot's latest finding, that outputs carrying a reinplace _share_alloc_with_arg_idx annotation were missing from the alias closure, is fixed in b7e0183 with a regression test. No Copilot findings remain open.
• I updated the branch from main, since it had fallen about 250 commits behind. The new CI runs are waiting on a maintainer to approve them, so a click there would get a fresh result. Two jobs are expected to stay red on any fork PR: "Build doc" fails at the AWS upload step because fork PRs do not receive the repo credentials, and the Cadence pull_request_target job refuses to check out fork code by design. Neither is a required check.

Nobody has measured the memory-plan saving end to end yet. The tests pin the placement rather than the plan size. If you want that number before landing, I can add a pybindings test that compares planned memory before and after on a model with mutable buffers.

If a smaller first landing would be easier to review, I can split the alias analysis into its own helper module.