From be4425793c6fb02486da12395a468318b37c3f7f Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 12 Aug 2026 10:48:27 -0700 Subject: [PATCH 1/2] Cortex-M: keep the max_pool2d reference off aten's channels-last int8 path ### Summary aten's channels-last max_pool2d buffers each window index in vec::int_same_size_t and guards it with TORCH_CHECK(input_depth * input_height * input_width <= numeric_limits::max()) (aten/src/ATen/native/cpu/MaxPoolKernel.cpp, cpu_max_pool_channels_last). For int8 that caps an image at 127 spatial elements -- H*W, with channels not counted. Reproduce with: torch.nn.functional.max_pool2d( torch.zeros(1, 1, 12, 12, dtype=torch.int8).to( memory_format=torch.channels_last), 2, 2) The Cortex-M reference implementation called F.max_pool2d on the tensor as given, so a channels-last graph tripped that ceiling and raised before producing a result. That is a limit of the eager kernel and not of arm_max_pool_s8, so it only ever struck the host dialect stage -- but it struck it hard, since the whole model test aborts. Pooling does not depend on the memory format, so the reference now pools a contiguous copy; the existing return already puts the result back in channels-last, and .contiguous() aliases an already-contiguous tensor, so NCHW graphs pay nothing. quantized_max_pool2d_impl is the only reference in this file that pools in the native int8 dtype: avg_pool2d dequantizes to float first and both convolutions promote to int32, so none of them can reach the check. backends/cortex_m/test/models/test_yolo11.py is what hit this -- at 640x640 yolo11n's SPPF pools its 20x20 P5 map, which is 400 -- but that test importorskips ultralytics and so never runs in CI. The suite came closer than it looks: test_nn_modules.py already pools a channels-last (1, 4, 8, 8), which escapes only because 8*8 is 64. The new op-level case is the coverage. ### Test plan pytest backends/cortex_m/test/ops/test_max_pool2d.py -- 16 passed, 2 xfailed, dialect and Corstone-300; the new case adds ~2.4 s, in line with the existing implementation cases. Full backends/cortex_m/test dialect run: 404 passed. Reverting input.contiguous() fails the new case with the TORCH_CHECK above. Authored with assistance from Claude Code. --- backends/cortex_m/ops/operators.py | 15 ++++++- backends/cortex_m/test/ops/test_max_pool2d.py | 40 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 0c731911e44..126c7a5bbd9 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -1495,8 +1495,21 @@ def quantized_max_pool2d_impl( if ceil_mode: raise RuntimeError("quantized_max_pool2d does not support ceil_mode=True") + # aten's 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. Pooling is + # layout-invariant, so pool a contiguous copy; the return below puts the + # result back in channels-last either 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..8abe18c45bc 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,32 @@ 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. + + The pool then sees a live view whose strides satisfy plain contiguity + while aten still reports it as channels-last, which is the one shape a + .contiguous() call cannot normalize. + """ + + # 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 +96,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),), + ), } From 2e66b6b1f68d60ea5fa2c2d4ddfb1633013ba818 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 14 Aug 2026 10:17:20 -0700 Subject: [PATCH 2/2] Cortex-M: say what the max_pool2d layout comments are for Review feedback. Both comments explained a mechanism without first saying which decision it justifies, so they read as trivia until the second pass. The one in quantized_max_pool2d_impl now opens with the conclusion -- aten's channels-last kernel caps how large an image it accepts, so pool a contiguous copy -- and keeps the TORCH_CHECK derivation underneath as the supporting detail. The permuted-view test's docstring described why .contiguous() cannot normalize a single-channel channels-last tensor without saying that the reference therefore does not use it, which left a reader wondering whether the case was meant to fail and what the pass had to do about it. It now states what the case pins: reverting quantized_max_pool2d_impl to .contiguous() fails this one at the dialect stage, and only this one. Verified both ways. Authored with assistance from Claude Code. --- backends/cortex_m/ops/operators.py | 10 ++++++---- backends/cortex_m/test/ops/test_max_pool2d.py | 8 +++++--- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 126c7a5bbd9..44e47087c11 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -1495,14 +1495,16 @@ def quantized_max_pool2d_impl( if ceil_mode: raise RuntimeError("quantized_max_pool2d does not support ceil_mode=True") - # aten's cpu_max_pool_channels_last buffers each window index in + # 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. Pooling is - # layout-invariant, so pool a contiguous copy; the return below puts the - # result back in channels-last either way. + # 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() diff --git a/backends/cortex_m/test/ops/test_max_pool2d.py b/backends/cortex_m/test/ops/test_max_pool2d.py index 8abe18c45bc..d394747dfb5 100644 --- a/backends/cortex_m/test/ops/test_max_pool2d.py +++ b/backends/cortex_m/test/ops/test_max_pool2d.py @@ -39,9 +39,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class CortexMMaxPool2dPermutedView(torch.nn.Module): """A single-channel NHWC image permuted to NCHW. - The pool then sees a live view whose strides satisfy plain contiguity - while aten still reports it as channels-last, which is the one shape a - .contiguous() call cannot normalize. + 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