Legalize large-stride non-overlapping maxpool on U55 (#21284)

[apullin]

Summary:

Teach the Arm partitioner to recognize exact U55 large-stride non-overlapping max-pool operations as safely decomposable, without mutating the graph during partitioning. Backend preprocessing then canonicalizes value-only max_pool2d_with_indices and rewrites eligible pools into supported stride-1 reductions. This preserves the partitioner contract, keeps unrelated TOSA/VGF/U85 paths unchanged, and also covers exported MaxPool1d.

Differential Revision: D113087299

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21284

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 0184392 with merge base e798c4d ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@apullin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113087299.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[rascani]

@claude review this PR

[digantdesai]

@claude review this pr, since the previous attempt didn't work, let's try again

[claude]

Claude finished @digantdesai's task in 12m 8s —— View job

---

PR Review: Legalize large-stride non-overlapping maxpool on U55 (#21284)

• Read changed files and trace partition → decompose → rewrite flow
• Verify decomposition math and the positive-check / negative-check interaction
• Review tests
• Post review

---

Overall this is a clean approach: a shape-only predicate (can_decompose_large_stride_maxpool2d) shared between the partitioner's positive check and the decomposition pass, with no graph mutation during partitioning. The tiling math is correct and the tests exercise numerics via run_method_and_compare_outputs. A few items below.

Blocking / should-fix

1. Copyright header (already flagged by @rascani). Both new files use Copyright 2026 Arm Limited and/or its affiliates. but should use the Meta header, matching e.g. backends/transforms/remove_getitem_op.py:

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

Applies to backends/arm/_passes/decompose_large_stride_maxpool2d_pass.py:1 and backends/arm/test/passes/test_decompose_large_stride_maxpool2d_pass.py:1. Fix this →

Correctness / robustness

2. U55 dimension limits are not enforced for these nodes. The standard MaxPool2dSupported positive check enforces dim_check (all non-batch dims ≤ 65536) and kernel_check (kernel_x·kernel_y ≤ 65536). Because positive_checks are combined with any_chain, a large-stride node that fails MaxPool2dSupported (stride > 3) is accepted purely via the new DecomposableLargeStrideMaxPool2dSupported check — which validates kernel ∈ [1,256] but not the 65536 dim/product limits. The decomposition then produces an intermediate reduction over a tensor of shape [n, c, output_h*output_w*kernel_h, kernel_w] (decompose_large_stride_maxpool2d_pass.py:119-124). Since decomposition runs after partitioning, nothing re-validates this against Vela's limits, so an oversized output_h*output_w*kernel_h would slip past the partitioner and fail later at Vela compile time rather than being cleanly rejected/kept on CPU.

Suggest having can_decompose_large_stride_maxpool2d also check that the original spatial dims and the intermediate output_h*output_w*kernel_h stay within U55 limits (reusing dim_check/kernel_check from pool_2d_support.py). If in-scope models are always small enough that this can't happen, a comment saying so would help.

3. Partitioner accepts max_pool2d_with_indices but the pass only rewrites max_pool2d.default. DecomposableLargeStrideMaxPool2dSupported accepts both max_pool2d.default and max_pool2d_with_indices.default (partitioner.py:72-76), but DecomposeLargeStrideMaxPool2dPass.target_ops is only max_pool2d.default. This works today only because RemoveGetItemPass runs earlier in arm_pass_manager and converts _with_indices (single user) → max_pool2d.default. If a _with_indices node ever reaches the pass with its indices actually used (2 users), RemoveGetItemPass raises AssertionError — after the partitioner already committed to delegating it. This mirrors the existing MaxPool2dSupported behavior so it isn't a regression, but a one-line comment noting the dependency on RemoveGetItemPass ordering would make the contract explicit.

Minor

4. Redundant SymInt / kernel re-computation. can_decompose_large_stride_maxpool2d already returns False when height/width are SymInt, so the re-check at decompose_large_stride_maxpool2d_pass.py:78-80 is dead code. Likewise kernel_h, kernel_w is recomputed at line 77 having already been unpacked inside the predicate. Not harmful, just removable.

5. qparams on intermediate pools. The intermediate max_pool2d ops use no_qparams_meta (lines 125-142) whereas the existing DecomposeMaxPool2dPass keeps full meta on its core pool. Since max-pool is quantization-invariant (no rescale) this is numerically fine, but the divergence from the sibling pass is worth a sanity check for the INT/U55 quantized path — the final view_copy returning the original meta (with output_qparams) is the piece that carries quantization, which looks intentional.

The decomposition tiling itself — crop → reshape to [n,c,oh,kh,ow,kw] → permute → reduce over kw then kh — is correct, and the four positive tests plus two negative tests (overlapping stride, supported scalar stride) give good coverage.
· branch export-D113087299

[apullin]

OK, I believe lint is finally satisfied, too.