Arm backend: Add count_program_io_kinds to Arm tester (#18405)

* Enables tests to check the ExportedProgram inputs/outputs are as
expected
* Re-enable fold_quantize path through ArmQuantizer to stop mutable
inputs in test_static_cache.py from being folded


Change-Id: I3b017923d6684aee784474d4894a81f4adc4c254



cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell

Signed-off-by: Tom Allsop <tom.allsop@arm.com>

Author: tom-arm

backends/arm/quantizer/arm_quantizer.py

@@ -808,17 +808,19 @@ def quantize_with_submodules(
         for _, submodule, _ in get_cond_while_submodules_nested(
             prepared, apply_quantization=True
         ):
-            converted = convert_pt2e(submodule)
+            converted = convert_pt2e(submodule, fold_quantize=fold_quantize)
             for submodule_node in submodule.graph.nodes:
                 if is_submodule_node(submodule_node):
                     for nested_name, nested_sub, _ in get_cond_while_submodules_nested(
                         submodule, apply_quantization=True
                     ):
                         converted.set_submodule(
-                            nested_name, convert_pt2e(nested_sub), strict=True
+                            nested_name,
+                            convert_pt2e(nested_sub, fold_quantize=fold_quantize),
+                            strict=True,
                         )
 
-        return convert_pt2e(prepared)
+        return convert_pt2e(prepared, fold_quantize=fold_quantize)
 
 
 class _TOSAQuantizerV1(Quantizer):

backends/arm/test/modules/test_static_cache.py

@@ -19,6 +19,8 @@
     TosaPipelineINT,
     VgfPipeline,
 )
+from torch.export.graph_signature import InputKind, OutputKind
+
 from transformers import LlamaConfig
 from transformers.cache_utils import StaticCache
 
@@ -34,6 +36,17 @@
 }
 
 
+EXPECTED_INPUT_COUNTS = {
+    InputKind.BUFFER: 2,
+    InputKind.USER_INPUT: 3,
+}
+
+EXPECTED_OUTPUT_COUNTS = {
+    OutputKind.BUFFER_MUTATION: 2,
+    OutputKind.USER_OUTPUT: 2,
+}
+
+
 @torch.no_grad()
 class StaticCacheModule(torch.nn.Module):
     def __init__(
@@ -121,15 +134,18 @@ def test_static_cache_tosa_FP(test_data):
         exir_op=[],
         transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()],
     )
+    pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS)
     pipeline.run()
 
 
+@pytest.mark.xfail(reason="BUFFER_MUTATION count mismatch: MLETORCH-1971")
 @common.parametrize("test_data", test_configs)
 def test_static_cache_tosa_INT(test_data):
     module = StaticCacheModule(test_data).eval()
     pipeline = TosaPipelineINT[input_t](
         module, module.get_inputs(), aten_op=[], exir_op=[], fold_quantize=False
     )
+    pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS)
     pipeline.run()
 
 
@@ -146,14 +162,22 @@ def test_static_cache_u55_INT(test_data):
     pipeline.run()
 
 
-@common.XfailIfNoCorstone320
 @common.parametrize(
     "test_data",
     test_configs,
     xfails={
-        "multihead_attention": "Incorrect numerical behavior: MLBEDSW-11589",
-        "grouped_query_attention": "Incorrect numerical behavior: MLBEDSW-11589",
-        "multi_query_attention": "Incorrect numerical behavior: MLBEDSW-11589",
+        "multihead_attention": (
+            "BUFFER_MUTATION count mismatch: MLETORCH-1971"
+            "Incorrect numerical behavior: MLBEDSW-11589"
+        ),
+        "grouped_query_attention": (
+            "BUFFER_MUTATION count mismatch: MLETORCH-1971"
+            "Incorrect numerical behavior: MLBEDSW-11589"
+        ),
+        "multi_query_attention": (
+            "BUFFER_MUTATION count mismatch: MLETORCH-1971"
+            "Incorrect numerical behavior: MLBEDSW-11589"
+        ),
     },
 )
 def test_static_cache_u85_INT(test_data):
@@ -166,6 +190,7 @@ def test_static_cache_u85_INT(test_data):
     )
     # U85: keep _to_dim_order_copy portable for int64->int32 cast of cache_position (not delegatable).
     pipeline.tester.use_portable_ops = True
+    pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS)
     pipeline.run()
 
 
@@ -181,10 +206,12 @@ def test_static_cache_vgf_no_quant(test_data):
         transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()],
         quantize=False,
     )
+    pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS)
     pipeline.run()
 
 
 @common.SkipIfNoModelConverter
+@pytest.mark.xfail(reason="BUFFER_MUTATION count mismatch: MLETORCH-1971")
 @common.parametrize("test_data", test_configs)
 def test_static_cache_vgf_quant(test_data):
     module = StaticCacheModule(test_data).eval()
@@ -197,4 +224,5 @@ def test_static_cache_vgf_quant(test_data):
         fold_quantize=False,
         tosa_spec="TOSA-1.0+INT",
     )
+    pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS)
     pipeline.run()

backends/arm/test/tester/arm_tester.py

@@ -90,7 +90,13 @@
 )
 from tabulate import tabulate  # type: ignore[import-untyped]
 
-from torch.export.graph_signature import ExportGraphSignature, InputSpec, OutputSpec
+from torch.export.graph_signature import (
+    ExportGraphSignature,
+    InputKind,
+    InputSpec,
+    OutputKind,
+    OutputSpec,
+)
 from torch.fx import Graph
 
 from torchao.quantization.pt2e.quantizer import QuantizationSpec, SharedQuantizationSpec
@@ -1272,3 +1278,39 @@ def count_tosa_ops(graph_module: torch.fx.GraphModule, expected_ops: Dict[str, i
                 raise AssertionError(
                     f"Expected {expected_count} occurrences of TOSA op {op} but found {actual_count}."
                 )
+
+
+def count_program_io_kinds(
+    exported_program: ExportedProgram,
+    expected_input_kinds: dict[InputKind, int] | None,
+    expected_output_kinds: dict[OutputKind, int] | None,
+):
+    """Checks that the number of InputKinds and OutputKinds in the final
+    ExportedProgram are equal to those in expected_inputs and
+    expected_outputs.
+    """
+
+    def check_spec_count(
+        spec: Sequence[InputSpec] | Sequence[OutputSpec],
+        expected_counts: dict[Any, int],
+    ):
+        kind = type(spec[0].kind)
+        counts: dict[Any, int] = Counter([s.kind for s in spec])
+
+        for e in kind:
+            if counts[e] != expected_counts.get(e, 0):
+                raise AssertionError(
+                    f"Expected to find {expected_counts[e]} for input/output kind {e}, but found {counts[e]}."
+                )
+
+    if expected_input_kinds:
+        check_spec_count(
+            exported_program.graph_signature.input_specs,
+            expected_input_kinds,
+        )
+
+    if expected_output_kinds:
+        check_spec_count(
+            exported_program.graph_signature.output_specs,
+            expected_output_kinds,
+        )

backends/arm/test/tester/test_pipeline.py

@@ -49,6 +49,7 @@
 from executorch.backends.test.harness.stages import StageType
 from executorch.exir.pass_base import ExportPass
 from torch._export.pass_base import PassType
+from torch.export.graph_signature import InputKind, OutputKind
 from torchao.quantization.pt2e.quantizer import QuantizationSpec
 
 logger = logging.getLogger(__name__)
@@ -322,6 +323,30 @@ def _count_tosa_ops():
         )
         return self
 
+    def count_program_io_kinds(
+        self,
+        expected_inputs: dict[InputKind, int] | None,
+        expected_outputs: dict[OutputKind, int] | None,
+    ):
+        """Assert the count of inputs/output specs of the ExportedProgram is
+        correct.
+        """
+        if not self.has_stage("to_edge_transform_and_lower"):
+            raise RuntimeError(
+                "count_program_io requires to_edge_transform_and_lower in the pipeline."
+            )
+
+        def _count_program_io_kinds():
+            stage = self.tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER)
+            arm_tester_module.count_program_io_kinds(
+                stage.exported_program(), expected_inputs, expected_outputs
+            )
+
+        self.add_stage_after(
+            "to_edge_transform_and_lower", _count_program_io_kinds, suffix="count_io"
+        )
+        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)