sync 2e7068faf55e

Browse files

Files changed (9) hide show

README.md +72 -0
build/webgpu/bench.json +66 -0
build/webgpu/manifest.json +591 -0
build/webgpu/metadata.json +21 -0
build/webgpu/norm-row-stats.wgsl.jinja +130 -0
build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja +98 -0
build/webgpu/rms-normalization-splitk-partials.wgsl.jinja +93 -0
build/webgpu/rms-normalization.wgsl.jinja +147 -0
build/webgpu/test.json +388 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.SimplifiedLayerNormalization
+`ai.onnx`  ·  ONNX Runtime compatibility operator  ·  default-domain since_version 1
+## Description
+Implements ONNX Runtime's legacy RMS normalization for models that serialize `SimplifiedLayerNormalization` in the default ONNX domain: `Y = (X / sqrt(mean(X^2) + epsilon)) * scale`. It shares kernels with [ONNX `RMSNormalization`](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html), but retains scale-before-output-cast ordering and optional `inv_std_var`. See the [compatibility discussion](https://github.com/onnx/onnx/issues/6582#issuecomment-3591862327). Supports float16/float32, float32 statistics, and a nonempty normalization suffix.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `X` | `x` | `T` | — | — | Input tensor; the mean square is taken over the suffix dimensions starting at `axis`. | required |
+| `scale` | `scale` | `V` | — | — | Scale tensor, unidirectionally broadcastable to `X`; output `Y` has this tensor's dtype. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `y` | `V` | same as `X` | same as `X` | Normalized and scaled output; same shape as `X` and same dtype as `scale`. | required |
+| `inv_std_var` | `invStdVar` | `U` | same as `X` | derived; see description | Optional inverse RMS statistic `1 / sqrt(mean(X^2) + epsilon)`; same rank as `X`, with the dimensions from `axis` onward set to one. | optional |
+## Attributes
+Default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `axis` | `-1` | The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes only the last dimension. |
+| `epsilon` | `0.00001` | Small constant added to the mean square before taking the square root to avoid division by zero. |
+| `stash_type` | `1` | TensorProto element type of `inv_std_var`; this package implements the legacy float32 value `1`. |
+| `keep_dims` | `1` | Legacy compatibility attribute accepted and ignored by ONNX Runtime. The supported value `1` records the runtime's fixed behavior: optional statistics retain the input rank with normalized dimensions set to one. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `V` | `float32`, `float16` |
+| `U` | `float32` |
+## 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
+- [`norm-row-stats.wgsl.jinja`](build/webgpu/norm-row-stats.wgsl.jinja)
+- [`rms-normalization-splitk-normalize.wgsl.jinja`](build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja)
+- [`rms-normalization-splitk-partials.wgsl.jinja`](build/webgpu/rms-normalization-splitk-partials.wgsl.jinja)
+- [`rms-normalization.wgsl.jinja`](build/webgpu/rms-normalization.wgsl.jinja)
+## 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.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/ai.onnx.SimplifiedLayerNormalization", { version: 1 });
+const { y } = await kernel({ x: { data: xData, shape: [8] }, scale: { data: scaleData, shape: [8] } });
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,66 @@

+{
+  "op": "ai.onnx.SimplifiedLayerNormalization",
+  "cases": [
+    {
+      "name": "simplifiedln-f32-4096x4096",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 4096, "dim": 4096 },
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 720, "scale": 0.5 },
+        "scale": { "shape": [4096], "dtype": "float32", "dist": "uniform", "seed": 721, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [4096, 4096], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "(args.rows * args.dim * 2 + args.dim) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "simplifiedln-f32-4096x4096-stats",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 4096, "dim": 4096 },
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 722, "scale": 0.5 },
+        "scale": { "shape": [4096], "dtype": "float32", "dist": "uniform", "seed": 723, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": {
+        "y": { "shape": [4096, 4096], "dtype": "float32" },
+        "invStdVar": { "shape": [4096, 1], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(args.rows * args.dim * 2 + args.dim) * dtypeBytes(args.dtype) + args.rows * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "simplifiedln-f16-8192x3072-stats",
+      "preset": "smoke",
+      "vars": { "dtype": "float16", "rows": 8192, "dim": 3072 },
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": { "shape": [8192, 3072], "dtype": "float16", "dist": "normal", "seed": 724, "scale": 0.5 },
+        "scale": { "shape": [3072], "dtype": "float16", "dist": "uniform", "seed": 725, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": {
+        "y": { "shape": [8192, 3072], "dtype": "float16" },
+        "invStdVar": { "shape": [8192, 1], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(args.rows * args.dim * 2 + args.dim) * dtypeBytes(args.dtype) + args.rows * 4"
+          }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,591 @@

+{
+  "domain": "ai.onnx",
+  "name": "SimplifiedLayerNormalization",
+  "conformance": "legacy-default-domain",
+  "sinceVersion": 1,
+  "description": "Implements ONNX Runtime's legacy RMS normalization for models that serialize `SimplifiedLayerNormalization` in the default ONNX domain: `Y = (X / sqrt(mean(X^2) + epsilon)) * scale`. It shares kernels with [ONNX `RMSNormalization`](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html), but retains scale-before-output-cast ordering and optional `inv_std_var`. See the [compatibility discussion](https://github.com/onnx/onnx/issues/6582#issuecomment-3591862327). Supports float16/float32, float32 statistics, and a nonempty normalization suffix.",
+  "inputs": [
+    {
+      "role": "X",
+      "dtype": "T",
+      "description": "Input tensor; the mean square is taken over the suffix dimensions starting at `axis`."
+    },
+    {
+      "role": "scale",
+      "dtype": "V",
+      "description": "Scale tensor, unidirectionally broadcastable to `X`; output `Y` has this tensor's dtype."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "Y",
+      "dtype": "V",
+      "rank": "ranks.X",
+      "shape": "shapes.X",
+      "description": "Normalized and scaled output; same shape as `X` and same dtype as `scale`."
+    },
+    {
+      "role": "inv_std_var",
+      "dtype": "U",
+      "rank": "ranks.X",
+      "shape": "prefix(shapes.X, axisNorm) + fill(1, ranks.X - axisNorm)",
+      "optional": true,
+      "description": "Optional inverse RMS statistic `1 / sqrt(mean(X^2) + epsilon)`; same rank as `X`, with the dimensions from `axis` onward set to one."
+    }
+  ],
+  "attributes": { "axis": -1, "epsilon": 0.00001, "stash_type": 1, "keep_dims": 1 },
+  "attributeDescriptions": {
+    "axis": "The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes only the last dimension.",
+    "epsilon": "Small constant added to the mean square before taking the square root to avoid division by zero.",
+    "stash_type": "TensorProto element type of `inv_std_var`; this package implements the legacy float32 value `1`.",
+    "keep_dims": "Legacy compatibility attribute accepted and ignored by ONNX Runtime. The supported value `1` records the runtime's fixed behavior: optional statistics retain the input rank with normalized dimensions set to one."
+  },
+  "attributeConstraints": { "stash_type": { "values": [1] }, "keep_dims": { "values": [1] } },
+  "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"], "U": ["float32"] },
+  "args": {
+    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
+    "scale": { "kind": "tensor", "semantic": "scale", "role": "input" },
+    "y": { "kind": "tensor", "semantic": "Y", "role": "output" },
+    "invStdVar": { "kind": "tensor", "semantic": "inv_std_var", "role": "output", "required": false }
+  },
+  "tunables": {
+    "WORKGROUP_SIZE": 256,
+    "SPLIT_MAX_ROWS": 256,
+    "SPLIT_MIN_HIDDEN": 16384,
+    "SPLIT_TARGET_ELEMENTS": 4096,
+    "MAX_SPLITS": 64
+  },
+  "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\"))",
+    "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
+    "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.X",
+    "normRows": "outer(shapes.X, axisNorm)",
+    "normHidden": "dim(shapes.X, axisNorm) * inner(shapes.X, axisNorm)",
+    "normRowStride": "max(1, min(normRows, device.limits.maxComputeWorkgroupsPerDimension))",
+    "rowWg": "min(normMaxWorkgroup, pow2ceil(max(1, normHidden)))",
+    "baseOk": "ranks.X >= 1 and sameShape(shapes.Y, shapes.X) and ranks.scale >= 0 and ranks.scale <= ranks.X and broadcastable(shapes.scale, shapes.X) and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and normHidden > 0 and attrs.stash_type == onnxDtypeCode(\"float32\") and f16Ok(dtypes.T) and f16Ok(dtypes.V)",
+    "lastAxisOk": "baseOk and (attrs.axis == -1 or attrs.axis == ranks.X - 1)",
+    "suffixAxisOk": "baseOk and ranks.X >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.X - 1)",
+    "noStats": "not present.invStdVar",
+    "statsOk": "present.invStdVar and ranks.inv_std_var == ranks.X and sameShape(prefix(shapes.inv_std_var, axisNorm), prefix(shapes.X, axisNorm)) and numel(suffix(shapes.inv_std_var, axisNorm)) == 1",
+    "sameDtype": "dtypes.T == dtypes.V",
+    "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normHidden, tunables.SPLIT_TARGET_ELEMENTS)))",
+    "splitScratchBytes": "normRows * splitCount * 4",
+    "splitFits": "normRows <= tunables.SPLIT_MAX_ROWS and splitCount <= device.limits.maxComputeWorkgroupsPerDimension and splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize"
+  },
+  "bindingSets": {
+    "rows": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "X",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$xElement"
+      },
+      {
+        "name": "scale",
+        "arg": "scale",
+        "semantic": "scale",
+        "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": "normRows" },
+            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
+          ]
+        }
+      }
+    ],
+    "rowsStats": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "X",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$xElement"
+      },
+      {
+        "name": "scale",
+        "arg": "scale",
+        "semantic": "scale",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$ioElement"
+      },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
+      {
+        "name": "inv_std_out",
+        "arg": "invStdVar",
+        "semantic": "inv_std_var",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "normRows" },
+            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
+          ]
+        }
+      }
+    ],
+    "splitPartials": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "X",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$xElement"
+      },
+      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "normRows" },
+            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
+          ]
+        }
+      }
+    ],
+    "splitNormalize": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "X",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$xElement"
+      },
+      {
+        "name": "scale",
+        "arg": "scale",
+        "semantic": "scale",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$ioElement"
+      },
+      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "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": "normRows" },
+            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
+          ]
+        }
+      }
+    ],
+    "splitNormalizeStats": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "X",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$xElement"
+      },
+      {
+        "name": "scale",
+        "arg": "scale",
+        "semantic": "scale",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$ioElement"
+      },
+      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
+      {
+        "name": "inv_std_out",
+        "arg": "invStdVar",
+        "semantic": "inv_std_var",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "normRows" },
+            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
+          ]
+        }
+      }
+    ]
+  },
+  "variants": [
+    {
+      "id": "last_axis",
+      "priority": 1,
+      "when": ["lastAxisOk", "noStats"],
+      "constants": {
+        "scalar": "dtypes.V",
+        "xElement": "dtypes.T",
+        "ioElement": "dtypes.V",
+        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
+        "hiddenSize": "normHidden",
+        "workgroupSize": "rowWg",
+        "epsilon": "attrs.epsilon"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.Row",
+          "source": {
+            "shader": "rms-normalization.wgsl.jinja",
+            "inputs": {
+              "xShape": "shapes.X",
+              "scaleShape": "shapes.scale",
+              "xRank": "ranks.X",
+              "scaleRank": "ranks.scale",
+              "writeStats": false,
+              "rmsScaleAfterCast": false
+            }
+          },
+          "bindings": "rows",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_stats",
+      "priority": 2,
+      "when": ["lastAxisOk", "statsOk"],
+      "constants": {
+        "scalar": "dtypes.V",
+        "xElement": "dtypes.T",
+        "ioElement": "dtypes.V",
+        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
+        "hiddenSize": "normHidden",
+        "workgroupSize": "rowWg",
+        "epsilon": "attrs.epsilon"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.Row",
+          "source": {
+            "shader": "rms-normalization.wgsl.jinja",
+            "inputs": {
+              "xShape": "shapes.X",
+              "scaleShape": "shapes.scale",
+              "xRank": "ranks.X",
+              "scaleRank": "ranks.scale",
+              "writeStats": true,
+              "rmsScaleAfterCast": false
+            }
+          },
+          "bindings": "rowsStats",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "suffix_axis",
+      "priority": 10,
+      "when": ["suffixAxisOk", "noStats"],
+      "constants": {
+        "scalar": "dtypes.V",
+        "xElement": "dtypes.T",
+        "ioElement": "dtypes.V",
+        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
+        "hiddenSize": "normHidden",
+        "workgroupSize": "rowWg",
+        "epsilon": "attrs.epsilon"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.Row",
+          "source": {
+            "shader": "rms-normalization.wgsl.jinja",
+            "inputs": {
+              "xShape": "shapes.X",
+              "scaleShape": "shapes.scale",
+              "xRank": "ranks.X",
+              "scaleRank": "ranks.scale",
+              "writeStats": false,
+              "rmsScaleAfterCast": false
+            }
+          },
+          "bindings": "rows",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "suffix_axis_stats",
+      "priority": 11,
+      "when": ["suffixAxisOk", "statsOk"],
+      "constants": {
+        "scalar": "dtypes.V",
+        "xElement": "dtypes.T",
+        "ioElement": "dtypes.V",
+        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
+        "hiddenSize": "normHidden",
+        "workgroupSize": "rowWg",
+        "epsilon": "attrs.epsilon"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.Row",
+          "source": {
+            "shader": "rms-normalization.wgsl.jinja",
+            "inputs": {
+              "xShape": "shapes.X",
+              "scaleShape": "shapes.scale",
+              "xRank": "ranks.X",
+              "scaleRank": "ranks.scale",
+              "writeStats": true,
+              "rmsScaleAfterCast": false
+            }
+          },
+          "bindings": "rowsStats",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "suffix_axis_splitk",
+      "priority": 15,
+      "when": ["baseOk", "ranks.X >= 2", "noStats", "splitFits"],
+      "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
+      "constants": {
+        "scalar": "dtypes.V",
+        "xElement": "dtypes.T",
+        "ioElement": "dtypes.V",
+        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
+        "hiddenSize": "normHidden",
+        "workgroupSize": "normMaxWorkgroup",
+        "split": "splitCount",
+        "epsilon": "attrs.epsilon"
+      },
+      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
+      "passes": [
+        {
+          "id": "partials",
+          "name": "SimplifiedLayerNormalization.SplitKPartials",
+          "source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
+          "bindings": "splitPartials",
+          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
+        },
+        {
+          "id": "normalize",
+          "name": "SimplifiedLayerNormalization.SplitKNormalize",
+          "source": {
+            "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
+            "inputs": {
+              "xShape": "shapes.X",
+              "scaleShape": "shapes.scale",
+              "xRank": "ranks.X",
+              "scaleRank": "ranks.scale",
+              "writeStats": false,
+              "rmsScaleAfterCast": false
+            }
+          },
+          "bindings": "splitNormalize",
+          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
+        }
+      ]
+    },
+    {
+      "id": "suffix_axis_splitk_stats",
+      "priority": 16,
+      "when": ["baseOk", "ranks.X >= 2", "statsOk", "splitFits"],
+      "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
+      "constants": {
+        "scalar": "dtypes.V",
+        "xElement": "dtypes.T",
+        "ioElement": "dtypes.V",
+        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
+        "hiddenSize": "normHidden",
+        "workgroupSize": "normMaxWorkgroup",
+        "split": "splitCount",
+        "epsilon": "attrs.epsilon"
+      },
+      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
+      "passes": [
+        {
+          "id": "partials",
+          "name": "SimplifiedLayerNormalization.SplitKPartials",
+          "source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
+          "bindings": "splitPartials",
+          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
+        },
+        {
+          "id": "normalize",
+          "name": "SimplifiedLayerNormalization.SplitKNormalize",
+          "source": {
+            "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
+            "inputs": {
+              "xShape": "shapes.X",
+              "scaleShape": "shapes.scale",
+              "xRank": "ranks.X",
+              "scaleRank": "ranks.scale",
+              "writeStats": true,
+              "rmsScaleAfterCast": false
+            }
+          },
+          "bindings": "splitNormalizeStats",
+          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_row_vec4",
+      "priority": 110,
+      "when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)", "dim(shapes.X, -1) % 4 == 0"],
+      "constants": {
+        "scalar": "dtypes.T",
+        "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "source": {
+            "shader": "norm-row-stats.wgsl.jinja",
+            "inputs": {
+              "mode": "\"rms\"",
+              "vec4": true,
+              "writeStats": false,
+              "rmsScaleAfterCast": false,
+              "scalar": "dtypes.T",
+              "usesF16": "dtypes.T == \"f16\"",
+              "hidden": "dim(shapes.X, -1)",
+              "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1) / 4)))",
+              "epsilon": "attrs.epsilon",
+              "hiddenVec": "dim(shapes.X, -1) / 4",
+              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+              "combineSubgroups": "hasSubgroupId"
+            }
+          },
+          "subgroupCollectivesWidth": "portable",
+          "bindings": "rows",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_row",
+      "priority": 100,
+      "when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)"],
+      "constants": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "source": {
+            "shader": "norm-row-stats.wgsl.jinja",
+            "inputs": {
+              "mode": "\"rms\"",
+              "vec4": false,
+              "writeStats": false,
+              "rmsScaleAfterCast": false,
+              "scalar": "dtypes.T",
+              "usesF16": "dtypes.T == \"f16\"",
+              "hidden": "dim(shapes.X, -1)",
+              "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1))))",
+              "epsilon": "attrs.epsilon",
+              "hiddenVec": 1,
+              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+              "combineSubgroups": "hasSubgroupId"
+            }
+          },
+          "subgroupCollectivesWidth": "portable",
+          "bindings": "rows",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_row_vec4_stats",
+      "priority": 112,
+      "when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)", "dim(shapes.X, -1) % 4 == 0"],
+      "constants": {
+        "scalar": "dtypes.T",
+        "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "source": {
+            "shader": "norm-row-stats.wgsl.jinja",
+            "inputs": {
+              "mode": "\"rms\"",
+              "vec4": true,
+              "writeStats": true,
+              "rmsScaleAfterCast": false,
+              "scalar": "dtypes.T",
+              "usesF16": "dtypes.T == \"f16\"",
+              "hidden": "dim(shapes.X, -1)",
+              "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1) / 4)))",
+              "epsilon": "attrs.epsilon",
+              "hiddenVec": "dim(shapes.X, -1) / 4",
+              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+              "combineSubgroups": "hasSubgroupId"
+            }
+          },
+          "subgroupCollectivesWidth": "portable",
+          "bindings": "rowsStats",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_row_stats",
+      "priority": 102,
+      "when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)"],
+      "constants": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "source": {
+            "shader": "norm-row-stats.wgsl.jinja",
+            "inputs": {
+              "mode": "\"rms\"",
+              "vec4": false,
+              "writeStats": true,
+              "rmsScaleAfterCast": false,
+              "scalar": "dtypes.T",
+              "usesF16": "dtypes.T == \"f16\"",
+              "hidden": "dim(shapes.X, -1)",
+              "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1))))",
+              "epsilon": "attrs.epsilon",
+              "hiddenVec": 1,
+              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+              "combineSubgroups": "hasSubgroupId"
+            }
+          },
+          "subgroupCollectivesWidth": "portable",
+          "bindings": "rowsStats",
+          "dispatch": { "workgroups": "normRows" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+  "name": "ai.onnx.SimplifiedLayerNormalization",
+  "id": "_ai_onnx_simplifiedlayernormalization_webgpu_e7aa570",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "pvZ/eyxjfpTLTo48nnQhDaKG61UPA6L0EPXu4QnH7MY=",
+      "manifest.json": "hfKp3YN8J74f7KzVbFXsOBq5wNbZm2A8QMJ3ww5MAdw=",
+      "norm-row-stats.wgsl.jinja": "GopNOsand0HoXzNtQnzGdvJKmHbCv96No/7heLcS9VA=",
+      "rms-normalization-splitk-normalize.wgsl.jinja": "YklXSmW/XGI1M5KVzfayKp0EQ+3iAYHeOnTAMA4UW4c=",
+      "rms-normalization-splitk-partials.wgsl.jinja": "GbYu7Cj/LLO9y2Rvhw30RJAUAduI0xfN5YS3H/4Z1k8=",
+      "rms-normalization.wgsl.jinja": "wNOsrtTIc7DmBifozpAm2fZ3iWznPJ04EhFMkOq6GqA=",
+      "test.json": "ycI7oQ1MKS87yo6QXLI36xLz840eYS69Y/anRQHQkUY="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.SimplifiedLayerNormalization" }
+}

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

	@@ -0,0 +1,130 @@

+{% if source.usesF16 %}
+enable f16;
+{% endif %}
+{% set combineSubgroups = source.combineSubgroups %}
+{% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
+{% set writeStats = source.writeStats if source.writeStats is defined else false %}
+{% set rmsScaleVec = "vec4<f32>(scale[i])" %}
+{% set rmsScaleScalar = "f32(scale[i])" %}
+{% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
+    if combineSubgroups else ", tid: u32" %}
+{% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
+    if combineSubgroups else ", tid" %}
+{% if combineSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Workgroup-parallel single-pass row statistics + fused normalize/affine.
+//
+// One workgroup owns one contiguous normalization span ("row": a last-axis
+// row, an instance plane, or a channel group). Threads stride the row once,
+// accumulating (sum, sum_sq) simultaneously. Partials are reduced either with
+// subgroupAdd plus a shared-memory combine or with a portable shared-memory
+// tree, then every thread applies the fused normalize + affine write.
+//
+// RMS mode uses sum_sq / HIDDEN without computing or subtracting a mean.
+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 }};
+{% if combineSubgroups %}
+var<workgroup> sg_partials: array<f32, WG>;
+fn reduce_scalar(value: f32{{ reduceThreadParameters }}) -> f32 {
+  let s = subgroupAdd(value);
+  if (num_sg == 1u) {
+    return s;
+  }
+  if (sg_lane == 0u) {
+    sg_partials[sg_id] = s;
+  }
+  workgroupBarrier();
+  var total = 0.0;
+  for (var i = 0u; i < num_sg; i++) {
+    total += sg_partials[i];
+  }
+  return total;
+}
+{% else %}
+// Each shared-memory tree reduction deliberately ends with a barrier. It keeps
+// lanes that have read the result from starting a later reduction and
+// overwriting scratch while slower lanes are still reading it.
+var<workgroup> tr0: array<f32, WG>;
+fn reduce_scalar(value: f32, tid: u32) -> f32 {
+  tr0[tid] = value;
+  workgroupBarrier();
+  var stride: u32 = WG / 2u;
+  loop {
+    if (stride == 0u) { break; }
+    if (tid < stride) { tr0[tid] = tr0[tid] + tr0[tid + stride]; }
+    stride = stride / 2u;
+    workgroupBarrier();
+  }
+  let reduced = tr0[0];
+  workgroupBarrier();
+  return reduced;
+}
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg_id: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
+  @builtin(subgroup_invocation_id) sg_lane: u32,
+  @builtin(subgroup_id) sg_id: u32,
+  @builtin(num_subgroups) num_sg: u32{% endif %}
+) {
+  let row = wg_id.x + wg_id.y * params.rowStride;
+  if (row >= params.rows) {
+    return;
+  }
+  let tid = lid.x;
+{% if source.vec4 and not scalarIo %}
+  let base = row * HIDDEN_V;
+{% else %}
+  let base = row * HIDDEN;
+{% endif %}
+  var acc = 0.0;
+{% if source.vec4 %}
+  for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let v = vec4<f32>(x[base + i]);
+    acc = acc + dot(v, v);
+  }
+{% else %}
+  for (var i = tid; i < HIDDEN; i = i + WG) {
+    let v = f32(x[base + i]);
+    acc = acc + v * v;
+  }
+{% endif %}
+  let total = reduce_scalar(acc{{ reduceThreadArguments }});
+  let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
+{% if writeStats %}
+  if (tid == 0u) {
+    inv_std_out[row] = inv;
+  }
+{% endif %}
+{% if source.vec4 %}
+  for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let idx = base + i;
+    let v = vec4<f32>(x[idx]);
+    y[idx] = {{ source.vecType }}(v * inv * {{ rmsScaleVec }});
+  }
+{% else %}
+  for (var i = tid; i < HIDDEN; i = i + WG) {
+    let idx = base + i;
+    let v = f32(x[idx]);
+    y[idx] = {{ source.scalar }}(v * inv * {{ rmsScaleScalar }});
+  }
+{% endif %}
+}

build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,98 @@

+// Split-K normalize pass. Each workgroup (row = wg.x, split index = wg.z)
+// folds the SPLIT per-row partial sums of squares into the RMS scale, then
+// normalizes its HIDDEN/SPLIT slice. SPLIT is small (<=64), so the serial fold
+// avoids a third combine pass. Scale offsets follow the suffix-axis broadcast
+// contract.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hiddenSize }}u;
+const EPSILON: f32 = {{ epsilon }};
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+{% if source.scaleRank > 0 %}
+const X_RANK: u32 = {{ source.xRank }}u;
+const SCALE_RANK: u32 = {{ source.scaleRank }}u;
+const X_SHAPE: array<u32, {{ source.xRank }}> = array<u32, {{ source.xRank }}>({% for d in source.xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+const SCALE_SHAPE: array<u32, {{ source.scaleRank }}> = array<u32, {{ source.scaleRank }}>({% for d in source.scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+fn x_stride(axis: u32) -> u32 {
+  var stride = 1u;
+  for (var i = axis + 1u; i < X_RANK; i += 1u) {
+    stride *= X_SHAPE[i];
+  }
+  return stride;
+}
+fn scale_stride(axis: u32) -> u32 {
+  var stride = 1u;
+  for (var i = axis + 1u; i < SCALE_RANK; i += 1u) {
+    stride *= SCALE_SHAPE[i];
+  }
+  return stride;
+}
+{% endif %}
+fn scale_offset({% if source.scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
+{% if source.scaleRank == 0 %}
+  return 0u;
+{% else %}
+  var rem = out_index;
+  var offset = 0u;
+  for (var axis = 0u; axis < X_RANK; axis += 1u) {
+    let stride = x_stride(axis);
+    let coord = rem / stride;
+    rem %= stride;
+    let scale_axis = i32(axis) - i32(X_RANK - SCALE_RANK);
+    if (scale_axis >= 0) {
+      let s_axis = u32(scale_axis);
+      if (SCALE_SHAPE[s_axis] != 1u) {
+        offset += coord * scale_stride(s_axis);
+      }
+    }
+  }
+  return offset;
+{% endif %}
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x + wg.y * params.rowStride;
+  if (row >= params.rows) {
+    return;
+  }
+  let k = wg.z;
+  let tid = lid.x;
+  var total = 0.0;
+  for (var i = 0u; i < SPLIT; i = i + 1u) {
+    total = total + partials[row * SPLIT + i];
+  }
+  let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
+{% if source.writeStats %}
+  // Every split workgroup folds the same partials, so one designated
+  // workgroup writes the row statistic.
+  if (k == 0u && tid == 0u) {
+    inv_std_out[row] = inv;
+  }
+{% endif %}
+  let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
+  let start = k * chunk;
+  var end = start + chunk;
+  if (end > HIDDEN) { end = HIDDEN; }
+  let base = row * HIDDEN;
+  var d = start + tid;
+  loop {
+    if (d >= end) { break; }
+    let index = base + d;
+    let value = f32(x[index]) * inv * f32(scale[scale_offset({% if source.scaleRank > 0 %}index{% endif %})]);
+    y[index] = {{ scalar }}(value);
+    d = d + WG;
+  }
+}

build/webgpu/rms-normalization-splitk-partials.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,93 @@

+{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" %}
+{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
+{%- else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
+{%- endif %}
+{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
+  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
+  loop {
+{% if form == "head" %}
+{% if breakInline %}
+    if ({{ svar }} == 0u) { break; }
+{% else %}
+    if ({{ svar }} == 0u) {
+      break;
+    }
+{% endif %}
+{% endif %}
+{% if bodyInline %}
+    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
+{% else %}
+    if ({{ idx }} < {{ svar }}) {
+{% for a in arrays %}
+      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
+{% endfor %}
+    }
+{% endif %}
+{% if form == "head" %}
+{% if barrierFirst %}
+    workgroupBarrier();
+    {{ svar }} = {{ svar }} / 2u;
+{% else %}
+    {{ svar }} = {{ svar }} / 2u;
+    workgroupBarrier();
+{% endif %}
+{% else %}
+    workgroupBarrier();
+    if ({{ svar }} == 1u) {
+      break;
+    }
+    {{ svar }} = {{ svar }} / 2u;
+{% endif %}
+  }
+{%- endmacro %}
+/* Split-K partial sum-of-squares for tensors with few rows and a large hidden
+   dimension. A workgroup-per-row kernel exposes too little parallelism in this
+   regime, so this pass splits each row across SPLIT workgroups
+   (row = wg.x, split index = wg.z). Each workgroup accumulates a partial
+   sum-of-squares over its HIDDEN/SPLIT slice and writes one partial to scratch.
+   The normalize pass folds the SPLIT partials per row. Split-K reassociates the
+   f32 sum, so this route is not bit-identical to the unsplit reduction. */
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hiddenSize }}u;
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+var<workgroup> red: array<f32, WG>;
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x + wg.y * params.rowStride;
+  if (row >= params.rows) {
+    return;
+  }
+  let k = wg.z;
+  let tid = lid.x;
+  let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
+  let start = k * chunk;
+  var end = start + chunk;
+  if (end > HIDDEN) { end = HIDDEN; }
+  let base = row * HIDDEN;
+  var acc = 0.0;
+  var d = start + tid;
+  loop {
+    if (d >= end) { break; }
+    let v = f32(x[base + d]);
+    acc = acc + v * v;
+    d = d + WG;
+  }
+  red[tid] = acc;
+  workgroupBarrier();
+{{ wgsl_tree_fold(["red"], idx="tid", wg="WG", typed=true, form="head", breakInline=true, bodyInline=true) }}
+  if (tid == 0u) {
+    partials[row * SPLIT + k] = red[0];
+  }
+}

build/webgpu/rms-normalization.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,147 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hiddenSize }}u;
+const EPSILON: f32 = {{ epsilon }};
+const WG: u32 = {{ workgroupSize }}u;
+var<workgroup> partial: array<f32, WG>;
+{% if source.scaleRank > 0 %}
+const X_RANK: u32 = {{ source.xRank }}u;
+const SCALE_RANK: u32 = {{ source.scaleRank }}u;
+const X_SHAPE: array<u32, {{ source.xRank }}> = array<u32, {{ source.xRank }}>({% for d in source.xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+const SCALE_SHAPE: array<u32, {{ source.scaleRank }}> = array<u32, {{ source.scaleRank }}>({% for d in source.scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+fn x_stride(axis: u32) -> u32 {
+  var stride = 1u;
+  for (var i = axis + 1u; i < X_RANK; i += 1u) {
+    stride *= X_SHAPE[i];
+  }
+  return stride;
+}
+fn scale_stride(axis: u32) -> u32 {
+  var stride = 1u;
+  for (var i = axis + 1u; i < SCALE_RANK; i += 1u) {
+    stride *= SCALE_SHAPE[i];
+  }
+  return stride;
+}
+{% endif %}
+fn scale_offset({% if source.scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
+{% if source.scaleRank == 0 %}
+  return 0u;
+{% else %}
+  var rem = out_index;
+  var offset = 0u;
+  for (var axis = 0u; axis < X_RANK; axis += 1u) {
+    let stride = x_stride(axis);
+    let coord = rem / stride;
+    rem %= stride;
+    let scale_axis = i32(axis) - i32(X_RANK - SCALE_RANK);
+    if (scale_axis >= 0) {
+      let s_axis = u32(scale_axis);
+      if (SCALE_SHAPE[s_axis] != 1u) {
+        offset += coord * scale_stride(s_axis);
+      }
+    }
+  }
+  return offset;
+{% endif %}
+}
+{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" %}
+{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
+{%- else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
+{%- endif %}
+{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
+  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
+  loop {
+{% if form == "head" %}
+{% if breakInline %}
+    if ({{ svar }} == 0u) { break; }
+{% else %}
+    if ({{ svar }} == 0u) {
+      break;
+    }
+{% endif %}
+{% endif %}
+{% if bodyInline %}
+    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
+{% else %}
+    if ({{ idx }} < {{ svar }}) {
+{% for a in arrays %}
+      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
+{% endfor %}
+    }
+{% endif %}
+{% if form == "head" %}
+{% if barrierFirst %}
+    workgroupBarrier();
+    {{ svar }} = {{ svar }} / 2u;
+{% else %}
+    {{ svar }} = {{ svar }} / 2u;
+    workgroupBarrier();
+{% endif %}
+{% else %}
+    workgroupBarrier();
+    if ({{ svar }} == 1u) {
+      break;
+    }
+    {{ svar }} = {{ svar }} / 2u;
+{% endif %}
+  }
+{%- endmacro %}
+// Reusing partial after this reduction requires a barrier between the read of
+// partial[0] and the next write, or the next round can race the prior readers.
+{% set trailingBarrier = trailingBarrier is defined and trailingBarrier %}
+fn reduce_sum(value: f32, tid: u32) -> f32 {
+  partial[tid] = value;
+  workgroupBarrier();
+{{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
+{% if trailingBarrier %}
+  let total = partial[0];
+  workgroupBarrier();
+  return total;
+{% else %}
+  return partial[0];
+{% endif %}
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x + wg.y * params.rowStride;
+  if (row >= params.rows) {
+    return;
+  }
+  let tid = lid.x;
+  let base = row * HIDDEN;
+  var local_sq = 0.0;
+  for (var d = tid; d < HIDDEN; d = d + WG) {
+    let value = f32(x[base + d]);
+    local_sq = local_sq + value * value;
+  }
+  let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(HIDDEN) + EPSILON);
+{% if source.writeStats %}
+  if (tid == 0u) {
+    inv_std_out[row] = inv;
+  }
+{% endif %}
+  for (var d = tid; d < HIDDEN; d = d + WG) {
+    let index = base + d;
+    let value = f32(x[index]) * inv * f32(scale[scale_offset({% if source.scaleRank > 0 %}index{% endif %})]);
+    y[base + d] = {{ scalar }}(value);
+  }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,388 @@

+{
+  "op": "ai.onnx.SimplifiedLayerNormalization",
+  "fixtureArrays": {
+    "f16_scalar_cast_x": [-1.1103515625, 2.982421875, 1.248046875, -1.8544921875],
+    "f16_scalar_cast_scale": [2.015625],
+    "f16_scalar_cast_y": [-1.1513671875, 3.091796875, 1.2939453125, -1.921875]
+  },
+  "cases": [
+    {
+      "name": "legacy_last_axis_vec4_no_stats_2x8",
+      "attrs": { "epsilon": 0.000001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.19, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.000002 } }
+    },
+    {
+      "name": "legacy_last_axis_scalar_no_stats_2x5",
+      "attrs": { "epsilon": 0.000001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.19, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 5], "tolerance": 0.000002 } }
+    },
+    {
+      "name": "legacy_last_axis_vec4_stats_2x8",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.21 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.11, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.000002 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "legacy_last_axis_scalar_stats_2x33",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.27 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.13, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 33], "tolerance": 0.000002 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "legacy_scalar_scale_generic_stats_2x8",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.09 }
+        },
+        "scale": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.75] } }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.000002 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "legacy_suffix_axis1_stats_2x3x4",
+      "attrs": { "epsilon": 0.00001, "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.29, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3, 4], "tolerance": 0.000002 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "legacy_mixed_f16_x_f32_scale_stats_2x4",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [2, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.21 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.11, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.002 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.002 }
+      }
+    },
+    {
+      "name": "legacy_mixed_f32_x_f16_scale_2x4",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.15, "cosStep": 0.25 }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.27, "cosStep": 0.09, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 4], "tolerance": 0.003 } }
+    },
+    {
+      "name": "legacy_f16_vec4_stats_4x32",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.27 }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.13, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float16", "shape": [4, 32], "tolerance": 0.005 },
+        "invStdVar": { "dtype": "float32", "shape": [4, 1], "tolerance": 0.002 }
+      }
+    },
+    {
+      "name": "legacy_suffix_splitk_stats_2x64x512",
+      "attrs": { "epsilon": 0.00001, "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 64, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.017 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [64, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.007, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 64, 512], "tolerance": 0.00001 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "legacy_suffix_splitk_no_stats_1x32x512",
+      "attrs": { "epsilon": 0.00001, "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 32, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [32, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.019, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 32, 512], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "legacy_suffix_splitk_scalar_scale_no_stats_1x32x512",
+      "attrs": { "epsilon": 0.00001, "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 32, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013 }
+        },
+        "scale": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.75] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 32, 512], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "legacy_f16_scale_before_output_cast_vec4_exact",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc",
+        "test": "SimplifiedLayerNormalization legacy f16 cast boundary",
+        "notes": "The legacy op applies scale before its single output cast. The first value differs by one f16 ULP from standardized RMSNormalization, which casts the normalized stage before scale."
+      },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "data": { "kind": "values", "values": [-1.4765625, -1.861328125, 1.4697265625, 1.0703125] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [4],
+          "data": { "kind": "values", "values": [0.316162109375, -2.791015625, -0.0877685546875, 2.0859375] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "tolerance": 0,
+          "relTolerance": 0,
+          "data": { "kind": "values", "values": [-0.31201171875, 3.47265625, -0.08624267578125, 1.4921875] }
+        }
+      }
+    },
+    {
+      "name": "legacy_f16_scale_before_output_cast_generic_exact",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc",
+        "test": "SimplifiedLayerNormalization scalar-scale cast boundary",
+        "notes": "A scalar scale forces the generic shared template and pins the same legacy cast ordering independently of the optimized row path."
+      },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/f16_scalar_cast_x" } }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/f16_scalar_cast_scale" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "tolerance": 0,
+          "relTolerance": 0,
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/f16_scalar_cast_y" } }
+        }
+      }
+    },
+    {
+      "name": "legacy_f16_scale_before_output_cast_splitk_exact",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc",
+        "test": "SimplifiedLayerNormalization split-K f16 cast boundary",
+        "notes": "Forces the split-K shared kernel on the exact scalar-scale boundary and pins the legacy scale-before-output-cast ordering."
+      },
+      "requires": { "features": ["shader-f16"] },
+      "tunables": { "SPLIT_MIN_HIDDEN": 1, "SPLIT_TARGET_ELEMENTS": 1 },
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/f16_scalar_cast_x" } }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/f16_scalar_cast_scale" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "tolerance": 0,
+          "relTolerance": 0,
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/f16_scalar_cast_y" } }
+        }
+      }
+    },
+    {
+      "name": "legacy_rank1_default_attrs_8",
+      "attrs": {},
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [8], "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.03 } },
+        "scale": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.41, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [8], "tolerance": 0.000002 } }
+    },
+    {
+      "name": "legacy_explicit_stash_type_float32",
+      "attrs": { "epsilon": 0.00001, "axis": -1, "stash_type": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.37, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.000002 } }
+    },
+    {
+      "name": "legacy_keep_dims_compatibility_attribute",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/layer_norm_test.cc",
+        "test": "SimplifiedLayerNormalization keep_dims=1",
+        "notes": "ONNX Runtime's legacy schema permits unchecked attributes and its own compatibility test emits keep_dims=1. The kernel ignores the attribute and always retains reduced dimensions for optional statistics."
+      },
+      "attrs": { "epsilon": 0.00001, "axis": 1, "keep_dims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.09, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3, 4], "tolerance": 0.000002 },
+        "invStdVar": { "dtype": "float32", "shape": [2, 1, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "legacy_zero_outer_rows_stats",
+      "attrs": { "epsilon": 0.00001, "axis": -1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [0, 4], "data": { "kind": "values", "values": [] } },
+        "scale": { "dtype": "float32", "shape": [4], "data": { "kind": "constant", "value": 1.0 } }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [0, 4], "tolerance": 0.000002 },
+        "invStdVar": { "dtype": "float32", "shape": [0, 1], "tolerance": 0.00001 }
+      }
+    }
+  ]
+}