Cortex-M: keep the max_pool2d reference off aten's channels-last int8 path

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

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