Cortex-M: keep the max_pool2d reference off aten's channels-last int8 path - #21820
Open
rascani wants to merge 2 commits into
Open
Cortex-M: keep the max_pool2d reference off aten's channels-last int8 path#21820rascani wants to merge 2 commits into
rascani wants to merge 2 commits into
Conversation
… path
### Summary
aten's channels-last max_pool2d buffers each window index in
vec::int_same_size_t<opmath_t> and guards it with
TORCH_CHECK(input_depth * input_height * input_width <= numeric_limits<integer_t>::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.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21820
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Pending, 1 Unrelated FailureAs of commit 2e66b6b with merge base a56af1c ( NEW FAILURE - The following job has failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
meta-cla
Bot
added
the
CLA Signed
rascani
marked this pull request as ready for review
August 13, 2026 18:33
Erik-Lundell
left a comment
Erik-Lundell
left a comment
Collaborator
There was a problem hiding this comment.
Minor comments on wording, but looks solid.
I love this wording :)
it only ever struck the host dialect stage -- but it struck it hard
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.
Contributor
Author
Claude never misses an opportunity to add some drama. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
aten's channels-last max_pool2d buffers each window index in vec::int_same_size_t<opmath_t> and guards it with
TORCH_CHECK(input_depth * input_height * input_width <= numeric_limits<integer_t>::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.
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.