Revert "Arm backend: Serialize controlflow submodules. (pytorch#15381)"

This reverts commit a4e7475.

Author: GregoryComer

backends/arm/_passes/arm_pass_manager.py

@@ -153,15 +153,15 @@ def _transform(self, graph_module: GraphModule):
         with TosaLoweringContext(self.tosa_spec):
             return self(graph_module).graph_module
 
-    def _tosa_INT_pipeline(
-        self, exported_program: ExportedProgram, graph_module: GraphModule
-    ) -> GraphModule:
+    def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(AnnotateOutputDimOrderPass())
         self.add_pass(FuseQuantizedActivationPass())
         self.add_pass(RemoveGetItemPass())
         self.add_pass(ConvertSplitToSlicePass())
         self.add_pass(ConvertMmToBmmPass())
-        self.add_pass(DecomposeMeanDimPass(graph_module, self.tosa_spec))
+        self.add_pass(
+            DecomposeMeanDimPass(exported_program.graph_module, self.tosa_spec)
+        )
         self.add_pass(ConvertFullLikeToFullPass())
         self.add_pass(ConvertToClampPass())
         self.add_pass(ConvertMinMaxPass())
@@ -218,11 +218,9 @@ def _tosa_INT_pipeline(
         self.add_pass(InsertRescalePass())
 
         self.validate_constraints_mandatory()
-        return self._transform(graph_module)
+        return self._transform(exported_program.graph_module)
 
-    def _tosa_FP_pipeline(
-        self, exported_program: ExportedProgram, graph_module: GraphModule
-    ) -> GraphModule:
+    def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(AnnotateOutputDimOrderPass())
         self.add_pass(DecomposeExpm1Pass())
         self.add_pass(DecomposeLogitPass())
@@ -257,7 +255,9 @@ def _tosa_FP_pipeline(
         self.add_pass(DecomposeLayerNormPass())
         self.add_pass(DecomposeBatchNormNoStatsPass())
         self.add_pass(DecomposeVarPass())
-        self.add_pass(DecomposeMeanDimPass(graph_module, self.tosa_spec))
+        self.add_pass(
+            DecomposeMeanDimPass(exported_program.graph_module, self.tosa_spec)
+        )
         self.add_pass(DecomposeNotEqualPass())
         self.add_pass(DecomposeDivPass())
         self.add_pass(DecomposeAddSubAlphaPass())
@@ -305,16 +305,14 @@ def _tosa_FP_pipeline(
         self.add_pass(InsertRescalePass())
 
         self.validate_constraints_mandatory()
-        return self._transform(graph_module)
+        return self._transform(exported_program.graph_module)
 
-    def transform_to_backend_pipeline(
-        self, exported_program: ExportedProgram, graph_module: GraphModule
-    ):
+    def transform_to_backend_pipeline(self, exported_program: ExportedProgram):
         """Apply passes before transforming program to backend"""
         if self.tosa_spec == TosaSpecification.create_from_string("TOSA-1.0+FP"):
-            return self._tosa_FP_pipeline(exported_program, graph_module)
+            return self._tosa_FP_pipeline(exported_program)
         elif self.tosa_spec == TosaSpecification.create_from_string("TOSA-1.0+INT"):
-            return self._tosa_INT_pipeline(exported_program, graph_module)
+            return self._tosa_INT_pipeline(exported_program)
         else:
             raise NotImplementedError(
                 f"No pass pipeline implemented for {self.tosa_spec=}"

backends/arm/process_node.py

@@ -158,7 +158,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=tosa_arg.name
+        buffer_values.shape, tosa_arg.dtype, buffer_values, name=node.name
     )
 
 
@@ -215,9 +215,11 @@ def process_placeholder(
         raise RuntimeError(f"Placeholder '{node.name}' is of unknown type.")
 
 
-def process_output(node: torch.fx.Node, tosa_graph: Any, tosa_spec: TosaSpecification):
+def process_output(
+    node: torch.fx.Node,
+    tosa_graph: Any,
+):
     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_arg.name]
+            tosa_graph.currRegion.currBasicBlock.tensors[output.name]
         )

backends/arm/tosa/backend.py

@@ -24,13 +24,10 @@
     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
-
+from torch.fx import Graph, Node
 
 # TOSA backend debug functionality
 logger = logging.getLogger(__name__)
@@ -55,39 +52,13 @@ def bfs_mark(start_nodes: List[Node], idx: int, seen: Set[Node]):
             # Walk backwards so we touch every producer
             q.extend(n.all_input_nodes)
 
-    out = ep_graph.output_node()
-    # First argument of output node is tuple of outputs
-    output_list = cast(tuple, out.args[0])
+    out = next(n for n in ep_graph.nodes if n.op == "output")
     seen: Set[Node] = set()
-    for idx, val in enumerate(output_list):
+    for idx, val in enumerate(out.args[0]):
         bfs_mark([val], idx, seen)
     return node2external_id
 
 
-def _sort_outputs(graph_module: GraphModule, node_to_id_map: dict[str, int]):
-    def _external_id(n: Node, node_2_id, fallback: int) -> int:
-        return node_2_id.get(n.name, fallback)
-
-    out_node = graph_module.graph.output_node()
-    out_list = cast(tuple, out_node.args[0])
-    _counter = count()
-
-    # sort nodes by the key that is id
-    def _sort_key(t: Node) -> int:
-        return _external_id(t, node_to_id_map, next(_counter))
-
-    orig_ord = tuple(sorted(out_list, key=_sort_key))
-
-    current_order = tuple(out_list)
-    if orig_ord != current_order:
-        replacement = list(orig_ord) if isinstance(out_node.args[0], list) else orig_ord
-        out_node.args = (replacement,)
-        graph_module.graph.lint()
-        graph_module.recompile()
-
-    return graph_module
-
-
 def arm_get_first_delegation_tag(graph_module) -> str:
     """Get the first delegation tag from the graph_module or return empty string."""
     for node in graph_module.graph.nodes:
@@ -122,9 +93,9 @@ def _preprocess(  # noqa: C901
         artifact_path = compile_spec.get_intermediate_path()
         tosa_spec = compile_spec.tosa_spec
         dump_debug_info = compile_spec.tosa_debug_mode
-        debug_hook = None
-        if dump_debug_info is not None:
-            debug_hook = DebugHook(dump_debug_info)
+
+        # Assign to every node external id
+        node_2_id = _annotate_external_ids(edge_program.graph)
 
         logger.info(f"Converting ExportedProgram to TOSA: {tosa_spec}")
 
@@ -145,66 +116,45 @@ def _preprocess(  # noqa: C901
                 f"doesn't match specification {tosa_spec}"
             )
 
-        TOSABackend._preprocess_module(
-            edge_program.graph_module,
-            edge_program,
-            compile_spec,
-            tosa_graph,
-            debug_hook,
-        )
-        # Serialize and return the TOSA flatbuffer.
-        binary = tosa_graph.serialize()
-
-        if artifact_path:
-            tag = arm_get_first_delegation_tag(edge_program.graph_module)
-            debug_tosa_dump(
-                binary,
-                artifact_path,
-                suffix="{}".format(f"_{tag}" if tag else "") + (f"_{tosa_spec}"),
-            )
-
-            if debug_hook is not None:
-                if debug_hook.mode == ArmCompileSpec.DebugMode.JSON:
-                    json_output = debug_hook.serialize()
-                    with open(f"{artifact_path}/debug.json", "w") as f:
-                        f.write(json_output)
-
-        return PreprocessResult(processed_bytes=binary)
-
-    @staticmethod
-    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
-        node_to_id_map = _annotate_external_ids(graph_module.graph)
-        artifact_path = compile_spec.get_intermediate_path()
-
         # TODO: Fix the need to lazily import this.
         from executorch.backends.arm._passes import ArmPassManager
 
         graph_module = ArmPassManager(tosa_spec).transform_to_backend_pipeline(  # type: ignore
-            exported_program=edge_program, graph_module=graph_module
+            exported_program=edge_program
         )
 
+        debug_hook = None
+        if dump_debug_info is not None:
+            debug_hook = DebugHook(dump_debug_info)
+
         # 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)
-        graph_module = _sort_outputs(graph_module, node_to_id_map)
 
-        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
+        # Re-shuffle output nodes to preserve author's order
+        def _external_id(n: Node, node_2_id, fallback: int) -> int:
+            return node_2_id.get(n.name, fallback)
+
+        out_node = next(n for n in graph_module.graph.nodes if n.op == "output")
+        _counter = count()
+
+        # sort nodes by the key that is id
+        def _sort_key(t: Node) -> int:
+            return _external_id(t, node_2_id, next(_counter))
 
+        orig_ord = tuple(sorted(out_node.args[0], key=_sort_key))
+
+        current_order = tuple(out_node.args[0])
+        if orig_ord != current_order:
+            replacement = (
+                list(orig_ord) if isinstance(out_node.args[0], list) else orig_ord
+            )
+            out_node.args = (replacement,)
+            graph_module.graph.lint()
+            graph_module.recompile()
+
+        input_count = 0
         for node in graph_module.graph.nodes:
             node = cast(Node, node)
             try:
@@ -214,27 +164,37 @@ def _preprocess_module(  # noqa: C901
                     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, tosa_spec)
+                    process_output(node, tosa_graph)
                 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, artifact_path)
+                debug_fail(node, graph_module, tosa_graph.serialize(), 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,
+        # Serialize and return the TOSA flatbuffer.
+        binary = tosa_graph.serialize()
+
+        if artifact_path:
+            tag = arm_get_first_delegation_tag(graph_module)
+            debug_tosa_dump(
+                binary,
+                artifact_path,
+                suffix="{}".format(f"_{tag}" if tag else "") + (f"_{tosa_spec}"),
             )
 
+            if debug_hook is not None:
+                if debug_hook.mode == ArmCompileSpec.DebugMode.JSON:
+                    json_output = debug_hook.serialize()
+                    with open(f"{artifact_path}/debug.json", "w") as f:
+                        f.write(json_output)
+
+        return PreprocessResult(processed_bytes=binary)
+
     @staticmethod
     def filter_tosa_compile_specs(
         compile_spec: ArmCompileSpec,

backends/arm/tosa/mapping.py

@@ -17,8 +17,6 @@
 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,
@@ -146,7 +144,7 @@ def __process_node(self, argument: torch.fx.Node):
             argument (torch.fx.Node): FX node to inspect.
 
         """
-        self.name = argument.name + argument.meta.get(TOSA_TENSOR_NAME_META, "")
+        self.name: str = argument.name
         output_dtype, self.shape, self.dim_order = extract_tensor_meta(
             argument.meta, self.tosa_spec
         )