xref: /aosp_15_r20/external/executorch/backends/arm/test/ops/test_conv1d.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
8
9from typing import List, Optional, Tuple, Union
10
11import torch
12from executorch.backends.arm.test import common
13
14from executorch.backends.arm.test.tester.arm_tester import ArmTester
15from executorch.exir.backend.backend_details import CompileSpec
16from parameterized import parameterized
17
18
19class Conv1d(torch.nn.Module):
20    """
21    Creates one or many chained 1D-convolutions. For multiple convolutions, the
22    respective parameteres are provided as lists.
23    """
24
25    def __init__(
26        self,
27        inputs: Optional[torch.Tensor] = None,
28        length=8,
29        nbr_conv=1,  # Number of chained convs
30        in_channels: Union[List, int, None] = None,
31        out_channels: Union[List, int, None] = None,
32        kernel_size: Union[List, Tuple, None] = None,
33        stride: Union[List, Tuple, None] = None,
34        padding: Union[List, Tuple, None] = None,
35        dilation: Union[List, Tuple, None] = None,
36        groups: Union[List, int, None] = None,
37        bias: Union[List, bool, None] = None,
38        padding_mode: Union[List, str, None] = None,
39        batches=1,
40        dtype=torch.float32,
41    ):
42        super().__init__()
43        self.nbr_convs = nbr_conv
44
45        # Handle default values
46        in_channels = [2] * nbr_conv if in_channels is None else in_channels
47        out_channels = [1 * nbr_conv] if out_channels is None else out_channels
48        kernel_size = [3] * nbr_conv if kernel_size is None else kernel_size
49        stride = [2] * nbr_conv if stride is None else stride
50        padding = [1] * nbr_conv if padding is None else padding
51        dilation = [1] * nbr_conv if dilation is None else dilation
52        groups = [1] * nbr_conv if groups is None else groups
53        bias = [True] * nbr_conv if bias is None else bias
54        padding_mode = ["zeros"] * nbr_conv if padding_mode is None else padding_mode
55
56        # This allows the input parameters to be either a single value or a list
57        # as type hint implies
58        if not isinstance(in_channels, List):
59            in_channels = [in_channels]
60        if not isinstance(out_channels, List):
61            out_channels = [out_channels]
62        if not isinstance(kernel_size, List):
63            kernel_size = [kernel_size]
64        if not isinstance(stride, List):
65            stride = [stride]
66        if not isinstance(padding, List):
67            padding = [padding]
68        if not isinstance(dilation, List):
69            dilation = [dilation]
70        if not isinstance(groups, List):
71            groups = [groups]
72        if not isinstance(bias, List):
73            bias = [bias]
74        if not isinstance(padding_mode, List):
75            padding_mode = [padding_mode]
76
77        # Generate test data if not provided
78        if inputs is None:
79            self.inputs = (torch.randn(batches, in_channels[0], length).to(dtype),)
80        else:
81            self.inputs = (inputs,)
82
83        # Build chain of convs
84        for i in range(self.nbr_convs):
85            setattr(
86                self,
87                f"conv_{i}",
88                torch.nn.Conv1d(
89                    in_channels=in_channels[i],
90                    out_channels=out_channels[i],
91                    kernel_size=kernel_size[i],
92                    stride=stride[i],
93                    padding=padding[i],
94                    dilation=dilation[i],
95                    groups=groups[i],
96                    bias=bias[i],
97                    padding_mode=padding_mode[i],
98                ).to(dtype),
99            )
100
101    def get_inputs(self):
102        return self.inputs
103
104    def forward(self, x):
105        for i in range(self.nbr_convs):
106            conv = getattr(self, f"conv_{i}")
107            x = conv(x)
108        return x
109
110
111conv1d_2_3x2x40_nobias = Conv1d(
112    in_channels=2,
113    out_channels=3,
114    kernel_size=2,
115    stride=1,
116    bias=False,
117    padding=0,
118    length=40,
119    batches=1,
120)
121
122conv1d_3_1x3x256_st1 = Conv1d(
123    in_channels=3,
124    out_channels=10,
125    kernel_size=3,
126    stride=1,
127    padding=0,
128    length=256,
129    batches=1,
130)
131
132conv1d_3_1x3x12_st2_pd1 = Conv1d(
133    in_channels=3,
134    out_channels=4,
135    kernel_size=3,
136    stride=2,
137    padding=1,
138    length=12,
139    batches=1,
140)
141
142conv1d_1_1x2x128_st1 = Conv1d(
143    in_channels=2,
144    out_channels=1,
145    kernel_size=1,
146    stride=1,
147    padding=0,
148    length=128,
149    batches=1,
150)
151
152conv1d_2_1x2x14_st2 = Conv1d(
153    in_channels=2,
154    out_channels=1,
155    kernel_size=2,
156    stride=2,
157    padding=0,
158    length=14,
159    batches=1,
160)
161
162conv1d_5_3x2x128_st1 = Conv1d(
163    in_channels=2,
164    out_channels=3,
165    kernel_size=5,
166    stride=1,
167    padding=0,
168    length=128,
169    batches=3,
170)
171
172conv1d_3_1x3x224_st2_pd1 = Conv1d(
173    in_channels=3,
174    out_channels=16,
175    kernel_size=3,
176    stride=2,
177    padding=1,
178    length=224,
179    batches=1,
180)
181
182two_conv1d_nobias = Conv1d(
183    nbr_conv=2,
184    length=256,
185    in_channels=[3, 10],
186    out_channels=[10, 15],
187    kernel_size=[5, 5],
188    stride=[1, 1],
189    padding=[0, 0],
190    bias=[False, False],
191    batches=1,
192)
193
194two_conv1d = Conv1d(
195    nbr_conv=2,
196    length=256,
197    in_channels=[3, 10],
198    out_channels=[10, 15],
199    kernel_size=[5, 5],
200    stride=[1, 1],
201    padding=[0, 0],
202    bias=[True, True],
203    batches=1,
204)
205
206# Shenanigan to get a nicer output when test fails. With unittest it looks like:
207# FAIL: test_conv1d_tosa_BI_2_3x3_1x3x12x12_st2_pd1
208testsuite = [
209    ("2_3x2x40_nobias", conv1d_2_3x2x40_nobias),
210    ("3_1x3x256_st1", conv1d_3_1x3x256_st1),
211    ("3_1x3x12_st2_pd1", conv1d_3_1x3x12_st2_pd1),
212    ("1_1x2x128_st1", conv1d_1_1x2x128_st1),
213    ("2_1x2x14_st2", conv1d_2_1x2x14_st2),
214    ("5_3x2x128_st1", conv1d_5_3x2x128_st1),
215    ("3_1x3x224_st2_pd1", conv1d_3_1x3x224_st2_pd1),
216    ("two_conv1d_nobias", two_conv1d_nobias),
217    ("two_conv1d", two_conv1d),
218]
219
220
221class TestConv1D(unittest.TestCase):
222    def _test_conv1d_tosa_MI_pipeline(
223        self, module: torch.nn.Module, test_data: Tuple[torch.Tensor]
224    ):
225        (
226            ArmTester(
227                module,
228                example_inputs=test_data,
229                compile_spec=common.get_tosa_compile_spec(
230                    "TOSA-0.80.0+MI", permute_memory_to_nhwc=True
231                ),
232            )
233            .export()
234            .to_edge()
235            .partition()
236            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
237            .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
238            .to_executorch()
239            .run_method_and_compare_outputs(inputs=test_data)
240        )
241
242    def _test_conv1d_tosa_BI_pipeline(
243        self,
244        module: torch.nn.Module,
245        test_data: Tuple[torch.Tensor],
246    ):
247        (
248            ArmTester(
249                module,
250                example_inputs=test_data,
251                compile_spec=common.get_tosa_compile_spec(
252                    "TOSA-0.80.0+BI", permute_memory_to_nhwc=True
253                ),
254            )
255            .quantize()
256            .export()
257            .to_edge()
258            .partition()
259            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
260            .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
261            .to_executorch()
262            .run_method_and_compare_outputs(inputs=test_data, qtol=1)
263        )
264
265    def _test_conv1d_ethosu_BI_pipeline(
266        self,
267        module: torch.nn.Module,
268        compile_spec: CompileSpec,
269        test_data: Tuple[torch.Tensor],
270    ):
271        (
272            ArmTester(module, example_inputs=test_data, compile_spec=compile_spec)
273            .quantize()
274            .export()
275            .to_edge()
276            .partition()
277            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
278            .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
279            .to_executorch()
280        )
281
282    @parameterized.expand(testsuite)
283    def test_conv1d_tosa_MI(self, test_name, model):
284        self._test_conv1d_tosa_MI_pipeline(model, model.get_inputs())
285
286    @parameterized.expand(testsuite)
287    def test_conv1d_tosa_BI(self, test_name, model):
288        self._test_conv1d_tosa_BI_pipeline(model, model.get_inputs())
289
290    # Expeted to fail as Conv1D requires transpoes which isn't supported on u55
291    @parameterized.expand(testsuite)
292    @unittest.expectedFailure
293    def test_conv1d_u55_BI(self, test_name, model):
294        self._test_conv1d_ethosu_BI_pipeline(
295            model, common.get_u55_compile_spec(), model.get_inputs()
296        )
297
298    @parameterized.expand(testsuite)
299    def test_conv1d_u85_BI(self, test_name, model):
300        self._test_conv1d_ethosu_BI_pipeline(
301            model, common.get_u85_compile_spec(), model.get_inputs()
302        )
303