Remove no-op copies so mutated inputs are reinplaced after to_edge_transform_and_lower

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(

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,
@@ -59,6 +60,31 @@ 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:
+                # 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 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
+
             if node.target not in (
                 torch.ops.aten.to.dtype,
                 torch.ops.aten.slice_copy.Tensor,

exir/tests/test_passes.py

@@ -302,6 +302,66 @@ 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 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 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)
+                return x + 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, 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, dynamic_shapes=dynamic_shapes, 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)
+        index = torch.tensor([1])
+        self.assertEqual(count(IndexUpdate(), (x, index), "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:
         class FooWithNoSlice(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor:

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
@@ -813,6 +813,69 @@ 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.
+
+        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):
+            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_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: