From b9e74bd85cf9d339f1cc958dab33f03a45de826a 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 1/2] Arm backend: Fix ceil_mode pooling when the last window is dropped PyTorch drops a last ceil_mode pooling window that would start in the right padding, so e.g. max_pool2d(kernel=2, stride=4, padding=0, ceil_mode=True) on a width of 7 gives 2 output columns, the same as floor mode. The input then has a remainder that no window uses. SizeAdjustInputPass skipped every ceil_mode pool, so that remainder was never sliced off and adjust_pooling_pad_if_needed returned a negative post padding (-1 above). MAX_POOL2D and AVG_POOL2D were emitted with a negative pad, which the TOSA reference model rejects ("At least one pad is smaller than zero"). For avg_pool2d with count_include_pad=True, DecomposeAvgPool2dPass fell back to padding with the full symmetric pad instead, which made ceil_mode keep the dropped window and produced one output element too many. Slice the unused input for ceil_mode pools with static spatial shapes in every dimension where ceil_mode does not add an output element, and do not add post padding in DecomposeAvgPool2dPass when the needed post padding is negative. Test plan: Lowered 400 ceil_mode max_pool2d/avg_pool2d configurations (kernel 1-3, stride 1-4, padding 0-1, count_include_pad, divisor_override, input height 4-7) through TosaPipelineFP: 169 failed before, the remaining 40 failures are all padding=0 with divisor_override, which is a separate issue. Added test cases to test_max_pool.py and test_avg_pool2d.py. Signed-off-by: Eylon Eliyahu Krause <64223519+EylonKrause@users.noreply.github.com> --- .../arm/_passes/decompose_avg_pool2d_pass.py | 8 +- .../arm/_passes/size_adjust_input_pass.py | 79 +++++++++++++++---- backends/arm/test/ops/test_avg_pool2d.py | 8 ++ backends/arm/test/ops/test_max_pool.py | 8 ++ 4 files changed, 84 insertions(+), 19 deletions(-) diff --git a/backends/arm/_passes/decompose_avg_pool2d_pass.py b/backends/arm/_passes/decompose_avg_pool2d_pass.py index e0c974598bc..95d2e42b757 100644 --- a/backends/arm/_passes/decompose_avg_pool2d_pass.py +++ b/backends/arm/_passes/decompose_avg_pool2d_pass.py @@ -53,10 +53,14 @@ def _compute_post_pad( 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. + # Padding must never be below 0; the above adjustment is negative when the + # last window ends before the end of the input. No post padding is needed + # then, SizeAdjustInputPass truncates the unused input instead. Padding + # with ``pad`` would make ceil_mode keep a last window that PyTorch drops + # because it starts in the padding. if pad_adjust >= 0: return pad_adjust - return pad + return 0 def _get_avgpool_post_pad( diff --git a/backends/arm/_passes/size_adjust_input_pass.py b/backends/arm/_passes/size_adjust_input_pass.py index e26d12e8660..c6eb1bc67c9 100644 --- a/backends/arm/_passes/size_adjust_input_pass.py +++ b/backends/arm/_passes/size_adjust_input_pass.py @@ -96,6 +96,42 @@ def pooling_remainder( return (input_size + 2 * pad - kernel_size) % stride +def _get_pooling_ceil_mode(pooling_node: torch.fx.Node) -> bool: + ceil_mode_index = ( + 5 + if pooling_node.target + in ( + max_pooling_op, + exir_ops.edge.aten.max_pool2d_with_indices.default, + ) + else 4 + ) + return bool( + pooling_node.args[ceil_mode_index] + if len(pooling_node.args) > ceil_mode_index + else False + ) + + +def _ceil_mode_adds_output( + input_size: int, pad: int, kernel_size: int, stride: int +) -> bool: + """Return whether ceil_mode gives one more output element than floor mode. + + PyTorch drops a last ceil_mode window that would start in the right + padding. Ceil mode then pools over the same windows as floor mode, and the + input has to be truncated in the same way. + + """ + ceil_post_pad = adjust_pooling_pad_if_needed( + input_size, kernel_size, stride, pad, True + ) + floor_post_pad = adjust_pooling_pad_if_needed( + input_size, kernel_size, stride, pad, False + ) + return ceil_post_pad != floor_post_pad + + def has_dynamic_pooling_padding(pooling_node: torch.fx.Node) -> bool: """Return whether rewriting a pool requires symbolic padding. @@ -114,20 +150,7 @@ def has_dynamic_pooling_padding(pooling_node: torch.fx.Node) -> bool: else kernel_size ) padding = pooling_node.args[3] if len(pooling_node.args) >= 4 else 0 - ceil_mode_index = ( - 5 - if pooling_node.target - in ( - max_pooling_op, - exir_ops.edge.aten.max_pool2d_with_indices.default, - ) - else 4 - ) - ceil_mode = ( - pooling_node.args[ceil_mode_index] - if len(pooling_node.args) > ceil_mode_index - else False - ) + ceil_mode = _get_pooling_ceil_mode(pooling_node) input_shape = input_node.meta["val"].shape kernel_sizes = expand_around_channel(cast(Sequence[int] | int, kernel_size), 2) strides = expand_around_channel(cast(Sequence[int] | int, stride), 2) @@ -243,6 +266,7 @@ def get_slices_pooling(pooling_node: torch.fx.Node) -> Slices: else kernel_size ) padding = pooling_node.args[3] if len(pooling_node.args) >= 4 else 0 + ceil_mode = _get_pooling_ceil_mode(pooling_node) input_shape = cast(torch.fx.Node, input_node).meta["val"].shape @@ -253,6 +277,16 @@ def get_slices_pooling(pooling_node: torch.fx.Node) -> Slices: for dim, (kernel_length, stride_length, pad_size) in enumerate( zip(kernel_sizes, strides, pads), start=2 ): + # When ceil_mode adds an output element, its window reaches the end of + # the input, so there is nothing to truncate in this dimension. + if ( + ceil_mode + and not isinstance(input_shape[dim], torch.SymInt) + and _ceil_mode_adds_output( + input_shape[dim], pad_size, kernel_length, stride_length + ) + ): + continue remainder = pooling_remainder( input_shape[dim], pad_size, kernel_length, stride_length ) @@ -274,6 +308,11 @@ def get_slices(node: torch.fx.Node) -> Slices: raise ValueError(f"Unsupported node target, was expecting {valid_operators}") +def _has_static_spatial_shape(node: torch.fx.Node) -> bool: + input_shape = cast(torch.fx.Node, node.args[0]).meta["val"].shape + return not any(isinstance(dim, torch.SymInt) for dim in input_shape[2:]) + + def is_valid_operator(node: torch.fx.Node) -> bool: if node.target == conv2d_op: return True @@ -294,14 +333,20 @@ def is_valid_operator(node: torch.fx.Node) -> bool: "Expected max_pool2d with dilation = [1, 1], has DecomposeMaxPool2dPass been run?" ) - # If using ceil mode for rounding, the input does not need adjusting - return not ceil_mode + # Ceil mode needs the same input adjustment as floor mode where it does + # not add an output element, see get_slices_pooling. This is only + # handled for static spatial shapes. + return not ceil_mode or _has_static_spatial_shape(node) elif node.target == avg_pooling_op: ceil_mode = node.args[4] if len(node.args) >= 5 else False count_include_pad = node.args[5] if len(node.args) >= 6 else True divisor_override = node.args[6] if len(node.args) >= 7 else None - return not ceil_mode and not count_include_pad and divisor_override is None + return ( + (not ceil_mode or _has_static_spatial_shape(node)) + and not count_include_pad + and divisor_override is None + ) return False diff --git a/backends/arm/test/ops/test_avg_pool2d.py b/backends/arm/test/ops/test_avg_pool2d.py index 9d2e863a31e..b662d1dbf7f 100644 --- a/backends/arm/test/ops/test_avg_pool2d.py +++ b/backends/arm/test/ops/test_avg_pool2d.py @@ -94,6 +94,14 @@ def forward(self, x: torch.Tensor): AvgPool2d(3, 3, 1, True, True), (torch.rand(1, 1, 5, 5),), ), + "ceil_mode_last_window_dropped": lambda: ( + AvgPool2d(1, 2, 0, True, False), + (torch.rand(1, 16, 56, 57),), + ), + "ceil_mode_last_window_dropped_count_include_pad": lambda: ( + AvgPool2d(2, 3, 1, True, True), + (torch.rand(1, 16, 56, 56),), + ), "divisor_override": lambda: ( AvgPool2d(3, 2, 1, False, False, divisor_override=2), (torch.rand(1, 1, 14, 14),), diff --git a/backends/arm/test/ops/test_max_pool.py b/backends/arm/test/ops/test_max_pool.py index 24af4857282..e49865a7577 100644 --- a/backends/arm/test/ops/test_max_pool.py +++ b/backends/arm/test/ops/test_max_pool.py @@ -45,6 +45,14 @@ torch.rand(1, 16, 56, 56), [3, (1, 2), 1, 1, True], ), + "ceil_mode_last_window_dropped": lambda: ( + torch.rand(1, 16, 56, 57), + [1, 2, 0, 1, True], + ), + "ceil_mode_last_window_dropped_padding": lambda: ( + torch.rand(1, 16, 56, 56), + [2, 3, 1, 1, True], + ), "non_divisible_window_adjust_padding": lambda: ( torch.rand(1, 16, 112, 112), [3, 2, 1], From 4f2217595ee22690dd7d120b395f34632f9c45cc Mon Sep 17 00:00:00 2001 From: Eylon Eliyahu Krause <64223519+EylonKrause@users.noreply.github.com> Date: Fri, 13 Nov 2026 15:18:09 +0200 Subject: [PATCH 2/2] Arm backend: Fix lint in SizeAdjustInputPass Return a plain bool from _ceil_mode_adds_output (mypy: the comparison of two adjust_pooling_pad_if_needed results is typed bool | SymBool), and rewrap its docstring as docformatter expects. Signed-off-by: Eylon Eliyahu Krause <64223519+EylonKrause@users.noreply.github.com> --- backends/arm/_passes/size_adjust_input_pass.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/backends/arm/_passes/size_adjust_input_pass.py b/backends/arm/_passes/size_adjust_input_pass.py index c6eb1bc67c9..3783888972c 100644 --- a/backends/arm/_passes/size_adjust_input_pass.py +++ b/backends/arm/_passes/size_adjust_input_pass.py @@ -118,9 +118,9 @@ def _ceil_mode_adds_output( ) -> bool: """Return whether ceil_mode gives one more output element than floor mode. - PyTorch drops a last ceil_mode window that would start in the right - padding. Ceil mode then pools over the same windows as floor mode, and the - input has to be truncated in the same way. + PyTorch drops a last ceil_mode window that would start in the right padding. + Ceil mode then pools over the same windows as floor mode, and the input has + to be truncated in the same way. """ ceil_post_pad = adjust_pooling_pad_if_needed( @@ -129,7 +129,7 @@ def _ceil_mode_adds_output( floor_post_pad = adjust_pooling_pad_if_needed( input_size, kernel_size, stride, pad, False ) - return ceil_post_pad != floor_post_pad + return bool(ceil_post_pad != floor_post_pad) def has_dynamic_pooling_padding(pooling_node: torch.fx.Node) -> bool: