You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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.
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"
meta-claBot
added
the
CLA Signed
This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
label
Jul 10, 2026
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
CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
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.
Stack from ghstack (oldest at bottom):
Adds
aten.max_pool2d_with_indices.defaultto the WebGPU backend, enabling the SAM2 Hiera q_pool path (values, with optional indices).Problem — max pooling (
F.max_pool2d, which decomposes tomax_pool2d_with_indices) had no WebGPU kernel, blocking the SAM2 Hieraq_pooldownsampling stages.Solution — Before — no handler; the multi-output
max_pool2d_with_indicesunsupported. 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:
kH x kWwindow with generalstride/padding/dilation, skips out-of-range (padding) cells, and tracks the running max;bestinitialises to-3.4e38(a large finite negative, because Dawn/Tint rejects the exact-FLT_MAXliteral).write_indicesoverride (a compile-time spec constant) gates only the finalout_idxstore, so the values-only and with-indices paths run one shader and stay bit-identical onout_vals. Indices are the flatih*IW+iwspatial-plane offset (matchingtorch.nn.functional.max_pool2d(return_indices=True), not an absolute NCHW offset).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 thedummy_affinepattern inNativeLayerNorm.cpp. When indices ARE requested it validates that the indices tensor isint32(4 bytes/elem) and shares the values shape, else throws (the kernel writesi32, so anint64indices tensor would mis-stride the write).kernel_size/stride/padding/dilationviautils::parse_hw(a single value broadcasts to both spatial dims;stridedefaults tokernel_sizewhen empty), derivesOH/OWfrom the pooling formula, and validates them against the serialized values output (loud-fail on aceil_mode/ layout mismatch).utils::compute_dispatch_grid.backends/vulkan/runtime/graph/ops/impl/Pool.cpp(theValueList[values, indices]shape, thewrite_indicesspec constant, unconditional argmax tracking, and 32-bit indices despite ATen's int64 schema).Constraints — fp32 values; 4D NCHW; general
kernel/stride/padding/dilation;ceil_modeunsupported (the output-shape validation assumes floor); the optional indices output must beint32; the output element count must fitu32.Co-authored-with: Claude Code.
Differential Revision: D110836680