Skip to content

[ExecuTorch][WebGPU] Port native_group_norm#21173

Open
JCNTH wants to merge 1 commit into
gh/JCNTH/133/basefrom
gh/JCNTH/133/head
Open

[ExecuTorch][WebGPU] Port native_group_norm#21173
JCNTH wants to merge 1 commit into
gh/JCNTH/133/basefrom
gh/JCNTH/133/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Ports aten.native_group_norm.default as a 2-pass, 3-output group normalization, the second Phase B (vision) op. Pass 1 reduces per group to mean/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)*HxW elements), writing mean = 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-local add_gn_dispatch helper builds each of the two dispatches from an explicit binding list sharing one params UBO. Registered aten.native_group_norm.default.
  • CMakeLists.txtruntime/ops/native_group_norm/GroupNorm.cpp in WEBGPU_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 because execute() runs one compute pass per dispatch (WebGPUGraph.cpp; proven by test_dispatch_order.cpp to 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

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 22, 2026

Copy link
Copy Markdown

🔗 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 SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 77 New Failures

As of commit 55b3df2 with merge base 266e0dc (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

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"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@meta-cla meta-cla Bot 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 22, 2026
This was referenced Jul 22, 2026
This was referenced Jul 22, 2026
@JCNTH
JCNTH requested a review from psiddh July 23, 2026 02:08

@psiddh psiddh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving full WebGPU stack

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants