Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.QuantizeLinear

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
ce7680e
·
verified ·
1 Parent(s): 8b49b15

sync 6fdf6301e2bb

Browse files
Files changed (8) hide show
  1. README.md +7 -4
  2. build/webgpu/bench.json +3 -3
  3. build/webgpu/manifest.json +99 -107
  4. build/webgpu/metadata.json +11 -11
  5. build/webgpu/quant-linear-blocked-axis.wgsl.jinja +14 -26
  6. build/webgpu/quant-linear-scalar.wgsl.jinja +6 -8
  7. build/webgpu/quant-linear-vec4.wgsl.jinja +9 -11
  8. build/webgpu/test.json +27 -6
README.md CHANGED
@@ -55,7 +55,7 @@ Default values (overridable per request):
55
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
56
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
57
  - [`test.json`](build/webgpu/test.json) — correctness cases
58
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
59
  - [`quant-linear-blocked-axis.wgsl.jinja`](build/webgpu/quant-linear-blocked-axis.wgsl.jinja)
60
  - [`quant-linear-scalar.wgsl.jinja`](build/webgpu/quant-linear-scalar.wgsl.jinja)
61
  - [`quant-linear-vec4.wgsl.jinja`](build/webgpu/quant-linear-vec4.wgsl.jinja)
@@ -63,7 +63,7 @@ Default values (overridable per request):
63
  ## Use with `@huggingface/kernels`
64
 
65
  ```sh
66
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
67
  ```
68
 
69
  Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
@@ -82,7 +82,10 @@ import { getKernel } from "@huggingface/kernels";
82
 
83
  const kernel = await getKernel("webgpu-kernels/ai.onnx.QuantizeLinear", { version: 1 });
84
  // Explicit destinations request optional results or supply metadata that cannot be inferred.
85
- const { y } = await kernel({ x: { data: xData, shape: [] }, y_scale: { data: y_scaleData, shape: [] } }, {
86
- outputs: { y: { shape: [], dtype: "uint8" } },
 
 
 
87
  });
88
  ```
 
55
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
56
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
57
  - [`test.json`](build/webgpu/test.json) — correctness cases
58
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
59
  - [`quant-linear-blocked-axis.wgsl.jinja`](build/webgpu/quant-linear-blocked-axis.wgsl.jinja)
60
  - [`quant-linear-scalar.wgsl.jinja`](build/webgpu/quant-linear-scalar.wgsl.jinja)
61
  - [`quant-linear-vec4.wgsl.jinja`](build/webgpu/quant-linear-vec4.wgsl.jinja)
 
63
  ## Use with `@huggingface/kernels`
64
 
65
  ```sh
66
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
67
  ```
68
 
69
  Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
 
82
 
83
  const kernel = await getKernel("webgpu-kernels/ai.onnx.QuantizeLinear", { version: 1 });
84
  // Explicit destinations request optional results or supply metadata that cannot be inferred.
85
+ const { y } = await kernel({
86
+ x: { data: xData, shape: [1] },
87
+ y_scale: { data: y_scaleData, shape: [1] },
88
+ }, {
89
+ outputs: { y: { shape: [1], dtype: "uint8" } },
90
  });
91
  ```
build/webgpu/bench.json CHANGED
@@ -123,7 +123,7 @@
123
  }
124
  },
125
  {
126
- "name": "pertensor_4096x4096_scalar_healthy_sibling",
127
  "preset": "smoke",
128
  "vars": { "inputDtype": "float32", "outputDtype": "uint8", "count": 16777216 },
129
  "inputs": {
@@ -168,7 +168,7 @@
168
  }
169
  },
170
  {
171
- "name": "per_axis0_inner2048_vec4_healthy_sibling",
172
  "preset": "smoke",
173
  "vars": { "inputDtype": "float32", "outputDtype": "uint8", "count": 1048576 },
174
  "attrs": { "axis": 0 },
@@ -193,7 +193,7 @@
193
  }
194
  },
195
  {
196
- "name": "f16_per_axis0_inner2048_vec4_unbenched",
197
  "preset": "smoke",
198
  "vars": { "inputDtype": "float16", "outputDtype": "uint8", "count": 1048576 },
199
  "attrs": { "axis": 0 },
 
123
  }
124
  },
125
  {
126
+ "name": "pertensor_4096x4096_scalar_control_sibling",
127
  "preset": "smoke",
128
  "vars": { "inputDtype": "float32", "outputDtype": "uint8", "count": 16777216 },
129
  "inputs": {
 
168
  }
169
  },
170
  {
171
+ "name": "per_axis0_inner2048_vec4_control_sibling",
172
  "preset": "smoke",
173
  "vars": { "inputDtype": "float32", "outputDtype": "uint8", "count": 1048576 },
174
  "attrs": { "axis": 0 },
 
193
  }
194
  },
195
  {
196
+ "name": "f16_per_axis0_inner2048_vec4",
197
  "preset": "smoke",
198
  "vars": { "inputDtype": "float16", "outputDtype": "uint8", "count": 1048576 },
199
  "attrs": { "axis": 0 },
build/webgpu/manifest.json CHANGED
@@ -22,18 +22,19 @@
22
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
23
  "workgroupOk": "tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
24
  "outputDispatchFits": "workgroupOk and ceilDiv(ceilDiv(numel(shapes.y), tunables.WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
 
 
 
 
25
  "outputDtypeOk": "(attrs.output_dtype == 0 and (present.y_zero_point or tensorDtypes.y == \"uint8\")) or attrs.output_dtype == onnxDtypeCode(logicalDtypes.TQ)",
26
  "quantizeDivisionF16": "attrs.precision == onnxDtypeCode(\"float16\") or (attrs.precision == 0 and tensorDtypes.y_scale == \"float16\")",
27
  "sameShapeOk": "ranks.y == ranks.x and numel(shapes.x) == numel(shapes.y) and outputDispatchFits and outputDtypeOk",
28
  "quantDtypesOk": "(tensorDtypes.x != \"float16\" and tensorDtypes.y_scale != \"float16\" and not quantizeDivisionF16) or device.features.has(\"shader-f16\")",
29
  "blockedScaleOk": "attrs.block_size > 0 and ranks.x >= 2 and ranks.y_scale == ranks.x",
30
- "elementCount": "numel(shapes.y)",
31
- "elementCount4": "floor(numel(shapes.y) / 4)",
32
  "scaleSize": "1 if ranks.y_scale == 0 else dim(shapes.y_scale, 0)",
33
  "scaleInner": "1 if ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 else inner(shapes.x, attrs.axis)",
34
- "blockedAxisDim": "dim(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
35
  "blockedScaleAxisDim": "dim(shapes.y_scale, attrs.axis) if ranks.y_scale >= 2 else 1",
36
- "blockedInner": "inner(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
37
  "xVec4": "\"vec4<\" ~ dtypes.TX ~ \">\"",
38
  "scaleVec4": "\"vec4<\" ~ dtypes.TS ~ \">\"",
39
  "yScalar": "dtypes.TQ",
@@ -41,51 +42,45 @@
41
  "yUnsigned": "tensorDtypes.y == \"uint8\"",
42
  "qMin": "0 if tensorDtypes.y == \"uint8\" else 0 - 128",
43
  "qMax": "255 if tensorDtypes.y == \"uint8\" else 127",
44
- "divisionF16": "quantizeDivisionF16"
 
 
 
 
 
 
 
 
 
45
  },
46
  "when": ["sameShapeOk"],
47
  "bindings": {
48
- "x": { "buffer": "read-only-storage", "elementType": "$xVec4" },
49
- "y_scale": { "buffer": "read-only-storage", "elementType": "$scaleVec4" },
50
- "y_zero_point": { "buffer": "read-only-storage", "elementType": "$yVec4" },
51
- "y": { "buffer": "storage", "elementType": "$yVec4" },
52
  "params": {
53
- "buffer": "uniform",
54
  "struct": [
55
  { "name": "count4", "type": "u32", "value": "elementCount4" },
56
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" }
57
  ]
58
  },
59
- "params_2": {
60
  "name": "params",
61
- "buffer": "uniform",
62
  "struct": [
63
  { "name": "count4", "type": "u32", "value": "elementCount4" },
64
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
65
  { "name": "inner", "type": "u32", "value": "scaleInner" }
66
  ]
67
  },
68
- "params_3": {
69
- "name": "params",
70
- "buffer": "uniform",
71
- "struct": [{ "name": "count4", "type": "u32", "value": "elementCount4" }]
72
- },
73
- "params_4": {
 
74
  "name": "params",
75
- "buffer": "uniform",
76
- "struct": [{ "name": "count", "type": "u32", "value": "elementCount" }]
77
- },
78
- "x_3": { "name": "x", "buffer": "read-only-storage", "elementType": "vec4<f32>" },
79
- "y_scale_3": { "name": "y_scale", "buffer": "read-only-storage", "elementType": "f32" },
80
- "params_5": {
81
- "name": "params",
82
- "buffer": "uniform",
83
- "struct": [{ "name": "count", "type": "u32", "value": "elementCount4" }]
84
- },
85
- "y_scale_4": { "name": "y_scale", "buffer": "read-only-storage", "elementType": "vec4<f32>" },
86
- "params_6": {
87
- "name": "params",
88
- "buffer": "uniform",
89
  "struct": [
90
  { "name": "count", "type": "u32", "value": "elementCount4" },
91
  { "name": "axisDim", "type": "u32", "value": "blockedAxisDim" },
@@ -94,9 +89,8 @@
94
  { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
95
  ]
96
  },
97
- "params_7": {
98
  "name": "params",
99
- "buffer": "uniform",
100
  "struct": [
101
  { "name": "count", "type": "u32", "value": "elementCount" },
102
  { "name": "axisDim", "type": "u32", "value": "blockedAxisDim" },
@@ -105,9 +99,8 @@
105
  { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
106
  ]
107
  },
108
- "params_8": {
109
  "name": "params",
110
- "buffer": "uniform",
111
  "struct": [
112
  { "name": "count", "type": "u32", "value": "elementCount" },
113
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
@@ -119,8 +112,7 @@
119
  {
120
  "id": "innermost_axis_vec4_with_zero_point",
121
  "priority": 20,
122
- "when": ["present.y_zero_point", "ranks.x >= 1", "attrs.axis == -1 or attrs.axis == ranks.x - 1", "ranks.y_scale == 1", "dim(shapes.y_scale, 0) == dim(shapes.x, ranks.x - 1)", "dim(shapes.y_scale, 0) % 4 == 0", "ranks.y_zero_point == 1", "dim(shapes.y_zero_point, 0) == dim(shapes.y_scale, 0)", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "quantDtypesOk"],
123
- "derive": { "hasZero": true },
124
  "passes": [
125
  {
126
  "id": "main",
@@ -139,8 +131,7 @@
139
  {
140
  "id": "innermost_axis_vec4_no_zero_point",
141
  "priority": 20,
142
- "when": ["not present.y_zero_point", "ranks.x >= 1", "attrs.axis == -1 or attrs.axis == ranks.x - 1", "ranks.y_scale == 1", "dim(shapes.y_scale, 0) == dim(shapes.x, ranks.x - 1)", "dim(shapes.y_scale, 0) % 4 == 0", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "quantDtypesOk"],
143
- "derive": { "hasZero": false },
144
  "passes": [
145
  {
146
  "id": "main",
@@ -159,15 +150,14 @@
159
  {
160
  "id": "vec4_cross_axis_with_zero_point",
161
  "priority": 16,
162
- "when": ["present.y_zero_point", "attrs.block_size == 0", "ranks.x >= 1", "attrs.axis >= 0", "attrs.axis < ranks.x", "ranks.y_scale == 1", "dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis)", "inner(shapes.x, attrs.axis) % 4 != 0", "ranks.y_zero_point == 1", "dim(shapes.y_zero_point, 0) == dim(shapes.y_scale, 0)", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "quantDtypesOk"],
163
- "derive": { "hasZero": true },
164
  "passes": [
165
  {
166
  "id": "main",
167
  "name": "QuantizeLinear.Vec4CrossAxis",
168
  "shader": "quant-linear-vec4.wgsl.jinja",
169
  "derive": { "op": "\"quantize\"", "perAxis": true, "crossingParams": true },
170
- "bindings": ["x", { "arg": "y_scale" }, { "arg": "y_zero_point" }, "y", "params_2"],
171
  "dispatch": {
172
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
173
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -179,15 +169,14 @@
179
  {
180
  "id": "vec4_with_zero_point",
181
  "priority": 15,
182
- "when": ["present.y_zero_point", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "(ranks.y_scale == 0 or ranks.y_scale == 1)", "(ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))", "(ranks.y_zero_point == 0 or ranks.y_zero_point == 1)", "(ranks.y_zero_point == 0 or dim(shapes.y_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_zero_point, 0) == dim(shapes.x, attrs.axis)))", "quantDtypesOk"],
183
- "derive": { "hasZero": true },
184
  "passes": [
185
  {
186
  "id": "main",
187
  "name": "QuantizeLinear.Vec4",
188
  "shader": "quant-linear-vec4.wgsl.jinja",
189
  "derive": { "op": "\"quantize\"", "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)" },
190
- "bindings": ["x", { "arg": "y_scale" }, { "arg": "y_zero_point" }, "y", "params_2"],
191
  "dispatch": {
192
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
193
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -199,15 +188,14 @@
199
  {
200
  "id": "vec4_no_zero_point",
201
  "priority": 15,
202
- "when": ["not present.y_zero_point", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "(ranks.y_scale == 0 or ranks.y_scale == 1)", "(ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))", "quantDtypesOk"],
203
- "derive": { "hasZero": false },
204
  "passes": [
205
  {
206
  "id": "main",
207
  "name": "QuantizeLinear.Vec4",
208
  "shader": "quant-linear-vec4.wgsl.jinja",
209
  "derive": { "op": "\"quantize\"", "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)" },
210
- "bindings": ["x", { "arg": "y_scale" }, "y", "params_2"],
211
  "dispatch": {
212
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
213
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -219,15 +207,14 @@
219
  {
220
  "id": "vec4_tail_with_zero_point",
221
  "priority": 12,
222
- "when": ["present.y_zero_point", "numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS", "numel(shapes.y) % 4 != 0", "(ranks.y_scale == 0 or (ranks.y_scale == 1 and dim(shapes.y_scale, 0) == 1))", "(ranks.y_zero_point == 0 or (ranks.y_zero_point == 1 and dim(shapes.y_zero_point, 0) == 1))", "quantDtypesOk"],
223
- "derive": { "hasZero": true },
224
  "passes": [
225
  {
226
  "id": "bulk",
227
  "name": "QuantizeLinear.Vec4Bulk",
228
  "shader": "quant-linear-vec4.wgsl.jinja",
229
  "derive": { "op": "\"quantize\"", "perAxis": false },
230
- "bindings": ["x", { "arg": "y_scale" }, { "arg": "y_zero_point" }, "y", "params_3"],
231
  "dispatch": {
232
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
233
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -239,7 +226,7 @@
239
  "name": "QuantizeLinear.ScalarTail",
240
  "shader": "quant-linear-scalar.wgsl.jinja",
241
  "derive": { "op": "\"quantize\"", "x4": true, "perAxis": false },
242
- "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y_zero_point" }, { "arg": "y" }, "params_4"],
243
  "dispatch": {
244
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
245
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -251,15 +238,14 @@
251
  {
252
  "id": "vec4_tail_no_zero_point",
253
  "priority": 12,
254
- "when": ["not present.y_zero_point", "numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS", "numel(shapes.y) % 4 != 0", "(ranks.y_scale == 0 or (ranks.y_scale == 1 and dim(shapes.y_scale, 0) == 1))", "quantDtypesOk"],
255
- "derive": { "hasZero": false },
256
  "passes": [
257
  {
258
  "id": "bulk",
259
  "name": "QuantizeLinear.Vec4Bulk",
260
  "shader": "quant-linear-vec4.wgsl.jinja",
261
  "derive": { "op": "\"quantize\"", "perAxis": false },
262
- "bindings": ["x", { "arg": "y_scale" }, "y", "params_3"],
263
  "dispatch": {
264
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
265
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -271,7 +257,7 @@
271
  "name": "QuantizeLinear.ScalarTail",
272
  "shader": "quant-linear-scalar.wgsl.jinja",
273
  "derive": { "op": "\"quantize\"", "x4": true, "perAxis": false },
274
- "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y" }, "params_4"],
275
  "dispatch": {
276
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
277
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -280,18 +266,60 @@
280
  }
281
  ]
282
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
283
  {
284
  "id": "blocked_last_axis_vec4_with_zero_point",
285
  "priority": 24,
286
  "when": ["present.y_zero_point", "blockedScaleOk", "attrs.block_size % 4 == 0", "ranks.y_zero_point == ranks.x", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "inner(shapes.x, attrs.axis) == 1", "dim(shapes.x, attrs.axis) % attrs.block_size == 0", "dim(shapes.y_scale, attrs.axis) == dim(shapes.x, attrs.axis) / attrs.block_size", "numel(shapes.y_zero_point) == numel(shapes.y_scale)", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
287
- "derive": { "hasZero": true, "blockVectors": "attrs.block_size / 4" },
288
  "passes": [
289
  {
290
  "id": "main",
291
  "name": "QuantizeLinear.BlockedLastAxisVec4",
292
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
293
  "derive": { "lastAxisVectorized": true },
294
- "bindings": ["x_3", "y_scale_3", { "arg": "y_zero_point" }, "y", "params_5"],
295
  "dispatch": {
296
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
297
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -304,14 +332,14 @@
304
  "id": "blocked_last_axis_vec4_no_zero_point",
305
  "priority": 23,
306
  "when": ["not present.y_zero_point", "blockedScaleOk", "attrs.block_size % 4 == 0", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "inner(shapes.x, attrs.axis) == 1", "dim(shapes.x, attrs.axis) % attrs.block_size == 0", "dim(shapes.y_scale, attrs.axis) == dim(shapes.x, attrs.axis) / attrs.block_size", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
307
- "derive": { "hasZero": false, "blockVectors": "attrs.block_size / 4" },
308
  "passes": [
309
  {
310
  "id": "main",
311
  "name": "QuantizeLinear.BlockedLastAxisVec4",
312
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
313
  "derive": { "lastAxisVectorized": true },
314
- "bindings": ["x_3", "y_scale_3", "y", "params_5"],
315
  "dispatch": {
316
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
317
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -324,14 +352,13 @@
324
  "id": "blocked_vec4_with_zero_point",
325
  "priority": 22,
326
  "when": ["present.y_zero_point", "blockedScaleOk", "ranks.y_zero_point == ranks.x", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "numel(shapes.y_scale) % 4 == 0", "numel(shapes.y_zero_point) == numel(shapes.y_scale)", "inner(shapes.x, attrs.axis) % 4 == 0", "dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)", "dim(shapes.y_zero_point, attrs.axis) == dim(shapes.y_scale, attrs.axis)", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
327
- "derive": { "hasZero": true },
328
  "passes": [
329
  {
330
  "id": "main",
331
  "name": "QuantizeLinear.BlockedVec4WithZeroPoint",
332
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
333
  "derive": { "vectorizedSpec": true },
334
- "bindings": ["x_3", "y_scale_4", "y_zero_point", "y", "params_6"],
335
  "dispatch": {
336
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
337
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -344,14 +371,13 @@
344
  "id": "blocked_vec4_no_zero_point",
345
  "priority": 21,
346
  "when": ["not present.y_zero_point", "blockedScaleOk", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "numel(shapes.y_scale) % 4 == 0", "inner(shapes.x, attrs.axis) % 4 == 0", "dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
347
- "derive": { "hasZero": false },
348
  "passes": [
349
  {
350
  "id": "main",
351
  "name": "QuantizeLinear.BlockedVec4",
352
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
353
  "derive": { "vectorizedSpec": true },
354
- "bindings": ["x_3", "y_scale_4", "y", "params_6"],
355
  "dispatch": {
356
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
357
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -360,49 +386,10 @@
360
  }
361
  ]
362
  },
363
- {
364
- "id": "blocked_with_zero_point",
365
- "priority": 19,
366
- "when": ["present.y_zero_point", "blockedScaleOk", "ranks.y_zero_point == ranks.x", "dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)", "dim(shapes.y_zero_point, attrs.axis) == dim(shapes.y_scale, attrs.axis)", "quantDtypesOk"],
367
- "derive": { "hasZero": true },
368
- "passes": [
369
- {
370
- "id": "main",
371
- "name": "QuantizeLinear.Blocked",
372
- "shader": "quant-linear-blocked-axis.wgsl.jinja",
373
- "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y_zero_point" }, { "arg": "y" }, "params_7"],
374
- "dispatch": {
375
- "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
376
- "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
377
- "z": 1
378
- }
379
- }
380
- ]
381
- },
382
- {
383
- "id": "blocked_no_zero_point",
384
- "priority": 19,
385
- "when": ["not present.y_zero_point", "blockedScaleOk", "dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)", "quantDtypesOk"],
386
- "derive": { "hasZero": false },
387
- "passes": [
388
- {
389
- "id": "main",
390
- "name": "QuantizeLinear.BlockedNoZero",
391
- "shader": "quant-linear-blocked-axis.wgsl.jinja",
392
- "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y" }, "params_7"],
393
- "dispatch": {
394
- "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
395
- "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
396
- "z": 1
397
- }
398
- }
399
- ]
400
- },
401
  {
402
  "id": "with_zero_point",
403
  "priority": 10,
404
  "when": ["present.y_zero_point", "(ranks.y_scale == 0 or ranks.y_scale == 1)", "(ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis)))", "(ranks.y_zero_point == 0 or ranks.y_zero_point == 1)", "(ranks.y_zero_point == 0 or dim(shapes.y_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_zero_point, 0) == dim(shapes.x, attrs.axis)))", "quantDtypesOk"],
405
- "derive": { "hasZero": true },
406
  "passes": [
407
  {
408
  "id": "main",
@@ -413,7 +400,13 @@
413
  "x4": "false",
414
  "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)"
415
  },
416
- "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y_zero_point" }, { "arg": "y" }, "params_8"],
 
 
 
 
 
 
417
  "dispatch": {
418
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
419
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -425,7 +418,6 @@
425
  {
426
  "id": "no_zero_point",
427
  "when": ["not present.y_zero_point", "(ranks.y_scale == 0 or ranks.y_scale == 1)", "(ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis)))", "quantDtypesOk"],
428
- "derive": { "hasZero": false },
429
  "passes": [
430
  {
431
  "id": "main",
@@ -436,7 +428,7 @@
436
  "x4": "false",
437
  "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)"
438
  },
439
- "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y" }, "params_8"],
440
  "dispatch": {
441
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
442
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
 
22
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
23
  "workgroupOk": "tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
24
  "outputDispatchFits": "workgroupOk and ceilDiv(ceilDiv(numel(shapes.y), tunables.WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
25
+ "elementCount": "numel(shapes.y)",
26
+ "elementCount4": "floor(numel(shapes.y) / 4)",
27
+ "blockedAxisDim": "dim(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
28
+ "blockedInner": "inner(shapes.x, attrs.axis) if ranks.x >= 2 else 1",
29
  "outputDtypeOk": "(attrs.output_dtype == 0 and (present.y_zero_point or tensorDtypes.y == \"uint8\")) or attrs.output_dtype == onnxDtypeCode(logicalDtypes.TQ)",
30
  "quantizeDivisionF16": "attrs.precision == onnxDtypeCode(\"float16\") or (attrs.precision == 0 and tensorDtypes.y_scale == \"float16\")",
31
  "sameShapeOk": "ranks.y == ranks.x and numel(shapes.x) == numel(shapes.y) and outputDispatchFits and outputDtypeOk",
32
  "quantDtypesOk": "(tensorDtypes.x != \"float16\" and tensorDtypes.y_scale != \"float16\" and not quantizeDivisionF16) or device.features.has(\"shader-f16\")",
33
  "blockedScaleOk": "attrs.block_size > 0 and ranks.x >= 2 and ranks.y_scale == ranks.x",
34
+ "hasZero": "present.y_zero_point",
 
35
  "scaleSize": "1 if ranks.y_scale == 0 else dim(shapes.y_scale, 0)",
36
  "scaleInner": "1 if ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 else inner(shapes.x, attrs.axis)",
 
37
  "blockedScaleAxisDim": "dim(shapes.y_scale, attrs.axis) if ranks.y_scale >= 2 else 1",
 
38
  "xVec4": "\"vec4<\" ~ dtypes.TX ~ \">\"",
39
  "scaleVec4": "\"vec4<\" ~ dtypes.TS ~ \">\"",
40
  "yScalar": "dtypes.TQ",
 
42
  "yUnsigned": "tensorDtypes.y == \"uint8\"",
43
  "qMin": "0 if tensorDtypes.y == \"uint8\" else 0 - 128",
44
  "qMax": "255 if tensorDtypes.y == \"uint8\" else 127",
45
+ "divisionF16": "quantizeDivisionF16",
46
+ "routeInnermostAxisVec4WithZeroPoint": "(present.y_zero_point) and (ranks.x >= 1) and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and (ranks.y_scale == 1) and (dim(shapes.y_scale, 0) == dim(shapes.x, ranks.x - 1)) and (dim(shapes.y_scale, 0) % 4 == 0) and (ranks.y_zero_point == 1) and (dim(shapes.y_zero_point, 0) == dim(shapes.y_scale, 0)) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and (quantDtypesOk)",
47
+ "routeInnermostAxisVec4NoZeroPoint": "(not present.y_zero_point) and (ranks.x >= 1) and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and (ranks.y_scale == 1) and (dim(shapes.y_scale, 0) == dim(shapes.x, ranks.x - 1)) and (dim(shapes.y_scale, 0) % 4 == 0) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and (quantDtypesOk)",
48
+ "routeVec4CrossAxisWithZeroPoint": "(present.y_zero_point) and (attrs.block_size == 0) and (ranks.x >= 1) and (attrs.axis >= 0) and (attrs.axis < ranks.x) and (ranks.y_scale == 1) and (dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis)) and (inner(shapes.x, attrs.axis) % 4 != 0) and (ranks.y_zero_point == 1) and (dim(shapes.y_zero_point, 0) == dim(shapes.y_scale, 0)) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and (quantDtypesOk)",
49
+ "routeVec4WithZeroPoint": "(present.y_zero_point) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and ((ranks.y_scale == 0 or ranks.y_scale == 1)) and ((ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))) and ((ranks.y_zero_point == 0 or ranks.y_zero_point == 1)) and ((ranks.y_zero_point == 0 or dim(shapes.y_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_zero_point, 0) == dim(shapes.x, attrs.axis)))) and (quantDtypesOk)",
50
+ "routeVec4NoZeroPoint": "(not present.y_zero_point) and (numel(shapes.y) > 0) and (numel(shapes.y) % 4 == 0) and ((ranks.y_scale == 0 or ranks.y_scale == 1)) and ((ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis) and inner(shapes.x, attrs.axis) % 4 == 0))) and (quantDtypesOk)",
51
+ "routeVec4TailWithZeroPoint": "(present.y_zero_point) and (numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS) and (numel(shapes.y) % 4 != 0) and ((ranks.y_scale == 0 or (ranks.y_scale == 1 and dim(shapes.y_scale, 0) == 1))) and ((ranks.y_zero_point == 0 or (ranks.y_zero_point == 1 and dim(shapes.y_zero_point, 0) == 1))) and (quantDtypesOk)",
52
+ "routeVec4TailNoZeroPoint": "(not present.y_zero_point) and (numel(shapes.y) >= tunables.VEC4_TAIL_MIN_ELEMENTS) and (numel(shapes.y) % 4 != 0) and ((ranks.y_scale == 0 or (ranks.y_scale == 1 and dim(shapes.y_scale, 0) == 1))) and (quantDtypesOk)",
53
+ "routeBlockedWithZeroPoint": "(present.y_zero_point) and (blockedScaleOk) and (ranks.y_zero_point == ranks.x) and (dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)) and (dim(shapes.y_zero_point, attrs.axis) == dim(shapes.y_scale, attrs.axis)) and (quantDtypesOk)",
54
+ "routeBlockedNoZeroPoint": "(not present.y_zero_point) and (blockedScaleOk) and (dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)) and (quantDtypesOk)"
55
  },
56
  "when": ["sameShapeOk"],
57
  "bindings": {
58
+ "x": { "elementType": "$xVec4" },
59
+ "y_scale": { "elementType": "$scaleVec4" },
60
+ "y_zero_point": { "elementType": "$yVec4" },
61
+ "y": { "elementType": "$yVec4" },
62
  "params": {
 
63
  "struct": [
64
  { "name": "count4", "type": "u32", "value": "elementCount4" },
65
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" }
66
  ]
67
  },
68
+ "params_main": {
69
  "name": "params",
 
70
  "struct": [
71
  { "name": "count4", "type": "u32", "value": "elementCount4" },
72
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
73
  { "name": "inner", "type": "u32", "value": "scaleInner" }
74
  ]
75
  },
76
+ "params_bulk": { "name": "params", "struct": [{ "name": "count4", "type": "u32", "value": "elementCount4" }] },
77
+ "params_tail": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "elementCount" }] },
78
+ "x_main": { "name": "x", "elementType": "vec4<f32>" },
79
+ "y_scale_f32": { "name": "y_scale", "elementType": "f32" },
80
+ "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "elementCount4" }] },
81
+ "y_scale_main": { "name": "y_scale", "elementType": "vec4<f32>" },
82
+ "params__uniform": {
83
  "name": "params",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  "struct": [
85
  { "name": "count", "type": "u32", "value": "elementCount4" },
86
  { "name": "axisDim", "type": "u32", "value": "blockedAxisDim" },
 
89
  { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
90
  ]
91
  },
92
+ "params_blocked_no_zero_point": {
93
  "name": "params",
 
94
  "struct": [
95
  { "name": "count", "type": "u32", "value": "elementCount" },
96
  { "name": "axisDim", "type": "u32", "value": "blockedAxisDim" },
 
99
  { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
100
  ]
101
  },
102
+ "params_no_zero_point": {
103
  "name": "params",
 
104
  "struct": [
105
  { "name": "count", "type": "u32", "value": "elementCount" },
106
  { "name": "scaleSize", "type": "u32", "value": "scaleSize" },
 
112
  {
113
  "id": "innermost_axis_vec4_with_zero_point",
114
  "priority": 20,
115
+ "when": ["routeInnermostAxisVec4WithZeroPoint"],
 
116
  "passes": [
117
  {
118
  "id": "main",
 
131
  {
132
  "id": "innermost_axis_vec4_no_zero_point",
133
  "priority": 20,
134
+ "when": ["routeInnermostAxisVec4NoZeroPoint"],
 
135
  "passes": [
136
  {
137
  "id": "main",
 
150
  {
151
  "id": "vec4_cross_axis_with_zero_point",
152
  "priority": 16,
153
+ "when": ["routeVec4CrossAxisWithZeroPoint"],
 
154
  "passes": [
155
  {
156
  "id": "main",
157
  "name": "QuantizeLinear.Vec4CrossAxis",
158
  "shader": "quant-linear-vec4.wgsl.jinja",
159
  "derive": { "op": "\"quantize\"", "perAxis": true, "crossingParams": true },
160
+ "bindings": ["x", { "arg": "y_scale" }, { "arg": "y_zero_point" }, "y", "params_main"],
161
  "dispatch": {
162
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
163
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
169
  {
170
  "id": "vec4_with_zero_point",
171
  "priority": 15,
172
+ "when": ["routeVec4WithZeroPoint"],
 
173
  "passes": [
174
  {
175
  "id": "main",
176
  "name": "QuantizeLinear.Vec4",
177
  "shader": "quant-linear-vec4.wgsl.jinja",
178
  "derive": { "op": "\"quantize\"", "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)" },
179
+ "bindings": ["x", { "arg": "y_scale" }, { "arg": "y_zero_point" }, "y", "params_main"],
180
  "dispatch": {
181
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
182
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
188
  {
189
  "id": "vec4_no_zero_point",
190
  "priority": 15,
191
+ "when": ["routeVec4NoZeroPoint"],
 
192
  "passes": [
193
  {
194
  "id": "main",
195
  "name": "QuantizeLinear.Vec4",
196
  "shader": "quant-linear-vec4.wgsl.jinja",
197
  "derive": { "op": "\"quantize\"", "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)" },
198
+ "bindings": ["x", { "arg": "y_scale" }, "y", "params_main"],
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
  {
208
  "id": "vec4_tail_with_zero_point",
209
  "priority": 12,
210
+ "when": ["routeVec4TailWithZeroPoint"],
 
211
  "passes": [
212
  {
213
  "id": "bulk",
214
  "name": "QuantizeLinear.Vec4Bulk",
215
  "shader": "quant-linear-vec4.wgsl.jinja",
216
  "derive": { "op": "\"quantize\"", "perAxis": false },
217
+ "bindings": ["x", { "arg": "y_scale" }, { "arg": "y_zero_point" }, "y", "params_bulk"],
218
  "dispatch": {
219
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
220
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
226
  "name": "QuantizeLinear.ScalarTail",
227
  "shader": "quant-linear-scalar.wgsl.jinja",
228
  "derive": { "op": "\"quantize\"", "x4": true, "perAxis": false },
229
+ "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y_zero_point" }, { "arg": "y" }, "params_tail"],
230
  "dispatch": {
231
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
232
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
 
238
  {
239
  "id": "vec4_tail_no_zero_point",
240
  "priority": 12,
241
+ "when": ["routeVec4TailNoZeroPoint"],
 
242
  "passes": [
243
  {
244
  "id": "bulk",
245
  "name": "QuantizeLinear.Vec4Bulk",
246
  "shader": "quant-linear-vec4.wgsl.jinja",
247
  "derive": { "op": "\"quantize\"", "perAxis": false },
248
+ "bindings": ["x", { "arg": "y_scale" }, "y", "params_bulk"],
249
  "dispatch": {
250
  "x": "min(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
251
  "y": "ceilDiv(ceilDiv((floor(numel(shapes.y) / 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
257
  "name": "QuantizeLinear.ScalarTail",
258
  "shader": "quant-linear-scalar.wgsl.jinja",
259
  "derive": { "op": "\"quantize\"", "x4": true, "perAxis": false },
260
+ "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y" }, "params_tail"],
261
  "dispatch": {
262
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
263
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
 
266
  }
267
  ]
268
  },
269
+ {
270
+ "id": "blocked_with_zero_point",
271
+ "priority": 19,
272
+ "when": ["routeBlockedWithZeroPoint"],
273
+ "passes": [
274
+ {
275
+ "id": "main",
276
+ "name": "QuantizeLinear.Blocked",
277
+ "shader": "quant-linear-blocked-axis.wgsl.jinja",
278
+ "bindings": [
279
+ { "arg": "x" },
280
+ { "arg": "y_scale" },
281
+ { "arg": "y_zero_point" },
282
+ { "arg": "y" },
283
+ "params_blocked_no_zero_point"
284
+ ],
285
+ "dispatch": {
286
+ "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
287
+ "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
288
+ "z": 1
289
+ }
290
+ }
291
+ ]
292
+ },
293
+ {
294
+ "id": "blocked_no_zero_point",
295
+ "priority": 19,
296
+ "when": ["routeBlockedNoZeroPoint"],
297
+ "passes": [
298
+ {
299
+ "id": "main",
300
+ "name": "QuantizeLinear.BlockedNoZero",
301
+ "shader": "quant-linear-blocked-axis.wgsl.jinja",
302
+ "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y" }, "params_blocked_no_zero_point"],
303
+ "dispatch": {
304
+ "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
305
+ "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
306
+ "z": 1
307
+ }
308
+ }
309
+ ]
310
+ },
311
  {
312
  "id": "blocked_last_axis_vec4_with_zero_point",
313
  "priority": 24,
314
  "when": ["present.y_zero_point", "blockedScaleOk", "attrs.block_size % 4 == 0", "ranks.y_zero_point == ranks.x", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "inner(shapes.x, attrs.axis) == 1", "dim(shapes.x, attrs.axis) % attrs.block_size == 0", "dim(shapes.y_scale, attrs.axis) == dim(shapes.x, attrs.axis) / attrs.block_size", "numel(shapes.y_zero_point) == numel(shapes.y_scale)", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
315
+ "derive": { "blockVectors": "attrs.block_size / 4" },
316
  "passes": [
317
  {
318
  "id": "main",
319
  "name": "QuantizeLinear.BlockedLastAxisVec4",
320
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
321
  "derive": { "lastAxisVectorized": true },
322
+ "bindings": ["x_main", "y_scale_f32", { "arg": "y_zero_point" }, "y", "params_count"],
323
  "dispatch": {
324
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
325
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
332
  "id": "blocked_last_axis_vec4_no_zero_point",
333
  "priority": 23,
334
  "when": ["not present.y_zero_point", "blockedScaleOk", "attrs.block_size % 4 == 0", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "inner(shapes.x, attrs.axis) == 1", "dim(shapes.x, attrs.axis) % attrs.block_size == 0", "dim(shapes.y_scale, attrs.axis) == dim(shapes.x, attrs.axis) / attrs.block_size", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
335
+ "derive": { "blockVectors": "attrs.block_size / 4" },
336
  "passes": [
337
  {
338
  "id": "main",
339
  "name": "QuantizeLinear.BlockedLastAxisVec4",
340
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
341
  "derive": { "lastAxisVectorized": true },
342
+ "bindings": ["x_main", "y_scale_f32", "y", "params_count"],
343
  "dispatch": {
344
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
345
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
352
  "id": "blocked_vec4_with_zero_point",
353
  "priority": 22,
354
  "when": ["present.y_zero_point", "blockedScaleOk", "ranks.y_zero_point == ranks.x", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "numel(shapes.y_scale) % 4 == 0", "numel(shapes.y_zero_point) == numel(shapes.y_scale)", "inner(shapes.x, attrs.axis) % 4 == 0", "dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)", "dim(shapes.y_zero_point, attrs.axis) == dim(shapes.y_scale, attrs.axis)", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
 
355
  "passes": [
356
  {
357
  "id": "main",
358
  "name": "QuantizeLinear.BlockedVec4WithZeroPoint",
359
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
360
  "derive": { "vectorizedSpec": true },
361
+ "bindings": ["x_main", "y_scale_main", "y_zero_point", "y", "params__uniform"],
362
  "dispatch": {
363
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
364
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
371
  "id": "blocked_vec4_no_zero_point",
372
  "priority": 21,
373
  "when": ["not present.y_zero_point", "blockedScaleOk", "numel(shapes.y) > 0", "numel(shapes.y) % 4 == 0", "numel(shapes.y_scale) % 4 == 0", "inner(shapes.x, attrs.axis) % 4 == 0", "dim(shapes.y_scale, attrs.axis) == ceil(dim(shapes.x, attrs.axis) / attrs.block_size)", "tensorDtypes.x == \"float32\"", "tensorDtypes.y_scale == \"float32\""],
 
374
  "passes": [
375
  {
376
  "id": "main",
377
  "name": "QuantizeLinear.BlockedVec4",
378
  "shader": "quant-linear-blocked-axis.wgsl.jinja",
379
  "derive": { "vectorizedSpec": true },
380
+ "bindings": ["x_main", "y_scale_main", "y", "params__uniform"],
381
  "dispatch": {
382
  "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
383
  "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
 
386
  }
387
  ]
388
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
389
  {
390
  "id": "with_zero_point",
391
  "priority": 10,
392
  "when": ["present.y_zero_point", "(ranks.y_scale == 0 or ranks.y_scale == 1)", "(ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis)))", "(ranks.y_zero_point == 0 or ranks.y_zero_point == 1)", "(ranks.y_zero_point == 0 or dim(shapes.y_zero_point, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_zero_point, 0) == dim(shapes.x, attrs.axis)))", "quantDtypesOk"],
 
393
  "passes": [
394
  {
395
  "id": "main",
 
400
  "x4": "false",
401
  "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)"
402
  },
403
+ "bindings": [
404
+ { "arg": "x" },
405
+ { "arg": "y_scale" },
406
+ { "arg": "y_zero_point" },
407
+ { "arg": "y" },
408
+ "params_no_zero_point"
409
+ ],
410
  "dispatch": {
411
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
412
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
 
418
  {
419
  "id": "no_zero_point",
420
  "when": ["not present.y_zero_point", "(ranks.y_scale == 0 or ranks.y_scale == 1)", "(ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1 or (ranks.x >= 1 and dim(shapes.y_scale, 0) == dim(shapes.x, attrs.axis)))", "quantDtypesOk"],
 
421
  "passes": [
422
  {
423
  "id": "main",
 
428
  "x4": "false",
429
  "perAxis": "not (ranks.y_scale == 0 or dim(shapes.y_scale, 0) == 1)"
430
  },
431
+ "bindings": [{ "arg": "x" }, { "arg": "y_scale" }, { "arg": "y" }, "params_no_zero_point"],
432
  "dispatch": {
433
  "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
434
  "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
build/webgpu/metadata.json CHANGED
@@ -1,23 +1,23 @@
1
  {
2
  "name": "ai.onnx.QuantizeLinear",
3
- "id": "_ai_onnx_quantizelinear_webgpu_bf60a95",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "B6CqLHGCbyXXUAAcatPURKRkCnRMOZCNu7ix44MZbZw=",
11
- "manifest.json": "X5UX5Db4+s1nUDXg5jez+yaYuFnt+y5vPqCc1mWXy3I=",
12
- "quant-linear-blocked-axis.wgsl.jinja": "fhTtN3vZ8jl4fTshwThxKYds9QrAZ9U2GzjPtWQs2OE=",
13
- "quant-linear-scalar.wgsl.jinja": "/OPcCfkXBiT0dEm1wEnImi5EBY8l17n6jpS5KKCjvQ4=",
14
- "quant-linear-vec4.wgsl.jinja": "qQrnnz6UMx0Pg/d7Cs2cjh1HBaurycYK0H24VIbq+Fs=",
15
- "test.json": "cniKZLfeclIx3sXUNY1qFbwal5ui2JqFKfODpgZEoCA="
16
  }
17
  },
18
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
19
  "webgpu": {
20
- "manifestSpec": "2.0",
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"],
@@ -26,12 +26,12 @@
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_last_axis_vec4_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
30
  "blocked_last_axis_vec4_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
31
  "blocked_vec4_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
32
  "blocked_vec4_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
33
- "blocked_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
34
- "blocked_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
35
  "with_zero_point": ["quant-linear-scalar.wgsl.jinja"],
36
  "no_zero_point": ["quant-linear-scalar.wgsl.jinja"]
37
  }
 
1
  {
2
  "name": "ai.onnx.QuantizeLinear",
3
+ "id": "_ai_onnx_quantizelinear_webgpu_70dd9f2",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "utMq67YYL5So5bM+YrNks8SuNNuaTZIbsrZTLwNFL3o=",
11
+ "manifest.json": "RQxE5Ec0ba1isqgrsZShPPhWw38GJpPpdT7kxPyFmn4=",
12
+ "quant-linear-blocked-axis.wgsl.jinja": "lkqnjbH6rAsEuQj5w62L/NulyNvjapRLOiAnQOgY4/Q=",
13
+ "quant-linear-scalar.wgsl.jinja": "oNmEXVPTyQ6UabRKITYEvFtVY5pqtDdxWzwTT974ryw=",
14
+ "quant-linear-vec4.wgsl.jinja": "m2aQCFDc38lScqYnEBk+DSurlIVvd8k3mlH0VVN2Wn8=",
15
+ "test.json": "pHs/HO6wZKy0olozXbnHwhQtkSN4+KnI4i1BAEx6Dks="
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"],
 
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
  "blocked_last_axis_vec4_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
32
  "blocked_last_axis_vec4_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
33
  "blocked_vec4_with_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
34
  "blocked_vec4_no_zero_point": ["quant-linear-blocked-axis.wgsl.jinja"],
 
 
35
  "with_zero_point": ["quant-linear-scalar.wgsl.jinja"],
36
  "no_zero_point": ["quant-linear-scalar.wgsl.jinja"]
37
  }
build/webgpu/quant-linear-blocked-axis.wgsl.jinja CHANGED
@@ -2,34 +2,29 @@
2
  {% set operation = op if op is defined else "quantize" %}
3
  {% set vectorized = vectorizedSpec if vectorizedSpec is defined else false %}
4
  {% set lastAxisVectorized = lastAxisVectorized if lastAxisVectorized is defined else false %}
5
- {% set xUnsigned = xUnsigned is defined and xUnsigned %}
6
- {% set yUnsigned = yUnsigned is defined and yUnsigned %}
7
- {% set qMin = qMin | default(0 if yUnsigned else 0 - 128) %}
8
- {% set qMax = qMax | default(255 if yUnsigned else 127) %}
9
  {% set blockVectors = blockVectors | default(1) %}
 
 
 
 
 
 
10
  {{ env.wgsl.resourceDeclarations }}
11
-
12
  {% if operation == "dequantize" %}
 
13
  fn read_x(index: u32) -> i32 {
14
- {% if xUnsigned %}
15
- return i32(x[index]);
16
- {% else %}
17
  return x[index];
18
- {% endif %}
19
  }
20
 
21
  fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
22
  {% if hasZero %}
23
- {% if xUnsigned %}
24
- return i32(x_zero_point[index]);
25
- {% else %}
26
  return x_zero_point[index];
27
- {% endif %}
28
  {% else %}
29
  return 0;
30
  {% endif %}
31
  }
32
  {% elif not vectorized and not lastAxisVectorized %}
 
33
  fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
34
  {% if hasZero %}
35
  {% if yUnsigned %}
@@ -50,18 +45,17 @@ fn write_quantized(index: u32, value: i32) {
50
  {% endif %}
51
  }
52
  {% endif %}
53
-
54
  {% if not lastAxisVectorized %}
 
55
  fn scale_index(index: u32) -> u32 {
56
  let inner_index = index % params.inner;
57
  let axis_index = (index / params.inner) % params.axisDim;
58
  let outer_index = index / (params.axisDim * params.inner);
59
  return (outer_index * params.scaleAxisDim + axis_index / params.blockSize) * params.inner + inner_index;
60
  }
61
-
62
  {% endif %}
63
  {% if operation == "quantize" %}
64
- {% set divisionF16 = divisionF16 is defined and divisionF16 %}
65
  // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of
66
  // infinities, saturation, and halfway values across every kernel route.
67
  {% if divisionF16 %}
@@ -106,16 +100,14 @@ fn round_scaled_half_to_even(value: f32, scale: f32) -> i32 {
106
  return i32(select(hi, lo, is_even));
107
  }
108
  {% endif %}
109
-
110
-
111
  {% endif %}
 
112
  {% if lastAxisVectorized %}
113
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
114
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
115
  // A vec4 cannot cross a block boundary on this route. Map vector indices
116
  // directly to their shared scale/zero-point entry.
117
- let i4 = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
118
- if (i4 >= params.count) { return; }
119
  let scale_index = i4 / {{ blockVectors }}u;
120
  let values = x[i4];
121
  let scale = f32(y_scale[scale_index]);
@@ -144,8 +136,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
144
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
145
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
146
  // Four adjacent inner elements map to four adjacent scale entries.
147
- let i4 = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
148
- if (i4 >= params.count) { return; }
149
  let scale4 = scale_index(i4 * 4u) / 4u;
150
  let values = x[i4];
151
  let scales = y_scale[scale4];
@@ -171,10 +162,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
171
  {% else %}
172
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
173
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
174
- // The flat dispatch is folded across x/y at a fixed per-axis workgroup
175
- // width; gid.y carries the high portion of the element index.
176
- let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
177
- if (index >= params.count) { return; }
178
  let scaleIndex = scale_index(index);
179
  {% if operation == "quantize" %}
180
  let rounded = round_scaled_half_to_even(f32(x[index]), f32(y_scale[scaleIndex])) + read_zero({% if hasZero %}scaleIndex{% endif %});
 
2
  {% set operation = op if op is defined else "quantize" %}
3
  {% set vectorized = vectorizedSpec if vectorizedSpec is defined else false %}
4
  {% set lastAxisVectorized = lastAxisVectorized if lastAxisVectorized is defined else false %}
 
 
 
 
5
  {% set blockVectors = blockVectors | default(1) %}
6
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
7
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
8
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
9
+ // per-axis workgroup fold width.
10
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
11
+ if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
12
  {{ env.wgsl.resourceDeclarations }}
 
13
  {% if operation == "dequantize" %}
14
+
15
  fn read_x(index: u32) -> i32 {
 
 
 
16
  return x[index];
 
17
  }
18
 
19
  fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
20
  {% if hasZero %}
 
 
 
21
  return x_zero_point[index];
 
22
  {% else %}
23
  return 0;
24
  {% endif %}
25
  }
26
  {% elif not vectorized and not lastAxisVectorized %}
27
+
28
  fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
29
  {% if hasZero %}
30
  {% if yUnsigned %}
 
45
  {% endif %}
46
  }
47
  {% endif %}
 
48
  {% if not lastAxisVectorized %}
49
+
50
  fn scale_index(index: u32) -> u32 {
51
  let inner_index = index % params.inner;
52
  let axis_index = (index / params.inner) % params.axisDim;
53
  let outer_index = index / (params.axisDim * params.inner);
54
  return (outer_index * params.scaleAxisDim + axis_index / params.blockSize) * params.inner + inner_index;
55
  }
 
56
  {% endif %}
57
  {% if operation == "quantize" %}
58
+
59
  // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of
60
  // infinities, saturation, and halfway values across every kernel route.
61
  {% if divisionF16 %}
 
100
  return i32(select(hi, lo, is_even));
101
  }
102
  {% endif %}
 
 
103
  {% endif %}
104
+
105
  {% if lastAxisVectorized %}
106
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
107
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
108
  // A vec4 cannot cross a block boundary on this route. Map vector indices
109
  // directly to their shared scale/zero-point entry.
110
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "i4", guardInline=true) }}
 
111
  let scale_index = i4 / {{ blockVectors }}u;
112
  let values = x[i4];
113
  let scale = f32(y_scale[scale_index]);
 
136
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
137
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
138
  // Four adjacent inner elements map to four adjacent scale entries.
139
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "i4", guardInline=true) }}
 
140
  let scale4 = scale_index(i4 * 4u) / 4u;
141
  let values = x[i4];
142
  let scales = y_scale[scale4];
 
162
  {% else %}
163
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
164
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
165
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "index", guardInline=true) }}
 
 
 
166
  let scaleIndex = scale_index(index);
167
  {% if operation == "quantize" %}
168
  let rounded = round_scaled_half_to_even(f32(x[index]), f32(y_scale[scaleIndex])) + read_zero({% if hasZero %}scaleIndex{% endif %});
build/webgpu/quant-linear-scalar.wgsl.jinja CHANGED
@@ -1,11 +1,12 @@
1
- {% set yUnsigned = yUnsigned is defined and yUnsigned %}
2
- {% set qMin = qMin | default(0 if yUnsigned else 0 - 128) %}
3
- {% set qMax = qMax | default(255 if yUnsigned else 127) %}
4
  // The x4 route handles widened int8/uint8 storage with scalar scale parameters;
5
  // the linear route supports both per-tensor and per-axis quantization.
 
 
 
 
 
6
  {{ env.wgsl.resourceDeclarations }}
7
 
8
- {% set divisionF16 = divisionF16 is defined and divisionF16 %}
9
  // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of
10
  // infinities, saturation, and halfway values across every kernel route.
11
  {% if divisionF16 %}
@@ -51,7 +52,6 @@ fn round_scaled_half_to_even(value: f32, scale: f32) -> i32 {
51
  }
52
  {% endif %}
53
 
54
-
55
  fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
56
  {% if hasZero %}
57
  {% if yUnsigned %}
@@ -79,9 +79,7 @@ fn transform_one(index: u32, scale: f32, zero_point: i32) {
79
 
80
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
81
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
82
- // Fold oversized flat dispatches into two dimensions; gid.y carries work
83
- // beyond the device-capped x dimension.
84
- let invocation = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
85
  {% if x4 %}
86
  if (invocation != 0u) {
87
  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
  // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of
11
  // infinities, saturation, and halfway values across every kernel route.
12
  {% if divisionF16 %}
 
52
  }
53
  {% endif %}
54
 
 
55
  fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
56
  {% if hasZero %}
57
  {% if yUnsigned %}
 
79
 
80
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
81
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
82
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "invocation", "") }}
 
 
83
  {% if x4 %}
84
  if (invocation != 0u) {
85
  return;
build/webgpu/quant-linear-vec4.wgsl.jinja CHANGED
@@ -1,12 +1,16 @@
1
- {% set yUnsigned = yUnsigned is defined and yUnsigned %}
2
- {% set qMin = qMin | default(0 if yUnsigned else 0 - 128) %}
3
- {% set qMax = qMax | default(255 if yUnsigned else 127) %}
4
  // int8/uint8 tensors are stored widened (one u32/i32 per element), so the vec4
5
  // binding gives 128-bit loads/stores of four elements. Per-component arithmetic
6
  // remains identical to the scalar quantize/dequantize paths.
 
 
 
 
 
 
 
 
7
  {{ env.wgsl.resourceDeclarations }}
8
 
9
- {% set divisionF16 = divisionF16 is defined and divisionF16 %}
10
  // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of
11
  // infinities, saturation, and halfway values across every kernel route.
12
  {% if divisionF16 %}
@@ -52,7 +56,6 @@ fn round_scaled_half_to_even(value: f32, scale: f32) -> i32 {
52
  }
53
  {% endif %}
54
 
55
-
56
  {% if vectorParams is defined and vectorParams %}
57
  fn read_zero4({% if hasZero %}index: u32{% endif %}) -> vec4<i32> {
58
  {% if hasZero %}
@@ -81,12 +84,7 @@ fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 {
81
 
82
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
83
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
84
- // The flat dispatch is folded across x/y at a fixed per-axis workgroup
85
- // width; gid.y carries the high portion of the vector index.
86
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
87
- if (i >= params.count4) {
88
- return;
89
- }
90
  {% if vectorParams is defined and vectorParams %}
91
  let scale_index = i % (params.scaleSize / 4u);
92
  {% 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
 
 
14
  // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of
15
  // infinities, saturation, and halfway values across every kernel route.
16
  {% if divisionF16 %}
 
56
  }
57
  {% endif %}
58
 
 
59
  {% if vectorParams is defined and vectorParams %}
60
  fn read_zero4({% if hasZero %}index: u32{% endif %}) -> vec4<i32> {
61
  {% if hasZero %}
 
84
 
85
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
86
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
87
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, bound="params.count4") }}
 
 
 
 
 
88
  {% if vectorParams is defined and vectorParams %}
89
  let scale_index = i % (params.scaleSize / 4u);
90
  {% elif crossingParams is defined and crossingParams %}
build/webgpu/test.json CHANGED
@@ -1128,7 +1128,7 @@
1128
  {
1129
  "name": "blocked_vec4_int8_negative_axis_partial_block",
1130
  "provenance": {
1131
- "notes": "Route lock for blocked vec4 quantization: negative axis -3 normalizes to axis 1, block_size=2 leaves a partial fifth-axis entry, inner=8 keeps each four-lane x/scale access aligned, and signed int8 output exercises both saturation bounds."
1132
  },
1133
  "attrs": { "axis": -3, "block_size": 2, "output_dtype": 3 },
1134
  "inputs": {
@@ -1148,7 +1148,7 @@
1148
  {
1149
  "name": "blocked_with_zp_int8_negative_axis_partial_block_vec4_candidate",
1150
  "provenance": {
1151
- "notes": "Asymmetric blocked-quantization lock adjacent to the no-zero-point vec4 route: negative axis -3, a partial final axis block, inner=8 alignment, lane-varying signed zero points, and saturating inputs distinguish signed ZP conversion from the symmetric path."
1152
  },
1153
  "attrs": { "axis": -3, "block_size": 2 },
1154
  "inputs": {
@@ -1173,7 +1173,7 @@
1173
  {
1174
  "name": "blocked_last_axis_vec4_no_zero_point_exact_blocks",
1175
  "provenance": {
1176
- "notes": "Route lock for the scalar-scale last-axis vec4 path: each aligned group of four values shares one scale, and signed output covers negative rounding and saturation semantics."
1177
  },
1178
  "attrs": { "axis": 1, "block_size": 4, "output_dtype": 3 },
1179
  "inputs": {
@@ -1224,7 +1224,7 @@
1224
  {
1225
  "name": "vec4_cross_axis_rows_width6_with_zero_point",
1226
  "provenance": {
1227
- "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."
1228
  },
1229
  "attrs": { "axis": 0 },
1230
  "inputs": {
@@ -1284,7 +1284,7 @@
1284
  {
1285
  "name": "f32_x_f16_scale_default_precision",
1286
  "provenance": {
1287
- "notes": "Independent T1/T2 type-variable witness: float32 x and float16 y_scale are a standard mixed route. With precision omitted, the scale type selects f16 division and the quotient rounds to integer 14."
1288
  },
1289
  "inputs": {
1290
  "x": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [35.25] } },
@@ -1297,7 +1297,7 @@
1297
  {
1298
  "name": "f16_x_f32_scale_default_precision",
1299
  "provenance": {
1300
- "notes": "Independent T1/T2 type-variable witness: float16 x and float32 y_scale are a standard mixed route. With precision omitted, the scale type selects f32 division and the quotient rounds to integer 13."
1301
  },
1302
  "inputs": {
1303
  "x": { "dtype": "float16", "shape": [1], "data": { "kind": "values", "values": [35.25] } },
@@ -1321,6 +1321,27 @@
1321
  "outputs": {
1322
  "y": { "dtype": "uint8", "shape": [1], "tolerance": 0, "data": { "kind": "values", "values": [14] } }
1323
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1324
  }
1325
  ]
1326
  }
 
1128
  {
1129
  "name": "blocked_vec4_int8_negative_axis_partial_block",
1130
  "provenance": {
1131
+ "notes": "Blocked quantization with axis -3, block size two and a partial fifth-axis entry checks axis normalization, signed int8 rounding and both saturation bounds."
1132
  },
1133
  "attrs": { "axis": -3, "block_size": 2, "output_dtype": 3 },
1134
  "inputs": {
 
1148
  {
1149
  "name": "blocked_with_zp_int8_negative_axis_partial_block_vec4_candidate",
1150
  "provenance": {
1151
+ "notes": "Negative axis -3, a partial final block, lane-varying signed zero points and saturating inputs distinguish asymmetric blocked quantization from symmetric quantization."
1152
  },
1153
  "attrs": { "axis": -3, "block_size": 2 },
1154
  "inputs": {
 
1173
  {
1174
  "name": "blocked_last_axis_vec4_no_zero_point_exact_blocks",
1175
  "provenance": {
1176
+ "notes": "Last-axis quantization with one scale per four values checks negative rounding and signed saturation."
1177
  },
1178
  "attrs": { "axis": 1, "block_size": 4, "output_dtype": 3 },
1179
  "inputs": {
 
1224
  {
1225
  "name": "vec4_cross_axis_rows_width6_with_zero_point",
1226
  "provenance": {
1227
+ "notes": "Four consecutive values cross a six-element row boundary, so they can require different per-axis scales and zero points."
1228
  },
1229
  "attrs": { "axis": 0 },
1230
  "inputs": {
 
1284
  {
1285
  "name": "f32_x_f16_scale_default_precision",
1286
  "provenance": {
1287
+ "notes": "Float32 X with a float16 Y scale divides in float16 by default: 35.25 / 2.611328125 rounds to 13.5 and quantizes to 14."
1288
  },
1289
  "inputs": {
1290
  "x": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [35.25] } },
 
1297
  {
1298
  "name": "f16_x_f32_scale_default_precision",
1299
  "provenance": {
1300
+ "notes": "Float16 X with a float32 Y scale divides in float32 by default: 35.25 / 2.611328125 = 13.499 quantizes to 13."
1301
  },
1302
  "inputs": {
1303
  "x": { "dtype": "float16", "shape": [1], "data": { "kind": "values", "values": [35.25] } },
 
1321
  "outputs": {
1322
  "y": { "dtype": "uint8", "shape": [1], "tolerance": 0, "data": { "kind": "values", "values": [14] } }
1323
  }
1324
+ },
1325
+ {
1326
+ "name": "ort_f16_int8_fractional_quotient_rounds_half_to_even",
1327
+ "provenance": {
1328
+ "source": "onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc",
1329
+ "test": "QuantizeLinearOpMLFloat16Test.Int8RoundsFractionalValues",
1330
+ "notes": "ONNX Runtime's CPU provider truncates the f16 quotient toward zero instead of rounding half-to-even. All four x values and the scale are f16-exact multiples of 2^-15 / 2^-14; the quotients are 0.5003052... and exactly 0.5, so the expected output is the same whether the division is evaluated in f16 or f32."
1331
+ },
1332
+ "inputs": {
1333
+ "x": {
1334
+ "dtype": "float16",
1335
+ "shape": [4],
1336
+ "data": {
1337
+ "kind": "values",
1338
+ "values": [0.050018310546875, -0.050018310546875, 0.04998779296875, -0.04998779296875]
1339
+ }
1340
+ },
1341
+ "y_scale": { "dtype": "float16", "shape": [], "data": { "kind": "values", "values": [0.0999755859375] } },
1342
+ "y_zero_point": { "dtype": "int8", "shape": [], "data": { "kind": "values", "values": [0] } }
1343
+ },
1344
+ "outputs": { "y": { "dtype": "int8", "shape": [4], "tolerance": 0 } }
1345
  }
1346
  ]
1347
  }