Revert "Implement dumping operator distribution for TOSA graph" (#5131)

Reverts #4970

The pull / unittest-arm (buck2) / linux-job and trunk / test-arm-reference-delegation / linux-job jobs started failing in trunk (dashboard) after this PR merged. See #4970 (comment) for more details.

Author: dbort

backends/arm/test/misc/test_debug_feats.py

@@ -126,67 +126,26 @@ def test_numerical_diff_prints(self):
             self.fail()
 
 
-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()
+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()
         )
-        .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
+        # Just test that there are no execeptions.

backends/arm/test/tester/arm_tester.py

@@ -13,7 +13,7 @@
 
 import numpy as np
 
-import torch.fx
+import torch
 
 from executorch.backends.arm.arm_backend import get_intermediate_path, is_permute_memory
 from executorch.backends.arm.arm_partitioner import ArmPartitioner
@@ -297,24 +297,23 @@ def get_graph(self, stage: str | None = None) -> Graph:
 
         return graph
 
-    def dump_operator_distribution(self, path_to_dump: Optional[str] = None):
+    def dump_operator_distribution(
+        self, path_to_dump: Optional[str] = None
+    ) -> ArmQuantizer:
         """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(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
-            )
+        to_print = self.cur + " operators: " + _format_dict(op_dist) + "\n"
         _dump_str(to_print, path_to_dump)
         return self
 
-    def dump_dtype_distribution(self, path_to_dump: Optional[str] = None):
+    def dump_dtype_distribution(
+        self, path_to_dump: Optional[str] = None
+    ) -> ArmQuantizer:
         """Dump a dictionary with {dtype: dtype count} for the dtypes of the nodes in the
         graph of the current stage.
 
@@ -422,36 +421,6 @@ 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: