sync 91d990483a17

Browse files

Files changed (10) hide show

README.md +14 -10
build/webgpu/bench.json +2 -3
build/webgpu/instance-normalization-apply.wgsl.jinja +2 -5
build/webgpu/instance-normalization-batched-planes-vec4.wgsl.jinja +1 -5
build/webgpu/instance-normalization-splitk-combine.wgsl.jinja +15 -6
build/webgpu/instance-normalization-splitk-partials.wgsl.jinja +14 -8
build/webgpu/manifest.json +248 -324
build/webgpu/metadata.json +21 -11
build/webgpu/norm-row-stats.wgsl.jinja +103 -18
build/webgpu/test.json +22 -4

README.md CHANGED Viewed

@@ -18,17 +18,17 @@ See the [ONNX `InstanceNormalization` spec](https://onnx.ai/onnx/operators/onnx_
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `input` | `T` | — | — | Input tensor of shape `(N x C x D1 x ... x Dn)`; at least 3-D. | required |
-| `scale` | `scale` | `T` | `1` | — | 1-D scale tensor of size C, one scale factor per channel. | required |
-| `B` | `b` | `T` | `1` | — | 1-D bias tensor of size C, one bias value per channel. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `output` | `output` | `T` | same as `input` | same as `input` | Normalized output tensor; same shape as the input. | required |
 ## Attributes
@@ -50,7 +50,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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 + tuning cases
@@ -62,10 +62,14 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `input` | — | `T` | — | — | Input tensor of shape `(N x C x D1 x ... x Dn)`; at least 3-D. | required |
+| `scale` | — | `T` | `1` | — | 1-D scale tensor of size C, one scale factor per channel. | required |
+| `b` | `B` | `T` | `1` | — | 1-D bias tensor of size C, one bias value per channel. | required |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | same as `input` | same as `input` | Normalized output tensor; same shape as the input. | required |
 ## Attributes
 ## 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 + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+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.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.InstanceNormalization",
   "cases": [
     {
       "name": "nchw_4x64x128x128",
@@ -209,7 +208,7 @@
       "name": "splitk-priority-cliff-c256-256x256",
       "preset": "stress",
       "provenance": {
-        "source": "authored for variant coverage",
         "notes": "Realistic 16.8M-element feature map that pins the selector boundary between plane_subgroup_vec4 and plane_splitk."
       },
       "vars": { "dtype": "float32", "batch": 1, "channels": 256, "spatial": 65536 },
@@ -233,7 +232,7 @@
       "name": "splitk-priority-cliff-c32-512x512",
       "preset": "stress",
       "provenance": {
-        "source": "authored for variant coverage",
         "notes": "Realistic 8.4M-element high-resolution feature map that pins the selector boundary between plane_subgroup_vec4 and plane_splitk."
       },
       "vars": { "dtype": "float32", "batch": 1, "channels": 32, "spatial": 262144 },

 {
   "cases": [
     {
       "name": "nchw_4x64x128x128",
       "name": "splitk-priority-cliff-c256-256x256",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic benchmark",
         "notes": "Realistic 16.8M-element feature map that pins the selector boundary between plane_subgroup_vec4 and plane_splitk."
       },
       "vars": { "dtype": "float32", "batch": 1, "channels": 256, "spatial": 65536 },
       "name": "splitk-priority-cliff-c32-512x512",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic benchmark",
         "notes": "Realistic 8.4M-element high-resolution feature map that pins the selector boundary between plane_subgroup_vec4 and plane_splitk."
       },
       "vars": { "dtype": "float32", "batch": 1, "channels": 32, "spatial": 262144 },

build/webgpu/instance-normalization-apply.wgsl.jinja CHANGED Viewed

@@ -2,9 +2,6 @@
 // The vectorized route packs four adjacent spatial values per invocation; each
 // packed load and store remains within one plane.
 {% set vectorized = vectorized if vectorized is defined else false %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {% set LOAD_OPEN = "vec4<f32>(" if usesF16 else "" %}
 {% set LOAD_CLOSE = ")" if usesF16 else "" %}
 {% set STORE_OPEN = "vec4<f16>(" if usesF16 else "" %}
@@ -16,11 +13,11 @@ enable f16;
 const WG: u32 = {{ applyWorkgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {% if not vectorized %}
   // 2D-folded flat index: gid.y carries the high bits after dispatch folding.
 {% endif %}
-  let index = gid.x + gid.y * nwg.x * WG;
   if (index >= params.count) {
     return;
   }

 // The vectorized route packs four adjacent spatial values per invocation; each
 // packed load and store remains within one plane.
 {% set vectorized = vectorized if vectorized is defined else false %}
 {% set LOAD_OPEN = "vec4<f32>(" if usesF16 else "" %}
 {% set LOAD_CLOSE = ")" if usesF16 else "" %}
 {% set STORE_OPEN = "vec4<f16>(" if usesF16 else "" %}
 const WG: u32 = {{ applyWorkgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {% if not vectorized %}
   // 2D-folded flat index: gid.y carries the high bits after dispatch folding.
 {% endif %}
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (index >= params.count) {
     return;
   }

build/webgpu/instance-normalization-batched-planes-vec4.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // A power-of-two lane cohort reduces one plane while several cohorts share a
@@ -19,13 +16,12 @@ var<workgroup> plane_shift: array<f32, PLANES_PER_WG>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
-  @builtin(num_workgroups) workgroup_count: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>
 ) {
   let tid = local.x;
   let plane_in_workgroup = tid / LANES;
   let lane = tid % LANES;
-  let group = workgroup.x + workgroup.y * workgroup_count.x;
   let row = group * PLANES_PER_WG + plane_in_workgroup;
   let is_active = row < params.rows;
   let base = row * HIDDEN_V4;

 {{ env.wgsl.resourceDeclarations }}
 // A power-of-two lane cohort reduces one plane while several cohorts share a
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>
 ) {
   let tid = local.x;
   let plane_in_workgroup = tid / LANES;
   let lane = tid % LANES;
+  let group = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let row = group * PLANES_PER_WG + plane_in_workgroup;
   let is_active = row < params.rows;
   let base = row * HIDDEN_V4;

build/webgpu/instance-normalization-splitk-combine.wgsl.jinja CHANGED Viewed

@@ -1,14 +1,15 @@
 // Fold SPLIT per-plane (sum, sum-of-squares) partials into mean and inverse
-// standard deviation. One thread handles each plane. Variance uses
-// E[x^2] - E[x]^2; max(value, 0) guards against negative rounding residue.
 {{ env.wgsl.resourceDeclarations }}
 const SPLIT: u32 = {{ split }}u;
 const COMBINE_WG: u32 = {{ combineWorkgroupSize }}u;
 @compute @workgroup_size(COMBINE_WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
-  let plane = gid.x + gid.y * nwg.x * COMBINE_WG;
   if (plane >= params.planes) {
     return;
   }
@@ -20,8 +21,16 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
     total_sq = total_sq + partials[(b + k) * 2u + 1u];
   }
   let n = f32(params.spatial);
-  let mean = total / n;
-  let variance = max(total_sq / n - mean * mean, 0.0);
   stats[plane * 2u] = mean;
   stats[plane * 2u + 1u] = inverseSqrt(variance + params.epsilon);
 }

 // Fold SPLIT per-plane (sum, sum-of-squares) partials into mean and inverse
+// standard deviation. One thread handles each plane. The partials are centred on
+// the plane's first element, so E[y^2] - E[y]^2 keeps the variance a raw second
+// moment would cancel away; max(value, 0) guards against negative rounding residue.
 {{ env.wgsl.resourceDeclarations }}
 const SPLIT: u32 = {{ split }}u;
 const COMBINE_WG: u32 = {{ combineWorkgroupSize }}u;
 @compute @workgroup_size(COMBINE_WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let plane = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * COMBINE_WG;
   if (plane >= params.planes) {
     return;
   }
     total_sq = total_sq + partials[(b + k) * 2u + 1u];
   }
   let n = f32(params.spatial);
+  // The partials are accumulated around the plane's first element; undo the shift
+  // on the mean and leave the variance, which the shift does not change.
+{% if vectorizedSpec %}
+  let shift = f32(input[plane * (params.spatial / 4u)].x);
+{% else %}
+  let shift = f32(input[plane * params.spatial]);
+{% endif %}
+  let centred_mean = total / n;
+  let mean = shift + centred_mean;
+  let variance = max(total_sq / n - centred_mean * centred_mean, 0.0);
   stats[plane * 2u] = mean;
   stats[plane * 2u + 1u] = inverseSqrt(variance + params.epsilon);
 }

build/webgpu/instance-normalization-splitk-partials.wgsl.jinja CHANGED Viewed

@@ -51,9 +51,6 @@
    standard deviation, and the apply pass normalizes. */
 {% set vectorized = vectorized if vectorized is defined else false %}
 {% set useSubgroups = useSubgroups if useSubgroups is defined else false %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
@@ -74,12 +71,11 @@ var<workgroup> red_sq: array<f32, WG>;
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>{% if useSubgroups %},
         @builtin(subgroup_invocation_id) subgroup_lane: u32,
         @builtin(subgroup_id) subgroup_id: u32,
         @builtin(num_subgroups) num_subgroups: u32{% endif %}) {
-  let plane = wg.x + wg.y * nwg.x;
   if (plane >= params.planes) {
     return;
   }
@@ -96,17 +92,27 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   if (end > spatial) { end = spatial; }
   let base = plane * spatial;
   var s = 0.0;
   var sq = 0.0;
   var i = start + tid;
   loop {
     if (i >= end) { break; }
 {% if vectorized %}
-    let v = {{ LOAD_OPEN }}input[base + i]{{ LOAD_CLOSE }};
     s = s + v.x + v.y + v.z + v.w;
     sq = sq + dot(v, v);
 {% else %}
-    let v = f32(input[base + i]);
     s = s + v;
     sq = sq + v * v;
 {% endif %}

    standard deviation, and the apply pass normalizes. */
 {% set vectorized = vectorized if vectorized is defined else false %}
 {% set useSubgroups = useSubgroups if useSubgroups is defined else false %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
         @builtin(subgroup_invocation_id) subgroup_lane: u32,
         @builtin(subgroup_id) subgroup_id: u32,
         @builtin(num_subgroups) num_subgroups: u32{% endif %}) {
+  let plane = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (plane >= params.planes) {
     return;
   }
   if (end > spatial) { end = spatial; }
   let base = plane * spatial;
+  // Raw second moments cancel: a plane centred on 8192 with unit variance loses
+  // the variance entirely in E[x^2] - E[x]^2, and the combine's max(.,0) then
+  // reports zero. Both passes accumulate around the plane's first element, which
+  // costs one broadcast load and leaves the squared term holding the residual.
+{% if vectorized %}
+  let shift = f32({{ LOAD_OPEN }}input[base]{{ LOAD_CLOSE }}.x);
+  let shift4 = vec4<f32>(shift);
+{% else %}
+  let shift = f32(input[base]);
+{% endif %}
   var s = 0.0;
   var sq = 0.0;
   var i = start + tid;
   loop {
     if (i >= end) { break; }
 {% if vectorized %}
+    let v = {{ LOAD_OPEN }}input[base + i]{{ LOAD_CLOSE }} - shift4;
     s = s + v.x + v.y + v.z + v.w;
     sq = sq + dot(v, v);
 {% else %}
+    let v = f32(input[base + i]) - shift;
     s = s + v;
     sq = sq + v * v;
 {% endif %}

build/webgpu/manifest.json CHANGED Viewed

@@ -2,51 +2,28 @@
   "domain": "ai.onnx",
   "name": "InstanceNormalization",
   "sinceVersion": 6,
-  "description": "Applies instance normalization to the input: `y = scale * (x - mean) / sqrt(variance + epsilon) + B`, where `mean` and `variance` are computed per instance per channel over the spatial dimensions. Equivalent to batch normalization with a batch size of one per channel.",
-  "inputs": [
-    { "role": "input", "dtype": "T", "description": "Input tensor of shape `(N x C x D1 x ... x Dn)`; at least 3-D." },
-    {
-      "role": "scale",
-      "dtype": "T",
-      "rank": 1,
-      "description": "1-D scale tensor of size C, one scale factor per channel."
-    },
-    { "role": "B", "dtype": "T", "rank": 1, "description": "1-D bias tensor of size C, one bias value per channel." }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "rank": "ranks.input",
-      "description": "Normalized output tensor; same shape as the input.",
-      "shape": "shapes.input"
-    }
-  ],
-  "attributes": { "epsilon": 0.00001 },
-  "attributeDescriptions": {
-    "epsilon": "Small constant added to the variance before taking the square root to avoid division by zero."
   },
   "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "input": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "scale": { "kind": "tensor", "semantic": "scale", "role": "input" },
-    "b": { "kind": "tensor", "semantic": "B", "role": "input" },
-    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "MAX_STATS_SPLITS": 256,
-    "STATS_VALUES_PER_SPLIT": 2048,
-    "SPLIT_STATS_MIN_SPATIAL": 65536,
-    "SPLIT_STATS_MAX_PLANES": 256,
-    "COMBINE_WORKGROUP_SIZE": 64,
-    "BATCHED_MIN_PLANES_PER_WORKGROUP": 8
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
-    "instanceContractOk": "f16Ok(dtypes.T) and ranks.input >= 3 and ranks.output == ranks.input and sameShape(shapes.output, shapes.input) and ranks.scale == 1 and ranks.B == 1 and dim(shapes.scale, 0) == dim(shapes.input, 1) and dim(shapes.B, 0) == dim(shapes.input, 1)",
     "instancePlanes": "dim(shapes.input, 0) * dim(shapes.input, 1)",
     "instanceSpatial": "inner(shapes.input, 1)",
     "normDeviceWorkgroupCap": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
@@ -65,226 +42,48 @@
     "instanceStatsBytes": "instancePlanes * 2 * 4",
     "instanceStatsFits": "instanceStatsBytes <= device.limits.maxStorageBufferBindingSize and instanceStatsBytes <= device.limits.maxBufferSize",
     "instanceRowCovered": "instanceContractOk and instanceRowWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize",
-    "instanceSplitCount": "min(tunables.MAX_STATS_SPLITS, device.limits.maxComputeWorkgroupsPerDimension, pow2ceil(ceilDiv(instanceSpatial, tunables.STATS_VALUES_PER_SPLIT)))",
     "instancePartialBytes": "instancePlanes * instanceSplitCount * 2 * 4",
-    "splitStatsCovered": "instanceRowCovered and instanceStatsFits and instancePlanes <= tunables.SPLIT_STATS_MAX_PLANES and instancePlanes <= device.limits.maxComputeWorkgroupsPerDimension and instanceSpatial >= tunables.SPLIT_STATS_MIN_SPATIAL and instancePartialBytes <= device.limits.maxStorageBufferBindingSize and instancePartialBytes <= device.limits.maxBufferSize",
     "splitStatsPreferred": "splitStatsCovered and instancePlanes < normSubgroupMax"
   },
-  "bindingSets": {
-    "planeIo": [
-      {
-        "name": "x",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$T"
-      },
-      { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      {
-        "name": "y",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$ioElement"
-      }
-    ],
-    "plane": [
-      {
-        "name": "x",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$T"
-      },
-      { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      {
-        "name": "y",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(dim(shapes.input, 0) * dim(shapes.input, 1), device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
-        }
-      }
-    ],
-    "planeBatched": [
-      {
-        "name": "x",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$T"
-      },
-      { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      {
-        "name": "y",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [{ "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" }]
-        }
-      }
-    ],
-    "applyScalar": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$T"
-      },
-      { "name": "stats", "semantic": "stats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$T"
-      },
-      { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output)" },
-            { "name": "channels", "type": "u32", "value": "dim(shapes.input, 1)" },
-            { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
-          ]
-        }
-      }
-    ],
-    "applyVec4": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "stats", "semantic": "stats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$T"
-      },
-      { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      {
-        "name": "output",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" },
-            { "name": "channels", "type": "u32", "value": "dim(shapes.input, 1)" },
-            { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
-          ]
-        }
-      }
-    ],
-    "splitPartials": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$splitInputElement"
-      },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "planes", "type": "u32", "value": "instancePlanes" },
-            { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
-          ]
-        }
-      }
-    ],
-    "splitCombine": [
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "stats", "semantic": "stats", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "planes", "type": "u32", "value": "instancePlanes" },
-            { "name": "spatial", "type": "u32", "value": "instanceSpatial" },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "plane_batched_vec4",
       "priority": 115,
       "when": ["instanceRowCovered", "instanceSpatial % 4 == 0", "instanceSpatial >= 4", "instancePlanes >= normWorkgroupCap", "instanceBatchedVec4PlanesPerWorkgroup >= tunables.BATCHED_MIN_PLANES_PER_WORKGROUP", "instanceBatchedVec4StorageBytes <= device.limits.maxComputeWorkgroupStorageSize"],
-      "demoteWhen": ["reportedNonWave32Adapter and instancePlanes <= device.limits.maxComputeWorkgroupsPerDimension"],
-      "constants": {
         "usesF16": "dtypes.T == \"f16\"",
         "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
@@ -301,103 +100,158 @@
           "id": "main",
           "name": "InstanceNormalization.PlaneBatchedVec4",
           "shader": "instance-normalization-batched-planes-vec4.wgsl.jinja",
-          "bindings": "planeBatched",
-          "dispatch": { "workgroups": "instanceBatchedVec4Workgroups" }
         }
       ]
     },
     {
       "id": "plane_subgroup_vec4",
       "priority": 110,
-      "requires": { "features": [] },
       "when": ["instanceRowCovered", "inner(shapes.input, 1) % 4 == 0", "instanceVec4SubgroupEfficient"],
-      "constants": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "InstanceNormalization.plane_subgroup_vec4",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"instance\"",
-              "vec4": true,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "instanceSpatial",
-              "wg": "instanceVec4Workgroup",
-              "epsilon": "attrs.epsilon",
-              "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1",
-              "hiddenVec": "instanceSpatial / 4",
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "combineSubgroups": "hasSubgroupId"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "plane",
-          "dispatch": { "workgroups": "instancePlanes" }
         }
       ]
     },
     {
       "id": "plane_subgroup_vec4_scalar_io",
       "priority": 111,
-      "requires": { "features": ["subgroups"] },
       "when": ["dtypes.T == \"f32\"", "wave32Adapter", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "instanceRowCovered", "instanceSpatial % 4 == 0", "instanceVec4SubgroupEfficient"],
-      "constants": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "InstanceNormalization.plane_subgroup_vec4_scalar_io",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"instance\"",
-              "vec4": true,
-              "scalarIo": true,
-              "scalar": "dtypes.T",
-              "usesF16": false,
-              "hidden": "instanceSpatial",
-              "wg": "instanceVec4Workgroup",
-              "epsilon": "attrs.epsilon",
-              "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1",
-              "hiddenVec": "instanceSpatial / 4",
-              "vecType": "\"vec4<f32>\"",
-              "combineSubgroups": true
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "plane",
-          "dispatch": { "workgroups": "instancePlanes" }
         }
       ]
     },
     {
       "id": "plane_subgroup",
       "priority": 100,
-      "requires": { "features": [] },
       "when": ["instanceRowCovered"],
-      "constants": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "InstanceNormalization.plane_subgroup",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"instance\"",
-              "vec4": false,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "instanceSpatial",
-              "wg": "instanceScalarWorkgroup",
-              "epsilon": "attrs.epsilon",
-              "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1",
-              "combineSubgroups": "hasSubgroupId"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "plane",
-          "dispatch": { "workgroups": "instancePlanes" }
         }
       ]
     },
@@ -405,7 +259,7 @@
       "id": "plane_splitk_vec4",
       "priority": 121,
       "when": ["splitStatsPreferred", "instanceSpatial % 4 == 0"],
-      "constants": {
         "vectorized": true,
         "usesF16": "dtypes.T == \"f16\"",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
@@ -425,23 +279,58 @@
           "id": "partials",
           "name": "InstanceNormalization.SplitKPartialsVec4",
           "shader": "instance-normalization-splitk-partials.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "splitPartials",
-          "dispatch": { "workgroups": "instancePlanes", "z": "instanceSplitCount" }
         },
         {
           "id": "combine",
           "name": "InstanceNormalization.SplitKCombine",
           "shader": "instance-normalization-splitk-combine.wgsl.jinja",
-          "bindings": "splitCombine",
-          "dispatch": { "threads": "instancePlanes", "workgroupSize": "constants.combineWorkgroupSize" }
         },
         {
           "id": "apply",
           "name": "InstanceNormalization.ApplyVec4",
           "shader": "instance-normalization-apply.wgsl.jinja",
-          "bindings": "applyVec4",
-          "dispatch": { "threads": "numel(shapes.output) / 4", "workgroupSize": "constants.applyWorkgroupSize" }
         }
       ]
     },
@@ -449,7 +338,7 @@
       "id": "plane_splitk",
       "priority": 120,
       "when": ["splitStatsPreferred"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "splitInputElement": "dtypes.T",
@@ -467,22 +356,57 @@
           "id": "partials",
           "name": "InstanceNormalization.SplitKPartials",
           "shader": "instance-normalization-splitk-partials.wgsl.jinja",
-          "bindings": "splitPartials",
-          "dispatch": { "workgroups": "instancePlanes", "z": "instanceSplitCount" }
         },
         {
           "id": "combine",
           "name": "InstanceNormalization.SplitKCombine",
           "shader": "instance-normalization-splitk-combine.wgsl.jinja",
-          "bindings": "splitCombine",
-          "dispatch": { "threads": "instancePlanes", "workgroupSize": "constants.combineWorkgroupSize" }
         },
         {
           "id": "apply",
           "name": "InstanceNormalization.Apply",
           "shader": "instance-normalization-apply.wgsl.jinja",
-          "bindings": "applyScalar",
-          "dispatch": { "threads": "numel(shapes.output)", "workgroupSize": "constants.applyWorkgroupSize" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "InstanceNormalization",
   "sinceVersion": 6,
+  "inputs": {
+    "input": { "dtype": "T" },
+    "scale": { "dtype": "T", "rank": 1 },
+    "b": { "onnx": "B", "dtype": "T", "rank": 1 }
   },
+  "outputs": { "output": { "dtype": "T", "rank": "ranks.input", "shape": "shapes.input" } },
+  "attributes": { "epsilon": { "default": 0.00001 } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "MAX_STATS_SPLITS": { "default": 256 },
+    "STATS_VALUES_PER_SPLIT": { "default": 2048 },
+    "SPLIT_STATS_MIN_SPATIAL": { "default": 65536 },
+    "SPLIT_STATS_MAX_PLANES": { "default": 256 },
+    "COMBINE_WORKGROUP_SIZE": { "default": 64 },
+    "BATCHED_MIN_PLANES_PER_WORKGROUP": { "default": 8 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
+    "instanceContractOk": "f16Ok(dtypes.T) and ranks.input >= 3 and ranks.output == ranks.input and sameShape(shapes.output, shapes.input) and ranks.scale == 1 and ranks.b == 1 and dim(shapes.scale, 0) == dim(shapes.input, 1) and dim(shapes.b, 0) == dim(shapes.input, 1)",
     "instancePlanes": "dim(shapes.input, 0) * dim(shapes.input, 1)",
     "instanceSpatial": "inner(shapes.input, 1)",
     "normDeviceWorkgroupCap": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "instanceStatsBytes": "instancePlanes * 2 * 4",
     "instanceStatsFits": "instanceStatsBytes <= device.limits.maxStorageBufferBindingSize and instanceStatsBytes <= device.limits.maxBufferSize",
     "instanceRowCovered": "instanceContractOk and instanceRowWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "instanceSplitCount": "min(tunables.MAX_STATS_SPLITS, min(device.limits.maxComputeWorkgroupsPerDimension, 65535), pow2ceil(ceilDiv(instanceSpatial, tunables.STATS_VALUES_PER_SPLIT)))",
     "instancePartialBytes": "instancePlanes * instanceSplitCount * 2 * 4",
+    "splitStatsCovered": "instanceRowCovered and instanceStatsFits and instancePlanes <= tunables.SPLIT_STATS_MAX_PLANES and instancePlanes <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and instanceSpatial >= tunables.SPLIT_STATS_MIN_SPATIAL and instancePartialBytes <= device.limits.maxStorageBufferBindingSize and instancePartialBytes <= device.limits.maxBufferSize",
     "splitStatsPreferred": "splitStatsCovered and instancePlanes < normSubgroupMax"
   },
+  "bindings": {
+    "x": { "arg": "input", "buffer": "read-only-storage", "elementType": "$ioElement" },
+    "scale": { "buffer": "read-only-storage", "elementType": "$T" },
+    "bias": { "arg": "b", "buffer": "read-only-storage", "elementType": "$T" },
+    "y": { "arg": "output", "buffer": "storage", "elementType": "$ioElement" },
+    "input": { "buffer": "read-only-storage", "elementType": "$splitInputElement" },
+    "input_2": { "name": "input", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "stats_2": { "name": "stats", "buffer": "read-only-storage", "elementType": "f32" },
+    "output": { "buffer": "storage", "elementType": "$vectorScalar" },
+    "params_5": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" },
+        { "name": "channels", "type": "u32", "value": "dim(shapes.input, 1)" },
+        { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
+      ]
+    },
+    "input_3": { "name": "input", "buffer": "read-only-storage", "elementType": "$T" },
+    "output_2": { "name": "output", "buffer": "storage", "elementType": "$T" },
+    "params_6": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output)" },
+        { "name": "channels", "type": "u32", "value": "dim(shapes.input, 1)" },
+        { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "plane_batched_vec4",
       "priority": 115,
       "when": ["instanceRowCovered", "instanceSpatial % 4 == 0", "instanceSpatial >= 4", "instancePlanes >= normWorkgroupCap", "instanceBatchedVec4PlanesPerWorkgroup >= tunables.BATCHED_MIN_PLANES_PER_WORKGROUP", "instanceBatchedVec4StorageBytes <= device.limits.maxComputeWorkgroupStorageSize"],
+      "demoteWhen": ["reportedNonWave32Adapter and instancePlanes <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": {
         "usesF16": "dtypes.T == \"f16\"",
         "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
           "id": "main",
           "name": "InstanceNormalization.PlaneBatchedVec4",
           "shader": "instance-normalization-batched-planes-vec4.wgsl.jinja",
+          "bindings": [
+            "x",
+            "scale",
+            "bias",
+            "y",
+            {
+              "name": "params",
+              "struct": [{ "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" }]
+            }
+          ],
+          "dispatch": {
+            "x": "min(instanceBatchedVec4Workgroups, 65535)",
+            "y": "ceilDiv(instanceBatchedVec4Workgroups, 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "plane_subgroup_vec4",
       "priority": 110,
       "when": ["instanceRowCovered", "inner(shapes.input, 1) % 4 == 0", "instanceVec4SubgroupEfficient"],
+      "requires": { "features": [] },
+      "derive": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "InstanceNormalization.plane_subgroup_vec4",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"instance\"",
+            "vec4": true,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "instanceSpatial",
+            "wg": "instanceVec4Workgroup",
+            "epsilon": "attrs.epsilon",
+            "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1",
+            "hiddenVec": "instanceSpatial / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId"
           },
+          "bindings": [
+            "x",
+            "scale",
+            "bias",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" },
+                {
+                  "name": "rowStride",
+                  "type": "u32",
+                  "value": "max(1, min(dim(shapes.input, 0) * dim(shapes.input, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+                }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(instancePlanes, 65535)", "y": "ceilDiv(instancePlanes, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "plane_subgroup_vec4_scalar_io",
       "priority": 111,
       "when": ["dtypes.T == \"f32\"", "wave32Adapter", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "instanceRowCovered", "instanceSpatial % 4 == 0", "instanceVec4SubgroupEfficient"],
+      "requires": { "features": ["subgroups"] },
+      "derive": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "InstanceNormalization.plane_subgroup_vec4_scalar_io",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"instance\"",
+            "vec4": true,
+            "scalarIo": true,
+            "scalar": "dtypes.T",
+            "usesF16Spec": false,
+            "hidden": "instanceSpatial",
+            "wg": "instanceVec4Workgroup",
+            "epsilon": "attrs.epsilon",
+            "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1",
+            "hiddenVec": "instanceSpatial / 4",
+            "vecType": "\"vec4<f32>\"",
+            "combineSubgroups": true
           },
+          "bindings": [
+            "x",
+            "scale",
+            "bias",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" },
+                {
+                  "name": "rowStride",
+                  "type": "u32",
+                  "value": "max(1, min(dim(shapes.input, 0) * dim(shapes.input, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+                }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(instancePlanes, 65535)", "y": "ceilDiv(instancePlanes, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "plane_subgroup",
       "priority": 100,
       "when": ["instanceRowCovered"],
+      "requires": { "features": [] },
+      "derive": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "InstanceNormalization.plane_subgroup",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"instance\"",
+            "vec4": false,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "instanceSpatial",
+            "wg": "instanceScalarWorkgroup",
+            "epsilon": "attrs.epsilon",
+            "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1",
+            "combineSubgroups": "hasSubgroupId"
           },
+          "bindings": [
+            "x",
+            "scale",
+            "bias",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" },
+                {
+                  "name": "rowStride",
+                  "type": "u32",
+                  "value": "max(1, min(dim(shapes.input, 0) * dim(shapes.input, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+                }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(instancePlanes, 65535)", "y": "ceilDiv(instancePlanes, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "plane_splitk_vec4",
       "priority": 121,
       "when": ["splitStatsPreferred", "instanceSpatial % 4 == 0"],
+      "derive": {
         "vectorized": true,
         "usesF16": "dtypes.T == \"f16\"",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
           "id": "partials",
           "name": "InstanceNormalization.SplitKPartialsVec4",
           "shader": "instance-normalization-splitk-partials.wgsl.jinja",
+          "bindings": [
+            "input",
+            { "name": "partials", "buffer": "storage", "elementType": "f32" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "planes", "type": "u32", "value": "instancePlanes" },
+                { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(instancePlanes, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(instancePlanes, DISPATCH_FOLD_WIDTH)",
+            "z": "instanceSplitCount"
+          },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "combine",
           "name": "InstanceNormalization.SplitKCombine",
           "shader": "instance-normalization-splitk-combine.wgsl.jinja",
+          "derive": { "vectorizedSpec": true },
+          "bindings": [
+            "input",
+            { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "stats", "buffer": "storage", "elementType": "f32" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "planes", "type": "u32", "value": "instancePlanes" },
+                { "name": "spatial", "type": "u32", "value": "instanceSpatial" },
+                { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((instancePlanes), (combineWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((instancePlanes), (combineWorkgroupSize)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "apply",
           "name": "InstanceNormalization.ApplyVec4",
           "shader": "instance-normalization-apply.wgsl.jinja",
+          "bindings": ["input_2", "stats_2", "scale", "bias", "output", "params_5"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / 4), (applyWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / 4), (applyWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "plane_splitk",
       "priority": 120,
       "when": ["splitStatsPreferred"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "splitInputElement": "dtypes.T",
           "id": "partials",
           "name": "InstanceNormalization.SplitKPartials",
           "shader": "instance-normalization-splitk-partials.wgsl.jinja",
+          "bindings": [
+            "input",
+            { "name": "partials", "buffer": "storage", "elementType": "f32" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "planes", "type": "u32", "value": "instancePlanes" },
+                { "name": "spatial", "type": "u32", "value": "instanceSpatial" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(instancePlanes, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(instancePlanes, DISPATCH_FOLD_WIDTH)",
+            "z": "instanceSplitCount"
+          }
         },
         {
           "id": "combine",
           "name": "InstanceNormalization.SplitKCombine",
           "shader": "instance-normalization-splitk-combine.wgsl.jinja",
+          "derive": { "vectorizedSpec": false },
+          "bindings": [
+            "input",
+            { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "stats", "buffer": "storage", "elementType": "f32" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "planes", "type": "u32", "value": "instancePlanes" },
+                { "name": "spatial", "type": "u32", "value": "instanceSpatial" },
+                { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((instancePlanes), (combineWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((instancePlanes), (combineWorkgroupSize)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "apply",
           "name": "InstanceNormalization.Apply",
           "shader": "instance-normalization-apply.wgsl.jinja",
+          "bindings": ["input_3", "stats_2", "scale", "bias", "output_2", "params_6"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output)), (applyWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output)), (applyWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,22 +1,32 @@
 {
   "name": "ai.onnx.InstanceNormalization",
-  "id": "_ai_onnx_instancenormalization_webgpu_1aea375",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "GgyPTi2dmoXyw7KRjHv05H93qU5QjSYSvl7Sq6uxIQ8=",
-      "instance-normalization-apply.wgsl.jinja": "xg8Pb+LQ5uMHQJNjPgxQAC9te2uiz1hfhLiKSxqD6WQ=",
-      "instance-normalization-batched-planes-vec4.wgsl.jinja": "fj7ftnvxAyb54O/TFqt5sezNAX0zVCyQNZfWZ34EiLE=",
-      "instance-normalization-splitk-combine.wgsl.jinja": "faaEHRq8I1cCglDAVTTpxsfaUtjDKjQOgDF6cobZWhY=",
-      "instance-normalization-splitk-partials.wgsl.jinja": "EPM/k6ud99fto6j+fbpwVBzJbg65sv3aiO4Cg6X9lLw=",
-      "manifest.json": "bM48AIoFxoG9TXkEwvxB57O6GrOiJMrH/MTEU39Fplk=",
-      "norm-row-stats.wgsl.jinja": "ZK1Wy+fDkHUODnGtjaF3FoacpqZidO/r2VxqrHtpBCg=",
-      "test.json": "TTDhrrElGsd6piLKBRdf7VIHBRzb1DwvWSgc5xB3iyk="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.InstanceNormalization" }
 }

 {
   "name": "ai.onnx.InstanceNormalization",
+  "id": "_ai_onnx_instancenormalization_webgpu_16b3576",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "HBqAImIaa6N6gGqA53oRQ/ZvaCAFNJPKltDsBnjF5wQ=",
+      "instance-normalization-apply.wgsl.jinja": "ss27/JnTr5XXoUd0lj8yyYcXfYImAO5jPB9BdkneLwU=",
+      "instance-normalization-batched-planes-vec4.wgsl.jinja": "cLkDhQOaM/T+im43mRMyLa+kEoIHfm8i31IdkiyeMcI=",
+      "instance-normalization-splitk-combine.wgsl.jinja": "z2uqoUYCw3foyDxmSSD9i8fjH3Z6HFpzVhBlNCV11BM=",
+      "instance-normalization-splitk-partials.wgsl.jinja": "SymImIFJ0cyKgoJZtsuTtPAEy2m8CtCKWn8ylnuBQJE=",
+      "manifest.json": "Eq6DFSvT5bXhoGN7gsrUL106yaDyqRNzf+r4t2e/Nh0=",
+      "norm-row-stats.wgsl.jinja": "pIQ85rOY7wxOVoImYTZF8yUjKpdvCKAH4h2EBmD1Kvc=",
+      "test.json": "AIj6JyZjhnniT+w2T4DZkIjVYrltvKmjY6cFQcNm+kY="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "plane_batched_vec4": ["instance-normalization-batched-planes-vec4.wgsl.jinja"],
+      "plane_subgroup_vec4": ["norm-row-stats.wgsl.jinja"],
+      "plane_subgroup_vec4_scalar_io": ["norm-row-stats.wgsl.jinja"],
+      "plane_subgroup": ["norm-row-stats.wgsl.jinja"],
+      "plane_splitk_vec4": ["instance-normalization-apply.wgsl.jinja", "instance-normalization-splitk-combine.wgsl.jinja", "instance-normalization-splitk-partials.wgsl.jinja"],
+      "plane_splitk": ["instance-normalization-apply.wgsl.jinja", "instance-normalization-splitk-combine.wgsl.jinja", "instance-normalization-splitk-partials.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/norm-row-stats.wgsl.jinja CHANGED Viewed

@@ -1,8 +1,15 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
-{% set combineSubgroups = source.combineSubgroups %}
-{% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -22,16 +29,54 @@ enable subgroups;
 //
 // Shifted moments avoid cancellation from a large common offset; clamp the
 // variance to zero before adding EPSILON and taking inverseSqrt.
-const HIDDEN: u32 = {{ source.hidden }}u;
-{% if source.vec4 %}
-const HIDDEN_V: u32 = {{ source.hiddenVec }}u;
 {% endif %}
-const WG: u32 = {{ source.wg }}u;
-const EPSILON: f32 = {{ source.epsilon }};
-const CHANNELS: u32 = {{ source.channels }}u;
-{% if source.vec4 and scalarIo %}
 fn load_vec4(index: u32) -> vec4<f32> {
   return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
 }
@@ -94,13 +139,20 @@ fn main(
     return;
   }
   let tid = lid.x;
-{% if source.vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
-{% if source.vec4 %}
 {% if scalarIo %}
   let shift = f32(x[base]);
 {% else %}
@@ -111,9 +163,12 @@ fn main(
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if scalarIo %}
     let v = load_vec4(base + i * 4u);
 {% else %}
     let v = vec4<f32>(x[base + i]);
@@ -124,7 +179,12 @@ fn main(
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
     let d = v - shift;
     acc.x = acc.x + d;
     acc.y = acc.y + d * d;
@@ -141,9 +201,15 @@ fn main(
   let ch_scale = f32(scale[c]);
   let ch_bias = f32(bias[c]);
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if scalarIo %}
     let idx = base + i * 4u;
     let v = load_vec4(idx);
 {% else %}
@@ -157,14 +223,33 @@ fn main(
     y[idx + 2u] = value.z;
     y[idx + 3u] = value.w;
 {% else %}
-    y[idx] = {{ source.vecType }}((v - vec4<f32>(row_mean)) * inv * vec4<f32>(ch_scale) + vec4<f32>(ch_bias));
 {% endif %}
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
-    y[idx] = {{ source.scalar }}((v - row_mean) * inv * ch_scale + ch_bias);
   }
 {% endif %}
 }

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
+{% set combineSubgroups = combineSubgroups %}
+{% set scalarIo = scalarIo if scalarIo is defined else false %}
+{% set packedBf16Embedding = packedBf16Embedding if packedBf16Embedding is defined else false %}
+{% set rmsChainNorm = rmsChainNorm if rmsChainNorm is defined else false %}
+{% set hiddenPairs = hiddenPairs | default(0) %}
+{% set numRows = numRows | default(0) %}
+{% set epsilon = epsilon | default("0.0") %}
+{% set epsilon2 = epsilon2 | default("0.0") %}
+{% set channels = channels | default(0) %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 //
 // Shifted moments avoid cancellation from a large common offset; clamp the
 // variance to zero before adding EPSILON and taking inverseSqrt.
+const HIDDEN: u32 = {{ hidden }}u;
+{% if vec4 %}
+const HIDDEN_V: u32 = {{ hiddenVec }}u;
+{% endif %}
+{% if packedBf16Embedding %}
+const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
+const NUM_ROWS: u32 = {{ numRows }}u;
+{% endif %}
+const WG: u32 = {{ wg }}u;
+const EPSILON: f32 = {{ epsilon }};
+{% if rmsChainNorm %}
+const EPSILON2: f32 = {{ epsilon2 }};
+{% endif %}
+const CHANNELS: u32 = {{ channels }}u;
+{% if packedBf16Embedding %}
+{% if vec4 %}
+fn unpack_bf16_pair(word: u32) -> vec2<f32> {
+  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
+  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
+}
+{% endif %}
+{% if not vec4 %}
+fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
+  if (source_row >= NUM_ROWS) {
+    return 0.0;
+  }
+  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
+  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
+  return bitcast<f32>(bits << 16u);
+}
 {% endif %}
+{% if vec4 %}
+fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
+  if (source_row >= NUM_ROWS) {
+    return vec4<f32>(0.0);
+  }
+  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
+  let low = unpack_bf16_pair(x[base]);
+  let high = unpack_bf16_pair(x[base + 1u]);
+  return vec4<f32>(low, high);
+}
+{% endif %}
+{% endif %}
+{% if vec4 and scalarIo %}
 fn load_vec4(index: u32) -> vec4<f32> {
   return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
 }
     return;
   }
   let tid = lid.x;
+{% if packedBf16Embedding %}
+  let source_row = indices[row];
+{% if vec4 %}
+  let base = row * HIDDEN_V;
+{% else %}
+  let base = row * HIDDEN;
+{% endif %}
+{% elif vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
+{% if vec4 %}
 {% if scalarIo %}
   let shift = f32(x[base]);
 {% else %}
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_vec4(source_row, i);
+    embedding_out[base + i] = v;
+{% elif scalarIo %}
     let v = load_vec4(base + i * 4u);
 {% else %}
     let v = vec4<f32>(x[base + i]);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_scalar(source_row, i);
+    embedding_out[base + i] = v;
+{% else %}
     let v = f32(x[base + i]);
+{% endif %}
     let d = v - shift;
     acc.x = acc.x + d;
     acc.y = acc.y + d * d;
   let ch_scale = f32(scale[c]);
   let ch_bias = f32(bias[c]);
+{% if rmsChainNorm %}
+  var acc2 = 0.0;
+{% endif %}
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let idx = base + i;
+    let v = embedding_vec4(source_row, i);
+{% elif scalarIo %}
     let idx = base + i * 4u;
     let v = load_vec4(idx);
 {% else %}
     y[idx + 2u] = value.z;
     y[idx + 3u] = value.w;
 {% else %}
+    y[idx] = {{ vecType }}((v - vec4<f32>(row_mean)) * inv * vec4<f32>(ch_scale) + vec4<f32>(ch_bias));
 {% endif %}
   }
+{% if rmsChainNorm %}
+  // The chained second norm reads the residual row this loop just stored. This
+  // barrier completes those stores and any preceding shared-scratch use before
+  // the next reduction reuses its scratch; each lane then re-reads only the
+  // elements it wrote itself.
+  workgroupBarrier();
+  let total2 = reduce_scalar(acc2{{ reduceThreadArguments }});
+  let inv2 = inverseSqrt(total2 / f32(HIDDEN) + EPSILON2);
+  for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let idx = base + i;
+    let hv = vec4<f32>(y[idx]);
+    normed2[idx] = {{ vecType }}(hv * inv2 * vec4<f32>(scale2[i]));
+  }
+{% endif %}
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
+{% if packedBf16Embedding %}
+    let v = embedding_scalar(source_row, i);
+{% else %}
     let v = f32(x[idx]);
+{% endif %}
+    y[idx] = {{ scalar }}((v - row_mean) * inv * ch_scale + ch_bias);
   }
 {% endif %}
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.InstanceNormalization",
   "fixtureArrays": {
     "ort_rank3_batch2_repeated_epsilon_point3_input_input": [3.1513367, 9.283596, 1.4546119, 5.4617004, 8.519701, 1.2382338, 1.7930176, 5.1099434, 7.9195533, 7.638727, 8.065445, 3.8082376, 3.1513367, 9.283596, 1.4546119, 5.4617004, 8.519701, 1.2382338, 1.7930176, 5.1099434, 7.9195533, 7.638727, 8.065445, 3.8082376]
   },
@@ -65,7 +64,7 @@
     {
       "name": "f32_yfold_many_channels_1x70000x4",
       "provenance": {
-        "notes": "Compact sibling for the y-fold InstanceNormalization benchmark; preserves C>65535 with much smaller spatial work."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -183,7 +182,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/instance_norm_op_test.cc",
         "test": "InstanceNormalizationOpTest.InstanceNormBatch1",
-        "notes": "Vec4 spatial companion for valid subnormal scale outputs."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -670,7 +669,7 @@
     {
       "name": "split_stats_spatial_65536_f16",
       "provenance": {
-        "notes": "float16 on the split-statistics rungs, which were gated to float32 until the partials pass learned to widen its packed load and the apply pass to narrow its packed store. A spatial extent of SPLIT_STATS_MIN_SPATIAL that is also a multiple of four leaves both the vec4 and the scalar rung eligible, so one case renders both."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -700,6 +699,25 @@
         "b": { "dtype": "float16", "shape": [128], "data": { "kind": "constant", "value": 0.0 } }
       },
       "outputs": { "output": { "dtype": "float16", "shape": [2, 128, 64], "tolerance": 0.03 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_rank3_batch2_repeated_epsilon_point3_input_input": [3.1513367, 9.283596, 1.4546119, 5.4617004, 8.519701, 1.2382338, 1.7930176, 5.1099434, 7.9195533, 7.638727, 8.065445, 3.8082376, 3.1513367, 9.283596, 1.4546119, 5.4617004, 8.519701, 1.2382338, 1.7930176, 5.1099434, 7.9195533, 7.638727, 8.065445, 3.8082376]
   },
     {
       "name": "f32_yfold_many_channels_1x70000x4",
       "provenance": {
+        "notes": "More than 65,535 channels exercise folded Y dispatch with a compact four-element spatial plane."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/instance_norm_op_test.cc",
         "test": "InstanceNormalizationOpTest.InstanceNormBatch1",
+        "notes": "The vec4 spatial path must preserve valid outputs produced by a subnormal scale."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
     {
       "name": "split_stats_spatial_65536_f16",
       "provenance": {
+        "notes": "A float16 input with spatial extent 65,536 exercises split statistics, widened accumulation, and narrowed output storage. The extent is divisible by four, allowing both packed and scalar split-statistics variants."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
         "b": { "dtype": "float16", "shape": [128], "data": { "kind": "constant", "value": 0.0 } }
       },
       "outputs": { "output": { "dtype": "float16", "shape": [2, 128, 64], "tolerance": 0.03 } }
+    },
+    {
+      "name": "splitk_large_offset_unit_variance",
+      "provenance": {
+        "notes": "A plane centred on 8192 alternating by one unit: the true variance is 1, but raw second moments cancel it away at this magnitude, the combine's max(.,0) clamps it to zero, and every output becomes +/-1/sqrt(epsilon). The split-K partials centre on the plane's first element instead."
+      },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 256],
+          "data": { "kind": "cycle", "values": [8191.0, 8193.0] }
+        },
+        "scale": { "dtype": "float32", "shape": [1], "data": { "kind": "constant", "value": 1.0 } },
+        "b": { "dtype": "float32", "shape": [1], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [1, 1, 256, 256], "tolerance": 0.002, "relTolerance": 0.002 }
+      }
     }
   ]
 }