Arm backend: Add squeeze-op (#5681)

Summary:
- Add squeeze op and unittest
- Update unsqueeze test

Pull Request resolved: #5681

Reviewed By: digantdesai

Differential Revision: D63637246

Pulled By: mergennachin

fbshipit-source-id: bf068d0305ea9f499f33ed27b5050eb09829b353

Author: oscarandersson8218

backends/arm/arm_partitioner.py

@@ -66,6 +66,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.clone.default,
             exir_ops.edge.aten.mean.dim,
             exir_ops.edge.aten.unsqueeze_copy.default,
+            exir_ops.edge.aten.squeeze_copy.dims,
             operator.getitem,
             exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
             exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,

backends/arm/operators/__init__.py

@@ -32,6 +32,7 @@
     op_sigmoid,
     op_slice,
     op_softmax,
+    op_squeeze,
     op_sub,
     op_unsqueeze,
     op_view,

backends/arm/operators/op_squeeze.py

@@ -0,0 +1,42 @@
+# Copyright 2024 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 List
+
+import serializer.tosa_serializer as ts
+import torch
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_utils import tosa_shape
+from serializer.tosa_serializer import TosaOp
+
+
+@register_node_visitor
+class SqueezeVisitor(NodeVisitor):
+    target = "aten.squeeze_copy.dims"
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+        shape = inputs[0].shape
+        rank = len(shape)
+        # In some cases, e.g. torch.randn((1, 5, 1, 5)).squeeze(),
+        # dims == [0, 1, 2, 3] even though all dims cannot be squeezed.
+        # We need to verify that shape[dim] == 1 before squeezing the dim.
+        dims = [dim % rank for dim in inputs[1].special if shape[dim] == 1]
+        new_shape = [shape[i] for i in range(rank) if i not in dims]
+        new_shape = tosa_shape(new_shape, output.dim_order)
+        attr = ts.TosaSerializerAttribute()
+        attr.ReshapeAttribute(new_shape)
+        tosa_graph.addOperator(
+            TosaOp.Op().RESHAPE, [inputs[0].name], [output.name], attr
+        )

backends/arm/quantizer/arm_quantizer_utils.py

@@ -146,7 +146,10 @@ def is_share_obs_or_fq_op(op: Callable) -> bool:
         torch.ops.aten.permute.default,
         torch.ops.aten.permute_copy.default,
         torch.ops.aten.squeeze.dim,
+        torch.ops.aten.squeeze.dims,
+        torch.ops.aten.squeeze.default,
         torch.ops.aten.squeeze_copy.dim,
+        torch.ops.aten.unsqueeze.default,
         # TODO: remove?
         torch.ops.aten.adaptive_avg_pool2d.default,
         torch.ops.aten.view_copy.default,

backends/arm/test/ops/test_squeeze.py

@@ -0,0 +1,221 @@
+# Copyright 2024 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.
+
+#
+# Tests the squeeze op which squeezes a given dimension with size 1 into a lower ranked tensor.
+#
+
+import unittest
+from typing import Optional, Tuple
+
+import torch
+
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    ArmQuantizer,
+    get_symmetric_quantization_config,
+)
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+
+from executorch.backends.xnnpack.test.tester.tester import Quantize
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from parameterized import parameterized
+
+
+class TestSqueeze(unittest.TestCase):
+    class SqueezeDim(torch.nn.Module):
+        test_parameters: list[tuple[torch.Tensor, int]] = [
+            (torch.randn(1, 1, 5), -2),
+            (torch.randn(1, 2, 3, 1), 3),
+            (torch.randn(1, 5, 1, 5), -2),
+        ]
+
+        def forward(self, x: torch.Tensor, dim: int):
+            return x.squeeze(dim)
+
+    class SqueezeDims(torch.nn.Module):
+        test_parameters: list[tuple[torch.Tensor, tuple[int]]] = [
+            (torch.randn(1, 5, 5, 1), (0, -1)),
+            (torch.randn(1, 5, 1, 5), (0, -2)),
+        ]
+
+        def forward(self, x: torch.Tensor, dims: tuple[int]):
+            return x.squeeze(dims)
+
+    class Squeeze(torch.nn.Module):
+        test_parameters: list[tuple[torch.Tensor]] = [
+            (torch.randn(1, 5, 5, 1),),
+            (torch.randn(1, 5, 1, 5),),
+        ]
+
+        def forward(self, x: torch.Tensor):
+            return x.squeeze()
+
+    def _test_squeeze_tosa_MI_pipeline(
+        self,
+        module: torch.nn.Module,
+        test_data: Tuple[torch.Tensor, Optional[tuple[int]]],
+        export_target: str,
+    ):
+        (
+            ArmTester(
+                module,
+                example_inputs=test_data,
+                compile_spec=common.get_tosa_compile_spec(permute_memory_to_nhwc=False),
+            )
+            .export()
+            .check_count({export_target: 1})
+            .to_edge()
+            .partition()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=test_data)
+        )
+
+    def _test_squeeze_tosa_BI_pipeline(
+        self,
+        module: torch.nn.Module,
+        test_data: Tuple[torch.Tensor, Optional[tuple[int]]],
+        export_target: str,
+    ):
+        quantizer = ArmQuantizer().set_io(get_symmetric_quantization_config())
+        (
+            ArmTester(
+                module,
+                example_inputs=test_data,
+                compile_spec=common.get_tosa_compile_spec(permute_memory_to_nhwc=False),
+            )
+            .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
+            .export()
+            .check_count({export_target: 1})
+            .to_edge()
+            .partition()
+            .dump_artifact()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .run_method_and_compare_outputs(inputs=test_data, qtol=1)
+        )
+
+    def _test_squeeze_ethosu_BI_pipeline(
+        self,
+        compile_spec: CompileSpec,
+        module: torch.nn.Module,
+        test_data: Tuple[torch.Tensor, Optional[tuple[int]]],
+        export_target: str,
+    ):
+        quantizer = ArmQuantizer().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({export_target: 1})
+            .to_edge()
+            .partition()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+        )
+
+    @parameterized.expand(Squeeze.test_parameters)
+    def test_squeeze_tosa_MI(
+        self,
+        test_tensor: torch.Tensor,
+    ):
+        self._test_squeeze_tosa_MI_pipeline(
+            self.Squeeze(), (test_tensor,), "torch.ops.aten.squeeze.default"
+        )
+
+    @parameterized.expand(Squeeze.test_parameters)
+    def test_squeeze_tosa_BI(
+        self,
+        test_tensor: torch.Tensor,
+    ):
+        self._test_squeeze_tosa_BI_pipeline(
+            self.Squeeze(), (test_tensor,), "torch.ops.aten.squeeze.default"
+        )
+
+    @parameterized.expand(Squeeze.test_parameters)
+    def test_squeeze_u55_BI(
+        self,
+        test_tensor: torch.Tensor,
+    ):
+        self._test_squeeze_ethosu_BI_pipeline(
+            common.get_u55_compile_spec(permute_memory_to_nhwc=False),
+            self.Squeeze(),
+            (test_tensor,),
+            "torch.ops.aten.squeeze.default",
+        )
+
+    @parameterized.expand(Squeeze.test_parameters)
+    def test_squeeze_u85_BI(
+        self,
+        test_tensor: torch.Tensor,
+    ):
+        self._test_squeeze_ethosu_BI_pipeline(
+            common.get_u85_compile_spec(permute_memory_to_nhwc=False),
+            self.Squeeze(),
+            (test_tensor,),
+            "torch.ops.aten.squeeze.default",
+        )
+
+    @parameterized.expand(SqueezeDim.test_parameters)
+    def test_squeeze_dim_tosa_MI(self, test_tensor: torch.Tensor, dim: int):
+        self._test_squeeze_tosa_MI_pipeline(
+            self.SqueezeDim(), (test_tensor, dim), "torch.ops.aten.squeeze.dim"
+        )
+
+    @parameterized.expand(SqueezeDim.test_parameters)
+    def test_squeeze_dim_tosa_BI(self, test_tensor: torch.Tensor, dim: int):
+        self._test_squeeze_tosa_BI_pipeline(
+            self.SqueezeDim(), (test_tensor, dim), "torch.ops.aten.squeeze.dim"
+        )
+
+    @parameterized.expand(SqueezeDim.test_parameters)
+    def test_squeeze_dim_u55_BI(self, test_tensor: torch.Tensor, dim: int):
+        self._test_squeeze_ethosu_BI_pipeline(
+            common.get_u55_compile_spec(permute_memory_to_nhwc=False),
+            self.SqueezeDim(),
+            (test_tensor, dim),
+            "torch.ops.aten.squeeze.dim",
+        )
+
+    @parameterized.expand(SqueezeDim.test_parameters)
+    def test_squeeze_dim_u85_BI(self, test_tensor: torch.Tensor, dim: int):
+        self._test_squeeze_ethosu_BI_pipeline(
+            common.get_u85_compile_spec(permute_memory_to_nhwc=False),
+            self.SqueezeDim(),
+            (test_tensor, dim),
+            "torch.ops.aten.squeeze.dim",
+        )
+
+    @parameterized.expand(SqueezeDims.test_parameters)
+    def test_squeeze_dims_tosa_MI(self, test_tensor: torch.Tensor, dims: tuple[int]):
+        self._test_squeeze_tosa_MI_pipeline(
+            self.SqueezeDims(), (test_tensor, dims), "torch.ops.aten.squeeze.dims"
+        )
+
+    @parameterized.expand(SqueezeDims.test_parameters)
+    def test_squeeze_dims_tosa_BI(self, test_tensor: torch.Tensor, dims: tuple[int]):
+        self._test_squeeze_tosa_BI_pipeline(
+            self.SqueezeDims(), (test_tensor, dims), "torch.ops.aten.squeeze.dims"
+        )
+
+    @parameterized.expand(SqueezeDims.test_parameters)
+    def test_squeeze_dims_u55_BI(self, test_tensor: torch.Tensor, dims: tuple[int]):
+        self._test_squeeze_ethosu_BI_pipeline(
+            common.get_u55_compile_spec(permute_memory_to_nhwc=False),
+            self.SqueezeDims(),
+            (test_tensor, dims),
+            "torch.ops.aten.squeeze.dims",
+        )
+
+    @parameterized.expand(SqueezeDims.test_parameters)
+    def test_squeeze_dims_u85_BI(self, test_tensor: torch.Tensor, dims: tuple[int]):
+        self._test_squeeze_ethosu_BI_pipeline(
+            common.get_u85_compile_spec(permute_memory_to_nhwc=False),
+            self.SqueezeDims(),
+            (test_tensor, dims),
+            "torch.ops.aten.squeeze.dims",
+        )

backends/arm/test/ops/test_unsqueeze.py

@@ -27,8 +27,8 @@
 
 class TestSimpleUnsqueeze(unittest.TestCase):
     class Unsqueeze(torch.nn.Module):
-        shapes: list[int | Sequence[int]] = [5, (5, 5), (5, 5), (5, 5, 5)]
-        test_parameters: list[tuple[torch.Tensor]] = [(torch.ones(n),) for n in shapes]
+        shapes: list[int | Sequence[int]] = [5, (5, 5), (5, 5), (5, 4, 3)]
+        test_parameters: list[tuple[torch.Tensor]] = [(torch.randn(n),) for n in shapes]
 
         def forward(self, x: torch.Tensor, dim):
             return x.unsqueeze(dim)
@@ -40,7 +40,7 @@ def _test_unsqueeze_tosa_MI_pipeline(
             ArmTester(
                 module,
                 example_inputs=test_data,
-                compile_spec=common.get_tosa_compile_spec(),
+                compile_spec=common.get_tosa_compile_spec(permute_memory_to_nhwc=False),
             )
             .export()
             .check_count({"torch.ops.aten.unsqueeze.default": 1})
@@ -59,7 +59,7 @@ def _test_unsqueeze_tosa_BI_pipeline(
             ArmTester(
                 module,
                 example_inputs=test_data,
-                compile_spec=common.get_tosa_compile_spec(),
+                compile_spec=common.get_tosa_compile_spec(permute_memory_to_nhwc=False),
             )
             .quantize(Quantize(quantizer, get_symmetric_quantization_config()))
             .export()
@@ -105,11 +105,15 @@ def test_unsqueeze_tosa_BI(self, test_tensor: torch.Tensor):
     @parameterized.expand(Unsqueeze.test_parameters)
     def test_unsqueeze_u55_BI(self, test_tensor: torch.Tensor):
         self._test_unsqueeze_ethosu_BI_pipeline(
-            common.get_u55_compile_spec(), self.Unsqueeze(), (test_tensor, 0)
+            common.get_u55_compile_spec(permute_memory_to_nhwc=False),
+            self.Unsqueeze(),
+            (test_tensor, 0),
         )
 
     @parameterized.expand(Unsqueeze.test_parameters)
     def test_unsqueeze_u85_BI(self, test_tensor: torch.Tensor):
         self._test_unsqueeze_ethosu_BI_pipeline(
-            common.get_u85_compile_spec(), self.Unsqueeze(), (test_tensor, 0)
+            common.get_u85_compile_spec(permute_memory_to_nhwc=False),
+            self.Unsqueeze(),
+            (test_tensor, 0),
         )

backends/arm/test/runner_utils.py

@@ -407,7 +407,7 @@ def run_tosa_ref_model(
 def prep_data_for_save(
     data, is_quantized: bool, input_name: str, quant_param: QuantizationParams
 ):
-    data_np = data.detach().numpy().astype(np.float32)
+    data_np = np.array(data.detach(), order="C").astype(np.float32)
 
     if is_quantized:
         assert (