Arm backend: Refactor pass tests for TOSA V1.0 (#10843)

PassPipelines will handle tosa_version aligned with other test
pipelines.

Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/test/passes/test_cast_int64_pass.py

@@ -6,37 +6,38 @@
 from typing import Tuple
 
 import torch
-from executorch.backends.arm._passes.cast_int64_pass import CastInt64BuffersToInt32Pass
+from executorch.backends.arm._passes import CastInt64BuffersToInt32Pass
 
+from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 input_t = Tuple[torch.Tensor]  # Input x
 
 
 class Int64Model(torch.nn.Module):
+    test_data = {
+        "rand": (torch.rand(4),),
+    }
 
     def forward(self, x: torch.Tensor):
         return x + 3
 
-    def get_inputs(self) -> input_t:
-        return (torch.rand(4),)
-
 
-def test_int64_model_tosa_BI():
+@common.parametrize("test_data", Int64Model.test_data)
+def test_int64_model(test_data: input_t):
     module = Int64Model()
     op_checks = {
         "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1,
         "executorch_exir_dialects_edge__ops_aten_add_Tensor": 1,
     }
     pipeline = PassPipeline[input_t](
         module,
-        module.get_inputs(),
-        tosa_version="TOSA-0.80+BI",
+        test_data,
+        quantize=False,
         ops_before_pass=op_checks,
         ops_after_pass=op_checks,
         passes_with_exported_program=[CastInt64BuffersToInt32Pass],
     )
-    pipeline.pop_stage("quantize")
     pipeline.run()
 
     exported_program = pipeline.tester.get_artifact("RunPasses").exported_program()

backends/arm/test/passes/test_convert_expand_copy_to_repeat.py

@@ -35,7 +35,7 @@ def test_expand_to_repeat_tosa_BI():
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+BI",
+        quantize=True,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten_expand_copy_default": 1,
         },

backends/arm/test/passes/test_convert_split_to_slice.py

@@ -49,7 +49,7 @@ def test_split_to_slice_tosa_BI(module):
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+BI",
+        quantize=True,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default": 1,
         },

backends/arm/test/passes/test_convert_to_clamp.py

@@ -3,18 +3,21 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import unittest
+
+from typing import Tuple
 
 import torch
 from executorch.backends.arm._passes.convert_to_clamp import ConvertToClampPass
 
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.arm_tester import ArmTester
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
-from executorch.backends.xnnpack.test.tester.tester import RunPasses
+input_t = Tuple[torch.Tensor]  # Input x
 
 
 class HardTanh(torch.nn.Module):
+    test_data = {"rand": (torch.rand(1, 64, 64, 3),)}
+
     def __init__(self):
         super().__init__()
 
@@ -23,11 +26,10 @@ def __init__(self):
     def forward(self, x):
         return self.hardtanh(x)
 
-    def get_inputs(self):
-        return (torch.rand(1, 64, 64, 3),)
-
 
 class ReLU(torch.nn.Module):
+    test_data = {"rand": (torch.rand(1, 64, 64, 3),)}
+
     def __init__(self):
         super().__init__()
 
@@ -36,45 +38,55 @@ def __init__(self):
     def forward(self, x):
         return self.relu(x)
 
-    def get_inputs(self):
-        return (torch.rand(1, 64, 64, 3),)
-
-
-class TestConvertToClampPass(unittest.TestCase):
-    """
-    Tests the ConvertToClampPass which converts hardtanh.default and relu.default to clamp.default
-    """
-
-    def test_tosa_MI_hardtahn(self):
-        module = HardTanh()
-        test_pass_stage = RunPasses([ConvertToClampPass])
-        (
-            ArmTester(
-                module,
-                example_inputs=module.get_inputs(),
-                compile_spec=common.get_tosa_compile_spec("TOSA-0.80+MI"),
-            )
-            .export()
-            .to_edge()
-            .check(["executorch_exir_dialects_edge__ops_aten_hardtanh_default"])
-            .run_passes(test_pass_stage)
-            .check(["executorch_exir_dialects_edge__ops_aten_clamp_default"])
-            .check_not(["executorch_exir_dialects_edge__ops_aten_hardtanh_default"])
-        )
-
-    def test_tosa_MI_relu(self):
-        module = ReLU()
-        test_pass_stage = RunPasses([ConvertToClampPass])
-        (
-            ArmTester(
-                module,
-                example_inputs=module.get_inputs(),
-                compile_spec=common.get_tosa_compile_spec("TOSA-0.80+MI"),
-            )
-            .export()
-            .to_edge()
-            .check(["executorch_exir_dialects_edge__ops_aten_relu_default"])
-            .run_passes(test_pass_stage)
-            .check(["executorch_exir_dialects_edge__ops_aten_clamp_default"])
-            .check_not(["executorch_exir_dialects_edge__ops_aten_relu_default"])
-        )
+
+"""
+Tests the ConvertToClampPass which converts hardtanh.default and relu.default to clamp.default
+"""
+
+
+@common.parametrize("test_data", HardTanh.test_data)
+def test_tosa_MI_hardtahn(test_data: input_t):
+    module = HardTanh()
+    op_checks_before_pass = {
+        "executorch_exir_dialects_edge__ops_aten_hardtanh_default": 1,
+    }
+    op_checks_after_pass = {
+        "executorch_exir_dialects_edge__ops_aten_clamp_default": 1,
+    }
+    op_checks_not_after_pass = [
+        "executorch_exir_dialects_edge__ops_aten_hardtanh_default",
+    ]
+    pipeline = PassPipeline[input_t](
+        module,
+        test_data,
+        quantize=False,
+        ops_before_pass=op_checks_before_pass,
+        ops_after_pass=op_checks_after_pass,
+        ops_not_after_pass=op_checks_not_after_pass,
+        pass_list=[ConvertToClampPass],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", ReLU.test_data)
+def test_tosa_MI_relu(test_data: input_t):
+    module = ReLU()
+    op_checks_before_pass = {
+        "executorch_exir_dialects_edge__ops_aten_relu_default": 1,
+    }
+    op_checks_after_pass = {
+        "executorch_exir_dialects_edge__ops_aten_clamp_default": 1,
+    }
+    op_checks_not_after_pass = [
+        "executorch_exir_dialects_edge__ops_aten_relu_default",
+    ]
+    pipeline = PassPipeline[input_t](
+        module,
+        test_data,
+        quantize=False,
+        ops_before_pass=op_checks_before_pass,
+        ops_after_pass=op_checks_after_pass,
+        ops_not_after_pass=op_checks_not_after_pass,
+        pass_list=[ConvertToClampPass],
+    )
+    pipeline.run()

backends/arm/test/passes/test_decompose_cosine_similarity_pass.py

@@ -42,11 +42,11 @@ def test_decompose_cosine_similarity_tosa_BI(module):
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+BI",
         ops_before_pass=None,
         ops_not_before_pass=None,
         ops_after_pass=ops_after_pass,
         ops_not_after_pass=None,
         pass_list=[DecomposeCosineSimilarityPass],
+        quantize=True,
     )
     pipeline.run()

backends/arm/test/passes/test_decompose_div_pass.py

@@ -47,7 +47,7 @@ def test_decompose_div_tosa_MI(module):
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten_div_Tensor": 1,
         },

backends/arm/test/passes/test_decompose_layernorm_pass.py

@@ -37,7 +37,7 @@ def test_decompose_layernorm_tosa_MI():
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten_native_layer_norm_default": 1,
         },

backends/arm/test/passes/test_decompose_meandim_pass.py

@@ -53,7 +53,7 @@ def test_decompose_meandim_tosa_MI(module):
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten_mean_dim": 1,
         },

backends/arm/test/passes/test_decompose_softmax_pass.py

@@ -52,7 +52,7 @@ def test_softmax_basic_tosa_MI():
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten__softmax_default": 1,
         },
@@ -79,7 +79,7 @@ def test_softmax_log_tosa_MI():
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten__log_softmax_default": 1,
         },

backends/arm/test/passes/test_decompose_var_pass.py

@@ -60,7 +60,7 @@ def test_decompose_var_tosa_MI(module):
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_aten_var_correction": 1,
         },

backends/arm/test/passes/test_fold_qdq_pass.py

@@ -7,21 +7,24 @@
 
 import torch
 from executorch.backends.arm._passes import FoldAndAnnotateQParamsPass
+from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 
 input_t = Tuple[torch.Tensor, torch.Tensor]  # Input x, y
 
 
 class SimpleQuantizeModel(torch.nn.Module):
+    test_data = {
+        "rand": (torch.rand(1, 1280, 7, 7), torch.rand(1, 1280, 7, 7)),
+    }
+
     def forward(self, x, y):
         return x + torch.max((x + x), (y + y))
 
-    def get_inputs(self) -> input_t:
-        return (torch.rand(1, 1280, 7, 7), torch.rand(1, 1280, 7, 7))
-
 
-def test_fold_qdq_pass_tosa_BI():
+@common.parametrize("test_data", SimpleQuantizeModel.test_data)
+def test_fold_qdq_pass_tosa_BI(test_data: input_t):
     """
     Tests the FoldAndAnnotateQParamsPass which folds dq/q nodes into
     the node and stores the quantization parameters in meta.
@@ -32,8 +35,8 @@ def test_fold_qdq_pass_tosa_BI():
     module = SimpleQuantizeModel()
     pipeline = PassPipeline[input_t](
         module,
-        module.get_inputs(),
-        tosa_version="TOSA-0.80+BI",
+        test_data,
+        quantize=True,
         ops_before_pass={
             "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 7,
             "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 6,

backends/arm/test/passes/test_fuse_batchnorm_pass.py

@@ -136,12 +136,12 @@ def forward(self, x):
 
 
 @common.parametrize("module", modules)
-def test_fuse_batchnorm_tosa_MI(module):
+def test_fuse_batchnorm_tosa_MI(module: torch.nn.Module):
     """Test various cases where the batchnorm should and shouldn't be fused."""
     pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass=module.ops_before_pass,
         ops_after_pass=module.ops_after_pass,
         passes_with_exported_program=[FuseBatchnorm2DPass],

backends/arm/test/passes/test_fuse_constant_ops_pass.py

@@ -98,11 +98,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 @common.parametrize("module", modules)
-def test_fuse_const_ops_tosa_MI(module):
+def test_fuse_const_ops_tosa_MI(module: torch.nn.Module):
     pipeline = PassPipeline[input_t](
         module=module,
         test_data=(torch.rand(1),),
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass=module.ops_before_pass,
         ops_after_pass=module.ops_after_pass,
         ops_not_after_pass=module.ops_not_after_pass,
@@ -113,8 +113,13 @@ def test_fuse_const_ops_tosa_MI(module):
 
 @unittest.skip("Test failing on internal CI")
 @common.parametrize("module", modules)
-def test_fuse_const_ops_tosa_BI(module):
+def test_fuse_const_ops_tosa_BI(module: torch.nn.Module):
     pipeline = TosaPipelineBI[input_t](
-        module, (torch.rand(10, 10),), [], [], use_to_edge_transform_and_lower=True
+        module,
+        (torch.rand(10, 10),),
+        [],
+        [],
+        quantize=True,
+        use_to_edge_transform_and_lower=True,
     )
     pipeline.run()

backends/arm/test/passes/test_fuse_equal_placeholders_ops_pass.py

@@ -60,7 +60,7 @@ def test_fuse_equal_placeholders_constants_tosa_MI():
     pipeline = PassPipeline[input_t](
         module,
         data,
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass=module.ops_before_pass,
         ops_after_pass=module.ops_after_pass,
         passes_with_exported_program=[FuseEqualPlaceholdersPass],
@@ -81,7 +81,7 @@ def test_fuse_equal_placeholders_state_dict_tosa_MI():
     pipeline = PassPipeline[input_t](
         module,
         data,
-        tosa_version="TOSA-0.80+MI",
+        quantize=False,
         ops_before_pass=module.ops_before_pass,
         ops_after_pass=module.ops_after_pass,
         passes_with_exported_program=[FuseEqualPlaceholdersPass],