Improve depthwise conv2d implementation with AOT optimizations and validations

Key changes:
- Move depth_multiplier calculation from runtime to AOT pass (eliminates
  runtime division by computing depth_multiplier = output_channels / input_channels
  in the graph transformation pass)
- Add critical defensive validations in validate_depthwise_conv2d_arguments():
  * Validate IHWO weight layout (dimension 0 must be 1)
  * Validate dilation == 1 (CMSIS-NN constraint)
  * Validate depth_multiplier consistency with channel counts
- Fix CMSIS-NN API usage:
  * Use arm_depthwise_conv_wrapper_s8_get_buffer_size() with correct parameters
  * Improve buffer allocation error handling with detailed error messages
- Add _compute_depthwise_conv2d_output_shape() to read channels from correct
  dimension (dim 3 for IHWO layout vs dim 0 for OHWI)
- Update operator schema to use depth_multiplier parameter instead of groups

This ensures proper validation of CMSIS-NN constraints and moves computation
to compile-time where possible.

Author: RJ Ascani

backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp

@@ -1,5 +1,6 @@
 /*
- * Copyright 2025 Arm Limited and/or its affiliates.
+ * 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.
@@ -28,7 +29,7 @@ bool validate_depthwise_conv2d_arguments(
     const IntArrayRef& stride,
     const IntArrayRef& padding,
     const IntArrayRef& dilation,
-    const int64_t groups,
+    const int64_t depth_multiplier,
     const Tensor& requantize_multipliers,
     const Tensor& requantize_shifts) {
   if (input.dim() != kConvDim || weight.dim() != kConvDim ||
@@ -38,6 +39,28 @@ bool validate_depthwise_conv2d_arguments(
     return false;
   }
 
+  // Validate weight is in IHWO layout: [1, H, W, C_OUT]
+  if (weight.size(0) != 1) {
+    ET_LOG(
+        Error,
+        "quantized_depthwise_conv2d_out: weight dim 0 must be 1 for IHWO layout, got %zd",
+        weight.size(0));
+    context.fail(Error::InvalidArgument);
+    return false;
+  }
+
+  const int64_t weight_output_channels = weight.size(3);
+  const int64_t output_channels = output.size(1);
+  if (weight_output_channels != output_channels) {
+    ET_LOG(
+        Error,
+        "quantized_depthwise_conv2d_out: weight output channels (%zd) must match output channels (%zd)",
+        weight_output_channels,
+        output_channels);
+    context.fail(Error::InvalidArgument);
+    return false;
+  }
+
   // Check for channels_last dim_order (NHWC: 0, 2, 3, 1)
   // Skip check if channels == 1, as dim_order is ambiguous in that case
   constexpr executorch::aten::DimOrderType kChannelsLastDimOrder[] = {
@@ -91,25 +114,36 @@ bool validate_depthwise_conv2d_arguments(
     return false;
   }
 
-  // Depthwise convolution constraint: groups must equal input channels
+  // CMSIS-NN depthwise convolution does not support dilation != 1
+  if (dilation[0] != 1 || dilation[1] != 1) {
+    ET_LOG(
+        Error,
+        "quantized_depthwise_conv2d_out: dilation != 1 not supported, got (%zd, %zd)",
+        dilation[0],
+        dilation[1]);
+    context.fail(Error::InvalidArgument);
+    return false;
+  }
+
   const int64_t input_channels = input.size(1);
-  if (groups != input_channels) {
+  // output_channels already extracted above for weight validation
+  if (output_channels != input_channels * depth_multiplier) {
     ET_LOG(
         Error,
-        "quantized_depthwise_conv2d_out: groups (%zd) must equal input channels (%zd) for depthwise convolution",
-        groups,
-        input_channels);
+        "quantized_depthwise_conv2d_out: output channels (%zd) must equal input channels (%zd) * depth_multiplier (%zd)",
+        output_channels,
+        input_channels,
+        depth_multiplier);
     context.fail(Error::InvalidArgument);
     return false;
   }
 
-  const int64_t out_channels = output.size(1);
-  if (requantize_multipliers.size(0) != out_channels ||
-      requantize_shifts.size(0) != out_channels) {
+  if (requantize_multipliers.size(0) != output_channels ||
+      requantize_shifts.size(0) != output_channels) {
     ET_LOG(
         Error,
         "quantized_depthwise_conv2d_out: per-channel params must match output channels (%zd)",
-        out_channels);
+        output_channels);
     context.fail(Error::InvalidArgument);
     return false;
   }
@@ -126,7 +160,7 @@ Tensor& quantized_depthwise_conv2d_out(
     const IntArrayRef stride,
     const IntArrayRef padding,
     const IntArrayRef dilation,
-    const int64_t groups,
+    const int64_t depth_multiplier,
     const int64_t input_offset,
     const int64_t output_offset,
     const Tensor& requantize_multipliers,
@@ -143,7 +177,7 @@ Tensor& quantized_depthwise_conv2d_out(
           stride,
           padding,
           dilation,
-          groups,
+          depth_multiplier,
           requantize_multipliers,
           requantize_shifts)) {
     return out;
@@ -154,15 +188,17 @@ Tensor& quantized_depthwise_conv2d_out(
   const int32_t input_height = static_cast<int32_t>(input.size(2));
   const int32_t input_width = static_cast<int32_t>(input.size(3));
 
-  const int32_t kernel_output_channels = static_cast<int32_t>(weight.size(0));
+  // Weight is in IHWO layout after permutation in the pass: [1, H, W, C_OUT]
+  // For depthwise conv, this matches CMSIS-NN's expected format
   const int32_t kernel_height = static_cast<int32_t>(weight.size(1));
   const int32_t kernel_width = static_cast<int32_t>(weight.size(2));
-  const int32_t depth_multiplier = static_cast<int32_t>(weight.size(3));
 
   const int32_t output_channels = static_cast<int32_t>(out.size(1));
   const int32_t output_height = static_cast<int32_t>(out.size(2));
   const int32_t output_width = static_cast<int32_t>(out.size(3));
 
+  const int32_t depth_multiplier_val = static_cast<int32_t>(depth_multiplier);
+
   const int32_t input_offset_val = static_cast<int32_t>(input_offset);
   const int32_t output_offset_val = static_cast<int32_t>(output_offset);
   const int32_t activation_min_val = static_cast<int32_t>(activation_min);
@@ -179,7 +215,7 @@ Tensor& quantized_depthwise_conv2d_out(
   cmsis_nn_dw_conv_params dw_conv_params;
   dw_conv_params.input_offset = input_offset_val;
   dw_conv_params.output_offset = output_offset_val;
-  dw_conv_params.ch_mult = depth_multiplier;
+  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]);
   dw_conv_params.padding.h = static_cast<const int32_t>(padding[0]);
@@ -203,24 +239,23 @@ Tensor& quantized_depthwise_conv2d_out(
   cmsis_context.buf = nullptr;
   cmsis_context.size = 0;
 
-  const size_t buffer_bytes = static_cast<size_t>(
-      arm_depthwise_conv_s8_get_buffer_size(&input_dims, &filter_dims));
+  const size_t buffer_bytes =
+      static_cast<size_t>(arm_depthwise_conv_wrapper_s8_get_buffer_size(
+          &dw_conv_params, &input_dims, &filter_dims, &output_dims));
   if (buffer_bytes > 0) {
     auto buffer_or_error =
         context.allocate_temp(buffer_bytes, alignof(int16_t));
     if (!buffer_or_error.ok()) {
-      if (buffer_or_error.error() != Error::NotFound) {
-        ET_LOG(
-            Error,
-            "quantized_depthwise_conv2d_out: failed to allocate scratch buffer (%d)",
-            static_cast<int>(buffer_or_error.error()));
-        context.fail(buffer_or_error.error());
-        return out;
-      }
-    } else {
-      cmsis_context.buf = buffer_or_error.get();
-      cmsis_context.size = buffer_bytes;
+      ET_LOG(
+          Error,
+          "quantized_depthwise_conv2d_out: failed to allocate scratch buffer (%d bytes, error %d)",
+          static_cast<int>(buffer_bytes),
+          static_cast<int>(buffer_or_error.error()));
+      context.fail(buffer_or_error.error());
+      return out;
     }
+    cmsis_context.buf = buffer_or_error.get();
+    cmsis_context.size = buffer_bytes;
   }
 
   const arm_cmsis_nn_status status = arm_depthwise_conv_wrapper_s8(

backends/cortex_m/ops/operators.py

@@ -488,6 +488,35 @@ def _compute_conv2d_output_shape(
     return torch.Size([batch, out_channels, out_height, out_width])
 
 
+def _compute_depthwise_conv2d_output_shape(
+    input_shape: torch.Size,
+    weight_shape: torch.Size,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+) -> torch.Size:
+    batch = input_shape[0]
+    in_height = input_shape[2]
+    in_width = input_shape[3]
+    # For depthwise conv, we store the weights in IHWO layout (1, kernel_h, kernel_w, out)
+    # where dimension 3 contains the output channels
+    kernel_height = weight_shape[1]
+    kernel_width = weight_shape[2]
+
+    stride_h, stride_w = stride
+    pad_h, pad_w = padding
+    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
+    ) // stride_h + 1
+    out_width = (
+        in_width + 2 * pad_w - dilation_w * (kernel_width - 1) - 1
+    ) // stride_w + 1
+    return torch.Size([batch, out_channels, out_height, out_width])
+
+
 @register_fake("cortex_m::quantized_conv2d")
 def quantized_conv2d_meta(
     input: torch.Tensor,
@@ -591,7 +620,7 @@ def quantized_conv2d_impl(
     "int[] stride, "
     "int[] padding, "
     "int[] dilation, "
-    "int groups, "
+    "int depth_multiplier, "
     "int input_offset, "
     "int output_offset, "
     "Tensor requantize_multipliers, "
@@ -610,7 +639,7 @@ def quantized_conv2d_impl(
     "int[] stride, "
     "int[] padding, "
     "int[] dilation, "
-    "int groups, "
+    "int depth_multiplier, "
     "int input_offset, "
     "int output_offset, "
     "Tensor requantize_multipliers, "
@@ -630,7 +659,7 @@ def quantized_depthwise_conv2d_meta(
     stride: Sequence[int],
     padding: Sequence[int],
     dilation: Sequence[int],
-    groups: int,
+    depth_multiplier: int,
     input_offset: int,
     output_offset: int,
     requantize_multipliers: torch.Tensor,
@@ -641,7 +670,7 @@ def quantized_depthwise_conv2d_meta(
     stride_vals = list(stride)
     padding_vals = list(padding)
     dilation_vals = list(dilation)
-    output_shape = _compute_conv2d_output_shape(
+    output_shape = _compute_depthwise_conv2d_output_shape(
         input.shape, weight.shape, stride_vals, padding_vals, dilation_vals
     )
     return torch.empty(
@@ -660,7 +689,7 @@ def quantized_depthwise_conv2d_impl(
     stride: Sequence[int],
     padding: Sequence[int],
     dilation: Sequence[int],
-    groups: int,
+    depth_multiplier: int,
     input_offset: int,
     output_offset: int,
     requantize_multipliers: torch.Tensor,
@@ -673,26 +702,26 @@ def quantized_depthwise_conv2d_impl(
             "quantized_depthwise_conv2d expects 4D input and weight tensors"
         )
 
-    # Validate depthwise convolution constraint: groups == input_channels
     input_channels = input.shape[1]
-    if groups != input_channels:
-        raise RuntimeError(
-            f"quantized_depthwise_conv2d: groups ({groups}) must equal input channels ({input_channels})"
-        )
+    groups = input_channels
 
     # Convert to int32 for accumulation and apply offsets
     input_int32 = input.to(torch.int32) + int(input_offset)
     weight_int32 = weight.to(torch.int32)
 
     if bias is None:
         bias_int32 = torch.zeros(
-            weight.shape[0], dtype=torch.int32, device=input.device
+            weight.shape[3],
+            dtype=torch.int32,
+            device=input.device,  # C_OUT is at dim 3 in IHWO
         )
     else:
         bias_int32 = bias.to(torch.int32)
 
-    # Convert weights back to OIHW layout expected by torch.nn.functional.conv2d
-    weight_oi_hw = weight_int32.permute(0, 3, 1, 2).contiguous()
+    # Weight is in IHWO layout: [1, H, W, C_OUT]
+    # Convert to OIHW layout expected by torch.nn.functional.conv2d
+    # IHWO [1, H, W, C_OUT] -> OIHW [C_OUT, 1, H, W]
+    weight_oi_hw = weight_int32.permute(3, 0, 1, 2).contiguous()
 
     # Depthwise convolution has groups == input_channels
     conv_acc = F.conv2d(

backends/cortex_m/ops/operators.yaml

@@ -59,7 +59,7 @@
     - arg_meta: null
       kernel_name: cortex_m::quantized_conv2d_out
 
-- func: cortex_m::quantized_depthwise_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int groups, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!)
+- func: cortex_m::quantized_depthwise_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:
     - arg_meta: null

backends/cortex_m/passes/convert_to_cortex_m_pass.py

@@ -156,11 +156,24 @@ def _get_convolution_replacement(self, node) -> int:
             quantized_multipliers.append(quantized_multiplier)
             quantized_shifts.append(quantized_shift)
 
-        # Permute the weight tensor to the OHWI layout expected by CMSIS-NN.
         weight_tensor = get_param_tensor(self.exported_program, weight)
-        weight_permuted = weight_tensor.permute(0, 2, 3, 1).contiguous(
-            memory_format=torch.channels_last
-        )
+
+        # Detect depthwise convolution:
+        # PyTorch depthwise weight is [out_ch, 1, H, W] where dimension 1 is 1
+        # and groups == input_channels (groups > 1)
+        is_depthwise = weight_tensor.shape[1] == 1 and groups > 1
+
+        if is_depthwise:
+            # For depthwise: OIHW -> IHWO which gives [1, H, W, C_OUT] for CMSIS-NN
+            # PyTorch depthwise weight is [out_ch, 1, H, W], permute to [1, H, W, out_ch]
+            weight_permuted = weight_tensor.permute(1, 2, 3, 0).contiguous(
+                memory_format=torch.channels_last
+            )
+        else:
+            # For regular conv: OIHW -> OHWI
+            weight_permuted = weight_tensor.permute(0, 2, 3, 1).contiguous(
+                memory_format=torch.channels_last
+            )
 
         with node.graph.inserting_after(weight):
             weight_nhwc = create_constant_placeholder(
@@ -187,21 +200,28 @@ def _get_convolution_replacement(self, node) -> int:
                 torch.tensor(quantized_shifts, dtype=torch.int32),
             )
 
-        # Detect depthwise convolution: groups == input_channels
-        input_tensor = get_first_fake_tensor(x)
-        input_channels = input_tensor.shape[1]
-        is_depthwise = groups == input_channels
-
         if is_depthwise:
-            # Use depthwise convolution operator
+            # Compute depth_multiplier for depthwise convolution
+            # For depthwise: output_channels = input_channels * depth_multiplier
+            # PyTorch depthwise weight is [C_OUT, 1, H, W]
+            output_channels = weight_tensor.shape[0]
+            input_channels = groups  # For depthwise, groups == input_channels
+
+            if output_channels % input_channels != 0:
+                raise ValueError(
+                    f"Depthwise conv: output_channels ({output_channels}) must be "
+                    f"divisible by input_channels ({input_channels})"
+                )
+            depth_multiplier = output_channels // input_channels
+
             new_args = (
                 x,
                 weight_nhwc,
                 bias,
                 stride,
                 padding,
                 dilation,
-                groups,
+                depth_multiplier,
                 -input_zero_point,
                 output_zero_point,
                 quantized_multiplier_tensor,