diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 0c731911e44..44e47087c11 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -1495,8 +1495,23 @@ def quantized_max_pool2d_impl( if ceil_mode: raise RuntimeError("quantized_max_pool2d does not support ceil_mode=True") + # aten's channels-last max_pool2d caps how large an image it will take, so + # pool a contiguous copy instead. Pooling is layout-invariant, and the + # return below puts the result back in channels-last either way. + # + # The cap: cpu_max_pool_channels_last buffers each window index in + # vec::int_same_size_t and guards it with + # TORCH_CHECK(input_depth * input_height * input_width <= max), so int8 + # rejects any image with more than 127 spatial elements -- H*W, with + # channels not counted. int16 hits the same wall at 32767, which a future + # quantized_max_pool2d_s16 will need to handle the same way. + # + # .to(memory_format=...) rather than .contiguous(): for C == 1 the + # channels-last strides also satisfy plain contiguity, so .contiguous() + # returns the same tensor while aten still dispatches on the memory-format + # hint and raises anyway. result = F.max_pool2d( - input, + input.to(memory_format=torch.contiguous_format), kernel, stride=stride_vals, padding=padding_vals, diff --git a/backends/cortex_m/test/ops/test_max_pool2d.py b/backends/cortex_m/test/ops/test_max_pool2d.py index a67dd6b6e01..d394747dfb5 100644 --- a/backends/cortex_m/test/ops/test_max_pool2d.py +++ b/backends/cortex_m/test/ops/test_max_pool2d.py @@ -1,3 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. # Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the @@ -34,6 +36,34 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.pool(x) +class CortexMMaxPool2dPermutedView(torch.nn.Module): + """A single-channel NHWC image permuted to NCHW. + + With C == 1 the channels-last strides also satisfy plain contiguity, so + .contiguous() hands the reference back the very tensor it was given while + aten keeps dispatching on the memory-format hint. This case is what pins + the .to(memory_format=...) in quantized_max_pool2d_impl: it fails at the + dialect stage if that is written as .contiguous(). + """ + + # The permute brings its own quant/dequant pair, so only the pool itself + # is pinned here. + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1, + } + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_max_pool2d_default": 1, + "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 0, + } + + def __init__(self, *args, **kwargs): + super().__init__() + self.pool = torch.nn.MaxPool2d(*args, **kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.pool(x.permute(0, 3, 1, 2)) + + class CortexMMaxPool2dIndices(torch.nn.Module): ops_before_transforms = CortexMMaxPool2d.ops_before_transforms ops_after_transforms = CortexMMaxPool2d.ops_after_transforms @@ -68,6 +98,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: CortexMMaxPool2dIndices(kernel_size=2, stride=2), (ramp_tensor(-50, 50, (1, 1, 6, 6)),), ), + # 576 spatial elements (24x24), past the 127 that aten's channels-last int8 + # max_pool2d accepts; the reference pools a contiguous copy to avoid it. + # randn rather than a ramp: a ramp over this many elements puts a whole + # pooling window inside one int8 code, so it cannot tell max from min. + "maxpool_2x2_large_channels_last": McuTestCase( + CortexMMaxPool2d(kernel_size=2, stride=2), + ((torch.randn(1, 64, 24, 24) * 30).to(memory_format=torch.channels_last),), + ), + "maxpool_2x2_single_channel_view": McuTestCase( + CortexMMaxPool2dPermutedView(kernel_size=2, stride=2), + ((torch.randn(1, 24, 24, 1) * 30),), + ), }