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(