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