[ExecuTorch][WebGPU] Port pixel_shuffle#21177
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21177
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: ❌ 77 New FailuresAs of commit 7b92361 with merge base 266e0dc ( NEW FAILURES - The following jobs have failed:
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
aten.pixel_shuffle.defaultas a channel->space rearrange, a Phase B (vision) op. Pure data movement: each output element gathers from one input element.Key changes:
runtime/ops/pixel_shuffle/{PixelShuffle.cpp, pixel_shuffle.wgsl}(+ generated_wgsl.h) — for output(n, c_out, h_out, w_out):h_in = h_out / r,w_in = w_out / r,c_in = c_out*r*r + (h_out%r)*r + (w_out%r), thenoutput[out] = input[in]. Leading dims collapse into a batch, so any rank >= 3 works. Registeredaten.pixel_shuffle.default.CMakeLists.txt—runtime/ops/pixel_shuffle/PixelShuffle.cppinWEBGPU_SRCS.Mirrors Vulkan
glsl/pixel_shuffle_buffer.glsl(samec_in/h_in/w_inmap) +impl/PixelShuffle.cpp(arg order[in, upscale_factor, out],r >= 1,in_channels % (r*r) == 0,resize_pixel_shuffle_node). No int dtype-guard (bit-preserving movement);nbytes % 4guard. A dynamic-shape resize hook recomputes the out shape (inwith the last 3 dims rescaled by r) + params + dispatch + cur_dims. Fail-loud guards: rank >= 3, r >= 1,in_channels % (r*r) == 0(mirrors Vulkan VK_CHECK_COND), fp32.Differential Revision: D112257663