Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.LayerNormalization

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
a3b790c
·
verified ·
1 Parent(s): 49b0523

sync 91d990483a17

Browse files
Files changed (7) hide show
  1. README.md +15 -11
  2. build/webgpu/bench.json +0 -1
  3. build/webgpu/layer-normalization.wgsl.jinja +7 -10
  4. build/webgpu/manifest.json +289 -625
  5. build/webgpu/metadata.json +33 -8
  6. build/webgpu/norm-row-stats.wgsl.jinja +120 -17
  7. build/webgpu/test.json +7 -5
README.md CHANGED
@@ -18,19 +18,19 @@ See the [ONNX `LayerNormalization` spec](https://onnx.ai/onnx/operators/onnx__La
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
- | `X` | `x` | `T` | — | — | Tensor to be normalized. | required |
24
- | `Scale` | `scale` | `T` | — | — | Scale tensor applied after normalization. | required |
25
- | `B` | `b` | `T` | — | — | Optional bias tensor added after scaling. | optional |
26
 
27
  ## Outputs
28
 
29
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
30
  | --- | --- | --- | --- | --- | --- | --- |
31
- | `Y` | `y` | `T` | same as `X` | same as `X` | Normalized and scaled output tensor; same shape as X. | required |
32
- | `Mean` | `mean` | `float32` | same as `X` | — | Per-normalization-group mean in the ONNX broadcastable keepdims shape: dimensions before `axis` are preserved and dimensions from `axis` onward are 1. | optional |
33
- | `InvStdDev` | `invStdDev` | `float32` | same as `X` | — | Per-normalization-group reciprocal standard deviation `1 / sqrt(variance + epsilon)`, returned in the same ONNX broadcastable keepdims shape as `Mean`. | optional |
34
 
35
  ## Attributes
36
 
@@ -50,7 +50,7 @@ Default values (overridable per request):
50
 
51
  ## Files
52
 
53
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
54
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
55
  - [`test.json`](build/webgpu/test.json) — correctness cases
56
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -59,10 +59,14 @@ Default values (overridable per request):
59
 
60
  ## Use with `@huggingface/kernels`
61
 
62
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
63
- It then allocates the result tensors automatically.
 
 
 
64
 
65
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
66
 
67
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
68
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
+ | `x` | `X` | `T` | — | — | Tensor to be normalized. | required |
24
+ | `scale` | `Scale` | `T` | — | — | Scale tensor applied after normalization. | required |
25
+ | `b` | `B` | `T` | — | — | Optional bias tensor added after scaling. | optional |
26
 
27
  ## Outputs
28
 
29
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
30
  | --- | --- | --- | --- | --- | --- | --- |
31
+ | `y` | `Y` | `T` | same as `x` | same as `x` | Normalized and scaled output tensor; same shape as X. | required |
32
+ | `mean` | `Mean` | `float32` | same as `x` | — | Per-normalization-group mean in the ONNX broadcastable keepdims shape: dimensions before `axis` are preserved and dimensions from `axis` onward are 1. | optional |
33
+ | `invStdDev` | `InvStdDev` | `float32` | same as `x` | — | Per-normalization-group reciprocal standard deviation `1 / sqrt(variance + epsilon)`, returned in the same ONNX broadcastable keepdims shape as `Mean`. | optional |
34
 
35
  ## Attributes
36
 
 
50
 
51
  ## Files
52
 
53
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
54
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
55
  - [`test.json`](build/webgpu/test.json) — correctness cases
56
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
59
 
60
  ## Use with `@huggingface/kernels`
61
 
62
+ ```sh
63
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
64
+ ```
65
+
66
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
67
 
68
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
69
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
70
 
71
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
72
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "ai.onnx.LayerNormalization",
3
  "cases": [
4
  {
5
  "name": "layernorm-f32-256x1024",
 
1
  {
 
2
  "cases": [
3
  {
4
  "name": "layernorm-f32-256x1024",
build/webgpu/layer-normalization.wgsl.jinja CHANGED
@@ -1,6 +1,3 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
5
  fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
6
  {% if out_numel == 0 %}
@@ -56,24 +53,24 @@ var<workgroup> row_mean: f32;
56
  var<workgroup> row_inv: f32;
57
 
58
  {% set xNumel = namespace(value=1) %}
59
- {% for dim in source.xShape %}
60
  {% set xNumel.value = xNumel.value * dim %}
61
  {% endfor %}
62
  {% set scaleNumel = namespace(value=1) %}
63
- {% for dim in source.scaleShape %}
64
  {% set scaleNumel.value = scaleNumel.value * dim %}
65
  {% endfor %}
66
  {% if scaleNumel.value != 1 %}
67
- {{ offset_fn("scale_offset", source.scaleShape, source.scaleShape | length, source.scaleShape == source.xShape, scaleNumel.value, source.xShape, source.xShape | length, xNumel.value) }}
68
  {% endif %}
69
 
70
  {% if hasBias %}
71
  {% set biasNumel = namespace(value=1) %}
72
- {% for dim in source.biasShape %}
73
  {% set biasNumel.value = biasNumel.value * dim %}
74
  {% endfor %}
75
  {% if biasNumel.value != 1 %}
76
- {{ offset_fn("bias_offset", source.biasShape, source.biasShape | length, source.biasShape == source.xShape, biasNumel.value, source.xShape, source.xShape | length, xNumel.value) }}
77
  {% endif %}
78
 
79
  {% endif %}
@@ -183,9 +180,9 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
183
  for (var d = tid; d < HIDDEN; d = d + WG) {
184
  let index = base + d;
185
  let normalized = (f32(x[index]) - row_mean) * row_inv;
186
- var value = normalized * f32(scale[{% if scaleNumel.value == 1 %}0u{% else %}{{ broadcast_offset_call("scale_offset", source.scaleShape, source.xShape, "index") }}{% endif %}]);
187
  {% if hasBias %}
188
- value = value + f32(bias[{% if biasNumel.value == 1 %}0u{% else %}{{ broadcast_offset_call("bias_offset", source.biasShape, source.xShape, "index") }}{% endif %}]);
189
  {% endif %}
190
  y[index] = {{ scalar }}(value);
191
  }
 
 
 
 
1
  {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
2
  fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
3
  {% if out_numel == 0 %}
 
53
  var<workgroup> row_inv: f32;
54
 
55
  {% set xNumel = namespace(value=1) %}
56
+ {% for dim in xShape %}
57
  {% set xNumel.value = xNumel.value * dim %}
58
  {% endfor %}
59
  {% set scaleNumel = namespace(value=1) %}
60
+ {% for dim in scaleShape %}
61
  {% set scaleNumel.value = scaleNumel.value * dim %}
62
  {% endfor %}
63
  {% if scaleNumel.value != 1 %}
64
+ {{ offset_fn("scale_offset", scaleShape, scaleShape | length, scaleShape == xShape, scaleNumel.value, xShape, xShape | length, xNumel.value) }}
65
  {% endif %}
66
 
67
  {% if hasBias %}
68
  {% set biasNumel = namespace(value=1) %}
69
+ {% for dim in biasShape %}
70
  {% set biasNumel.value = biasNumel.value * dim %}
71
  {% endfor %}
72
  {% if biasNumel.value != 1 %}
73
+ {{ offset_fn("bias_offset", biasShape, biasShape | length, biasShape == xShape, biasNumel.value, xShape, xShape | length, xNumel.value) }}
74
  {% endif %}
75
 
76
  {% endif %}
 
180
  for (var d = tid; d < HIDDEN; d = d + WG) {
181
  let index = base + d;
182
  let normalized = (f32(x[index]) - row_mean) * row_inv;
183
+ var value = normalized * f32(scale[{% if scaleNumel.value == 1 %}0u{% else %}{{ broadcast_offset_call("scale_offset", scaleShape, xShape, "index") }}{% endif %}]);
184
  {% if hasBias %}
185
+ value = value + f32(bias[{% if biasNumel.value == 1 %}0u{% else %}{{ broadcast_offset_call("bias_offset", biasShape, xShape, "index") }}{% endif %}]);
186
  {% endif %}
187
  y[index] = {{ scalar }}(value);
188
  }
build/webgpu/manifest.json CHANGED
@@ -2,655 +2,354 @@
2
  "domain": "ai.onnx",
3
  "name": "LayerNormalization",
4
  "sinceVersion": 17,
5
- "description": "Normalizes a tensor along a suffix of axes starting at `axis` by subtracting the mean and dividing by the square root of the variance plus `epsilon`, then scales and optionally shifts the result with learnable `Scale` and `B` tensors. The output `Y` has the same shape as `X`; optional outputs `Mean` and `InvStdDev` expose the per-normalization-group statistics computed during normalization.",
6
- "inputs": [
7
- { "role": "X", "dtype": "T", "description": "Tensor to be normalized." },
8
- { "role": "Scale", "dtype": "T", "description": "Scale tensor applied after normalization." },
9
- { "role": "B", "dtype": "T", "optional": true, "description": "Optional bias tensor added after scaling." }
10
- ],
11
- "outputs": [
12
- {
13
- "role": "Y",
14
- "dtype": "T",
15
- "rank": "ranks.X",
16
- "description": "Normalized and scaled output tensor; same shape as X.",
17
- "shape": "shapes.X"
18
- },
19
- {
20
- "role": "Mean",
21
- "dtype": "float32",
22
- "optional": true,
23
- "description": "Per-normalization-group mean in the ONNX broadcastable keepdims shape: dimensions before `axis` are preserved and dimensions from `axis` onward are 1.",
24
- "rank": "ranks.X"
25
- },
26
- {
27
- "role": "InvStdDev",
28
- "dtype": "float32",
29
- "optional": true,
30
- "description": "Per-normalization-group reciprocal standard deviation `1 / sqrt(variance + epsilon)`, returned in the same ONNX broadcastable keepdims shape as `Mean`.",
31
- "rank": "ranks.X"
32
- }
33
- ],
34
- "attributes": { "axis": -1, "epsilon": 0.00001, "stash_type": 1 },
35
- "attributeDescriptions": {
36
- "axis": "The first axis of the normalization range; all axes from `axis` to the last are normalized together. Negative values count from the end; the default `-1` normalizes only the last dimension.",
37
- "epsilon": "Small constant added to the variance before taking the square root to avoid division by zero.",
38
- "stash_type": "TensorProto element type used for the normalization stage and optional statistics; the implemented ONNX route supports the standard float32 value (`1`)."
39
  },
 
40
  "attributeConstraints": { "stash_type": { "values": [1] } },
41
  "typeConstraints": { "T": ["float32", "float16"] },
42
- "args": {
43
- "x": { "kind": "tensor", "semantic": "X", "role": "input" },
44
- "scale": { "kind": "tensor", "semantic": "Scale", "role": "input" },
45
- "b": { "kind": "tensor", "semantic": "B", "role": "input", "required": false },
46
- "y": { "kind": "tensor", "semantic": "Y", "role": "output" },
47
- "mean": { "kind": "tensor", "semantic": "Mean", "role": "output", "required": false },
48
- "invStdDev": { "kind": "tensor", "semantic": "InvStdDev", "role": "output", "required": false }
49
- },
50
- "tunables": { "MAX_WORKGROUP_SIZE": 256, "SCALAR_FAST_MAX_HIDDEN": 1024 },
51
  "derive": {
52
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
53
  "normWorkgroupCap": "min(tunables.MAX_WORKGROUP_SIZE, deviceWorkgroupCap)",
54
  "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
55
- "lastAxisWg": "min(normWorkgroupCap, pow2ceil(dim(shapes.X, -1)))",
56
- "lastAxisWgVec4": "min(normWorkgroupCap, pow2ceil(dim(shapes.X, -1) / 4))",
57
- "lastAxisContractOk": "ranks.X >= 1 and ranks.Y == ranks.X and numel(shapes.X) == numel(shapes.Y) and (attrs.axis == -1 or attrs.axis == ranks.X - 1)",
58
- "suffixAxisContractOk": "ranks.X >= 2 and ranks.Y == ranks.X and numel(shapes.X) == numel(shapes.Y) and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and not (attrs.axis == -1 or attrs.axis == ranks.X - 1)",
59
- "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.X",
60
- "normRows": "numel(shapes.X) / max(1, dim(shapes.X, -1)) if lastAxisContractOk else outer(shapes.X, axisNorm)",
61
- "normRowStride": "max(1, min(normRows, device.limits.maxComputeWorkgroupsPerDimension))",
62
- "suffixAxisSize": "numel(shapes.X) / max(1, outer(shapes.X, axisNorm))",
63
  "suffixAxisWg": "min(normWorkgroupCap, pow2ceil(suffixAxisSize))",
64
  "suffixAxisWgVec4": "min(normWorkgroupCap, pow2ceil(suffixAxisSize / 4))",
65
- "genericHiddenSize": "dim(shapes.X, -1) if lastAxisContractOk else suffixAxisSize",
66
  "genericWorkgroupSize": "lastAxisWg if lastAxisContractOk else suffixAxisWg",
67
- "scaleExactOk": "ranks.X >= 1 and ranks.Scale >= 1 and numel(shapes.Scale) == dim(shapes.X, -1) and dim(shapes.Scale, -1) == dim(shapes.X, -1)",
68
- "scaleBroadcastOk": "ranks.Scale >= 0 and ranks.Scale <= ranks.X and broadcastable(shapes.Scale, shapes.X)",
69
- "biasExactOk": "present.b and ranks.X >= 1 and ranks.B >= 1 and numel(shapes.B) == dim(shapes.X, -1) and dim(shapes.B, -1) == dim(shapes.X, -1)",
70
- "biasBroadcastOk": "present.b and ranks.B >= 0 and ranks.B <= ranks.X and broadcastable(shapes.B, shapes.X)",
71
- "suffixScaleExactOk": "suffixAxisContractOk and scaleBroadcastOk and numel(shapes.Scale) == suffixAxisSize",
72
- "suffixBiasExactOk": "present.b and suffixAxisContractOk and biasBroadcastOk and numel(shapes.B) == suffixAxisSize",
73
  "lastAxisExactScaleOk": "lastAxisContractOk and scaleExactOk",
74
  "lastAxisBroadcastScaleOk": "lastAxisContractOk and scaleBroadcastOk",
75
  "suffixAxisBroadcastScaleOk": "suffixAxisContractOk and scaleBroadcastOk",
76
  "suffixAxisExactAffineOk": "suffixAxisContractOk and suffixScaleExactOk and suffixBiasExactOk",
77
- "lastAxisScalarFastOk": "dtypes.T == \"f16\" or dim(shapes.X, -1) <= tunables.SCALAR_FAST_MAX_HIDDEN",
78
  "noStatsOutputs": "not present.mean and not present.invStdDev",
79
  "meanOnlyOutputs": "present.mean and not present.invStdDev",
80
  "invStdOnlyOutputs": "not present.mean and present.invStdDev",
81
  "fullStatsOutputs": "present.mean and present.invStdDev",
82
- "statsRowsOk": "fullStatsOutputs and ranks.X >= 1 and numel(shapes.Mean) == normRows and numel(shapes.InvStdDev) == normRows",
83
- "meanRowsOk": "present.mean and ranks.X >= 1 and numel(shapes.Mean) == normRows",
84
- "invStdRowsOk": "present.invStdDev and ranks.X >= 1 and numel(shapes.InvStdDev) == normRows",
85
- "statsOuterOk": "fullStatsOutputs and ranks.X >= 2 and numel(shapes.Mean) == normRows and numel(shapes.InvStdDev) == normRows"
86
  },
87
- "bindingSets": {
88
- "vec4Affine": [
89
- {
90
- "name": "x",
91
- "arg": "x",
92
- "semantic": "X",
93
- "buffer": { "type": "read-only-storage" },
94
- "elementType": "$vectorScalar"
95
- },
96
- {
97
- "name": "scale",
98
- "arg": "scale",
99
- "semantic": "Scale",
100
- "buffer": { "type": "read-only-storage" },
101
- "elementType": "$vectorScalar"
102
- },
103
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
104
- {
105
- "name": "params",
106
- "semantic": "kernel.params",
107
- "buffer": { "type": "uniform" },
108
- "struct": {
109
- "name": "Params",
110
- "fields": [
111
- { "name": "rows", "type": "u32", "value": "normRows" },
112
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
113
- ]
114
- }
115
- }
116
- ],
117
- "vec4AffineBias": [
118
- {
119
- "name": "x",
120
- "arg": "x",
121
- "semantic": "X",
122
- "buffer": { "type": "read-only-storage" },
123
- "elementType": "$vectorScalar"
124
- },
125
- {
126
- "name": "scale",
127
- "arg": "scale",
128
- "semantic": "Scale",
129
- "buffer": { "type": "read-only-storage" },
130
- "elementType": "$vectorScalar"
131
- },
132
- {
133
- "name": "bias",
134
- "arg": "b",
135
- "semantic": "B",
136
- "buffer": { "type": "read-only-storage" },
137
- "elementType": "$vectorScalar"
138
- },
139
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
140
- {
141
- "name": "params",
142
- "semantic": "kernel.params",
143
- "buffer": { "type": "uniform" },
144
- "struct": {
145
- "name": "Params",
146
- "fields": [
147
- { "name": "rows", "type": "u32", "value": "normRows" },
148
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
149
- ]
150
- }
151
- }
152
- ],
153
- "vec4AffineStats": [
154
- {
155
- "name": "x",
156
- "arg": "x",
157
- "semantic": "X",
158
- "buffer": { "type": "read-only-storage" },
159
- "elementType": "$vectorScalar"
160
- },
161
- {
162
- "name": "scale",
163
- "arg": "scale",
164
- "semantic": "Scale",
165
- "buffer": { "type": "read-only-storage" },
166
- "elementType": "$vectorScalar"
167
- },
168
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
169
- { "name": "mean_out", "arg": "mean", "semantic": "Mean", "buffer": { "type": "storage" }, "elementType": "f32" },
170
- {
171
- "name": "inv_std_out",
172
- "arg": "invStdDev",
173
- "semantic": "InvStdDev",
174
- "buffer": { "type": "storage" },
175
- "elementType": "f32"
176
- },
177
- {
178
- "name": "params",
179
- "semantic": "kernel.params",
180
- "buffer": { "type": "uniform" },
181
- "struct": {
182
- "name": "Params",
183
- "fields": [
184
- { "name": "rows", "type": "u32", "value": "normRows" },
185
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
186
- ]
187
- }
188
- }
189
- ],
190
- "vec4AffineBiasStats": [
191
- {
192
- "name": "x",
193
- "arg": "x",
194
- "semantic": "X",
195
- "buffer": { "type": "read-only-storage" },
196
- "elementType": "$vectorScalar"
197
- },
198
- {
199
- "name": "scale",
200
- "arg": "scale",
201
- "semantic": "Scale",
202
- "buffer": { "type": "read-only-storage" },
203
- "elementType": "$vectorScalar"
204
- },
205
- {
206
- "name": "bias",
207
- "arg": "b",
208
- "semantic": "B",
209
- "buffer": { "type": "read-only-storage" },
210
- "elementType": "$vectorScalar"
211
- },
212
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
213
- { "name": "mean_out", "arg": "mean", "semantic": "Mean", "buffer": { "type": "storage" }, "elementType": "f32" },
214
- {
215
- "name": "inv_std_out",
216
- "arg": "invStdDev",
217
- "semantic": "InvStdDev",
218
- "buffer": { "type": "storage" },
219
- "elementType": "f32"
220
- },
221
- {
222
- "name": "params",
223
- "semantic": "kernel.params",
224
- "buffer": { "type": "uniform" },
225
- "struct": {
226
- "name": "Params",
227
- "fields": [
228
- { "name": "rows", "type": "u32", "value": "normRows" },
229
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
230
- ]
231
- }
232
- }
233
- ],
234
- "scalarAffineMean": [
235
- { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
236
- {
237
- "name": "scale",
238
- "arg": "scale",
239
- "semantic": "Scale",
240
- "buffer": { "type": "read-only-storage" },
241
- "elementType": "$scalar"
242
- },
243
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
244
- { "name": "mean_out", "arg": "mean", "semantic": "Mean", "buffer": { "type": "storage" }, "elementType": "f32" },
245
- {
246
- "name": "params",
247
- "semantic": "kernel.params",
248
- "buffer": { "type": "uniform" },
249
- "struct": {
250
- "name": "Params",
251
- "fields": [
252
- { "name": "rows", "type": "u32", "value": "normRows" },
253
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
254
- ]
255
- }
256
- }
257
- ],
258
- "scalarAffineInvStd": [
259
- { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
260
- {
261
- "name": "scale",
262
- "arg": "scale",
263
- "semantic": "Scale",
264
- "buffer": { "type": "read-only-storage" },
265
- "elementType": "$scalar"
266
- },
267
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
268
- {
269
- "name": "inv_std_out",
270
- "arg": "invStdDev",
271
- "semantic": "InvStdDev",
272
- "buffer": { "type": "storage" },
273
- "elementType": "f32"
274
- },
275
- {
276
- "name": "params",
277
- "semantic": "kernel.params",
278
- "buffer": { "type": "uniform" },
279
- "struct": {
280
- "name": "Params",
281
- "fields": [
282
- { "name": "rows", "type": "u32", "value": "normRows" },
283
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
284
- ]
285
- }
286
- }
287
- ],
288
- "scalarAffineBiasMean": [
289
- { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
290
- {
291
- "name": "scale",
292
- "arg": "scale",
293
- "semantic": "Scale",
294
- "buffer": { "type": "read-only-storage" },
295
- "elementType": "$scalar"
296
- },
297
- {
298
- "name": "bias",
299
- "arg": "b",
300
- "semantic": "B",
301
- "buffer": { "type": "read-only-storage" },
302
- "elementType": "$scalar"
303
- },
304
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
305
- { "name": "mean_out", "arg": "mean", "semantic": "Mean", "buffer": { "type": "storage" }, "elementType": "f32" },
306
- {
307
- "name": "params",
308
- "semantic": "kernel.params",
309
- "buffer": { "type": "uniform" },
310
- "struct": {
311
- "name": "Params",
312
- "fields": [
313
- { "name": "rows", "type": "u32", "value": "normRows" },
314
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
315
- ]
316
- }
317
- }
318
- ],
319
- "scalarAffineBiasInvStd": [
320
- { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
321
- {
322
- "name": "scale",
323
- "arg": "scale",
324
- "semantic": "Scale",
325
- "buffer": { "type": "read-only-storage" },
326
- "elementType": "$scalar"
327
- },
328
- {
329
- "name": "bias",
330
- "arg": "b",
331
- "semantic": "B",
332
- "buffer": { "type": "read-only-storage" },
333
- "elementType": "$scalar"
334
- },
335
- { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
336
- {
337
- "name": "inv_std_out",
338
- "arg": "invStdDev",
339
- "semantic": "InvStdDev",
340
- "buffer": { "type": "storage" },
341
- "elementType": "f32"
342
- },
343
- {
344
- "name": "params",
345
- "semantic": "kernel.params",
346
- "buffer": { "type": "uniform" },
347
- "struct": {
348
- "name": "Params",
349
- "fields": [
350
- { "name": "rows", "type": "u32", "value": "normRows" },
351
- { "name": "rowStride", "type": "u32", "value": "normRowStride" }
352
- ]
353
- }
354
- }
355
- ]
356
  },
357
  "variants": [
358
  {
359
  "id": "last_axis_row_vec4",
360
  "priority": 110,
361
- "when": ["f16Ok(dtypes.T)", "not present.b and noStatsOutputs", "lastAxisExactScaleOk", "dim(shapes.X, -1) % 4 == 0"],
362
- "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
363
  "passes": [
364
  {
365
  "id": "main",
366
  "name": "LayerNormalization.LastAxisRowVec4",
367
- "source": {
368
- "shader": "norm-row-stats.wgsl.jinja",
369
- "inputs": {
370
- "mode": "\"layer\"",
371
- "vec4": true,
372
- "hasBias": false,
373
- "writeStats": false,
374
- "scalar": "dtypes.T",
375
- "usesF16": "dtypes.T == \"f16\"",
376
- "hidden": "dim(shapes.X, -1)",
377
- "wg": "lastAxisWgVec4",
378
- "epsilon": "attrs.epsilon",
379
- "hiddenVec": "dim(shapes.X, -1) / 4",
380
- "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
381
- "combineSubgroups": "hasSubgroupId"
382
- }
383
  },
384
- "subgroupCollectivesWidth": "portable",
385
- "bindings": "vec4Affine",
386
- "dispatch": { "workgroups": "normRows" }
387
  }
388
  ]
389
  },
390
  {
391
  "id": "last_axis_row",
392
  "priority": 100,
393
- "when": ["not present.b and noStatsOutputs", "lastAxisExactScaleOk", "f16Ok(dtypes.T)"],
394
  "demoteWhen": ["not lastAxisScalarFastOk"],
395
- "constants": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
396
  "passes": [
397
  {
398
  "id": "main",
399
  "name": "LayerNormalization.LastAxisRow",
400
- "source": {
401
- "shader": "norm-row-stats.wgsl.jinja",
402
- "inputs": {
403
- "mode": "\"layer\"",
404
- "vec4": false,
405
- "hasBias": false,
406
- "writeStats": false,
407
- "scalar": "dtypes.T",
408
- "usesF16": "dtypes.T == \"f16\"",
409
- "hidden": "dim(shapes.X, -1)",
410
- "wg": "lastAxisWg",
411
- "epsilon": "attrs.epsilon",
412
- "combineSubgroups": "hasSubgroupId"
413
- }
414
  },
415
- "subgroupCollectivesWidth": "portable",
416
- "bindings": "vec4Affine",
417
- "dispatch": { "workgroups": "normRows" }
418
  }
419
  ]
420
  },
421
  {
422
  "id": "last_axis_bias_row_vec4",
423
  "priority": 111,
424
- "when": ["f16Ok(dtypes.T)", "present.b and noStatsOutputs and biasExactOk", "lastAxisExactScaleOk", "dim(shapes.X, -1) % 4 == 0"],
425
- "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
426
  "passes": [
427
  {
428
  "id": "main",
429
  "name": "LayerNormalization.LastAxisRowVec4",
430
- "source": {
431
- "shader": "norm-row-stats.wgsl.jinja",
432
- "inputs": {
433
- "mode": "\"layer\"",
434
- "vec4": true,
435
- "hasBias": true,
436
- "writeStats": false,
437
- "scalar": "dtypes.T",
438
- "usesF16": "dtypes.T == \"f16\"",
439
- "hidden": "dim(shapes.X, -1)",
440
- "wg": "lastAxisWgVec4",
441
- "epsilon": "attrs.epsilon",
442
- "hiddenVec": "dim(shapes.X, -1) / 4",
443
- "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
444
- "combineSubgroups": "hasSubgroupId"
445
- }
446
  },
447
- "subgroupCollectivesWidth": "portable",
448
- "bindings": "vec4AffineBias",
449
- "dispatch": { "workgroups": "normRows" }
450
  }
451
  ]
452
  },
453
  {
454
  "id": "last_axis_bias_row",
455
  "priority": 101,
456
- "when": ["present.b and noStatsOutputs and biasExactOk", "lastAxisExactScaleOk", "f16Ok(dtypes.T)"],
457
  "demoteWhen": ["not lastAxisScalarFastOk"],
458
- "constants": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
459
  "passes": [
460
  {
461
  "id": "main",
462
  "name": "LayerNormalization.LastAxisRow",
463
- "source": {
464
- "shader": "norm-row-stats.wgsl.jinja",
465
- "inputs": {
466
- "mode": "\"layer\"",
467
- "vec4": false,
468
- "hasBias": true,
469
- "writeStats": false,
470
- "scalar": "dtypes.T",
471
- "usesF16": "dtypes.T == \"f16\"",
472
- "hidden": "dim(shapes.X, -1)",
473
- "wg": "lastAxisWg",
474
- "epsilon": "attrs.epsilon",
475
- "combineSubgroups": "hasSubgroupId"
476
- }
477
  },
478
- "subgroupCollectivesWidth": "portable",
479
- "bindings": "vec4AffineBias",
480
- "dispatch": { "workgroups": "normRows" }
481
  }
482
  ]
483
  },
484
  {
485
  "id": "last_axis_stats_row_vec4",
486
  "priority": 112,
487
- "when": ["f16Ok(dtypes.T)", "not present.b and fullStatsOutputs and statsRowsOk", "lastAxisExactScaleOk", "dim(shapes.X, -1) % 4 == 0"],
488
- "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
489
  "passes": [
490
  {
491
  "id": "main",
492
  "name": "LayerNormalization.LastAxisRowVec4",
493
- "source": {
494
- "shader": "norm-row-stats.wgsl.jinja",
495
- "inputs": {
496
- "mode": "\"layer\"",
497
- "vec4": true,
498
- "hasBias": false,
499
- "writeStats": true,
500
- "scalar": "dtypes.T",
501
- "usesF16": "dtypes.T == \"f16\"",
502
- "hidden": "dim(shapes.X, -1)",
503
- "wg": "lastAxisWgVec4",
504
- "epsilon": "attrs.epsilon",
505
- "hiddenVec": "dim(shapes.X, -1) / 4",
506
- "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
507
- "combineSubgroups": "hasSubgroupId"
508
- }
509
  },
510
- "subgroupCollectivesWidth": "portable",
511
- "bindings": "vec4AffineStats",
512
- "dispatch": { "workgroups": "normRows" }
513
  }
514
  ]
515
  },
516
  {
517
  "id": "last_axis_stats_row",
518
  "priority": 102,
519
- "when": ["not present.b and fullStatsOutputs and statsRowsOk", "lastAxisExactScaleOk", "f16Ok(dtypes.T)"],
520
  "demoteWhen": ["not lastAxisScalarFastOk"],
521
- "constants": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
522
  "passes": [
523
  {
524
  "id": "main",
525
  "name": "LayerNormalization.LastAxisRow",
526
- "source": {
527
- "shader": "norm-row-stats.wgsl.jinja",
528
- "inputs": {
529
- "mode": "\"layer\"",
530
- "vec4": false,
531
- "hasBias": false,
532
- "writeStats": true,
533
- "scalar": "dtypes.T",
534
- "usesF16": "dtypes.T == \"f16\"",
535
- "hidden": "dim(shapes.X, -1)",
536
- "wg": "lastAxisWg",
537
- "epsilon": "attrs.epsilon",
538
- "combineSubgroups": "hasSubgroupId"
539
- }
540
  },
541
- "subgroupCollectivesWidth": "portable",
542
- "bindings": "vec4AffineStats",
543
- "dispatch": { "workgroups": "normRows" }
544
  }
545
  ]
546
  },
547
  {
548
  "id": "last_axis_bias_stats_row_vec4",
549
  "priority": 113,
550
- "when": ["f16Ok(dtypes.T)", "present.b and fullStatsOutputs and biasExactOk and statsRowsOk", "lastAxisExactScaleOk", "dim(shapes.X, -1) % 4 == 0"],
551
- "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
552
  "passes": [
553
  {
554
  "id": "main",
555
  "name": "LayerNormalization.LastAxisRowVec4",
556
- "source": {
557
- "shader": "norm-row-stats.wgsl.jinja",
558
- "inputs": {
559
- "mode": "\"layer\"",
560
- "vec4": true,
561
- "hasBias": true,
562
- "writeStats": true,
563
- "scalar": "dtypes.T",
564
- "usesF16": "dtypes.T == \"f16\"",
565
- "hidden": "dim(shapes.X, -1)",
566
- "wg": "lastAxisWgVec4",
567
- "epsilon": "attrs.epsilon",
568
- "hiddenVec": "dim(shapes.X, -1) / 4",
569
- "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
570
- "combineSubgroups": "hasSubgroupId"
571
- }
572
  },
573
- "subgroupCollectivesWidth": "portable",
574
- "bindings": "vec4AffineBiasStats",
575
- "dispatch": { "workgroups": "normRows" }
576
  }
577
  ]
578
  },
579
  {
580
  "id": "last_axis_bias_stats_row",
581
  "priority": 103,
582
- "when": ["present.b and fullStatsOutputs and biasExactOk and statsRowsOk", "lastAxisExactScaleOk", "f16Ok(dtypes.T)"],
583
  "demoteWhen": ["not lastAxisScalarFastOk"],
584
- "constants": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
585
  "passes": [
586
  {
587
  "id": "main",
588
  "name": "LayerNormalization.LastAxisRow",
589
- "source": {
590
- "shader": "norm-row-stats.wgsl.jinja",
591
- "inputs": {
592
- "mode": "\"layer\"",
593
- "vec4": false,
594
- "hasBias": true,
595
- "writeStats": true,
596
- "scalar": "dtypes.T",
597
- "usesF16": "dtypes.T == \"f16\"",
598
- "hidden": "dim(shapes.X, -1)",
599
- "wg": "lastAxisWg",
600
- "epsilon": "attrs.epsilon",
601
- "combineSubgroups": "hasSubgroupId"
602
- }
603
  },
604
- "subgroupCollectivesWidth": "portable",
605
- "bindings": "vec4AffineBiasStats",
606
- "dispatch": { "workgroups": "normRows" }
607
  }
608
  ]
609
  },
610
  {
611
  "id": "suffix_axis_bias_exact_row_vec4",
612
  "priority": 121,
613
- "when": ["f16Ok(dtypes.T)", "present.b", "noStatsOutputs", "suffixAxisExactAffineOk", "suffixAxisSize % 4 == 0"],
614
- "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
615
  "passes": [
616
  {
617
  "id": "main",
618
  "name": "LayerNormalization.SuffixAxisRowVec4",
619
- "source": {
620
- "shader": "norm-row-stats.wgsl.jinja",
621
- "inputs": {
622
- "mode": "\"layer\"",
623
- "vec4": true,
624
- "hasBias": true,
625
- "writeStats": false,
626
- "scalar": "dtypes.T",
627
- "usesF16": "dtypes.T == \"f16\"",
628
- "hidden": "suffixAxisSize",
629
- "wg": "suffixAxisWgVec4",
630
- "epsilon": "attrs.epsilon",
631
- "hiddenVec": "suffixAxisSize / 4",
632
- "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
633
- "combineSubgroups": "hasSubgroupId"
634
- }
635
  },
636
- "subgroupCollectivesWidth": "portable",
637
- "bindings": "vec4AffineBias",
638
- "dispatch": { "workgroups": "normRows" }
639
  }
640
  ]
641
  },
642
  {
643
  "id": "last_axis",
644
  "priority": 0,
645
- "when": ["not present.b", "noStatsOutputs", "lastAxisBroadcastScaleOk", "f16Ok(dtypes.T)"],
646
- "constants": {
647
  "hasBias": false,
648
  "writeMean": false,
649
  "writeInvStdDev": false,
650
  "scalar": "dtypes.T",
651
  "vectorScalar": "dtypes.T",
652
- "usesF16": "dtypes.T == \"f16\"",
653
- "hiddenSize": "dim(shapes.X, -1)",
654
  "workgroupSize": "lastAxisWg",
655
  "epsilon": "attrs.epsilon"
656
  },
@@ -658,27 +357,24 @@
658
  {
659
  "id": "main",
660
  "name": "LayerNormalization",
661
- "source": {
662
- "shader": "layer-normalization.wgsl.jinja",
663
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale" }
664
- },
665
- "bindings": "vec4Affine",
666
- "dispatch": { "workgroups": "normRows" }
667
  }
668
  ]
669
  },
670
  {
671
  "id": "last_axis_bias",
672
  "priority": 10,
673
- "when": ["present.b", "noStatsOutputs", "lastAxisBroadcastScaleOk", "biasBroadcastOk", "f16Ok(dtypes.T)"],
674
- "constants": {
675
  "hasBias": true,
676
  "writeMean": false,
677
  "writeInvStdDev": false,
678
  "scalar": "dtypes.T",
679
  "vectorScalar": "dtypes.T",
680
- "usesF16": "dtypes.T == \"f16\"",
681
- "hiddenSize": "dim(shapes.X, -1)",
682
  "workgroupSize": "lastAxisWg",
683
  "epsilon": "attrs.epsilon"
684
  },
@@ -686,27 +382,24 @@
686
  {
687
  "id": "main",
688
  "name": "LayerNormalization",
689
- "source": {
690
- "shader": "layer-normalization.wgsl.jinja",
691
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale", "biasShape": "shapes.B" }
692
- },
693
- "bindings": "vec4AffineBias",
694
- "dispatch": { "workgroups": "normRows" }
695
  }
696
  ]
697
  },
698
  {
699
  "id": "last_axis_stats",
700
  "priority": 20,
701
- "when": ["not present.b", "fullStatsOutputs", "lastAxisBroadcastScaleOk", "statsRowsOk", "f16Ok(dtypes.T)"],
702
- "constants": {
703
  "hasBias": false,
704
  "writeMean": true,
705
  "writeInvStdDev": true,
706
  "scalar": "dtypes.T",
707
  "vectorScalar": "dtypes.T",
708
- "usesF16": "dtypes.T == \"f16\"",
709
- "hiddenSize": "dim(shapes.X, -1)",
710
  "workgroupSize": "lastAxisWg",
711
  "epsilon": "attrs.epsilon"
712
  },
@@ -714,27 +407,24 @@
714
  {
715
  "id": "main",
716
  "name": "LayerNormalization",
717
- "source": {
718
- "shader": "layer-normalization.wgsl.jinja",
719
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale" }
720
- },
721
- "bindings": "vec4AffineStats",
722
- "dispatch": { "workgroups": "normRows" }
723
  }
724
  ]
725
  },
726
  {
727
  "id": "last_axis_bias_stats",
728
  "priority": 30,
729
- "when": ["present.b", "fullStatsOutputs", "lastAxisBroadcastScaleOk", "biasBroadcastOk", "statsRowsOk", "f16Ok(dtypes.T)"],
730
- "constants": {
731
  "hasBias": true,
732
  "writeMean": true,
733
  "writeInvStdDev": true,
734
  "scalar": "dtypes.T",
735
  "vectorScalar": "dtypes.T",
736
- "usesF16": "dtypes.T == \"f16\"",
737
- "hiddenSize": "dim(shapes.X, -1)",
738
  "workgroupSize": "lastAxisWg",
739
  "epsilon": "attrs.epsilon"
740
  },
@@ -742,26 +432,23 @@
742
  {
743
  "id": "main",
744
  "name": "LayerNormalization",
745
- "source": {
746
- "shader": "layer-normalization.wgsl.jinja",
747
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale", "biasShape": "shapes.B" }
748
- },
749
- "bindings": "vec4AffineBiasStats",
750
- "dispatch": { "workgroups": "normRows" }
751
  }
752
  ]
753
  },
754
  {
755
  "id": "suffix_axis",
756
  "priority": 40,
757
- "when": ["not present.b", "noStatsOutputs", "suffixAxisBroadcastScaleOk", "f16Ok(dtypes.T)"],
758
- "constants": {
759
  "hasBias": false,
760
  "writeMean": false,
761
  "writeInvStdDev": false,
762
  "scalar": "dtypes.T",
763
  "vectorScalar": "dtypes.T",
764
- "usesF16": "dtypes.T == \"f16\"",
765
  "hiddenSize": "suffixAxisSize",
766
  "workgroupSize": "suffixAxisWg",
767
  "epsilon": "attrs.epsilon"
@@ -770,26 +457,23 @@
770
  {
771
  "id": "main",
772
  "name": "LayerNormalization.SuffixAxis",
773
- "source": {
774
- "shader": "layer-normalization.wgsl.jinja",
775
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale" }
776
- },
777
- "bindings": "vec4Affine",
778
- "dispatch": { "workgroups": "normRows" }
779
  }
780
  ]
781
  },
782
  {
783
  "id": "suffix_axis_bias",
784
  "priority": 50,
785
- "when": ["present.b", "noStatsOutputs", "suffixAxisBroadcastScaleOk", "biasBroadcastOk", "f16Ok(dtypes.T)"],
786
- "constants": {
787
  "hasBias": true,
788
  "writeMean": false,
789
  "writeInvStdDev": false,
790
  "scalar": "dtypes.T",
791
  "vectorScalar": "dtypes.T",
792
- "usesF16": "dtypes.T == \"f16\"",
793
  "hiddenSize": "suffixAxisSize",
794
  "workgroupSize": "suffixAxisWg",
795
  "epsilon": "attrs.epsilon"
@@ -798,26 +482,23 @@
798
  {
799
  "id": "main",
800
  "name": "LayerNormalization.SuffixAxisBias",
801
- "source": {
802
- "shader": "layer-normalization.wgsl.jinja",
803
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale", "biasShape": "shapes.B" }
804
- },
805
- "bindings": "vec4AffineBias",
806
- "dispatch": { "workgroups": "normRows" }
807
  }
808
  ]
809
  },
810
  {
811
  "id": "suffix_axis_stats",
812
  "priority": 45,
813
- "when": ["not present.b", "fullStatsOutputs", "suffixAxisBroadcastScaleOk", "statsOuterOk", "f16Ok(dtypes.T)"],
814
- "constants": {
815
  "hasBias": false,
816
  "writeMean": true,
817
  "writeInvStdDev": true,
818
  "scalar": "dtypes.T",
819
  "vectorScalar": "dtypes.T",
820
- "usesF16": "dtypes.T == \"f16\"",
821
  "hiddenSize": "suffixAxisSize",
822
  "workgroupSize": "suffixAxisWg",
823
  "epsilon": "attrs.epsilon"
@@ -826,26 +507,23 @@
826
  {
827
  "id": "main",
828
  "name": "LayerNormalization.SuffixAxisStats",
829
- "source": {
830
- "shader": "layer-normalization.wgsl.jinja",
831
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale" }
832
- },
833
- "bindings": "vec4AffineStats",
834
- "dispatch": { "workgroups": "normRows" }
835
  }
836
  ]
837
  },
838
  {
839
  "id": "suffix_axis_bias_stats",
840
  "priority": 55,
841
- "when": ["present.b", "fullStatsOutputs", "suffixAxisBroadcastScaleOk", "biasBroadcastOk", "statsOuterOk", "f16Ok(dtypes.T)"],
842
- "constants": {
843
  "hasBias": true,
844
  "writeMean": true,
845
  "writeInvStdDev": true,
846
  "scalar": "dtypes.T",
847
  "vectorScalar": "dtypes.T",
848
- "usesF16": "dtypes.T == \"f16\"",
849
  "hiddenSize": "suffixAxisSize",
850
  "workgroupSize": "suffixAxisWg",
851
  "epsilon": "attrs.epsilon"
@@ -854,25 +532,22 @@
854
  {
855
  "id": "main",
856
  "name": "LayerNormalization.SuffixAxisBiasStats",
857
- "source": {
858
- "shader": "layer-normalization.wgsl.jinja",
859
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale", "biasShape": "shapes.B" }
860
- },
861
- "bindings": "vec4AffineBiasStats",
862
- "dispatch": { "workgroups": "normRows" }
863
  }
864
  ]
865
  },
866
  {
867
  "id": "mean_only",
868
  "priority": 31,
869
- "when": ["not present.b and meanOnlyOutputs and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk", "f16Ok(dtypes.T)"],
870
- "constants": {
871
  "hasBias": false,
872
  "writeMean": true,
873
  "writeInvStdDev": false,
874
  "scalar": "dtypes.T",
875
- "usesF16": "dtypes.T == \"f16\"",
876
  "hiddenSize": "genericHiddenSize",
877
  "workgroupSize": "genericWorkgroupSize",
878
  "epsilon": "attrs.epsilon"
@@ -881,25 +556,22 @@
881
  {
882
  "id": "main",
883
  "name": "LayerNormalization.MeanOnly",
884
- "source": {
885
- "shader": "layer-normalization.wgsl.jinja",
886
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale" }
887
- },
888
- "bindings": "scalarAffineMean",
889
- "dispatch": { "workgroups": "normRows" }
890
  }
891
  ]
892
  },
893
  {
894
  "id": "bias_mean_only",
895
  "priority": 32,
896
- "when": ["present.b and meanOnlyOutputs and biasBroadcastOk and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk", "f16Ok(dtypes.T)"],
897
- "constants": {
898
  "hasBias": true,
899
  "writeMean": true,
900
  "writeInvStdDev": false,
901
  "scalar": "dtypes.T",
902
- "usesF16": "dtypes.T == \"f16\"",
903
  "hiddenSize": "genericHiddenSize",
904
  "workgroupSize": "genericWorkgroupSize",
905
  "epsilon": "attrs.epsilon"
@@ -908,25 +580,22 @@
908
  {
909
  "id": "main",
910
  "name": "LayerNormalization.BiasMeanOnly",
911
- "source": {
912
- "shader": "layer-normalization.wgsl.jinja",
913
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale", "biasShape": "shapes.B" }
914
- },
915
- "bindings": "scalarAffineBiasMean",
916
- "dispatch": { "workgroups": "normRows" }
917
  }
918
  ]
919
  },
920
  {
921
  "id": "inv_std_dev_only",
922
  "priority": 33,
923
- "when": ["not present.b and invStdOnlyOutputs and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk", "f16Ok(dtypes.T)"],
924
- "constants": {
925
  "hasBias": false,
926
  "writeMean": false,
927
  "writeInvStdDev": true,
928
  "scalar": "dtypes.T",
929
- "usesF16": "dtypes.T == \"f16\"",
930
  "hiddenSize": "genericHiddenSize",
931
  "workgroupSize": "genericWorkgroupSize",
932
  "epsilon": "attrs.epsilon"
@@ -935,25 +604,22 @@
935
  {
936
  "id": "main",
937
  "name": "LayerNormalization.InvStdDevOnly",
938
- "source": {
939
- "shader": "layer-normalization.wgsl.jinja",
940
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale" }
941
- },
942
- "bindings": "scalarAffineInvStd",
943
- "dispatch": { "workgroups": "normRows" }
944
  }
945
  ]
946
  },
947
  {
948
  "id": "bias_inv_std_dev_only",
949
  "priority": 34,
950
- "when": ["present.b and invStdOnlyOutputs and biasBroadcastOk and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk", "f16Ok(dtypes.T)"],
951
- "constants": {
952
  "hasBias": true,
953
  "writeMean": false,
954
  "writeInvStdDev": true,
955
  "scalar": "dtypes.T",
956
- "usesF16": "dtypes.T == \"f16\"",
957
  "hiddenSize": "genericHiddenSize",
958
  "workgroupSize": "genericWorkgroupSize",
959
  "epsilon": "attrs.epsilon"
@@ -962,12 +628,10 @@
962
  {
963
  "id": "main",
964
  "name": "LayerNormalization.BiasInvStdDevOnly",
965
- "source": {
966
- "shader": "layer-normalization.wgsl.jinja",
967
- "inputs": { "xShape": "shapes.X", "scaleShape": "shapes.Scale", "biasShape": "shapes.B" }
968
- },
969
- "bindings": "scalarAffineBiasInvStd",
970
- "dispatch": { "workgroups": "normRows" }
971
  }
972
  ]
973
  }
 
2
  "domain": "ai.onnx",
3
  "name": "LayerNormalization",
4
  "sinceVersion": 17,
5
+ "inputs": {
6
+ "x": { "onnx": "X", "dtype": "T" },
7
+ "scale": { "onnx": "Scale", "dtype": "T" },
8
+ "b": { "onnx": "B", "dtype": "T", "optional": true }
9
+ },
10
+ "outputs": {
11
+ "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" },
12
+ "mean": { "onnx": "Mean", "dtype": "float32", "rank": "ranks.x", "optional": true },
13
+ "invStdDev": { "onnx": "InvStdDev", "dtype": "float32", "rank": "ranks.x", "optional": true }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  },
15
+ "attributes": { "axis": { "default": -1 }, "epsilon": { "default": 0.00001 }, "stash_type": { "default": 1 } },
16
  "attributeConstraints": { "stash_type": { "values": [1] } },
17
  "typeConstraints": { "T": ["float32", "float16"] },
18
+ "tunables": { "MAX_WORKGROUP_SIZE": { "default": 256 }, "SCALAR_FAST_MAX_HIDDEN": { "default": 1024 } },
 
 
 
 
 
 
 
 
19
  "derive": {
20
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
21
  "normWorkgroupCap": "min(tunables.MAX_WORKGROUP_SIZE, deviceWorkgroupCap)",
22
  "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
23
+ "lastAxisWg": "min(normWorkgroupCap, pow2ceil(dim(shapes.x, -1)))",
24
+ "lastAxisWgVec4": "min(normWorkgroupCap, pow2ceil(dim(shapes.x, -1) / 4))",
25
+ "lastAxisContractOk": "ranks.x >= 1 and ranks.y == ranks.x and numel(shapes.x) == numel(shapes.y) and (attrs.axis == -1 or attrs.axis == ranks.x - 1)",
26
+ "suffixAxisContractOk": "ranks.x >= 2 and ranks.y == ranks.x and numel(shapes.x) == numel(shapes.y) and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)",
27
+ "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
28
+ "normRows": "numel(shapes.x) / max(1, dim(shapes.x, -1)) if lastAxisContractOk else outer(shapes.x, axisNorm)",
29
+ "normRowStride": "max(1, min(normRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))",
30
+ "suffixAxisSize": "numel(shapes.x) / max(1, outer(shapes.x, axisNorm))",
31
  "suffixAxisWg": "min(normWorkgroupCap, pow2ceil(suffixAxisSize))",
32
  "suffixAxisWgVec4": "min(normWorkgroupCap, pow2ceil(suffixAxisSize / 4))",
33
+ "genericHiddenSize": "dim(shapes.x, -1) if lastAxisContractOk else suffixAxisSize",
34
  "genericWorkgroupSize": "lastAxisWg if lastAxisContractOk else suffixAxisWg",
35
+ "scaleExactOk": "ranks.x >= 1 and ranks.scale >= 1 and numel(shapes.scale) == dim(shapes.x, -1) and dim(shapes.scale, -1) == dim(shapes.x, -1)",
36
+ "scaleBroadcastOk": "ranks.scale >= 0 and ranks.scale <= ranks.x and broadcastable(shapes.scale, shapes.x)",
37
+ "biasExactOk": "present.b and ranks.x >= 1 and ranks.b >= 1 and numel(shapes.b) == dim(shapes.x, -1) and dim(shapes.b, -1) == dim(shapes.x, -1)",
38
+ "biasBroadcastOk": "present.b and ranks.b >= 0 and ranks.b <= ranks.x and broadcastable(shapes.b, shapes.x)",
39
+ "suffixScaleExactOk": "suffixAxisContractOk and scaleBroadcastOk and numel(shapes.scale) == suffixAxisSize",
40
+ "suffixBiasExactOk": "present.b and suffixAxisContractOk and biasBroadcastOk and numel(shapes.b) == suffixAxisSize",
41
  "lastAxisExactScaleOk": "lastAxisContractOk and scaleExactOk",
42
  "lastAxisBroadcastScaleOk": "lastAxisContractOk and scaleBroadcastOk",
43
  "suffixAxisBroadcastScaleOk": "suffixAxisContractOk and scaleBroadcastOk",
44
  "suffixAxisExactAffineOk": "suffixAxisContractOk and suffixScaleExactOk and suffixBiasExactOk",
45
+ "lastAxisScalarFastOk": "dtypes.T == \"f16\" or dim(shapes.x, -1) <= tunables.SCALAR_FAST_MAX_HIDDEN",
46
  "noStatsOutputs": "not present.mean and not present.invStdDev",
47
  "meanOnlyOutputs": "present.mean and not present.invStdDev",
48
  "invStdOnlyOutputs": "not present.mean and present.invStdDev",
49
  "fullStatsOutputs": "present.mean and present.invStdDev",
50
+ "statsRowsOk": "fullStatsOutputs and ranks.x >= 1 and numel(shapes.mean) == normRows and numel(shapes.invStdDev) == normRows",
51
+ "meanRowsOk": "present.mean and ranks.x >= 1 and numel(shapes.mean) == normRows",
52
+ "invStdRowsOk": "present.invStdDev and ranks.x >= 1 and numel(shapes.invStdDev) == normRows",
53
+ "statsOuterOk": "fullStatsOutputs and ranks.x >= 2 and numel(shapes.mean) == normRows and numel(shapes.invStdDev) == normRows"
54
  },
55
+ "when": ["f16Ok(dtypes.T)"],
56
+ "bindings": {
57
+ "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
58
+ "scale": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
59
+ "y": { "buffer": "storage", "elementType": "$vectorScalar" },
60
+ "params": {
61
+ "buffer": "uniform",
62
+ "struct": [
63
+ { "name": "rows", "type": "u32", "value": "normRows" },
64
+ { "name": "rowStride", "type": "u32", "value": "normRowStride" }
65
+ ]
66
+ },
67
+ "bias": { "arg": "b", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
68
+ "mean_out": { "arg": "mean", "buffer": "storage", "elementType": "f32" },
69
+ "inv_std_out": { "arg": "invStdDev", "buffer": "storage", "elementType": "f32" },
70
+ "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$scalar" },
71
+ "scale_2": { "name": "scale", "buffer": "read-only-storage", "elementType": "$scalar" },
72
+ "y_2": { "name": "y", "buffer": "storage", "elementType": "$scalar" },
73
+ "bias_2": { "arg": "b", "name": "bias", "buffer": "read-only-storage", "elementType": "$scalar" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74
  },
75
  "variants": [
76
  {
77
  "id": "last_axis_row_vec4",
78
  "priority": 110,
79
+ "when": ["not present.b and noStatsOutputs", "lastAxisExactScaleOk", "dim(shapes.x, -1) % 4 == 0"],
80
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
81
  "passes": [
82
  {
83
  "id": "main",
84
  "name": "LayerNormalization.LastAxisRowVec4",
85
+ "shader": "norm-row-stats.wgsl.jinja",
86
+ "derive": {
87
+ "modeSpec": "\"layer\"",
88
+ "vec4": true,
89
+ "hasBias": false,
90
+ "writeStats": false,
91
+ "scalar": "dtypes.T",
92
+ "usesF16Spec": "dtypes.T == \"f16\"",
93
+ "hidden": "dim(shapes.x, -1)",
94
+ "wg": "lastAxisWgVec4",
95
+ "epsilon": "attrs.epsilon",
96
+ "hiddenVec": "dim(shapes.x, -1) / 4",
97
+ "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
98
+ "combineSubgroups": "hasSubgroupId"
 
 
99
  },
100
+ "bindings": ["x", "scale", "y", "params"],
101
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
102
+ "subgroupCollectivesWidth": "portable"
103
  }
104
  ]
105
  },
106
  {
107
  "id": "last_axis_row",
108
  "priority": 100,
109
+ "when": ["not present.b and noStatsOutputs", "lastAxisExactScaleOk"],
110
  "demoteWhen": ["not lastAxisScalarFastOk"],
111
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
112
  "passes": [
113
  {
114
  "id": "main",
115
  "name": "LayerNormalization.LastAxisRow",
116
+ "shader": "norm-row-stats.wgsl.jinja",
117
+ "derive": {
118
+ "modeSpec": "\"layer\"",
119
+ "vec4": false,
120
+ "hasBias": false,
121
+ "writeStats": false,
122
+ "scalar": "dtypes.T",
123
+ "usesF16Spec": "dtypes.T == \"f16\"",
124
+ "hidden": "dim(shapes.x, -1)",
125
+ "wg": "lastAxisWg",
126
+ "epsilon": "attrs.epsilon",
127
+ "combineSubgroups": "hasSubgroupId"
 
 
128
  },
129
+ "bindings": ["x", "scale", "y", "params"],
130
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
131
+ "subgroupCollectivesWidth": "portable"
132
  }
133
  ]
134
  },
135
  {
136
  "id": "last_axis_bias_row_vec4",
137
  "priority": 111,
138
+ "when": ["present.b and noStatsOutputs and biasExactOk", "lastAxisExactScaleOk", "dim(shapes.x, -1) % 4 == 0"],
139
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
140
  "passes": [
141
  {
142
  "id": "main",
143
  "name": "LayerNormalization.LastAxisRowVec4",
144
+ "shader": "norm-row-stats.wgsl.jinja",
145
+ "derive": {
146
+ "modeSpec": "\"layer\"",
147
+ "vec4": true,
148
+ "hasBias": true,
149
+ "writeStats": false,
150
+ "scalar": "dtypes.T",
151
+ "usesF16Spec": "dtypes.T == \"f16\"",
152
+ "hidden": "dim(shapes.x, -1)",
153
+ "wg": "lastAxisWgVec4",
154
+ "epsilon": "attrs.epsilon",
155
+ "hiddenVec": "dim(shapes.x, -1) / 4",
156
+ "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
157
+ "combineSubgroups": "hasSubgroupId"
 
 
158
  },
159
+ "bindings": ["x", "scale", "bias", "y", "params"],
160
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
161
+ "subgroupCollectivesWidth": "portable"
162
  }
163
  ]
164
  },
165
  {
166
  "id": "last_axis_bias_row",
167
  "priority": 101,
168
+ "when": ["present.b and noStatsOutputs and biasExactOk", "lastAxisExactScaleOk"],
169
  "demoteWhen": ["not lastAxisScalarFastOk"],
170
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
171
  "passes": [
172
  {
173
  "id": "main",
174
  "name": "LayerNormalization.LastAxisRow",
175
+ "shader": "norm-row-stats.wgsl.jinja",
176
+ "derive": {
177
+ "modeSpec": "\"layer\"",
178
+ "vec4": false,
179
+ "hasBias": true,
180
+ "writeStats": false,
181
+ "scalar": "dtypes.T",
182
+ "usesF16Spec": "dtypes.T == \"f16\"",
183
+ "hidden": "dim(shapes.x, -1)",
184
+ "wg": "lastAxisWg",
185
+ "epsilon": "attrs.epsilon",
186
+ "combineSubgroups": "hasSubgroupId"
 
 
187
  },
188
+ "bindings": ["x", "scale", "bias", "y", "params"],
189
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
190
+ "subgroupCollectivesWidth": "portable"
191
  }
192
  ]
193
  },
194
  {
195
  "id": "last_axis_stats_row_vec4",
196
  "priority": 112,
197
+ "when": ["not present.b and fullStatsOutputs and statsRowsOk", "lastAxisExactScaleOk", "dim(shapes.x, -1) % 4 == 0"],
198
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
199
  "passes": [
200
  {
201
  "id": "main",
202
  "name": "LayerNormalization.LastAxisRowVec4",
203
+ "shader": "norm-row-stats.wgsl.jinja",
204
+ "derive": {
205
+ "modeSpec": "\"layer\"",
206
+ "vec4": true,
207
+ "hasBias": false,
208
+ "writeStats": true,
209
+ "scalar": "dtypes.T",
210
+ "usesF16Spec": "dtypes.T == \"f16\"",
211
+ "hidden": "dim(shapes.x, -1)",
212
+ "wg": "lastAxisWgVec4",
213
+ "epsilon": "attrs.epsilon",
214
+ "hiddenVec": "dim(shapes.x, -1) / 4",
215
+ "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
216
+ "combineSubgroups": "hasSubgroupId"
 
 
217
  },
218
+ "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
219
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
220
+ "subgroupCollectivesWidth": "portable"
221
  }
222
  ]
223
  },
224
  {
225
  "id": "last_axis_stats_row",
226
  "priority": 102,
227
+ "when": ["not present.b and fullStatsOutputs and statsRowsOk", "lastAxisExactScaleOk"],
228
  "demoteWhen": ["not lastAxisScalarFastOk"],
229
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
230
  "passes": [
231
  {
232
  "id": "main",
233
  "name": "LayerNormalization.LastAxisRow",
234
+ "shader": "norm-row-stats.wgsl.jinja",
235
+ "derive": {
236
+ "modeSpec": "\"layer\"",
237
+ "vec4": false,
238
+ "hasBias": false,
239
+ "writeStats": true,
240
+ "scalar": "dtypes.T",
241
+ "usesF16Spec": "dtypes.T == \"f16\"",
242
+ "hidden": "dim(shapes.x, -1)",
243
+ "wg": "lastAxisWg",
244
+ "epsilon": "attrs.epsilon",
245
+ "combineSubgroups": "hasSubgroupId"
 
 
246
  },
247
+ "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
248
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
249
+ "subgroupCollectivesWidth": "portable"
250
  }
251
  ]
252
  },
253
  {
254
  "id": "last_axis_bias_stats_row_vec4",
255
  "priority": 113,
256
+ "when": ["present.b and fullStatsOutputs and biasExactOk and statsRowsOk", "lastAxisExactScaleOk", "dim(shapes.x, -1) % 4 == 0"],
257
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
258
  "passes": [
259
  {
260
  "id": "main",
261
  "name": "LayerNormalization.LastAxisRowVec4",
262
+ "shader": "norm-row-stats.wgsl.jinja",
263
+ "derive": {
264
+ "modeSpec": "\"layer\"",
265
+ "vec4": true,
266
+ "hasBias": true,
267
+ "writeStats": true,
268
+ "scalar": "dtypes.T",
269
+ "usesF16Spec": "dtypes.T == \"f16\"",
270
+ "hidden": "dim(shapes.x, -1)",
271
+ "wg": "lastAxisWgVec4",
272
+ "epsilon": "attrs.epsilon",
273
+ "hiddenVec": "dim(shapes.x, -1) / 4",
274
+ "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
275
+ "combineSubgroups": "hasSubgroupId"
 
 
276
  },
277
+ "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
278
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
279
+ "subgroupCollectivesWidth": "portable"
280
  }
281
  ]
282
  },
283
  {
284
  "id": "last_axis_bias_stats_row",
285
  "priority": 103,
286
+ "when": ["present.b and fullStatsOutputs and biasExactOk and statsRowsOk", "lastAxisExactScaleOk"],
287
  "demoteWhen": ["not lastAxisScalarFastOk"],
288
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "dtypes.T" },
289
  "passes": [
290
  {
291
  "id": "main",
292
  "name": "LayerNormalization.LastAxisRow",
293
+ "shader": "norm-row-stats.wgsl.jinja",
294
+ "derive": {
295
+ "modeSpec": "\"layer\"",
296
+ "vec4": false,
297
+ "hasBias": true,
298
+ "writeStats": true,
299
+ "scalar": "dtypes.T",
300
+ "usesF16Spec": "dtypes.T == \"f16\"",
301
+ "hidden": "dim(shapes.x, -1)",
302
+ "wg": "lastAxisWg",
303
+ "epsilon": "attrs.epsilon",
304
+ "combineSubgroups": "hasSubgroupId"
 
 
305
  },
306
+ "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
307
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
308
+ "subgroupCollectivesWidth": "portable"
309
  }
310
  ]
311
  },
312
  {
313
  "id": "suffix_axis_bias_exact_row_vec4",
314
  "priority": 121,
315
+ "when": ["present.b", "noStatsOutputs", "suffixAxisExactAffineOk", "suffixAxisSize % 4 == 0"],
316
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
317
  "passes": [
318
  {
319
  "id": "main",
320
  "name": "LayerNormalization.SuffixAxisRowVec4",
321
+ "shader": "norm-row-stats.wgsl.jinja",
322
+ "derive": {
323
+ "modeSpec": "\"layer\"",
324
+ "vec4": true,
325
+ "hasBias": true,
326
+ "writeStats": false,
327
+ "scalar": "dtypes.T",
328
+ "usesF16Spec": "dtypes.T == \"f16\"",
329
+ "hidden": "suffixAxisSize",
330
+ "wg": "suffixAxisWgVec4",
331
+ "epsilon": "attrs.epsilon",
332
+ "hiddenVec": "suffixAxisSize / 4",
333
+ "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
334
+ "combineSubgroups": "hasSubgroupId"
 
 
335
  },
336
+ "bindings": ["x", "scale", "bias", "y", "params"],
337
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
338
+ "subgroupCollectivesWidth": "portable"
339
  }
340
  ]
341
  },
342
  {
343
  "id": "last_axis",
344
  "priority": 0,
345
+ "when": ["not present.b", "noStatsOutputs", "lastAxisBroadcastScaleOk"],
346
+ "derive": {
347
  "hasBias": false,
348
  "writeMean": false,
349
  "writeInvStdDev": false,
350
  "scalar": "dtypes.T",
351
  "vectorScalar": "dtypes.T",
352
+ "hiddenSize": "dim(shapes.x, -1)",
 
353
  "workgroupSize": "lastAxisWg",
354
  "epsilon": "attrs.epsilon"
355
  },
 
357
  {
358
  "id": "main",
359
  "name": "LayerNormalization",
360
+ "shader": "layer-normalization.wgsl.jinja",
361
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
362
+ "bindings": ["x", "scale", "y", "params"],
363
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
364
  }
365
  ]
366
  },
367
  {
368
  "id": "last_axis_bias",
369
  "priority": 10,
370
+ "when": ["present.b", "noStatsOutputs", "lastAxisBroadcastScaleOk", "biasBroadcastOk"],
371
+ "derive": {
372
  "hasBias": true,
373
  "writeMean": false,
374
  "writeInvStdDev": false,
375
  "scalar": "dtypes.T",
376
  "vectorScalar": "dtypes.T",
377
+ "hiddenSize": "dim(shapes.x, -1)",
 
378
  "workgroupSize": "lastAxisWg",
379
  "epsilon": "attrs.epsilon"
380
  },
 
382
  {
383
  "id": "main",
384
  "name": "LayerNormalization",
385
+ "shader": "layer-normalization.wgsl.jinja",
386
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
387
+ "bindings": ["x", "scale", "bias", "y", "params"],
388
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
389
  }
390
  ]
391
  },
392
  {
393
  "id": "last_axis_stats",
394
  "priority": 20,
395
+ "when": ["not present.b", "fullStatsOutputs", "lastAxisBroadcastScaleOk", "statsRowsOk"],
396
+ "derive": {
397
  "hasBias": false,
398
  "writeMean": true,
399
  "writeInvStdDev": true,
400
  "scalar": "dtypes.T",
401
  "vectorScalar": "dtypes.T",
402
+ "hiddenSize": "dim(shapes.x, -1)",
 
403
  "workgroupSize": "lastAxisWg",
404
  "epsilon": "attrs.epsilon"
405
  },
 
407
  {
408
  "id": "main",
409
  "name": "LayerNormalization",
410
+ "shader": "layer-normalization.wgsl.jinja",
411
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
412
+ "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
413
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
414
  }
415
  ]
416
  },
417
  {
418
  "id": "last_axis_bias_stats",
419
  "priority": 30,
420
+ "when": ["present.b", "fullStatsOutputs", "lastAxisBroadcastScaleOk", "biasBroadcastOk", "statsRowsOk"],
421
+ "derive": {
422
  "hasBias": true,
423
  "writeMean": true,
424
  "writeInvStdDev": true,
425
  "scalar": "dtypes.T",
426
  "vectorScalar": "dtypes.T",
427
+ "hiddenSize": "dim(shapes.x, -1)",
 
428
  "workgroupSize": "lastAxisWg",
429
  "epsilon": "attrs.epsilon"
430
  },
 
432
  {
433
  "id": "main",
434
  "name": "LayerNormalization",
435
+ "shader": "layer-normalization.wgsl.jinja",
436
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
437
+ "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
438
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
439
  }
440
  ]
441
  },
442
  {
443
  "id": "suffix_axis",
444
  "priority": 40,
445
+ "when": ["not present.b", "noStatsOutputs", "suffixAxisBroadcastScaleOk"],
446
+ "derive": {
447
  "hasBias": false,
448
  "writeMean": false,
449
  "writeInvStdDev": false,
450
  "scalar": "dtypes.T",
451
  "vectorScalar": "dtypes.T",
 
452
  "hiddenSize": "suffixAxisSize",
453
  "workgroupSize": "suffixAxisWg",
454
  "epsilon": "attrs.epsilon"
 
457
  {
458
  "id": "main",
459
  "name": "LayerNormalization.SuffixAxis",
460
+ "shader": "layer-normalization.wgsl.jinja",
461
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
462
+ "bindings": ["x", "scale", "y", "params"],
463
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
464
  }
465
  ]
466
  },
467
  {
468
  "id": "suffix_axis_bias",
469
  "priority": 50,
470
+ "when": ["present.b", "noStatsOutputs", "suffixAxisBroadcastScaleOk", "biasBroadcastOk"],
471
+ "derive": {
472
  "hasBias": true,
473
  "writeMean": false,
474
  "writeInvStdDev": false,
475
  "scalar": "dtypes.T",
476
  "vectorScalar": "dtypes.T",
 
477
  "hiddenSize": "suffixAxisSize",
478
  "workgroupSize": "suffixAxisWg",
479
  "epsilon": "attrs.epsilon"
 
482
  {
483
  "id": "main",
484
  "name": "LayerNormalization.SuffixAxisBias",
485
+ "shader": "layer-normalization.wgsl.jinja",
486
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
487
+ "bindings": ["x", "scale", "bias", "y", "params"],
488
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
489
  }
490
  ]
491
  },
492
  {
493
  "id": "suffix_axis_stats",
494
  "priority": 45,
495
+ "when": ["not present.b", "fullStatsOutputs", "suffixAxisBroadcastScaleOk", "statsOuterOk"],
496
+ "derive": {
497
  "hasBias": false,
498
  "writeMean": true,
499
  "writeInvStdDev": true,
500
  "scalar": "dtypes.T",
501
  "vectorScalar": "dtypes.T",
 
502
  "hiddenSize": "suffixAxisSize",
503
  "workgroupSize": "suffixAxisWg",
504
  "epsilon": "attrs.epsilon"
 
507
  {
508
  "id": "main",
509
  "name": "LayerNormalization.SuffixAxisStats",
510
+ "shader": "layer-normalization.wgsl.jinja",
511
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
512
+ "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
513
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
514
  }
515
  ]
516
  },
517
  {
518
  "id": "suffix_axis_bias_stats",
519
  "priority": 55,
520
+ "when": ["present.b", "fullStatsOutputs", "suffixAxisBroadcastScaleOk", "biasBroadcastOk", "statsOuterOk"],
521
+ "derive": {
522
  "hasBias": true,
523
  "writeMean": true,
524
  "writeInvStdDev": true,
525
  "scalar": "dtypes.T",
526
  "vectorScalar": "dtypes.T",
 
527
  "hiddenSize": "suffixAxisSize",
528
  "workgroupSize": "suffixAxisWg",
529
  "epsilon": "attrs.epsilon"
 
532
  {
533
  "id": "main",
534
  "name": "LayerNormalization.SuffixAxisBiasStats",
535
+ "shader": "layer-normalization.wgsl.jinja",
536
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
537
+ "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
538
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
539
  }
540
  ]
541
  },
542
  {
543
  "id": "mean_only",
544
  "priority": 31,
545
+ "when": ["not present.b and meanOnlyOutputs and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
546
+ "derive": {
547
  "hasBias": false,
548
  "writeMean": true,
549
  "writeInvStdDev": false,
550
  "scalar": "dtypes.T",
 
551
  "hiddenSize": "genericHiddenSize",
552
  "workgroupSize": "genericWorkgroupSize",
553
  "epsilon": "attrs.epsilon"
 
556
  {
557
  "id": "main",
558
  "name": "LayerNormalization.MeanOnly",
559
+ "shader": "layer-normalization.wgsl.jinja",
560
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
561
+ "bindings": ["x_2", "scale_2", "y_2", "mean_out", "params"],
562
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
563
  }
564
  ]
565
  },
566
  {
567
  "id": "bias_mean_only",
568
  "priority": 32,
569
+ "when": ["present.b and meanOnlyOutputs and biasBroadcastOk and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
570
+ "derive": {
571
  "hasBias": true,
572
  "writeMean": true,
573
  "writeInvStdDev": false,
574
  "scalar": "dtypes.T",
 
575
  "hiddenSize": "genericHiddenSize",
576
  "workgroupSize": "genericWorkgroupSize",
577
  "epsilon": "attrs.epsilon"
 
580
  {
581
  "id": "main",
582
  "name": "LayerNormalization.BiasMeanOnly",
583
+ "shader": "layer-normalization.wgsl.jinja",
584
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
585
+ "bindings": ["x_2", "scale_2", "bias_2", "y_2", "mean_out", "params"],
586
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
587
  }
588
  ]
589
  },
590
  {
591
  "id": "inv_std_dev_only",
592
  "priority": 33,
593
+ "when": ["not present.b and invStdOnlyOutputs and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
594
+ "derive": {
595
  "hasBias": false,
596
  "writeMean": false,
597
  "writeInvStdDev": true,
598
  "scalar": "dtypes.T",
 
599
  "hiddenSize": "genericHiddenSize",
600
  "workgroupSize": "genericWorkgroupSize",
601
  "epsilon": "attrs.epsilon"
 
604
  {
605
  "id": "main",
606
  "name": "LayerNormalization.InvStdDevOnly",
607
+ "shader": "layer-normalization.wgsl.jinja",
608
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
609
+ "bindings": ["x_2", "scale_2", "y_2", "inv_std_out", "params"],
610
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
611
  }
612
  ]
613
  },
614
  {
615
  "id": "bias_inv_std_dev_only",
616
  "priority": 34,
617
+ "when": ["present.b and invStdOnlyOutputs and biasBroadcastOk and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
618
+ "derive": {
619
  "hasBias": true,
620
  "writeMean": false,
621
  "writeInvStdDev": true,
622
  "scalar": "dtypes.T",
 
623
  "hiddenSize": "genericHiddenSize",
624
  "workgroupSize": "genericWorkgroupSize",
625
  "epsilon": "attrs.epsilon"
 
628
  {
629
  "id": "main",
630
  "name": "LayerNormalization.BiasInvStdDevOnly",
631
+ "shader": "layer-normalization.wgsl.jinja",
632
+ "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
633
+ "bindings": ["x_2", "scale_2", "bias_2", "y_2", "inv_std_out", "params"],
634
+ "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
 
 
635
  }
636
  ]
637
  }
build/webgpu/metadata.json CHANGED
@@ -1,19 +1,44 @@
1
  {
2
  "name": "ai.onnx.LayerNormalization",
3
- "id": "_ai_onnx_layernormalization_webgpu_af1721d",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "fLkWXUeB9lfy8s03WLvcd2LRTBWce01Yzzv5lBHMvig=",
11
- "layer-normalization.wgsl.jinja": "3BXN4VIAbNR6se39KXKK5jV514N8v756GeH37Tsuzi0=",
12
- "manifest.json": "j8ogYGGt7G1Pr0St9Gyhpff035GfQKHmuz6IBktIrV8=",
13
- "norm-row-stats.wgsl.jinja": "82e5r5vFGd0ylf/r3n+fucFbRPPMH3II1VBaTTKSkp4=",
14
- "test.json": "hjj+HEnuQ0NKto1gCH5+81D7bWaH+hgQWy5OlySa7C0="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
18
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.LayerNormalization" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  }
 
1
  {
2
  "name": "ai.onnx.LayerNormalization",
3
+ "id": "_ai_onnx_layernormalization_webgpu_7b13eb1",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "IPZzLq64+ycfl0fAzym0hDLorrSrZ5YDwLnpVEGgGHc=",
11
+ "layer-normalization.wgsl.jinja": "NJ1/CeeYHnToR+Ki5VHv4U4gxuHKa9zYg9PG2QOndME=",
12
+ "manifest.json": "aZkl41hMG0XVB5WthNHRcxLOw39oCNULp/kF2RczI8Y=",
13
+ "norm-row-stats.wgsl.jinja": "HUUqntKH7vqbffRpSu33tnmSU7PudcOCtl5QhV1xlGk=",
14
+ "test.json": "qJxDvb9POxT3rI+vLcahsq2mqIwAfen9b53MgOSG5Go="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
+ "webgpu": {
19
+ "manifestSpec": "2.0",
20
+ "variants": {
21
+ "last_axis_row_vec4": ["norm-row-stats.wgsl.jinja"],
22
+ "last_axis_row": ["norm-row-stats.wgsl.jinja"],
23
+ "last_axis_bias_row_vec4": ["norm-row-stats.wgsl.jinja"],
24
+ "last_axis_bias_row": ["norm-row-stats.wgsl.jinja"],
25
+ "last_axis_stats_row_vec4": ["norm-row-stats.wgsl.jinja"],
26
+ "last_axis_stats_row": ["norm-row-stats.wgsl.jinja"],
27
+ "last_axis_bias_stats_row_vec4": ["norm-row-stats.wgsl.jinja"],
28
+ "last_axis_bias_stats_row": ["norm-row-stats.wgsl.jinja"],
29
+ "suffix_axis_bias_exact_row_vec4": ["norm-row-stats.wgsl.jinja"],
30
+ "last_axis": ["layer-normalization.wgsl.jinja"],
31
+ "last_axis_bias": ["layer-normalization.wgsl.jinja"],
32
+ "last_axis_stats": ["layer-normalization.wgsl.jinja"],
33
+ "last_axis_bias_stats": ["layer-normalization.wgsl.jinja"],
34
+ "suffix_axis": ["layer-normalization.wgsl.jinja"],
35
+ "suffix_axis_bias": ["layer-normalization.wgsl.jinja"],
36
+ "suffix_axis_stats": ["layer-normalization.wgsl.jinja"],
37
+ "suffix_axis_bias_stats": ["layer-normalization.wgsl.jinja"],
38
+ "mean_only": ["layer-normalization.wgsl.jinja"],
39
+ "bias_mean_only": ["layer-normalization.wgsl.jinja"],
40
+ "inv_std_dev_only": ["layer-normalization.wgsl.jinja"],
41
+ "bias_inv_std_dev_only": ["layer-normalization.wgsl.jinja"]
42
+ }
43
+ }
44
  }
build/webgpu/norm-row-stats.wgsl.jinja CHANGED
@@ -1,9 +1,16 @@
1
- {% if source.usesF16 %}
2
  enable f16;
3
  {% endif %}
4
- {% set combineSubgroups = source.combineSubgroups %}
5
- {% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
6
- {% set writeStats = source.writeStats if source.writeStats is defined else false %}
 
 
 
 
 
 
 
7
  {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
8
  if combineSubgroups else ", tid: u32" %}
9
  {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -23,14 +30,57 @@ enable subgroups;
23
  //
24
  // Shifted moments avoid cancellation from a large common offset; scaling uses
25
  // inverseSqrt(variance + EPSILON).
26
- const HIDDEN: u32 = {{ source.hidden }}u;
27
- {% if source.vec4 %}
28
- const HIDDEN_V: u32 = {{ source.hiddenVec }}u;
 
 
 
 
 
 
 
 
 
29
  {% endif %}
30
- const WG: u32 = {{ source.wg }}u;
31
- const EPSILON: f32 = {{ source.epsilon }};
32
 
 
 
 
 
 
 
 
33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
34
 
35
  {% if combineSubgroups %}
36
  var<workgroup> sg_partials: array<vec2<f32>, WG>;
@@ -89,29 +139,52 @@ fn main(
89
  return;
90
  }
91
  let tid = lid.x;
92
- {% if source.vec4 and not scalarIo %}
 
 
 
 
 
 
 
93
  let base = row * HIDDEN_V;
94
  {% else %}
95
  let base = row * HIDDEN;
96
  {% endif %}
97
 
98
- {% if source.vec4 %}
 
 
 
99
  let shift = f32(x[base].x);
 
100
  {% else %}
101
  let shift = f32(x[base]);
102
  {% endif %}
103
 
104
  var acc = vec2<f32>(0.0, 0.0);
105
- {% if source.vec4 %}
106
  for (var i = tid; i < HIDDEN_V; i = i + WG) {
 
 
 
 
 
 
107
  let v = vec4<f32>(x[base + i]);
 
108
  let d = v - vec4<f32>(shift);
109
  acc.x = acc.x + d.x + d.y + d.z + d.w;
110
  acc.y = acc.y + dot(d, d);
111
  }
112
  {% else %}
113
  for (var i = tid; i < HIDDEN; i = i + WG) {
 
 
 
 
114
  let v = f32(x[base + i]);
 
115
  let d = v - shift;
116
  acc.x = acc.x + d;
117
  acc.y = acc.y + d * d;
@@ -131,25 +204,55 @@ fn main(
131
  }
132
  {% endif %}
133
 
134
- {% if source.vec4 %}
 
 
 
135
  for (var i = tid; i < HIDDEN_V; i = i + WG) {
 
 
 
 
 
 
 
136
  let idx = base + i;
137
  let v = vec4<f32>(x[idx]);
 
138
  var value = (v - vec4<f32>(row_mean)) * inv * vec4<f32>(scale[i]);
139
- {% if source.hasBias %}
140
  value = value + vec4<f32>(bias[i]);
141
  {% endif %}
142
- y[idx] = {{ source.vecType }}(value);
143
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
144
  {% else %}
145
  for (var i = tid; i < HIDDEN; i = i + WG) {
146
  let idx = base + i;
 
 
 
147
  let v = f32(x[idx]);
 
148
  var value = (v - row_mean) * inv * f32(scale[i]);
149
- {% if source.hasBias %}
150
  value = value + f32(bias[i]);
151
  {% endif %}
152
- y[idx] = {{ source.scalar }}(value);
153
  }
154
  {% endif %}
155
  }
 
1
+ {% if usesF16Spec %}
2
  enable f16;
3
  {% endif %}
4
+ {% set combineSubgroups = combineSubgroups %}
5
+ {% set scalarIo = scalarIo if scalarIo is defined else false %}
6
+ {% set packedBf16Embedding = packedBf16Embedding if packedBf16Embedding is defined else false %}
7
+ {% set writeStats = writeStats if writeStats is defined else false %}
8
+ {% set rmsChainNorm = rmsChainNorm if rmsChainNorm is defined else false %}
9
+ {% set hiddenPairs = hiddenPairs | default(0) %}
10
+ {% set numRows = numRows | default(0) %}
11
+ {% set epsilon = epsilon | default("0.0") %}
12
+ {% set epsilon2 = epsilon2 | default("0.0") %}
13
+ {% set hasBias = hasBias is defined and hasBias %}
14
  {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
15
  if combineSubgroups else ", tid: u32" %}
16
  {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 
30
  //
31
  // Shifted moments avoid cancellation from a large common offset; scaling uses
32
  // inverseSqrt(variance + EPSILON).
33
+ const HIDDEN: u32 = {{ hidden }}u;
34
+ {% if vec4 %}
35
+ const HIDDEN_V: u32 = {{ hiddenVec }}u;
36
+ {% endif %}
37
+ {% if packedBf16Embedding %}
38
+ const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
39
+ const NUM_ROWS: u32 = {{ numRows }}u;
40
+ {% endif %}
41
+ const WG: u32 = {{ wg }}u;
42
+ const EPSILON: f32 = {{ epsilon }};
43
+ {% if rmsChainNorm %}
44
+ const EPSILON2: f32 = {{ epsilon2 }};
45
  {% endif %}
 
 
46
 
47
+ {% if packedBf16Embedding %}
48
+ {% if vec4 %}
49
+ fn unpack_bf16_pair(word: u32) -> vec2<f32> {
50
+ let bits = vec2<u32>(word & 0xffffu, word >> 16u);
51
+ return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
52
+ }
53
+ {% endif %}
54
 
55
+ {% if not vec4 %}
56
+ fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
57
+ if (source_row >= NUM_ROWS) {
58
+ return 0.0;
59
+ }
60
+ let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
61
+ let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
62
+ return bitcast<f32>(bits << 16u);
63
+ }
64
+ {% endif %}
65
+
66
+ {% if vec4 %}
67
+ fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
68
+ if (source_row >= NUM_ROWS) {
69
+ return vec4<f32>(0.0);
70
+ }
71
+ let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
72
+ let low = unpack_bf16_pair(x[base]);
73
+ let high = unpack_bf16_pair(x[base + 1u]);
74
+ return vec4<f32>(low, high);
75
+ }
76
+ {% endif %}
77
+ {% endif %}
78
+
79
+ {% if vec4 and scalarIo %}
80
+ fn load_vec4(index: u32) -> vec4<f32> {
81
+ return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
82
+ }
83
+ {% endif %}
84
 
85
  {% if combineSubgroups %}
86
  var<workgroup> sg_partials: array<vec2<f32>, WG>;
 
139
  return;
140
  }
141
  let tid = lid.x;
142
+ {% if packedBf16Embedding %}
143
+ let source_row = indices[row];
144
+ {% if vec4 %}
145
+ let base = row * HIDDEN_V;
146
+ {% else %}
147
+ let base = row * HIDDEN;
148
+ {% endif %}
149
+ {% elif vec4 and not scalarIo %}
150
  let base = row * HIDDEN_V;
151
  {% else %}
152
  let base = row * HIDDEN;
153
  {% endif %}
154
 
155
+ {% if vec4 %}
156
+ {% if scalarIo %}
157
+ let shift = f32(x[base]);
158
+ {% else %}
159
  let shift = f32(x[base].x);
160
+ {% endif %}
161
  {% else %}
162
  let shift = f32(x[base]);
163
  {% endif %}
164
 
165
  var acc = vec2<f32>(0.0, 0.0);
166
+ {% if vec4 %}
167
  for (var i = tid; i < HIDDEN_V; i = i + WG) {
168
+ {% if packedBf16Embedding %}
169
+ let v = embedding_vec4(source_row, i);
170
+ embedding_out[base + i] = v;
171
+ {% elif scalarIo %}
172
+ let v = load_vec4(base + i * 4u);
173
+ {% else %}
174
  let v = vec4<f32>(x[base + i]);
175
+ {% endif %}
176
  let d = v - vec4<f32>(shift);
177
  acc.x = acc.x + d.x + d.y + d.z + d.w;
178
  acc.y = acc.y + dot(d, d);
179
  }
180
  {% else %}
181
  for (var i = tid; i < HIDDEN; i = i + WG) {
182
+ {% if packedBf16Embedding %}
183
+ let v = embedding_scalar(source_row, i);
184
+ embedding_out[base + i] = v;
185
+ {% else %}
186
  let v = f32(x[base + i]);
187
+ {% endif %}
188
  let d = v - shift;
189
  acc.x = acc.x + d;
190
  acc.y = acc.y + d * d;
 
204
  }
205
  {% endif %}
206
 
207
+ {% if rmsChainNorm %}
208
+ var acc2 = 0.0;
209
+ {% endif %}
210
+ {% if vec4 %}
211
  for (var i = tid; i < HIDDEN_V; i = i + WG) {
212
+ {% if packedBf16Embedding %}
213
+ let idx = base + i;
214
+ let v = embedding_vec4(source_row, i);
215
+ {% elif scalarIo %}
216
+ let idx = base + i * 4u;
217
+ let v = load_vec4(idx);
218
+ {% else %}
219
  let idx = base + i;
220
  let v = vec4<f32>(x[idx]);
221
+ {% endif %}
222
  var value = (v - vec4<f32>(row_mean)) * inv * vec4<f32>(scale[i]);
223
+ {% if hasBias %}
224
  value = value + vec4<f32>(bias[i]);
225
  {% endif %}
226
+ y[idx] = {{ vecType }}(value);
227
  }
228
+ {% if rmsChainNorm %}
229
+
230
+ // The chained second norm reads the residual row this loop just stored. This
231
+ // barrier completes those stores and any preceding shared-scratch use before
232
+ // the next reduction reuses its scratch; each lane then re-reads only the
233
+ // elements it wrote itself.
234
+ workgroupBarrier();
235
+ let total2 = reduce_scalar(acc2{{ reduceThreadArguments }});
236
+ let inv2 = inverseSqrt(total2 / f32(HIDDEN) + EPSILON2);
237
+ for (var i = tid; i < HIDDEN_V; i = i + WG) {
238
+ let idx = base + i;
239
+ let hv = vec4<f32>(y[idx]);
240
+ normed2[idx] = {{ vecType }}(hv * inv2 * vec4<f32>(scale2[i]));
241
+ }
242
+ {% endif %}
243
  {% else %}
244
  for (var i = tid; i < HIDDEN; i = i + WG) {
245
  let idx = base + i;
246
+ {% if packedBf16Embedding %}
247
+ let v = embedding_scalar(source_row, i);
248
+ {% else %}
249
  let v = f32(x[idx]);
250
+ {% endif %}
251
  var value = (v - row_mean) * inv * f32(scale[i]);
252
+ {% if hasBias %}
253
  value = value + f32(bias[i]);
254
  {% endif %}
255
+ y[idx] = {{ scalar }}(value);
256
  }
257
  {% endif %}
258
  }
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "ai.onnx.LayerNormalization",
3
  "fixtureArrays": {
4
  "onnx_backend_layer_normalization_3d_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358],
5
  "onnx_backend_layer_normalization_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902, -0.6724604368209839, -0.35955315828323364, -0.8131462931632996, -1.7262825965881348, 0.17742614448070526, -0.4017809331417084, -1.630198359489441, 0.46278226375579834, -0.9072983860969543, 0.05194539576768875, 0.7290905714035034, 0.12898291647434235, 1.1394007205963135, -1.234825849533081, 0.4023416340351105, -0.6848101019859314, -0.8707971572875977, -0.5788496732711792, -0.3115525245666504, 0.056165341287851334, -1.1651498079299927, 0.9008265137672424, 0.4656624495983124, -1.5362436771392822, 1.4882521629333496, 1.895889163017273, 1.1787796020507812, -0.1799248307943344, -1.0707526206970215, 1.0544517040252686, -0.4031769335269928, 1.222445011138916, 0.2082749754190445, 0.9766390323638916, 0.3563663959503174, 0.7065731883049011, 0.01050002034753561, 1.7858705520629883, 0.12691208720207214, 0.4019893705844879, 1.8831506967544556, -1.3477590084075928, -1.2704850435256958, 0.969396710395813, -1.1731233596801758, 1.9436211585998535, -0.4136189818382263, -0.747454822063446, 1.922942042350769, 1.4805147647857666, 1.8675589561462402, 0.9060446619987488, -0.8612256646156311, 1.910064935684204, -0.26800337433815, 0.8024563789367676, 0.9472519755363464, -0.15501008927822113, 0.6140793561935425, 0.922206699848175]
@@ -165,7 +164,7 @@
165
  "provenance": {
166
  "source": "onnxruntime/test/contrib_ops/layer_norm_op_test.cc",
167
  "test": "LayerNormalization",
168
- "notes": "Odd hidden-size companion for subnormal scale values; this exercises the non-vec4 last-axis path."
169
  },
170
  "attrs": { "epsilon": 0.00001, "axis": -1 },
171
  "inputs": {
@@ -226,14 +225,14 @@
226
  "provenance": {
227
  "source": "onnxruntime/test/contrib_ops/layer_norm_op_test.cc",
228
  "test": "LayerNormTest.LayerNorm17_opset",
229
- "notes": "Valid epsilon=0 edge: normal inputs produce subnormal variance but finite order-one normalized outputs and finite large InvStdDev."
230
  },
231
  "attrs": { "epsilon": 0, "axis": -1 },
232
  "inputs": {
233
  "x": {
234
  "dtype": "float32",
235
  "shape": [2, 2],
236
- "data": { "kind": "values", "values": [1e-20, -1e-20, 2e-20, -2e-20] }
237
  },
238
  "scale": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [1.0, 1.0] } },
239
  "b": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [0.0, 0.0] } }
@@ -1582,7 +1581,7 @@
1582
  "x": {
1583
  "dtype": "float16",
1584
  "shape": [8, 2048],
1585
- "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.5 }
1586
  },
1587
  "scale": {
1588
  "dtype": "float16",
@@ -1599,6 +1598,9 @@
1599
  "y": { "dtype": "float16", "shape": [8, 2048], "tolerance": 0.005 },
1600
  "mean": { "dtype": "float32", "shape": [8, 1], "tolerance": 0.005 },
1601
  "invStdDev": { "dtype": "float32", "shape": [8, 1], "tolerance": 0.02 }
 
 
 
1602
  }
1603
  },
1604
  {
 
1
  {
 
2
  "fixtureArrays": {
3
  "onnx_backend_layer_normalization_3d_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358],
4
  "onnx_backend_layer_normalization_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902, -0.6724604368209839, -0.35955315828323364, -0.8131462931632996, -1.7262825965881348, 0.17742614448070526, -0.4017809331417084, -1.630198359489441, 0.46278226375579834, -0.9072983860969543, 0.05194539576768875, 0.7290905714035034, 0.12898291647434235, 1.1394007205963135, -1.234825849533081, 0.4023416340351105, -0.6848101019859314, -0.8707971572875977, -0.5788496732711792, -0.3115525245666504, 0.056165341287851334, -1.1651498079299927, 0.9008265137672424, 0.4656624495983124, -1.5362436771392822, 1.4882521629333496, 1.895889163017273, 1.1787796020507812, -0.1799248307943344, -1.0707526206970215, 1.0544517040252686, -0.4031769335269928, 1.222445011138916, 0.2082749754190445, 0.9766390323638916, 0.3563663959503174, 0.7065731883049011, 0.01050002034753561, 1.7858705520629883, 0.12691208720207214, 0.4019893705844879, 1.8831506967544556, -1.3477590084075928, -1.2704850435256958, 0.969396710395813, -1.1731233596801758, 1.9436211585998535, -0.4136189818382263, -0.747454822063446, 1.922942042350769, 1.4805147647857666, 1.8675589561462402, 0.9060446619987488, -0.8612256646156311, 1.910064935684204, -0.26800337433815, 0.8024563789367676, 0.9472519755363464, -0.15501008927822113, 0.6140793561935425, 0.922206699848175]
 
164
  "provenance": {
165
  "source": "onnxruntime/test/contrib_ops/layer_norm_op_test.cc",
166
  "test": "LayerNormalization",
167
+ "notes": "An odd hidden size with subnormal scale values exercises the scalar last-axis path."
168
  },
169
  "attrs": { "epsilon": 0.00001, "axis": -1 },
170
  "inputs": {
 
225
  "provenance": {
226
  "source": "onnxruntime/test/contrib_ops/layer_norm_op_test.cc",
227
  "test": "LayerNormTest.LayerNorm17_opset",
228
+ "notes": "Diverges from the upstream test's inputs (inputs.x values [1e-20, -1e-20, 2e-20, -2e-20] -> values [1e-20, -1e-20, 4e-20, 0.0]); the expected output is recomputed by the CPU reference for the new inputs. Valid epsilon=0 edge: normal inputs produce subnormal variance but finite order-one normalized outputs and finite large InvStdDev."
229
  },
230
  "attrs": { "epsilon": 0, "axis": -1 },
231
  "inputs": {
232
  "x": {
233
  "dtype": "float32",
234
  "shape": [2, 2],
235
+ "data": { "kind": "values", "values": [1e-20, -1e-20, 4e-20, 0.0] }
236
  },
237
  "scale": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [1.0, 1.0] } },
238
  "b": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [0.0, 0.0] } }
 
1581
  "x": {
1582
  "dtype": "float16",
1583
  "shape": [8, 2048],
1584
+ "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.5, "offset": 1.0 }
1585
  },
1586
  "scale": {
1587
  "dtype": "float16",
 
1598
  "y": { "dtype": "float16", "shape": [8, 2048], "tolerance": 0.005 },
1599
  "mean": { "dtype": "float32", "shape": [8, 1], "tolerance": 0.005 },
1600
  "invStdDev": { "dtype": "float32", "shape": [8, 1], "tolerance": 0.02 }
1601
+ },
1602
+ "provenance": {
1603
+ "notes": "A 2,048-wide row with a scale and a bias. The input oscillates about 1.0, so the reported mean is O(1) and a rescaled or mis-divided mean is observable."
1604
  }
1605
  },
1606
  {