[ExecuTorch][WebGPU] Port amax + amin (last-dim reductions)#21164
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21164
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ You can merge normally! (1 Unrelated Failure)As of commit f7e00b4 with merge base 266e0dc ( 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
|
psiddh
left a comment
There was a problem hiding this comment.
Approving full WebGPU stack
Stack from ghstack (oldest at bottom):
Ports the Vulkan delegate's
amax/aminas last-dim per-row reductions — the buffer-storage path Vulkan itself restricts reductions to (impl/Reduce.cppadd_reduce_per_row_node; its texture path throws on buffer storage). One thread per output row serially reduces over the contiguous last dim.Key changes:
runtime/ops/amax/{Reduce.cpp, amax.wgsl}+runtime/ops/amin/{Reduce.cpp, amin.wgsl}(+ generated_wgsl.h) —out[row] = reduce(input[row*W .. +W])whereW= last-dim extent; args[in, dim, keepdim, out]; fp32 int-guard; keepdim-aware resize hook; registeredaten.amax.default/aten.amin.default.CMakeLists.txt— both inWEBGPU_SRCS.Constraint: only a single last-dim reduction (
dim == -1orndim-1) is supported on buffer storage; any other reduce dim throws (mirrors Vulkan's buffer-per-row gate). The accumulator seedsinput[base](not 0), so all-negative rows reduce correctly.Differential Revision: D112257619