From 0a8bc0a7ec36bf06fc0f3186987d4399ade3e6ca Mon Sep 17 00:00:00 2001
From: Andrew Pullin <pullinandrew@meta.com>
Date: Fri, 28 Aug 2026 08:56:48 -0700
Subject: [PATCH] Cache param/buffer/constant names in
 RemovePermutesAroundElementwiseTosaOps (#22164)

Summary:

`RemovePermutesAroundElementwiseTosaOps._is_constant` called `is_param_node()` for every node reached during recursive `visit()` walk. Each `is_param_node` rebuilds immutable dict over full input_specs via uncached graph_signature properties. This made pass O(nodes * inputs). Precompute the union of param/buffer/lifted-constant placeholder names once per pass invocation in `call()` and use O(1) set membership. Refreshing it at invocation time tracks the current `ExportedProgram` when a pass instance is reused. Semantics unchanged.

On an ensemble network of ~1M parameters, lowering produced identical output while running 10.5 seconds faster (2.7%).

Differential Revision: D114224782
---
 ...ve_permutes_around_elementwise_tosa_ops.py | 18 ++++++++++---
 ...ve_permutes_around_elementwise_tosa_ops.py | 27 +++++++++++++++++--
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py
index b241038f7a9..eb5ac30f890 100644
--- a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py
+++ b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py
@@ -5,13 +5,13 @@
 
 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,
 )
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
 
 
 class RemovePermutesAroundElementwiseTosaOps(RemovePermutesAroundElementwiseOps):
@@ -25,10 +25,22 @@ def __init__(self, exported_program: ExportedProgram) -> None:
             }
         )
         self.exported_program = exported_program
+        self._constant_input_names: set[str] = set()
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        # The ARM pass manager updates exported_program before each invocation.
+        gs = self.exported_program.graph_signature
+        self._constant_input_names = (
+            set(gs.inputs_to_parameters)
+            | set(gs.inputs_to_buffers)
+            | set(gs.inputs_to_lifted_tensor_constants)
+        )
+        return super().call(graph_module)
 
     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.
diff --git a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py
index 35e02b52573..864b6c669f9 100644
--- a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py
+++ b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py
@@ -29,14 +29,14 @@
 ERF_TARGET = exir_ops.edge.aten.erf.default
 
 
-def _fake_exported_program() -> ExportedProgram:
+def _fake_exported_program(*constant_input_names: str) -> ExportedProgram:
     return cast(
         ExportedProgram,
         SimpleNamespace(
             graph_signature=SimpleNamespace(
                 inputs_to_buffers={},
                 inputs_to_lifted_tensor_constants={},
-                inputs_to_parameters={},
+                inputs_to_parameters={name: name for name in constant_input_names},
             )
         ),
     )
@@ -50,6 +50,29 @@ def _count_nodes(graph_module: torch.fx.GraphModule, target) -> int:
     )
 
 
+def test_constant_input_cache_refreshes_for_reused_pass() -> None:
+    first_graph = torch.fx.Graph()
+    first_constant = first_graph.placeholder("first_constant")
+    first_constant.meta["val"] = torch.randn(1)
+    first_graph.output(first_constant)
+
+    second_graph = torch.fx.Graph()
+    second_constant = second_graph.placeholder("second_constant")
+    second_constant.meta["val"] = torch.randn(1)
+    second_graph.output(second_constant)
+
+    remove_permutes = RemovePermutesAroundElementwiseTosaOps(
+        _fake_exported_program("first_constant")
+    )
+    remove_permutes.call(torch.fx.GraphModule({}, first_graph))
+    assert remove_permutes._is_constant(first_constant)
+
+    remove_permutes.exported_program = _fake_exported_program("second_constant")
+    remove_permutes.call(torch.fx.GraphModule({}, second_graph))
+    assert remove_permutes._is_constant(second_constant)
+    assert not remove_permutes._is_constant(first_constant)
+
+
 def test_extra_permutable_ops_makes_op_permutable() -> None:
     """Ops in extra_permutable_ops are permutable in the base pass."""