Do not move a view's shape earlier than the nodes that compute it

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
 

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()