From 68fe421979c5819dd5c4806ef2175d9b0d038001 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 1 Sep 2026 12:27:59 +0200 Subject: [PATCH] Do not move a view's shape earlier than the nodes that compute it `_merge_view_copy_chains` gives the first view in a chain the shape of the last one: final_shape = view_nodes_to_remove[-1].args[1] node.args = (node.args[0], final_shape) Under dynamic shapes a view's shape is not a list of ints, it is a list that can contain nodes computing symbolic sizes. Those nodes may be defined anywhere between the first and the last view, so assigning the later shape to the earlier node can make that node read a value that does not exist yet. The result is a graph that fails fx's own `lint()`, which the pass then hits itself on the very next line via `eliminate_dead_code()`: RuntimeError: Argument 'mul' of Node 'aten_view_copy_default_21' was used before it has been defined! Please check that Nodes in the graph are topologically ordered seen while lowering a TTS model to Vulkan, where the chain ends on a `view_copy(x, [2*s, 512])` whose `2*s` is produced one node after the first view. Fuse the longest prefix of the chain whose shape is available where the first view sits, and leave the rest alone. Chains with static shapes, which is the common case, are unaffected. --- backends/transforms/fuse_view_copy.py | 31 +++++++++ .../transforms/test/test_fuse_view_copy.py | 63 +++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 backends/transforms/test/test_fuse_view_copy.py 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()