Cache graph_signature property lookups in partitioning hot loops (#18352)

inputs_to_parameters, inputs_to_buffers, etc. are @Property methods
that rebuild a dict from input_specs on every access. In partitioning
loops that check each graph node, this causes O(nodes × params) dict
constructions. For a 40-layer MoE model with ~26K nodes and ~1,671
parameters, this results in ~260M iterations of pure overhead that
showed up as multiple minutes during to_edge_transform_and_lower.

Cache the dicts before the loops in:
- backend/utils.py: tag_constant_data (2 loops × 3 dict lookups each)
- backend/utils.py: tag_mutated_buffer (1 loop × 2 dict lookups)
- lowered_backend_module.py: __call__ (2 passes × 2 dict lookups)
- lowered_backend_module.py: arrange_graph_placeholders (1 loop × 2)

No behavioral change — the graph signature is not mutated between
caching and use in any of these functions.

Author: mergennachin

exir/backend/utils.py

@@ -24,7 +24,6 @@
 
 from executorch.exir.lowered_backend_module import create_submodule_from_nodes
 from tabulate import tabulate
-from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 from torch.fx.experimental.symbolic_shapes import has_free_symbols
 from torch.fx.node import Node
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -351,15 +350,22 @@ def tag_constant_data(edge_program: ExportedProgram) -> None:
     subgraph. Throw error when const/param/buffers is used across different partitions. That is the
     underlying data will be owned by multiple delegates.
     """
+    # Cache signature lookups to avoid rebuilding dicts on every access.
+    sig = edge_program.graph_signature
+    params_map = sig.inputs_to_parameters
+    buffers_map = sig.inputs_to_buffers
+    constants_map = sig.inputs_to_lifted_tensor_constants
+    buffers_to_mutate = sig.buffers_to_mutate
+
     mutated_buffer = set()
     for node in edge_program.graph.nodes:
         if node.op == "placeholder" and (
-            is_param(edge_program, node)
-            or is_buffer(edge_program, node)
-            or is_lifted_tensor_constant(edge_program, node)
+            node.name in params_map
+            or node.name in buffers_map
+            or node.name in constants_map
         ):
             for node_user in node.users:
-                if node_user.name in edge_program.graph_signature.buffers_to_mutate:
+                if node_user.name in buffers_to_mutate:
                     logging.info(
                         "The buffer node is a mutated buffer node, which is not constant."
                     )
@@ -368,9 +374,9 @@ def tag_constant_data(edge_program: ExportedProgram) -> None:
     for node in edge_program.graph.nodes:
         # go through const/param/buffer nodes, if all users of const/param/buffer nodes are partitioned then partition
         if node.op == "placeholder" and (
-            is_param(edge_program, node)
-            or is_buffer(edge_program, node)
-            or is_lifted_tensor_constant(edge_program, node)
+            node.name in params_map
+            or node.name in buffers_map
+            or node.name in constants_map
         ):
             if node not in mutated_buffer:
                 user_tags = set()
@@ -397,12 +403,17 @@ def tag_mutated_buffer(edge_program: ExportedProgram) -> None:
     subgraph. Throw error when buffers is used across different partitions. That is the
     underlying data will be owned by multiple delegates.
     """
+    # Cache signature lookups to avoid rebuilding dicts on every access.
+    sig = edge_program.graph_signature
+    buffers_map = sig.inputs_to_buffers
+    buffers_to_mutate = sig.buffers_to_mutate
+
     for node in edge_program.graph.nodes:
         # Determine whether this node is a mutated buffer
         is_mutated_buffer_node = False
-        if node.op == "placeholder" and is_buffer(edge_program, node):
+        if node.op == "placeholder" and node.name in buffers_map:
             for node_user in node.users:
-                if node_user.name in edge_program.graph_signature.buffers_to_mutate:
+                if node_user.name in buffers_to_mutate:
                     is_mutated_buffer_node = True
                     break
         # This node is mutated buffer, tag it

exir/lowered_backend_module.py

@@ -223,16 +223,19 @@ def program(
 
         lowered_exported_program = copy.deepcopy(self._original_exported_program)
 
+        # Cache these properties to avoid rebuilding the dict on each access.
+        sig = lowered_exported_program.graph_signature
+        params_map = sig.inputs_to_parameters
+        buffers_map = sig.inputs_to_buffers
+
         # The real input nodes are the ones not buffer or parameter
         all_input_nodes = [
             node
             for node in lowered_exported_program.graph.nodes
             if (
                 node.op == "placeholder"
-                and node.name
-                not in lowered_exported_program.graph_signature.inputs_to_buffers
-                and node.name
-                not in lowered_exported_program.graph_signature.inputs_to_parameters
+                and node.name not in buffers_map
+                and node.name not in params_map
             )
         ]
 
@@ -250,9 +253,7 @@ def program(
         # Find placeholders that are parameters or buffers, remove them from the main graph
         for node in lowered_exported_program.graph.nodes:
             if node.op == "placeholder" and (
-                node.name in lowered_exported_program.graph_signature.inputs_to_buffers
-                or node.name
-                in lowered_exported_program.graph_signature.inputs_to_parameters
+                node.name in buffers_map or node.name in params_map
             ):
                 lowered_exported_program.graph.erase_node(node)
 
@@ -402,22 +403,19 @@ def arrange_graph_placeholders(
     graph_sign = owning_program.graph_signature
 
     # Add all placeholders into the graph first:
+    # Cache these properties — each call rebuilds the dict from input_specs.
+    params_map = graph_sign.inputs_to_parameters
+    buffers_map = graph_sign.inputs_to_buffers
     param_nodes = []
     buffer_nodes = []
     input_nodes = []
     for node in gm.graph.nodes:
         if node.op != "placeholder":
             continue
 
-        if (
-            node.name in graph_sign.inputs_to_parameters
-            and node.meta.get("delegation_tag", None) == tag
-        ):
+        if node.name in params_map and node.meta.get("delegation_tag", None) == tag:
             param_nodes.append(node)
-        elif (
-            node.name in graph_sign.inputs_to_buffers
-            and node.meta.get("delegation_tag", None) == tag
-        ):
+        elif node.name in buffers_map and node.meta.get("delegation_tag", None) == tag:
             buffer_nodes.append(node)
         else:
             input_nodes.append(node)