[IlyasMoutawwakil]

Summary

A model that mutates one of its inputs, like a static KV cache passed in by the caller and updated with index_copy_, previously failed to lower on MLX: the INDEX_COPY pattern pointed the update at the input's slot after the builder had already given it an output slot. This PR lowers such updates, and lets MLX write them straight into the input instead of copying the cache back on every call.

Lowering: index_copy into anything other than a mutated buffer now goes through a functional aten.index_copy.default handler that emits an IndexCopyNode. Like ATen, IndexCopyNode now rejects negative indices instead of wrapping them.

In place: once a delegate returns, the method writes a mutated input back with copy_(input, output), which copies the whole cache on every call. BackendDetails gets a writes_mutated_inputs_in_place opt-in. With run_reinplace_pass=True, reinplace_delegate_input_mutations gives such a delegate output the input's TensorSpec and drops the copy_, so the method passes one EValue as both input and output and returns the input. MLX opts in. MLXBackend::execute detects an input that is also an output, releases the state's reference to it once the graph is built so MLX can donate the wrapped input buffer to the update, waits for every output before writing any, and skips the copy when the result is already in the output's memory. Backends that don't opt in are unchanged.

Combined with #23183, which passes an unplanned caller tensor to the method without copying and returns it uncloned, a caller-owned static cache is then updated in place in one MLX graph, with no write-back.

Suggested review order: exir/backend/backend_details.py and reinplace_delegate_input_mutations in exir/passes/reinplace.py (with its call in to_executorch), then backends/mlx/runtime/MLXBackend.cpp, then the MLX lowering in backends/mlx/ops.py.

Test plan

test_delegate_writes_mutated_input_in_place in exir/tests/test_reinplace_pass.py lowers a cache update to a demo backend that opts in and to one that doesn't. It checks that only the opted-in delegate gets the input's spec for its output and loses the copy_. TestMLXPartitionerMutatedUserInput in backends/mlx/test/test_partitioner.py covers index_copy_ into an input along both axes, rejection of a negative index, and the in-place path with run_reinplace_pass=True over several calls: no copy_ in the program, the method output is the input value, and the values match eager. These run through the Python runtime rather than op_test_runner, because without the in-place path the write-back copy_ needs portable kernels that op_test_runner doesn't link.

python -m pytest exir/tests/test_reinplace_pass.py -k delegate
python -m pytest backends/mlx/test/test_partitioner.py -k "MutatedUserInput"

The MLX runtime part of these tests needs Apple silicon and runs in test-mlx.

This PR was authored with Claude Code.

🤖 Generated with Claude Code

cc @nil-is-all @metascroy

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 17 Awaiting Approval

As of commit f628e6c with merge base 0549e1e ():

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

• Apple (gh)
• 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 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]

Copilot review overview

🟢 Approval recommended

The targeted lowering change is consistent with MLX slot ownership and is covered across both tested axes.

Review effort: Balanced
Findings: None

What changed in this PR

This PR makes MLX lower index_copy on mutated user inputs functionally while preserving in-place lowering for mutable buffers.

Changes:

• Restricts the in-place pattern to mutable buffers.
• Adds functional aten.index_copy.default lowering.
• Tests mutated inputs along axes 0 and 1.

File	Description
backends/​mlx/​patterns.py	Limits in-place index copy to buffers.
backends/​mlx/​ops.py	Adds functional index-copy lowering.
backends/​mlx/​test/​test_ops.py	Adds mutation and output regression coverage.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[nil-is-all]

Hi @IlyasMoutawwakil, could you fix the failing MLX CI test and lintrunner? Do ping us when ready for review.

[Copilot]

🟡 Changes recommended

The generic handler silently wraps negative indices although ATen rejects them.

1 open finding

• Reject negative indices in aten.index_copy handler · New

🧠 Review effort: Balanced

---

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

[Copilot]

🔵 Needs a closer look

It introduces cross-cutting tensor aliasing and asynchronous GPU buffer-donation behavior requiring final human validation.

2 open findings

• Reject negative indices in aten.index_copy handler
• Update test plan to reference new index_copy tests · New

🧠 Review effort: Balanced

---

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

[Copilot]

🔵 Needs a closer look

Cross-layer aliasing and asynchronous buffer donation require the missing ordering and safety-fallback tests before final human approval.

2 open findings

• In-place test misses read-before-index_copy ordering hazard · New
• Missing test for intervening reads before delegated write-back · New

2 resolved since last review

• Reject negative indices in aten.index_copy handler
• Update test plan to reference new index_copy tests

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

Cross-layer aliasing and asynchronous MLX buffer-donation semantics require final human validation despite the added tests.

0 open findings

🧠 Review effort: Balanced

[IlyasMoutawwakil]

@nil-is-all

[nil-is-all]

@IlyasMoutawwakil seems the branch went out-of-date over the weekend. Could you update & give us a ping?

[Copilot]

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

[metascroy]

@IlyasMoutawwakil did you run a model through this change? What was the perf like?

I think there is a correctness issue that can show up. The new aliased-output handling can overwrite a mutated input’s storage before copying another output that still references its pre-mutation contents, causing that output to return updated data. Waiting for all MLX outputs to finish does not prevent this sequential write-back hazard; output sources overlapping mutation destinations must be preserved before those destinations are overwritten.

Consider adding a regression test like this:

class SnapshotBeforeMutation(torch.nn.Module):
    def forward(self, cache, indices, update):
        before = cache.clone()
        cache.index_copy_(0, indices, update)
        return before, cache

cache = torch.tensor([1.0, 2.0, 3.0, 4.0])
indices = torch.tensor([1, 3], dtype=torch.int64)
update = torch.tensor([20.0, 40.0])

[metascroy]

Could we express this as a partitioner annotation on the specific input mutation, analogous to how we tag mutated buffers, rather than a backend-wide writes_mutated_inputs_in_place flag? The annotation would declare that this partition supports aliasing the mutation output with its corresponding input; EXIR would still perform its reinplacement safety checks, establish the alias, and remove the external write-back only when safe, retaining the existing path otherwise.

Unlike buffer delegation, this would not transfer ownership: the user input remains an external argument. This gives us per-mutation opt-in, with the annotation preserved through lowering, instead of requiring every partition produced by a backend to support the behavior.

Maybe @JacobSzwejbka has thoughts here on how we should handle user input mutation inside delegates as well.

[IlyasMoutawwakil]

makes sense, a per-mutation annotation is more precise than a backend-wide flag and keeps the safety checks in EXIR. I'd like to hear what @JacobSzwejbka thinks about delegates mutating user inputs in general before reworking it, then I'll update the PR to match.

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

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[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. the workflows are waiting for approval though, could someone approve them? the MLX job is the one that actually runs the new tests on apple silicon.

@metascroy good catch on the write back order, fixed in 64e84e0. when an input is written in place, MLXBackend now snapshots any output whose data overlaps memory written at or before it, before writing anything back, so an output that's still the input's own buffer gets read before the update lands. added test_snapshot_before_index_copy_with_reinplace with your SnapshotBeforeMutation module plus a variant that returns before.view(2, 2). funnily enough your exact example doesn't hit it, the delegate returns the snapshot before the cache so it's written first, the view variant flips that order and does.

on perf, I haven't run a full model through it on apple silicon yet so no numbers so far. the change just drops the copy back into the mutated input, so the gain should scale with whatever is written in place (the kv cache in practice).

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 34m 8s —— View job

---

PR Review: #23508

• Understand the PR context and diff (base 73c65da3)
• Run three independent reviewers
• Consolidate and deduplicate findings
• Fact-check findings
• Post review

Request Changes: with run_reinplace_pass, this PR has a delegate write a mutated input in place instead of copying it back (MLX opts in), but older MLX runtimes can run the resulting programs and return wrong results with no error, the new index_copy handler delegates inputs the runtime always rejects, and two of the three in-place tests still pass if the pass doesn't run.

Full review

Summary

The PR adds BackendDetails.writes_mutated_inputs_in_place and a to_executorch pass. When run_reinplace_pass is set, the pass gives a delegate output the TensorSpec of the input it updates and removes the write-back copy_. MLX opts in, lowers index_copy functionally, and execute now handles an input that is also an output. Problems found:

• Old MLX runtimes can run the new programs and return wrong results, and nothing in the payload lets them reject the programs.
• The backend contract the flag creates is not written down.
• The new handler delegates index_copy nodes that the runtime always rejects.
• Two of the three in-place tests don't show that the pass ran.

Code Quality

• 0-d index or 0-d dst (backends/mlx/ops.py:2411). _index_copy_handler accepts every functional aten.index_copy. ATen allows a 0-d index and a 0-d dst, but the runtime rejects both:

  • a 0-d index throws "indices must be 1D" in exec_index_copy (MLXInterpreter.h:1317);
  • a 0-d dst throws in normalize_axis (:1320).

  The partitioner treats a node as supported unless its handler raises. So y = x.index_copy(0, torch.tensor(2), src), which ran on portable before this PR, is now delegated and fails on every execute. Raise in the handler, based on the nodes' meta["val"], so these nodes stay on portable. While there, ops.py:2422 should use require_static_int(axis, "dim", "aten.index_copy.default") (ops.py:190) instead of a hand-written ValueError, like the other handlers in the file. Fix this →

• Alias passed to the delegate itself (exir/passes/reinplace.py:595-599). The guard only rejects readers strictly between the delegate and the copy_. It misses the delegate itself reading the input through a second argument.

  • Trigger: an alias of mutated created outside the delegate is also a delegate argument. Examples are a memory.view from a view op the partitioner didn't take, or an in-place result that shares the input's spec.
  • Effect: the delegate reads that alias while it writes the input in place. The MLX runtime recognizes the aliasing by EValue* identity. It could donate the input's buffer and write it while the second argument, wrapped separately over the same memory, still needs to be read.
  • Confidence: the gap in the guard is confirmed. The MLX race is suspected, not reproduced.
  • Fix: skip the rewrite when any node of _storage_aliases(mutated) other than mutated itself is in delegate.args.

  Fix this →

Testing

• Two tests don't show the pass ran (backends/mlx/test/test_partitioner.py:509 and :544). test_read_before_index_copy_with_reinplace and test_snapshot_before_index_copy_with_reinplace never check that the write-back copy_ is gone.

  • Their plan.outputs[0] == plan.inputs[0] assertion (:519, :562) holds without the pass too. The copy_ shares the input's TensorSpec, and the emitter assigns values per spec (spec2id_dict, exir/emit/_emitter.py:671-688).
  • So if the guard rejects these graphs, both tests still pass on the old copy-back path, which is exactly the path they are meant to rule out.
  • Fix: assert there is no copy_, as :487-490 does.

  Fix this →

• The runtime's new code may not be exercised. It's unclear whether any test reaches the release, wait-all and snapshot code that execute now has. That code only matters when MLX wraps the input's memory without copying it.

  • The test inputs are memory-planned (alloc_graph_input defaults to True) and small; the snapshot test uses 4 floats.
  • Zero-copy wrapping may require page alignment. This is low confidence, because the MLX source isn't in this checkout.

  Can you confirm that the snapshot test fails with the snapshot block from 64e84e0 reverted? If it doesn't, add a case with unplanned inputs (MemoryPlanningPass(alloc_graph_input=False)) and a page-aligned caller tensor, and assert that the caller's tensor holds the result after execute.

• No test mutates two inputs through one delegate (for example separate K and V caches). Nothing covers several aliased outputs in one execute, including the snapshot ordering between them.

API Design

• The flag's contract is undocumented (exir/backend/backend_details.py:101). writes_mutated_inputs_in_place is a contract any backend can opt into. Its comment only says the backend "computes its outputs correctly" when one tensor is both an input and an output. An implementer needs the requirements that the MLX runtime turned out to need:

  • Aliasing is by identity: the backend receives the same EValue* as an input argument and as an output argument.
  • The backend must finish reading the input before it writes the aliased output.
  • It must preserve any other output whose storage overlaps the input. This is what 64e84e0 fixes.
  • It must be told which inputs can be aliased. The pass accepts any placeholder (reinplace.py:583), so mutable buffers are included, but the PR only describes user inputs.

  The user-facing side and the runtime comment are out of date too:

  • run_reinplace_pass is still documented as "a pass to reinplace ops in the graph" (exir/capture/_config.py:124-125), but it now also changes how delegates are called.
  • Setting only reinplace_extra_ops runs reinplace_pass (exir/program/_program.py:1742) but not the new pass (:1771).
  • The SAFETY comment in tensor_to_mlx (backends/mlx/runtime/MLXBackend.cpp:108-110) still says MLX only reads the wrapped memory. That stops being true once MLX can donate the input's buffer.

Backward Compatibility

• Old MLX runtimes return wrong results for the new programs.
  • After this PR, a program exported with run_reinplace_pass=True passes the same EValue as a delegate input and output.
  • An MLX runtime from before this PR has none of the new handling in execute: it doesn't release the input, wait for all outputs, or snapshot overlapping outputs.
  • When MLX wraps the input without copying, that runtime can write the updated input over memory that another output still reads. The read-before and snapshot tests cover exactly these cases. The runtime then returns wrong results with no error.
• The payload can't prevent this.
  • The alias is created in to_executorch, after preprocess has already serialized the payload.
  • So the alias can't be recorded in the payload or gated on its version. kMaxSupportedVersion (MLXBackend.cpp:348) is unchanged.
• Fix: decide the alias at partition time, as metascroy suggested. The payload could then record it and bump the schema version so older runtimes reject the program. If not, state explicitly that forward compatibility isn't promised here.

Performance

• Question (low confidence, since the MLX source isn't in this checkout): does MLX actually donate the input's buffer here?

  • The input is wrapped with a no-op deleter (MLXBackend.cpp:116-117).
  • If MLX won't donate external buffers, or copies inputs it can't wrap without copying, the result goes into a new buffer. write_output then still copies the whole input back on every call (MLXBackend.cpp:196-197).
  • In that case the PR removes the portable copy_ but not the full-size write-back.

  Checking whether that memcpy is skipped for a cache of realistic size would answer this and metascroy's performance question.

Recommendation

Request Changes

What blocks merging:

• _index_copy_handler must reject the inputs the runtime can't run.
• The in-place tests must fail when the pass doesn't run.
• It needs confirming that the tests reach the aliased runtime path.
• Forward compatibility needs a decision, ideally by deciding the alias at partition time so the payload can record it.

The contract for writes_mutated_inputs_in_place and the run_reinplace_pass documentation should land together with the flag.

[IlyasMoutawwakil]

Thanks for the review. Addressed in the latest commit:

• 0-d index / 0-d dst: _index_copy_handler now raises for both, so those nodes stay on portable, and the dim goes through require_static_int.
• Alias passed to the delegate: the rewrite is skipped when the delegate also takes the mutated input through an alias.
• Tests that pass without the pass: test_read_before_index_copy_with_reinplace and test_snapshot_before_index_copy_with_reinplace now assert the write-back copy_ is gone.
• Runtime coverage: new test_unplanned_input_written_in_place_with_reinplace (an unplanned, page-aligned cache exported with alloc_graph_input=False, so MLX can wrap it without a copy; the update must land in the caller's tensor) and test_two_inputs_written_in_place_with_reinplace (K and V caches mutated by one delegate, both aliased).
• Contract: the writes_mutated_inputs_in_place comment now says what opting in requires (read the input before writing the aliased output, keep any overlapping output's contents) and which aliases EXIR refuses, and run_reinplace_pass mentions the delegate behavior.

Forward compatibility with older MLX runtimes is tied to @metascroy's suggestion to express this as a per-mutation partitioner annotation instead of a backend-wide flag: deciding the alias at partition time is what lets the payload record it and bump its version. I'm waiting on @JacobSzwejbka's view on delegates mutating user inputs before reworking it that way, so I've left the flag as is for now.

exir/tests/test_reinplace_pass.py passes locally. The MLX tests, the two new ones included, need Apple silicon, so test-mlx is the first place they run. Whether MLX then skips the write-back copy for a realistic cache (your perf question) also needs measuring on a Mac.

[IlyasMoutawwakil]

CI on f4c81be, for the record:

• test-pybindings-no-aten: ImportError: cannot import name 'ModuleAddBFloat16' from ...make_test, the CI script expecting a test helper newer than this branch's October 14 base.
• nxp-mcuxpresso-test, unittest-nxp-neutron: ModuleNotFoundError: model_explorer.pytorch_exported_program_adater_impl, a dependency that moved under that base.
• test-arm-cortex-m-size-test (bare_metal): "Credentials could not be loaded", a runner issue for fork PRs.

These pass on main and don't touch the changed code. I merged main to pick up the current CI setup.