Cache param/buffer/constant names in RemovePermutesAroundElementwiseTosaOps (#22164)

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.

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."""