From e6c7ba906416aa6bbca0c836108352b3b789fa77 Mon Sep 17 00:00:00 2001 From: Eylon Eliyahu Krause <64223519+EylonKrause@users.noreply.github.com> Date: Wed, 11 Nov 2026 13:49:57 +0200 Subject: [PATCH] Arm backend: Fix avg_pool2d divisor_override with ceil_mode and no padding DecomposeAvgPool2dPass lowers avg_pool2d with a divisor_override to avg_pool2d(count_include_pad=False) on a zero padded input, followed by a multiplication by kernel_size / divisor_override. That is only exact when every window is full. _compute_post_pad returned early for padding=0, so with ceil_mode=True the last, partial window was averaged over its valid elements only. For avg_pool2d(kernel=(1, 2), stride=(1, 2), padding=0, ceil_mode=True, divisor_override=2) on [1, 2, 3, 4, 5], PyTorch gives [1.5, 3.5, 2.5] (5 / 2) and the decomposition gave [1.5, 3.5, 5.0]. Compute the post padding for ceil_mode with a divisor_override also when padding=0, and insert the pad whenever any pad value is non-zero rather than only when the symmetric pre-padding is. Test plan: Ran the decomposition on the example above, and lowered avg_pool2d configurations with ceil_mode and divisor_override through TosaPipelineFP. Added test cases to test_avg_pool2d.py. Signed-off-by: Eylon Eliyahu Krause <64223519+EylonKrause@users.noreply.github.com> --- backends/arm/_passes/decompose_avg_pool2d_pass.py | 9 +++++++-- backends/arm/test/ops/test_avg_pool2d.py | 8 ++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/backends/arm/_passes/decompose_avg_pool2d_pass.py b/backends/arm/_passes/decompose_avg_pool2d_pass.py index e0c974598bc..f19c4fb515c 100644 --- a/backends/arm/_passes/decompose_avg_pool2d_pass.py +++ b/backends/arm/_passes/decompose_avg_pool2d_pass.py @@ -46,7 +46,12 @@ def _compute_post_pad( divisor_override, ) -> int | torch.SymInt: - if pad == 0: + # Without padding, count_include_pad=True divides a window that runs past + # the end of the input by its valid elements only, as the decomposed + # count_include_pad=False avg_pool2d does. A divisor_override needs full + # windows though, so the part of a ceil_mode window outside the input must + # still be padded with zeros. + if pad == 0 and not (ceil_mode and divisor_override is not None): return pad pad_adjust = adjust_pooling_pad_if_needed(size, kernel, stride, pad, ceil_mode) @@ -141,7 +146,7 @@ def call_operator(self, op, args, kwargs, meta): divisor_override, ) - if count_include_pad and (pad_h > 0 or pad_w > 0): + if count_include_pad and any(p != 0 for p in pad): if op in aten_avg_pool2d: pad_args = (x, pad, "constant", 0.0) else: diff --git a/backends/arm/test/ops/test_avg_pool2d.py b/backends/arm/test/ops/test_avg_pool2d.py index 9d2e863a31e..7ec9e3cb77d 100644 --- a/backends/arm/test/ops/test_avg_pool2d.py +++ b/backends/arm/test/ops/test_avg_pool2d.py @@ -110,6 +110,14 @@ def forward(self, x: torch.Tensor): AvgPool2d(3, 2, 1, True, True, divisor_override=2), (torch.rand(1, 1, 14, 14),), ), + "divisor_override_ceil_mode_no_padding": lambda: ( + AvgPool2d(3, 2, 0, True, False, divisor_override=2), + (torch.rand(1, 1, 15, 14),), + ), + "divisor_override_ceil_mode_no_padding_count_include_pad": lambda: ( + AvgPool2d(3, 2, 0, True, True, divisor_override=4), + (torch.rand(1, 1, 14, 14),), + ), "non_divisible_no_padding": lambda: ( AvgPool2d(3, 2, 0), (torch.rand(1, 16, 56, 56),),