Speed up permute propagation cleanup (#22161)

backends/arm/test/passes/test_propagate_permutes_views_pass.py

@@ -1168,6 +1168,39 @@ def test_up_pass_fuses_equivalent_output_permutations_before_fan_out() -> None:
     assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(ADD)
 
 
+def test_up_pass_fusion_enables_further_propagation() -> 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))
+    first_permute = graph.call_function(PERMUTE, args=(neg, [0, 2, 3, 1]))
+    first_permute.meta["val"] = torch.empty((1, 3, 4, 2))
+    second_permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1]))
+    second_permute.meta["val"] = torch.empty((1, 3, 4, 2))
+    add = graph.call_function(ADD, args=(first_permute, second_permute))
+    add.meta["val"] = torch.empty((1, 3, 4, 2))
+    graph.output(add)
+
+    graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass)
+    call_nodes = [
+        node for node in graph_module.graph.nodes if node.op == "call_function"
+    ]
+    permutes = [node for node in call_nodes if node.target == PERMUTE]
+    x = next(node for node in graph_module.graph.nodes if node.op == "placeholder")
+    relu = next(node for node in call_nodes if node.target == RELU)
+    neg = next(node for node in call_nodes if node.target == NEG)
+    add = next(node for node in call_nodes if node.target == ADD)
+
+    assert len(permutes) == 1
+    assert permutes[0].args[0] is x
+    assert relu.args[0] is permutes[0]
+    assert neg.args[0] is relu
+    assert set(add.all_input_nodes) == {relu, neg}
+
+
 def test_propagate_moves_before_dtype_changing_rescale() -> None:
     graph = torch.fx.Graph()
     x = graph.placeholder("x")

backends/transforms/canonicalize_view_copy_permute_pass.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections import deque
 from typing import Any, cast, Iterable, Sequence, Set, Type
 
 import torch
@@ -104,20 +105,25 @@ def _collect_chains(self, graph_module: GraphModule) -> list[list[Node]]:
         """Returns a list of linear chains of view/permutes in the graph."""
         chains: list[list[Node]] = []
 
-        view_permute_nodes = [
+        view_permute_nodes = deque(
             node for node in graph_module.graph.nodes if node.target in self._targets
-        ]
+        )
+        remaining = set(view_permute_nodes)
 
         while view_permute_nodes:
-            node = view_permute_nodes.pop(0)
+            node = view_permute_nodes.popleft()
+            if node not in remaining:
+                continue
+            remaining.remove(node)
+
             chain = [node]
             current = node
 
             while len(current.users) == 1:
                 user = next(iter(current.users))
-                if user.target not in self._targets:
+                if user.target not in self._targets or user not in remaining:
                     break
-                view_permute_nodes.remove(user)
+                remaining.remove(user)
                 chain.append(user)
                 current = user
 

backends/transforms/fuse_duplicate_users_pass.py

@@ -37,17 +37,26 @@
         # so backends name those here rather than the pass guessing.
         self._excluded_targets = excluded_targets or frozenset()
 
     def call(self, graph_module: GraphModule) -> PassResult:
         graph = graph_module.graph
         modified = False
 
         node_order = {node: index for index, node in enumerate(graph.nodes)}
         producers: Deque[Node] = deque(node for node in graph.nodes)
+        queued_producers: Set[Node] = set(producers)
+        erased_nodes: Set[Node] = set()
+
+        def enqueue_producer(node: Node) -> None:
+            if node in erased_nodes or node in queued_producers:
+                return
+            producers.append(node)
+            queued_producers.add(node)
 
         while producers:
             producer = producers.popleft()
+            queued_producers.discard(producer)
 
-            if producer.graph is None:
+            if producer in erased_nodes:
                 # Node was deleted by a previous rewrite while still queued.
                 continue
 
@@ -56,7 +65,9 @@
             if len(user_nodes) < 2:
                 continue
 
-            candidate_groups = self._get_candidate_groups(node_order, user_nodes)
+            candidate_groups = self._get_candidate_groups(
+                node_order, user_nodes, erased_nodes
+            )
 
             signature_to_user: Dict[Tuple[Hashable, ...], Node] = {}
             for group in candidate_groups:
@@ -77,13 +88,14 @@
 
                     user.replace_all_uses_with(representative)
                     graph.erase_node(user)
+                    erased_nodes.add(user)
                     modified = True
 
                     # Revisit the current producer and the surviving user so that
                     # newly formed duplicate chains can be fused in later
                     # iterations.
-                    producers.append(producer)
-                    producers.append(representative)
+                    enqueue_producer(producer)
+                    enqueue_producer(representative)
 
         if modified:
             if self._recompile_before_retrace:
@@ -93,10 +105,10 @@
 
         return PassResult(graph_module, modified)
 
-    def _get_candidate_groups(self, node_order, user_nodes):
+    def _get_candidate_groups(self, node_order, user_nodes, erased_nodes):
         users_by_target: Dict[Tuple[str, Hashable], List[Node]] = {}
         for user in user_nodes:
-            if user.graph is None:
+            if user in erased_nodes:
                 # User might already have been removed by a prior rewrite.
                 continue