[ExecuTorch][WebGPU] Add max_pool2d op#20861
Open
JCNTH wants to merge 3 commits into
Open
Conversation
🔗 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. ❌ 57 New Failures, 3 Unrelated Failures, 7 Unclassified FailuresAs of commit 2704061 with merge base c2b273e ( NEW FAILURES - The following jobs have failed:
UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
This was referenced Jul 10, 2026
[ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline)
#20863
Open
This was referenced Jul 10, 2026
meta-cla
Bot
added
the
CLA Signed
This was referenced Jul 16, 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
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.
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.
@exported-using-ghexport
Differential Revision: D110836680
Differential Revision: D110836680