Add a16w8 per-op test for bmm

backends/arm/test/ops/test_bmm.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-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.
@@ -10,11 +10,12 @@
 
 import torch
 
+from executorch.backends.arm.quantizer import get_symmetric_a16w8_quantization_config
 from executorch.backends.arm.test import common
-
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
     EthosU85PipelineINT,
+    OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
@@ -23,9 +24,6 @@
 aten_op_bmm = "torch.ops.aten.bmm.default"
 exir_op_bmm = "executorch_exir_dialects_edge__ops_aten_bmm_default"
 
-aten_op_mm = "torch.ops.aten.matmul.default"
-exir_op_mm = "executorch_exir_dialects_edge__ops_aten_matmul_default"
-
 input_t1 = Tuple[torch.Tensor, torch.Tensor]  # Input x
 
 
@@ -191,3 +189,52 @@ def test_bmm_vgf_quant_single_input(test_data: input_t1):
         quantize=True,
     )
     pipeline.run()
+
+
+a16w8_bmm_test_parameters = {
+    "rand_same": lambda: (torch.rand(2, 1, 1), torch.rand(2, 1, 1)),
+    "rand_diff": lambda: (torch.rand(5, 3, 5), torch.rand(5, 5, 2)),
+    "rand_rect": lambda: (torch.rand(1, 55, 3), torch.rand(1, 3, 44)),
+    "rand_batch10": lambda: (torch.rand(10, 1, 10), torch.rand(10, 10, 5)),
+    "rand_neg": lambda: (
+        -10 * torch.randn(2, 32, 64),
+        5 + 5 * torch.randn(2, 64, 32),
+    ),
+}
+
+
+@common.parametrize("test_data", a16w8_bmm_test_parameters)
+@common.XfailIfNoCorstone300
+def test_bmm_a16w8_u55_INT(test_data: input_t1):
+    """U55 does not support bmm with INT16 inputs.
+
+    Verify bmm is rejected.
+
+    """
+    pipeline = OpNotSupportedPipeline[input_t1](
+        BMM(),
+        test_data(),
+        non_delegated_ops={exir_op_bmm: 1},
+        n_expected_delegates=0,
+        u55_subset=True,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
+@common.parametrize("test_data", a16w8_bmm_test_parameters)
+@common.XfailIfNoCorstone320
+def test_bmm_a16w8_u85_INT(test_data: input_t1):
+    pipeline = EthosU85PipelineINT[input_t1](
+        BMM(),
+        test_data(),
+        aten_op_bmm,
+        exir_op_bmm,
+        a16w8_quantization=True,
+        symmetric_io_quantization=True,
+        qtol=1,
+        epsilon=2**-16,
+    )
+    pipeline.run()

backends/arm/test/targets.bzl

@@ -42,6 +42,7 @@ def define_arm_tests():
         "ops/test_var.py",
         "ops/test_conv1d.py",
         "ops/test_gelu.py",
+        "ops/test_bmm.py",
     ]
 
     # Quantization