Arm backend: Added 8 new unit tests for testing various passes. (#9037)

Refactored tests cases to use TestPassPipeline instead of ArmTester.

Signed-off-by: Michiel Olieslagers <[email protected]>

Author: Michiel-Olieslagers

backends/arm/test/passes/test_convert_expand_copy_to_repeat.py

@@ -0,0 +1,51 @@
+# Copyright 2025 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.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes.convert_expand_copy_to_repeat import (
+    ConvertExpandCopyToRepeatPass,
+)
+
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+
+input_t = Tuple[torch.Tensor]  # Input x
+
+
+class Expand(torch.nn.Module):
+    """
+    Basic expand model using torch.Tensor.expand function
+    """
+
+    def __init__(self):
+        super(Expand, self).__init__()
+
+    def forward(self, x):
+        return x.expand(3, 4)
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(3, 1),)
+
+
+def test_expand_to_repeat_tosa_BI():
+    module = Expand()
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        tosa_version="TOSA-0.80+BI",
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_expand_copy_default": 1,
+        },
+        ops_not_before_pass=["executorch_exir_dialects_edge__ops_aten_repeat_default"],
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_repeat_default": 1,
+        },
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_aten_expand_copy_default"
+        ],
+        pass_list=[ConvertExpandCopyToRepeatPass],
+    )
+    pipeline.run()

backends/arm/test/passes/test_convert_split_to_slice.py

@@ -0,0 +1,67 @@
+# Copyright 2025 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.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes.convert_split_to_slice import (
+    ConvertSplitToSlicePass,
+)
+
+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 Split(torch.nn.Module):
+    """
+    Basic split model using torch.split function
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(10),)
+
+    def forward(self, x):
+        return torch.split(x, 2)
+
+
+class SplitTensor(torch.nn.Module):
+    """
+    Basic split model using torch.Tensor.split function
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(10),)
+
+    def forward(self, x):
+        return x.split(2)
+
+
+modules = {"split_basic": Split(), "split_tensor": SplitTensor()}
+
+
+@common.parametrize("module", modules)
+def test_split_to_slice_tosa_BI(module):
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        tosa_version="TOSA-0.80+BI",
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default": 1,
+        },
+        ops_not_before_pass=[
+            "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"
+        ],
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 5,
+        },
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default"
+        ],
+        pass_list=[ConvertSplitToSlicePass],
+    )
+    pipeline.run()

backends/arm/test/passes/test_decompose_div_pass.py

@@ -0,0 +1,65 @@
+# Copyright 2025 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.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes.decompose_div_pass import DecomposeDivPass
+
+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 Div(torch.nn.Module):
+    """
+    Basic div model using torch.div
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(10),)
+
+    def forward(self, x):
+        return torch.div(x, 2)
+
+
+class DivTensor(torch.nn.Module):
+    """
+    Basic div model using torch.Tensor.div
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(10),)
+
+    def forward(self, x):
+        return x.div(2)
+
+
+modules = {"div_basic": Div(), "div_tensor": DivTensor()}
+
+
+@common.parametrize("module", modules)
+def test_decompose_div_tosa_MI(module):
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        tosa_version="TOSA-0.80+MI",
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_div_Tensor": 1,
+        },
+        ops_not_before_pass=[
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor",
+            "executorch_exir_dialects_edge__ops_aten_reciprocal_default",
+        ],
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1,
+            "executorch_exir_dialects_edge__ops_aten_reciprocal_default": 1,
+        },
+        ops_not_after_pass=["executorch_exir_dialects_edge__ops_aten_div_Tensor"],
+        pass_list=[DecomposeDivPass],
+    )
+    pipeline.run()

backends/arm/test/passes/test_decompose_layernorm_pass.py

@@ -0,0 +1,69 @@
+# Copyright 2025 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.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes.decompose_layernorm_pass import (
+    DecomposeLayerNormPass,
+)
+
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+
+input_t = Tuple[torch.Tensor]  # Input x
+
+
+class LayerNorm(torch.nn.Module):
+    """
+    Basic layer_norm model using torch.nn.layer_norm layer
+    """
+
+    def __init__(self):
+        super(LayerNorm, self).__init__()
+        self.layer_norm = torch.nn.LayerNorm(10)
+
+    def forward(self, x):
+        x = self.layer_norm(x)
+        return x
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(10),)
+
+
+def test_decompose_layernorm_tosa_MI():
+    module = LayerNorm()
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        tosa_version="TOSA-0.80+MI",
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_native_layer_norm_default": 1,
+        },
+        ops_not_before_pass=[
+            "executorch_exir_dialects_edge__ops_aten_add_Tensor",
+            "executorch_exir_dialects_edge__ops_aten_view_copy_default",
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor",
+            "executorch_exir_dialects_edge__ops_aten_full_default",
+            "executorch_exir_dialects_edge__ops_aten_rsqrt_default",
+            "executorch_exir_dialects_edge__ops_aten_var_correction",
+            "executorch_exir_dialects_edge__ops_aten_sub_Tensor",
+            "executorch_exir_dialects_edge__ops_aten_mean_dim",
+        ],
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_add_Tensor": 2,
+            "executorch_exir_dialects_edge__ops_aten_view_copy_default": 2,
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 2,
+            "executorch_exir_dialects_edge__ops_aten_full_default": 1,
+            "executorch_exir_dialects_edge__ops_aten_rsqrt_default": 1,
+            "executorch_exir_dialects_edge__ops_aten_var_correction": 1,
+            "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1,
+            "executorch_exir_dialects_edge__ops_aten_mean_dim": 1,
+        },
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_aten_expand_copy_default"
+        ],
+        pass_list=[DecomposeLayerNormPass],
+    )
+    pipeline.run()

backends/arm/test/passes/test_decompose_meandim_pass.py

@@ -0,0 +1,73 @@
+# Copyright 2025 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.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes.decompose_meandim_pass import DecomposeMeanDimPass
+
+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 MeanDim(torch.nn.Module):
+    """
+    Basic mean model using torch.mean function making sure keepdim=True (keepdim=False doesnt work for this pass for some reason)
+    """
+
+    def __init__(self):
+        super(MeanDim, self).__init__()
+
+    def forward(self, x):
+        return torch.mean(x, 1, True)
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(4, 4),)
+
+
+class MeanDimTensor(torch.nn.Module):
+    """
+    Basic mean model using torch.Tensor.mean function making sure keepdim=True (keepdim=False doesnt work for this pass for some reason)
+    """
+
+    def __init__(self):
+        super(MeanDimTensor, self).__init__()
+
+    def forward(self, x):
+        return x.mean(1, True)
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(4, 4),)
+
+
+modules = {"meandim_basic": MeanDim(), "meandim_tensor": MeanDimTensor()}
+
+
+@common.parametrize("module", modules)
+def test_decompose_meandim_tosa_MI(module):
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        tosa_version="TOSA-0.80+MI",
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_mean_dim": 1,
+        },
+        ops_not_before_pass=[
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor",
+            "executorch_exir_dialects_edge__ops_aten_full_default",
+            "executorch_exir_dialects_edge__ops_aten_sum_dim_IntList",
+        ],
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1,
+            "executorch_exir_dialects_edge__ops_aten_full_default": 1,
+            "executorch_exir_dialects_edge__ops_aten_sum_dim_IntList": 1,
+        },
+        ops_not_after_pass=["executorch_exir_dialects_edge__ops_aten_mean_dim"],
+        pass_list=[DecomposeMeanDimPass],
+    )
+    pipeline.run()