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 aeb226443f0..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, @@ -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, diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 81c5298aab1..547f9ecdf68 100644 --- a/exir/tests/test_passes.py +++ b/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: diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 52675fedb59..0d869952520 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 @@ -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: