Arm backend: Refactor test infra to enable TOSA V1.0 testing (part 3)

Signed-off-by: Saoirse Stewart <[email protected]>

Change-Id: I7a1896ba1d40d2bd0e313da0c33ff21ed8c0738c

Author: SaoirseARM

backends/arm/test/ops/test_lshift.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+from executorch.backends.arm.test import common
 from executorch.backends.arm.test.common import (
     XfailIfNoCorstone300,
     XfailIfNoCorstone320,
@@ -14,7 +15,6 @@
     TosaPipelineBI,
     TosaPipelineMI,
 )
-from parameterized import parameterized
 
 scalar_input_t = tuple[torch.Tensor, int]
 
@@ -23,11 +23,20 @@ class LshiftScalar(torch.nn.Module):
     torch_op_MI = "torch.ops.aten.__lshift__.Scalar"
     torch_op_BI = "torch.ops.aten.bitwise_left_shift.Tensor"
     exir_op = "executorch_exir_dialects_edge__ops_aten_bitwise_left_shift_Tensor"
-    test_data = [
-        ((torch.randint(-8, 8, (1, 12, 3, 4), dtype=torch.int8), 1),),
-        ((torch.randint(-100, 100, (1, 5, 3, 4), dtype=torch.int16), 5),),
-        ((torch.randint(-100, 100, (1, 5, 3, 4), dtype=torch.int32), 2),),
-    ]
+    test_data = {
+        "randint_neg_8_int8": (
+            torch.randint(-8, 8, (1, 12, 3, 4), dtype=torch.int8),
+            1,
+        ),
+        "randint_neg_100_int16": (
+            torch.randint(-100, 100, (1, 5, 3, 4), dtype=torch.int16),
+            5,
+        ),
+        "randint_neg_100_int32": (
+            torch.randint(-100, 100, (1, 5, 3, 4), dtype=torch.int32),
+            2,
+        ),
+    }
 
     def forward(self, x: torch.Tensor, shift: int):
         return x << shift
@@ -39,33 +48,27 @@ def forward(self, x: torch.Tensor, shift: int):
 class LshiftTensor(torch.nn.Module):
     torch_op = "torch.ops.aten.bitwise_left_shift.Tensor"
     exir_op = "executorch_exir_dialects_edge__ops_aten_bitwise_left_shift_Tensor"
-    test_data = [
-        (
-            (
-                torch.randint(-8, 8, (3, 3), dtype=torch.int8),
-                torch.randint(0, 4, (3, 3), dtype=torch.int8),
-            ),
+    test_data = {
+        "randint_neg_8_tensor_int8": (
+            torch.randint(-8, 8, (3, 3), dtype=torch.int8),
+            torch.randint(0, 4, (3, 3), dtype=torch.int8),
         ),
-        (
-            (
-                torch.randint(-1024, 1024, (3, 3, 3), dtype=torch.int16),
-                torch.randint(0, 5, (3, 3, 3), dtype=torch.int16),
-            ),
+        "randint_neg_1024_tensor_int16": (
+            torch.randint(-1024, 1024, (3, 3, 3), dtype=torch.int16),
+            torch.randint(0, 5, (3, 3, 3), dtype=torch.int16),
         ),
-        (
-            (
-                torch.randint(0, 127, (1, 2, 3, 3), dtype=torch.int32),
-                torch.randint(0, 5, (1, 2, 3, 3), dtype=torch.int32),
-            ),
+        "randint_0_tensor_int16": (
+            torch.randint(0, 127, (1, 2, 3, 3), dtype=torch.int32),
+            torch.randint(0, 5, (1, 2, 3, 3), dtype=torch.int32),
         ),
-    ]
+    }
 
     def forward(self, x: torch.Tensor, shift: torch.Tensor):
         return x.bitwise_left_shift(shift)
 
 
-@parameterized.expand(LshiftScalar.test_data)
-def test_lshift_scalar_tosa_MI(test_data):
+@common.parametrize("test_data", LshiftScalar.test_data)
+def test_lshift_scalar_tosa_MI_scalar(test_data):
     TosaPipelineMI[scalar_input_t](
         LshiftScalar(),
         test_data,
@@ -74,18 +77,21 @@ def test_lshift_scalar_tosa_MI(test_data):
     ).run()
 
 
-@parameterized.expand(LshiftScalar.test_data)
-def test_lshift_scalar_tosa_BI(test_data):
+@common.parametrize("test_data", LshiftScalar.test_data)
+def test_bitwise_left_shift_tensor_tosa_BI_scalar(test_data):
     pipeline = TosaPipelineBI[scalar_input_t](
-        LshiftScalar(), test_data, LshiftScalar.torch_op_BI, LshiftScalar.exir_op
+        LshiftScalar(),
+        test_data,
+        LshiftScalar.torch_op_BI,
+        LshiftScalar.exir_op,
     )
     pipeline.pop_stage("check.quant_nodes")
     pipeline.run()
 
 
-@parameterized.expand(LshiftScalar.test_data)
+@common.parametrize("test_data", LshiftScalar.test_data)
 @XfailIfNoCorstone300
-def test_lshift_scalar_tosa_u55(test_data):
+def test_bitwise_left_shift_tensor_u55_BI_scalar(test_data):
     pipeline = EthosU55PipelineBI[scalar_input_t](
         LshiftScalar(),
         test_data,
@@ -97,9 +103,9 @@ def test_lshift_scalar_tosa_u55(test_data):
     pipeline.run()
 
 
-@parameterized.expand(LshiftScalar.test_data)
+@common.parametrize("test_data", LshiftScalar.test_data)
 @XfailIfNoCorstone320
-def test_lshift_scalar_tosa_u85(test_data):
+def test_bitwise_left_shift_tensor_u85_BI_scalar(test_data):
     pipeline = EthosU85PipelineBI[scalar_input_t](
         LshiftScalar(),
         test_data,
@@ -111,8 +117,8 @@ def test_lshift_scalar_tosa_u85(test_data):
     pipeline.run()
 
 
-@parameterized.expand(LshiftTensor.test_data)
-def test_lshift_tensor_tosa_MI(test_data):
+@common.parametrize("test_data", LshiftTensor.test_data)
+def test_lshift_scalar_tosa_MI(test_data):
     TosaPipelineMI[scalar_input_t](
         LshiftTensor(),
         test_data,
@@ -121,18 +127,21 @@ def test_lshift_tensor_tosa_MI(test_data):
     ).run()
 
 
-@parameterized.expand(LshiftTensor.test_data)
-def test_lshift_tensor_tosa_BI(test_data):
+@common.parametrize("test_data", LshiftTensor.test_data)
+def test_bitwise_left_shift_tensor_tosa_BI(test_data):
     pipeline = TosaPipelineBI[scalar_input_t](
-        LshiftTensor(), test_data, LshiftTensor.torch_op, LshiftTensor.exir_op
+        LshiftTensor(),
+        test_data,
+        LshiftTensor.torch_op,
+        LshiftTensor.exir_op,
     )
     pipeline.pop_stage("check.quant_nodes")
     pipeline.run()
 
 
-@parameterized.expand(LshiftTensor.test_data)
+@common.parametrize("test_data", LshiftTensor.test_data)
 @XfailIfNoCorstone300
-def test_lshift_tensor_tosa_u55(test_data):
+def test_bitwise_left_shift_tensor_u55_BI(test_data):
     pipeline = EthosU55PipelineBI[scalar_input_t](
         LshiftTensor(),
         test_data,
@@ -144,9 +153,9 @@ def test_lshift_tensor_tosa_u55(test_data):
     pipeline.run()
 
 
-@parameterized.expand(LshiftTensor.test_data)
+@common.parametrize("test_data", LshiftTensor.test_data)
 @XfailIfNoCorstone320
-def test_lshift_tensor_tosa_u85(test_data):
+def test_bitwise_left_shift_tensor_u85_BI(test_data):
     pipeline = EthosU85PipelineBI[scalar_input_t](
         LshiftTensor(),
         test_data,

backends/arm/test/ops/test_repeat.py

@@ -1,5 +1,4 @@
 # Copyright 2024-2025 Arm Limited and/or its affiliates.
-# All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -8,115 +7,83 @@
 # Tests the repeat op which copies the data of the input tensor (possibly with new data format)
 #
 
-import unittest
+
 from typing import Sequence, Tuple
 
 import torch
 
-from executorch.backends.arm.quantizer import (
-    EthosUQuantizer,
-    get_symmetric_quantization_config,
-    TOSAQuantizer,
-)
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.arm_tester import ArmTester
-from executorch.backends.arm.tosa_specification import TosaSpecification
-
-from executorch.backends.xnnpack.test.tester.tester import Quantize
-from executorch.exir.backend.backend_details import CompileSpec
-from parameterized import parameterized
-
-
-class TestSimpleRepeat(unittest.TestCase):
-    """Tests Tensor.repeat for different ranks and dimensions."""
-
-    class Repeat(torch.nn.Module):
-        # (input tensor, multiples)
-        test_parameters = [
-            (torch.randn(3), (2,)),
-            (torch.randn(3, 4), (2, 1)),
-            (torch.randn(1, 1, 2, 2), (1, 2, 3, 4)),
-            (torch.randn(3), (2, 2)),
-            (torch.randn(3), (1, 2, 3)),
-            (torch.randn((3, 3)), (2, 2, 2)),
-            (torch.randn((3, 3, 3)), (2, 1, 2, 4)),
-        ]
-
-        def forward(self, x: torch.Tensor, multiples: Sequence):
-            return x.repeat(multiples)
-
-    def _test_repeat_tosa_MI_pipeline(self, module: torch.nn.Module, test_data: Tuple):
-        (
-            ArmTester(
-                module,
-                example_inputs=test_data,
-                compile_spec=common.get_tosa_compile_spec("TOSA-0.80+MI"),
-            )
-            .export()
-            .check_count({"torch.ops.aten.repeat.default": 1})
-            .to_edge()
-            .partition()
-            .check_not(["torch.ops.aten.repeat.default"])
-            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
-            .to_executorch()
-            .run_method_and_compare_outputs(inputs=test_data)
-        )
-
-    def _test_repeat_tosa_BI_pipeline(self, module: torch.nn.Module, test_data: Tuple):
-        tosa_spec = TosaSpecification.create_from_string("TOSA-0.80+BI")
-        compile_spec = common.get_tosa_compile_spec(tosa_spec)
-        quantizer = TOSAQuantizer(tosa_spec).set_io(get_symmetric_quantization_config())
-        (
-            ArmTester(module, example_inputs=test_data, compile_spec=compile_spec)
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
-            .export()
-            .check_count({"torch.ops.aten.repeat.default": 1})
-            .to_edge()
-            .partition()
-            .check_not(["torch.ops.aten.repeat.default"])
-            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
-            .to_executorch()
-            .run_method_and_compare_outputs(inputs=test_data, qtol=1)
-        )
-
-    def _test_repeat_ethosu_pipeline(
-        self, compile_spec: CompileSpec, module: torch.nn.Module, test_data: Tuple
-    ):
-        quantizer = EthosUQuantizer(compile_spec).set_io(
-            get_symmetric_quantization_config()
-        )
-        (
-            ArmTester(
-                module,
-                example_inputs=test_data,
-                compile_spec=compile_spec,
-            )
-            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
-            .export()
-            .check_count({"torch.ops.aten.repeat.default": 1})
-            .to_edge()
-            .partition()
-            .check_not(["torch.ops.aten.repeat.default"])
-            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
-            .to_executorch()
-        )
-
-    @parameterized.expand(Repeat.test_parameters)
-    def test_repeat_tosa_MI(self, test_input, multiples):
-        self._test_repeat_tosa_MI_pipeline(self.Repeat(), (test_input, multiples))
-
-    @parameterized.expand(Repeat.test_parameters)
-    def test_repeat_tosa_BI(self, test_input, multiples):
-        self._test_repeat_tosa_BI_pipeline(self.Repeat(), (test_input, multiples))
-
-    @parameterized.expand(Repeat.test_parameters)
-    def test_repeat_u55_BI(self, test_input, multiples):
-        self._test_repeat_ethosu_pipeline(
-            common.get_u55_compile_spec(), self.Repeat(), (test_input, multiples)
-        )
-
-    @parameterized.expand(Repeat.test_parameters)
-    def test_repeat_u85_BI(self, test_input, multiples):
-        self._test_repeat_ethosu_pipeline(
-            common.get_u85_compile_spec(), self.Repeat(), (test_input, multiples)
-        )
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineBI,
+    EthosU85PipelineBI,
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+input_t1 = Tuple[torch.Tensor, torch.Tensor]  # Input x, Input y
+aten_op = "torch.ops.aten.repeat.default"
+
+
+"""Tests Tensor.repeat for different ranks and dimensions."""
+
+
+class Repeat(torch.nn.Module):
+    # (input tensor, multiples)
+    test_parameters = {
+        "1_x_1": lambda: (torch.randn(3), (2,)),
+        "2_x_2": lambda: (torch.randn(3, 4), (2, 1)),
+        "4_x_4": lambda: (torch.randn(1, 1, 2, 2), (1, 2, 3, 4)),
+        "1_x_2": lambda: (torch.randn(3), (2, 2)),
+        "1_x_3": lambda: (torch.randn(3), (1, 2, 3)),
+        "2_x_3": lambda: (torch.randn((3, 3)), (2, 2, 2)),
+        "1_x_4": lambda: (torch.randn((3, 3, 3)), (2, 1, 2, 4)),
+    }
+
+    def forward(self, x: torch.Tensor, multiples: Sequence):
+        return x.repeat(multiples)
+
+
+@common.parametrize("test_data", Repeat.test_parameters)
+def test_repeat_tosa_MI(test_data: Tuple):
+    pipeline = TosaPipelineMI[input_t1](
+        Repeat(),
+        test_data(),
+        aten_op,
+        exir_op=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Repeat.test_parameters)
+def test_repeat_tosa_BI(test_data: Tuple):
+    pipeline = TosaPipelineBI[input_t1](
+        Repeat(),
+        test_data(),
+        aten_op,
+        exir_op=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Repeat.test_parameters)
+def test_repeat_u55_BI(test_data: Tuple):
+    pipeline = EthosU55PipelineBI[input_t1](
+        Repeat(),
+        test_data(),
+        aten_op,
+        exir_ops=[],
+        run_on_fvp=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Repeat.test_parameters)
+def test_repeat_u85_BI(test_data: Tuple):
+    pipeline = EthosU85PipelineBI[input_t1](
+        Repeat(),
+        test_data(),
+        aten_op,
+        exir_ops=[],
+        run_on_fvp=False,
+    )
+    pipeline.run()