Arm backend: Avoid repeated retracing in view-permute propagation (pytorch#23055)

Refresh fake-tensor metadata only for nodes crossed by each view or
permute rewrite. This allows dependent rewrites to continue in the same
scan instead of retracing the full graph after every rewrite wave.

Visit nodes in reverse graph order during downward propagation so later
transforms move first instead of blocking earlier transforms. Keep one
final full retrace to validate and reconstruct the graph.

On the SmolLM2 VGF export, avoiding repeated retracing reduced AOT
compilation from 2h08m to 37m23s in the original measurement.

Add coverage for metadata refresh, retrace scheduling, and reverse
propagation order.

AI-Assisted-by: Codex

Signed-off-by: Xingguo Li <xingguo.li@arm.com>

Author: xingguo01

backends/arm/test/passes/test_propagate_permutes_views_pass.py

@@ -7,6 +7,7 @@
 from typing import Tuple
 
 import pytest
+import sympy  # type: ignore[import-untyped]
 import torch
 from executorch.backends.arm._passes import (
     MoveDataMovementOpsToSmallerDtypePass,
@@ -23,6 +24,8 @@
 )
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
 input_t = Tuple[torch.Tensor]
 
@@ -334,6 +337,134 @@ def _run_pass_on_graph_module(
     return result.graph_module
 
 
+class _CountingDownPass(PropagateViewCopyPermuteDownPass):
+    def __init__(self) -> None:
+        super().__init__()
+        self.retrace_count = 0
+        self.local_refresh_count = 0
+
+    def _refresh_propagation_meta(self, moving_node, propagation_path):
+        self.local_refresh_count += 1
+        return super()._refresh_propagation_meta(moving_node, propagation_path)
+
+    def _retrace(self, graph_module):
+        self.retrace_count += 1
+        return super()._retrace(graph_module)
+
+
+def test_propagation_retraces_once_after_reaching_fixed_point() -> None:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.empty((1, 2, 3, 4))
+    permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1]))
+    permute.meta["val"] = torch.empty((1, 3, 4, 2))
+    relu = graph.call_function(RELU, args=(permute,))
+    relu.meta["val"] = torch.empty((1, 3, 4, 2))
+    neg = graph.call_function(NEG, args=(relu,))
+    neg.meta["val"] = torch.empty((1, 3, 4, 2))
+    graph.output(neg)
+    graph_module = torch.fx.GraphModule(torch.nn.Module(), graph)
+    propagation_pass = _CountingDownPass()
+
+    result = propagation_pass.call(graph_module)
+
+    assert result.modified
+    assert propagation_pass.retrace_count == 1
+    assert propagation_pass.local_refresh_count == 1
+
+
+@pytest.mark.parametrize("target", [VIEW, PERMUTE], ids=["view", "permute"])
+def test_local_refresh_preserves_symbolic_view_permute_meta(target) -> None:
+    shape_env = ShapeEnv()
+    batch = shape_env.create_symintnode(sympy.Symbol("batch"), hint=2)
+    assert isinstance(batch, torch.SymInt)
+    shape_env.constrain_symbol_range(batch.node.expr, compiler_min=1, compiler_max=8)
+
+    with FakeTensorMode(shape_env=shape_env, allow_non_fake_inputs=True):
+        graph = torch.fx.Graph()
+        x = graph.placeholder("x")
+        x.meta["val"] = torch.empty((batch, 2, 3))
+        if target == VIEW:
+            transform = graph.call_function(target, args=(x, [batch, 6]))
+            symbolic_axis = 0
+        else:
+            transform = graph.call_function(target, args=(x, [1, 0, 2]))
+            symbolic_axis = 1
+        transform.meta["val"] = torch.empty((2, 6))
+        graph.output(transform)
+
+        PropagateViewCopyPermuteDownPass()._refresh_node_meta(transform)
+
+    output_dim = transform.meta["val"].shape[symbolic_axis]
+    assert isinstance(output_dim, torch.SymInt)
+    assert output_dim.node.expr == batch.node.expr
+
+
+def test_down_propagation_visits_nodes_in_reverse_order() -> None:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.empty((1, 2, 3, 4))
+    view = graph.call_function(VIEW, args=(x, [1, 2, 3, 4]))
+    view.meta["val"] = torch.empty((1, 2, 3, 4))
+    relu = graph.call_function(RELU, args=(view,))
+    relu.meta["val"] = torch.empty((1, 2, 3, 4))
+    permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1]))
+    permute.meta["val"] = torch.empty((1, 3, 4, 2))
+    neg = graph.call_function(NEG, args=(permute,))
+    neg.meta["val"] = torch.empty((1, 3, 4, 2))
+    graph.output(neg)
+    graph_module = torch.fx.GraphModule(torch.nn.Module(), graph)
+    propagation_pass = PropagateViewCopyPermuteDownPass()
+
+    assert propagation_pass._propagate_iteration(graph_module)
+
+    targets = [
+        node.target for node in graph_module.graph.nodes if node.op == "call_function"
+    ]
+    assert targets == [RELU, NEG, VIEW, PERMUTE]
+    assert not propagation_pass._propagate_iteration(graph_module)
+
+
+def test_down_propagation_refreshes_metadata_in_dependency_order() -> None:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.empty((1, 2, 3, 4))
+    permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1]))
+    permute.meta["val"] = torch.empty((1, 3, 4, 2))
+    relu = graph.call_function(RELU, args=(permute,))
+    relu.meta["val"] = torch.empty((1, 3, 4, 2))
+    neg = graph.call_function(NEG, args=(relu,))
+    neg.meta["val"] = torch.empty((1, 3, 4, 2))
+    graph.output(neg)
+    torch.fx.GraphModule(torch.nn.Module(), graph)
+
+    assert PropagateViewCopyPermuteDownPass()._propagate(permute)
+
+    assert relu.meta["val"].shape == torch.Size((1, 2, 3, 4))
+    assert neg.meta["val"].shape == torch.Size((1, 2, 3, 4))
+    assert permute.meta["val"].shape == torch.Size((1, 3, 4, 2))
+
+
+def test_up_propagation_refreshes_metadata_in_dependency_order() -> None:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.empty((1, 2, 3, 4))
+    relu = graph.call_function(RELU, args=(x,))
+    relu.meta["val"] = torch.empty((1, 2, 3, 4))
+    neg = graph.call_function(NEG, args=(relu,))
+    neg.meta["val"] = torch.empty((1, 2, 3, 4))
+    permute = graph.call_function(PERMUTE, args=(neg, [0, 2, 3, 1]))
+    permute.meta["val"] = torch.empty((1, 3, 4, 2))
+    graph.output(permute)
+    torch.fx.GraphModule(torch.nn.Module(), graph)
+
+    assert PropagateViewCopyPermuteUpPass()._propagate(permute)
+
+    assert permute.meta["val"].shape == torch.Size((1, 3, 4, 2))
+    assert relu.meta["val"].shape == torch.Size((1, 3, 4, 2))
+    assert neg.meta["val"].shape == torch.Size((1, 3, 4, 2))
+
+
 def _run_pass_on_graph(
     graph: torch.fx.Graph,
     pass_cls: type[ArmPass] = PropagateViewCopyPermuteUpPass,

backends/transforms/propagate_view_copy_permute_pass.py

@@ -24,9 +24,12 @@
 from executorch.backends.transforms.fuse_identical_input_transforms_pass import (
     FuseIdenticalInputTransformsPass,
 )
+from executorch.backends.transforms.permute_view_meta import refresh_permute_view_meta
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
+from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
+from torch.fx.node import map_arg
 
 
 @dataclass(frozen=True)
@@ -113,22 +116,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             graph_module = self._retrace(graph_module)
 
         while True:
-            iteration_modified = False
-            # A rewrite invalidates metadata along its path and downstream.
-            # Defer that region while still batching independent branches.
-            stale_nodes: set[torch.fx.Node] = set()
-            for node in list(graph_module.graph.nodes):
-                if node in stale_nodes:
-                    continue
-                if node.target in self._targets:
-                    if len(node.users) == 0:
-                        continue
-                    if self._propagate(node, stale_nodes):
-                        iteration_modified = True
+            iteration_modified = self._propagate_iteration(graph_module)
 
             if iteration_modified:
                 modified = True
-                graph_module = self._retrace(graph_module)
                 continue
 
             result = self.fuse_horizontal(graph_module)
@@ -147,36 +138,69 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
         return PassResult(graph_module, modified)
 
-    @staticmethod
-    def _mark_downstream_nodes_stale(
-        node: torch.fx.Node, stale_nodes: set[torch.fx.Node]
-    ) -> None:
-        pending = [node]
-        while pending:
-            current = pending.pop()
-            if current in stale_nodes:
+    def _propagate_iteration(
+        self,
+        graph_module: torch.fx.GraphModule,
+    ) -> bool:
+        modified = False
+        for node in self._propagation_order(graph_module):
+            if node.graph is None:
                 continue
-            stale_nodes.add(current)
-            pending.extend(current.users)
+            if node.target not in self._targets or len(node.users) == 0:
+                continue
+            modified |= self._propagate(node)
+        return modified
 
-    def _mark_stale_region(
-        self,
-        nodes: Iterable[torch.fx.Node],
-        stale_nodes: set[torch.fx.Node],
-    ) -> None:
-        for node in nodes:
-            self._mark_downstream_nodes_stale(node, stale_nodes)
+    def _propagation_order(
+        self, graph_module: torch.fx.GraphModule
+    ) -> Iterable[torch.fx.Node]:
+        return list(graph_module.graph.nodes)
 
     def _retrace(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
         graph_module.graph.eliminate_dead_code()
         graph_module.graph.lint()
         return super().call(graph_module).graph_module
 
+    def _refresh_node_meta(self, node: torch.fx.Node) -> None:
+        if not node.all_input_nodes:
+            return
+        if node.target in {self._VIEW_TARGET, self._PERMUTE_TARGET}:
+            # Derive copy-op shapes from the input so existing SymInts retain
+            # their identity instead of being recreated by fake execution.
+            refresh_permute_view_meta(node)
+            return
+
+        args = map_arg(node.args, lambda input_node: input_node.meta["val"])
+        kwargs = map_arg(node.kwargs, lambda input_node: input_node.meta["val"])
+        fake_mode = next(
+            (
+                value.fake_mode
+                for input_node in node.all_input_nodes
+                if isinstance((value := input_node.meta.get("val")), FakeTensor)
+            ),
+            None,
+        )
+        if fake_mode is None:
+            fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
+            args = torch.utils._pytree.tree_map_only(
+                torch.Tensor, fake_mode.from_tensor, args
+            )
+            kwargs = torch.utils._pytree.tree_map_only(
+                torch.Tensor, fake_mode.from_tensor, kwargs
+            )
+        with fake_mode:
+            node.meta["val"] = node.target(*args, **kwargs)  # type: ignore[operator]
+
+    @abstractmethod
+    def _refresh_propagation_meta(
+        self, moving_node: torch.fx.Node, propagation_path: Sequence[torch.fx.Node]
+    ) -> None:
+        pass
+
     def _validated_next_nodes(
         self,
         node: torch.fx.Node,
         frontier: torch.fx.Node,
-        stale_nodes: set[torch.fx.Node],
     ) -> list[torch.fx.Node] | None:
         next_nodes = list(self._get_next_nodes(frontier))
         if not next_nodes:
@@ -186,9 +210,6 @@ def _validated_next_nodes(
             ), f"{self.__class__.__name__} reached an endpoint node which is not a placeholder or output: {frontier}"
             return None
 
-        if any(next_node in stale_nodes for next_node in next_nodes):
-            return None
-
         if not self._can_cross_next_nodes(frontier, next_nodes):
             return None
 
@@ -217,17 +238,15 @@ def _advance_through_next_node(
         next_node.args = swapped_args[1]
         return True
 
-    def _propagate(self, node: torch.fx.Node, stale_nodes: set[torch.fx.Node]) -> bool:
-        """Propagate one node without consulting metadata invalidated this
-        scan.
-        """
+    def _propagate(self, node: torch.fx.Node) -> bool:
+        """Propagate one node and refresh metadata changed by the rewrite."""
 
         frontier = node
         previous_frontier = None
         propagation_path = [node]
         moved = False
         while True:
-            next_nodes = self._validated_next_nodes(node, frontier, stale_nodes)
+            next_nodes = self._validated_next_nodes(node, frontier)
             if next_nodes is None:
                 break
 
@@ -236,7 +255,7 @@ def _propagate(self, node: torch.fx.Node, stale_nodes: set[torch.fx.Node]) -> bo
                     node, frontier, previous_frontier, next_nodes
                 ):
                     break
-                self._mark_stale_region((*propagation_path, *next_nodes), stale_nodes)
+                self._refresh_propagation_meta(node, propagation_path)
                 return True
 
             next_node = next_nodes[0]
@@ -251,7 +270,7 @@ def _propagate(self, node: torch.fx.Node, stale_nodes: set[torch.fx.Node]) -> bo
             # Perform the swap directly in this case and return.
             # Otherwise break and move the node before the concat
             if self._maybe_split_upwards_cat_fanout(node, next_node):
-                self._mark_stale_region((*propagation_path, next_node), stale_nodes)
+                self._refresh_propagation_meta(node, propagation_path)
                 return True
 
             # Unhandled case, stop propagation
@@ -262,7 +281,7 @@ def _propagate(self, node: torch.fx.Node, stale_nodes: set[torch.fx.Node]) -> bo
 
         assert previous_frontier is not None
         self._move_node(node, frontier, previous_frontier)
-        self._mark_stale_region(propagation_path, stale_nodes)
+        self._refresh_propagation_meta(node, propagation_path)
         return True
 
     def duplicate_user_fusion_exclusions(self) -> frozenset:
@@ -513,6 +532,15 @@ def _shape_seen_by_next_node(self, moving_node: torch.fx.Node) -> Any:
         val = moving_node.meta.get("val")
         return getattr(val, "shape", None)
 
+    def _refresh_propagation_meta(self, moving_node, propagation_path) -> None:
+        # After moving up, dependencies run from the transform through the
+        # crossed nodes in reverse propagation order.
+        if moving_node.graph is not None:
+            self._refresh_node_meta(moving_node)
+        for node in reversed(propagation_path[1:]):
+            if node.graph is not None:
+                self._refresh_node_meta(node)
+
     def fuse_horizontal(self, graph_module):
         modified = False
         result = FuseDuplicateUsersPass(self.duplicate_user_fusion_exclusions()).call(
@@ -815,6 +843,20 @@ def _shape_seen_by_next_node(self, moving_node: torch.fx.Node) -> Any:
         val = cast(torch.fx.Node, moving_node.args[0]).meta.get("val")
         return getattr(val, "shape", None)
 
+    def _propagation_order(
+        self, graph_module: torch.fx.GraphModule
+    ) -> Iterable[torch.fx.Node]:
+        return reversed(list(graph_module.graph.nodes))
+
+    def _refresh_propagation_meta(self, moving_node, propagation_path) -> None:
+        # After moving down, dependencies run through the crossed nodes in
+        # propagation order and end at the transform.
+        for node in propagation_path[1:]:
+            if node.graph is not None:
+                self._refresh_node_meta(node)
+        if moving_node.graph is not None:
+            self._refresh_node_meta(moving_node)
+
     def fuse_horizontal(self, graph_module):
         modified = False
         result = FuseIdenticalInputTransformsPass(