diff --git a/backends/cortex_m/ops/op_quantized_conv2d.cpp b/backends/cortex_m/ops/op_quantized_conv2d.cpp index 91cc893fba7..4a316f9420e 100644 --- a/backends/cortex_m/ops/op_quantized_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_conv2d.cpp @@ -77,10 +77,11 @@ bool validate_conv2d_arguments( return false; } - if (stride.size() != 2 || padding.size() != 2 || dilation.size() != 2) { + if (stride.size() != 2 || (padding.size() != 2 && padding.size() != 4) || + dilation.size() != 2) { ET_LOG( Error, - "quantized_conv2d_out: stride, padding, and dilation must have length 2"); + "quantized_conv2d_out: stride/dilation must have length 2; padding must have length 2 or 4"); context.fail(Error::InvalidArgument); return false; } @@ -176,6 +177,7 @@ static Tensor& quantized_conv2d_out_impl( conv_params.output_offset = output_offset_val; conv_params.stride.h = static_cast(stride[0]); conv_params.stride.w = static_cast(stride[1]); + // Trailing padding is encoded in the planned output dimensions. conv_params.padding.h = static_cast(padding[0]); conv_params.padding.w = static_cast(padding[1]); conv_params.dilation.h = static_cast(dilation[0]); diff --git a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp index 296fda24b56..a78f844f5b3 100644 --- a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp @@ -112,10 +112,11 @@ bool validate_depthwise_conv2d_arguments( return false; } - if (stride.size() != 2 || padding.size() != 2 || dilation.size() != 2) { + if (stride.size() != 2 || (padding.size() != 2 && padding.size() != 4) || + dilation.size() != 2) { ET_LOG( Error, - "quantized_depthwise_conv2d_out: stride/padding/dilation must have length 2"); + "quantized_depthwise_conv2d_out: stride/dilation must have length 2; padding must have length 2 or 4"); context.fail(Error::InvalidArgument); return false; } @@ -224,6 +225,7 @@ static Tensor& quantized_depthwise_conv2d_out_impl( dw_conv_params.ch_mult = depth_multiplier_val; dw_conv_params.stride.h = static_cast(stride[0]); dw_conv_params.stride.w = static_cast(stride[1]); + // Trailing padding is encoded in the planned output dimensions. dw_conv_params.padding.h = static_cast(padding[0]); dw_conv_params.padding.w = static_cast(padding[1]); dw_conv_params.dilation.h = static_cast(dilation[0]); diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 56af4cbe9fb..15522bb8d00 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -769,6 +769,37 @@ def pad_impl( ) +def _conv2d_padding( + padding: Sequence[int], + input_shape: torch.Size, + weight_shape: torch.Size, + stride: Sequence[int], + dilation: Sequence[int], +) -> tuple[int, int, int, int]: + # Four values mean top/left/bottom/right for supported SAME convolutions. + if len(padding) not in (2, 4) or any(p < 0 for p in padding): + raise ValueError(f"Expected 2 or 4 nonnegative padding values, got {padding}") + top, left = padding[:2] + bottom, right = padding[2:] if len(padding) == 4 else padding + if len(padding) == 4: + height, width = input_shape[2:4] + kernel = weight_shape[1:3] + if height == 1 and kernel[0] == 1: + raise ValueError("Four-value padding is unsupported for 1xN convolution") + total = [ + max(((size + step - 1) // step - 1) * step + dil * (k - 1) + 1 - size, 0) + for size, step, dil, k in zip((height, width), stride, dilation, kernel) + ] + if (top, left, bottom, right) != ( + total[0] // 2, + total[1] // 2, + total[0] - total[0] // 2, + total[1] - total[1] // 2, + ): + raise ValueError("Four-value padding must match SAME convolution geometry") + return top, left, bottom, right + + def _compute_conv2d_output_shape( input_shape: torch.Size, weight_shape: torch.Size, @@ -784,15 +815,17 @@ def _compute_conv2d_output_shape( kernel_width = weight_shape[2] stride_h, stride_w = stride - pad_h, pad_w = padding + pad_h, pad_w, pad_bottom, pad_right = _conv2d_padding( + padding, input_shape, weight_shape, stride, dilation + ) dilation_h, dilation_w = dilation out_channels = weight_shape[0] out_height = ( - in_height + 2 * pad_h - dilation_h * (kernel_height - 1) - 1 + in_height + pad_h + pad_bottom - dilation_h * (kernel_height - 1) - 1 ) // stride_h + 1 out_width = ( - in_width + 2 * pad_w - dilation_w * (kernel_width - 1) - 1 + in_width + pad_w + pad_right - dilation_w * (kernel_width - 1) - 1 ) // stride_w + 1 return torch.Size([batch, out_channels, out_height, out_width]) @@ -813,15 +846,17 @@ def _compute_depthwise_conv2d_output_shape( kernel_width = weight_shape[2] stride_h, stride_w = stride - pad_h, pad_w = padding + pad_h, pad_w, pad_bottom, pad_right = _conv2d_padding( + padding, input_shape, weight_shape, stride, dilation + ) dilation_h, dilation_w = dilation out_channels = weight_shape[3] # IHWO format: output channels at dimension 3 out_height = ( - in_height + 2 * pad_h - dilation_h * (kernel_height - 1) - 1 + in_height + pad_h + pad_bottom - dilation_h * (kernel_height - 1) - 1 ) // stride_h + 1 out_width = ( - in_width + 2 * pad_w - dilation_w * (kernel_width - 1) - 1 + in_width + pad_w + pad_right - dilation_w * (kernel_width - 1) - 1 ) // stride_w + 1 return torch.Size([batch, out_channels, out_height, out_width]) @@ -876,6 +911,12 @@ def quantized_conv2d_impl( raise RuntimeError("quantized_conv2d expects 4D input and weight tensors") # Convert to int32 for accumulation and apply offsets input_int32 = input.to(torch.int32) + int(input_offset) + if len(padding) == 4: + top, left, bottom, right = _conv2d_padding( + padding, input.shape, weight.shape, stride, dilation + ) + input_int32 = F.pad(input_int32, (left, right, top, bottom)) + padding = (0, 0) weight_int32 = weight.to(torch.int32) if bias is None: @@ -1110,6 +1151,12 @@ def quantized_depthwise_conv2d_impl( # Convert to int32 for accumulation and apply offsets input_int32 = input.to(torch.int32) + int(input_offset) + if len(padding) == 4: + top, left, bottom, right = _conv2d_padding( + padding, input.shape, weight.shape, stride, dilation + ) + input_int32 = F.pad(input_int32, (left, right, top, bottom)) + padding = (0, 0) weight_int32 = weight.to(torch.int32) if bias is None: diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index d2bda57a491..9dfc46856a0 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -36,6 +36,7 @@ fbcode_target(_kind = runtime.python_library, "decompose_hardswish_pass.py", "decompose_mean_pass.py", "explicit_layout_pass.py", + "fuse_conv_padding_pass.py", "initialize_scratch_buffers_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 6efd551c0a9..1e007c32cc9 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -41,6 +41,7 @@ CortexMReplaceOpsWithChannelsLastVariants, ValidateCortexMExplicitLayoutPass, ) +from .fuse_conv_padding_pass import FuseConvPaddingPass from .initialize_scratch_buffers_pass import InitializeScratchBuffersPass from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass @@ -60,6 +61,7 @@ class CortexMPassManager(PassManager): QuantizedClampActivationPass, DecomposeHardswishPass, AtenToCortexMPass, + FuseConvPaddingPass, InitializeScratchBuffersPass, ] @@ -73,12 +75,13 @@ class CortexMPassManager(PassManager): ConvertConv1dToConv2dPass, CortexMReplaceOpsWithChannelsLastVariants, ReplaceSqueezeAndUnsqueezeWithViewPass, - CortexMCanonicalizeViewCopyPermutePass, + # Move layout copies across pads before singleton permutations become views. RemovePermutesAroundElementwiseOps, CortexMCanonicalizeViewCopyPermutePass, ValidateCortexMExplicitLayoutPass, ReplaceQuantNodesPass, AtenToCortexMPass, + FuseConvPaddingPass, InitializeScratchBuffersPass, ] diff --git a/backends/cortex_m/passes/fuse_conv_padding_pass.py b/backends/cortex_m/passes/fuse_conv_padding_pass.py new file mode 100644 index 00000000000..2ae7f60192f --- /dev/null +++ b/backends/cortex_m/passes/fuse_conv_padding_pass.py @@ -0,0 +1,81 @@ +# Copyright (c) Meta Platforms, Inc. and 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. + +import torch +from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + + +class FuseConvPaddingPass(ExportPass): + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + conv_ops = { + exir_ops.edge.cortex_m.quantized_conv2d.default: (False, 6), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: (True, 6), + exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default: (False, 7), + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default: (True, 7), + } + modified = False + for node in graph_module.graph.nodes: + if node.target not in conv_ops: + continue + pad = node.args[0] + if pad.target != exir_ops.edge.cortex_m.pad.default: + continue + explicit_nhwc, offset_index = conv_ops[node.target] + source, before, after, value = pad.args + tensor = get_first_fake_tensor(source) + if tensor.dim() != 4 or ( + not explicit_nhwc + and not tensor.is_contiguous(memory_format=torch.channels_last) + ): + continue + if ( + before[0] + or before[3] + or after[0] + or after[3] + or any(p < 0 for p in (*before, *after)) + or value != -node.args[offset_index] + ): + continue + height, width = tensor.shape[1:3] if explicit_nhwc else tensor.shape[2:4] + kernel = get_first_fake_tensor(node.args[1]).shape[1:3] + # The pinned CMSIS-NN MVE 1xN kernel mishandles asymmetric boundaries. + if height == 1 and kernel[0] == 1: + continue + stride, padding, dilation = node.args[3:6] + if len(padding) != 2: + continue + fused_padding = [ + before[1] + padding[0], + before[2] + padding[1], + after[1] + padding[0], + after[2] + padding[1], + ] + # Limit fusion to SAME padding, as used by the CMSIS-NN wrappers. + total = [ + max( + ((size + step - 1) // step - 1) * step + dil * (k - 1) + 1 - size, 0 + ) + for size, step, dil, k in zip((height, width), stride, dilation, kernel) + ] + if fused_padding != [ + total[0] // 2, + total[1] // 2, + total[0] - total[0] // 2, + total[1] - total[1] // 2, + ]: + continue + args = list(node.args) + args[0] = source + args[4] = fused_padding + node.args = tuple(args) + modified = True + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/cortex_m/test/models/test_mobilenet_v1_025.py b/backends/cortex_m/test/models/test_mobilenet_v1_025.py index c4b2554c7db..358c7e23a98 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v1_025.py +++ b/backends/cortex_m/test/models/test_mobilenet_v1_025.py @@ -23,6 +23,7 @@ ops_after_transforms: dict[str, int] = { "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 0, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 14, diff --git a/backends/cortex_m/test/models/test_resnet8.py b/backends/cortex_m/test/models/test_resnet8.py index fafa6c801a2..7b6f4cf0c4e 100644 --- a/backends/cortex_m/test/models/test_resnet8.py +++ b/backends/cortex_m/test/models/test_resnet8.py @@ -24,6 +24,7 @@ ops_after_transforms: dict[str, int] = { "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 0, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 3, "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1, diff --git a/backends/cortex_m/test/ops/test_nhwc_conv.py b/backends/cortex_m/test/ops/test_nhwc_conv.py index da7602447c2..6ab7c4b5918 100644 --- a/backends/cortex_m/test/ops/test_nhwc_conv.py +++ b/backends/cortex_m/test/ops/test_nhwc_conv.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import pytest import torch from executorch.backends.cortex_m.test.ops.nhwc_test_utils import ( @@ -16,8 +17,9 @@ class Conv2dNhwc(torch.nn.Module): - def __init__(self, grouped=False): + def __init__(self, grouped=False, padding=(1, 0)): super().__init__() + self.padding = padding in_channels = 4 if grouped else 3 self.register_buffer( "weight", int8_values((4, 2, 3, 2 if grouped else in_channels)) @@ -34,7 +36,7 @@ def forward(self, x, scratch): self.weight, self.bias, [2, 1], - [1, 0], + self.padding, [1, 1], 0, 0, @@ -47,8 +49,9 @@ def forward(self, x, scratch): class DepthwiseConv2dNhwc(torch.nn.Module): - def __init__(self): + def __init__(self, padding=(1, 0)): super().__init__() + self.padding = padding self.register_buffer("weight", int8_values((1, 3, 2, 4))) self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) self.register_buffer( @@ -62,7 +65,7 @@ def forward(self, x, scratch): self.weight, self.bias, [2, 1], - [1, 0], + self.padding, [1, 1], 1, 0, @@ -105,9 +108,10 @@ def forward(self, x, scratch, output_scratch): ) -def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): +@pytest.mark.parametrize("padding", [(1, 0), (0, 1, 1, 1)]) +def test_conv2d_nhwc_runs_on_fvp(cortex_m_target, padding): run_on_fvp( - Conv2dNhwc(), + Conv2dNhwc(padding=padding), int8_values((1, 7, 10, 3)), exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, cortex_m_target, @@ -115,9 +119,10 @@ def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): ) -def test_grouped_conv2d_nhwc_runs_on_fvp(cortex_m_target): +@pytest.mark.parametrize("padding", [(1, 0), (0, 1, 1, 1)]) +def test_grouped_conv2d_nhwc_runs_on_fvp(cortex_m_target, padding): run_on_fvp( - Conv2dNhwc(grouped=True), + Conv2dNhwc(grouped=True, padding=padding), int8_values((1, 7, 10, 4)), exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, cortex_m_target, @@ -125,9 +130,10 @@ def test_grouped_conv2d_nhwc_runs_on_fvp(cortex_m_target): ) -def test_depthwise_conv2d_nhwc_runs_on_fvp(cortex_m_target): +@pytest.mark.parametrize("padding", [(1, 0), (1, 0, 1, 1)]) +def test_depthwise_conv2d_nhwc_runs_on_fvp(cortex_m_target, padding): run_on_fvp( - DepthwiseConv2dNhwc(), + DepthwiseConv2dNhwc(padding=padding), int8_values((1, 7, 10, 4)), exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, cortex_m_target, diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index 100d4ca59c7..a215823164c 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -102,3 +102,18 @@ def define_common_targets(is_fbcode = False): ) + + python_pytest( + name = "test_fuse_conv_padding", + srcs = ["test_fuse_conv_padding.py"], + compile = "with-source", + typing = False, + env = {"PYTEST_ADDOPTS": "-k 'not test_implementation'"}, + deps = [ + "//caffe2:torch", + "//executorch/backends/test/harness:tester", + "//executorch/exir/dialects:lib", + ":tester", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index a278209f658..ce01ceba8b9 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -55,6 +55,19 @@ def forward(self, x): ) +class NHWCPaddedConv(torch.nn.Module): + def __init__(self, channels, shared_pad): + super().__init__() + out_channels = 64 if channels == 1 else 8 + self.conv = torch.nn.Conv2d(channels, out_channels, (10, 4), stride=2) + self.shared_pad = shared_pad + + def forward(self, x): + padded = torch.nn.functional.pad(x.permute(0, 3, 1, 2), (1, 1, 4, 5)) + output = self.conv(padded).permute(0, 2, 3, 1) + return (output, padded) if self.shared_pad else output + + def _count(exported_program, target) -> int: return sum(node.target == target for node in exported_program.graph.nodes) @@ -148,6 +161,37 @@ def test_explicit_layout_reuses_pad(): assert _count(program, exir_ops.edge.cortex_m.pad.default) == 1 +def _lower_nhwc_padded_conv(channels, shared_pad): + torch.manual_seed(7) + tester = _run_explicit_layout_passes( + CortexMTester( + NHWCPaddedConv(channels, shared_pad).eval(), + (torch.randn(1, 49, 10, channels),), + ) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.pad.default) == int(shared_pad) + if not shared_pad: + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 0 + return tester + + +@pytest.mark.parametrize("channels", [1, 3]) +@pytest.mark.parametrize("shared_pad", [False, True]) +def test_explicit_layout_fuses_same_padding(channels, shared_pad): + tester = _lower_nhwc_padded_conv(channels, shared_pad) + tester.run_method_and_compare_outputs(inputs=tester.example_inputs, qtol=1) + + +@pytest.mark.parametrize("channels", [1, 3]) +@pytest.mark.parametrize("shared_pad", [False, True]) +def test_implementation_explicit_layout_fuses_same_padding(channels, shared_pad): + tester = _lower_nhwc_padded_conv(channels, shared_pad) + tester.to_executorch().serialize() + tester.run_method_and_compare_outputs(inputs=tester.example_inputs, qtol=1) + + @pytest.mark.parametrize("hardtanh", [False, True]) def test_implementation_transpose_conv2d_strided_pointwise(hardtanh): torch.manual_seed(0) diff --git a/backends/cortex_m/test/test_fuse_conv_padding.py b/backends/cortex_m/test/test_fuse_conv_padding.py new file mode 100644 index 00000000000..56dd39d3dfe --- /dev/null +++ b/backends/cortex_m/test/test_fuse_conv_padding.py @@ -0,0 +1,165 @@ +# Copyright (c) Meta Platforms, Inc. and 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. + +from contextlib import nullcontext + +import pytest +import torch +from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.test.harness.stages import StageType +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode + + +@pytest.mark.parametrize("fake", [False, True]) +@pytest.mark.parametrize("depthwise", [False, True]) +@pytest.mark.parametrize("nhwc", [False, True]) +@pytest.mark.parametrize( + "size,kernel,stride,padding,error", + [ + ((2, 2), (1, 1), (1, 1), (0, 0, 1, 1), "SAME convolution geometry"), + ((1, 8), (1, 5), (1, 2), (0, 1, 0, 2), "unsupported for 1xN"), + ], +) +def test_unsupported_four_value_padding( + fake, depthwise, nhwc, size, kernel, stride, padding, error +): + ops = { + (False, False): torch.ops.cortex_m.quantized_conv2d.default, + (False, True): torch.ops.cortex_m.quantized_conv2d_nhwc.default, + (True, False): torch.ops.cortex_m.quantized_depthwise_conv2d.default, + (True, True): torch.ops.cortex_m.quantized_depthwise_conv2d_nhwc.default, + } + x = torch.ones((1, 1, *size), dtype=torch.int8).to( + memory_format=torch.channels_last + ) + if nhwc: + x = x.permute(0, 2, 3, 1).contiguous() + weight = torch.ones((1, *kernel, 1), dtype=torch.int8) + args = (x, weight, None, stride, padding, (1, 1)) + if depthwise: + args += (1,) + args += ( + 0, + 0, + torch.tensor([1 << 30], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + -128, + 127, + torch.empty(0, dtype=torch.uint8), + ) + with FakeTensorMode(allow_non_fake_inputs=True) if fake else nullcontext(): + with pytest.raises(ValueError, match=error): + ops[depthwise, nhwc](*args) + + +class PaddedConv(torch.nn.Module): + def __init__( + self, + channels, + out_channels, + groups, + kernel, + stride, + padding, + value=0, + shared=False, + ): + super().__init__() + self.conv = torch.nn.Conv2d( + channels, out_channels, kernel, stride=stride, groups=groups + ) + self.padding = padding + self.value = value + self.shared = shared + + def forward(self, x): + padded = torch.nn.functional.pad(x, self.padding, value=self.value) + result = self.conv(padded) + return (result, padded) if self.shared else result + + +pad_cases = [ + pytest.param( + 1, 64, 1, (10, 4), (2, 2), (1, 1, 4, 5), 0, False, True, id="ds_cnn_stem" + ), + pytest.param( + 1, + 8, + 1, + (10, 4), + (2, 2), + (1, 1, 4, 5), + 0, + False, + True, + id="single_channel_depthwise", + ), + pytest.param(4, 8, 1, (3, 3), (2, 2), (0, 1, 0, 1), 0, False, True, id="conv"), + pytest.param(4, 8, 4, (3, 3), (2, 2), (0, 1, 0, 1), 0, False, True, id="depthwise"), + pytest.param(4, 8, 2, (3, 3), (2, 2), (0, 1, 0, 1), 0, False, True, id="grouped"), + pytest.param(4, 8, 1, (1, 5), (1, 2), (1, 2, 0, 0), 0, False, False, id="1xn"), + pytest.param( + 4, 8, 1, (3, 3), (2, 2), (0, 1, 0, 1), 1.0, False, False, id="nonzero_pad" + ), + pytest.param( + 4, 8, 1, (3, 3), (2, 2), (1, 0, 1, 0), 0, False, False, id="same_lower" + ), + pytest.param(4, 8, 1, (3, 3), (2, 2), (0, 1, 0, 1), 0, True, True, id="shared_pad"), +] + + +@pytest.mark.parametrize( + "channels,out_channels,groups,kernel,stride,padding,value,shared,fused", pad_cases +) +def test_fuse_conv_padding( + channels, out_channels, groups, kernel, stride, padding, value, shared, fused +): + model = PaddedConv( + channels, out_channels, groups, kernel, stride, padding, value, shared + ).eval() + size = (49, 10) if channels == 1 else (8, 8) + if kernel[0] == 1: + size = (1, 8) + x = (torch.rand(1, channels, *size) * 5 - 1).to(memory_format=torch.channels_last) + tester = CortexMTester(model, (x,)).quantize().export().to_edge().run_passes() + tester.run_method_and_compare_outputs(inputs=(x,), qtol=1) + graph = tester.get_artifact(StageType.RUN_PASSES).exported_program().graph + convs = [n for n in graph.nodes if "conv2d" in str(n.target)] + assert len(convs) == 1 + conv = convs[0] + assert ("depthwise" in str(conv.target)) == ( + groups == channels and not (channels == 1 and out_channels > 8) + ) + assert (len(conv.args[4]) == 4) == fused + pads = [n for n in graph.nodes if n.target == exir_ops.edge.cortex_m.pad.default] + assert len(pads) == (not fused or shared) + + +@pytest.mark.parametrize( + "channels,out_channels,groups,kernel,stride,padding,value,shared,fused", + pad_cases[:6], +) +def test_implementation_fused_conv_padding( + channels, + out_channels, + groups, + kernel, + stride, + padding, + value, + shared, + fused, + cortex_m_target, +): + model = PaddedConv(channels, out_channels, groups, kernel, stride, padding).eval() + size = (49, 10) if channels == 1 else (8, 8) + if kernel[0] == 1: + size = (1, 8) + x = (torch.rand(1, channels, *size) * 5 - 1).to(memory_format=torch.channels_last) + CortexMTester(model, (x,), target_config=cortex_m_target).test_implementation( + qtol=1 + )