Arm backend: Decompose avg_pool2d count_include_pad via pad

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

Author: Baris Demir

backends/cortex_m/ops/op_quantized_avg_pool2d.cpp

@@ -31,8 +31,6 @@ Tensor& quantized_avg_pool2d_out(
 
   const int64_t dilation_values[2] = {1, 1};
   const Int64ArrayRef dilation(dilation_values, 2);
-  const bool ceil_mode = false;
-
   CmsisPool2DConfig pool_config;
   if (!prepare_cmsis_pool2d_config(
           context,
@@ -43,7 +41,7 @@ Tensor& quantized_avg_pool2d_out(
           stride,
           padding,
           dilation,
-          ceil_mode,
+          false,
           activation_min,
           activation_max,
           pool_config)) {
@@ -57,7 +55,7 @@ Tensor& quantized_avg_pool2d_out(
   const int8_t* input_data = input.const_data_ptr<int8_t>();
   int8_t* output_data = out.mutable_data_ptr<int8_t>();
 
-  arm_cmsis_nn_status status = arm_avgpool_s8(
+  const arm_cmsis_nn_status status = arm_avgpool_s8(
       &cmsis_ctx,
       &pool_config.pool_params,
       &pool_config.input_dims,

backends/cortex_m/ops/operators.py

@@ -1110,16 +1110,16 @@ def quantized_avg_pool2d_meta(
     multiplier: int,
     shift: int,
 ) -> torch.Tensor:
-    kernel = _ensure_tuple2(kernel_size)
-    stride_vals = _ensure_tuple2(stride)
-    padding_vals = _ensure_tuple2(padding)
-    dilation_vals = (1, 1)
-
-    output_shape = _compute_max_pool2d_output_shape(
-        input.shape, kernel, stride_vals, padding_vals, dilation_vals
+    output = F.avg_pool2d(
+        input.to(torch.float),
+        kernel_size,
+        stride=stride,
+        padding=padding,
+        ceil_mode=False,
+        count_include_pad=False,
     )
     return torch.empty(
-        output_shape,
+        output.shape,
         dtype=torch.int8,
         device=input.device,
         memory_format=torch.channels_last,
@@ -1136,21 +1136,20 @@ def quantized_avg_pool2d_impl(
     multiplier: int,
     shift: int,
 ) -> torch.Tensor:
-
     dequant_input = dequantize_per_tensor_cmsis(input, zero_point, multiplier, shift)
 
     kernel = _ensure_tuple2(kernel_size)
     stride_vals = _ensure_tuple2(stride)
     padding_vals = _ensure_tuple2(padding)
 
-    # TODO: implement count_include_pad=True, ceil_mode=True, dilation != 1.
+    # TODO: implement dilation != 1.
     result = F.avg_pool2d(
         dequant_input,
         kernel,
         stride=stride_vals,
         padding=padding_vals,
-        count_include_pad=False,
         ceil_mode=False,
+        count_include_pad=False,
     )
     result = quantize_per_tensor_cmsis(result, zero_point, multiplier, shift)
     output = torch.clamp(result, -128, 127)

backends/cortex_m/passes/quantized_op_fusion_pass.py

@@ -360,14 +360,29 @@ def _get_avg_pool2d_replacement(self, args, meta):
         divisor_override = args[6] if len(args) > 6 else None
         divisor_override_val = self._unwrap_argument(divisor_override)
 
-        if ceil_mode or count_include_pad or divisor_override_val is not None:
+        if ceil_mode or divisor_override_val is not None:
             return exir_ops.edge.aten.avg_pool2d.default, args
 
+        input_arg = args[0]
+        avg_padding = padding
+        if count_include_pad:
+            # Decompose count_include_pad=True into explicit input padding.
+            pad_h, pad_w = padding
+            pre_pad = [0, 0, pad_h, pad_w]
+            post_pad = [0, 0, pad_h, pad_w]
+            input_arg = super().call_operator(
+                exir_ops.edge.cortex_m.pad.default,
+                (input_arg, pre_pad, post_pad, int(zero_point)),
+                {},
+                NodeMetadata({}),
+            )
+            avg_padding = [0, 0]
+
         args = (
-            args[0],
+            input_arg,
             kernel_size,
             stride,
-            padding,
+            avg_padding,
             zero_point,
             output_mult,
             output_shift,

backends/cortex_m/quantizer/pattern_checkers.py

@@ -291,8 +291,7 @@ def check_pattern(cls, pattern):
             return False
         node = pattern[0]
         ceil_mode = cast(bool, node.args[4]) if len(node.args) > 4 else False
-        count_include_pad = cast(bool, node.args[5]) if len(node.args) > 5 else True
-        return not (ceil_mode or count_include_pad)
+        return not ceil_mode
 
     @classmethod
     def check_quantization_config(

backends/cortex_m/test/ops/test_avg_pool2d.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -59,13 +59,6 @@ def forward(self, x):  # noqa: D102
         CortexMAvgPool2d(kernel_size=3, stride=2, padding=1),
         (ramp_tensor(0, 15, (1, 1, 4, 4)),),
     ),
-}
-
-test_cases_fp = {
-    "avgpool_3x3_s2_pad1_ceil": 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": McuTestCase(
         CortexMAvgPool2d(kernel_size=3, stride=2, padding=1, count_include_pad=True),
         (ramp_tensor(0, 15, (1, 1, 4, 4)),),
@@ -76,19 +69,12 @@ def forward(self, x):  # noqa: D102
 @parametrize("test_case", test_cases)
 def test_dialect_avg_pool2d(test_case):
     tester = CortexMTester(test_case.model, test_case.example_inputs)
+    ops_after = dict(test_case.model.ops_after_transforms)
+    if test_case.model.pool.count_include_pad:
+        ops_after["executorch_exir_dialects_edge__ops_cortex_m_pad_default"] = 1
     tester.test_dialect(
         test_case.model.ops_before_transforms,
-        test_case.model.ops_after_transforms,
-        qtol=1,
-    )
-
-
-@parametrize("test_case", test_cases_fp)
-def test_dialect_avg_pool2d_fp(test_case):
-    tester = CortexMTester(test_case.model, test_case.example_inputs)
-    tester.test_dialect(
-        {"executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1},
-        {"executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1},
+        ops_after,
         qtol=1,
     )