From 1361ee4ae57f0e3d848676528ca5420ad11fa16e Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 25 Aug 2026 14:36:15 -0700 Subject: [PATCH] Cortex-M: lower a max pool whose ceil_mode changes nothing The kernel floors, so CortexMMaxPool2DCheck refused every ceil_mode pool. Most do not need the ceiling: it adds a row or column only when the last window would hang off the end of the padded input, and all three of SqueezeNet's pools divide exactly. The checker and the lowering now compare the two sizes and accept when they agree, taking SqueezeNet from 29 lowered nodes to 32 with nothing but the concatenations and two layout no-ops left. Pools whose ceiling genuinely differs still decline. The test is one-sided on purpose. Aten also drops a last window that starts beyond the input, which can bring the ceiling back down to the floor; those configurations keep falling back rather than earn a second rule. Refusing was worse than staying in fp32. The checker annotates the node before tagging it, so the surrounding dq/q fold away and the fallback runs portable max_pool2d_with_indices over an int8 channels_last tensor -- whose index type is templated on the input dtype, so aten refuses a plane above 127 elements. Pools that lower no longer reach it. Both gates now read the pool's arguments through one helper. The lowering used to index node.args by hand and compare the results against tuple literals, which a list never matches, so no pool that spells out its dilation could lower -- and naming ceil_mode is what puts dilation in the graph. Nor did it handle stride's schema default of [], meaning "same as the kernel". Authored with Claude Code. --- .../cortex_m/passes/aten_to_cortex_m_pass.py | 50 +++++------- backends/cortex_m/passes/passes_utils.py | 49 ++++++++++- .../cortex_m/quantizer/pattern_checkers.py | 13 ++- backends/cortex_m/test/ops/test_max_pool2d.py | 81 +++++++++++++++++++ 4 files changed, 157 insertions(+), 36 deletions(-) diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index af60df686f5..168b73d5198 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/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 diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py index fe18cc1b141..adef032ed8f 100644 --- a/backends/cortex_m/passes/passes_utils.py +++ b/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: diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index ef2a91e6c2c..71f22113eab 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/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", {}) diff --git a/backends/cortex_m/test/ops/test_max_pool2d.py b/backends/cortex_m/test/ops/test_max_pool2d.py index d394747dfb5..bed4bcb20dc 100644 --- a/backends/cortex_m/test/ops/test_max_pool2d.py +++ b/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(