Arm backend: Enable capture of debug information during TOSA serialization (#13908)

* Add _serialize_operator() function to NodeVisitor
* Enable passing of DebugHook to NodeVisitor if enabled
* Add option dump_debug_info to ArmCompileSpecBuilder

Author: tom-arm

backends/arm/arm_backend.py

@@ -10,6 +10,7 @@
 # backends. Converts via TOSA as an intermediate form supported by AoT and
 # JIT compiler flows.
 #
+from enum import Enum
 from typing import List, Optional
 
 from executorch.backends.arm.tosa_specification import (  # type: ignore[import-not-found]
@@ -22,12 +23,16 @@
 
 
 class ArmCompileSpecBuilder:
+    class DebugMode(Enum):
+        JSON = 1
+
     def __init__(self):
         self.compile_spec: List[CompileSpec] = []
         self.compiler_flags = []
         self.output_format = None
         self.path_for_intermediates = None
         self.tosa_spec = None
+        self.tosa_debug_mode = None
 
     def vgf_compile_spec(
         self,
@@ -163,6 +168,13 @@ def dump_intermediate_artifacts_to(
         self.path_for_intermediates = output_path
         return self
 
+    def dump_debug_info(self, debug_mode: DebugMode) -> "ArmCompileSpecBuilder":
+        """
+        Dump debugging information into the intermediates path
+        """
+        self.tosa_debug_mode = debug_mode.name
+        return self
+
     def build(self) -> List[CompileSpec]:
         """
         Generate a list of compile spec objects from the builder
@@ -188,6 +200,16 @@ def build(self) -> List[CompileSpec]:
                 CompileSpec("debug_artifact_path", self.path_for_intermediates.encode())
             )
 
+        if self.tosa_debug_mode is not None:
+            if not self.path_for_intermediates:
+                raise ValueError(
+                    "dump_debug_info() must be used in conjunction with dump_intermediate_artifacts_to()"
+                )
+
+            self.compile_spec.append(
+                CompileSpec("dump_debug_info", self.tosa_debug_mode.encode())
+            )
+
         return self.compile_spec
 
 

backends/arm/operators/node_visitor.py

@@ -5,10 +5,11 @@
 
 # pyre-unsafe
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import torch
 
+from executorch.backends.arm.debug.schema import DebugHook
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from torch.export import ExportedProgram
@@ -29,9 +30,38 @@ class NodeVisitor:
         TosaSpecification.create_from_string("TOSA-1.0+FP"),
     ]
 
-    def __init__(self, exported_program: ExportedProgram, tosa_spec: TosaSpecification):
+    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
+
+    def _serialize_operator(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        tosa_op: Any,
+        inputs: List[str],
+        outputs: List[str],
+        attributes: Optional[Any] = None,
+    ) -> None:
+        tosa_graph.addOperator(
+            tosa_op,
+            inputs=inputs,
+            outputs=outputs,
+            attributes=attributes,
+        )
+
+        if self.debug_hook:
+            self.debug_hook.add(
+                node,
+                tosa_op=outputs[0],
+                tosa_op_id=tosa_op,
+            )
 
     def define_node(
         self,

backends/arm/operators/op_abs.py

@@ -123,7 +123,9 @@ def define_node(
             )
 
             # MI lowering
-            tosa_graph.addOperator(
+            self._serialize_operator(
+                node,
+                tosa_graph,
                 ts.TosaOp.Op().ABS,
                 [inputs[0].name],
                 [output.name],

backends/arm/operators/op_add.py

@@ -73,7 +73,9 @@ def define_node(
         input1, input2 = rescaled_inputs
 
         # Do the INT32 Add
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().ADD,
             [input1.name, input2.name],
             [add_output.name],
@@ -127,7 +129,9 @@ def define_node(
             input1, input2 = inputs
 
             # FP lowering
-            tosa_graph.addOperator(
+            self._serialize_operator(
+                node,
+                tosa_graph,
                 ts.TosaOp.Op().ADD,
                 [input1.name, input2.name],
                 [output.name],

backends/arm/operators/op_amax.py

@@ -61,6 +61,11 @@ def define_node(
 
         attr = ts.TosaSerializerAttribute()
         attr.ReduceMaxAttribute(axis=input.dim_order.index(dim), nan_mode=1)
-        tosa_graph.addOperator(
-            ts.TosaOp.Op().REDUCE_MAX, [input.name], [output.name], attr
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.TosaOp.Op().REDUCE_MAX,
+            [input.name],
+            [output.name],
+            attr,
         )

backends/arm/operators/op_amin.py

@@ -61,6 +61,11 @@ def define_node(
 
         attr = ts.TosaSerializerAttribute()
         attr.ReduceMinAttribute(axis=input.dim_order.index(dim), nan_mode=1)
-        tosa_graph.addOperator(
-            ts.TosaOp.Op().REDUCE_MIN, [input.name], [output.name], attr
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.TosaOp.Op().REDUCE_MIN,
+            [input.name],
+            [output.name],
+            attr,
         )

backends/arm/operators/op_any.py

@@ -52,6 +52,11 @@ def define_node(
         attr = ts.TosaSerializerAttribute()
         attr.ReduceAnyAttribute(inputs[0].dim_order.index(dim))
 
-        tosa_graph.addOperator(
-            ts.TosaOp.Op().REDUCE_ANY, [inputs[0].name], [output.name], attr
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.TosaOp.Op().REDUCE_ANY,
+            [inputs[0].name],
+            [output.name],
+            attr,
         )

backends/arm/operators/op_avg_pool2d.py

@@ -100,7 +100,9 @@ def _build_generic_avgpool2d(
             shape=[1], dtype=output.dtype, vals=[output_zp]
         )
 
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().AVG_POOL2D,
             [input_tensor.name, input_zp_tensor.name, output_zp_tensor.name],
             [output.name],

backends/arm/operators/op_bmm.py

@@ -78,7 +78,9 @@ def define_node(
         tosa_graph.addConst([1], inputs[1].dtype, [input1_zp], name=f"{node.name}_B_ZP")
 
         # Add the MATMUL to the TOSA graph.
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().MATMUL,
             [
                 inputs[0].name,

backends/arm/operators/op_cat.py

@@ -47,7 +47,9 @@ def define_node(
         attr = ts.TosaSerializerAttribute()
         attr.ConcatAttribute(dim)
 
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().CONCAT,
             [tensor.name for tensor in tensors],
             [output.name],