Cortex-M: Support avg_pool2d ceil_mode lowering (#21039)

Cortex-M avg_pool2d quantized lowering rejected ceil_mode=True and left
the graph on the fallback aten.avg_pool2d path. That preserved
correctness but prevented optimized Cortex-M lowering for cases that
CMSIS-NN can execute from statically known output shapes.

This patch threads ceil_mode through the cortex_m::quantized_avg_pool2d
schema, fake implementation, lowering, and C++ kernel wrapper. The
wrapper validates the static output shape against ceil-mode pooling
semantics before dispatching to CMSIS-NN.

The avg_pool2d tests now cover optimized ceil_mode lowering, padded
ceil_mode, count_include_pad with ceil_mode, non-square multi-channel
ceil_mode, and divisor_override fallback behavior.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

---------

Signed-off-by: Baris Demir <baris.demir@arm.com>

Author: bdemirb

backends/arm/_passes/decompose_avg_pool2d_pass.py

@@ -48,13 +48,13 @@ def _compute_post_pad(
 
     if pad == 0:
         return pad
-    if ceil_mode and divisor_override is None:
-        return pad
-
     pad_adjust = adjust_pooling_pad_if_needed(size, kernel, stride, pad, ceil_mode)
 
-    # Padding must always be above 0, the above adjustment may return -1
-    if pad_adjust > 0:
+    if ceil_mode and divisor_override is None and pad_adjust > pad:
+        return pad
+
+    # Padding must never be below 0; the above adjustment may return -1.
+    if pad_adjust >= 0:
         return pad_adjust
     return pad
 

backends/arm/operators/operator_validation_utils.py

@@ -199,6 +199,8 @@ def adjust_pooling_pad_if_needed(
     numerator = input_size - kernel_size + 2 * pad
     if ceil_mode:
         output_size = (numerator + stride - 1) // stride + 1
+        if (output_size - 1) * stride >= input_size + pad:
+            output_size -= 1
     else:
         output_size = numerator // stride + 1
 

backends/arm/test/ops/test_avg_pool2d.py

@@ -91,6 +91,10 @@ def forward(self, x: torch.Tensor):
         AvgPool2d(3, (1, 2), 1, True, True),
         (torch.rand(1, 1, 14, 14),),
     ),
+    "ceil_mode_count_include_pad_kernel_stride": lambda: (
+        AvgPool2d(3, 3, 1, True, True),
+        (torch.rand(1, 1, 5, 5),),
+    ),
     "divisor_override": lambda: (
         AvgPool2d(3, 2, 1, False, False, divisor_override=2),
         (torch.rand(1, 1, 14, 14),),

backends/cortex_m/ops/op_quantized_avg_pool2d.cpp

@@ -12,13 +12,67 @@ namespace native {
 
 using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
 
+namespace {
+
+int32_t pooling_output_size(
+    int32_t input,
+    int32_t kernel,
+    int32_t stride,
+    int32_t padding,
+    bool ceil_mode) {
+  const int32_t numerator = input + 2 * padding - kernel;
+  int32_t output = numerator / stride + 1;
+  if (ceil_mode && numerator % stride != 0) {
+    output += 1;
+  }
+  if (ceil_mode && (output - 1) * stride >= input + padding) {
+    output -= 1;
+  }
+  return output;
+}
+
+bool validate_avg_pool2d_output_size(
+    KernelRuntimeContext& context,
+    const CmsisPool2DConfig& pool_config,
+    bool ceil_mode) {
+  const int32_t expected_h = pooling_output_size(
+      pool_config.input_dims.h,
+      pool_config.filter_dims.h,
+      pool_config.pool_params.stride.h,
+      pool_config.pool_params.padding.h,
+      ceil_mode);
+  const int32_t expected_w = pooling_output_size(
+      pool_config.input_dims.w,
+      pool_config.filter_dims.w,
+      pool_config.pool_params.stride.w,
+      pool_config.pool_params.padding.w,
+      ceil_mode);
+
+  if (pool_config.output_dims.h != expected_h ||
+      pool_config.output_dims.w != expected_w) {
+    ET_LOG(
+        Error,
+        "quantized_avg_pool2d_out: output shape mismatch - actual: (%d, %d) expected: (%d, %d)",
+        pool_config.output_dims.h,
+        pool_config.output_dims.w,
+        expected_h,
+        expected_w);
+    context.fail(Error::InvalidArgument);
+    return false;
+  }
+  return true;
+}
+
+} // namespace
+
 // cppcheck-suppress unusedFunction
 Tensor& quantized_avg_pool2d_out(
     KernelRuntimeContext& context,
     const Tensor& input,
     const Int64ArrayRef kernel_size,
     const Int64ArrayRef stride,
     const Int64ArrayRef padding,
+    const bool ceil_mode,
     const int64_t zero_point,
     const int64_t multiplier,
     const int64_t shift,
@@ -39,10 +93,16 @@ Tensor& quantized_avg_pool2d_out(
           stride,
           padding,
           dilation,
-          false,
+          ceil_mode,
           activation_min,
           activation_max,
-          pool_config)) {
+          pool_config,
+          true,
+          true)) {
+    return out;
+  }
+
+  if (!validate_avg_pool2d_output_size(context, pool_config, ceil_mode)) {
     return out;
   }
 

backends/cortex_m/ops/operators.py

@@ -1215,6 +1215,7 @@ def quantized_transpose_conv2d_impl(
     "int[] kernel_size, "
     "int[] stride, "
     "int[] padding, "
+    "bool ceil_mode, "
     "int zero_point, "
     "int multiplier, "
     "int shift, "
@@ -1227,6 +1228,7 @@ def quantized_transpose_conv2d_impl(
     "int[] kernel_size, "
     "int[] stride, "
     "int[] padding, "
+    "bool ceil_mode, "
     "int zero_point, "
     "int multiplier, "
     "int shift, "
@@ -1241,6 +1243,7 @@ def quantized_avg_pool2d_meta(
     kernel_size: Sequence[int],
     stride: Sequence[int],
     padding: Sequence[int],
+    ceil_mode: bool,
     zero_point: int,
     multiplier: int,
     shift: int,
@@ -1254,7 +1257,7 @@ def quantized_avg_pool2d_meta(
         kernel,
         stride=stride_vals,
         padding=padding_vals,
-        ceil_mode=False,
+        ceil_mode=ceil_mode,
         count_include_pad=False,
     )
     return torch.empty(
@@ -1271,6 +1274,7 @@ def quantized_avg_pool2d_impl(
     kernel_size: Sequence[int],
     stride: Sequence[int],
     padding: Sequence[int],
+    ceil_mode: bool,
     zero_point: int,
     multiplier: int,
     shift: int,
@@ -1288,7 +1292,7 @@ def quantized_avg_pool2d_impl(
         kernel,
         stride=stride_vals,
         padding=padding_vals,
-        ceil_mode=False,
+        ceil_mode=ceil_mode,
         count_include_pad=False,
     )
     result = quantize_per_tensor_cmsis(result, zero_point, multiplier, shift)

backends/cortex_m/ops/operators.yaml

@@ -90,7 +90,7 @@
     - arg_meta: null
       kernel_name: cortex_m::quantized_transpose_conv2d_out
 
-- func: cortex_m::quantized_avg_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)
+- func: cortex_m::quantized_avg_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:
     - arg_meta: null

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -820,7 +820,7 @@ def _get_avg_pool2d_replacement(
     count_include_pad = cast(bool, pool_args[5]) if len(pool_args) > 5 else True
     divisor_override = pool_args[6] if len(pool_args) > 6 else None
 
-    if ceil_mode or divisor_override is not None:
+    if divisor_override is not None:
         return None
 
     input_node = cast(Node, pool_args[0])
@@ -848,6 +848,7 @@ def _get_avg_pool2d_replacement(
         kernel_size,
         stride,
         avg_padding,
+        ceil_mode,
         int(input_zp),
         int(output_mult),
         int(output_shift),

backends/cortex_m/quantizer/pattern_checkers.py

@@ -3,8 +3,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import cast
-
 import torch
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.quantizer.arm_quantizer_utils import PatternCheck
@@ -250,8 +248,8 @@ def check_pattern(cls, pattern):
         if not pattern:
             return False
         node = pattern[0]
-        ceil_mode = cast(bool, node.args[4]) if len(node.args) > 4 else False
-        return not ceil_mode
+        divisor_override = node.args[6] if len(node.args) > 6 else None
+        return divisor_override is None
 
     @classmethod
     def check_quantization_config(

backends/cortex_m/test/ops/test_avg_pool2d.py

@@ -28,7 +28,13 @@ class CortexMAvgPool2d(torch.nn.Module):
     }
 
     def __init__(
-        self, kernel_size, stride, padding=0, ceil_mode=False, count_include_pad=False
+        self,
+        kernel_size,
+        stride,
+        padding=0,
+        ceil_mode=False,
+        count_include_pad=False,
+        divisor_override=None,
     ):
         super().__init__()
         self.pool = torch.nn.AvgPool2d(
@@ -37,6 +43,7 @@ def __init__(
             padding,
             ceil_mode=ceil_mode,
             count_include_pad=count_include_pad,
+            divisor_override=divisor_override,
         )
 
     def forward(self, x):  # noqa: D102
@@ -68,14 +75,37 @@ def forward(self, x):  # noqa: D102
 }
 
 
-# ceil_mode=True is not supported by the CMSIS-NN avg_pool kernel; the convert
-# pass leaves aten.avg_pool2d in the graph for a portable kernel to handle. The
-# Cortex-M runner does not register aten.avg_pool2d, so this is dialect-only.
-fallback_test_cases = {
+ceil_mode_test_cases = {
     "avgpool_2x2_ceil_mode": McuTestCase(
         CortexMAvgPool2d(kernel_size=2, stride=2, ceil_mode=True),
         (ramp_tensor(0, 24, (1, 1, 5, 5)),),
     ),
+    "avgpool_3x3_s2_pad1_ceil_mode": McuTestCase(
+        CortexMAvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True),
+        (ramp_tensor(0, 15, (1, 1, 4, 4)),),
+    ),
+    "avgpool_3x3_s2_pad1_countinc_ceil_mode": McuTestCase(
+        CortexMAvgPool2d(
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            ceil_mode=True,
+            count_include_pad=True,
+        ),
+        (ramp_tensor(0, 15, (1, 1, 4, 4)),),
+    ),
+    "avgpool_2x3_s2x3_ceil_mode": McuTestCase(
+        CortexMAvgPool2d(kernel_size=(2, 3), stride=(2, 3), ceil_mode=True),
+        (ramp_tensor(0, 79, (1, 2, 5, 8)).to(memory_format=torch.channels_last),),
+    ),
+}
+
+
+fallback_test_cases = {
+    "avgpool_2x2_divisor_override": McuTestCase(
+        CortexMAvgPool2d(kernel_size=2, stride=2, divisor_override=2),
+        (ramp_tensor(0, 15, (1, 1, 4, 4)),),
+    ),
 }
 
 
@@ -119,15 +149,33 @@ def test_dialect_avg_pool2d(test_case, cortex_m_target):
     ), f"scratch buffer size mismatch: got {scratch_size}, expected {expected_size}"
 
 
+@parametrize("test_case", ceil_mode_test_cases)
+def test_dialect_avg_pool2d_ceil_mode(test_case, cortex_m_target):
+    tester = CortexMTester(
+        test_case.model, test_case.example_inputs, target_config=cortex_m_target
+    )
+    tester.test_dialect(
+        test_case.model.ops_before_transforms,
+        test_case.model.ops_after_transforms,
+        qtol=1,
+    )
+
+
+@parametrize("test_case", ceil_mode_test_cases)
+def test_implementation_avg_pool2d_ceil_mode(test_case, cortex_m_target):
+    tester = CortexMTester(
+        test_case.model, test_case.example_inputs, target_config=cortex_m_target
+    )
+    tester.test_implementation(qtol=1)
+
+
 @parametrize("test_case", fallback_test_cases)
-def test_dialect_avg_pool2d_fallback(test_case):
-    tester = CortexMTester(test_case.model, test_case.example_inputs)
+def test_dialect_avg_pool2d_fallback(test_case, cortex_m_target):
+    tester = CortexMTester(
+        test_case.model, test_case.example_inputs, target_config=cortex_m_target
+    )
     tester.test_dialect(
-        {
-            "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
-            "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
-            "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 2,
-        },
+        test_case.model.ops_before_transforms,
         {
             "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
             "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2,