Implement dumping operator distribution for TOSA graph

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

Author: Erik-Lundell

backends/arm/test/misc/test_debug_feats.py

@@ -126,26 +126,67 @@ def test_numerical_diff_prints(self):
             self.fail()
 
 
-class TestDumpOperatorsAndDtypes(unittest.TestCase):
-    def test_dump_ops_and_dtypes(self):
-        model = Linear(20, 30)
-        (
-            ArmTester(
-                model,
-                example_inputs=model.get_inputs(),
-                compile_spec=common.get_tosa_compile_spec(),
-            )
-            .quantize()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
-            .export()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
-            .to_edge()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
-            .partition()
-            .dump_dtype_distribution()
-            .dump_operator_distribution()
+def test_dump_ops_and_dtypes():
+    model = Linear(20, 30)
+    (
+        ArmTester(
+            model,
+            example_inputs=model.get_inputs(),
+            compile_spec=common.get_tosa_compile_spec(),
+        )
+        .quantize()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+        .export()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+        .to_edge()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+        .partition()
+        .dump_dtype_distribution()
+        .dump_operator_distribution()
+    )
+    # Just test that there are no execptions.
+
+
+def test_dump_tosa_ops(capsys):
+    model = Linear(20, 30)
+    (
+        ArmTester(
+            model,
+            example_inputs=model.get_inputs(),
+            compile_spec=common.get_tosa_compile_spec(),
+        )
+        .quantize()
+        .export()
+        .to_edge()
+        .partition()
+        .dump_operator_distribution()
+    )
+    captured = capsys.readouterr()
+    assert "Partition operators:" in captured.out
+    assert "Tosa operators:" in captured.out
+
+
+def test_fail_dump_tosa_ops(capsys):
+    class Add(torch.nn.Module):
+        def forward(self, x):
+            return x + x
+
+    model = Add()
+    compile_spec = common.get_tosa_compile_spec_unbuilt()
+    compile_spec.output_format = "vela"
+    (
+        ArmTester(
+            model, example_inputs=(torch.ones(5),), compile_spec=compile_spec.build()
         )
-        # Just test that there are no execeptions.
+        .quantize()
+        .export()
+        .to_edge()
+        .partition()
+        .dump_operator_distribution()
+    )
+    captured = capsys.readouterr()
+    assert "Partition operators:" in captured.out
+    assert "Can not get operator distribution for vela command stream." in captured.out

backends/arm/test/tester/arm_tester.py

@@ -13,7 +13,7 @@
 
 import numpy as np
 
-import torch
+import torch.fx
 
 from executorch.backends.arm.arm_backend import get_intermediate_path, is_permute_memory
 from executorch.backends.arm.arm_partitioner import ArmPartitioner
@@ -297,23 +297,24 @@ def get_graph(self, stage: str | None = None) -> Graph:
 
         return graph
 
-    def dump_operator_distribution(
-        self, path_to_dump: Optional[str] = None
-    ) -> ArmQuantizer:
+    def dump_operator_distribution(self, path_to_dump: Optional[str] = None):
         """Dump a dictionary with {operator: operator count} for the operators in the
         graph of the current stage.
 
         Returns self for daisy-chaining.
         """
         graph = self.get_graph(self.cur)
         op_dist = _get_operator_distribution(graph)
-        to_print = self.cur + " operators: " + _format_dict(op_dist) + "\n"
+        to_print = self.cur + " operators: " + _format_dict(dict(op_dist)) + "\n"
+
+        if self.cur == self.stage_name(tester.Partition):
+            to_print += _get_tosa_operator_distribution(
+                self.get_artifact(self.cur).exported_program().graph_module
+            )
         _dump_str(to_print, path_to_dump)
         return self
 
-    def dump_dtype_distribution(
-        self, path_to_dump: Optional[str] = None
-    ) -> ArmQuantizer:
+    def dump_dtype_distribution(self, path_to_dump: Optional[str] = None):
         """Dump a dictionary with {dtype: dtype count} for the dtypes of the nodes in the
         graph of the current stage.
 
@@ -421,6 +422,36 @@ def _get_operator_distribution(graph: Graph) -> dict[str, int]:
     )
 
 
+def _get_tosa_operator_distribution(graph_module: torch.fx.GraphModule) -> str:
+    """Counts the occurences of operator names of all lowered modules containing
+    a TOSA flatbuffer.
+    The result is a string with the operator distribution or an error message.
+    """
+    op_list = []
+    id = 0
+    while lowered_module := getattr(graph_module, f"lowered_module_{id}", None):
+        for spec in lowered_module.compile_specs:
+            if spec.key != "output_format":
+                continue
+            if spec.value == b"tosa":
+                tosa_fb = lowered_module.processed_bytes
+                tosa_json = dbg_tosa_fb_to_json(tosa_fb)
+                for region in tosa_json["regions"]:
+                    for block in region["blocks"]:
+                        op_list.extend(
+                            [operator["op"] for operator in block["operators"]]
+                        )
+                break
+            elif spec.value == b"vela":
+                return "Can not get operator distribution for vela command stream."
+            else:
+                return f"Unknown output format '{spec.value}'."
+        id += 1
+    if id == 0:
+        return "No delegate with name 'lowered_module_0 found in graph module."
+    return "TOSA operators: " + _format_dict(dict(Counter(op_list)))
+
+
 def _dump_str(to_print: str, path_to_dump: Optional[str] = None):
     if path_to_dump:
         with open(path_to_dump, "a") as fp: