Cortex-M: lower a max pool whose ceil_mode changes nothing

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -7,7 +7,7 @@
 
 import math
 import operator
-from typing import cast, Optional
+from typing import cast
 
 import executorch.backends.cortex_m.ops.operators  # noqa
 import executorch.exir as exir
@@ -18,7 +18,9 @@
 
 from executorch.backends.cortex_m.passes.passes_utils import (
     build_activation_lut,
+    ceil_mode_is_redundant,
     is_foldable_alpha,
+    max_pool2d_params,
     quantize_multiplier_aot,
     quantize_val,
     SHIFT_INT8,
@@ -67,9 +69,9 @@ def __init__(
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         result = super().call(graph_module)
-        # RemoveGetItemPass converts max_pool2d_with_indices to max_pool2d. Models
-        # such as GoogleNet and SqueezeNet use configurations unsupported by the
-        # Cortex-M kernel, so restore the portable with-indices fallback for them.
+        # RemoveGetItemPass converts max_pool2d_with_indices to max_pool2d, so
+        # any pool that did not lower has to be put back into the spelling the
+        # portable kernel is registered under.
         max_pool_modified = _restore_max_pool2d_with_indices_fallback(
             result.graph_module
         )
@@ -188,28 +190,8 @@ def _restore_max_pool2d_with_indices_fallback(
 _SOFTMAX_INPUT_INTEGER_BITS = 5
 
 
-def _to_int_pair(
-    value: Argument, default: Optional[tuple[int, int]]
-) -> tuple[int, int]:
-    if value is None:
-        assert default is not None, "Expected default sequence for normalization"
-        return (default[0], default[1])
-
-    try:
-        int_pair = cast(tuple[int, int], value)
-        return int_pair
-    except Exception as exc:
-        raise ValueError(f"Expected a tuple of two integers, got {value}") from exc
-
-
 def _to_bool(value: Argument, default: bool) -> bool:
-    if value is None:
-        return default
-    try:
-        bool_value = cast(bool, value)
-        return bool_value
-    except Exception as exc:
-        raise ValueError(f"Expected a boolean value, got {value}") from exc
+    return default if value is None else bool(value)
 
 
 def _is_quant_per_tensor_qualified(node: Node) -> bool:
@@ -1140,16 +1122,20 @@ def _get_max_pool2d_replacement(
         activation_min = torch.iinfo(torch.int8).min
         activation_max = torch.iinfo(torch.int8).max
 
-    kernel_size = _to_int_pair(node.args[1], None)
-    stride_arg = node.args[2] if len(node.args) > 2 else None
-    stride = _to_int_pair(stride_arg, kernel_size)
-    padding_arg = node.args[3] if len(node.args) > 3 else None
-    padding = _to_int_pair(padding_arg, (0, 0))
-    dilation_arg = node.args[4] if len(node.args) > 4 else None
-    dilation = _to_int_pair(dilation_arg, (1, 1))
+    kernel_size, stride, padding, dilation = max_pool2d_params(node)
     ceil_mode_arg = node.args[5] if len(node.args) > 5 else False
     ceil_mode = _to_bool(ceil_mode_arg, False)
 
+    if ceil_mode and ceil_mode_is_redundant(
+        get_first_fake_tensor(node.all_input_nodes[0]).shape[-2:],
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+    ):
+        # The kernel always floors, which here lands on the same output size.
+        ceil_mode = False
+
     if dilation != (1, 1) or ceil_mode:
         return None
 

backends/cortex_m/passes/passes_utils.py

@@ -6,7 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
-from typing import Any, Callable, TypeGuard
+from typing import Any, Callable, NamedTuple, TypeGuard
 
 import torch
 
@@ -185,6 +185,53 @@ def is_foldable_alpha(alpha: Any) -> TypeGuard[int]:
     return isinstance(alpha, int)
 
 
+class MaxPool2dParams(NamedTuple):
+    kernel: tuple[int, int]
+    stride: tuple[int, int]
+    padding: tuple[int, int]
+    dilation: tuple[int, int]
+
+
+def max_pool2d_params(node: Node) -> MaxPool2dParams:
+    """The kernel, stride, padding and dilation of a max_pool2d node.
+
+    stride's schema default is the empty list, meaning "same as the kernel",
+    and export leaves it in place whenever a later argument is named -- which
+    ceil_mode always is. int[1] is a legal spelling of an int[2] argument too.
+    """
+    args = node.args
+    # kernel_size is the one required argument, so it is always present.
+    kernel = coerce_int_pair(args[1], (1, 1))
+    return MaxPool2dParams(
+        kernel,
+        coerce_int_pair(args[2] if len(args) > 2 else None, kernel),
+        coerce_int_pair(args[3] if len(args) > 3 else None, (0, 0)),
+        coerce_int_pair(args[4] if len(args) > 4 else None, (1, 1)),
+    )
+
+
+def ceil_mode_is_redundant(
+    input_hw: torch.Size,
+    kernel: tuple[int, int],
+    stride: tuple[int, int],
+    padding: tuple[int, int],
+    dilation: tuple[int, int],
+) -> bool:
+    """Whether ceil_mode picks the same output size that floor_mode would.
+
+    Rounding up adds a row or column whenever the stride does not divide the
+    span the windows have to cover. The converse does not hold -- aten drops a
+    last window that starts inside the right padding, which can bring the
+    ceiling back to the floor -- so this refuses some it could accept.
+    """
+    if any(not isinstance(n, int) for n in input_hw):
+        return False
+    return all(
+        (n + 2 * p - d * (k - 1) - 1) % s == 0
+        for n, k, s, p, d in zip(input_hw, kernel, stride, padding, dilation)
+    )
+
+
 def is_qualified_int8_node(args) -> bool:
     try:
         if len(args) < 6:

backends/cortex_m/quantizer/pattern_checkers.py

@@ -8,10 +8,11 @@
 from executorch.backends.arm.quantizer.arm_quantizer_utils import PatternCheck
 from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
 from executorch.backends.cortex_m.passes.passes_utils import (
-    coerce_int_pair,
+    ceil_mode_is_redundant,
     is_channel_broadcast,
     is_channels_last,
     is_foldable_alpha,
+    max_pool2d_params,
 )
 from executorch.backends.cortex_m.quantizer.quantization_configs import (
     CMSIS_SOFTMAX_SCALE,
@@ -358,14 +359,20 @@ def _pool_arg_as_bool(cls, node: Node, index: int, default: bool) -> bool:
             return default
         return bool(raw)
 
+    @classmethod
+    def _ceil_mode_is_redundant(cls, node: Node) -> bool:
+        shape = get_first_fake_tensor(node.all_input_nodes[0]).shape
+        return ceil_mode_is_redundant(shape[-2:], *max_pool2d_params(node))
+
     @classmethod
     def check_pattern(cls, pattern):
         if not pattern:
             return False
         node = pattern[0]
-        raw_dilation = node.args[4] if len(node.args) > 4 else (1, 1)
-        dilation = coerce_int_pair(raw_dilation, (1, 1))
+        dilation = max_pool2d_params(node).dilation
         ceil_mode = cls._pool_arg_as_bool(node, 5, False)
+        if ceil_mode and cls._ceil_mode_is_redundant(node):
+            ceil_mode = False
         if dilation != (1, 1) or ceil_mode:
             meta_custom = node.meta.get("custom", {})
             cortex_m_meta = meta_custom.get("cortex_m", {})

backends/cortex_m/test/ops/test_max_pool2d.py

@@ -8,6 +8,7 @@
 
 import torch
 from executorch.backends.arm.test.common import parametrize, xfail_type
+from executorch.backends.cortex_m.passes.passes_utils import ceil_mode_is_redundant
 from executorch.backends.cortex_m.test.tester import (
     CortexMTester,
     McuTestCase,
@@ -64,6 +65,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.pool(x.permute(0, 3, 1, 2))
 
 
+class CortexMMaxPool2dFunctional(torch.nn.Module):
+    ops_before_transforms = CortexMMaxPool2d.ops_before_transforms
+    ops_after_transforms = CortexMMaxPool2d.ops_after_transforms
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.kwargs = kwargs
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.max_pool2d(x, **self.kwargs)
+
+
 class CortexMMaxPool2dIndices(torch.nn.Module):
     ops_before_transforms = CortexMMaxPool2d.ops_before_transforms
     ops_after_transforms = CortexMMaxPool2d.ops_after_transforms
@@ -110,6 +123,26 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         CortexMMaxPool2dPermutedView(kernel_size=2, stride=2),
         ((torch.randn(1, 24, 24, 1) * 30),),
     ),
+    # SqueezeNet's pool: ceil_mode is set but 11 - 3 divides by 2, so rounding
+    # up picks the same output size.
+    "maxpool_3x3_s2_ceil_redundant": McuTestCase(
+        CortexMMaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+        (ramp_tensor(-20, 20, (1, 1, 11, 11)),),
+    ),
+    # An odd stride, so the padding term decides the answer, over a 169-element
+    # plane -- past the 127 that aten's channels-last int8 pool accepts, so
+    # this one could not have been written as a fallback case.
+    "maxpool_3x3_s3_pad1_ceil_redundant": McuTestCase(
+        CortexMMaxPool2d(kernel_size=3, stride=3, padding=1, ceil_mode=True),
+        ((torch.randn(1, 8, 13, 13) * 30).to(memory_format=torch.channels_last),),
+    ),
+    # nn.MaxPool2d always fills in a stride; the functional spelling leaves the
+    # schema default of [] in the graph, which only appears once a later
+    # argument is named.
+    "maxpool_3x3_default_stride_ceil": McuTestCase(
+        CortexMMaxPool2dFunctional(kernel_size=3, ceil_mode=True),
+        (ramp_tensor(-20, 20, (1, 1, 12, 12)),),
+    ),
 }
 
 
@@ -118,6 +151,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         CortexMMaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True),
         (ramp_tensor(-10, 10, (1, 1, 4, 4)),),
     ),
+    # The height divides but the width does not, so the ceiling still adds a
+    # column and the whole pool has to decline.
+    "maxpool_3x3_s2_ceil_one_axis": McuTestCase(
+        CortexMMaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+        (ramp_tensor(-20, 20, (1, 1, 11, 12)),),
+    ),
     "maxpool_dilation": McuTestCase(
         CortexMMaxPool2d(kernel_size=2, stride=1, padding=0, dilation=2),
         (ramp_tensor(-25, 25, (1, 1, 6, 6)),),
@@ -132,6 +171,48 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 }
 
 
+def test_ceil_mode_redundancy_reads_every_term():
+    """Each end-to-end case exercises one configuration, leaving every term it
+    does not vary unconstrained."""
+    square = torch.Size([11, 11])
+    assert ceil_mode_is_redundant(square, (3, 3), (2, 2), (0, 0), (1, 1))
+
+    # Height and width are read as themselves: with a per-axis stride, only one
+    # of the two orders divides.
+    assert not ceil_mode_is_redundant(
+        torch.Size([11, 9]), (3, 3), (2, 4), (0, 0), (1, 1)
+    )
+    assert ceil_mode_is_redundant(torch.Size([9, 11]), (3, 3), (2, 4), (0, 0), (1, 1))
+
+    # Padding and dilation both move the span, which only an odd stride notices.
+    ten = torch.Size([10, 10])
+    assert not ceil_mode_is_redundant(ten, (3, 3), (3, 3), (0, 0), (1, 1))
+    assert ceil_mode_is_redundant(ten, (3, 3), (3, 3), (1, 1), (1, 1))
+    assert not ceil_mode_is_redundant(ten, (3, 3), (3, 3), (1, 1), (2, 2))
+
+
+class Identity(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + x
+
+
+def test_ceil_mode_redundancy_declines_a_symbolic_shape():
+    """A symbolic dimension divides symbolically rather than raising, so
+    without the guard the predicate would answer for whichever size happened to
+    be traced and bake that into the kernel. A max pool cannot supply one --
+    export rejects a dynamic spatial dim on it -- so borrow a shape.
+    """
+    exported = torch.export.export(
+        Identity().eval(),
+        (torch.randn(1, 1, 11, 11),),
+        dynamic_shapes=({2: torch.export.Dim("height", min=4, max=32)},),
+    )
+    (placeholder,) = [n for n in exported.graph.nodes if n.op == "placeholder"]
+    shape = placeholder.meta["val"].shape[-2:]
+    assert not all(isinstance(n, int) for n in shape), shape
+    assert not ceil_mode_is_redundant(shape, (3, 3), (2, 2), (0, 0), (1, 1))
+
+
 @parametrize("test_case", test_cases, xfails=xfails_max_pool2d)
 def test_dialect_max_pool2d(test_case, cortex_m_target):
     tester = CortexMTester(