From 09e3e27dc96c723b60753fd06bdf8693448e3b98 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 8 Oct 2026 11:00:11 +0200 Subject: [PATCH 1/4] Remove no-op copies so mutated inputs are reinplaced after to_edge_transform_and_lower to_edge_transform_and_lower functionalizes before decomposing, so index_copy_ into an input decomposes into index_put followed by aten.copy nodes into that input. Those copies take the input's sizes, strides and dtype, which already match, so they are no-ops, but they make the input look read after the update and keep the reinplace pass from rewriting index_put in place. RemoveNoopPass now drops them, so a static KV cache passed in as an input is updated with a single index_put_ and no copy or write-back, for XNNPACK and any other flow through to_edge_transform_and_lower. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- exir/passes/remove_noop_pass.py | 28 +++++++++++++++++++--- exir/tests/test_passes.py | 31 ++++++++++++++++++++++++ exir/tests/test_reinplace_pass.py | 40 ++++++++++++++++++++++++++++++- 3 files changed, 95 insertions(+), 4 deletions(-) diff --git a/exir/passes/remove_noop_pass.py b/exir/passes/remove_noop_pass.py index aeb226443f0..d3514a68c3e 100644 --- a/exir/passes/remove_noop_pass.py +++ b/exir/passes/remove_noop_pass.py @@ -43,6 +43,10 @@ def eliminate_dq_q( user.replace_all_uses_with(node.args[0]) # pyre-fixme[6] +def _is_static(tensor: torch.Tensor) -> bool: + return all(isinstance(dim, int) for dim in tensor.size()) + + class RemoveNoopPass(ExportPass): """ Removes noops that pass through arguments. @@ -59,6 +63,26 @@ def call(self, graph_module: GraphModule) -> PassResult: if node.op != "call_function": continue + if node.target == torch.ops.aten.copy.default and len(node.args) == 2: + # copy(self, src) takes self's sizes, strides and dtype and src's + # values, so it is src when those already match. Decomposing a + # functionalized index_copy leaves such copies between a mutated + # input and its update, which keeps the update from being + # reinplaced. + dst, src = node.args[0].meta["val"], node.args[1].meta["val"] + if ( + _is_static(dst) + and _is_static(src) + and ( + dst.shape == src.shape + and dst.stride() == src.stride() + and dst.dtype == src.dtype + and dst.device == src.device + ) + ): + node.replace_all_uses_with(node.args[1]) + continue + if node.target not in ( torch.ops.aten.to.dtype, torch.ops.aten.slice_copy.Tensor, @@ -81,9 +105,7 @@ def call(self, graph_module: GraphModule) -> PassResult: # The output may contain unbacked SymInts (e.g. from # data-dependent slicing with .item()), so we must check # both tensors before comparing shapes. - if all(isinstance(dim, int) for dim in orig_tensor.size()) and all( - isinstance(dim, int) for dim in output_tensor.size() - ): + if _is_static(orig_tensor) and _is_static(output_tensor): if orig_tensor.shape == output_tensor.shape: # If the graph is quantized, we must remove the entire pattern consisting of dq->op->q. # Otherwise, removing only the op will suffice. diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 0ebd71c1690..019e9269fa5 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -302,6 +302,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if node.op == "call_function": self.assertNotEqual(node.target, torch.ops.aten.to.dtype) + def test_redundant_copy_removal(self) -> None: + class NoopCopy(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = torch.empty_like(x) + y.copy_(x) + return y + 1 + + class CastingCopy(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = torch.empty(x.shape, dtype=torch.float64) + y.copy_(x) + return y + 1 + + def count(module: torch.nn.Module, x: torch.Tensor, op: str) -> int: + # Functionalization turns copy_ into copy, which RemoveNoopPass sees + # while to_edge_transform_and_lower decomposes. + gm = ( + to_edge_transform_and_lower(export(module, (x,), strict=True)) + .exported_program() + .graph_module + ) + return sum( + n.op == "call_function" and n.name.startswith(op) + for n in gm.graph.nodes + ) + + x = torch.randn(3, 4) + self.assertEqual(count(NoopCopy(), x, "aten_copy_default"), 0) + # A copy that converts the dtype is not a no-op. + self.assertEqual(count(CastingCopy(), x, "aten_copy_default"), 1) + def test_redundant_slice_copy_removal(self) -> None: class FooWithNoSlice(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index b453bfd120e..96500a66cb9 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -10,7 +10,7 @@ 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.capture._config import ExecutorchBackendConfig from executorch.exir.dialects._ops import ops as edge_ops from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass @@ -773,6 +773,44 @@ def forward( # program). edge.to_executorch() + def test_mutated_input_cache_reinplaced_after_transform_and_lower(self) -> None: + """A static KV cache passed in as an input and updated with + ``index_copy_`` is updated in place, with no copy of the cache. + + ``to_edge_transform_and_lower`` functionalizes before decomposing, + which turns the update into ``index_put`` followed by no-op ``copy`` + nodes into the cache. Those must be removed for the reinplace pass to + see that ``index_put`` is the last use of the cache. + """ + + 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) + + k, k_cache, pos = ( + torch.ones(1, 2, 1, 4), + torch.zeros(1, 2, 8, 4), + torch.tensor([3]), + ) + program = to_edge_transform_and_lower( + export(CacheUpdate(), (k, k_cache, pos), strict=True) + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + ep = program.exported_program() + + self.assertEqual(len(_find_nodes(ep, "index_put_")), 1) + self.assertEqual(len(_find_nodes(ep, "index_put", excludes="index_put_")), 0) + self.assertEqual(len(_find_nodes(ep, "copy")), 0) + + expected_cache = k_cache.clone().index_copy_(2, pos, k) + outputs = _load_for_executorch_from_buffer(program.buffer).forward( + (k, k_cache, pos) + ) + self.assertTrue(torch.equal(outputs[0], expected_cache)) + self.assertTrue(torch.equal(outputs[1], expected_cache.sum(-1))) + def test_chain_of_inplaceable_ops(self) -> None: """A chain of safe-to-reinplace ops gets fully rewritten in topological order. Exercises: From 1a3daa96eb939c8689b3ac6abc96f66ae5a9e70c Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 8 Oct 2026 11:48:24 +0200 Subject: [PATCH 2/4] Only remove the copy back of an index update Tensor metadata cannot show that two tensors share quantization parameters: ExportPass dequantizes legacy quantized values when it builds fake tensors, so two qint8 tensors with different scales can look identical, and removing a copy between them would skip a requantization. Remove a copy only when it writes index_put(dst, ...) back into dst. index_put keeps its self's representation, quantization included, so that copy is a no-op, and it is the one that keeps a mutated input's update from being reinplaced. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- exir/passes/remove_noop_pass.py | 29 ++++++++++++++--------------- exir/tests/test_passes.py | 26 ++++++++++++++++---------- 2 files changed, 30 insertions(+), 25 deletions(-) diff --git a/exir/passes/remove_noop_pass.py b/exir/passes/remove_noop_pass.py index d3514a68c3e..fdcb685a0ce 100644 --- a/exir/passes/remove_noop_pass.py +++ b/exir/passes/remove_noop_pass.py @@ -64,23 +64,22 @@ def call(self, graph_module: GraphModule) -> PassResult: continue if node.target == torch.ops.aten.copy.default and len(node.args) == 2: - # copy(self, src) takes self's sizes, strides and dtype and src's - # values, so it is src when those already match. Decomposing a - # functionalized index_copy leaves such copies between a mutated - # input and its update, which keeps the update from being - # reinplaced. - dst, src = node.args[0].meta["val"], node.args[1].meta["val"] + # Decomposing a functionalized index_copy leaves copy(dst, + # index_put(dst, ...)) between a mutated input and its update, + # which keeps the update from being reinplaced. index_put keeps + # its self's representation, quantization included, so copying + # its result back into that same tensor is a no-op. Copies of + # anything else are kept: their metadata alone cannot show that + # two tensors share quantization parameters. + dst, src = node.args if ( - _is_static(dst) - and _is_static(src) - and ( - dst.shape == src.shape - and dst.stride() == src.stride() - and dst.dtype == src.dtype - and dst.device == src.device - ) + src.target == torch.ops.aten.index_put.default + and src.args[0] is dst + and _is_static(dst.meta["val"]) + and dst.meta["val"].shape == src.meta["val"].shape + and dst.meta["val"].stride() == src.meta["val"].stride() ): - node.replace_all_uses_with(node.args[1]) + node.replace_all_uses_with(src) continue if node.target not in ( diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 019e9269fa5..5f9e5ff204f 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -303,11 +303,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.assertNotEqual(node.target, torch.ops.aten.to.dtype) def test_redundant_copy_removal(self) -> None: - class NoopCopy(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - y = torch.empty_like(x) - y.copy_(x) - return y + 1 + class IndexUpdate(torch.nn.Module): + def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + x.index_copy_(0, index, torch.ones(1, x.shape[1])) + return x + 1 + + class PlainCopy(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + x.copy_(y) + return x + 1 class CastingCopy(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -315,11 +319,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y.copy_(x) return y + 1 - def count(module: torch.nn.Module, x: torch.Tensor, op: str) -> int: + def count(module: torch.nn.Module, inputs: tuple, op: str) -> int: # Functionalization turns copy_ into copy, which RemoveNoopPass sees # while to_edge_transform_and_lower decomposes. gm = ( - to_edge_transform_and_lower(export(module, (x,), strict=True)) + to_edge_transform_and_lower(export(module, inputs, strict=True)) .exported_program() .graph_module ) @@ -329,9 +333,11 @@ def count(module: torch.nn.Module, x: torch.Tensor, op: str) -> int: ) x = torch.randn(3, 4) - self.assertEqual(count(NoopCopy(), x, "aten_copy_default"), 0) - # A copy that converts the dtype is not a no-op. - self.assertEqual(count(CastingCopy(), x, "aten_copy_default"), 1) + index = torch.tensor([1]) + self.assertEqual(count(IndexUpdate(), (x, index), "aten_copy_default"), 0) + # Only the copy back of an index update is known to be a no-op. + self.assertGreater(count(PlainCopy(), (x, x.clone()), "aten_copy_default"), 0) + self.assertEqual(count(CastingCopy(), (x,), "aten_copy_default"), 1) def test_redundant_slice_copy_removal(self) -> None: class FooWithNoSlice(torch.nn.Module): From 97f6d5235a75127c0116f089cd8429398e8ad17f Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 23:20:23 +0100 Subject: [PATCH 3/4] Keep mutation outputs in place, and remove the write-back copy at dynamic shapes - `_insert_copy` put every inserted copy_ first and the other outputs after, and the output signature was then renamed by position. With one mutated input updated in place and another written back, their outputs swapped. Each copy_ now replaces the output it writes back, and each mutation spec is renamed from its own position. - The copy rule compares strides with `statically_known_true` instead of requiring static shapes, so an input with a dynamic dim is reinplaced too. - The comment names where the copy comes from: the write-back of every mutated user input that unlifting adds on each run_decompositions. - Tests: a copy of another tensor's update is kept, a dynamic dim, exact copy counts, and two mutated inputs where one is reinplaced. Co-Authored-By: Claude Opus 5.5 --- .../insert_write_back_for_buffers_pass.py | 39 ++++++++----------- exir/passes/remove_noop_pass.py | 35 ++++++++++------- exir/tests/test_passes.py | 31 +++++++++++++-- exir/tests/test_reinplace_pass.py | 35 +++++++++++++++-- 4 files changed, 95 insertions(+), 45 deletions(-) diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py index bf170a48eea..e9bd6566a14 100644 --- a/exir/passes/insert_write_back_for_buffers_pass.py +++ b/exir/passes/insert_write_back_for_buffers_pass.py @@ -284,13 +284,13 @@ def _insert_copy( outputs = pytree.tree_flatten(output_node.args)[0] assert len(outputs) == len(mutated_outputs) - # Pair up the returns with the nodes they mutate. - copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = [] - user_output_nodes = [] - for return_node, mutated_node_name in zip(outputs, mutated_outputs): - # User output, leave alone + # Pair up the returns with the nodes they mutate, by output position. + copies: List[Tuple[int, torch.fx.Node, torch.fx.Node]] = [] + for position, (return_node, mutated_node_name) in enumerate( + zip(outputs, mutated_outputs) + ): + # User output, or a mutation already done in place: leave alone if mutated_node_name is None: - user_output_nodes.append(return_node) continue # Mutable buffer grab the node @@ -300,10 +300,10 @@ def _insert_copy( raise RuntimeError( f"Could not find {mutated_node_name} in either buffer or input nodes" ) - copies.append((mutated_node, return_node)) + copies.append((position, mutated_node, return_node)) - # insert the copies - buffer_output_nodes: List[torch.fx.Node] = [] + # insert the copies, each one replacing the output it writes back + new_outputs: List[torch.fx.Node] = list(outputs) # The alias analysis is only needed to place copies, so graphs with no # write-backs (the common case for models without mutable state) skip its # cost entirely. @@ -325,7 +325,7 @@ def _insert_copy( if len(copies) > 1: mutated_aliases: Set[torch.fx.Node] = set() return_aliases: Set[torch.fx.Node] = set() - for mutated_node, return_node in copies: + for _, mutated_node, return_node in copies: mutated_alias = alias_index.aliases(mutated_node) return_alias = alias_index.aliases(return_node) if ( @@ -340,7 +340,7 @@ def _insert_copy( mutated_aliases |= mutated_alias return_aliases |= return_alias - for mutated_node, return_node in copies: + for position, mutated_node, return_node in copies: if independent: insert_after = _insertion_point( mutated_node, return_node, node_order, last_placeholder, alias_index @@ -349,19 +349,16 @@ def _insert_copy( else: insertion = gm.graph.inserting_before(output_node) with insertion: - buffer_output = gm.graph.call_function( + new_outputs[position] = gm.graph.call_function( torch.ops.aten.copy_.default, (mutated_node, return_node) ) - # add output of copy to graph outputs - buffer_output_nodes.append(buffer_output) with gm.graph.inserting_before(output_node): - buffer_output_nodes.extend(user_output_nodes) # Remove old outputs - new_output = gm.graph.output(tuple(buffer_output_nodes)) + new_output = gm.graph.output(tuple(new_outputs)) output_node.replace_all_uses_with(new_output) gm.graph.erase_node(output_node) - return buffer_output_nodes + return new_outputs def _is_inplace_node(node: torch.fx.Node) -> bool: @@ -456,21 +453,19 @@ def insert_write_back_for_buffers_pass( mutated_outputs.append(None) # insert the copy ops and update the outputs - buffer_output_nodes = _insert_copy(gm, mutated_outputs, input_name_to_node) + new_outputs = _insert_copy(gm, mutated_outputs, input_name_to_node) gm.graph.lint() gm.graph.eliminate_dead_code() gm.recompile() # patch the output signature to point to the new updated outputs new_output_specs: List[OutputSpec] = [] - i = 0 - for output_spec in ep.graph_signature.output_specs: + for output_spec, new_output in zip(ep.graph_signature.output_specs, new_outputs): if output_spec.kind in ( OutputKind.BUFFER_MUTATION, OutputKind.USER_INPUT_MUTATION, ): - output_spec.arg.name = buffer_output_nodes[i].name - i += 1 + output_spec.arg.name = new_output.name new_output_specs.append(output_spec) signature = ExportGraphSignature( diff --git a/exir/passes/remove_noop_pass.py b/exir/passes/remove_noop_pass.py index fdcb685a0ce..4229a35c08f 100644 --- a/exir/passes/remove_noop_pass.py +++ b/exir/passes/remove_noop_pass.py @@ -12,6 +12,7 @@ from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult from torch.fx import GraphModule +from torch.fx.experimental.symbolic_shapes import statically_known_true _DEQUANT_OPS: Tuple[torch._ops.OpOverload] = ( torch.ops.quantized_decomposed.dequantize_per_tensor.default, @@ -43,10 +44,6 @@ def eliminate_dq_q( user.replace_all_uses_with(node.args[0]) # pyre-fixme[6] -def _is_static(tensor: torch.Tensor) -> bool: - return all(isinstance(dim, int) for dim in tensor.size()) - - class RemoveNoopPass(ExportPass): """ Removes noops that pass through arguments. @@ -64,20 +61,26 @@ def call(self, graph_module: GraphModule) -> PassResult: continue if node.target == torch.ops.aten.copy.default and len(node.args) == 2: - # Decomposing a functionalized index_copy leaves copy(dst, - # index_put(dst, ...)) between a mutated input and its update, - # which keeps the update from being reinplaced. index_put keeps - # its self's representation, quantization included, so copying - # its result back into that same tensor is a no-op. Copies of - # anything else are kept: their metadata alone cannot show that - # two tensors share quantization parameters. + # Each run_decompositions in to_edge_transform_and_lower unlifts + # the program, which writes every mutated user input back with + # copy_(input, new_value); refunctionalized, that is + # copy(input, new_value). Upstream drops it for buffers but not + # for user inputs, so it keeps an index update of an input from + # being reinplaced. index_put keeps its self's representation, + # quantization included, so copying its result back into that + # same tensor is a no-op. Other copies are kept: their metadata + # alone cannot show that two tensors share quantization + # parameters. dst, src = node.args if ( src.target == torch.ops.aten.index_put.default and src.args[0] is dst - and _is_static(dst.meta["val"]) - and dst.meta["val"].shape == src.meta["val"].shape - and dst.meta["val"].stride() == src.meta["val"].stride() + and all( + statically_known_true(dst_stride == src_stride) + for dst_stride, src_stride in zip( + dst.meta["val"].stride(), src.meta["val"].stride() + ) + ) ): node.replace_all_uses_with(src) continue @@ -104,7 +107,9 @@ def call(self, graph_module: GraphModule) -> PassResult: # The output may contain unbacked SymInts (e.g. from # data-dependent slicing with .item()), so we must check # both tensors before comparing shapes. - if _is_static(orig_tensor) and _is_static(output_tensor): + if all(isinstance(dim, int) for dim in orig_tensor.size()) and all( + isinstance(dim, int) for dim in output_tensor.size() + ): if orig_tensor.shape == output_tensor.shape: # If the graph is quantized, we must remove the entire pattern consisting of dq->op->q. # Otherwise, removing only the op will suffice. diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index d8c2d448d6b..d7514273adc 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -308,6 +308,11 @@ def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: x.index_copy_(0, index, torch.ones(1, x.shape[1])) return x + 1 + class CopyOfAnotherUpdate(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + y.copy_(x.index_put((torch.tensor([1]),), torch.ones(x.shape[1]))) + return y + 1 + class PlainCopy(torch.nn.Module): def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: x.copy_(y) @@ -319,11 +324,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y.copy_(x) return y + 1 - def count(module: torch.nn.Module, inputs: tuple, op: str) -> int: + def count( + module: torch.nn.Module, inputs: tuple, op: str, dynamic_shapes=None + ) -> int: # Functionalization turns copy_ into copy, which RemoveNoopPass sees # while to_edge_transform_and_lower decomposes. gm = ( - to_edge_transform_and_lower(export(module, inputs, strict=True)) + to_edge_transform_and_lower( + export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True) + ) .exported_program() .graph_module ) @@ -335,8 +344,22 @@ def count(module: torch.nn.Module, inputs: tuple, op: str) -> int: x = torch.randn(3, 4) index = torch.tensor([1]) self.assertEqual(count(IndexUpdate(), (x, index), "aten_copy_default"), 0) - # Only the copy back of an index update is known to be a no-op. - self.assertGreater(count(PlainCopy(), (x, x.clone()), "aten_copy_default"), 0) + rows = torch.export.Dim("rows", min=2) + self.assertEqual( + count( + IndexUpdate(), + (x, index), + "aten_copy_default", + dynamic_shapes={"x": {0: rows}, "index": None}, + ), + 0, + ) + # Only the copy of an index update back into its own input is known to be + # a no-op: the copy into y and its write-back both stay. + self.assertEqual( + count(CopyOfAnotherUpdate(), (x, x.clone()), "aten_copy_default"), 2 + ) + self.assertEqual(count(PlainCopy(), (x, x.clone()), "aten_copy_default"), 2) self.assertEqual(count(CastingCopy(), (x,), "aten_copy_default"), 1) def test_redundant_slice_copy_removal(self) -> None: diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 89c5f4c331c..373f9ffb499 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -817,10 +817,10 @@ def test_mutated_input_cache_reinplaced_after_transform_and_lower(self) -> None: """A static KV cache passed in as an input and updated with ``index_copy_`` is updated in place, with no copy of the cache. - ``to_edge_transform_and_lower`` functionalizes before decomposing, - which turns the update into ``index_put`` followed by no-op ``copy`` - nodes into the cache. Those must be removed for the reinplace pass to - see that ``index_put`` is the last use of the cache. + Each ``run_decompositions`` in ``to_edge_transform_and_lower`` writes + the mutated cache back with a ``copy`` of the ``index_put`` result into + the cache. Those copies must be removed for the reinplace pass to see + that ``index_put`` is the last use of the cache. """ class CacheUpdate(torch.nn.Module): @@ -851,6 +851,33 @@ def forward( self.assertTrue(torch.equal(outputs[0], expected_cache)) self.assertTrue(torch.equal(outputs[1], expected_cache.sum(-1))) + def test_reinplaced_and_written_back_inputs_keep_their_outputs(self) -> None: + """With one mutated input updated in place and another written back, + each mutation output stays in its position and names its own input.""" + + class TwoMutations(torch.nn.Module): + def forward( + self, a: torch.Tensor, b: torch.Tensor, pos: torch.Tensor + ) -> torch.Tensor: + a.index_copy_(0, pos, torch.ones(1, 4)) + b.add_(1) + return a.sum() + b.sum() + + a, b, pos = torch.zeros(3, 4), torch.zeros(2), torch.tensor([1]) + program = to_edge_transform_and_lower( + export(TwoMutations(), (a, b, pos), strict=True) + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + specs = program.exported_program().graph_signature.output_specs + self.assertEqual([spec.target for spec in specs[:2]], ["a", "b"]) + + expected_a = a.clone().index_copy_(0, pos, torch.ones(1, 4)) + outputs = _load_for_executorch_from_buffer(program.buffer).forward( + (a, b, pos) + ) + self.assertTrue(torch.equal(outputs[0], expected_a)) + self.assertTrue(torch.equal(outputs[1], torch.ones(2))) + self.assertTrue(torch.equal(outputs[2], expected_a.sum() + 2)) + def test_chain_of_inplaceable_ops(self) -> None: """A chain of safe-to-reinplace ops gets fully rewritten in topological order. Exercises: From da00fbd09e265f399b75846b7d4cc7782ea4c1e4 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 23:55:33 +0100 Subject: [PATCH 4/4] Format the output order test Co-Authored-By: Claude Opus 5.5 --- exir/tests/test_reinplace_pass.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 373f9ffb499..0d869952520 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -871,9 +871,7 @@ def forward( self.assertEqual([spec.target for spec in specs[:2]], ["a", "b"]) expected_a = a.clone().index_copy_(0, pos, torch.ones(1, 4)) - outputs = _load_for_executorch_from_buffer(program.buffer).forward( - (a, b, pos) - ) + outputs = _load_for_executorch_from_buffer(program.buffer).forward((a, b, pos)) self.assertTrue(torch.equal(outputs[0], expected_a)) self.assertTrue(torch.equal(outputs[1], torch.ones(2))) self.assertTrue(torch.equal(outputs[2], expected_a.sum() + 2))