• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 Tuple
10
11import pytest
12
13import torch
14from executorch.backends.arm.test import common
15from executorch.backends.arm.test.ops.test_conv1d import Conv1d
16from executorch.backends.arm.test.ops.test_conv2d import Conv2d
17
18from executorch.backends.arm.test.tester.arm_tester import ArmTester
19from executorch.exir.backend.backend_details import CompileSpec
20from parameterized import parameterized
21
22
23"""
24The configuration when
25  groups == in_channels and
26  out_channels = K * in_channels
27  where K is a positive integer
28is termed in literature as depthwise convolution.
29"""
30
31dw_conv1d_3_1x3x14_gp3_st1 = Conv1d(
32    in_channels=3,
33    out_channels=3,
34    kernel_size=7,
35    stride=1,
36    groups=3,
37    length=14,
38    batches=1,
39    padding=3,
40)
41
42dw_conv1d_2_1x6x4_gp6_st1 = Conv1d(
43    in_channels=6,
44    out_channels=12,
45    kernel_size=2,
46    stride=1,
47    groups=6,
48    padding=0,
49    length=4,
50    batches=1,
51)
52
53dw_conv2d_2x2_1x6x4x4_gp6_st1 = Conv2d(
54    in_channels=6,
55    out_channels=12,
56    kernel_size=(2, 2),
57    stride=(1, 1),
58    groups=6,
59    padding=0,
60    width=4,
61    height=4,
62    batches=1,
63)
64
65dw_conv1d_3_1x3x256_gp3_st1 = Conv1d(
66    in_channels=3,
67    out_channels=3,
68    kernel_size=3,
69    stride=1,
70    groups=3,
71    padding=0,
72    length=256,
73    batches=1,
74)
75
76dw_conv2d_3x3_1x3x256x256_gp3_st1 = Conv2d(
77    in_channels=3,
78    out_channels=3,
79    kernel_size=(3, 3),
80    stride=(1, 1),
81    groups=3,
82    padding=0,
83    width=256,
84    height=256,
85    batches=1,
86)
87
88dw_conv2d_3x3_1x4x256x256_gp4_st1 = Conv2d(
89    in_channels=4,
90    out_channels=8,
91    kernel_size=(3, 3),
92    stride=(1, 1),
93    groups=4,
94    padding=0,
95    width=256,
96    height=256,
97    batches=1,
98)
99
100dw_conv2d_3x3_2x8x198x198_gp8_st3 = Conv2d(
101    in_channels=8,
102    out_channels=16,
103    kernel_size=(3, 3),
104    stride=3,
105    groups=8,
106    padding=0,
107    width=198,
108    height=198,
109    batches=2,
110)
111
112dw_conv2d_3x3_1x4x256x256_gp4_nobias = Conv2d(
113    in_channels=4,
114    out_channels=8,
115    kernel_size=(3, 3),
116    stride=1,
117    groups=4,
118    bias=False,
119    width=256,
120    height=256,
121    batches=1,
122)
123
124two_dw_conv1d = Conv1d(
125    nbr_conv=2,
126    length=64,
127    in_channels=[4, 8],
128    out_channels=[8, 24],
129    kernel_size=[3, 3],
130    stride=[1, 1],
131    padding=[0, 0],
132    groups=[4, 8],
133    bias=[True, True],
134    batches=1,
135)
136
137two_dw_conv2d = Conv2d(
138    nbr_conv=2,
139    width=64,
140    height=64,
141    in_channels=[4, 8],
142    out_channels=[8, 24],
143    kernel_size=[(3, 3), (3, 3)],
144    stride=[1, 1],
145    padding=[0, 0],
146    groups=[4, 8],
147    bias=[True, True],
148    batches=2,
149)
150
151# Shenanigan to get a nicer output when test fails.
152testsuite_conv2d = [
153    ("2x2_1x6x4x4_gp6_st1", dw_conv2d_2x2_1x6x4x4_gp6_st1),
154    ("3x3_1x3x256x256_gp3_st1", dw_conv2d_3x3_1x3x256x256_gp3_st1),
155    ("3x3_1x4x256x256_gp4_st1", dw_conv2d_3x3_1x4x256x256_gp4_st1),
156    ("3x3_2x8x198x198_gp8_st3", dw_conv2d_3x3_2x8x198x198_gp8_st3),
157    ("3x3_1x4x256x256_gp4_nobias", dw_conv2d_3x3_1x4x256x256_gp4_nobias),
158    ("two_dw_conv2d", two_dw_conv2d),
159]
160
161testsuite_conv1d = [
162    ("2_1x6x4_gp6_st1", dw_conv1d_2_1x6x4_gp6_st1),
163    ("3_1x3x256_gp3_st1", dw_conv1d_3_1x3x256_gp3_st1),
164    ("two_dw_conv1d", two_dw_conv1d),
165    ("3_1x3x14_gp3_st1", dw_conv1d_3_1x3x14_gp3_st1),
166]
167
168
169class TestDepthwiseConv(unittest.TestCase):
170    """Tests Conv1D and Conv2D where groups == in_channels and out_channels = K * in_channels. This
171    is a special case enables depthwise convolution."""
172
173    def _test_dw_conv_tosa_MI_pipeline(
174        self, module: torch.nn.Module, test_data: Tuple[torch.Tensor]
175    ):
176        (
177            ArmTester(
178                module,
179                example_inputs=test_data,
180                compile_spec=common.get_tosa_compile_spec(
181                    "TOSA-0.80.0+MI", permute_memory_to_nhwc=True
182                ),
183            )
184            .export()
185            .to_edge()
186            .partition()
187            .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
188            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
189            .to_executorch()
190            .run_method_and_compare_outputs(inputs=test_data)
191        )
192
193    def _test_dw_conv_tosa_BI_pipeline(
194        self, module: torch.nn.Module, test_data: Tuple[torch.Tensor]
195    ):
196        (
197            ArmTester(
198                module,
199                example_inputs=test_data,
200                compile_spec=common.get_tosa_compile_spec(
201                    "TOSA-0.80.0+BI", permute_memory_to_nhwc=True
202                ),
203            )
204            .quantize()
205            .export()
206            .to_edge()
207            .partition()
208            .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
209            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
210            .to_executorch()
211            .run_method_and_compare_outputs(inputs=test_data, qtol=1)
212        )
213
214    def _test_dw_conv_ethos_BI_pipeline(
215        self,
216        module: torch.nn.Module,
217        compile_spec: CompileSpec,
218        test_data: Tuple[torch.Tensor],
219    ):
220        (
221            ArmTester(
222                module,
223                example_inputs=test_data,
224                compile_spec=compile_spec,
225            )
226            .quantize()
227            .export()
228            .to_edge()
229            .partition()
230            .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"])
231            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
232            .to_executorch()
233        )
234
235    @parameterized.expand(testsuite_conv1d + testsuite_conv2d)
236    def test_dw_conv_tosa_MI(self, test_name: str, model: torch.nn.Module):
237        self._test_dw_conv_tosa_MI_pipeline(model, model.get_inputs())
238
239    # TODO: Investigate flakyness (MLTORCH-307)
240    @parameterized.expand(testsuite_conv1d + testsuite_conv2d)
241    @pytest.mark.flaky(reruns=3)
242    def test_dw_conv_tosa_BI(self, test_name: str, model: torch.nn.Module):
243        self._test_dw_conv_tosa_BI_pipeline(model, model.get_inputs())
244
245    @parameterized.expand(testsuite_conv2d, skip_on_empty=True)
246    def test_dw_conv2d_u55_BI(
247        self, test_name: str, model: torch.nn.Module, set_quantize_io: bool = False
248    ):
249        self._test_dw_conv_ethos_BI_pipeline(
250            model,
251            common.get_u55_compile_spec(
252                permute_memory_to_nhwc=True, quantize_io=set_quantize_io
253            ),
254            model.get_inputs(),
255        )
256
257    # Expected to fail as conv1d needs transpose which is not supported
258    # on u55.
259    @parameterized.expand(testsuite_conv1d, skip_on_empty=True)
260    @unittest.expectedFailure
261    def test_dw_conv1d_u55_BI(
262        self, test_name: str, model: torch.nn.Module, set_quantize_io: bool = False
263    ):
264        self._test_dw_conv_ethos_BI_pipeline(
265            model,
266            common.get_u55_compile_spec(
267                permute_memory_to_nhwc=True, quantize_io=set_quantize_io
268            ),
269            model.get_inputs(),
270        )
271
272    @parameterized.expand(testsuite_conv1d + testsuite_conv2d)
273    def test_dw_conv_u85_BI(
274        self, test_name: str, model: torch.nn.Module, set_quantize_io: bool = False
275    ):
276        self._test_dw_conv_ethos_BI_pipeline(
277            model,
278            common.get_u85_compile_spec(
279                permute_memory_to_nhwc=True, quantize_io=set_quantize_io
280            ),
281            model.get_inputs(),
282        )
283