sync 91d990483a17

Browse files

Files changed (10) hide show

README.md +11 -7
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +111 -179
build/webgpu/mean-variance-normalization-serial-rows.wgsl.jinja +14 -15
build/webgpu/mean-variance-normalization-subgroup.wgsl.jinja +49 -46
build/webgpu/metadata.json +20 -11
build/webgpu/norm-flat-apply.wgsl.jinja +5 -6
build/webgpu/norm-flat-splitk-combine.wgsl.jinja +2 -2
build/webgpu/norm-flat-splitk-partials.wgsl.jinja +3 -3
build/webgpu/test.json +21 -5

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `MeanVarianceNormalization` spec](https://onnx.ai/onnx/operators/o
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Input tensor to normalize. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `y` | `T` | same as `X` | same as `X` | Normalized tensor with the same shape as `X`. | required |
 ## Attributes
@@ -44,7 +44,7 @@ Default values (overridable per request):
 ## 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
@@ -57,10 +57,14 @@ Default values (overridable per request):
 ## 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 |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Input tensor to normalize. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `T` | same as `x` | same as `x` | Normalized tensor with the same shape as `X`. | 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.MeanVarianceNormalization",
   "cases": [
     {
       "name": "1x32x32x32_default_axes",

 {
   "cases": [
     {
       "name": "1x32x32x32_default_axes",

build/webgpu/manifest.json CHANGED Viewed

@@ -2,44 +2,30 @@
   "domain": "ai.onnx",
   "name": "MeanVarianceNormalization",
   "sinceVersion": 13,
-  "description": "Normalizes each group as `(X - mean) / sqrt(variance)`, reducing over `axes` (default `[0, 2, 3]`).",
-  "inputs": [{ "role": "X", "dtype": "T", "description": "Input tensor to normalize." }],
-  "outputs": [
-    {
-      "role": "Y",
-      "dtype": "T",
-      "rank": "ranks.X",
-      "description": "Normalized tensor with the same shape as `X`.",
-      "shape": "shapes.X"
-    }
-  ],
-  "attributes": { "axes": [0, 2, 3] },
-  "attributeDescriptions": { "axes": "Axes that share a mean and variance; negative values count from the back." },
   "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "Y", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "SERIAL_WORKGROUP_SIZE": 256,
-    "SERIAL_TINY_WORKGROUP_SIZE": 64,
-    "SERIAL_MAX_REDUCTION": 128,
-    "SERIAL_MIN_ROWS": 256,
-    "TREE_MEDIUM_WORKGROUP_SIZE": 64,
-    "VEC4_MIN_REDUCTION": 8,
-    "FLAT_SPLIT_MIN_ELEMENTS": 65536,
-    "FLAT_SPLIT_TARGET_ELEMENTS": 4096,
-    "MAX_FLAT_SPLITS": 256
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
-    "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
-    "shapeContract": "ranks.X >= 3 and ranks.Y == ranks.X and sameShape(shapes.Y, shapes.X) and f16Ok(dtypes.T)",
-    "reduceCount": "(dim(shapes.X, 0) if hasAxis(attrs.axes, 0, ranks.X) else 1) * (dim(shapes.X, 1) if hasAxis(attrs.axes, 1, ranks.X) else 1) * (dim(shapes.X, 2) if hasAxis(attrs.axes, 2, ranks.X) else 1) * (dim(shapes.X, 3) if ranks.X >= 4 and hasAxis(attrs.axes, 3, ranks.X) else 1) * (dim(shapes.X, 4) if ranks.X >= 5 and hasAxis(attrs.axes, 4, ranks.X) else 1) * (dim(shapes.X, 5) if ranks.X >= 6 and hasAxis(attrs.axes, 5, ranks.X) else 1)",
-    "rowCount": "numel(shapes.X) / max(1, reduceCount)",
-    "allAxesReduced": "hasAxis(attrs.axes, 0, ranks.X) and hasAxis(attrs.axes, 1, ranks.X) and hasAxis(attrs.axes, 2, ranks.X) and (ranks.X < 4 or hasAxis(attrs.axes, 3, ranks.X)) and (ranks.X < 5 or hasAxis(attrs.axes, 4, ranks.X)) and (ranks.X < 6 or hasAxis(attrs.axes, 5, ranks.X))",
-    "vec4Eligible": "((ranks.X == 3 and hasAxis(attrs.axes, 2, 3) and (dim(shapes.X, 2) % 4 == 0 or (hasAxis(attrs.axes, 1, 3) and dim(shapes.X, 1) * dim(shapes.X, 2) % 4 == 0) or (hasAxis(attrs.axes, 0, 3) and hasAxis(attrs.axes, 1, 3) and numel(shapes.X) % 4 == 0))) or (ranks.X == 4 and hasAxis(attrs.axes, 3, 4) and (dim(shapes.X, 3) % 4 == 0 or (hasAxis(attrs.axes, 2, 4) and dim(shapes.X, 2) * dim(shapes.X, 3) % 4 == 0) or (hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and dim(shapes.X, 1) * dim(shapes.X, 2) * dim(shapes.X, 3) % 4 == 0) or (hasAxis(attrs.axes, 0, 4) and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and numel(shapes.X) % 4 == 0))) or (ranks.X == 5 and hasAxis(attrs.axes, 4, 5) and (dim(shapes.X, 4) % 4 == 0 or (hasAxis(attrs.axes, 3, 5) and dim(shapes.X, 3) * dim(shapes.X, 4) % 4 == 0) or (hasAxis(attrs.axes, 2, 5) and hasAxis(attrs.axes, 3, 5) and dim(shapes.X, 2) * dim(shapes.X, 3) * dim(shapes.X, 4) % 4 == 0) or (hasAxis(attrs.axes, 1, 5) and hasAxis(attrs.axes, 2, 5) and hasAxis(attrs.axes, 3, 5) and dim(shapes.X, 1) * dim(shapes.X, 2) * dim(shapes.X, 3) * dim(shapes.X, 4) % 4 == 0) or (allAxesReduced and numel(shapes.X) % 4 == 0))) or (ranks.X == 6 and hasAxis(attrs.axes, 5, 6) and (dim(shapes.X, 5) % 4 == 0 or (hasAxis(attrs.axes, 4, 6) and dim(shapes.X, 4) * dim(shapes.X, 5) % 4 == 0) or (hasAxis(attrs.axes, 3, 6) and hasAxis(attrs.axes, 4, 6) and dim(shapes.X, 3) * dim(shapes.X, 4) * dim(shapes.X, 5) % 4 == 0) or (hasAxis(attrs.axes, 2, 6) and hasAxis(attrs.axes, 3, 6) and hasAxis(attrs.axes, 4, 6) and dim(shapes.X, 2) * dim(shapes.X, 3) * dim(shapes.X, 4) * dim(shapes.X, 5) % 4 == 0) or (hasAxis(attrs.axes, 1, 6) and hasAxis(attrs.axes, 2, 6) and hasAxis(attrs.axes, 3, 6) and hasAxis(attrs.axes, 4, 6) and dim(shapes.X, 1) * dim(shapes.X, 2) * dim(shapes.X, 3) * dim(shapes.X, 4) * dim(shapes.X, 5) % 4 == 0) or (allAxesReduced and numel(shapes.X) % 4 == 0))))",
     "maxWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "minSubgroupSize": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 1",
     "useSubgroups": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and minSubgroupSize > 0",
@@ -48,93 +34,36 @@
     "serialWorkgroupSize": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, (tunables.SERIAL_TINY_WORKGROUP_SIZE if reduceCount <= 4 else tunables.SERIAL_WORKGROUP_SIZE))",
     "rowDispatchFits": "rowCount <= foldedDispatchCapacity",
     "serialDispatchFits": "ceilDiv(rowCount, serialWorkgroupSize) <= foldedDispatchCapacity",
-    "applyDispatchFits": "ceilDiv(numel(shapes.Y), maxWorkgroupSize) <= foldedDispatchCapacity",
     "scalarStorageFits": "scalarWorkgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "vectorStorageFits": "vectorWorkgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
-    "flatSplit": "min(tunables.MAX_FLAT_SPLITS, pow2ceil(ceilDiv(numel(shapes.X), tunables.FLAT_SPLIT_TARGET_ELEMENTS)))",
     "flatScratchBytes": "flatSplit * 8",
-    "flatPathFits": "flatSplit <= device.limits.maxComputeWorkgroupsPerDimension and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize and maxWorkgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize and applyDispatchFits"
   },
-  "bindingSets": {
-    "rows": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "X",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "rowCount" }] }
-      }
-    ],
-    "flatPartials": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "vec2<f32>" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
-      }
-    ],
-    "flatCombine": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      {
-        "name": "partials",
-        "semantic": "partials",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "vec2<f32>"
-      },
-      { "name": "stats", "semantic": "stats", "buffer": { "type": "storage" }, "elementType": "f32", "length": 2 },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
-      }
-    ],
-    "flatApply": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      {
-        "name": "stats",
-        "semantic": "stats",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32",
-        "length": 2
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }] }
-      }
-    ],
-    "noopParams": [
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "dummy", "type": "u32", "value": 0 }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "empty_noop",
       "priority": 200,
-      "when": ["shapeContract", "numel(shapes.X) == 0"],
       "passes": [
         {
           "id": "noop",
           "name": "MeanVarianceNormalization.Empty",
           "shader": "noop.wgsl.jinja",
-          "bindings": "noopParams",
           "dispatch": { "x": 0 }
         }
       ]
@@ -142,8 +71,8 @@
     {
       "id": "all_axes_flat_split",
       "priority": 120,
-      "when": ["shapeContract", "numel(shapes.X) > 0", "allAxesReduced", "numel(shapes.X) >= tunables.FLAT_SPLIT_MIN_ELEMENTS", "flatPathFits"],
-      "constants": { "scalar": "dtypes.T" },
       "intermediates": [
         { "id": "partials", "dtype": "float32", "shape": "[flatSplit, 2]" },
         { "id": "stats", "dtype": "float32", "shape": "[2]" }
@@ -152,120 +81,123 @@
         {
           "id": "partials",
           "name": "MeanVarianceNormalization.FlatPartials",
-          "source": {
-            "shader": "norm-flat-splitk-partials.wgsl.jinja",
-            "inputs": { "workgroupSize": "maxWorkgroupSize", "split": "flatSplit", "usesF16": "dtypes.T == \"f16\"" }
           },
-          "bindings": "flatPartials",
-          "dispatch": { "workgroups": "flatSplit" }
         },
         {
           "id": "combine",
           "name": "MeanVarianceNormalization.FlatCombine",
-          "source": {
-            "shader": "norm-flat-splitk-combine.wgsl.jinja",
-            "inputs": { "split": "flatSplit", "usesF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "flatCombine",
-          "dispatch": { "workgroups": 1 }
         },
         {
           "id": "apply",
           "name": "MeanVarianceNormalization.FlatApply",
-          "source": {
-            "shader": "norm-flat-apply.wgsl.jinja",
-            "inputs": { "workgroupSize": "maxWorkgroupSize", "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" }
           },
-          "bindings": "flatApply",
-          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "maxWorkgroupSize" }
         }
       ]
     },
     {
       "id": "serial_rows",
       "priority": 115,
-      "when": ["shapeContract", "numel(shapes.X) > 0", "reduceCount <= tunables.SERIAL_MAX_REDUCTION", "rowCount >= tunables.SERIAL_MIN_ROWS", "serialDispatchFits"],
-      "constants": { "scalar": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "MeanVarianceNormalization.SerialRows",
-          "source": {
-            "shader": "mean-variance-normalization-serial-rows.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "reduce": ["hasAxis(attrs.axes, 0, ranks.X)", "hasAxis(attrs.axes, 1, ranks.X)", "hasAxis(attrs.axes, 2, ranks.X)", "hasAxis(attrs.axes, 3, ranks.X)", "hasAxis(attrs.axes, 4, ranks.X)", "hasAxis(attrs.axes, 5, ranks.X)"],
-              "reduceCount": "reduceCount",
-              "workgroupSize": "serialWorkgroupSize",
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "threads": "rowCount", "workgroupSize": "serialWorkgroupSize" }
         }
       ]
     },
     {
       "id": "cooperative_vec4",
       "priority": 110,
-      "when": ["shapeContract", "numel(shapes.X) > 0", "reduceCount >= tunables.VEC4_MIN_REDUCTION", "vec4Eligible", "rowDispatchFits", "vectorStorageFits"],
-      "constants": {
-        "scalar": "dtypes.T",
-        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
-      },
       "passes": [
         {
           "id": "main",
           "name": "MeanVarianceNormalization.CooperativeVec4",
-          "source": {
-            "shader": "mean-variance-normalization-subgroup.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "reduce": ["hasAxis(attrs.axes, 0, ranks.X)", "hasAxis(attrs.axes, 1, ranks.X)", "hasAxis(attrs.axes, 2, ranks.X)", "hasAxis(attrs.axes, 3, ranks.X)", "hasAxis(attrs.axes, 4, ranks.X)", "hasAxis(attrs.axes, 5, ranks.X)"],
-              "reduceCount": "reduceCount",
-              "wg": "vectorWorkgroupSize",
-              "minSubgroupSize": "minSubgroupSize",
-              "maxSubgroups": "ceilDiv(vectorWorkgroupSize, minSubgroupSize)",
-              "scalar": "dtypes.T",
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "usesF16": "dtypes.T == \"f16\"",
-              "useSubgroups": "useSubgroups",
-              "vectorized": true
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "cooperative_scalar",
       "priority": 100,
-      "when": ["shapeContract", "numel(shapes.X) > 0", "rowDispatchFits", "scalarStorageFits"],
-      "constants": { "scalar": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "MeanVarianceNormalization.CooperativeScalar",
-          "source": {
-            "shader": "mean-variance-normalization-subgroup.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "reduce": ["hasAxis(attrs.axes, 0, ranks.X)", "hasAxis(attrs.axes, 1, ranks.X)", "hasAxis(attrs.axes, 2, ranks.X)", "hasAxis(attrs.axes, 3, ranks.X)", "hasAxis(attrs.axes, 4, ranks.X)", "hasAxis(attrs.axes, 5, ranks.X)"],
-              "reduceCount": "reduceCount",
-              "wg": "scalarWorkgroupSize",
-              "minSubgroupSize": "minSubgroupSize",
-              "maxSubgroups": "ceilDiv(scalarWorkgroupSize, minSubgroupSize)",
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "useSubgroups": "useSubgroups",
-              "vectorized": false
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "MeanVarianceNormalization",
   "sinceVersion": 13,
+  "inputs": { "x": { "onnx": "X", "dtype": "T" } },
+  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
+  "attributes": { "axes": { "default": [0, 2, 3] } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "SERIAL_WORKGROUP_SIZE": { "default": 256 },
+    "SERIAL_TINY_WORKGROUP_SIZE": { "default": 64 },
+    "SERIAL_MAX_REDUCTION": { "default": 128 },
+    "SERIAL_MIN_ROWS": { "default": 256 },
+    "TREE_MEDIUM_WORKGROUP_SIZE": { "default": 64 },
+    "VEC4_MIN_REDUCTION": { "default": 8 },
+    "FLAT_SPLIT_MIN_ELEMENTS": { "default": 65536 },
+    "FLAT_SPLIT_TARGET_ELEMENTS": { "default": 4096 },
+    "MAX_FLAT_SPLITS": { "default": 256 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "shapeContract": "ranks.x >= 3 and ranks.x <= 8 and ranks.y == ranks.x and sameShape(shapes.y, shapes.x) and f16Ok(dtypes.T)",
+    "reduceCount": "(dim(shapes.x, 0) if hasAxis(attrs.axes, 0, ranks.x) else 1) * (dim(shapes.x, 1) if hasAxis(attrs.axes, 1, ranks.x) else 1) * (dim(shapes.x, 2) if hasAxis(attrs.axes, 2, ranks.x) else 1) * (dim(shapes.x, 3) if ranks.x >= 4 and hasAxis(attrs.axes, 3, ranks.x) else 1) * (dim(shapes.x, 4) if ranks.x >= 5 and hasAxis(attrs.axes, 4, ranks.x) else 1) * (dim(shapes.x, 5) if ranks.x >= 6 and hasAxis(attrs.axes, 5, ranks.x) else 1) * (dim(shapes.x, 6) if ranks.x >= 7 and hasAxis(attrs.axes, 6, ranks.x) else 1) * (dim(shapes.x, 7) if ranks.x >= 8 and hasAxis(attrs.axes, 7, ranks.x) else 1)",
+    "rowCount": "numel(shapes.x) / max(1, reduceCount)",
+    "allAxesReduced": "hasAxis(attrs.axes, 0, ranks.x) and hasAxis(attrs.axes, 1, ranks.x) and hasAxis(attrs.axes, 2, ranks.x) and (ranks.x < 4 or hasAxis(attrs.axes, 3, ranks.x)) and (ranks.x < 5 or hasAxis(attrs.axes, 4, ranks.x)) and (ranks.x < 6 or hasAxis(attrs.axes, 5, ranks.x)) and (ranks.x < 7 or hasAxis(attrs.axes, 6, ranks.x)) and (ranks.x < 8 or hasAxis(attrs.axes, 7, ranks.x))",
+    "vec4Eligible": "((ranks.x == 3 and hasAxis(attrs.axes, 2, 3) and (dim(shapes.x, 2) % 4 == 0 or (hasAxis(attrs.axes, 1, 3) and dim(shapes.x, 1) * dim(shapes.x, 2) % 4 == 0) or (hasAxis(attrs.axes, 0, 3) and hasAxis(attrs.axes, 1, 3) and numel(shapes.x) % 4 == 0))) or (ranks.x == 4 and hasAxis(attrs.axes, 3, 4) and (dim(shapes.x, 3) % 4 == 0 or (hasAxis(attrs.axes, 2, 4) and dim(shapes.x, 2) * dim(shapes.x, 3) % 4 == 0) or (hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) % 4 == 0) or (hasAxis(attrs.axes, 0, 4) and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and numel(shapes.x) % 4 == 0))) or (ranks.x == 5 and hasAxis(attrs.axes, 4, 5) and (dim(shapes.x, 4) % 4 == 0 or (hasAxis(attrs.axes, 3, 5) and dim(shapes.x, 3) * dim(shapes.x, 4) % 4 == 0) or (hasAxis(attrs.axes, 2, 5) and hasAxis(attrs.axes, 3, 5) and dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4) % 4 == 0) or (hasAxis(attrs.axes, 1, 5) and hasAxis(attrs.axes, 2, 5) and hasAxis(attrs.axes, 3, 5) and dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4) % 4 == 0) or (allAxesReduced and numel(shapes.x) % 4 == 0))) or (ranks.x == 6 and hasAxis(attrs.axes, 5, 6) and (dim(shapes.x, 5) % 4 == 0 or (hasAxis(attrs.axes, 4, 6) and dim(shapes.x, 4) * dim(shapes.x, 5) % 4 == 0) or (hasAxis(attrs.axes, 3, 6) and hasAxis(attrs.axes, 4, 6) and dim(shapes.x, 3) * dim(shapes.x, 4) * dim(shapes.x, 5) % 4 == 0) or (hasAxis(attrs.axes, 2, 6) and hasAxis(attrs.axes, 3, 6) and hasAxis(attrs.axes, 4, 6) and dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4) * dim(shapes.x, 5) % 4 == 0) or (hasAxis(attrs.axes, 1, 6) and hasAxis(attrs.axes, 2, 6) and hasAxis(attrs.axes, 3, 6) and hasAxis(attrs.axes, 4, 6) and dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4) * dim(shapes.x, 5) % 4 == 0) or (allAxesReduced and numel(shapes.x) % 4 == 0))))",
     "maxWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "minSubgroupSize": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 1",
     "useSubgroups": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and minSubgroupSize > 0",
     "serialWorkgroupSize": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, (tunables.SERIAL_TINY_WORKGROUP_SIZE if reduceCount <= 4 else tunables.SERIAL_WORKGROUP_SIZE))",
     "rowDispatchFits": "rowCount <= foldedDispatchCapacity",
     "serialDispatchFits": "ceilDiv(rowCount, serialWorkgroupSize) <= foldedDispatchCapacity",
+    "applyDispatchFits": "ceilDiv(numel(shapes.y), maxWorkgroupSize) <= foldedDispatchCapacity",
     "scalarStorageFits": "scalarWorkgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "vectorStorageFits": "vectorWorkgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
+    "flatSplit": "min(tunables.MAX_FLAT_SPLITS, pow2ceil(ceilDiv(numel(shapes.x), tunables.FLAT_SPLIT_TARGET_ELEMENTS)))",
     "flatScratchBytes": "flatSplit * 8",
+    "flatPathFits": "flatSplit <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize and maxWorkgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize and applyDispatchFits"
   },
+  "when": ["shapeContract"],
+  "bindings": {
+    "x": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.x)" }] },
+    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$ioElement" },
+    "y_2": { "name": "y", "buffer": "storage", "elementType": "$ioElement" },
+    "params_4": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "rows", "type": "u32", "value": "rowCount" }]
+    }
   },
   "variants": [
     {
       "id": "empty_noop",
       "priority": 200,
+      "when": ["numel(shapes.x) == 0"],
       "passes": [
         {
           "id": "noop",
           "name": "MeanVarianceNormalization.Empty",
           "shader": "noop.wgsl.jinja",
+          "bindings": [{ "name": "params", "struct": [{ "name": "dummy", "type": "u32", "value": 0 }] }],
           "dispatch": { "x": 0 }
         }
       ]
     {
       "id": "all_axes_flat_split",
       "priority": 120,
+      "when": ["numel(shapes.x) > 0", "allAxesReduced", "numel(shapes.x) >= tunables.FLAT_SPLIT_MIN_ELEMENTS", "flatPathFits"],
+      "derive": { "scalar": "dtypes.T" },
       "intermediates": [
         { "id": "partials", "dtype": "float32", "shape": "[flatSplit, 2]" },
         { "id": "stats", "dtype": "float32", "shape": "[2]" }
         {
           "id": "partials",
           "name": "MeanVarianceNormalization.FlatPartials",
+          "shader": "norm-flat-splitk-partials.wgsl.jinja",
+          "derive": {
+            "workgroupSizeSpec": "maxWorkgroupSize",
+            "splitSpec": "flatSplit",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "vec2<f32>" }, "params"],
+          "dispatch": { "x": "min(flatSplit, 65535)", "y": "ceilDiv(flatSplit, 65535)", "z": 1 }
         },
         {
           "id": "combine",
           "name": "MeanVarianceNormalization.FlatCombine",
+          "shader": "norm-flat-splitk-combine.wgsl.jinja",
+          "derive": { "splitSpec": "flatSplit", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": [
+            "x",
+            { "name": "partials", "buffer": "read-only-storage", "elementType": "vec2<f32>" },
+            { "name": "stats", "buffer": "storage", "elementType": "f32", "length": 2 },
+            "params"
+          ],
+          "dispatch": { "x": "min(1, 65535)", "y": "ceilDiv(1, 65535)", "z": 1 }
         },
         {
           "id": "apply",
           "name": "MeanVarianceNormalization.FlatApply",
+          "shader": "norm-flat-apply.wgsl.jinja",
+          "derive": {
+            "workgroupSizeSpec": "maxWorkgroupSize",
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": [
+            "x",
+            { "name": "stats", "buffer": "read-only-storage", "elementType": "f32", "length": 2 },
+            { "arg": "y", "elementType": "$scalar" },
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (maxWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (maxWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "serial_rows",
       "priority": 115,
+      "when": ["numel(shapes.x) > 0", "reduceCount <= tunables.SERIAL_MAX_REDUCTION", "rowCount >= tunables.SERIAL_MIN_ROWS", "serialDispatchFits"],
+      "derive": { "scalar": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "MeanVarianceNormalization.SerialRows",
+          "shader": "mean-variance-normalization-serial-rows.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "reduce": ["hasAxis(attrs.axes, 0, ranks.x)", "hasAxis(attrs.axes, 1, ranks.x)", "hasAxis(attrs.axes, 2, ranks.x)", "hasAxis(attrs.axes, 3, ranks.x)", "hasAxis(attrs.axes, 4, ranks.x)", "hasAxis(attrs.axes, 5, ranks.x)", "hasAxis(attrs.axes, 6, ranks.x)", "hasAxis(attrs.axes, 7, ranks.x)"],
+            "workgroupSizeSpec": "serialWorkgroupSize",
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y_2", "params_4"],
+          "dispatch": {
+            "x": "min(ceilDiv((rowCount), (serialWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((rowCount), (serialWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "cooperative_vec4",
       "priority": 110,
+      "when": ["numel(shapes.x) > 0", "reduceCount >= tunables.VEC4_MIN_REDUCTION", "vec4Eligible", "rowDispatchFits", "vectorStorageFits"],
+      "derive": { "scalar": "dtypes.T", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "MeanVarianceNormalization.CooperativeVec4",
+          "shader": "mean-variance-normalization-subgroup.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": {
+            "xShape": "shapes.x",
+            "reduce": ["hasAxis(attrs.axes, 0, ranks.x)", "hasAxis(attrs.axes, 1, ranks.x)", "hasAxis(attrs.axes, 2, ranks.x)", "hasAxis(attrs.axes, 3, ranks.x)", "hasAxis(attrs.axes, 4, ranks.x)", "hasAxis(attrs.axes, 5, ranks.x)", "hasAxis(attrs.axes, 6, ranks.x)", "hasAxis(attrs.axes, 7, ranks.x)"],
+            "wg": "vectorWorkgroupSize",
+            "scalar": "dtypes.T",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "vectorizedSpec": true
           },
+          "bindings": ["x_2", "y_2", "params_4"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "cooperative_scalar",
       "priority": 100,
+      "when": ["numel(shapes.x) > 0", "rowDispatchFits", "scalarStorageFits"],
+      "derive": { "scalar": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "MeanVarianceNormalization.CooperativeScalar",
+          "shader": "mean-variance-normalization-subgroup.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": {
+            "xShape": "shapes.x",
+            "reduce": ["hasAxis(attrs.axes, 0, ranks.x)", "hasAxis(attrs.axes, 1, ranks.x)", "hasAxis(attrs.axes, 2, ranks.x)", "hasAxis(attrs.axes, 3, ranks.x)", "hasAxis(attrs.axes, 4, ranks.x)", "hasAxis(attrs.axes, 5, ranks.x)", "hasAxis(attrs.axes, 6, ranks.x)", "hasAxis(attrs.axes, 7, ranks.x)"],
+            "wg": "scalarWorkgroupSize",
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "vectorizedSpec": false
           },
+          "bindings": ["x_2", "y_2", "params_4"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     }

build/webgpu/mean-variance-normalization-serial-rows.wgsl.jinja CHANGED Viewed

@@ -1,4 +1,4 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
@@ -10,42 +10,41 @@ enable f16;
 // The reduction is shifted by its first value before accumulating moments. This
 // stabilizes the variance calculation while retaining the zero-variance ONNX
 // result as 0/0 = NaN.
-{% macro nd_offset(target, index, reduced, indent) %}
-{% set rank = source.xShape | length %}
 {% for i in range(rank) %}
-{% if source.reduce[i] if reduced else not source.reduce[i] %}
 {% set pa = namespace(v=1) %}
 {% set st = namespace(v=1) %}
 {% for j in range(i + 1, rank) %}
-{% set st.v = st.v * source.xShape[j] %}
-{% if source.reduce[j] if reduced else not source.reduce[j] %}{% set pa.v = pa.v * source.xShape[j] %}{% endif %}
 {% endfor %}
-{{ indent }}{{ target }} = {{ target }} + (({{ index }} / {{ pa.v }}u) % {{ source.xShape[i] }}u) * {{ st.v }}u;
 {% endif %}
 {% endfor %}
 {% endmacro %}
-const WG: u32 = {{ source.workgroupSize }}u;
-const R: u32 = {{ source.reduceCount }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   // `threads` dispatches fold past the WebGPU x-dimension limit into y.
-  let row = gid.x + gid.y * nwg.x * WG;
   if (row >= params.rows) {
     return;
   }
   var base_off = 0u;
-{{ nd_offset("base_off", "row", 0, "  ") }}
   let shift = f32(x[base_off]);
   var sum_d = 0.0;
   var sum_d2 = 0.0;
   for (var r = 0u; r < R; r = r + 1u) {
     var off = base_off;
-{{ nd_offset("off", "r", 1, "    ") }}    let d = f32(x[off]) - shift;
     sum_d = sum_d + d;
     sum_d2 = sum_d2 + d * d;
   }
@@ -57,6 +56,6 @@ fn main(
   for (var r = 0u; r < R; r = r + 1u) {
     var off = base_off;
-{{ nd_offset("off", "r", 1, "    ") }}    y[off] = {{ source.scalar }}((f32(x[off]) - mean) / denom);
   }
 }

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // The reduction is shifted by its first value before accumulating moments. This
 // stabilizes the variance calculation while retaining the zero-variance ONNX
 // result as 0/0 = NaN.
+{% macro nd_offset(target, index, reduced) %}
+{% set rank = xShape | length %}
 {% for i in range(rank) %}
+{% if reduce[i] if reduced else not reduce[i] %}
 {% set pa = namespace(v=1) %}
 {% set st = namespace(v=1) %}
 {% for j in range(i + 1, rank) %}
+{% set st.v = st.v * xShape[j] %}
+{% if reduce[j] if reduced else not reduce[j] %}{% set pa.v = pa.v * xShape[j] %}{% endif %}
 {% endfor %}
+{{ target }} = {{ target }} + (({{ index }} / {{ pa.v }}u) % {{ xShape[i] }}u) * {{ st.v }}u;
 {% endif %}
 {% endfor %}
 {% endmacro %}
+const WG: u32 = {{ workgroupSizeSpec }}u;
+const R: u32 = {{ reduceCount }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>
 ) {
   // `threads` dispatches fold past the WebGPU x-dimension limit into y.
+  let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (row >= params.rows) {
     return;
   }
   var base_off = 0u;
+{{ nd_offset("base_off", "row", 0) | indent(2, true) }}
   let shift = f32(x[base_off]);
   var sum_d = 0.0;
   var sum_d2 = 0.0;
   for (var r = 0u; r < R; r = r + 1u) {
     var off = base_off;
+{{ nd_offset("off", "r", 1) | indent(4, true) }}    let d = f32(x[off]) - shift;
     sum_d = sum_d + d;
     sum_d2 = sum_d2 + d * d;
   }
   for (var r = 0u; r < R; r = r + 1u) {
     var off = base_off;
+{{ nd_offset("off", "r", 1) | indent(4, true) }}    y[off] = {{ scalar }}((f32(x[off]) - mean) / denom);
   }
 }

build/webgpu/mean-variance-normalization-subgroup.wgsl.jinja CHANGED Viewed

@@ -1,58 +1,67 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
-{% if source.useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-{% macro nd_offset(target, index, reduced, indent) %}
-{% set rank = source.xShape | length %}
 {% for i in range(rank) %}
-{% if source.reduce[i] if reduced else not source.reduce[i] %}
 {% set pa = namespace(v=1) %}
 {% set st = namespace(v=1) %}
 {% for j in range(i + 1, rank) %}
-{% set st.v = st.v * source.xShape[j] %}
-{% if source.reduce[j] if reduced else not source.reduce[j] %}{% set pa.v = pa.v * source.xShape[j] %}{% endif %}
 {% endfor %}
-{{ indent }}{{ target }} = {{ target }} + (({{ index }} / {{ pa.v }}u) % {{ source.xShape[i] }}u) * {{ st.v }}u;
 {% endif %}
 {% endfor %}
 {% endmacro %}
-const WG: u32 = {{ source.wg }}u;
-const R: u32 = {{ source.reduceCount }}u;
-{% if source.vectorized %}
 const RV: u32 = R / 4u;
 {% endif %}
-{% if source.useSubgroups %}
-{% if source.wg > source.minSubgroupSize %}
-// The device-reported minimum subgroup width bounds cross-subgroup partials.
-const MAX_SG: u32 = {{ source.maxSubgroups }}u;
-var<workgroup> sg_partials: array<vec2<f32>, MAX_SG>;
-{% endif %}
 {% else %}
 var<workgroup> wg_red: array<vec2<f32>, WG>;
 {% endif %}
-fn reduce_pair(value: vec2<f32>{% if not source.useSubgroups or source.wg > source.minSubgroupSize %}, tid: u32{% endif %}{% if source.useSubgroups and source.wg > source.minSubgroupSize %}, sg_lane: u32, sg_size: u32{% endif %}) -> vec2<f32> {
-{% if source.useSubgroups %}
   let s = vec2<f32>(subgroupAdd(value.x), subgroupAdd(value.y));
-{% if source.wg > source.minSubgroupSize %}
-  if (sg_lane == 0u) {
-    sg_partials[tid / sg_size] = s;
   }
   workgroupBarrier();
-  let num_sg = (WG + sg_size - 1u) / sg_size;
   var total = vec2<f32>(0.0, 0.0);
-  for (var i = 0u; i < num_sg; i = i + 1u) {
     total = total + sg_partials[i];
   }
-  return total;
-{% else %}
-  return s;
-{% endif %}
 {% else %}
   wg_red[tid] = value;
   workgroupBarrier();
@@ -69,14 +78,12 @@ fn reduce_pair(value: vec2<f32>{% if not source.useSubgroups or source.wg > sour
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg_id: vec3<u32>,
-  @builtin(local_invocation_id) lid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>
-{%- if source.useSubgroups and source.wg > source.minSubgroupSize %},
-  @builtin(subgroup_invocation_id) sg_lane: u32,
   @builtin(subgroup_size) sg_size: u32
 {%- endif %}
 ) {
-  let row = wg_id.x + wg_id.y * nwg.x;
   if (row >= params.rows) {
     return;
   }
@@ -84,21 +91,21 @@ fn main(
   // Base offset from the kept-axis coordinates.
   var base_off = 0u;
-{{ nd_offset("base_off", "row", 0, "  ") }}
-{% if source.vectorized %}
   let shift = f32(x[base_off / 4u].x);
 {% else %}
   let shift = f32(x[base_off]);
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
-{% if source.vectorized %}
   for (var q = tid; q < RV; q = q + WG) {
     let r = q * 4u;
 {% else %}
   for (var r = tid; r < R; r = r + WG) {
 {% endif %}
     var off = base_off;
-{{ nd_offset("off", "r", 1, "    ") }}{% if source.vectorized %}
     let d = vec4<f32>(x[off / 4u]) - vec4<f32>(shift);
     acc.x = acc.x + d.x + d.y + d.z + d.w;
     acc.y = acc.y + dot(d, d);
@@ -109,28 +116,24 @@ fn main(
 {% endif %}
   }
-{% if source.useSubgroups %}
-  let totals = reduce_pair(acc{% if source.wg > source.minSubgroupSize %}, tid, sg_lane, sg_size{% endif %});
-{% else %}
-  let totals = reduce_pair(acc, tid);
-{% endif %}
   let mean_d = totals.x / f32(R);
   let variance = max(totals.y / f32(R) - mean_d * mean_d, 0.0);
   let mean = shift + mean_d;
   let denom = sqrt(variance);
-{% if source.vectorized %}
   for (var q = tid; q < RV; q = q + WG) {
     let r = q * 4u;
 {% else %}
   for (var r = tid; r < R; r = r + WG) {
 {% endif %}
     var off = base_off;
-{{ nd_offset("off", "r", 1, "    ") }}{% if source.vectorized %}
     let v = vec4<f32>(x[off / 4u]);
-    y[off / 4u] = {{ source.vecType }}((v - vec4<f32>(mean)) / vec4<f32>(denom));
 {% else %}
-    y[off] = {{ source.scalar }}((f32(x[off]) - mean) / denom);
 {% endif %}
   }
 }

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
+{% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro nd_offset(target, index, reduced) %}
+{% set rank = xShape | length %}
 {% for i in range(rank) %}
+{% if reduce[i] if reduced else not reduce[i] %}
 {% set pa = namespace(v=1) %}
 {% set st = namespace(v=1) %}
 {% for j in range(i + 1, rank) %}
+{% set st.v = st.v * xShape[j] %}
+{% if reduce[j] if reduced else not reduce[j] %}{% set pa.v = pa.v * xShape[j] %}{% endif %}
 {% endfor %}
+{{ target }} = {{ target }} + (({{ index }} / {{ pa.v }}u) % {{ xShape[i] }}u) * {{ st.v }}u;
 {% endif %}
 {% endfor %}
 {% endmacro %}
+const WG: u32 = {{ wg }}u;
+const R: u32 = {{ reduceCount }}u;
+{% if vectorizedSpec %}
 const RV: u32 = R / 4u;
 {% endif %}
+{% if useSubgroups %}
+// One slot per invocation: the cross-subgroup fold indexes by
+// local_invocation_id, never by a subgroup ordinal, so no subgroup width or
+// partition can overflow or alias it.
+var<workgroup> sg_partials: array<vec2<f32>, WG>;
 {% else %}
 var<workgroup> wg_red: array<vec2<f32>, WG>;
 {% endif %}
+fn reduce_pair(value: vec2<f32>, tid: u32{% if useSubgroups %}, sg_size: u32{% endif %}) -> vec2<f32> {
+{% if useSubgroups %}
   let s = vec2<f32>(subgroupAdd(value.x), subgroupAdd(value.y));
+  // The whole workgroup is one subgroup: the collective above already covers it
+  // (no shared memory, no barrier). The test reads the `subgroup_size` builtin,
+  // which is uniform; a collective's result is not uniform to WGSL's analysis
+  // and may not guard a barrier.
+  if (sg_size == WG) {
+    return s;
   }
+  // Cross-subgroup fold that assumes nothing about which invocations share a
+  // subgroup, how many subgroups there are, or which of a subgroup's lanes are
+  // active: every invocation owns the slot at its own index, the elected lane
+  // publishes its subgroup pair there and every other lane publishes the sum
+  // identity. Each subgroup then folds all WG slots — lane `rank`, its dense
+  // position among the active lanes, walks slots rank, rank + count, ... — and
+  // one more collective merges the lane partials, so every slot is added
+  // exactly once at any legal width and partition.
+  let rank = subgroupExclusiveAdd(1u);
+  let count = subgroupAdd(1u);
+  sg_partials[tid] = select(vec2<f32>(0.0, 0.0), s, rank == 0u);
   workgroupBarrier();
   var total = vec2<f32>(0.0, 0.0);
+  for (var i = rank; i < WG; i = i + count) {
     total = total + sg_partials[i];
   }
+  return vec2<f32>(subgroupAdd(total.x), subgroupAdd(total.y));
 {% else %}
   wg_red[tid] = value;
   workgroupBarrier();
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg_id: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+{%- if useSubgroups %},
   @builtin(subgroup_size) sg_size: u32
 {%- endif %}
 ) {
+  let row = wg_id.x + wg_id.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }
   // Base offset from the kept-axis coordinates.
   var base_off = 0u;
+{{ nd_offset("base_off", "row", 0) | indent(2, true) }}
+{% if vectorizedSpec %}
   let shift = f32(x[base_off / 4u].x);
 {% else %}
   let shift = f32(x[base_off]);
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
+{% if vectorizedSpec %}
   for (var q = tid; q < RV; q = q + WG) {
     let r = q * 4u;
 {% else %}
   for (var r = tid; r < R; r = r + WG) {
 {% endif %}
     var off = base_off;
+{{ nd_offset("off", "r", 1) | indent(4, true) }}{% if vectorizedSpec %}
     let d = vec4<f32>(x[off / 4u]) - vec4<f32>(shift);
     acc.x = acc.x + d.x + d.y + d.z + d.w;
     acc.y = acc.y + dot(d, d);
 {% endif %}
   }
+  let totals = reduce_pair(acc, tid{% if useSubgroups %}, sg_size{% endif %});
   let mean_d = totals.x / f32(R);
   let variance = max(totals.y / f32(R) - mean_d * mean_d, 0.0);
   let mean = shift + mean_d;
   let denom = sqrt(variance);
+{% if vectorizedSpec %}
   for (var q = tid; q < RV; q = q + WG) {
     let r = q * 4u;
 {% else %}
   for (var r = tid; r < R; r = r + WG) {
 {% endif %}
     var off = base_off;
+{{ nd_offset("off", "r", 1) | indent(4, true) }}{% if vectorizedSpec %}
     let v = vec4<f32>(x[off / 4u]);
+    y[off / 4u] = {{ vecType }}((v - vec4<f32>(mean)) / vec4<f32>(denom));
 {% else %}
+    y[off] = {{ scalar }}((f32(x[off]) - mean) / denom);
 {% endif %}
   }
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,23 +1,32 @@
 {
   "name": "ai.onnx.MeanVarianceNormalization",
-  "id": "_ai_onnx_meanvariancenormalization_webgpu_09532e8",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "Z7dtgwBXEEaVuYjwvsi+nzdPnqr+T3GbPX6F4A4qUyQ=",
-      "manifest.json": "8549q1hb8wFCRzXlrsGTXzEyEVO5WtP64OtKI3Hq6Ug=",
-      "mean-variance-normalization-serial-rows.wgsl.jinja": "C/TVvmYhaLeq+yKB4nAK535avgH9kcYaeSZKnRY/Df0=",
-      "mean-variance-normalization-subgroup.wgsl.jinja": "eFY8+N5f0/V/npXvheysb002bNEWtJLxNE//2aGzIr0=",
       "noop.wgsl.jinja": "k/5BMD6UO81N7XlF+t4iSKyt3dbtcqNMCru5aUKNBKE=",
-      "norm-flat-apply.wgsl.jinja": "O8G9eyv748OmW3GMQjqSFbgsiJ/pW/plJAzAbFGvGNc=",
-      "norm-flat-splitk-combine.wgsl.jinja": "uBFuaqUxbf6Qr2uE5huMcdP2+S2nHF+zM3RH0tzAx9o=",
-      "norm-flat-splitk-partials.wgsl.jinja": "YwepfH5ztse2kwIz4xhl80Gw5GEnWNMue0oGZUD47p8=",
-      "test.json": "OAgHM6FyfQCVSV8h47D92KbimR901cmqk1fTTm4g3q8="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.MeanVarianceNormalization" }
 }

 {
   "name": "ai.onnx.MeanVarianceNormalization",
+  "id": "_ai_onnx_meanvariancenormalization_webgpu_b716a78",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "zH0PcLLzlKTcTGM2WzPi8keFUjvDcCnttgcQLxJ9WuU=",
+      "manifest.json": "Snfz0B5Kkd7enbnP+ETHCh/vyME+C3vu4puxVB1xF3c=",
+      "mean-variance-normalization-serial-rows.wgsl.jinja": "qcoGxwmc0jShVeozos6da6awjOJbbIIKLfoMCTyzNuk=",
+      "mean-variance-normalization-subgroup.wgsl.jinja": "GTbYa5JP9Z186OoP6/BZ6w9CyZD2iihXHHl2Cil91pE=",
       "noop.wgsl.jinja": "k/5BMD6UO81N7XlF+t4iSKyt3dbtcqNMCru5aUKNBKE=",
+      "norm-flat-apply.wgsl.jinja": "fJkwUp+SOc5+EV2M+skzaoOyf4cCAj6RE/l9EFpRzlw=",
+      "norm-flat-splitk-combine.wgsl.jinja": "qrWEZGCTfaNnkDSBX7p6j+pnJHKCaAwFmHIXe0DJQC0=",
+      "norm-flat-splitk-partials.wgsl.jinja": "mQFvRpYjvQsoFwdFiShP0nMYFrWbaGmWxYIMXrSi4Sc=",
+      "test.json": "bGHJmclwK4mv21Rm4JYKnWsrwe5GeHdDnBXy+q8fIj8="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "empty_noop": ["noop.wgsl.jinja"],
+      "all_axes_flat_split": ["norm-flat-apply.wgsl.jinja", "norm-flat-splitk-combine.wgsl.jinja", "norm-flat-splitk-partials.wgsl.jinja"],
+      "serial_rows": ["mean-variance-normalization-serial-rows.wgsl.jinja"],
+      "cooperative_vec4": ["mean-variance-normalization-subgroup.wgsl.jinja"],
+      "cooperative_scalar": ["mean-variance-normalization-subgroup.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/norm-flat-apply.wgsl.jinja CHANGED Viewed

@@ -1,14 +1,13 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-@compute @workgroup_size({{ source.workgroupSize }}, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  let i = gid.x + gid.y * nwg.x * {{ source.workgroupSize }}u;
   if (i >= params.count) {
     return;
   }
-  y[i] = {{ source.scalar }}((f32(x[i]) - stats[0]) / stats[1]);
 }

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+@compute @workgroup_size({{ workgroupSizeSpec }}, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ workgroupSizeSpec }}u;
   if (i >= params.count) {
     return;
   }
+  y[i] = {{ scalar }}((f32(x[i]) - stats[0]) / stats[1]);
 }

build/webgpu/norm-flat-splitk-combine.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,9 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const SPLIT: u32 = {{ source.split }}u;
 @compute @workgroup_size(1, 1, 1)
 fn main() {

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+const SPLIT: u32 = {{ splitSpec }}u;
 @compute @workgroup_size(1, 1, 1)
 fn main() {

build/webgpu/norm-flat-splitk-partials.wgsl.jinja CHANGED Viewed

@@ -1,10 +1,10 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const WG: u32 = {{ source.workgroupSize }}u;
-const SPLIT: u32 = {{ source.split }}u;
 var<workgroup> reduction: array<vec2<f32>, WG>;

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSizeSpec }}u;
+const SPLIT: u32 = {{ splitSpec }}u;
 var<workgroup> reduction: array<vec2<f32>, WG>;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.MeanVarianceNormalization",
   "fixtureArrays": {
     "ort_axes_subset_rank5_0_2_4_input_x": [0.6369617, 0.2697867, 0.0409735, 0.0165276, 0.8132702, 0.9127556, 0.6066358, 0.7294966, 0.543625, 0.9350724, 0.8158536, 0.0027385, 0.8574043, 0.0335856, 0.7296554, 0.1756556, 0.8631789, 0.5414612, 0.2997119, 0.4226872, 0.0283197, 0.1242833, 0.6706244, 0.6471895, 0.6153851, 0.3836776, 0.9972099, 0.9808353, 0.685542, 0.6504593, 0.6884467, 0.3889214]
   },
@@ -144,7 +143,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc",
         "test": "MeanVarianceNormalizationTest.DefaultAxes",
-        "notes": "Rank-4 default/spatial axes companion: the variance is positive but subnormal, so the normalized output should be finite."
       },
       "attrs": { "axes": [0, 2, 3] },
       "inputs": {
@@ -165,7 +164,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc",
         "test": "MeanVarianceNormalizationTest.DefaultAxes",
-        "notes": "Rank-5 companion: the reduced variance is positive but subnormal, so normalization should produce finite values."
       },
       "attrs": { "axes": [0, 2, 3, 4] },
       "inputs": {
@@ -259,7 +258,7 @@
     {
       "name": "f32_spatial_axes_2x8x32x32",
       "provenance": {
-        "notes": "Compact sibling for the spatial-axes MVN benchmark; preserves axes=[2,3] over many channel planes without benchmark-scale tensors."
       },
       "attrs": { "axes": [2, 3] },
       "inputs": {
@@ -566,7 +565,7 @@
     {
       "name": "rank5_serial_rows_channel_axis_f16",
       "provenance": {
-        "notes": "float16 on the coalesced serial-row route, which needs at least SERIAL_MIN_ROWS independent rows and a reduction within SERIAL_MAX_REDUCTION. Only f32 cases and a bench had reached it."
       },
       "attrs": { "axes": [1] },
       "inputs": {
@@ -577,6 +576,23 @@
         }
       },
       "outputs": { "y": { "dtype": "float16", "shape": [2, 32, 8, 8, 4], "tolerance": 0.02 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_axes_subset_rank5_0_2_4_input_x": [0.6369617, 0.2697867, 0.0409735, 0.0165276, 0.8132702, 0.9127556, 0.6066358, 0.7294966, 0.543625, 0.9350724, 0.8158536, 0.0027385, 0.8574043, 0.0335856, 0.7296554, 0.1756556, 0.8631789, 0.5414612, 0.2997119, 0.4226872, 0.0283197, 0.1242833, 0.6706244, 0.6471895, 0.6153851, 0.3836776, 0.9972099, 0.9808353, 0.685542, 0.6504593, 0.6884467, 0.3889214]
   },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc",
         "test": "MeanVarianceNormalizationTest.DefaultAxes",
+        "notes": "With rank-4 default spatial axes, a positive subnormal variance must produce a finite normalized output."
       },
       "attrs": { "axes": [0, 2, 3] },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc",
         "test": "MeanVarianceNormalizationTest.DefaultAxes",
+        "notes": "With a rank-5 input, a positive subnormal reduced variance must produce finite normalized values."
       },
       "attrs": { "axes": [0, 2, 3, 4] },
       "inputs": {
     {
       "name": "f32_spatial_axes_2x8x32x32",
       "provenance": {
+        "notes": "A compact input exercises axes [2,3] over many channel planes without benchmark-scale tensors."
       },
       "attrs": { "axes": [2, 3] },
       "inputs": {
     {
       "name": "rank5_serial_rows_channel_axis_f16",
       "provenance": {
+        "notes": "Float16 data exercises the coalesced serial-row path with at least SERIAL_MIN_ROWS independent rows and a reduction within SERIAL_MAX_REDUCTION."
       },
       "attrs": { "axes": [1] },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "float16", "shape": [2, 32, 8, 8, 4], "tolerance": 0.02 } }
+    },
+    {
+      "name": "rank7_last_axis_above_axis_table",
+      "provenance": {
+        "notes": "A rank-7 input reduced along its final axis exercises axis 6. Each independent 64-element row must be normalized across that last dimension."
+      },
+      "attrs": { "axes": [-1] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 4, 2, 2, 2, 2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 1.0 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [1, 4, 2, 2, 2, 2, 64], "tolerance": 0.00002, "relTolerance": 0.00002 }
+      }
     }
   ]
 }