Cortex-M: Support avg_pool2d ceil_mode lowering

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

@@ -819,7 +819,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])
@@ -847,6 +847,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,