[IlyasMoutawwakil]

Summary

to_edge_transform_and_lower runs run_decompositions again on an already functional program, and each run unlifts it, which writes every mutated user input back with copy_(input, new_value). Refunctionalized, that is aten.copy(input, new_value). Upstream removes this copy for buffers but not for user inputs, so a static KV cache passed in as an input and updated with index_copy_ (copy(cache, index_put(cache, ...))) looks read after the update, and the reinplace pass can't rewrite index_put in place.

RemoveNoopPass now drops aten.copy(dst, index_put(dst, ...)) when the strides match (compared with statically_known_true, so dynamic dims are covered). index_put keeps its self's representation, quantization included, so copying its result back into the same tensor is a no-op. Other copies are kept: their metadata alone can't show that two tensors share quantization parameters.

With run_reinplace_pass=True, the cache is then updated in place, with no copy of it.

Making user inputs reinplaceable here exposed an ordering bug in insert_write_back_for_buffers_pass: _insert_copy put every inserted copy_ first and the other outputs after, and the output signature was renamed by position, so with one mutated input updated in place and another written back, their outputs swapped. Each copy_ now replaces the output it writes back.

Test plan

python -m pytest exir/tests/test_passes.py -k copy_removal
python -m pytest exir/tests/test_reinplace_pass.py

• test_redundant_copy_removal: the copy back of an index update is removed, also with a dynamic dim; a copy of another tensor's update, a plain copy and a casting copy are kept (exact counts).
• test_mutated_input_cache_reinplaced_after_transform_and_lower: a cache input updated with index_copy_ ends up as one index_put_ with no copy, and the program returns the right values.
• test_reinplaced_and_written_back_inputs_keep_their_outputs: two mutated inputs, one reinplaced and one written back, keep their output positions and signature targets.

Authored with Claude Code.

cc @JacobSzwejbka @angelayi @nil-is-all

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 21 Awaiting Approval

As of commit 8632dce with merge base 18e0cc0 ():

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

• Build Aarch64 Linux Native Wheels (gh)
• Build documentation (gh)
• Build Linux Native Wheels (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 CPU Backend (gh)
• Test CUDA Builds (gh)
• Test Native Engines (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (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 @IlyasMoutawwakil!

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!

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Copilot]

🟡 Changes recommended

The copy-elision checks can incorrectly remove copies between quantized tensors with different quantization parameters.

1 open finding

• Avoid replacing quantized copies without matching quantization parameters · New

What changed in this PR

Removes redundant functional copies so mutated inputs can be reinplaced after edge lowering.

Changes:

• Detects and removes matching aten.copy operations.
• Adds copy-removal and KV-cache reinplacement tests.

File	Description
exir/​passes/​remove_noop_pass.py	Adds copy elimination logic.
exir/​tests/​test_passes.py	Tests copy removal and dtype conversion.
exir/​tests/​test_reinplace_pass.py	Tests mutated KV-cache reinplacement.

🧠 Review effort: Balanced

---

Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.

[Copilot]

🟢 Approval recommended

The narrowly scoped rewrite is guarded by metadata checks and covered by focused and end-to-end tests.

0 open findings

1 resolved since last review

• Avoid replacing quantized copies without matching quantization parameters

🧠 Review effort: Balanced

---

Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.

[Copilot]

🟢 Approval recommended

The implementation is narrowly scoped and covered by structural and runtime regression tests.

0 open findings

🧠 Review effort: Balanced

[IlyasMoutawwakil]

@nil-is-all

[nil-is-all]

Thanks @IlyasMoutawwakil, could you update the branch?

[Copilot]

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

[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!

[Copilot]

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

[IlyasMoutawwakil]

@nil-is-all merged main, should be good now

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 38m 2s —— View job

---

Request Changes: the new rule matches on the index_put that produced the update, not on the redundant write-back copy, which every mutated user input gets. Its comment also names the wrong source for that copy. And making user inputs reinplaceable here newly exposes an output-ordering bug in insert_write_back_for_buffers_pass.

Detailed review

PR Review: #23583

Summary

This PR makes RemoveNoopPass drop aten.copy(dst, index_put(dst, …)), so that a mutated user input can be reinplaced after to_edge_transform_and_lower. The rewrite doesn't change any values. The problems:

• It matches on index_put rather than on the redundant copy itself.
• Its comment gives the wrong source for that copy.
• It newly exposes an output-ordering bug for programs with more than one mutated input.
• The unit test doesn't reach most of the new checks.

All findings come from reading the code. I couldn't run anything, because Python execution isn't allowed in this job.

Code Quality

• The rule matches on index_put instead of the redundant copy, and the comment gives the wrong source. exir/passes/remove_noop_pass.py:66-83.
  • Where the copy comes from. It isn't left over from decomposing index_copy. _gen_edge_manager_for_partitioners runs run_decompositions again on a program that is already functional (exir/program/_program.py:1174 without partitioners; 1199, or 1213 and 1226, with them). Each run unlifts the program through ep.module(). That step's _insert_copy_for_mutations writes every mutated input back with copy_(input, new_value), and re-functionalizing turns that into aten.copy(input, new_value), whatever op produced new_value. Upstream, _remove_unneccessary_copy_op_pass (torch/export/exported_program.py) removes this copy for buffers and parameters but not for user inputs, which is why only user inputs hit the problem. (The torch side is from my knowledge of the PyTorch source; I didn't re-check it in this job.)
  • What stays broken. Because the rule only matches index_put, these cases keep the copy (a full copy of the input on every call) and still can't be reinplaced:
    • The documented extension path: index_copy kept with EdgeCompileConfig(preserve_ops=[aten.index_copy.default]) and reinplaced with reinplace_extra_ops gives copy(x, index_copy(x, …)). The relevant code is _program.py:1172-1173 (where preserve_ops are popped from the second-pass table) and _program.py:1738-1743, with an example at exir/passes/reinplace.py:357-364.
    • Two updates to one input, copy(x, index_put(index_put(x, …), …)), which fails src.args[0] is dst.
    • Any other producer, for example the outer copy copy(x, copy(x, y)) in the PR's own PlainCopy case.
  • Suggested fix. Identify the copy through the graph signature, in _gen_edge_manager_for_partitioners before _generate_edge_program. For each USER_INPUT_MUTATION output that is aten.copy(p, v), where p is that input's placeholder and v has matching dtype, shape and stride, point the output at v. Repeat while the output is still such a copy.
  • The quantization caveat. It is the only stated reason for limiting the rule to index_put (lines 71-73), but no test or concrete case backs it. Quantization in this IR uses explicit q/dq nodes. The Vulkan pass already drops copy when src matches self in dtype and shape (backends/vulkan/_passes/remove_redundant_ops.py:37-38, 55-73).
  • Docs. However this is fixed, the comment, the test docstring (exir/tests/test_reinplace_pass.py:820-823) and the PR description should name the real source of the copy. Fix this →
• _is_static(dst) silently skips inputs with a symbolic dim. exir/passes/remove_noop_pass.py:78-79. A cache exported with a dynamic dim, such as batch, keeps the copy and is never reinplaced. The shape check on line 79 is always true, because index_put returns a tensor with self's shape. Comparing strides with statically_known_true, as _is_inplace_shape_dtype_compatible does (exir/passes/reinplace.py:237), would cover dynamic inputs. Fix this →
• The PR description is out of date.
  • The [PLEASE REMOVE] template text is still there.
  • There is no test plan, which CONTRIBUTING.md:403-404 asks for.
  • "RemoveNoopPass now drops them" no longer matches the rule after it was narrowed in "Only remove the copy back of an index update".
  • "No copy or write-back" only holds with run_reinplace_pass=True, which defaults to False (exir/capture/_config.py:125).
  • "For XNNPACK and any other flow" isn't tested (see Testing).

Testing

• test_redundant_copy_removal never reaches most of the new checks. exir/tests/test_passes.py:311-320, 337-340.
  • Checks never reached. In both negative cases no copy has an index_put as its src (it is a placeholder, or PlainCopy's inner copy), so the first condition rejects them. The src.args[0] is dst, _is_static, shape and stride checks are never exercised, and the tests still pass with all four deleted. The name CastingCopy also suggests a dtype check that the rule doesn't have.
  • Missing cases. y.copy_(x.index_put(...)), where index_put runs on a different tensor and the copy must stay. A non-dense input, to exercise the stride check.
  • Loose assertions. assertGreater(..., 0) at line 339 should be an exact count. Nothing checks that the copy existed before the pass ran, so the == 0 at line 337 would still pass if the pipeline stopped producing the copy. Fix this →
• The end-to-end test only covers the no-partitioner path. exir/tests/test_reinplace_pass.py:838-840.
  • That path runs run_decompositions once more. The EDGE_DO_NOT_DECOMP partitioner path runs it twice more (_program.py:1213, 1226), which should produce copy(x, copy(x, index_put(x, …))).
  • The rule only removes both copies because replace_all_uses_with on the inner copy rewrites the outer copy's arguments before the loop reaches it.
  • A case with a test partitioner would back up the "XNNPACK and any other flow" claim.

Backward Compatibility

• Programs with one reinplaced input and one written-back input return their mutated inputs in the wrong order.
  • How. exir/passes/remove_noop_pass.py:82 makes mutated user inputs reinplaceable under to_edge_transform_and_lower for the first time. In insert_write_back_for_buffers_pass, an output that is already updated in place is marked None (exir/passes/insert_write_back_for_buffers_pass.py:440-456) and handled as a user output (:290-294). _insert_copy then puts all copy_ outputs first and the rest after them (:358-361), and the loop that patches the output signature renames the mutation entries by position (:467-473).
  • Trigger. to_edge_transform_and_lower with run_reinplace_pass=True, and two mutated user inputs in this order: a updated by index_copy_ (now reinplaced), then b updated by add_(1).
  • Result. Before this PR both inputs got a copy_ and the outputs were (a', b', out). After it they are (b', a', out), and user_inputs_to_mutate maps each output to the other input.
  • Scope. The bug predates this PR (it can already be hit through to_edge), but this PR newly exposes it in the flow it targets.
  • Ask. Please add this case to the end-to-end test, and make _insert_copy swap each return node for its copy_ where it stands instead of putting all copies first. I found this by reading the code; I didn't run it. Fix this →

Recommendation

Request Changes

The output reordering that this PR newly exposes needs a test and a fix before it lands. The copy should also be matched through the graph signature rather than by checking for index_put. If that isn't done, the comment and PR description should at least say where the copy really comes from and which cases the rule still misses.

Review steps

• Gather context (diff, commits, checklist)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review
• Re-check cited line numbers after posting

Not done: I didn't run any tests or exports, because Python execution isn't allowed in this job. All findings come from reading the code.

[IlyasMoutawwakil]

Thanks for the review. Addressed in 97f6d52:

• Output ordering: _insert_copy now puts each copy_ at the position of the output it writes back, and the signature is renamed from each spec's own position. New test test_reinplaced_and_written_back_inputs_keep_their_outputs (one input updated with index_copy_, one with add_, under run_reinplace_pass=True) checks the outputs and their targets.
• Dynamic dims: strides are compared with statically_known_true, and the always-true shape check is gone. Covered by a dynamic-dim case.
• Where the copy comes from: the comment, the test docstring and the description now say it's the write-back that unlifting adds on each run_decompositions.
• Tests: a copy of another tensor's index_put (copy kept, reaches the src.args[0] is dst check), exact counts instead of assertGreater, and the PR description has a test plan instead of the template.

I kept the rule narrow (copy(dst, index_put(dst, ...))) rather than dropping every user-input write-back through the graph signature: that's a broader change to _gen_edge_manager_for_partitioners, and the narrow rule is what the KV-cache case needs. Happy to generalize it in a follow-up if that's preferred.

[IlyasMoutawwakil]

CI on 97f6d52, for the record:

• lintrunner: a ufmt line in the new test, fixed in da00fbd.
• test-model-vulkan-e2e, unittest-cuda (system), test-arm-cortex-m-size-test (bare_metal): the jobs stop on "Credentials could not be loaded" before running anything, a runner issue for fork PRs.
• nxp-mcuxpresso-test, unittest-nxp-neutron: ModuleNotFoundError: model_explorer.pytorch_exported_program_adater_impl, a dependency that moved under the October 14 base. These pass on main.

None of these touch the changed code. I merged main to pick up the current CI setup.

[IlyasMoutawwakil]

The failures on the latest run are all the same checkout error: fatal: repository 'https://github.com/facebook/buck2-shims-meta/' not found (and buck2-prelude) while cloning submodules around 18:27 UTC on Oct 28, before anything built. The URLs are unchanged on main, and main's runs after that pass, so it looks like a transient GitHub outage. Could someone re-run the failed jobs? I don't have rights to.

[nil-is-all]

@IlyasMoutawwakil could you update the branch with main? Recently we put up fixes to resolve CI failures due to dependencies and this might also get resolved if you rebase. Do give a ping after updating the branch.

[Copilot]

🟢 Approval recommended

No active review issues remain for the changes enabling in-place updates of mutated user inputs.

No open issues 🎉

_{Review effort: Balanced}

[IlyasMoutawwakil]

@nil-is-all done

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.