From 5654233d7fa9200c50227a7562cb8d8c41396d95 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 6 Oct 2026 09:36:15 +0200 Subject: [PATCH 01/11] Lower index_copy into a mutated user input functionally in MLX A mutated user input is not the delegate's to own across calls, so only buffers the method mutates take the in-place INDEX_COPY pattern; index_copy into a user input lowers to a functional IndexCopyNode and the method writes the result back. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/ops.py | 28 ++++++++++++++ backends/mlx/patterns.py | 10 +++-- backends/mlx/test/test_ops.py | 72 +++++++++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 4 deletions(-) diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 3973f1c80d7..0cd379846b4 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -94,6 +94,7 @@ GreaterEqualNode, GreaterNode, IdCopyNode, + IndexCopyNode, IntOrVid, IntOrVidOrTid, ItemIntNode, @@ -2406,6 +2407,33 @@ def _roll_handler(P: MLXProgramBuilder, n: Node) -> Slot: return out +@REGISTRY.register(target=[torch.ops.aten.index_copy.default]) +def _index_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot: + """index_copy into a new tensor, leaving `dst` as it was. + + The INDEX_COPY pattern writes in place when `dst` is a buffer the method + mutates; anything else, a mutated user input included, lands here and the + caller writes the result back. + """ + args = P.args(n) + require_args(args, 4, 4, "aten.index_copy.default") + require_kwargs(P.kwargs(n), set(), "aten.index_copy.default") + dst, axis, indices, update = args + if not isinstance(axis, int): + raise ValueError(f"aten.index_copy.default requires a literal int dim, got {axis}") + out = P.make_or_get_slot(n) + P.emit( + IndexCopyNode( + dst=P.slot_to_tid(dst), + update=P.slot_to_tid(update), + indices=P.slot_to_tid(indices), + out=P.slot_to_tid(out), + axis=axis, + ) + ) + return out + + @REGISTRY.register(target=[torch.ops.aten.index.Tensor]) def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: args = P.args(n) diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 8abd8933e4f..c1d682f9a1c 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -123,6 +123,10 @@ def _int4_qparams_lower(qdata_node: Node, scale_node: Node, group_size: int) -> class IndexCopyHandler(PatternHandler): """ Pattern for index-based updates on mutable buffers. + + Only a buffer the method mutates is written in place: the delegate owns it + across calls. A mutated user input is not the delegate's to keep, so its + index_copy is left to the functional handler and written back by the caller. """ def __init__( @@ -148,10 +152,8 @@ def maybe_create( # noqa: C901 if not match_target(index_copy_node, torch.ops.aten.index_copy.default): return None - # index_copy should write to a mutable input/buffer to be an index update. - if (index_copy_node.name not in ep.graph_signature.buffers_to_mutate) and ( - index_copy_node.name not in ep.graph_signature.user_inputs_to_mutate - ): + # index_copy should write to a mutable buffer to be an index update. + if index_copy_node.name not in ep.graph_signature.buffers_to_mutate: return None # index_copy(dst, axis, indices, update) diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 6942cd7ae1b..9b71fc1fbfb 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -4750,6 +4750,78 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (indices, update) +class IndexCopyInputModel(nn.Module): + """Model that performs index_copy on a tensor passed in as an input. + + Unlike a buffer, a user input is not the delegate's to own: the MLX graph + returns the updated tensor and the method writes it back into the input. + """ + + def __init__(self, axis: int = 0): + super().__init__() + self.axis = axis + + def forward( + self, data: torch.Tensor, indices: torch.Tensor, update: torch.Tensor + ) -> torch.Tensor: + data.index_copy_(self.axis, indices, update) + return data.sum(self.axis) + + +@register_test +class IndexCopyInputTest(OpTestCase): + """Test case for index_copy into a mutated user input (e.g. a static KV + cache passed to the method), lowered to a functional IndexCopyNode.""" + + name = "index_copy_input" + rtol = 1e-5 + atol = 1e-5 + expected_node_counts = {"IndexCopyNode": 1} + + def __init__( + self, + buffer_size: int = 128, + feature_dim: int = 64, + num_indices: int = 8, + axis: int = 0, + ): + self.buffer_size = buffer_size + self.feature_dim = feature_dim + self.num_indices = num_indices + self.axis = axis + self.name = ( + f"index_copy_input_axis{axis}_{buffer_size}x{feature_dim}_idx{num_indices}" + ) + + @classmethod + def get_test_configs(cls) -> List["IndexCopyInputTest"]: + return [ + cls(buffer_size=128, feature_dim=64, num_indices=8, axis=0), + cls(buffer_size=64, feature_dim=32, num_indices=8, axis=1), + ] + + def create_model(self) -> nn.Module: + return IndexCopyInputModel(axis=self.axis) + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + shape = ( + (self.buffer_size, self.feature_dim) + if self.axis == 0 + else (self.feature_dim, self.buffer_size) + ) + update_shape = list(shape) + update_shape[self.axis] = self.num_indices + indices = torch.randperm(self.buffer_size)[: self.num_indices].to(torch.int64) + return (torch.randn(shape), indices, torch.randn(update_shape)) + + def compute_expected_outputs(self, model, test_inputs): + # The method returns the mutated input first, then the model's own + # output. Mutate a copy, so the saved inputs stay as they were fed. + data, indices, update = test_inputs + data = data.clone() + return [data, model(data, indices, update)] + + class SplitWithSizesModel(nn.Module): """Model that splits a tensor into chunks with specified sizes.""" From 70c4e3b23405b371f71ff8523d7b42e995dec666 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 8 Oct 2026 11:08:01 +0200 Subject: [PATCH 02/11] Test index_copy into a mutated input through the Python runtime The method writes the mutated input back with a portable copy_ outside the delegate, and op_test_runner links no portable kernels, so the op test could not load. Cover it in test_partitioner instead, which runs through the Python runtime, and format ops.py for ufmt. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/ops.py | 4 +- backends/mlx/test/test_ops.py | 72 --------------------------- backends/mlx/test/test_partitioner.py | 43 ++++++++++++++++ 3 files changed, 46 insertions(+), 73 deletions(-) diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 0cd379846b4..50bab2f1f4c 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -2420,7 +2420,9 @@ def _index_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot: require_kwargs(P.kwargs(n), set(), "aten.index_copy.default") dst, axis, indices, update = args if not isinstance(axis, int): - raise ValueError(f"aten.index_copy.default requires a literal int dim, got {axis}") + raise ValueError( + f"aten.index_copy.default requires a literal int dim, got {axis}" + ) out = P.make_or_get_slot(n) P.emit( IndexCopyNode( diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 9b71fc1fbfb..6942cd7ae1b 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -4750,78 +4750,6 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (indices, update) -class IndexCopyInputModel(nn.Module): - """Model that performs index_copy on a tensor passed in as an input. - - Unlike a buffer, a user input is not the delegate's to own: the MLX graph - returns the updated tensor and the method writes it back into the input. - """ - - def __init__(self, axis: int = 0): - super().__init__() - self.axis = axis - - def forward( - self, data: torch.Tensor, indices: torch.Tensor, update: torch.Tensor - ) -> torch.Tensor: - data.index_copy_(self.axis, indices, update) - return data.sum(self.axis) - - -@register_test -class IndexCopyInputTest(OpTestCase): - """Test case for index_copy into a mutated user input (e.g. a static KV - cache passed to the method), lowered to a functional IndexCopyNode.""" - - name = "index_copy_input" - rtol = 1e-5 - atol = 1e-5 - expected_node_counts = {"IndexCopyNode": 1} - - def __init__( - self, - buffer_size: int = 128, - feature_dim: int = 64, - num_indices: int = 8, - axis: int = 0, - ): - self.buffer_size = buffer_size - self.feature_dim = feature_dim - self.num_indices = num_indices - self.axis = axis - self.name = ( - f"index_copy_input_axis{axis}_{buffer_size}x{feature_dim}_idx{num_indices}" - ) - - @classmethod - def get_test_configs(cls) -> List["IndexCopyInputTest"]: - return [ - cls(buffer_size=128, feature_dim=64, num_indices=8, axis=0), - cls(buffer_size=64, feature_dim=32, num_indices=8, axis=1), - ] - - def create_model(self) -> nn.Module: - return IndexCopyInputModel(axis=self.axis) - - def create_inputs(self) -> Tuple[torch.Tensor, ...]: - shape = ( - (self.buffer_size, self.feature_dim) - if self.axis == 0 - else (self.feature_dim, self.buffer_size) - ) - update_shape = list(shape) - update_shape[self.axis] = self.num_indices - indices = torch.randperm(self.buffer_size)[: self.num_indices].to(torch.int64) - return (torch.randn(shape), indices, torch.randn(update_shape)) - - def compute_expected_outputs(self, model, test_inputs): - # The method returns the mutated input first, then the model's own - # output. Mutate a copy, so the saved inputs stay as they were fed. - data, indices, update = test_inputs - data = data.clone() - return [data, model(data, indices, update)] - - class SplitWithSizesModel(nn.Module): """Model that splits a tensor into chunks with specified sizes.""" diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index d8b7b8e485a..8cf261329a3 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -428,5 +428,48 @@ def forward(self, x): self.assertLess(err, 1e-4) +class TestMLXPartitionerMutatedUserInput(unittest.TestCase): + """index_copy_ into a user input lowers functionally and is written back.""" + + class IndexCopyInput(nn.Module): + def __init__(self, axis): + super().__init__() + self.axis = axis + + def forward(self, data, indices, update): + data.index_copy_(self.axis, indices, update) + return data.sum(self.axis) + + def test_index_copy_into_input_lowers_and_executes(self): + # The write-back of the mutated input runs outside the delegate, as a + # portable copy_, so this goes through the Python runtime rather than + # op_test_runner, which links no portable kernels. + for axis, shape, update_shape in ( + (0, (16, 8), (4, 8)), + (1, (8, 16), (8, 4)), + ): + with self.subTest(axis=axis): + model = self.IndexCopyInput(axis).eval() + inputs = ( + torch.randn(shape), + torch.tensor([1, 5, 9, 12]), + torch.randn(update_shape), + ) + expected = inputs[0].clone() + with torch.no_grad(): + expected_sum = model(expected, *inputs[1:]) + program = _lower(model, inputs) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + counts = get_mlx_node_counts(path) + method = Runtime.get().load_program(path).load_method("forward") + data, summed = method.execute(list(inputs)) + self.assertEqual(_delegate_count(program), 1) + self.assertEqual(counts.get("IndexCopyNode", 0), 1) + self.assertTrue(torch.allclose(data, expected)) + self.assertTrue(torch.allclose(summed, expected_sum, atol=1e-5)) + + if __name__ == "__main__": unittest.main() From 78f44d46191ef61f883073ccf45db8b92b972ac4 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 8 Oct 2026 11:34:53 +0200 Subject: [PATCH 03/11] Let MLX write a mutated user input in place The method writes an input mutated inside a delegate back with copy_(input, output) once the delegate returns, so a static KV cache passed in as an input is copied in full on every call. With run_reinplace_pass, a backend that sets BackendDetails.writes_mutated_inputs_in_place now gets that output planned as the input itself: reinplace_delegate_input_mutations gives the delegate output the input's TensorSpec and drops the copy, so the method passes one EValue as both and returns the input. MLX opts in. When an input is also an output, MLXBackend::execute 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 already lives in the output's memory. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/preprocess.py | 5 +++ backends/mlx/runtime/MLXBackend.cpp | 29 +++++++++++++- backends/mlx/runtime/MLXExecutor.h | 4 ++ backends/mlx/test/test_partitioner.py | 32 ++++++++++++++- exir/backend/backend_details.py | 6 +++ exir/passes/reinplace.py | 55 ++++++++++++++++++++++++++ exir/program/_program.py | 9 ++++- exir/tests/test_reinplace_pass.py | 57 ++++++++++++++++++++++++++- 8 files changed, 193 insertions(+), 4 deletions(-) diff --git a/backends/mlx/preprocess.py b/backends/mlx/preprocess.py index b2e26a4d2b2..81e89743018 100644 --- a/backends/mlx/preprocess.py +++ b/backends/mlx/preprocess.py @@ -57,6 +57,11 @@ class MLXBackend(BackendDetails): EXPECTED_MAGIC: ClassVar[bytes] = MAGIC EXPECTED_LENGTH: ClassVar[int] = HEADER_LENGTH + # MLXBackend::execute releases an input that is also an output before + # evaluating, so MLX can update it in its own memory, and only writes + # outputs once everything reading inputs is done. + writes_mutated_inputs_in_place: ClassVar[bool] = True + @staticmethod def preprocess( edge_program: ExportedProgram, diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index d4455d8090c..f4aefd76a9a 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/backends/mlx/runtime/MLXBackend.cpp @@ -32,6 +32,7 @@ #include +#include #include #include #include @@ -176,7 +177,11 @@ void write_output(array& arr, ETTensor& out) { throw std::runtime_error( "write_output: arr.data() is null after wait()"); } - std::memcpy(out.mutable_data_ptr(), src, out_nbytes); + // An output that is also an input can be computed into the input's own + // memory when MLX donates its buffer, and is then already in place. + if (src != out.const_data_ptr()) { + std::memcpy(out.mutable_data_ptr(), src, out_nbytes); + } } } // namespace @@ -472,6 +477,7 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { }; std::vector tensor_output_info; + std::vector inputs_written_in_place; size_t arg_idx = 0; auto* h = static_cast(handle); @@ -504,6 +510,13 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { ")"); } if (slot.slot_type == SlotType::TensorSlot) { + // The method passes an input it mutates in place as the output + // that carries the mutation too, as one and the same EValue. + EValue* arg = args[arg_idx]; + if (std::find(args.begin() + n_inputs, args.end(), arg) != + args.end()) { + inputs_written_in_place.push_back(Tid{slot.idx}); + } const ETTensor& tensor = args[arg_idx++]->toTensor(); Tid tid{slot.idx}; std::optional expected_meta = std::nullopt; @@ -572,6 +585,12 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { } } + // With nothing but the graph left holding an input that is written in + // place, MLX can donate the input's buffer to its update. + for (Tid tid : inputs_written_in_place) { + h->state.release_tensor(tid); + } + // Submit all output work to GPU asynchronously // async_eval encodes Metal commands and returns immediately. // The GPU will signal events on completion. @@ -581,6 +600,14 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { } // Lock released — GPU is still executing + // An output written into an input's memory must not land while work + // still reading that input is in flight. + if (!inputs_written_in_place.empty()) { + for (auto& arr : prepared_outputs) { + arr.wait(); + } + } + for (auto& info : tensor_output_info) { ETTensor& out_tensor = args[info.arg_idx]->toTensor(); diff --git a/backends/mlx/runtime/MLXExecutor.h b/backends/mlx/runtime/MLXExecutor.h index 95414ae8e5c..46f85acdb22 100644 --- a/backends/mlx/runtime/MLXExecutor.h +++ b/backends/mlx/runtime/MLXExecutor.h @@ -320,6 +320,10 @@ struct ExecutionState { return idx; } + void release_tensor(Tid id) { + tensors[tensor_index(id)] = std::nullopt; + } + void reset() { // Clear per-execution tensors (inputs, outputs, temps) // Constants and mutable buffers are not in this vector diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 8cf261329a3..271b19dd038 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -19,7 +19,12 @@ from executorch.backends.mlx.partitioner import MLXPartitioner from executorch.backends.mlx.passes import get_default_passes from executorch.backends.mlx.test.test_utils import get_mlx_node_counts -from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge, + to_edge_transform_and_lower, +) from executorch.runtime import Runtime from torch.export import export @@ -470,6 +475,31 @@ def test_index_copy_into_input_lowers_and_executes(self): self.assertTrue(torch.allclose(data, expected)) self.assertTrue(torch.allclose(summed, expected_sum, atol=1e-5)) + def test_index_copy_into_input_written_in_place_with_reinplace(self): + # With run_reinplace_pass the delegate writes the mutated input itself: + # the method has no write-back copy_, and returns the input value. + model = self.IndexCopyInput(0).eval() + cache = torch.zeros(16, 8) + program = to_edge_transform_and_lower( + export(model, (cache, torch.tensor([1]), torch.randn(1, 8)), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[0], plan.inputs[0]) + + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected = cache.clone() + for row in (3, 7, 11): + indices, update = torch.tensor([row]), torch.randn(1, 8) + expected.index_copy_(0, indices, update) + cache, summed = method.execute([cache, indices, update]) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-5)) + if __name__ == "__main__": unittest.main() diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index 67cf59e0f59..4cbb00006f1 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -58,6 +58,12 @@ class BackendDetails(ABC): """ + # Whether the runtime backend computes its outputs correctly when the method + # passes it the same tensor as one of its inputs and one of its outputs. + # With run_reinplace_pass, an input the delegate mutates is then written in + # place, instead of copied back after the delegate returns. + writes_mutated_inputs_in_place: bool = False + def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index b5da4cc605c..f33b6de666e 100644 --- a/exir/passes/reinplace.py +++ b/exir/passes/reinplace.py @@ -6,6 +6,7 @@ # pyre-strict +import operator from typing import Any, Dict, FrozenSet, Iterable, Optional, Set, Tuple, Union import torch @@ -15,6 +16,7 @@ unwrap_op_overload, ) from torch.export import ExportedProgram +from torch.export.graph_signature import ExportGraphSignature # --------------------------------------------------------------------------- @@ -530,3 +532,56 @@ def reinplace_pass( # noqa: C901 spec.arg.name = arg.name return ep + + +def reinplace_delegate_input_mutations( + graph_module: torch.fx.GraphModule, graph_signature: ExportGraphSignature +) -> None: + """Let a delegate write an input it mutates in place. + + The method writes a mutated input back with ``copy_(input, output)`` once + the delegate that computed ``output`` returns. When that delegate also + reads ``input`` and its backend sets ``writes_mutated_inputs_in_place``, + give ``output`` the input's TensorSpec, so the two are one value the + delegate writes directly, and drop the copy. + + Runs on a graph with tensor specs, before memory planning. + """ + from executorch.exir.backend.backend_details import BackendDetails + from executorch.exir.delegate import executorch_call_delegate + + backends = {cls.__name__: cls for cls in BackendDetails.__subclasses__()} + graph = graph_module.graph + position = {node: i for i, node in enumerate(graph.nodes)} + for node in list(graph.nodes): + if node.target != torch.ops.aten.copy_.default: + continue + mutated, update = node.args[0], node.args[1] + if mutated.op != "placeholder" or update.target != operator.getitem: + continue + delegate = update.args[0] + if delegate.target != executorch_call_delegate or mutated not in delegate.args: + continue + backend = backends.get( + getattr(graph_module, delegate.args[0].target).backend_id + ) + if backend is None or not backend.writes_mutated_inputs_in_place: + continue + # The input now changes while the delegate runs, not at the copy, so + # nothing in between may read it. + if any( + position[delegate] < position[user] < position[node] + for user in mutated.users + ): + continue + # The emitter takes the delegate's outputs from its own specs. + delegate_specs = list(delegate.meta["spec"]) + delegate_specs[update.args[1]] = mutated.meta["spec"] + delegate.meta["spec"] = type(delegate.meta["spec"])(delegate_specs) + update.meta["spec"] = mutated.meta["spec"] + for spec in graph_signature.output_specs: + if spec.arg.name == node.name: + spec.arg.name = update.name + node.replace_all_uses_with(update) + graph.erase_node(node) + graph_module.recompile() diff --git a/exir/program/_program.py b/exir/program/_program.py index ea9b09b847f..137ddc4010c 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -66,7 +66,11 @@ from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig from executorch.exir.passes.propagate_device_pass import PropagateDevicePass from executorch.exir.passes.quant_fusion_pass import quant_fusion_and_const_prop_pass -from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass +from executorch.exir.passes.reinplace import ( + DEFAULT_INPLACEABLE_OPS, + reinplace_delegate_input_mutations, + reinplace_pass, +) from executorch.exir.passes.remove_graph_asserts_pass import ( RemoveGraphAssertsPass, RemoveNonCoreAtenOpGraphAssertsPass, @@ -1761,6 +1765,9 @@ def to_executorch( # noqa (FLAKE8) C901 # TODO(who?) p.update_placeholder_tensor_specs(program, new_gm) + if config.run_reinplace_pass: + reinplace_delegate_input_mutations(new_gm, new_signature) + # Tag constant weights. if ( isinstance(config.external_constants, bool) diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index b453bfd120e..139df1533f8 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -10,8 +10,13 @@ from typing import List, Optional import torch -from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult +from executorch.exir.backend.canonical_partitioners.all_node_partitioner import ( + AllNodePartitioner, +) from executorch.exir.capture._config import ExecutorchBackendConfig +from executorch.exir.delegate import executorch_call_delegate from executorch.exir.dialects._ops import ops as edge_ops from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass from executorch.extension.pybindings.portable_lib import ( # @manual=//executorch/extension/pybindings:portable_lib @@ -40,6 +45,20 @@ def _find_nodes( ] +class _InPlaceDemoBackend(BackendDetails): + writes_mutated_inputs_in_place = True + + @staticmethod + def preprocess(edge_program, compile_specs) -> PreprocessResult: + return PreprocessResult(processed_bytes=b"") + + +class _CopyingDemoBackend(BackendDetails): + @staticmethod + def preprocess(edge_program, compile_specs) -> PreprocessResult: + return PreprocessResult(processed_bytes=b"") + + class TestReinplacePass(unittest.TestCase): def test_index_put_reinplace(self) -> None: """Test that index_put on a mutable buffer can be reinplaced.""" @@ -773,6 +792,42 @@ def forward( # program). edge.to_executorch() + def test_delegate_writes_mutated_input_in_place(self) -> None: + """A delegate that mutates one of its inputs writes it in place when + its backend opts in: its output shares the input's spec and the + write-back copy_ is gone. Other backends keep the copy_.""" + + class CacheUpdate(torch.nn.Module): + def forward( + self, k: torch.Tensor, k_cache: torch.Tensor, pos: torch.Tensor + ) -> torch.Tensor: + k_cache.index_copy_(2, pos, k) + return k_cache.sum(-1) + + inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) + for backend, in_place in ( + ("_InPlaceDemoBackend", True), + ("_CopyingDemoBackend", False), + ): + with self.subTest(backend=backend): + program = to_edge_transform_and_lower( + export(CacheUpdate(), inputs, strict=True), + partitioner=[AllNodePartitioner(backend, [])], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + k_cache = next(n for n in graph.nodes if n.name == "k_cache") + delegate = next( + n for n in graph.nodes if n.target == executorch_call_delegate + ) + self.assertEqual( + len(_find_nodes(program.exported_program(), "copy_")), + 0 if in_place else 1, + ) + self.assertEqual( + any(spec is k_cache.meta["spec"] for spec in delegate.meta["spec"]), + in_place, + ) + def test_chain_of_inplaceable_ops(self) -> None: """A chain of safe-to-reinplace ops gets fully rewritten in topological order. Exercises: From 5ff5435f4370bcca35294632b6f3089cbb2d2da8 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 8 Oct 2026 11:32:50 +0200 Subject: [PATCH 04/11] Reject negative indices in MLX IndexCopyNode ATen's index_copy raises on a negative index, but exec_index_copy wrapped it to dst_dim + idx. A program exported with valid example indices could then silently write the wrong row when given a negative index at runtime instead of failing the way eager PyTorch does. The runtime now throws for any index outside [0, size) with a message mirroring ATen's, and a partitioner test checks that executing with a negative index raises. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/runtime/MLXInterpreter.h | 13 +++++-------- backends/mlx/test/test_partitioner.py | 12 ++++++++++++ 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 4c68bd40022..239e206c730 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -1321,7 +1321,7 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) { const size_t uaxis = static_cast(axis); const int dst_dim = static_cast(dst.shape()[uaxis]); - // Get indices as a vector of ints, handling negative indices + // Like ATen, negative indices are rejected rather than wrapped. // Note: PyTorch uses int64 for indices, so we read as int64_t eval(indices); // Ensure indices are materialized before accessing data if (indices.dtype() != ::mlx::core::int64) { @@ -1332,15 +1332,12 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) { std::vector idx_vec(indices.size()); auto idx_data = indices.data(); for (size_t i = 0; i < indices.size(); ++i) { - int64_t idx = idx_data[i]; - if (idx < 0) { - idx += dst_dim; - } + const int64_t idx = idx_data[i]; if (idx < 0 || idx >= dst_dim) { throw std::out_of_range( - "IndexCopyNode: index " + std::to_string(idx_data[i]) + - " out of range for axis " + std::to_string(axis) + " with size " + - std::to_string(dst_dim)); + "index_copy: index " + std::to_string(idx) + + " is out of bounds for dimension " + std::to_string(axis) + + " with size " + std::to_string(dst_dim)); } if (idx > std::numeric_limits::max()) { throw std::out_of_range( diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 271b19dd038..4e7c04893b5 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -500,6 +500,18 @@ def test_index_copy_into_input_written_in_place_with_reinplace(self): self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-5)) + def test_index_copy_rejects_negative_index(self): + # ATen raises on negative indices, so the delegate must not wrap them. + model = self.IndexCopyInput(0).eval() + inputs = (torch.randn(4, 8), torch.tensor([1]), torch.randn(1, 8)) + program = _lower(model, inputs) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + method = Runtime.get().load_program(path).load_method("forward") + with self.assertRaises(RuntimeError): + method.execute([inputs[0], torch.tensor([-1]), inputs[2]]) + if __name__ == "__main__": unittest.main() From 40ee44644d053251aa9f48b5cd78293c26e3de31 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 9 Oct 2026 08:57:54 +0200 Subject: [PATCH 05/11] Test reads of a mutated input around its in-place write Cover the two orderings in-place writes must respect. In the pass, a read of the input between the delegate and its write-back keeps the copy_ and leaves the specs separate, since writing in place would move the mutation ahead of that read. In MLX, a value computed from the input before index_copy_ must still see the old contents across repeated in-place calls, which relies on MLXBackend waiting for every output before writing into the input. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/test/test_partitioner.py | 29 ++++++++++++++++ exir/tests/test_reinplace_pass.py | 49 +++++++++++++++++++++------ 2 files changed, 68 insertions(+), 10 deletions(-) diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 4e7c04893b5..0e1cfb747af 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -500,6 +500,35 @@ def test_index_copy_into_input_written_in_place_with_reinplace(self): self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-5)) + class ReadThenIndexCopy(nn.Module): + def forward(self, data, indices, update): + before = data.sum(0) + data.index_copy_(0, indices, update) + return before, data.sum(0) + + def test_read_before_index_copy_with_reinplace(self): + # A value read from the input before it is updated in place must be + # computed from the old contents, not race the write into the input. + model = self.ReadThenIndexCopy().eval() + cache = torch.randn(16, 8) + program = to_edge_transform_and_lower( + export(model, (cache, torch.tensor([1]), torch.randn(1, 8)), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[0], plan.inputs[0]) + + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected = cache.clone() + for row in (3, 7, 11): + indices, update = torch.tensor([row]), torch.randn(1, 8) + expected_before = expected.sum(0) + expected.index_copy_(0, indices, update) + cache, before, after = method.execute([cache, indices, update]) + self.assertTrue(torch.allclose(before, expected_before, atol=1e-5)) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.allclose(after, expected.sum(0), atol=1e-5)) + def test_index_copy_rejects_negative_index(self): # ATen raises on negative indices, so the delegate must not wrap them. model = self.IndexCopyInput(0).eval() diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 139df1533f8..b80ccc2d278 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -18,7 +18,11 @@ from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.delegate import executorch_call_delegate from executorch.exir.dialects._ops import ops as edge_ops -from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass +from executorch.exir.passes.reinplace import ( + DEFAULT_INPLACEABLE_OPS, + reinplace_delegate_input_mutations, + reinplace_pass, +) from executorch.extension.pybindings.portable_lib import ( # @manual=//executorch/extension/pybindings:portable_lib _load_for_executorch_from_buffer, ) @@ -59,6 +63,14 @@ def preprocess(edge_program, compile_specs) -> PreprocessResult: return PreprocessResult(processed_bytes=b"") +class _CacheUpdate(torch.nn.Module): + def forward( + self, k: torch.Tensor, k_cache: torch.Tensor, pos: torch.Tensor + ) -> torch.Tensor: + k_cache.index_copy_(2, pos, k) + return k_cache.sum(-1) + + class TestReinplacePass(unittest.TestCase): def test_index_put_reinplace(self) -> None: """Test that index_put on a mutable buffer can be reinplaced.""" @@ -796,14 +808,6 @@ def test_delegate_writes_mutated_input_in_place(self) -> None: """A delegate that mutates one of its inputs writes it in place when its backend opts in: its output shares the input's spec and the write-back copy_ is gone. Other backends keep the copy_.""" - - class CacheUpdate(torch.nn.Module): - def forward( - self, k: torch.Tensor, k_cache: torch.Tensor, pos: torch.Tensor - ) -> torch.Tensor: - k_cache.index_copy_(2, pos, k) - return k_cache.sum(-1) - inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) for backend, in_place in ( ("_InPlaceDemoBackend", True), @@ -811,7 +815,7 @@ def forward( ): with self.subTest(backend=backend): program = to_edge_transform_and_lower( - export(CacheUpdate(), inputs, strict=True), + export(_CacheUpdate(), inputs, strict=True), partitioner=[AllNodePartitioner(backend, [])], ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) graph = program.exported_program().graph @@ -828,6 +832,31 @@ def forward( in_place, ) + def test_delegate_input_mutation_kept_when_read_before_write_back( + self, + ) -> None: + """Writing the input in place moves its mutation from the copy_ to the + delegate, so anything reading the input in between would see the new + values. The copy_ then stays and the specs stay separate.""" + inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) + program = to_edge_transform_and_lower( + export(_CacheUpdate(), inputs, strict=True), + partitioner=[AllNodePartitioner("_InPlaceDemoBackend", [])], + ).to_executorch() + ep = program.exported_program() + graph = ep.graph + k_cache = next(n for n in graph.nodes if n.name == "k_cache") + delegate = next(n for n in graph.nodes if n.target == executorch_call_delegate) + with graph.inserting_after(delegate): + graph.call_function(torch.ops.aten.clone.default, (k_cache,)) + + reinplace_delegate_input_mutations(ep.graph_module, ep.graph_signature) + + self.assertEqual(len(_find_nodes(ep, "copy_")), 1) + self.assertFalse( + any(spec is k_cache.meta["spec"] for spec in delegate.meta["spec"]) + ) + def test_chain_of_inplaceable_ops(self) -> None: """A chain of safe-to-reinplace ops gets fully rewritten in topological order. Exercises: From f06ccfb8f56919d123b44f522bd9ff88b57672b8 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 9 Oct 2026 09:06:38 +0200 Subject: [PATCH 06/11] Keep the write-back when an alias of the input is read before it ReplaceViewCopyWithViewPass has already turned safe view_copy nodes into memory.view aliases of their base by the time the delegate's input mutation is written in place, so a view of the input taken before the delegate and read after it would see the update early. The guard now checks reads of every node sharing the input's storage, views of it and results planned onto its spec, not only its direct users. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- exir/passes/reinplace.py | 25 ++++++++++++-- exir/tests/test_reinplace_pass.py | 54 +++++++++++++++++++------------ 2 files changed, 57 insertions(+), 22 deletions(-) diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index f33b6de666e..1c447085270 100644 --- a/exir/passes/reinplace.py +++ b/exir/passes/reinplace.py @@ -534,6 +534,26 @@ def reinplace_pass( # noqa: C901 return ep +def _storage_aliases(node: torch.fx.Node) -> Set[torch.fx.Node]: + """``node`` and every node sharing its storage: views of it, which + ReplaceViewCopyWithViewPass has turned into ``memory.view``, and results + planned onto its TensorSpec, like those of in-place ops.""" + from executorch.exir import memory + + aliases = {node} + frontier = [node] + while frontier: + for user in frontier.pop().users: + if user in aliases: + continue + if (user.target == memory.view and user.args[0] in aliases) or ( + user.meta.get("spec") is node.meta["spec"] + ): + aliases.add(user) + frontier.append(user) + return aliases + + def reinplace_delegate_input_mutations( graph_module: torch.fx.GraphModule, graph_signature: ExportGraphSignature ) -> None: @@ -568,10 +588,11 @@ def reinplace_delegate_input_mutations( if backend is None or not backend.writes_mutated_inputs_in_place: continue # The input now changes while the delegate runs, not at the copy, so - # nothing in between may read it. + # nothing in between may read it, directly or through an alias. if any( position[delegate] < position[user] < position[node] - for user in mutated.users + for alias in _storage_aliases(mutated) + for user in alias.users ): continue # The emitter takes the delegate's outputs from its own specs. diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index b80ccc2d278..9a0a7f6e44b 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -10,7 +10,12 @@ from typing import List, Optional import torch -from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.exir import ( + EdgeCompileConfig, + memory, + to_edge, + to_edge_transform_and_lower, +) from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.canonical_partitioners.all_node_partitioner import ( AllNodePartitioner, @@ -836,26 +841,35 @@ def test_delegate_input_mutation_kept_when_read_before_write_back( self, ) -> None: """Writing the input in place moves its mutation from the copy_ to the - delegate, so anything reading the input in between would see the new - values. The copy_ then stays and the specs stay separate.""" + delegate, so anything reading the input in between, directly or through + a view taken before the delegate, would see the new values. The copy_ + then stays and the specs stay separate.""" inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) - program = to_edge_transform_and_lower( - export(_CacheUpdate(), inputs, strict=True), - partitioner=[AllNodePartitioner("_InPlaceDemoBackend", [])], - ).to_executorch() - ep = program.exported_program() - graph = ep.graph - k_cache = next(n for n in graph.nodes if n.name == "k_cache") - delegate = next(n for n in graph.nodes if n.target == executorch_call_delegate) - with graph.inserting_after(delegate): - graph.call_function(torch.ops.aten.clone.default, (k_cache,)) - - reinplace_delegate_input_mutations(ep.graph_module, ep.graph_signature) - - self.assertEqual(len(_find_nodes(ep, "copy_")), 1) - self.assertFalse( - any(spec is k_cache.meta["spec"] for spec in delegate.meta["spec"]) - ) + for through_view in (False, True): + with self.subTest(through_view=through_view): + program = to_edge_transform_and_lower( + export(_CacheUpdate(), inputs, strict=True), + partitioner=[AllNodePartitioner("_InPlaceDemoBackend", [])], + ).to_executorch() + ep = program.exported_program() + graph = ep.graph + k_cache = next(n for n in graph.nodes if n.name == "k_cache") + delegate = next( + n for n in graph.nodes if n.target == executorch_call_delegate + ) + read = k_cache + if through_view: + with graph.inserting_before(delegate): + read = graph.call_function(memory.view, (k_cache, [16, 4])) + with graph.inserting_after(delegate): + graph.call_function(torch.ops.aten.clone.default, (read,)) + + reinplace_delegate_input_mutations(ep.graph_module, ep.graph_signature) + + self.assertEqual(len(_find_nodes(ep, "copy_")), 1) + self.assertFalse( + any(spec is k_cache.meta["spec"] for spec in delegate.meta["spec"]) + ) def test_chain_of_inplaceable_ops(self) -> None: """A chain of safe-to-reinplace ops gets fully rewritten in From 4a1ff49066d367981068a7c593b19347db84fa83 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 9 Oct 2026 09:19:45 +0200 Subject: [PATCH 07/11] Accept int32 indices in MLX IndexCopyNode torch.export accepts int32 indices for index_copy, and ExecuTorch's portable path runs them through index_put, but the MLX interpreter only read int64 indices, so such a model lowered to MLX and then failed at runtime. Widen int32 indices to int64 before reading them; other dtypes are still rejected. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/runtime/MLXInterpreter.h | 19 +++++++++++-------- backends/mlx/test/test_partitioner.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 239e206c730..410c619227f 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -1321,17 +1321,20 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) { const size_t uaxis = static_cast(axis); const int dst_dim = static_cast(dst.shape()[uaxis]); - // Like ATen, negative indices are rejected rather than wrapped. - // Note: PyTorch uses int64 for indices, so we read as int64_t - eval(indices); // Ensure indices are materialized before accessing data - if (indices.dtype() != ::mlx::core::int64) { + // Like ATen, negative indices are rejected rather than wrapped. Export + // accepts int32 indices too, so they are widened rather than rejected. + const array idx_arr = indices.dtype() == ::mlx::core::int32 + ? astype(indices, ::mlx::core::int64, s) + : indices; + eval(idx_arr); // Ensure indices are materialized before accessing data + if (idx_arr.dtype() != ::mlx::core::int64) { throw std::invalid_argument( - std::string("IndexCopyNode: expected int64 indices, got ") + + std::string("IndexCopyNode: expected int64 or int32 indices, got ") + ExecutionState::dtype_str(indices.dtype())); } - std::vector idx_vec(indices.size()); - auto idx_data = indices.data(); - for (size_t i = 0; i < indices.size(); ++i) { + std::vector idx_vec(idx_arr.size()); + auto idx_data = idx_arr.data(); + for (size_t i = 0; i < idx_arr.size(); ++i) { const int64_t idx = idx_data[i]; if (idx < 0 || idx >= dst_dim) { throw std::out_of_range( diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 0e1cfb747af..353b415881a 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -529,6 +529,22 @@ def test_read_before_index_copy_with_reinplace(self): self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.allclose(after, expected.sum(0), atol=1e-5)) + def test_index_copy_accepts_int32_indices(self): + # Export accepts int32 indices for index_copy, so the delegate must too. + # Eager ATen does not, so the reference uses int64 indices. + model = self.IndexCopyInput(0).eval() + data, update = torch.randn(16, 8), torch.randn(2, 8) + indices = torch.tensor([2, 9], dtype=torch.int32) + expected = data.clone().index_copy_(0, indices.long(), update) + program = _lower(model, (data.clone(), indices, update)) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + method = Runtime.get().load_program(path).load_method("forward") + result, summed = method.execute([data.clone(), indices, update]) + self.assertTrue(torch.equal(result, expected)) + self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-5)) + def test_index_copy_rejects_negative_index(self): # ATen raises on negative indices, so the delegate must not wrap them. model = self.IndexCopyInput(0).eval() From 98035644b260bec84b591787db5be5c9b1efe374 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 9 Oct 2026 09:28:16 +0200 Subject: [PATCH 08/11] Follow aliases of aliases when guarding in-place delegate writes An in-place op on a view of the input is planned onto the view's spec, not the input's, so matching specs only against the input missed it, and a read of its result between the delegate and the write-back would see the update early. Each alias's users are now matched against that alias's own spec, so every node sharing the input's storage is found transitively. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- exir/passes/reinplace.py | 13 ++++++++----- exir/tests/test_reinplace_pass.py | 16 ++++++++++++---- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index 1c447085270..e2f620e2eda 100644 --- a/exir/passes/reinplace.py +++ b/exir/passes/reinplace.py @@ -535,19 +535,22 @@ def reinplace_pass( # noqa: C901 def _storage_aliases(node: torch.fx.Node) -> Set[torch.fx.Node]: - """``node`` and every node sharing its storage: views of it, which + """``node`` and every node sharing its storage, transitively: views, which ReplaceViewCopyWithViewPass has turned into ``memory.view``, and results - planned onto its TensorSpec, like those of in-place ops.""" + planned onto the TensorSpec of the alias they consume, like those of + in-place ops, including on views.""" from executorch.exir import memory aliases = {node} frontier = [node] while frontier: - for user in frontier.pop().users: + alias = frontier.pop() + spec = alias.meta.get("spec") + for user in alias.users: if user in aliases: continue - if (user.target == memory.view and user.args[0] in aliases) or ( - user.meta.get("spec") is node.meta["spec"] + if (user.target == memory.view and user.args[0] is alias) or ( + spec is not None and user.meta.get("spec") is spec ): aliases.add(user) frontier.append(user) diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 9a0a7f6e44b..c14fa13f40a 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -845,8 +845,8 @@ def test_delegate_input_mutation_kept_when_read_before_write_back( a view taken before the delegate, would see the new values. The copy_ then stays and the specs stay separate.""" inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) - for through_view in (False, True): - with self.subTest(through_view=through_view): + for read_through in ("input", "view", "in-place op on a view"): + with self.subTest(read_through=read_through): program = to_edge_transform_and_lower( export(_CacheUpdate(), inputs, strict=True), partitioner=[AllNodePartitioner("_InPlaceDemoBackend", [])], @@ -858,9 +858,17 @@ def test_delegate_input_mutation_kept_when_read_before_write_back( n for n in graph.nodes if n.target == executorch_call_delegate ) read = k_cache - if through_view: - with graph.inserting_before(delegate): + with graph.inserting_before(delegate): + if read_through != "input": read = graph.call_function(memory.view, (k_cache, [16, 4])) + read.meta["spec"] = object() + if read_through == "in-place op on a view": + # Planned onto the view's spec, not the input's. + update_view = graph.call_function( + torch.ops.aten.add_.Tensor, (read, read) + ) + update_view.meta["spec"] = read.meta["spec"] + read = update_view with graph.inserting_after(delegate): graph.call_function(torch.ops.aten.clone.default, (read,)) From 64e84e0676ec035fd2236d99aa504e7880054d39 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Wed, 14 Oct 2026 14:08:43 -0700 Subject: [PATCH 09/11] Read outputs that alias an in-place input before writing any output MLXBackend writes outputs one after another. An output written into an input's memory then also changes any output still reading that memory and written after it, e.g. a clone of the input taken before the update, which MLX leaves as the input's own buffer. Waiting for every output does not prevent that. When an input is written in place, copy any output whose data overlaps memory written at or before it out of the way before writing any output. test_snapshot_before_index_copy_with_reinplace covers the snapshot from the review, and a view of it, which the delegate returns after the update and so is the order that read updated data. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/runtime/MLXBackend.cpp | 59 +++++++++++++++++++++++---- backends/mlx/test/test_partitioner.py | 45 ++++++++++++++++++++ 2 files changed, 97 insertions(+), 7 deletions(-) diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index f4aefd76a9a..e4afdd4997c 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/backends/mlx/runtime/MLXBackend.cpp @@ -33,10 +33,12 @@ #include #include +#include #include #include #include #include +#include #include @@ -129,9 +131,21 @@ array prepare_output( return result; } -// Wait for a prepared output array and copy its data to an ET tensor. +// Whether [a, a + a_nbytes) and [b, b + b_nbytes) share any byte. +bool overlaps(const void* a, size_t a_nbytes, const void* b, size_t b_nbytes) { + const auto a_begin = reinterpret_cast(a); + const auto b_begin = reinterpret_cast(b); + return a_nbytes > 0 && b_nbytes > 0 && a_begin < b_begin + b_nbytes && + b_begin < a_begin + a_nbytes; +} + +// Wait for a prepared output array and copy its data to an ET tensor, or copy +// `snapshot` instead when it holds the array's data taken earlier. // The array must have been submitted via async_eval before calling this. -void write_output(array& arr, ETTensor& out) { +void write_output( + array& arr, + ETTensor& out, + const std::vector& snapshot = {}) { arr.wait(); // Resize output tensor if shape doesn't match (dynamic shapes) @@ -172,7 +186,7 @@ void write_output(array& arr, ETTensor& out) { " bytes, output has " + std::to_string(out_nbytes) + " bytes"); } - const void* src = arr.data(); + const void* src = snapshot.empty() ? arr.data() : snapshot.data(); if (!src) { throw std::runtime_error( "write_output: arr.data() is null after wait()"); @@ -600,19 +614,50 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { } // Lock released — GPU is still executing - // An output written into an input's memory must not land while work - // still reading that input is in flight. + // Per tensor output, a copy of its data when it must not be read from + // the array at the time it is written. + std::vector> snapshots(tensor_output_info.size()); if (!inputs_written_in_place.empty()) { + // An output written into an input's memory must not land while work + // still reading that input is in flight. for (auto& arr : prepared_outputs) { arr.wait(); } + // Outputs are then written one after another, so one written into an + // input's memory would also change any output still reading that + // memory and written after it, e.g. a clone of the input taken before + // the update, which MLX leaves as the input's own buffer. Copy such + // outputs out before writing any. + for (size_t i = 0; i < tensor_output_info.size(); ++i) { + const array& arr = + prepared_outputs[tensor_output_info[i].prepared_idx]; + // Outputs written up to and including this one. + for (size_t w = 0; w <= i; ++w) { + const auto& written = tensor_output_info[w]; + const array& written_arr = prepared_outputs[written.prepared_idx]; + const void* dst = + args[written.arg_idx]->toTensor().const_data_ptr(); + if (written_arr.data() != dst && + overlaps( + arr.data(), + arr.nbytes(), + dst, + written_arr.nbytes())) { + const auto* src = static_cast(arr.data()); + snapshots[i].assign(src, src + arr.nbytes()); + break; + } + } + } } - for (auto& info : tensor_output_info) { + for (size_t i = 0; i < tensor_output_info.size(); ++i) { + const auto& info = tensor_output_info[i]; ETTensor& out_tensor = args[info.arg_idx]->toTensor(); // write_output waits on arr to be ready - write_output(prepared_outputs[info.prepared_idx], out_tensor); + write_output( + prepared_outputs[info.prepared_idx], out_tensor, snapshots[i]); } h->state.reset(); // Release temp GPU buffers back to MLX cache diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 353b415881a..ccae1a3473e 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -529,6 +529,51 @@ def test_read_before_index_copy_with_reinplace(self): self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.allclose(after, expected.sum(0), atol=1e-5)) + class SnapshotBeforeMutation(nn.Module): + def forward(self, cache, indices, update): + before = cache.clone() + cache.index_copy_(0, indices, update) + return before, cache + + class ViewOfSnapshotBeforeMutation(nn.Module): + def forward(self, cache, indices, update): + before = cache.clone() + cache.index_copy_(0, indices, update) + return before.view(2, 2), cache + + def test_snapshot_before_index_copy_with_reinplace(self): + # MLX can return the clone as the input's own memory, so it must be + # read before the update is written into that memory. The view puts + # the snapshot after the update in the delegate's outputs, so the + # update is written first. + for model in ( + self.SnapshotBeforeMutation().eval(), + self.ViewOfSnapshotBeforeMutation().eval(), + ): + with self.subTest(model=type(model).__name__): + 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]) + program = to_edge_transform_and_lower( + export(model, (cache.clone(), indices, update), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[0], plan.inputs[0]) + + method = ( + Runtime.get().load_program(program.buffer).load_method("forward") + ) + expected = cache.clone() + for step in range(3): + update = torch.tensor([20.0, 40.0]) * (step + 1) + expected_before = expected.clone() + expected.index_copy_(0, indices, update) + cache, before, after = method.execute([cache, indices, update]) + self.assertTrue(torch.equal(before.flatten(), expected_before)) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.equal(after, expected)) + def test_index_copy_accepts_int32_indices(self): # Export accepts int32 indices for index_copy, so the delegate must too. # Eager ATen does not, so the reference uses int64 indices. From ec2a26577596c2a19e4db65c131bad4e07d8d2b9 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 23:22:14 +0100 Subject: [PATCH 10/11] Keep 0-d index_copy off MLX, guard aliases passed to the delegate, document the opt-in - `_index_copy_handler` leaves a 0-d index or a 0-d destination to portable, since the runtime rejects both, and checks the dim with `require_static_int`. - The in-place rewrite also skips a delegate that takes the mutated input through an alias as well. - The read-before and snapshot tests assert the write-back `copy_` is gone, so they fail if the pass doesn't run. - `writes_mutated_inputs_in_place` and `run_reinplace_pass` say what opting in requires and does. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/ops.py | 9 ++++++--- backends/mlx/test/test_partitioner.py | 8 ++++++++ exir/backend/backend_details.py | 8 ++++++-- exir/capture/_config.py | 4 +++- exir/passes/reinplace.py | 8 +++++--- 5 files changed, 28 insertions(+), 9 deletions(-) diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 8845b6e328b..edd796e0b3f 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -2419,9 +2419,12 @@ def _index_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot: require_args(args, 4, 4, "aten.index_copy.default") require_kwargs(P.kwargs(n), set(), "aten.index_copy.default") dst, axis, indices, update = args - if not isinstance(axis, int): - raise ValueError( - f"aten.index_copy.default requires a literal int dim, got {axis}" + require_static_int(axis, "dim", "aten.index_copy.default") + # The runtime takes a 1-D index into a tensor with at least one dim; ATen also + # allows 0-d ones, which stay off the delegate. + if n.args[0].meta["val"].dim() == 0 or n.args[2].meta["val"].dim() != 1: + raise NotImplementedError( + "aten.index_copy.default needs a 1-D index into a tensor of rank >= 1" ) out = P.make_or_get_slot(n) P.emit( diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index ccae1a3473e..2d0cc7381c3 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -515,6 +515,10 @@ def test_read_before_index_copy_with_reinplace(self): export(model, (cache, torch.tensor([1]), torch.randn(1, 8)), strict=False), partitioner=[MLXPartitioner()], ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) plan = program.executorch_program.execution_plan[0] self.assertEqual(plan.outputs[0], plan.inputs[0]) @@ -558,6 +562,10 @@ def test_snapshot_before_index_copy_with_reinplace(self): export(model, (cache.clone(), indices, update), strict=False), partitioner=[MLXPartitioner()], ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) plan = program.executorch_program.execution_plan[0] self.assertEqual(plan.outputs[0], plan.inputs[0]) diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index ddc937e8341..77f9a3481d9 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -95,9 +95,13 @@ class BackendDetails(ABC): """ # Whether the runtime backend computes its outputs correctly when the method - # passes it the same tensor as one of its inputs and one of its outputs. + # passes it the same EValue as one of its inputs and one of its outputs. # With run_reinplace_pass, an input the delegate mutates is then written in - # place, instead of copied back after the delegate returns. + # place, instead of copied back after the delegate returns. Opting in means + # the backend reads that input before writing the aliased output, and keeps + # the contents of any other output whose storage overlaps it. EXIR still only + # aliases a mutated placeholder that nothing reads between the delegate and + # its write-back, and that the delegate doesn't also take through an alias. writes_mutated_inputs_in_place: bool = False def __init_subclass__(cls, **kwargs): diff --git a/exir/capture/_config.py b/exir/capture/_config.py index 13ab5b3f46c..7cfa6dfefca 100644 --- a/exir/capture/_config.py +++ b/exir/capture/_config.py @@ -121,7 +121,9 @@ class ExecutorchBackendConfig: # If set to true, we run quant fusion and constant propagation passes do_quant_fusion_and_const_prop: bool = False - # Experimental: If set to true, we run a pass to reinplace ops in the graph. + # Experimental: If set to true, we run a pass to reinplace ops in the graph, + # and let a delegate whose backend sets writes_mutated_inputs_in_place + # write a mutated input in place instead of copying it back. run_reinplace_pass: bool = False # When True, memory planning partitions specs by device and runs the diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index e2f620e2eda..51465ed8477 100644 --- a/exir/passes/reinplace.py +++ b/exir/passes/reinplace.py @@ -591,12 +591,14 @@ def reinplace_delegate_input_mutations( if backend is None or not backend.writes_mutated_inputs_in_place: continue # The input now changes while the delegate runs, not at the copy, so - # nothing in between may read it, directly or through an alias. + # nothing in between may read it, directly or through an alias, and the + # delegate may not also take it through an alias. + aliases = _storage_aliases(mutated) if any( position[delegate] < position[user] < position[node] - for alias in _storage_aliases(mutated) + for alias in aliases for user in alias.users - ): + ) or any(alias is not mutated and alias in delegate.args for alias in aliases): continue # The emitter takes the delegate's outputs from its own specs. delegate_specs = list(delegate.meta["spec"]) From f4c81be77fe179a286724ab1cd913f793efa39ea Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 23:38:58 +0100 Subject: [PATCH 11/11] Test MLX writing an unplanned input and two inputs in place - An unplanned, page-aligned cache (alloc_graph_input=False), which MLX can wrap without a copy, so the update must land in the caller's own tensor. - Two caches mutated by one delegate, both aliased to their inputs. Co-Authored-By: Claude Opus 5.5 --- backends/mlx/test/test_partitioner.py | 69 +++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 2d0cc7381c3..9649769de87 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -25,6 +25,7 @@ to_edge, to_edge_transform_and_lower, ) +from executorch.exir.passes import MemoryPlanningPass from executorch.runtime import Runtime from torch.export import export @@ -533,6 +534,74 @@ def test_read_before_index_copy_with_reinplace(self): self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.allclose(after, expected.sum(0), atol=1e-5)) + def test_unplanned_input_written_in_place_with_reinplace(self): + # An unplanned input is passed to MLX without a copy when its memory can + # be wrapped, so the update lands in the caller's own tensor. A large + # cache is page-aligned, which wrapping without a copy may require. + model = self.IndexCopyInput(0).eval() + cache = torch.zeros(1024, 1024) + program = to_edge_transform_and_lower( + export( + model, + (cache.clone(), torch.tensor([1]), torch.randn(1, 1024)), + strict=False, + ), + partitioner=[MLXPartitioner()], + ).to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + ) + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected = cache.clone() + for row in (3, 700): + indices, update = torch.tensor([row]), torch.randn(1, 1024) + expected.index_copy_(0, indices, update) + _, summed = method.execute([cache, indices, update]) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-3)) + + class TwoCaches(nn.Module): + def forward(self, k_cache, v_cache, indices, k, v): + k_cache.index_copy_(0, indices, k) + v_cache.index_copy_(0, indices, v) + return k_cache.sum(0) + v_cache.sum(0) + + def test_two_inputs_written_in_place_with_reinplace(self): + # Two caches mutated by one delegate are both written in place. + model = self.TwoCaches().eval() + k_cache, v_cache = torch.zeros(16, 8), torch.zeros(16, 8) + example = ( + k_cache, + v_cache, + torch.tensor([1]), + torch.randn(1, 8), + torch.randn(1, 8), + ) + program = to_edge_transform_and_lower( + export(model, tuple(t.clone() for t in example), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[:2], plan.inputs[:2]) + + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected_k, expected_v = k_cache.clone(), v_cache.clone() + for row in (3, 7): + indices, k, v = torch.tensor([row]), torch.randn(1, 8), torch.randn(1, 8) + expected_k.index_copy_(0, indices, k) + expected_v.index_copy_(0, indices, v) + k_cache, v_cache, summed = method.execute([k_cache, v_cache, indices, k, v]) + self.assertTrue(torch.equal(k_cache, expected_k)) + self.assertTrue(torch.equal(v_cache, expected_v)) + expected = expected_k.sum(0) + expected_v.sum(0) + self.assertTrue(torch.allclose(summed, expected, atol=1e-5)) + class SnapshotBeforeMutation(nn.Module): def forward(self, cache, indices, update): before = cache.clone()