Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.DequantizeLinear

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
62f8cc0
·
verified ·
1 Parent(s): 1ddacab

sync 6fdf6301e2bb

Browse files
Files changed (8) hide show
  1. README.md +2 -2
  2. build/webgpu/bench.json +1 -1
  3. build/webgpu/manifest.json +94 -111
  4. build/webgpu/metadata.json +11 -11
  5. build/webgpu/quant-linear-blocked-axis.wgsl.jinja +7 -45
  6. build/webgpu/quant-linear-scalar.wgsl.jinja +6 -4
  7. build/webgpu/quant-linear-vec4.wgsl.jinja +9 -7
  8. build/webgpu/test.json +2 -2
README.md CHANGED
@@ -52,7 +52,7 @@ Default values (overridable per request):
52
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
53
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
54
  - [`test.json`](build/webgpu/test.json) — correctness cases
55
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
56
  - [`quant-linear-blocked-axis.wgsl.jinja`](build/webgpu/quant-linear-blocked-axis.wgsl.jinja)
57
  - [`quant-linear-scalar.wgsl.jinja`](build/webgpu/quant-linear-scalar.wgsl.jinja)
58
  - [`quant-linear-vec4.wgsl.jinja`](build/webgpu/quant-linear-vec4.wgsl.jinja)
@@ -60,7 +60,7 @@ Default values (overridable per request):
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.
 
52
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
53
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
54
  - [`test.json`](build/webgpu/test.json) — correctness cases
55
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
56
  - [`quant-linear-blocked-axis.wgsl.jinja`](build/webgpu/quant-linear-blocked-axis.wgsl.jinja)
57
  - [`quant-linear-scalar.wgsl.jinja`](build/webgpu/quant-linear-scalar.wgsl.jinja)
58
  - [`quant-linear-vec4.wgsl.jinja`](build/webgpu/quant-linear-vec4.wgsl.jinja)
 
60
  ## Use with `@huggingface/kernels`
61
 
62
  ```sh
63
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
64
  ```
65
 
66
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/bench.json CHANGED
@@ -123,7 +123,7 @@
123
  }
124
  },
125
  {
126
- "name": "per_channel_axis0_vec4_healthy_2048x2048",
127
  "preset": "smoke",
128
  "vars": { "inputDtype": "int8", "outputDtype": "float32", "count": 4194304 },
129
  "attrs": { "axis": 0 },
 
123
  }
124
  },
125
  {
126
+ "name": "per_channel_axis0_vec4_control_2048x2048",
127
  "preset": "smoke",
128
  "vars": { "inputDtype": "int8", "outputDtype": "float32", "count": 4194304 },
129
  "attrs": { "axis": 0 },
build/webgpu/manifest.json CHANGED
@@ -15,26 +15,35 @@
15
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
16
  "workgroupOk": "tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
17
  "outputDispatchFits": "workgroupOk and ceilDiv(ceilDiv(numel(shapes.y), tunables.WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
 
 
 
 
18
  "sameShapeOk": "ranks.y == ranks.x and numel(shapes.x) == numel(shapes.y) and outputDispatchFits",
19
  "dequantDtypesOk": "(attrs.output_dtype == 0 or attrs.output_dtype == onnxDtypeCode(logicalDtypes.TF)) and f16Ok(dtypes.TF)",
20
  "blockedScaleOk": "attrs.block_size > 0 and ranks.x >= 2 and ranks.x_scale == ranks.x and sameShapeOk and dim(shapes.x_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)",
21
- "elementCount": "numel(shapes.y)",
22
- "elementCount4": "floor(numel(shapes.y) / 4)",
23
  "scaleSize": "1 if ranks.x_scale == 0 else dim(shapes.x_scale, 0)",
24
  "scaleInner": "1 if ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 else inner(shapes.x, attrs.axis)",
25
- "blockedAxisDim": "dim(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
26
  "blockedScaleAxisDim": "dim(shapes.x_scale, attrs.axis) if ranks.x_scale >= 2 else 1",
27
- "blockedInner": "inner(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
28
  "xVec4": "\"vec4<\" ~ dtypes.TQ ~ \">\"",
29
  "scaleVec4": "\"vec4<\" ~ dtypes.TF ~ \">\"",
30
  "yScalar": "dtypes.TF",
31
  "yVec4": "\"vec4<\" ~ dtypes.TF ~ \">\"",
32
- "xUnsigned": "tensorDtypes.x == \"uint8\""
 
 
 
 
 
 
 
 
 
33
  },
34
  "when": ["dequantDtypesOk"],
35
  "bindings": {
36
  "params": {
37
- "buffer": "uniform",
38
  "struct": [
39
  { "name": "count", "type": "u32", "value": "elementCount" },
40
  { "name": "axisDim", "type": "u32", "value": "blockedAxisDim" },
@@ -43,39 +52,28 @@
43
  { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
44
  ]
45
  },
46
- "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$xVec4" },
47
- "x_scale_2": { "name": "x_scale", "buffer": "read-only-storage", "elementType": "$scaleVec4" },
48
- "y_2": { "name": "y", "buffer": "storage", "elementType": "$yVec4" },
49
- "params_2": {
50
  "name": "params",
51
- "buffer": "uniform",
52
  "struct": [
53
  { "name": "count4", "type": "u32", "value": "elementCount4" },
54
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" }
55
  ]
56
  },
57
- "params_3": {
58
  "name": "params",
59
- "buffer": "uniform",
60
  "struct": [
61
  { "name": "count4", "type": "u32", "value": "elementCount4" },
62
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
63
  { "name": "inner", "type": "u32", "value": "scaleInner" }
64
  ]
65
  },
66
- "params_4": {
67
- "name": "params",
68
- "buffer": "uniform",
69
- "struct": [{ "name": "count4", "type": "u32", "value": "elementCount4" }]
70
- },
71
- "params_5": {
72
  "name": "params",
73
- "buffer": "uniform",
74
- "struct": [{ "name": "count", "type": "u32", "value": "elementCount" }]
75
- },
76
- "params_6": {
77
- "name": "params",
78
- "buffer": "uniform",
79
  "struct": [
80
  { "name": "count", "type": "u32", "value": "elementCount" },
81
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
@@ -84,58 +82,23 @@
84
  }
85
  },
86
  "variants": [
87
- {
88
- "id": "blocked_with_zero_point",
89
- "priority": 19,
90
- "when": ["present.x_zero_point", "blockedScaleOk", "ranks.x_zero_point == ranks.x", "dim(shapes.x_zero_point, attrs.axis) == dim(shapes.x_scale, attrs.axis)"],
91
- "derive": { "hasZero": true },
92
- "passes": [
93
- {
94
- "id": "main",
95
- "name": "DequantizeLinear.Blocked",
96
- "shader": "quant-linear-blocked-axis.wgsl.jinja",
97
- "derive": { "op": "\"dequantize\"" },
98
- "bindings": ["x", "x_scale", "x_zero_point", "y", "params"],
99
- "dispatch": {
100
- "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
101
- "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
102
- "z": 1
103
- }
104
- }
105
- ]
106
- },
107
- {
108
- "id": "blocked_no_zero_point",
109
- "priority": 19,
110
- "when": ["not present.x_zero_point", "blockedScaleOk"],
111
- "derive": { "hasZero": false },
112
- "passes": [
113
- {
114
- "id": "main",
115
- "name": "DequantizeLinear.BlockedNoZero",
116
- "shader": "quant-linear-blocked-axis.wgsl.jinja",
117
- "derive": { "op": "\"dequantize\"" },
118
- "bindings": ["x", "x_scale", "y", "params"],
119
- "dispatch": {
120
- "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
121
- "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
122
- "z": 1
123
- }
124
- }
125
- ]
126
- },
127
  {
128
  "id": "innermost_axis_vec4_with_zero_point",
129
  "priority": 20,
130
- "when": ["present.x_zero_point", "ranks.x >= 1", "attrs.axis == -1 or attrs.axis == ranks.x - 1", "ranks.x_scale == 1", "dim(shapes.x_scale, 0) == dim(shapes.x, ranks.x - 1)", "dim(shapes.x_scale, 0) % 4 == 0", "ranks.x_zero_point == 1", "dim(shapes.x_zero_point, 0) == dim(shapes.x_scale, 0)", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "sameShapeOk"],
131
- "derive": { "hasZero": true },
132
  "passes": [
133
  {
134
  "id": "main",
135
  "name": "DequantizeLinear.InnermostAxisVec4",
136
  "shader": "quant-linear-vec4.wgsl.jinja",
137
- "derive": { "op": "\"dequantize\"", "perAxis": true, "vectorParams": true },
138
- "bindings": ["x_2", "x_scale_2", { "arg": "x_zero_point", "elementType": "$xVec4" }, "y_2", "params_2"],
 
 
 
 
 
 
139
  "dispatch": {
140
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
141
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -147,15 +110,14 @@
147
  {
148
  "id": "innermost_axis_vec4_no_zero_point",
149
  "priority": 20,
150
- "when": ["not present.x_zero_point", "ranks.x >= 1", "attrs.axis == -1 or attrs.axis == ranks.x - 1", "ranks.x_scale == 1", "dim(shapes.x_scale, 0) == dim(shapes.x, ranks.x - 1)", "dim(shapes.x_scale, 0) % 4 == 0", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "sameShapeOk"],
151
- "derive": { "hasZero": false },
152
  "passes": [
153
  {
154
  "id": "main",
155
  "name": "DequantizeLinear.InnermostAxisVec4NoZero",
156
  "shader": "quant-linear-vec4.wgsl.jinja",
157
- "derive": { "op": "\"dequantize\"", "perAxis": true, "vectorParams": true },
158
- "bindings": ["x_2", "x_scale_2", "y_2", "params_2"],
159
  "dispatch": {
160
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
161
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -167,15 +129,14 @@
167
  {
168
  "id": "vec4_cross_axis_with_zero_point",
169
  "priority": 16,
170
- "when": ["present.x_zero_point", "attrs.block_size == 0", "ranks.x >= 1", "attrs.axis >= 0", "attrs.axis < ranks.x", "ranks.x_scale == 1", "dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis)", "inner(shapes.x, attrs.axis) % 4 != 0", "ranks.x_zero_point == 1", "dim(shapes.x_zero_point, 0) == dim(shapes.x_scale, 0)", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "sameShapeOk"],
171
- "derive": { "hasZero": true },
172
  "passes": [
173
  {
174
  "id": "main",
175
  "name": "DequantizeLinear.Vec4CrossAxis",
176
  "shader": "quant-linear-vec4.wgsl.jinja",
177
- "derive": { "op": "\"dequantize\"", "perAxis": true, "crossingParams": true },
178
- "bindings": ["x_2", "x_scale", "x_zero_point", "y_2", "params_3"],
179
  "dispatch": {
180
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
181
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -187,15 +148,14 @@
187
  {
188
  "id": "vec4_with_zero_point",
189
  "priority": 15,
190
- "when": ["present.x_zero_point", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "(ranks.x_scale == 0 or ranks.x_scale == 1)", "sameShapeOk", "(ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))", "(ranks.x_zero_point == 0 or ranks.x_zero_point == 1)", "(ranks.x_zero_point == 0 or dim(shapes.x_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_zero_point, 0) == dim(shapes.x, attrs.axis)))"],
191
- "derive": { "hasZero": true },
192
  "passes": [
193
  {
194
  "id": "main",
195
  "name": "DequantizeLinear.Vec4",
196
  "shader": "quant-linear-vec4.wgsl.jinja",
197
- "derive": { "op": "\"dequantize\"", "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)" },
198
- "bindings": ["x_2", "x_scale", "x_zero_point", "y_2", "params_3"],
199
  "dispatch": {
200
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
201
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -207,15 +167,14 @@
207
  {
208
  "id": "vec4_no_zero_point",
209
  "priority": 15,
210
- "when": ["not present.x_zero_point", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "(ranks.x_scale == 0 or ranks.x_scale == 1)", "sameShapeOk", "(ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))"],
211
- "derive": { "hasZero": false },
212
  "passes": [
213
  {
214
  "id": "main",
215
  "name": "DequantizeLinear.Vec4",
216
  "shader": "quant-linear-vec4.wgsl.jinja",
217
- "derive": { "op": "\"dequantize\"", "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)" },
218
- "bindings": ["x_2", "x_scale", "y_2", "params_3"],
219
  "dispatch": {
220
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
221
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -227,15 +186,14 @@
227
  {
228
  "id": "vec4_tail_with_zero_point",
229
  "priority": 10,
230
- "when": ["present.x_zero_point", "numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS", "numel(shapes.y) % 4 != 0", "(ranks.x_scale == 0 or (ranks.x_scale == 1 and dim(shapes.x_scale, 0) == 1))", "(ranks.x_zero_point == 0 or (ranks.x_zero_point == 1 and dim(shapes.x_zero_point, 0) == 1))", "sameShapeOk"],
231
- "derive": { "hasZero": true },
232
  "passes": [
233
  {
234
  "id": "bulk",
235
  "name": "DequantizeLinear.Vec4Bulk",
236
  "shader": "quant-linear-vec4.wgsl.jinja",
237
- "derive": { "op": "\"dequantize\"", "perAxis": false },
238
- "bindings": ["x_2", "x_scale", "x_zero_point", "y_2", "params_4"],
239
  "dispatch": {
240
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
241
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -246,8 +204,8 @@
246
  "id": "tail",
247
  "name": "DequantizeLinear.ScalarTail",
248
  "shader": "quant-linear-scalar.wgsl.jinja",
249
- "derive": { "op": "\"dequantize\"", "x4": true, "perAxis": false },
250
- "bindings": ["x", "x_scale", "x_zero_point", "y", "params_5"],
251
  "dispatch": {
252
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
253
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -259,15 +217,14 @@
259
  {
260
  "id": "vec4_tail_no_zero_point",
261
  "priority": 10,
262
- "when": ["not present.x_zero_point", "numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS", "numel(shapes.y) % 4 != 0", "(ranks.x_scale == 0 or (ranks.x_scale == 1 and dim(shapes.x_scale, 0) == 1))", "sameShapeOk"],
263
- "derive": { "hasZero": false },
264
  "passes": [
265
  {
266
  "id": "bulk",
267
  "name": "DequantizeLinear.Vec4Bulk",
268
  "shader": "quant-linear-vec4.wgsl.jinja",
269
- "derive": { "op": "\"dequantize\"", "perAxis": false },
270
- "bindings": ["x_2", "x_scale", "y_2", "params_4"],
271
  "dispatch": {
272
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
273
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -278,8 +235,8 @@
278
  "id": "tail",
279
  "name": "DequantizeLinear.ScalarTail",
280
  "shader": "quant-linear-scalar.wgsl.jinja",
281
- "derive": { "op": "\"dequantize\"", "x4": true, "perAxis": false },
282
- "bindings": ["x", "x_scale", "y", "params_5"],
283
  "dispatch": {
284
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
285
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -288,21 +245,52 @@
288
  }
289
  ]
290
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
291
  {
292
  "id": "linear_scalar_or_axis_zero",
293
  "when": ["present.x_zero_point", "(ranks.x_scale == 0 or ranks.x_scale == 1)", "sameShapeOk", "(ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis)))", "(ranks.x_zero_point == 0 or ranks.x_zero_point == 1)", "(ranks.x_zero_point == 0 or dim(shapes.x_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_zero_point, 0) == dim(shapes.x, attrs.axis)))"],
294
- "derive": { "hasZero": true },
295
  "passes": [
296
  {
297
  "id": "main",
298
  "name": "DequantizeLinear",
299
  "shader": "quant-linear-scalar.wgsl.jinja",
300
- "derive": {
301
- "op": "\"dequantize\"",
302
- "x4": "false",
303
- "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)"
304
- },
305
- "bindings": ["x", "x_scale", "x_zero_point", "y", "params_6"],
306
  "dispatch": {
307
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
308
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -314,18 +302,13 @@
314
  {
315
  "id": "linear_scalar_or_axis_no_zero",
316
  "when": ["not present.x_zero_point", "(ranks.x_scale == 0 or ranks.x_scale == 1)", "sameShapeOk", "(ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis)))"],
317
- "derive": { "hasZero": false },
318
  "passes": [
319
  {
320
  "id": "main",
321
  "name": "DequantizeLinear",
322
  "shader": "quant-linear-scalar.wgsl.jinja",
323
- "derive": {
324
- "op": "\"dequantize\"",
325
- "x4": "false",
326
- "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)"
327
- },
328
- "bindings": ["x", "x_scale", "y", "params_6"],
329
  "dispatch": {
330
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
331
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
 
15
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
16
  "workgroupOk": "tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
17
  "outputDispatchFits": "workgroupOk and ceilDiv(ceilDiv(numel(shapes.y), tunables.WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
18
+ "elementCount": "numel(shapes.y)",
19
+ "elementCount4": "floor(numel(shapes.y) / 4)",
20
+ "blockedAxisDim": "dim(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
21
+ "blockedInner": "inner(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
22
  "sameShapeOk": "ranks.y == ranks.x and numel(shapes.x) == numel(shapes.y) and outputDispatchFits",
23
  "dequantDtypesOk": "(attrs.output_dtype == 0 or attrs.output_dtype == onnxDtypeCode(logicalDtypes.TF)) and f16Ok(dtypes.TF)",
24
  "blockedScaleOk": "attrs.block_size > 0 and ranks.x >= 2 and ranks.x_scale == ranks.x and sameShapeOk and dim(shapes.x_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)",
25
+ "hasZero": "present.x_zero_point",
 
26
  "scaleSize": "1 if ranks.x_scale == 0 else dim(shapes.x_scale, 0)",
27
  "scaleInner": "1 if ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 else inner(shapes.x, attrs.axis)",
 
28
  "blockedScaleAxisDim": "dim(shapes.x_scale, attrs.axis) if ranks.x_scale >= 2 else 1",
 
29
  "xVec4": "\"vec4<\" ~ dtypes.TQ ~ \">\"",
30
  "scaleVec4": "\"vec4<\" ~ dtypes.TF ~ \">\"",
31
  "yScalar": "dtypes.TF",
32
  "yVec4": "\"vec4<\" ~ dtypes.TF ~ \">\"",
33
+ "xUnsigned": "tensorDtypes.x == \"uint8\"",
34
+ "routeInnermostAxisVec4WithZeroPoint": "(present.x_zero_point) and (ranks.x >= 1) and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and (ranks.x_scale == 1) and (dim(shapes.x_scale, 0) == dim(shapes.x, ranks.x - 1)) and (dim(shapes.x_scale, 0) % 4 == 0) and (ranks.x_zero_point == 1) and (dim(shapes.x_zero_point, 0) == dim(shapes.x_scale, 0)) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and (sameShapeOk)",
35
+ "routeInnermostAxisVec4NoZeroPoint": "(not present.x_zero_point) and (ranks.x >= 1) and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and (ranks.x_scale == 1) and (dim(shapes.x_scale, 0) == dim(shapes.x, ranks.x - 1)) and (dim(shapes.x_scale, 0) % 4 == 0) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and (sameShapeOk)",
36
+ "routeVec4CrossAxisWithZeroPoint": "(present.x_zero_point) and (attrs.block_size == 0) and (ranks.x >= 1) and (attrs.axis >= 0) and (attrs.axis < ranks.x) and (ranks.x_scale == 1) and (dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis)) and (inner(shapes.x, attrs.axis) % 4 != 0) and (ranks.x_zero_point == 1) and (dim(shapes.x_zero_point, 0) == dim(shapes.x_scale, 0)) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and (sameShapeOk)",
37
+ "routeVec4WithZeroPoint": "(present.x_zero_point) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and ((ranks.x_scale == 0 or ranks.x_scale == 1)) and (sameShapeOk) and ((ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))) and ((ranks.x_zero_point == 0 or ranks.x_zero_point == 1)) and ((ranks.x_zero_point == 0 or dim(shapes.x_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_zero_point, 0) == dim(shapes.x, attrs.axis))))",
38
+ "routeVec4NoZeroPoint": "(not present.x_zero_point) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and ((ranks.x_scale == 0 or ranks.x_scale == 1)) and (sameShapeOk) and ((ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0)))",
39
+ "routeVec4TailWithZeroPoint": "(present.x_zero_point) and (numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS) and (numel(shapes.y) % 4 != 0) and ((ranks.x_scale == 0 or (ranks.x_scale == 1 and dim(shapes.x_scale, 0) == 1))) and ((ranks.x_zero_point == 0 or (ranks.x_zero_point == 1 and dim(shapes.x_zero_point, 0) == 1))) and (sameShapeOk)",
40
+ "routeVec4TailNoZeroPoint": "(not present.x_zero_point) and (numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS) and (numel(shapes.y) % 4 != 0) and ((ranks.x_scale == 0 or (ranks.x_scale == 1 and dim(shapes.x_scale, 0) == 1))) and (sameShapeOk)",
41
+ "routeBlockedWithZeroPoint": "(present.x_zero_point) and (blockedScaleOk) and (ranks.x_zero_point == ranks.x) and (dim(shapes.x_zero_point, attrs.axis) == dim(shapes.x_scale, attrs.axis))",
42
+ "routeBlockedNoZeroPoint": "(not present.x_zero_point) and (blockedScaleOk)"
43
  },
44
  "when": ["dequantDtypesOk"],
45
  "bindings": {
46
  "params": {
 
47
  "struct": [
48
  { "name": "count", "type": "u32", "value": "elementCount" },
49
  { "name": "axisDim", "type": "u32", "value": "blockedAxisDim" },
 
52
  { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
53
  ]
54
  },
55
+ "x_main": { "name": "x", "elementType": "$xVec4" },
56
+ "x_scale_main": { "name": "x_scale", "elementType": "$scaleVec4" },
57
+ "y_main": { "name": "y", "elementType": "$yVec4" },
58
+ "params_main": {
59
  "name": "params",
 
60
  "struct": [
61
  { "name": "count4", "type": "u32", "value": "elementCount4" },
62
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" }
63
  ]
64
  },
65
+ "params__uniform": {
66
  "name": "params",
 
67
  "struct": [
68
  { "name": "count4", "type": "u32", "value": "elementCount4" },
69
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
70
  { "name": "inner", "type": "u32", "value": "scaleInner" }
71
  ]
72
  },
73
+ "params_bulk": { "name": "params", "struct": [{ "name": "count4", "type": "u32", "value": "elementCount4" }] },
74
+ "params_tail": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "elementCount" }] },
75
+ "params_quant_linear_scalar": {
 
 
 
76
  "name": "params",
 
 
 
 
 
 
77
  "struct": [
78
  { "name": "count", "type": "u32", "value": "elementCount" },
79
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
 
82
  }
83
  },
84
  "variants": [
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85
  {
86
  "id": "innermost_axis_vec4_with_zero_point",
87
  "priority": 20,
88
+ "when": ["routeInnermostAxisVec4WithZeroPoint"],
 
89
  "passes": [
90
  {
91
  "id": "main",
92
  "name": "DequantizeLinear.InnermostAxisVec4",
93
  "shader": "quant-linear-vec4.wgsl.jinja",
94
+ "derive": { "perAxis": true, "vectorParams": true },
95
+ "bindings": [
96
+ "x_main",
97
+ "x_scale_main",
98
+ { "arg": "x_zero_point", "elementType": "$xVec4" },
99
+ "y_main",
100
+ "params_main"
101
+ ],
102
  "dispatch": {
103
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
104
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
110
  {
111
  "id": "innermost_axis_vec4_no_zero_point",
112
  "priority": 20,
113
+ "when": ["routeInnermostAxisVec4NoZeroPoint"],
 
114
  "passes": [
115
  {
116
  "id": "main",
117
  "name": "DequantizeLinear.InnermostAxisVec4NoZero",
118
  "shader": "quant-linear-vec4.wgsl.jinja",
119
+ "derive": { "perAxis": true, "vectorParams": true },
120
+ "bindings": ["x_main", "x_scale_main", "y_main", "params_main"],
121
  "dispatch": {
122
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
123
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
129
  {
130
  "id": "vec4_cross_axis_with_zero_point",
131
  "priority": 16,
132
+ "when": ["routeVec4CrossAxisWithZeroPoint"],
 
133
  "passes": [
134
  {
135
  "id": "main",
136
  "name": "DequantizeLinear.Vec4CrossAxis",
137
  "shader": "quant-linear-vec4.wgsl.jinja",
138
+ "derive": { "perAxis": true, "crossingParams": true },
139
+ "bindings": ["x_main", "x_scale", "x_zero_point", "y_main", "params__uniform"],
140
  "dispatch": {
141
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
142
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
148
  {
149
  "id": "vec4_with_zero_point",
150
  "priority": 15,
151
+ "when": ["routeVec4WithZeroPoint"],
 
152
  "passes": [
153
  {
154
  "id": "main",
155
  "name": "DequantizeLinear.Vec4",
156
  "shader": "quant-linear-vec4.wgsl.jinja",
157
+ "derive": { "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)" },
158
+ "bindings": ["x_main", "x_scale", "x_zero_point", "y_main", "params__uniform"],
159
  "dispatch": {
160
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
161
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
167
  {
168
  "id": "vec4_no_zero_point",
169
  "priority": 15,
170
+ "when": ["routeVec4NoZeroPoint"],
 
171
  "passes": [
172
  {
173
  "id": "main",
174
  "name": "DequantizeLinear.Vec4",
175
  "shader": "quant-linear-vec4.wgsl.jinja",
176
+ "derive": { "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)" },
177
+ "bindings": ["x_main", "x_scale", "y_main", "params__uniform"],
178
  "dispatch": {
179
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
180
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
186
  {
187
  "id": "vec4_tail_with_zero_point",
188
  "priority": 10,
189
+ "when": ["routeVec4TailWithZeroPoint"],
 
190
  "passes": [
191
  {
192
  "id": "bulk",
193
  "name": "DequantizeLinear.Vec4Bulk",
194
  "shader": "quant-linear-vec4.wgsl.jinja",
195
+ "derive": { "perAxis": false },
196
+ "bindings": ["x_main", "x_scale", "x_zero_point", "y_main", "params_bulk"],
197
  "dispatch": {
198
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
199
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
204
  "id": "tail",
205
  "name": "DequantizeLinear.ScalarTail",
206
  "shader": "quant-linear-scalar.wgsl.jinja",
207
+ "derive": { "x4": true, "perAxis": false },
208
+ "bindings": ["x", "x_scale", "x_zero_point", "y", "params_tail"],
209
  "dispatch": {
210
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
211
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
 
217
  {
218
  "id": "vec4_tail_no_zero_point",
219
  "priority": 10,
220
+ "when": ["routeVec4TailNoZeroPoint"],
 
221
  "passes": [
222
  {
223
  "id": "bulk",
224
  "name": "DequantizeLinear.Vec4Bulk",
225
  "shader": "quant-linear-vec4.wgsl.jinja",
226
+ "derive": { "perAxis": false },
227
+ "bindings": ["x_main", "x_scale", "y_main", "params_bulk"],
228
  "dispatch": {
229
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
230
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
235
  "id": "tail",
236
  "name": "DequantizeLinear.ScalarTail",
237
  "shader": "quant-linear-scalar.wgsl.jinja",
238
+ "derive": { "x4": true, "perAxis": false },
239
+ "bindings": ["x", "x_scale", "y", "params_tail"],
240
  "dispatch": {
241
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
242
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
 
245
  }
246
  ]
247
  },
248
+ {
249
+ "id": "blocked_with_zero_point",
250
+ "priority": 19,
251
+ "when": ["routeBlockedWithZeroPoint"],
252
+ "passes": [
253
+ {
254
+ "id": "main",
255
+ "name": "DequantizeLinear.Blocked",
256
+ "shader": "quant-linear-blocked-axis.wgsl.jinja",
257
+ "bindings": ["x", "x_scale", "x_zero_point", "y", "params"],
258
+ "dispatch": {
259
+ "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
260
+ "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
261
+ "z": 1
262
+ }
263
+ }
264
+ ]
265
+ },
266
+ {
267
+ "id": "blocked_no_zero_point",
268
+ "priority": 19,
269
+ "when": ["routeBlockedNoZeroPoint"],
270
+ "passes": [
271
+ {
272
+ "id": "main",
273
+ "name": "DequantizeLinear.BlockedNoZero",
274
+ "shader": "quant-linear-blocked-axis.wgsl.jinja",
275
+ "bindings": ["x", "x_scale", "y", "params"],
276
+ "dispatch": {
277
+ "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
278
+ "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
279
+ "z": 1
280
+ }
281
+ }
282
+ ]
283
+ },
284
  {
285
  "id": "linear_scalar_or_axis_zero",
286
  "when": ["present.x_zero_point", "(ranks.x_scale == 0 or ranks.x_scale == 1)", "sameShapeOk", "(ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis)))", "(ranks.x_zero_point == 0 or ranks.x_zero_point == 1)", "(ranks.x_zero_point == 0 or dim(shapes.x_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_zero_point, 0) == dim(shapes.x, attrs.axis)))"],
 
287
  "passes": [
288
  {
289
  "id": "main",
290
  "name": "DequantizeLinear",
291
  "shader": "quant-linear-scalar.wgsl.jinja",
292
+ "derive": { "x4": "false", "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)" },
293
+ "bindings": ["x", "x_scale", "x_zero_point", "y", "params_quant_linear_scalar"],
 
 
 
 
294
  "dispatch": {
295
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
296
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
 
302
  {
303
  "id": "linear_scalar_or_axis_no_zero",
304
  "when": ["not present.x_zero_point", "(ranks.x_scale == 0 or ranks.x_scale == 1)", "sameShapeOk", "(ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.x_scale, 0) == dim(shapes.x, attrs.axis)))"],
 
305
  "passes": [
306
  {
307
  "id": "main",
308
  "name": "DequantizeLinear",
309
  "shader": "quant-linear-scalar.wgsl.jinja",
310
+ "derive": { "x4": "false", "perAxis": "not (ranks.x_scale == 0 or dim(shapes.x_scale, 0) == 1)" },
311
+ "bindings": ["x", "x_scale", "y", "params_quant_linear_scalar"],
 
 
 
 
312
  "dispatch": {
313
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
314
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
build/webgpu/metadata.json CHANGED
@@ -1,26 +1,24 @@
1
  {
2
  "name": "ai.onnx.DequantizeLinear",
3
- "id": "_ai_onnx_dequantizelinear_webgpu_90d93ed",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "fPBI9jQFz3Fksyma9pLyaW1Ivekq1sXErhIhAlvM8Mw=",
11
- "manifest.json": "omGyGN9P5BpbXt0tvvl0AGNEBPfsAb1Zj82wlH2H/2A=",
12
- "quant-linear-blocked-axis.wgsl.jinja": "Fe40jY40I0VeefzHQ3D11z7+i4brKhQktHgjb4oUl54=",
13
- "quant-linear-scalar.wgsl.jinja": "AezdMTq8nS5MZ2qAxNCCleeOYs4NK0taL6reMDMFscQ=",
14
- "quant-linear-vec4.wgsl.jinja": "ictcyAzu8sMWImMrZ6hLS3drt4yttxbIn/MjotR8RmE=",
15
- "test.json": "Cm0FRA9IYSR+GReOqHoCO/Y4ay5K5NFV3iW4eEBokWY="
16
  }
17
  },
18
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
19
  "webgpu": {
20
- "manifestSpec": "2.0",
21
  "variants": {
22
- "blocked_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
23
- "blocked_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
24
  "innermost_axis_vec4_with_zero_point": ["quant-linear-vec4.wgsl.jinja"],
25
  "innermost_axis_vec4_no_zero_point": ["quant-linear-vec4.wgsl.jinja"],
26
  "vec4_cross_axis_with_zero_point": ["quant-linear-vec4.wgsl.jinja"],
@@ -28,6 +26,8 @@
28
  "vec4_no_zero_point": ["quant-linear-vec4.wgsl.jinja"],
29
  "vec4_tail_with_zero_point": ["quant-linear-scalar.wgsl.jinja", "quant-linear-vec4.wgsl.jinja"],
30
  "vec4_tail_no_zero_point": ["quant-linear-scalar.wgsl.jinja", "quant-linear-vec4.wgsl.jinja"],
 
 
31
  "linear_scalar_or_axis_zero": ["quant-linear-scalar.wgsl.jinja"],
32
  "linear_scalar_or_axis_no_zero": ["quant-linear-scalar.wgsl.jinja"]
33
  }
 
1
  {
2
  "name": "ai.onnx.DequantizeLinear",
3
+ "id": "_ai_onnx_dequantizelinear_webgpu_d91b472",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "uGFDTwvqf+xiM9KA340XUzCNls85W0NdfNK11RkOv5I=",
11
+ "manifest.json": "+nZgBcFSOLVyTiqFHUX7mbVHkWaK1sJU+TorK3j4pcc=",
12
+ "quant-linear-blocked-axis.wgsl.jinja": "ppwPl0wzYvhilArhWPxHyYP0iRcQQIGEjuwOOw6hxiE=",
13
+ "quant-linear-scalar.wgsl.jinja": "CUMvooQjegnIraNbmb/MrkAY932M9g9EiEmgIVni1y8=",
14
+ "quant-linear-vec4.wgsl.jinja": "xcjv8IaRYMpZIP2sg/HGe0zrdWhRLq8aoDMk2SB2swc=",
15
+ "test.json": "rq8/XlKSoSjs0pG6M+snPJ4qbIJpXfdl7wFnr3p7eBU="
16
  }
17
  },
18
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
19
  "webgpu": {
20
+ "manifestSpec": "2.1",
21
  "variants": {
 
 
22
  "innermost_axis_vec4_with_zero_point": ["quant-linear-vec4.wgsl.jinja"],
23
  "innermost_axis_vec4_no_zero_point": ["quant-linear-vec4.wgsl.jinja"],
24
  "vec4_cross_axis_with_zero_point": ["quant-linear-vec4.wgsl.jinja"],
 
26
  "vec4_no_zero_point": ["quant-linear-vec4.wgsl.jinja"],
27
  "vec4_tail_with_zero_point": ["quant-linear-scalar.wgsl.jinja", "quant-linear-vec4.wgsl.jinja"],
28
  "vec4_tail_no_zero_point": ["quant-linear-scalar.wgsl.jinja", "quant-linear-vec4.wgsl.jinja"],
29
+ "blocked_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
30
+ "blocked_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
31
  "linear_scalar_or_axis_zero": ["quant-linear-scalar.wgsl.jinja"],
32
  "linear_scalar_or_axis_no_zero": ["quant-linear-scalar.wgsl.jinja"]
33
  }
build/webgpu/quant-linear-blocked-axis.wgsl.jinja CHANGED
@@ -1,10 +1,10 @@
1
  // Blocked-axis scale and zero-point indexing for linear quantization.
2
- {% set lastAxisVectorized = lastAxisVectorized if lastAxisVectorized is defined else false %}
3
- {% set xUnsigned = xUnsigned is defined and xUnsigned %}
4
- {% set yUnsigned = yUnsigned is defined and yUnsigned %}
5
- {% set qMin = qMin | default(0 if yUnsigned else 0 - 128) %}
6
- {% set qMax = qMax | default(255 if yUnsigned else 127) %}
7
- {% set blockVectors = blockVectors | default(1) %}
8
  {{ env.wgsl.resourceDeclarations }}
9
 
10
  fn read_x(index: u32) -> i32 {
@@ -27,7 +27,6 @@ fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
27
  {% endif %}
28
  }
29
 
30
- {% if not lastAxisVectorized %}
31
  fn scale_index(index: u32) -> u32 {
32
  let inner_index = index % params.inner;
33
  let axis_index = (index / params.inner) % params.axisDim;
@@ -35,47 +34,10 @@ fn scale_index(index: u32) -> u32 {
35
  return (outer_index * params.scaleAxisDim + axis_index / params.blockSize) * params.inner + inner_index;
36
  }
37
 
38
- {% endif %}
39
- {% if lastAxisVectorized %}
40
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
41
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
42
- // A vec4 cannot cross a block boundary on this route. Map vector indices
43
- // directly to their shared scale/zero-point entry.
44
- let i4 = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
45
- if (i4 >= params.count) { return; }
46
- let scale_index = i4 / {{ blockVectors }}u;
47
- let values = x[i4];
48
- let scale = f32(y_scale[scale_index]);
49
- {% if hasZero %}
50
- {% if yUnsigned %}
51
- let zero = i32(y_zero_point[scale_index]);
52
- {% else %}
53
- let zero = y_zero_point[scale_index];
54
- {% endif %}
55
- {% else %}
56
- let zero = 0;
57
- {% endif %}
58
- let quantized = vec4<i32>(
59
- clamp(round_scaled_half_to_even(f32(values.x), scale) + zero, {{ qMin }}, {{ qMax }}),
60
- clamp(round_scaled_half_to_even(f32(values.y), scale) + zero, {{ qMin }}, {{ qMax }}),
61
- clamp(round_scaled_half_to_even(f32(values.z), scale) + zero, {{ qMin }}, {{ qMax }}),
62
- clamp(round_scaled_half_to_even(f32(values.w), scale) + zero, {{ qMin }}, {{ qMax }})
63
- );
64
- {% if yUnsigned %}
65
- y[i4] = vec4<u32>(quantized);
66
- {% else %}
67
- y[i4] = quantized;
68
- {% endif %}
69
- }
70
- {% else %}
71
- @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
72
- fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
73
- // The flat dispatch is folded across x/y at a fixed per-axis workgroup
74
- // width; gid.y carries the high portion of the element index.
75
- let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
76
- if (index >= params.count) { return; }
77
  let scaleIndex = scale_index(index);
78
  let result = f32(read_x(index) - read_zero({% if hasZero %}scaleIndex{% endif %})) * f32(x_scale[scaleIndex]);
79
  y[index] = {{ yScalar }}(result);
80
  }
81
- {% endif %}
 
1
  // Blocked-axis scale and zero-point indexing for linear quantization.
2
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
3
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
4
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
5
+ // per-axis workgroup fold width.
6
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
7
+ if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
8
  {{ env.wgsl.resourceDeclarations }}
9
 
10
  fn read_x(index: u32) -> i32 {
 
27
  {% endif %}
28
  }
29
 
 
30
  fn scale_index(index: u32) -> u32 {
31
  let inner_index = index % params.inner;
32
  let axis_index = (index / params.inner) % params.axisDim;
 
34
  return (outer_index * params.scaleAxisDim + axis_index / params.blockSize) * params.inner + inner_index;
35
  }
36
 
 
 
37
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
38
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
39
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "index", guardInline=true) }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
  let scaleIndex = scale_index(index);
41
  let result = f32(read_x(index) - read_zero({% if hasZero %}scaleIndex{% endif %})) * f32(x_scale[scaleIndex]);
42
  y[index] = {{ yScalar }}(result);
43
  }
 
build/webgpu/quant-linear-scalar.wgsl.jinja CHANGED
@@ -1,6 +1,10 @@
1
- {% set xUnsigned = xUnsigned is defined and xUnsigned %}
2
  // The x4 route handles widened int8/uint8 storage with scalar scale parameters;
3
  // the linear route supports both per-tensor and per-axis quantization.
 
 
 
 
 
4
  {{ env.wgsl.resourceDeclarations }}
5
 
6
  fn read_x(index: u32) -> i32 {
@@ -30,9 +34,7 @@ fn transform_one(index: u32, scale: f32, zero_point: i32) {
30
 
31
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
32
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
33
- // Fold oversized flat dispatches into two dimensions; gid.y carries work
34
- // beyond the device-capped x dimension.
35
- let invocation = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
36
  {% if x4 %}
37
  if (invocation != 0u) {
38
  return;
 
 
1
  // The x4 route handles widened int8/uint8 storage with scalar scale parameters;
2
  // the linear route supports both per-tensor and per-axis quantization.
3
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
4
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
5
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
6
+ // per-axis workgroup fold width.
7
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
8
  {{ env.wgsl.resourceDeclarations }}
9
 
10
  fn read_x(index: u32) -> i32 {
 
34
 
35
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
36
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
37
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "invocation", "") }}
 
 
38
  {% if x4 %}
39
  if (invocation != 0u) {
40
  return;
build/webgpu/quant-linear-vec4.wgsl.jinja CHANGED
@@ -1,7 +1,14 @@
1
- {% set xUnsigned = xUnsigned is defined and xUnsigned %}
2
  // int8/uint8 tensors are stored widened (one u32/i32 per element), so the vec4
3
  // binding gives 128-bit loads/stores of four elements. Per-component arithmetic
4
  // remains identical to the scalar quantize/dequantize paths.
 
 
 
 
 
 
 
 
5
  {{ env.wgsl.resourceDeclarations }}
6
  {% if vectorParams is defined and vectorParams %}
7
  fn read_zero4({% if hasZero %}index: u32{% endif %}) -> vec4<i32> {
@@ -31,12 +38,7 @@ fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
31
 
32
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
33
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
34
- // The flat dispatch is folded across x/y at a fixed per-axis workgroup
35
- // width; gid.y carries the high portion of the vector index.
36
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
37
- if (i >= params.count4) {
38
- return;
39
- }
40
  {% if vectorParams is defined and vectorParams %}
41
  let scale_index = i % (params.scaleSize / 4u);
42
  {% elif crossingParams is defined and crossingParams %}
 
 
1
  // int8/uint8 tensors are stored widened (one u32/i32 per element), so the vec4
2
  // binding gives 128-bit loads/stores of four elements. Per-component arithmetic
3
  // remains identical to the scalar quantize/dequantize paths.
4
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
5
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
6
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
7
+ // per-axis workgroup fold width.
8
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
9
+ if ({{ name }} >= {{ bound }}) {
10
+ return;
11
+ }{% endmacro %}
12
  {{ env.wgsl.resourceDeclarations }}
13
  {% if vectorParams is defined and vectorParams %}
14
  fn read_zero4({% if hasZero %}index: u32{% endif %}) -> vec4<i32> {
 
38
 
39
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
40
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
41
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, bound="params.count4") }}
 
 
 
 
 
42
  {% if vectorParams is defined and vectorParams %}
43
  let scale_index = i % (params.scaleSize / 4u);
44
  {% elif crossingParams is defined and crossingParams %}
build/webgpu/test.json CHANGED
@@ -531,7 +531,7 @@
531
  {
532
  "name": "int8_no_zero_scalar_explicit_float32_output_dtype",
533
  "provenance": {
534
- "notes": "Explicit output_dtype=FLOAT is the non-default spelling of the same float32 route selected from x_scale when output_dtype is omitted."
535
  },
536
  "attrs": { "axis": 0, "output_dtype": 1 },
537
  "inputs": {
@@ -970,7 +970,7 @@
970
  {
971
  "name": "vec4_cross_axis_rows_width6_with_zero_point",
972
  "provenance": {
973
- "notes": "A vec4 crosses each six-element row boundary, so its four lanes may use two different per-axis scale and zero-point entries. Locks the lane-specific parameter-index path."
974
  },
975
  "attrs": { "axis": 0 },
976
  "inputs": {
 
531
  {
532
  "name": "int8_no_zero_scalar_explicit_float32_output_dtype",
533
  "provenance": {
534
+ "notes": "Explicit output_dtype=FLOAT matches the float32 output inferred from x_scale when the attribute is omitted."
535
  },
536
  "attrs": { "axis": 0, "output_dtype": 1 },
537
  "inputs": {
 
970
  {
971
  "name": "vec4_cross_axis_rows_width6_with_zero_point",
972
  "provenance": {
973
+ "notes": "Four consecutive values cross a six-element row boundary, so they can require different per-axis scales and zero points."
974
  },
975
  "attrs": { "axis": 0 },
976
  "inputs": {