Skip to content

[ExecuTorch][WebGPU] Add max_pool2d op#20861

Open
JCNTH wants to merge 1 commit into
gh/JCNTH/36/basefrom
gh/JCNTH/36/head
Open

[ExecuTorch][WebGPU] Add max_pool2d op#20861
JCNTH wants to merge 1 commit into
gh/JCNTH/36/basefrom
gh/JCNTH/36/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 10, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Adds aten.max_pool2d_with_indices.default to the WebGPU backend, enabling the SAM2 Hiera q_pool path (values, with optional indices).

Problem — max pooling (F.max_pool2d, which decomposes to max_pool2d_with_indices) had no WebGPU kernel, blocking the SAM2 Hiera q_pool downsampling stages.

Solution — Before — no handler; the multi-output max_pool2d_with_indices unsupported. After — one thread per output element (n, c, oh, ow) gathers its pooling window and writes the max (and, when the graph requests them, the argmax indices).

Implementation:

  • The kernel iterates the kH x kW window with general stride/padding/dilation, skips out-of-range (padding) cells, and tracks the running max; best initialises to -3.4e38 (a large finite negative, because Dawn/Tint rejects the exact -FLT_MAX literal).
  • Argmax is ALWAYS tracked; a write_indices override (a compile-time spec constant) gates only the final out_idx store, so the values-only and with-indices paths run one shader and stay bit-identical on out_vals. Indices are the flat ih*IW+iw spatial-plane offset (matching torch.nn.functional.max_pool2d(return_indices=True), not an absolute NCHW offset).
  • The out is a ValueList [values, indices]; when indices are not requested the handler binds a tiny dummy storage buffer to slot 3 (the shader never writes it) rather than allocating a real indices tensor — mirroring the dummy_affine pattern in NativeLayerNorm.cpp. When indices ARE requested it validates that the indices tensor is int32 (4 bytes/elem) and shares the values shape, else throws (the kernel writes i32, so an int64 indices tensor would mis-stride the write).
  • The handler parses kernel_size/stride/padding/dilation via utils::parse_hw (a single value broadcasts to both spatial dims; stride defaults to kernel_size when empty), derives OH/OW from the pooling formula, and validates them against the serialized values output (loud-fail on a ceil_mode / layout mismatch).
  • Adaptive 1D->2D dispatch via utils::compute_dispatch_grid.
  • Mirrors Vulkan backends/vulkan/runtime/graph/ops/impl/Pool.cpp (the ValueList[values, indices] shape, the write_indices spec constant, unconditional argmax tracking, and 32-bit indices despite ATen's int64 schema).

Constraints — fp32 values; 4D NCHW; general kernel/stride/padding/dilation; ceil_mode unsupported (the output-shape validation assumes floor); the optional indices output must be int32; the output element count must fit u32.

Co-authored-with: Claude Code.

Differential Revision: D110836680

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 10, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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 4f64523 with merge base aceeb40 (image):

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.

@github-actions

Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants