Arm backend: Fix ceil_mode pooling when the last window is dropped

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(

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 bool(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
 

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),),

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],