xref: /aosp_15_r20/external/executorch/backends/arm/test/ops/test_slice.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright 2024 Arm Limited and/or its affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import unittest
8from typing import Tuple
9
10import torch
11
12from executorch.backends.arm.test import common
13from executorch.backends.arm.test.tester.arm_tester import ArmTester
14from executorch.exir.backend.compile_spec_schema import CompileSpec
15from parameterized import parameterized
16
17
18class TestSimpleSlice(unittest.TestCase):
19
20    class Slice(torch.nn.Module):
21
22        sizes = [(10), (10, 10), (10, 10, 10), ((1, 12, 10, 10))]
23        test_tensors = [(torch.ones(n),) for n in sizes]
24
25        def forward(self, x: torch.Tensor):
26            if x.dim() == 1:
27                return x[3:-3]
28            elif x.dim() == 2:
29                return x[1:3, 3:]
30            elif x.dim() == 3:
31                return x[0:7, 0:, 0:8]
32            elif x.dim() == 4:
33                return x[:, :5, 3:5, 4:10]
34
35    def _test_slice_tosa_MI_pipeline(
36        self, module: torch.nn.Module, test_data: torch.Tensor
37    ):
38        (
39            ArmTester(
40                module,
41                example_inputs=test_data,
42                compile_spec=common.get_tosa_compile_spec("TOSA-0.80.0+MI"),
43            )
44            .export()
45            .check(["torch.ops.aten.slice.Tensor"])
46            .to_edge()
47            .check(["executorch_exir_dialects_edge__ops_aten_slice_copy"])
48            .partition()
49            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
50            .to_executorch()
51            .run_method_and_compare_outputs(inputs=test_data)
52        )
53
54    def _test_slice_tosa_BI_pipeline(
55        self, module: torch.nn.Module, test_data: Tuple[torch.Tensor], permute: bool
56    ):
57
58        (
59            ArmTester(
60                module,
61                example_inputs=test_data,
62                compile_spec=common.get_tosa_compile_spec(
63                    "TOSA-0.80.0+BI", permute_memory_to_nhwc=permute
64                ),
65            )
66            .quantize()
67            .export()
68            .check(["torch.ops.aten.slice.Tensor"])
69            .to_edge()
70            .partition()
71            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
72            .to_executorch()
73            .run_method_and_compare_outputs(inputs=test_data, qtol=1)
74        )
75
76    def _test_slice_ethos_BI_pipeline(
77        self,
78        compile_spec: list[CompileSpec],
79        module: torch.nn.Module,
80        test_data: Tuple[torch.Tensor],
81    ):
82        (
83            ArmTester(
84                module,
85                example_inputs=test_data,
86                compile_spec=common.get_u55_compile_spec(),
87            )
88            .quantize()
89            .export()
90            .check(["torch.ops.aten.slice.Tensor"])
91            .to_edge()
92            .partition()
93            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
94            .to_executorch()
95        )
96
97    def _test_slice_u55_BI_pipeline(
98        self, module: torch.nn.Module, test_data: Tuple[torch.Tensor]
99    ):
100        self._test_slice_ethos_BI_pipeline(
101            common.get_u55_compile_spec(), module, test_data
102        )
103
104    def _test_slice_u85_BI_pipeline(
105        self, module: torch.nn.Module, test_data: Tuple[torch.Tensor]
106    ):
107        self._test_slice_ethos_BI_pipeline(
108            common.get_u85_compile_spec(), module, test_data
109        )
110
111    @parameterized.expand(Slice.test_tensors)
112    def test_slice_tosa_MI(self, tensor):
113        self._test_slice_tosa_MI_pipeline(self.Slice(), (tensor,))
114
115    @parameterized.expand(Slice.test_tensors[:2])
116    def test_slice_nchw_tosa_BI(self, test_tensor: torch.Tensor):
117        self._test_slice_tosa_BI_pipeline(self.Slice(), (test_tensor,), False)
118
119    @parameterized.expand(Slice.test_tensors[2:])
120    def test_slice_nhwc_tosa_BI(self, test_tensor: torch.Tensor):
121        self._test_slice_tosa_BI_pipeline(self.Slice(), (test_tensor,), True)
122
123    @parameterized.expand(Slice.test_tensors)
124    def test_slice_u55_BI(self, test_tensor: torch.Tensor):
125        self._test_slice_u55_BI_pipeline(self.Slice(), (test_tensor,))
126
127    @parameterized.expand(Slice.test_tensors)
128    def test_slice_u85_BI(self, test_tensor: torch.Tensor):
129        self._test_slice_u85_BI_pipeline(self.Slice(), (test_tensor,))
130