Cortex-M: fuse SAME padding into quantized convolutions

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<const int32_t>(stride[0]);
   conv_params.stride.w = static_cast<const int32_t>(stride[1]);
+  // Trailing padding is encoded in the planned output dimensions.
   conv_params.padding.h = static_cast<const int32_t>(padding[0]);
   conv_params.padding.w = static_cast<const int32_t>(padding[1]);
   conv_params.dilation.h = static_cast<const int32_t>(dilation[0]);

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<const int32_t>(stride[0]);
   dw_conv_params.stride.w = static_cast<const int32_t>(stride[1]);
+  // Trailing padding is encoded in the planned output dimensions.
   dw_conv_params.padding.h = static_cast<const int32_t>(padding[0]);
   dw_conv_params.padding.w = static_cast<const int32_t>(padding[1]);
   dw_conv_params.dilation.h = static_cast<const int32_t>(dilation[0]);

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:

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",

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,
     ]
 

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)

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,

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,

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,29 +108,32 @@ 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,
         1,
     )
 
 
-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,
         1,
     )
 
 
-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,

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",
+            ],
+        )