[ExecuTorch][WebGPU] Port native_group_norm#21173
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21173
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 55b3df2 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.native_group_norm.defaultas a 2-pass, 3-output group normalization, the second Phase B (vision) op. Pass 1 reduces per group tomean/rstd; pass 2 applies the per-channel affine.Key changes:
runtime/ops/native_group_norm/{GroupNorm.cpp, group_norm_reduce.wgsl, group_norm.wgsl}(+ generated_wgsl.h) — reduce: one thread per(n, group)serially scans that group's contiguous NCHW block (base = n*C*HxW + g*(C/G)*HxW,group_size = (C/G)*HxWelements), writingmean = s/count,rstd = inverseSqrt(ss/count - mean^2 + eps); normalize: one thread per element,g = c / (C/G),out = (x - mean[n*G+g]) * rstd[n*G+g] * weight[c] + bias[c]. A file-localadd_gn_dispatchhelper builds each of the two dispatches from an explicit binding list sharing one params UBO. Registeredaten.native_group_norm.default.CMakeLists.txt—runtime/ops/native_group_norm/GroupNorm.cppinWEBGPU_SRCS.Mirrors Vulkan
impl/GroupNorm.cpp(2 nodes: reduce -> mean/rstd, then normalize) +glsl/group_norm_reduce_texture.glsl(mean=sum/count,variance=sumsq/count - mean^2,rstd=1/sqrt(var+eps)). Buffer NCHW re-derivation (Vulkan is texture/channels-packed), so a serial per-group reduce (a group's channels x HxW are contiguous in NCHW) replaces the cooperative texture reduction. All 3 outputs are real. The inter-pass RAW ordering (normalize reads the reduce's mean/rstd) is guaranteed becauseexecute()runs one compute pass per dispatch (WebGPUGraph.cpp; proven bytest_dispatch_order.cppto 1M elems). A dynamic-shape resize hook recomputes params + both dispatch counts + propagates cur_dims for out/mean/rstd. Fail-loud guards: 4D,C % group == 0, fp32, weight/bias len== C, mean/rstd size== N*group.Differential Revision: D112257596