[executorch][arm] Speed up lowering: cache constant names, fix quadratic permute chain collection, dedup producer queue, defer cleanup, and drop redundant recompiles

backends/arm/_passes/canonicalize_view_copy_permute_pass.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections import deque
 from typing import cast, Sequence, Set, Type
 
 import torch
@@ -82,7 +83,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
 
         if modified:
             graph_module.graph.eliminate_dead_code()
-            graph_module.recompile()
+            graph_module.graph.lint()
             graph_module = super().call(graph_module).graph_module
 
         return PassResult(graph_module, modified)
@@ -91,20 +92,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/arm/_passes/fuse_duplicate_users_pass.py

@@ -45,9 +45,17 @@ def call(self, graph_module: GraphModule) -> PassResult:
 
         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)
+
+        def enqueue_producer(node: Node) -> None:
+            if node.graph is None 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:
                 # Node was deleted by a previous rewrite while still queued.
@@ -84,11 +92,10 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     # 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:
-            graph_module.recompile()
             graph_module.graph.lint()
             graph_module = super().call(graph_module).graph_module
 

backends/arm/_passes/fuse_identical_input_transforms_pass.py

@@ -180,7 +180,6 @@ def call(self, graph_module: GraphModule) -> PassResult:
         if modified:
             graph_module.graph.eliminate_dead_code()
             graph_module.graph.lint()
-            graph_module.recompile()
             graph_module = super().call(graph_module).graph_module
 
         return PassResult(graph_module, modified)

backends/arm/_passes/propagate_view_copy_permute_pass.py

@@ -108,24 +108,28 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                         continue
                     if self._propagate(node):
                         iteration_modified = True
+                        graph_module = self._retrace(graph_module)
                         break
 
             if iteration_modified:
-                graph_module = self._retrace(graph_module)
-                result = self.fuse_horizontal(graph_module)
-                graph_module = result.graph_module
-                iteration_modified |= result.modified
-                result = self.fuse_vertical(graph_module)
-                graph_module = result.graph_module
-                iteration_modified |= result.modified
+                modified = True
+                continue
+
+            result = self.fuse_horizontal(graph_module)
+            graph_module = result.graph_module
+            iteration_modified |= result.modified
+            result = self.fuse_vertical(graph_module)
+            graph_module = result.graph_module
+            iteration_modified |= result.modified
 
             modified |= iteration_modified
-            if not iteration_modified:
-                break
+            if iteration_modified:
+                graph_module = self._retrace(graph_module)
+                continue
+            break
 
         if modified:
             graph_module = self._retrace(graph_module)
-            graph_module.recompile()
 
         return PassResult(graph_module, modified)
 

backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py

@@ -5,7 +5,6 @@
 
 import torch
 
-from executorch.backends.arm._passes.arm_pass_utils import is_param_node
 from executorch.backends.arm._passes.insert_table_ops import TableOps
 from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
     RemovePermutesAroundElementwiseOps,
@@ -25,10 +24,20 @@ def __init__(self, exported_program: ExportedProgram) -> None:
             }
         )
         self.exported_program = exported_program
+        # Precompute parameter/buffer/lifted-constant placeholder names once.
+        # is_param_node() rebuilds these graph_signature maps on every call, so
+        # calling it per node made the visit() walk O(nodes * inputs).
+        gs = exported_program.graph_signature
+        self._constant_input_names: set[str] = (
+            set(gs.inputs_to_parameters)
+            | set(gs.inputs_to_buffers)
+            | set(gs.inputs_to_lifted_tensor_constants)
+        )
 
     def _is_constant(self, node: torch.fx.Node) -> bool:
-        # Override fragile string match check with exported program check
-        return super()._is_constant(node) or is_param_node(self.exported_program, node)
+        # get_attr nodes are handled by super()._is_constant; set membership
+        # here is equivalent to is_param_node for placeholder inputs.
+        return super()._is_constant(node) or node.name in self._constant_input_names
 
     def permute_subgraph(self, subgraph) -> bool:
         # TABLE lookup inputs are already tied to the table layout.