Arm backend: Add count_tosa_ops test helper (pytorch#18006)

Signed-off-by: Adrian Lundell <adrian.lundell@arm.com>

Author: AdrianLundell

backends/arm/test/misc/test_count_tosa_ops.py

@@ -0,0 +1,41 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import pytest
+import torch
+
+from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP
+
+
+class AddModule(torch.nn.Module):
+    def forward(self, x, y):
+        return x + y
+
+
+def test_count_tosa_ops_add_no_target():
+    model = AddModule()
+    test_data = (torch.randn(1, 8, 8, 8), torch.randn(1, 8, 8, 8))
+    pipeline = TosaPipelineFP[type(test_data)](
+        model,
+        test_data,
+        ["torch.ops.aten.add.Tensor"],
+        run_on_tosa_ref_model=False,
+    )
+    pipeline.count_tosa_ops({"ADD": 1, "SUB": 0})
+    pipeline.run()
+
+
+def test_count_tosa_ops_2_adds_no_target():
+    model = AddModule()
+    test_data = (torch.randn(1, 8, 8, 8), torch.randn(1, 8, 8, 8))
+    pipeline = TosaPipelineFP[type(test_data)](
+        model,
+        test_data,
+        ["torch.ops.aten.add.Tensor"],
+        run_on_tosa_ref_model=False,
+    )
+    pipeline.count_tosa_ops({"ADD": 2})
+    with pytest.raises(AssertionError, match="Expected 2 occurrences of TOSA op ADD"):
+        pipeline.run()

backends/arm/test/tester/arm_tester.py

@@ -1182,3 +1182,24 @@ def _format_dict(to_print: dict, print_table: bool = True) -> str:
         )
     else:
         return pformat(to_print, compact=True, indent=1)
+
+
+def count_tosa_ops(graph_module: torch.fx.GraphModule, expected_ops: Dict[str, int]):
+    """Asserts that the number of occurrences of TOSA operators in the graph of
+    a partitioned module matches the expected counts.
+    """
+    op_counts = dict(_get_tosa_operator_distribution(graph_module))
+    for op, expected_count in expected_ops.items():
+        actual_count = op_counts.get(op, 0)
+
+        if expected_count != actual_count:
+            if expected_count == 0:
+                raise AssertionError(
+                    f"Expected no occurrences of TOSA op {op} but found {actual_count}."
+                )
+            elif actual_count == 0:
+                raise AssertionError(f"Expected TOSA op {op} but it was not found.")
+            else:
+                raise AssertionError(
+                    f"Expected {expected_count} occurrences of TOSA op {op} but found {actual_count}."
+                )

backends/arm/test/tester/test_pipeline.py

@@ -20,6 +20,8 @@
     Union,
 )
 
+import executorch.backends.arm.test.tester.arm_tester as arm_tester_module
+
 import torch
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.ethosu import EthosUCompileSpec
@@ -301,6 +303,25 @@ def visualize(self, stage_id: str, suffix: str | None = None):
         self.add_stage_after(stage_id, self.tester.visualize, suffix=suffix)
         return self
 
+    def count_tosa_ops(self, expected_ops: Dict[str, int]):
+        """Assert the number of TOSA ops in the graph,"""
+        if not self.has_stage("to_edge_transform_and_lower"):
+            raise RuntimeError(
+                "count_tosa_ops requires to_edge_transform_and_lower in the pipeline."
+            )
+
+        def _count_tosa_ops():
+            stage = self.tester.stages[StageType.TO_EDGE_TRANSFORM_AND_LOWER]
+            graph_module = stage.graph_module
+            arm_tester_module.count_tosa_ops(graph_module, expected_ops)
+
+        self.add_stage_after(
+            "to_edge_transform_and_lower",
+            _count_tosa_ops,
+            suffix="tosa_ops",
+        )
+        return self
+
     def change_args(self, stage_id: str, *args, **kwargs):
         """Updates the args to the given stage id."""
         pos = self.find_pos(stage_id)