Arm backend: remove exported_program arg from NodeVisitor (#19370)

It was only used for a single check in op_tosa_table. Node visitors
should only require local knowledge.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell

Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/arm/operators/node_visitor.py

@@ -30,7 +30,6 @@
     TosaSpecification,
     TosaSpecMapping,
 )
-from torch.export import ExportedProgram
 
 logger = logging.getLogger(__name__)
 
@@ -39,7 +38,6 @@ class NodeVisitor:
     """Provide a visitor pattern to lower edge IR to TOSA.
 
     Attributes:
-        _exported_program (torch.export.ExportedProgram): Source program being lowered.
         tosa_spec (TosaSpecification): Active TOSA specification for lowering.
         debug_hook (Optional[DebugHook]): Optional hook for debug metadata.
 
@@ -54,11 +52,9 @@ class NodeVisitor:
 
     def __init__(
         self,
-        exported_program: ExportedProgram,
         tosa_spec: TosaSpecification,
         debug_hook: Optional[DebugHook] = None,
     ):
-        self._exported_program = exported_program
         self.tosa_spec = tosa_spec
         self.debug_hook = debug_hook
 

backends/arm/operators/op_tosa_table.py

@@ -52,21 +52,17 @@ def define_node(
         # The name of the table constant is a bit complex.
         # The name of the pytorch buffer will be the target of last node argument.
         # However, when it is serialized to TOSA, a submodule suffix might be added. The TOSA buffer name thus
-        # needs to be taken from the last TosaArg.
-        pytorch_table_buffer_name = node.args[-1].target  # type: ignore[union-attr]
-        tosa_table_buffer_name = inputs[-1].name
-        if pytorch_table_buffer_name not in self._exported_program.state_dict.keys():
-            raise RuntimeError(
-                f"Did not find key {node.name} in state_dict {self._exported_program.state_dict.keys()}."
-            )
+        # needs to be taken from the buffer TosaArg.
+        input, table_buffer = inputs
+        tosa_table_buffer_name = table_buffer.name
 
         attr = ts.TosaSerializerAttribute()
         attr.TableAttribute()
         self._serialize_operator(
             node,
             tosa_graph,
             ts.Op.TABLE,
-            [inputs[0].name, tosa_table_buffer_name],
+            [input.name, tosa_table_buffer_name],
             [output.name],
             attr,
         )

backends/arm/tosa/backend.py

@@ -340,7 +340,7 @@ def _preprocess_module(  # noqa: C901
         # TODO: Fix the need to lazily import this.
         from executorch.backends.arm.operators.node_visitor import get_node_visitors
 
-        node_visitors = get_node_visitors(edge_program, tosa_spec, debug_hook)
+        node_visitors = get_node_visitors(tosa_spec, debug_hook)
 
         if output_order_workaround:
             logger.debug("Re-sorting outputs during TOSA lowering.")