sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +89 -0
build/webgpu/bench.json +425 -0
build/webgpu/group-norm-combine.wgsl.jinja +50 -0
build/webgpu/group-norm-fused.wgsl.jinja +136 -0
build/webgpu/group-norm-slab.wgsl.jinja +233 -0
build/webgpu/manifest.json +226 -0
build/webgpu/metadata.json +26 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -1,3 +1,92 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.GroupNorm
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Group normalization over a 4-D image tensor with an optional fused SiLU: `Y = gamma * (X - mean) / sqrt(variance + epsilon) + beta`, with mean and variance taken per batch item over each group of `C / groups` channels and every spatial position. Both layouts are implemented: `channels_last = 1` reads `(N, H, W, C)` and `channels_last = 0` reads `(N, C, H, W)`. `X`/`Y` and `gamma`/`beta` carry independent element types, so float16 activations with float32 weights is a first-class case. Statistics and the affine are float32 whatever the tensor type; only the store narrows.
+See the [ONNX Runtime `GroupNorm` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.GroupNorm) for the reference semantics.
+## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `X` | `T` | `4` | — | Input image tensor, `(N, H, W, C)` when `channels_last` is 1 and `(N, C, H, W)` otherwise. | required |
+| `gamma` | `M` | `1` | — | Per-channel affine scale of shape `(C)`. | required |
+| `beta` | `M` | `1` | — | Per-channel affine offset of shape `(C)`. | required |
+## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `Y` | `T` | `4` | same as `X` | Normalized tensor with the same shape and element type as `X`. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `activation` | — | Activation applied after the affine transform: `0` for none, `1` for SiLU. The request must supply it; the operator declares no default. |
+| `channels_last` | `1` | `1` when `X` and `Y` are `(N, H, W, C)`, `0` when they are `(N, C, H, W)`. Defaults to `1`. Any non-zero value selects the channels-last layout. |
+| `epsilon` | `0.00001` | Value added to the group variance before the reciprocal square root. Defaults to `1e-5`. |
+| `groups` | — | Number of channel groups the statistics are taken over. It must divide `C`. The request must supply it; the operator declares no default. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `M` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `split` — Three dispatches: partial group statistics, combine mean and reciprocal standard deviation, then normalize and apply affine weights. Channels-last statistics use four, two, or one aligned values per lane; group blocks and spatial rows follow device limits.
+- `fused` — One workgroup per batch item and group, walking the group twice inside a single dispatch: once for the statistics and once to write the affine result. It avoids the split route's two extra launches on small tensors and is the route that needs no scratch buffer.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
+- [`group-norm-combine.wgsl.jinja`](build/webgpu/group-norm-combine.wgsl.jinja)
+- [`group-norm-fused.wgsl.jinja`](build/webgpu/group-norm-fused.wgsl.jinja)
+- [`group-norm-slab.wgsl.jinja`](build/webgpu/group-norm-slab.wgsl.jinja)
+## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.GroupNorm", { version: 1 });
+const { Y } = await kernel({
+  X: { data: XData, shape: [1, 2, 2, 16] },
+  gamma: { data: gammaData, shape: [16] },
+  beta: { data: betaData, shape: [16] },
+}, {
+  attrs: { groups: 4, activation: 0 },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,425 @@

+{
+  "tunableSpace": { "WORKGROUP_SIZE": [128, 256] },
+  "cases": [
+    {
+      "name": "flux-vae-mid-1x64x64x512-g32-f16",
+      "preset": "smoke",
+      "provenance": { "note": "Flux.2 VAE decoder mid/up-block group norm, NHWC, SiLU" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 2097152, "channels": 512 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [1, 64, 64, 512], "dist": "normal", "seed": 100, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [512], "dist": "uniform", "seed": 101, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [512], "dist": "normal", "seed": 102, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 64, 64, 512], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ],
+        "primary": true
+      }
+    },
+    {
+      "name": "flux-vae-mid-1x64x64x512-g32-f32",
+      "preset": "smoke",
+      "provenance": { "note": "Flux.2 VAE decoder mid/up-block group norm at float32" },
+      "vars": { "dtype": "float32", "weightDtype": "float32", "count": 2097152, "channels": 512 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "X": { "dtype": "float32", "shape": [1, 64, 64, 512], "dist": "normal", "seed": 110, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [512], "dist": "uniform", "seed": 111, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [512], "dist": "normal", "seed": 112, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 64, 64, 512], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sd15-unet-2x64x64x320-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Stable-Diffusion-1.5 UNet resnet at the classifier-free-guidance batch; 10 channels per group"
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 2621440, "channels": 320 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 64, 64, 320], "dist": "normal", "seed": 120, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [320], "dist": "uniform", "seed": 121, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [320], "dist": "normal", "seed": 122, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64, 320], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sd15-unet-2x32x32x640-g32-f16",
+      "preset": "smoke",
+      "provenance": { "note": "Stable-Diffusion-1.5 UNet mid-resolution resnet; 20 channels per group" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 1310720, "channels": 640 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 32, 32, 640], "dist": "normal", "seed": 130, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [640], "dist": "uniform", "seed": 131, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [640], "dist": "normal", "seed": 132, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 32, 32, 640], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sd15-unet-2x8x8x1280-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Stable-Diffusion-1.5 UNet lowest-resolution resnet; the few-element, low-parallelism corner"
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 163840, "channels": 1280 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 8, 8, 1280], "dist": "normal", "seed": 140, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [1280], "dist": "uniform", "seed": 141, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [1280], "dist": "normal", "seed": 142, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 8, 8, 1280], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "odd-cpg-1x64x64x96-g32-f32",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Three channels per group: the odd channels-per-group case the CUDA kernel refuses outright"
+      },
+      "vars": { "dtype": "float32", "weightDtype": "float32", "count": 393216, "channels": 96 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float32", "shape": [1, 64, 64, 96], "dist": "normal", "seed": 150, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [96], "dist": "uniform", "seed": 151, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [96], "dist": "normal", "seed": 152, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 64, 64, 96], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nchw-4x64x128x128-g32-f32",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Channels-first at the shape ai.onnx.GroupNormalization benches, for a direct comparison"
+      },
+      "vars": { "dtype": "float32", "weightDtype": "float32", "count": 4194304, "channels": 64 },
+      "attrs": { "groups": 32, "activation": 0, "channels_last": 0 },
+      "inputs": {
+        "X": { "dtype": "float32", "shape": [4, 64, 128, 128], "dist": "normal", "seed": 160, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [64], "dist": "uniform", "seed": 161, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [64], "dist": "normal", "seed": 162, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 64, 128, 128], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nchw-2x320x64x64-g32-f16",
+      "preset": "smoke",
+      "provenance": { "note": "Channels-first Stable-Diffusion-1.5 resnet width with the fused SiLU" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 2621440, "channels": 320 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 0 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 320, 64, 64], "dist": "normal", "seed": 170, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [320], "dist": "uniform", "seed": 171, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [320], "dist": "normal", "seed": 172, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 320, 64, 64], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "micro-2x2x2x128-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "ORT: onnxruntime/test/contrib_ops/group_norm_op_test.cc GroupNormTest.GroupNorm_128 geometry"
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 1024, "channels": 128 },
+      "attrs": { "groups": 32, "activation": 0, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 2, 2, 128], "dist": "normal", "seed": 180, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [128], "dist": "uniform", "seed": 181, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [128], "dist": "normal", "seed": 182, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 2, 2, 128], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "flux-vae-up-1x128x128x512-g32-f16",
+      "preset": "scale",
+      "provenance": { "note": "Flux.2 VAE decoder after the first 2x upsample" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 8388608, "channels": 512 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [1, 128, 128, 512], "dist": "normal", "seed": 200, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [512], "dist": "uniform", "seed": 201, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [512], "dist": "normal", "seed": 202, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 128, 128, 512], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "flux-vae-up-1x256x256x256-g32-f16",
+      "preset": "scale",
+      "provenance": { "note": "Flux.2 VAE decoder after the second 2x upsample" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 16777216, "channels": 256 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [1, 256, 256, 256], "dist": "normal", "seed": 210, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [256], "dist": "uniform", "seed": 211, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [256], "dist": "normal", "seed": 212, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 256, 256, 256], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "flux-vae-convout-1x512x512x128-g32-f16",
+      "preset": "model",
+      "provenance": { "note": "Flux.2 VAE conv_norm_out: the largest single group-norm dispatch in the decoder" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 33554432, "channels": 128 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [1, 512, 512, 128], "dist": "normal", "seed": 220, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [128], "dist": "uniform", "seed": 221, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [128], "dist": "normal", "seed": 222, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 512, 512, 128], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nchw-1x512x64x64-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Channels-first at the Flux.2 VAE mid-block width: 32 groups, the fewest reduction outputs this op sees"
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 2097152, "channels": 512 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 0 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [1, 512, 64, 64], "dist": "normal", "seed": 300, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [512], "dist": "uniform", "seed": 301, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [512], "dist": "normal", "seed": 302, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 512, 64, 64], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nchw-8x64x64x64-g32-f32",
+      "preset": "smoke",
+      "provenance": { "note": "Channels-first with the batch axis carrying the parallelism: 256 reduction outputs" },
+      "vars": { "dtype": "float32", "weightDtype": "float32", "count": 2097152, "channels": 64 },
+      "attrs": { "groups": 32, "activation": 0, "channels_last": 0 },
+      "inputs": {
+        "X": { "dtype": "float32", "shape": [8, 64, 64, 64], "dist": "normal", "seed": 310, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [64], "dist": "uniform", "seed": 311, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [64], "dist": "normal", "seed": 312, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 64, 64, 64], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nchw-2x320x128x128-g32-f16",
+      "preset": "scale",
+      "provenance": { "note": "Channels-first Stable-Diffusion resnet at a working set past cache" },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 10485760, "channels": 320 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 0 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 320, 128, 128], "dist": "normal", "seed": 320, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [320], "dist": "uniform", "seed": 321, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [320], "dist": "normal", "seed": 322, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 320, 128, 128], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nhwc-even-groups-2x64x64x320-g32-f32",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Even channel-group widths spanning one, multiple, and capacity-limited channel rows per workgroup, including float32 activations. Same distributions as the authored large NHWC case."
+      },
+      "vars": { "dtype": "float32", "weightDtype": "float32", "count": 2621440, "channels": 320 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float32", "shape": [2, 64, 64, 320], "dist": "normal", "seed": 120, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [320], "dist": "uniform", "seed": 121, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [320], "dist": "normal", "seed": 122, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 64, 64, 320], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nhwc-even-groups-2x64x64x64-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Even channel-group widths spanning one, multiple, and capacity-limited channel rows per workgroup, including float32 activations. Same distributions as the authored large NHWC case."
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 524288, "channels": 64 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 64, 64, 64], "dist": "normal", "seed": 120, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [64], "dist": "uniform", "seed": 121, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [64], "dist": "normal", "seed": 122, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64, 64], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nhwc-even-groups-2x64x64x192-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Even channel-group widths spanning one, multiple, and capacity-limited channel rows per workgroup, including float32 activations. Same distributions as the authored large NHWC case."
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 1572864, "channels": 192 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 64, 64, 192], "dist": "normal", "seed": 120, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [192], "dist": "uniform", "seed": 121, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [192], "dist": "normal", "seed": 122, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64, 192], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "nhwc-even-groups-2x64x64x448-g32-f16",
+      "preset": "smoke",
+      "provenance": {
+        "note": "Even channel-group widths spanning one, multiple, and capacity-limited channel rows per workgroup, including float32 activations. Same distributions as the authored large NHWC case."
+      },
+      "vars": { "dtype": "float16", "weightDtype": "float32", "count": 3670016, "channels": 448 },
+      "attrs": { "groups": 32, "activation": 1, "channels_last": 1 },
+      "inputs": {
+        "X": { "dtype": "float16", "shape": [2, 64, 64, 448], "dist": "normal", "seed": 120, "scale": 0.5 },
+        "gamma": { "dtype": "float32", "shape": [448], "dist": "uniform", "seed": 121, "scale": 0.25, "offset": 1 },
+        "beta": { "dtype": "float32", "shape": [448], "dist": "normal", "seed": 122, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64, 448], "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.count * 2 * dtypeBytes(args.dtype) + args.channels * 2 * dtypeBytes(args.weightDtype)"
+          }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/group-norm-combine.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,50 @@

+{% set comp0 = ".x" if vecSpec > 1 else "" %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// The partials hold sum(value - shift) and sum((value - shift)^2) over disjoint blocks
+// of the group, all taken against the same `shift` — the group's first element — so
+// they add directly. Recovering the mean adds the shift back; the variance is formed
+// from the shifted second moment and clamped at zero, which is where a group whose
+// values are all equal would otherwise land a few ulp below it.
+const GROUPS: u32 = {{ groupsSpec }}u;
+const CPG: u32 = {{ cpgSpec }}u;
+const ROWS: u32 = {{ normRowsSpec }}u;
+const PARTIALS_PER_ROW: u32 = {{ partialsPerRowSpec }}u;
+const GROUP_SIZE: f32 = {{ groupSizeSpec }}.0;
+const EPSILON: f32 = {{ epsilonSpec }};
+const WG: u32 = {{ combineWgSpec }}u;
+{% if phaseLayout %}
+const HW: u32 = {{ dimHWSpec }}u;
+const CHUNKS: u32 = {{ chunksSpec }}u;
+const VEC: u32 = {{ vecSpec }}u;
+{% else %}
+const CHANNELS: u32 = {{ dimCSpec }}u;
+const PLANE_V: u32 = {{ planeVecsSpec }}u;
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "row", "ROWS") }}
+  let n = row / GROUPS;
+  let g = row % GROUPS;
+{% if phaseLayout %}
+  let shift = f32(src[n * HW * CHUNKS + g * (CPG / VEC)]{{ comp0 }});
+{% else %}
+  let shift = f32(src[(n * CHANNELS + g * CPG) * PLANE_V]{{ comp0 }});
+{% endif %}
+  var total = vec2<f32>(0.0);
+  for (var k = 0u; k < PARTIALS_PER_ROW; k = k + 1u) {
+    total = total + partials[row * PARTIALS_PER_ROW + k];
+  }
+  let mean_d = total.x / GROUP_SIZE;
+  let variance = max(total.y / GROUP_SIZE - mean_d * mean_d, 0.0);
+  stats[row] = vec2<f32>(shift + mean_d, inverseSqrt(variance + EPSILON));
+}

build/webgpu/group-norm-fused.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,136 @@

+{% set comp0 = ".x" if useVec4 else "" %}
+{% macro pre_norm(idx, sidx) %}
+{% set x = vecF32 ~ "(X[" ~ idx ~ "])" %}
+{% set skipTerm = "skip_v" if broadcastSkip else vecF32 ~ "(skip[" ~ sidx ~ "])" %}
+{{ vecF32 ~ "(" ~ storeElement ~ "(" ~ x ~ " + " ~ skipTerm ~ " + bias_v))" if hasSkip else x }}{% endmacro %}
+{% macro shift_at(idx, sidx, ch) %}
+{% set x = "f32(X[" ~ idx ~ "]" ~ comp0 ~ ")" %}
+{% set skipTerm = "f32(skip[" ~ sidx ~ "]" ~ ("" if broadcastSkip else comp0) ~ ")" %}
+{% set biasTerm = " + f32(bias[" ~ ch ~ "])" if hasBias else "" %}
+{{ "f32(" ~ tScalar ~ "(" ~ x ~ " + " ~ skipTerm ~ biasTerm ~ "))" if hasSkip else x }}{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// Group normalization in one dispatch. The workgroup accumulates the group's shifted
+// moments, folds them through workgroup memory, then walks the same elements again to
+// write `gamma * (value - mean) * inv_std + beta`. Every intermediate is f32 whatever
+// the tensor type; only the store narrows. The reduction is a plain halving tree over
+// a power-of-two-padded scratch array, so it is deterministic and needs no subgroup
+// support.
+const CPG: u32 = {{ cpgSpec }}u;
+const GROUPS: u32 = {{ groupsSpec }}u;
+const ROWS: u32 = {{ normRowsSpec }}u;
+const GROUP_SIZE: f32 = {{ groupSizeSpec }}.0;
+const EPSILON: f32 = {{ epsilonSpec }};
+const WG: u32 = {{ wgSpec }}u;
+const LANES: u32 = {{ ldsSpec }}u;
+{% if phaseLayout %}
+const HW: u32 = {{ dimHWSpec }}u;
+const VEC: u32 = {{ vecSpec }}u;
+const CHUNKS: u32 = {{ chunksSpec }}u;
+const SEGMENT: u32 = {{ fusedChunksSpec }}u;
+const SUB_ROWS: u32 = {{ fusedRowsSpec }}u;
+{% else %}
+const CHANNELS: u32 = {{ dimCSpec }}u;
+const PLANE_V: u32 = {{ planeVecsSpec }}u;
+const SPAN_V: u32 = {{ spanVecsSpec }}u;
+{% endif %}
+// LANES is WG rounded up to a power of two; the lanes past WG are zeroed so the
+// halving tree below stays uniform.
+var<workgroup> lane_sums: array<vec2<f32>, LANES>;
+{% if activationSpec == 1 %}
+fn silu_e(v: {{ vecF32 }}) -> {{ vecF32 }} {
+  return v / ({{ vecF32 }}(1.0) + exp(-v));
+}
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-folded row index: wg.y carries the high bits past the dispatch's per-axis
+  // workgroup fold width.
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  if (row >= ROWS) {
+    return;
+  }
+  let n = row / GROUPS;
+  let g = row - n * GROUPS;
+{% if phaseLayout %}
+  let chunk = lid.x % SEGMENT;
+  let sub_row = lid.x / SEGMENT;
+  let c0 = g * CPG + chunk * VEC;
+  let base_v = n * HW * CHUNKS + g * SEGMENT;
+  let shift_index = base_v;
+{% else %}
+  let base_v = (n * CHANNELS + g * CPG) * PLANE_V;
+  let shift_index = base_v;
+{% endif %}
+  let shift = {{ shift_at("shift_index", "shift_skip", "g * CPG") }};
+  var acc = vec2<f32>(0.0);
+{% if phaseLayout %}
+  for (var hw = sub_row; hw < HW; hw = hw + SUB_ROWS) {
+    let idx = base_v + hw * CHUNKS + chunk;
+    let v = {{ pre_norm("idx", "skip_base + hw * CHUNKS + chunk") }};
+{% else %}
+  for (var i = lid.x; i < SPAN_V; i = i + WG) {
+    let idx = base_v + i;
+    let v = {{ pre_norm("idx", "skip_base + i") }};
+{% endif %}
+    let d = v - {{ vecF32 }}(shift);
+{% if useVec4 %}
+    acc = acc + vec2<f32>(d.x + d.y + d.z + d.w, dot(d, d));
+{% else %}
+    acc = acc + vec2<f32>(d, d * d);
+{% endif %}
+  }
+  lane_sums[lid.x] = acc;
+  if (lid.x + WG < LANES) {
+    lane_sums[lid.x + WG] = vec2<f32>(0.0);
+  }
+  workgroupBarrier();
+  var stride = LANES / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lid.x < stride) {
+      lane_sums[lid.x] = lane_sums[lid.x] + lane_sums[lid.x + stride];
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let total = lane_sums[0];
+  let mean_d = total.x / GROUP_SIZE;
+  let variance = max(total.y / GROUP_SIZE - mean_d * mean_d, 0.0);
+  let mean = {{ vecF32 }}(shift + mean_d);
+  let inv_std = {{ vecF32 }}(inverseSqrt(variance + EPSILON));
+{% if phaseLayout %}
+{% if useVec4 %}
+  let gamma_v = vec4<f32>(f32(gamma[c0]), f32(gamma[c0 + 1u]), f32(gamma[c0 + 2u]), f32(gamma[c0 + 3u]));
+  let beta_v = vec4<f32>(f32(beta[c0]), f32(beta[c0 + 1u]), f32(beta[c0 + 2u]), f32(beta[c0 + 3u]));
+{% else %}
+  let gamma_v = f32(gamma[c0]);
+  let beta_v = f32(beta[c0]);
+{% endif %}
+  for (var hw = sub_row; hw < HW; hw = hw + SUB_ROWS) {
+    let idx = base_v + hw * CHUNKS + chunk;
+{% else %}
+  for (var i = lid.x; i < SPAN_V; i = i + WG) {
+    let idx = base_v + i;
+    let channel = g * CPG + i / PLANE_V;
+{% if useVec4 %}
+    let gamma_v = vec4<f32>(f32(gamma[channel]));
+    let beta_v = vec4<f32>(f32(beta[channel]));
+{% else %}
+    let gamma_v = f32(gamma[channel]);
+    let beta_v = f32(beta[channel]);
+{% endif %}
+{% endif %}
+    let normalized = ({{ vecF32 }}({{ fusedSrc }}[idx]) - mean) * inv_std;
+{% if activationSpec == 1 %}
+    Y[idx] = {{ storeElement }}(silu_e(gamma_v * normalized + beta_v));
+{% else %}
+    Y[idx] = {{ storeElement }}(gamma_v * normalized + beta_v);
+{% endif %}
+  }
+}

build/webgpu/group-norm-slab.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,233 @@

+{% set blockedStats = phase == "stats" and phaseLayout and (statsBlockedSpec | default(false)) %}
+{% set useVector = vecSpec > 1 %}
+{% set comp0 = ".x" if useVector else "" %}
+{% macro channel_values(buffer, channel) %}
+{{ vecF32 }}({% for i in range(vecSpec) %}f32({{ buffer }}[{{ channel }}{% if i > 0 %} + {{ i }}u{% endif %}]){% if not loop.last %}, {% endif %}{% endfor %}){% endmacro %}
+{% macro component_sum(value) %}
+{% for i in range(vecSpec) %}{{ value }}.{{ "xyzw"[i] }}{% if not loop.last %} + {% endif %}{% endfor %}{% endmacro %}
+{% macro pre_norm(idx, sidx) %}
+{% set x = vecF32 ~ "(X[" ~ idx ~ "])" %}
+{% set skipTerm = "skip_v" if broadcastSkip else vecF32 ~ "(skip[" ~ sidx ~ "])" %}
+{{ vecF32 ~ "(" ~ storeElement ~ "(" ~ x ~ " + " ~ skipTerm ~ " + bias_v))" if hasSkip else x }}{% endmacro %}
+{% macro shift_at(idx, sidx, ch) %}
+{% set x = "f32(X[" ~ idx ~ "]" ~ comp0 ~ ")" %}
+{% set skipTerm = "f32(skip[" ~ sidx ~ "]" ~ ("" if broadcastSkip else comp0) ~ ")" %}
+{% set biasTerm = " + f32(bias[" ~ ch ~ "])" if hasBias else "" %}
+{{ "f32(" ~ tScalar ~ "(" ~ x ~ " + " ~ skipTerm ~ biasTerm ~ "))" if hasSkip else x }}{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// One phase of group normalization over a 4-D image tensor.
+//
+// A group is every spatial position of `CPG` adjacent channels of one batch item. The
+// statistics phase accumulates the sum and the sum of squares of `value - shift`,
+// where `shift` is the group's first element: subtracting one member of the group
+// removes the common offset before the second moment is formed, so the variance the
+// combine pass derives does not cancel against a large mean. The apply phase reads the
+// combined mean and reciprocal standard deviation back and writes
+// `gamma * (value - mean) * inv_std + beta`, optionally through SiLU.
+//
+// Every accumulation and the whole affine are f32 whatever the tensor type; only the
+// store narrows.
+const CPG: u32 = {{ cpgSpec }}u;
+const GROUPS: u32 = {{ groupsSpec }}u;
+const WG: u32 = {{ wgSpec }}u;
+const UNITS: u32 = {{ unitsSpec }}u;
+{% if phaseLayout %}
+const HW: u32 = {{ dimHWSpec }}u;
+const VEC: u32 = {{ vecSpec }}u;
+const CHUNKS: u32 = {{ chunksSpec }}u;
+const ROWS_PER_STEP: u32 = {{ rowsPerStepSpec }}u;
+const ROWS_PER_SLAB: u32 = {{ rowsPerSlabSpec }}u;
+const SLABS: u32 = {{ splitSpec }}u;
+{% if blockedStats %}
+const GROUP_BLOCKS: u32 = {{ statsGroupBlocksSpec }}u;
+const GROUPS_PER_BLOCK: u32 = {{ statsGroupsPerBlockSpec }}u;
+const TILE_CHUNKS: u32 = {{ statsTileChunksSpec }}u;
+{% endif %}
+{% else %}
+const CHANNELS: u32 = {{ dimCSpec }}u;
+const PLANE_V: u32 = {{ planeVecsSpec }}u;
+const CHUNK_V: u32 = {{ chunkVecsSpec }}u;
+const SPLITS: u32 = {{ splitSpec }}u;
+{% endif %}
+{% if phase == "stats" %}
+{% if not phaseLayout %}
+const PARTIALS_PER_ROW: u32 = {{ partialsPerRowSpec }}u;
+{% endif %}
+// Per-lane moments, combined below: the channels-last arm folds one contiguous
+// segment per group out of this array, the channels-first arm folds all of it.
+var<workgroup> lane_sums: array<vec2<f32>, WG>;
+{% endif %}
+{% if phase == "apply" and activationSpec == 1 %}
+fn silu_e(v: {{ vecF32 }}) -> {{ vecF32 }} {
+  return v / ({{ vecF32 }}(1.0) + exp(-v));
+}
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-folded work-unit index: wg.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let unit = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  if (unit >= UNITS) {
+    return;
+  }
+{% if phaseLayout %}
+  // Channels last. The workgroup owns one batch item and one slab of spatial rows and
+  // covers every channel of those rows: lane `lid.x` keeps channel chunk
+  // `lid.x % CHUNKS` and steps through the slab `ROWS_PER_STEP` rows at a time, so the
+  // workgroup reads whole contiguous channel rows and each lane's group is fixed.
+{% if blockedStats %}
+  let spatial_unit = unit / GROUP_BLOCKS;
+  let group_begin = (unit % GROUP_BLOCKS) * GROUPS_PER_BLOCK;
+  let n = spatial_unit / SLABS;
+  let slab = spatial_unit - n * SLABS;
+  let chunk = group_begin * (CPG / VEC) + lid.x % TILE_CHUNKS;
+  let sub_row = lid.x / TILE_CHUNKS;
+{% else %}
+  let n = unit / SLABS;
+  let slab = unit - n * SLABS;
+  let chunk = lid.x % CHUNKS;
+  let sub_row = lid.x / CHUNKS;
+{% endif %}
+  let c0 = chunk * VEC;
+{% if phase == "stats" or groupUniform %}
+  let g = c0 / CPG;
+{% endif %}
+  let image_base = n * HW * CHUNKS;
+  let hw_begin = slab * ROWS_PER_SLAB;
+  let hw_end = min(hw_begin + ROWS_PER_SLAB, HW);
+{% else %}
+  // Channels first. The workgroup owns one contiguous chunk of one (batch item,
+  // channel) plane, so every lane shares the channel and the group and the loads are
+  // contiguous across the workgroup. Adjacent units are adjacent chunks of one plane.
+  let plane = unit / SPLITS;
+  let part = unit - plane * SPLITS;
+  let n = plane / CHANNELS;
+  let c = plane % CHANNELS;
+  let g = c / CPG;
+  let base_v = plane * PLANE_V;
+  let chunk_begin = part * CHUNK_V;
+  let chunk_end = min(chunk_begin + CHUNK_V, PLANE_V);
+{% endif %}
+{% if phase == "stats" %}
+{% if blockedStats %}
+  var acc = vec2<f32>(0.0);
+  // Keep tail lanes participating in the barrier without reading outside the group.
+  if (g < GROUPS) {
+{% endif %}
+{% if phaseLayout %}
+  let shift_index = image_base + g * (CPG / VEC);
+{% else %}
+  let shift_index = (n * CHANNELS + g * CPG) * PLANE_V;
+{% endif %}
+  let shift = {{ shift_at("shift_index", "shift_skip", "g * CPG") }};
+{% if not blockedStats %}
+  var acc = vec2<f32>(0.0);
+{% endif %}
+{% if phaseLayout %}
+  for (var hw = hw_begin + sub_row; hw < hw_end; hw = hw + ROWS_PER_STEP) {
+    let idx = image_base + hw * CHUNKS + chunk;
+    let v = {{ pre_norm("idx", "skip_base + hw * CHUNKS + chunk") }};
+{% else %}
+  for (var i = chunk_begin + lid.x; i < chunk_end; i = i + WG) {
+    let idx = base_v + i;
+    let v = {{ pre_norm("idx", "skip_base + i") }};
+{% endif %}
+    let d = v - {{ vecF32 }}(shift);
+{% if useVector %}
+    acc = acc + vec2<f32>({{ component_sum("d") }}, dot(d, d));
+{% else %}
+    acc = acc + vec2<f32>(d, d * d);
+{% endif %}
+  }
+{% if blockedStats %}
+  }
+{% endif %}
+  lane_sums[lid.x] = acc;
+  workgroupBarrier();
+{% if phaseLayout %}
+  // One aligned segment of `CPG / VEC` lanes per group per sub-row. Each group's
+  // partial is summed in a fixed order by a single lane, so the value does not depend
+  // on how the workgroup was scheduled.
+  let segment = CPG / VEC;
+  for (var sel = lid.x; sel < {% if blockedStats %}min(GROUPS_PER_BLOCK, GROUPS - group_begin){% else %}GROUPS{% endif %}; sel = sel + WG) {
+    var total = vec2<f32>(0.0);
+    for (var r = 0u; r < ROWS_PER_STEP; r = r + 1u) {
+      let seg_base = r * {% if blockedStats %}TILE_CHUNKS{% else %}CHUNKS{% endif %} + sel * segment;
+      for (var t = 0u; t < segment; t = t + 1u) {
+        total = total + lane_sums[seg_base + t];
+      }
+    }
+    partials[(n * GROUPS + {% if blockedStats %}group_begin + {% endif %}sel) * SLABS + slab] = total;
+  }
+{% else %}
+  // WG is a power of two on this arm, so the halving stride is uniform across the
+  // workgroup and every lane reaches every barrier.
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lid.x < stride) {
+      lane_sums[lid.x] = lane_sums[lid.x] + lane_sums[lid.x + stride];
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  if (lid.x == 0u) {
+    partials[(n * GROUPS + g) * PARTIALS_PER_ROW + (c - g * CPG) * SPLITS + part] = lane_sums[0];
+  }
+{% endif %}
+{% else %}
+{% if groupUniform %}
+  let st = stats[n * GROUPS + g];
+  let mean = {{ vecF32 }}(st.x);
+  let inv_std = {{ vecF32 }}(st.y);
+{% else %}
+  // The four channels this lane holds do not all belong to one group, which is what
+  // `channels_per_group % 4 != 0` means. The statistics are still read once per lane,
+  // as a component-wise gather over at most two adjacent groups.
+  let row_base = n * GROUPS;
+  let s0 = stats[row_base + c0 / CPG];
+  let s1 = stats[row_base + (c0 + 1u) / CPG];
+  let s2 = stats[row_base + (c0 + 2u) / CPG];
+  let s3 = stats[row_base + (c0 + 3u) / CPG];
+  let mean = vec4<f32>(s0.x, s1.x, s2.x, s3.x);
+  let inv_std = vec4<f32>(s0.y, s1.y, s2.y, s3.y);
+{% endif %}
+{% if useVector %}
+{% if phaseLayout %}
+  let gamma_v = {{ channel_values("gamma", "c0") }};
+  let beta_v = {{ channel_values("beta", "c0") }};
+{% else %}
+  let gamma_v = {{ vecF32 }}(f32(gamma[c]));
+  let beta_v = {{ vecF32 }}(f32(beta[c]));
+{% endif %}
+{% else %}
+{% if phaseLayout %}
+  let gamma_v = f32(gamma[c0]);
+  let beta_v = f32(beta[c0]);
+{% else %}
+  let gamma_v = f32(gamma[c]);
+  let beta_v = f32(beta[c]);
+{% endif %}
+{% endif %}
+{% if phaseLayout %}
+  for (var hw = hw_begin + sub_row; hw < hw_end; hw = hw + ROWS_PER_STEP) {
+    let idx = image_base + hw * CHUNKS + chunk;
+{% else %}
+  for (var i = chunk_begin + lid.x; i < chunk_end; i = i + WG) {
+    let idx = base_v + i;
+{% endif %}
+    let normalized = ({{ vecF32 }}({{ applySrc }}[idx]) - mean) * inv_std;
+{% if activationSpec == 1 %}
+    Y[idx] = {{ storeElement }}(silu_e(gamma_v * normalized + beta_v));
+{% else %}
+    Y[idx] = {{ storeElement }}(gamma_v * normalized + beta_v);
+{% endif %}
+  }
+{% endif %}
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,226 @@

+{
+  "domain": "com.microsoft",
+  "name": "GroupNorm",
+  "sinceVersion": 1,
+  "inputs": {
+    "X": { "dtype": "T", "rank": 4 },
+    "gamma": { "dtype": "M", "rank": 1 },
+    "beta": { "dtype": "M", "rank": 1 }
+  },
+  "outputs": { "Y": { "dtype": "T", "rank": 4, "shape": "shapes.X" } },
+  "attributes": { "activation": {}, "groups": {}, "channels_last": { "default": 1 }, "epsilon": { "default": 0.00001 } },
+  "attributeConstraints": { "activation": { "required": true, "values": [0, 1] }, "groups": { "required": true } },
+  "typeConstraints": { "T": ["float32", "float16"], "M": ["float32", "float16"] },
+  "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "MAX_SPLITS": { "default": 256 },
+    "VALUES_PER_SPLIT": { "default": 4096 },
+    "SPLIT_MIN_ELEMENTS": { "default": 262144 },
+    "APPLY_SLAB_ELEMENTS": { "default": 8192 },
+    "FUSED_MIN_ROWS": { "default": 64 },
+    "FUSED_MAX_GROUP_BYTES": { "default": 131072 }
+  },
+  "derive": {
+    "attrsPresent": "has(attrs, \"groups\") and has(attrs, \"activation\")",
+    "numGroups": "attrs.groups if attrsPresent else 0",
+    "activationKind": "attrs.activation if attrsPresent else 0",
+    "channelsLast": "attrs.channels_last != 0",
+    "rankOk": "ranks.X == 4 and ranks.Y == 4 and ranks.gamma == 1 and ranks.beta == 1",
+    "dimN": "dim(shapes.X, 0) if rankOk else 0",
+    "dimC": "(dim(shapes.X, 3) if channelsLast else dim(shapes.X, 1)) if rankOk else 0",
+    "dimHW": "(dim(shapes.X, 1) * dim(shapes.X, 2) if channelsLast else dim(shapes.X, 2) * dim(shapes.X, 3)) if rankOk else 0",
+    "groupsOk": "numGroups >= 1 and dimC >= 1 and dimC % numGroups == 0",
+    "channelsPerGroup": "dimC / numGroups if groupsOk else 0",
+    "planeCount": "dimN * dimC",
+    "groupSize": "channelsPerGroup * dimHW",
+    "normRows": "dimN * numGroups",
+    "elementCount": "numel(shapes.X) if rankOk else 0",
+    "shapeOk": "rankOk and groupsOk and dimN > 0 and dimHW > 0 and groupSize > 0 and sameShape(shapes.Y, shapes.X) and dim(shapes.gamma, 0) == dimC and dim(shapes.beta, 0) == dimC",
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "statsVec4Ok": "(channelsPerGroup % 4 == 0 and dimC % 4 == 0) if channelsLast else (dimHW % 4 == 0)",
+    "applyVec4Ok": "(dimC % 4 == 0) if channelsLast else (dimHW % 4 == 0)",
+    "fusedVecWidth": "4 if statsVec4Ok else 1",
+    "statsVecWidth": "4 if statsVec4Ok else (2 if channelsLast and channelsPerGroup % 2 == 0 else 1)",
+    "applyVecWidth": "4 if applyVec4Ok else 1",
+    "workgroupCap": "max(1, floor(pow2ceil(min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap) + 1) / 2))",
+    "statsChunks": "max(1, floor(dimC / statsVecWidth))",
+    "fusedRowChunks": "max(1, floor(dimC / fusedVecWidth))",
+    "statsSegment": "max(1, floor(channelsPerGroup / statsVecWidth))",
+    "statsGroupsPerBlock": "max(1, min(numGroups, floor(workgroupCap / statsSegment))) if channelsLast and statsChunks > deviceWorkgroupCap and statsSegment <= workgroupCap else max(1, numGroups)",
+    "statsGroupBlocks": "ceilDiv(max(1, numGroups), statsGroupsPerBlock)",
+    "statsTileChunks": "statsGroupsPerBlock * statsSegment if statsGroupBlocks > 1 else statsChunks",
+    "applyChunks": "max(1, floor(dimC / applyVecWidth))",
+    "statsRowsPerStep": "max(1, floor(workgroupCap / statsTileChunks))",
+    "applyRowsPerStep": "max(1, floor(workgroupCap / applyChunks))",
+    "planeVecs": "max(1, floor(dimHW / statsVecWidth))",
+    "slabCount": "ceilDiv(dimHW, ceilDiv(dimHW, max(1, min(tunables.MAX_SPLITS, dimHW, ceilDiv(dimHW * dimC, tunables.VALUES_PER_SPLIT)))))",
+    "planeSplits": "max(1, min(tunables.MAX_SPLITS, planeVecs, ceilDiv(dimHW, tunables.VALUES_PER_SPLIT)))",
+    "splitCount": "slabCount if channelsLast else planeSplits",
+    "partialsPerRow": "slabCount if channelsLast else max(1, channelsPerGroup) * planeSplits",
+    "partialElements": "max(1, normRows * partialsPerRow)",
+    "statsElements": "max(1, normRows)",
+    "statsWorkgroup": "(statsTileChunks * statsRowsPerStep) if channelsLast else min(workgroupCap, max(32, pow2ceil(ceilDiv(planeVecs, planeSplits))))",
+    "applyWorkgroup": "(applyChunks * applyRowsPerStep) if channelsLast else statsWorkgroup",
+    "applySlabCount": "ceilDiv(dimHW, ceilDiv(dimHW, max(1, min(dimHW, min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(elementCount, tunables.APPLY_SLAB_ELEMENTS)))))",
+    "applySplit": "applySlabCount if channelsLast else planeSplits",
+    "applyUnits": "(dimN * applySlabCount) if channelsLast else (planeCount * planeSplits)",
+    "workUnits": "(dimN * slabCount * statsGroupBlocks) if channelsLast else (planeCount * planeSplits)",
+    "splitFits": "statsWorkgroup <= deviceWorkgroupCap and applyWorkgroup <= deviceWorkgroupCap and statsWorkgroup * 8 <= device.limits.maxComputeWorkgroupStorageSize and partialElements * 8 <= storageBufferLimit and workUnits > 0 and workUnits <= foldedDispatchCapacity and applyUnits <= foldedDispatchCapacity",
+    "fusedChunks": "max(1, floor(channelsPerGroup / fusedVecWidth)) if channelsLast else 1",
+    "fusedWorkgroup": "(fusedChunks * max(1, floor(workgroupCap / fusedChunks))) if channelsLast else min(workgroupCap, max(32, pow2ceil(max(1, floor(groupSize / fusedVecWidth)))))",
+    "fusedLanes": "pow2ceil(fusedWorkgroup)",
+    "fusedFits": "fusedWorkgroup <= deviceWorkgroupCap and fusedLanes * 8 <= device.limits.maxComputeWorkgroupStorageSize and normRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "combineWorkgroup": "min(workgroupCap, 64)",
+    "statsElement": "(\"vec\" ~ statsVecWidth ~ \"<\" ~ dtypes.T ~ \">\") if statsVecWidth > 1 else dtypes.T",
+    "fusedElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if statsVec4Ok else dtypes.T",
+    "applyElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if applyVec4Ok else dtypes.T",
+    "mScalar": "dtypes.M",
+    "tScalar": "dtypes.T",
+    "contractOk": "shapeOk and f16Ok(dtypes.T) and f16Ok(dtypes.M)"
+  },
+  "when": ["contractOk"],
+  "bindings": {
+    "xIn": { "arg": "X", "name": "X", "elementType": "$statsElement" },
+    "srcX": { "arg": "X", "name": "src", "elementType": "$statsElement" },
+    "gamma": { "elementType": "$mScalar" },
+    "beta": { "elementType": "$mScalar" },
+    "yOut": { "arg": "Y", "name": "Y", "elementType": "$applyElement" },
+    "partialsW": { "name": "partials", "buffer": "storage", "elementType": "vec2<f32>" },
+    "partialsR": { "name": "partials", "buffer": "read-only-storage", "elementType": "vec2<f32>" },
+    "statsW": { "name": "stats", "buffer": "storage", "elementType": "vec2<f32>" },
+    "statsR": { "name": "stats", "buffer": "read-only-storage", "elementType": "vec2<f32>" },
+    "srcApply": { "arg": "X", "name": "src", "elementType": "$applyElement" },
+    "yFused": { "arg": "Y", "name": "Y", "elementType": "$statsElement" }
+  },
+  "variants": [
+    {
+      "id": "split",
+      "priority": 20,
+      "when": ["splitFits"],
+      "demoteWhen": ["elementCount < tunables.SPLIT_MIN_ELEMENTS * statsGroupBlocks", "(not channelsLast) and normRows >= tunables.FUSED_MIN_ROWS and groupSize * dtypeBytes(tensorDtypes.X) <= tunables.FUSED_MAX_GROUP_BYTES"],
+      "derive": {
+        "phaseLayout": "channelsLast",
+        "hasSkip": "false",
+        "hasBias": "false",
+        "broadcastSkip": "false",
+        "vecSpec": "statsVecWidth",
+        "dimCSpec": "dimC",
+        "dimHWSpec": "dimHW",
+        "cpgSpec": "channelsPerGroup",
+        "groupsSpec": "numGroups",
+        "groupSizeSpec": "groupSize",
+        "normRowsSpec": "normRows",
+        "epsilonSpec": "attrs.epsilon",
+        "activationSpec": "activationKind",
+        "chunksSpec": "statsChunks",
+        "rowsPerStepSpec": "statsRowsPerStep",
+        "splitSpec": "splitCount",
+        "rowsPerSlabSpec": "ceilDiv(dimHW, slabCount)",
+        "planeVecsSpec": "planeVecs",
+        "chunkVecsSpec": "ceilDiv(planeVecs, planeSplits)",
+        "partialsPerRowSpec": "partialsPerRow",
+        "wgSpec": "statsWorkgroup",
+        "combineWgSpec": "combineWorkgroup",
+        "vecF32": "(\"vec\" ~ statsVecWidth ~ \"<f32>\") if statsVecWidth > 1 else \"f32\"",
+        "storeElement": "statsElement",
+        "applySrc": "\"src\"",
+        "groupUniform": "true"
+      },
+      "intermediates": [
+        { "id": "partials", "dtype": "float32", "shape": "[partialElements, 2]" },
+        { "id": "stats", "dtype": "float32", "shape": "[statsElements, 2]" }
+      ],
+      "passes": [
+        {
+          "id": "stats",
+          "name": "GroupNorm.Partials",
+          "shader": "group-norm-slab.wgsl.jinja",
+          "derive": {
+            "phase": "\"stats\"",
+            "unitsSpec": "workUnits",
+            "statsBlockedSpec": "channelsLast and statsGroupBlocks > 1",
+            "statsGroupsPerBlockSpec": "statsGroupsPerBlock",
+            "statsGroupBlocksSpec": "statsGroupBlocks",
+            "statsTileChunksSpec": "statsTileChunks"
+          },
+          "bindings": ["xIn", "partialsW"],
+          "dispatch": { "x": "min(workUnits, 65535)", "y": "ceilDiv(workUnits, 65535)", "z": 1 }
+        },
+        {
+          "id": "combine",
+          "name": "GroupNorm.Combine",
+          "shader": "group-norm-combine.wgsl.jinja",
+          "bindings": ["srcX", "partialsR", "statsW"],
+          "dispatch": {
+            "x": "min(ceilDiv((normRows), (combineWorkgroup)), 65535)",
+            "y": "ceilDiv(ceilDiv((normRows), (combineWorkgroup)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupNorm.Apply",
+          "shader": "group-norm-slab.wgsl.jinja",
+          "derive": {
+            "phase": "\"apply\"",
+            "vecSpec": "applyVecWidth",
+            "chunksSpec": "applyChunks",
+            "rowsPerStepSpec": "applyRowsPerStep",
+            "wgSpec": "applyWorkgroup",
+            "splitSpec": "applySplit",
+            "rowsPerSlabSpec": "ceilDiv(dimHW, applySlabCount)",
+            "groupUniform": "(channelsPerGroup % applyVecWidth == 0) if channelsLast else true",
+            "storeElement": "applyElement",
+            "vecF32": "\"vec4<f32>\" if applyVec4Ok else \"f32\"",
+            "unitsSpec": "applyUnits"
+          },
+          "bindings": ["srcApply", "gamma", "beta", "statsR", "yOut"],
+          "dispatch": { "x": "min(applyUnits, 65535)", "y": "ceilDiv(applyUnits, 65535)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "fused",
+      "priority": 10,
+      "when": ["fusedFits"],
+      "derive": {
+        "statsElement": "fusedElement",
+        "phaseLayout": "channelsLast",
+        "hasSkip": "false",
+        "hasBias": "false",
+        "broadcastSkip": "false",
+        "useVec4": "statsVec4Ok",
+        "vecSpec": "fusedVecWidth",
+        "dimCSpec": "dimC",
+        "dimHWSpec": "dimHW",
+        "cpgSpec": "channelsPerGroup",
+        "groupsSpec": "numGroups",
+        "groupSizeSpec": "groupSize",
+        "normRowsSpec": "normRows",
+        "epsilonSpec": "attrs.epsilon",
+        "activationSpec": "activationKind",
+        "chunksSpec": "fusedRowChunks",
+        "fusedChunksSpec": "fusedChunks",
+        "fusedRowsSpec": "max(1, floor(fusedWorkgroup / fusedChunks))",
+        "spanVecsSpec": "max(1, floor(groupSize / fusedVecWidth))",
+        "wgSpec": "fusedWorkgroup",
+        "ldsSpec": "fusedLanes",
+        "vecF32": "\"vec4<f32>\" if statsVec4Ok else \"f32\"",
+        "storeElement": "fusedElement",
+        "fusedSrc": "\"X\"",
+        "planeVecsSpec": "planeVecs"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GroupNorm.Fused",
+          "shader": "group-norm-fused.wgsl.jinja",
+          "bindings": ["xIn", "gamma", "beta", "yFused"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+  "name": "com.microsoft.GroupNorm",
+  "id": "_com_microsoft_groupnorm_webgpu_3aa81a3",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "jld33+mA1k/b0oKawSw9bnQxUcFLroArkzBwCnOdyCg=",
+      "group-norm-combine.wgsl.jinja": "weqcoPzdY5lQ8K/lZVEF03/sdbu1c/2jwdrW/KT9xr4=",
+      "group-norm-fused.wgsl.jinja": "uL3EIKVa0AnJbW7kB5MRPipXhZDr9KZ62soBHeu40K0=",
+      "group-norm-slab.wgsl.jinja": "T6ccs9TsauqsOjpC2hBzy7mRE/66CmLVI1F9JcBhk+w=",
+      "manifest.json": "U5sLYCEeiEBL842xP9mLAEgUQaJ3aioLBFcvJYpPMok=",
+      "test.json": "pEkAgZoOGQcSf9k/rHscjf21d18AlN/baC/QpVhOX08="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "split": ["group-norm-combine.wgsl.jinja", "group-norm-slab.wgsl.jinja"],
+      "fused": ["group-norm-fused.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/test.json ADDED Viewed

The diff for this file is too large to render. See raw diff