Arm backend: Enable serializing to different regions.

Each conditional submodule in the graph_module gets its own region.
The TOSA reference model requires all tensor names in one model
to be unique, regardless of region. Pytorch's naming semantics,
however don't guarantee this. To fix this, attach a suffix containing
the submodule name to tensors in submodules.

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: I910a7d71f0b5da2d2d9219746efd012f9bd251fb

Author: Erik-Lundell

backends/arm/process_node.py

@@ -159,7 +159,7 @@ def process_inputs_to_buffers(
     buffer_values = np.transpose(buffer_values, tosa_arg.dim_order)
 
     tosa_graph.addConst(
-        buffer_values.shape, tosa_arg.dtype, buffer_values, name=node.name
+        buffer_values.shape, tosa_arg.dtype, buffer_values, name=tosa_arg.name
     )
 
 
@@ -216,11 +216,9 @@ def process_placeholder(
         raise RuntimeError(f"Placeholder '{node.name}' is of unknown type.")
 
 
-def process_output(
-    node: torch.fx.Node,
-    tosa_graph: Any,
-):
+def process_output(node: torch.fx.Node, tosa_graph: Any, tosa_spec: TosaSpecification):
     for output in cast(tuple[torch.fx.Node, ...], node.args[0]):
+        output_arg = TosaArg(output, tosa_spec)
         tosa_graph.addOutputTensor(
-            tosa_graph.currRegion.currBasicBlock.tensors[output.name]
+            tosa_graph.currRegion.currBasicBlock.tensors[output_arg.name]
         )

backends/arm/tosa/backend.py

@@ -24,11 +24,14 @@
     process_placeholder,
 )
 from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec
+from executorch.backends.arm.tosa.mapping import TOSA_TENSOR_NAME_META
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.exir.graph_module import get_control_flow_submodules
 from torch.export.exported_program import ExportedProgram
 from torch.fx import Graph, GraphModule, Node
 
+
 # TOSA backend debug functionality
 logger = logging.getLogger(__name__)
 
@@ -169,12 +172,13 @@ def _preprocess(  # noqa: C901
         return PreprocessResult(processed_bytes=binary)
 
     @staticmethod
-    def _preprocess_module(
+    def _preprocess_module(  # noqa: C901
         graph_module: GraphModule,
         edge_program: ExportedProgram,
         compile_spec: TosaCompileSpec,
         tosa_graph: ts.TosaSerializer,
         debug_hook: DebugHook | None,
+        submodule_name: str | None = None,
     ):
         """Convert 'graph_module' to a tosa_graph"""
         tosa_spec = compile_spec.tosa_spec
@@ -194,7 +198,13 @@ def _preprocess_module(
         node_visitors = get_node_visitors(edge_program, tosa_spec, debug_hook)
         graph_module = _sort_outputs(graph_module, node_to_id_map)
 
-        input_count = 0
+        if submodule_name is not None:
+            tosa_graph.startRegion(submodule_name)
+            tosa_graph.currRegion.addBasicBlock(submodule_name)
+            suffix = f"_{submodule_name}"
+            for loop_node in graph_module.graph.nodes:
+                loop_node.meta[TOSA_TENSOR_NAME_META] = suffix
+
         for node in graph_module.graph.nodes:
             node = cast(Node, node)
             try:
@@ -204,18 +214,27 @@ def _preprocess_module(
                     if len(node.users) == 0:
                         continue
                     process_placeholder(node, tosa_graph, edge_program, tosa_spec)
-                    if node.name in edge_program.graph_signature.user_inputs:
-                        input_count += 1
                 elif node.op == "output":
-                    process_output(node, tosa_graph)
+                    process_output(node, tosa_graph, tosa_spec)
                 else:
                     # This will only happen if an unpartitioned graph is passed without
                     # any checking of compatibility.
                     raise RuntimeError(f"{node.name} is unsupported op {node.op}")
             except Exception:
-                debug_fail(node, graph_module, tosa_graph.serialize(), artifact_path)
+                debug_fail(node, graph_module, tosa_graph, artifact_path)
                 raise
 
+        # Recursively preprocess controlflow submodules.
+        for name, submodule, _ in get_control_flow_submodules(graph_module):
+            TOSABackend._preprocess_module(
+                submodule,
+                edge_program,
+                compile_spec,
+                tosa_graph,
+                debug_hook,
+                submodule_name=name,
+            )
+
     @staticmethod
     def filter_tosa_compile_specs(
         compile_spec: ArmCompileSpec,

backends/arm/tosa/mapping.py

@@ -17,6 +17,8 @@
 import tosa_serializer as ts
 from executorch.backends.arm.tosa.specification import TosaSpecification
 
+TOSA_TENSOR_NAME_META = "tosa_tensor_name"
+
 UNSUPPORTED_DTYPES = (
     torch.float64,
     torch.double,
@@ -144,7 +146,7 @@ def __process_node(self, argument: torch.fx.Node):
             argument (torch.fx.Node): FX node to inspect.
 
         """
-        self.name: str = argument.name
+        self.name = argument.name + argument.meta.get(TOSA_TENSOR_NAME_META, "")
         output_dtype, self.shape, self.dim_order = extract_tensor_meta(
             argument.meta, self.tosa_spec
         )