diff --git a/backends/transforms/fuse_view_copy.py b/backends/transforms/fuse_view_copy.py index 009019389c4..8b5d04ac827 100644 --- a/backends/transforms/fuse_view_copy.py +++ b/backends/transforms/fuse_view_copy.py @@ -90,6 +90,26 @@ def _find_view_copy_chain( return view_nodes + @staticmethod + def _shape_is_available_at( + shape: object, + node: torch.fx.Node, + order: dict[torch.fx.Node, int], + ) -> bool: + """Whether every element of ``shape`` is already defined before ``node``. + + A view's shape is not always a list of ints: under dynamic shapes a + dimension is a node computing a symbolic size. Giving ``node`` the shape + of a later view moves that shape's arguments backwards in the graph, so + it is only valid when they are available at ``node`` to begin with. + """ + if not isinstance(shape, (list, tuple)): + return True + node_pos = order[node] + return all( + order[dim] < node_pos for dim in shape if isinstance(dim, torch.fx.Node) + ) + def _merge_view_copy_chains( self, graph: torch.fx.Graph ) -> tuple[torch.fx.Graph, bool]: @@ -110,9 +130,20 @@ def _merge_view_copy_chains( """ modified = False ops: list[EdgeOpOverload] = self.UNARY_ELEMENTWISE_OPS + [self.VIEW_OP] + # Nothing below inserts or moves nodes, so a single ordering snapshot + # stays valid for the whole sweep. + order = {n: i for i, n in enumerate(graph.nodes)} for node in graph.find_nodes(op="call_function", target=self.VIEW_OP): view_nodes_to_remove = self._find_view_copy_chain(node, ops) + # Fuse the longest prefix of the chain whose shape `node` can + # actually take on; a later view whose shape is computed after + # `node` has to stay where it is. + while view_nodes_to_remove and not self._shape_is_available_at( + view_nodes_to_remove[-1].args[1], node, order + ): + view_nodes_to_remove.pop() + if len(view_nodes_to_remove) > 0: modified = True diff --git a/backends/transforms/test/test_fuse_view_copy.py b/backends/transforms/test/test_fuse_view_copy.py new file mode 100644 index 00000000000..4cac83ccc3d --- /dev/null +++ b/backends/transforms/test/test_fuse_view_copy.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class TestFuseViewCopyTransform(unittest.TestCase): + def _fuse(self, model, example_inputs, dynamic_shapes=None): + ep = torch.export.export( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=False + ) + program = to_edge(ep).exported_program() + gm = FuseViewCopyTransform()(program.graph_module).graph_module + # lint() is what catches a node whose argument is defined after it. + gm.graph.lint() + return gm.graph + + @staticmethod + def _count_views(graph: torch.fx.Graph) -> int: + return sum( + 1 + for n in graph.nodes + if n.op == "call_function" + and n.target == exir_ops.edge.aten.view_copy.default + ) + + def test_static_view_chain_is_fused(self): + class Model(torch.nn.Module): + def forward(self, x): + return x.view(4, 8).relu().view(32).sqrt().view(2, 16) + + graph = self._fuse(Model().eval(), (torch.rand(2, 16) + 1.0,)) + self.assertEqual(self._count_views(graph), 1) + + def test_chain_ending_in_a_later_computed_shape_stays_ordered(self): + # The final view's shape is only known after `y` has been reduced, so + # the first view cannot take that shape on: it runs earlier. + class Model(torch.nn.Module): + def forward(self, x): + n = x.shape[0] + y = x.view(n * 4, 8).relu() + return y.view(n * 2, 16) + float(0) + + dim = torch.export.Dim("n", min=2, max=64) + graph = self._fuse( + Model().eval(), + (torch.rand(8, 32),), + dynamic_shapes={"x": {0: dim}}, + ) + # No assertion on the view count: the point is that lint() above passes, + # i.e. the pass never leaves an argument used before it is defined. + + +if __name__ == "__main__": + unittest.main()