Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Conv

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
70967a1
·
verified ·
1 Parent(s): 12f6661

sync c928d21e6cc1

Browse files
Files changed (3) hide show
  1. build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja +104 -39
  2. build/webgpu/manifest.json +120 -1
  3. build/webgpu/metadata.json +4 -4
build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja CHANGED
@@ -13,6 +13,7 @@
13
  {% set fusedNarrowProjection = fusedNarrowProjection if fusedNarrowProjection is defined else false %}
14
  {% set polyphase = polyphaseConvTranspose is defined and polyphaseConvTranspose %}
15
  {% set implicit = implicitIm2col is defined and implicitIm2col %}
 
16
  {% set splitKValue = splitK if splitK is defined else 1 %}
17
  {% set splitKPartial = splitKValue > 1 %}
18
  {% set kLoopVar = "K_LOOP" if padded else "K" %}
@@ -53,7 +54,18 @@ const K_LOOP: u32 = {{ kPadded }}u;
53
  {% if implicit %}
54
  const CONV_KERNEL_H: u32 = {{ convKernelH }}u;
55
  const CONV_KERNEL_W: u32 = {{ convKernelW }}u;
 
 
 
 
 
 
 
 
 
 
56
  const CONV_KSIZE: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
 
57
  const CONV_STRIDE_H: u32 = {{ convStrideH }}u;
58
  const CONV_STRIDE_W: u32 = {{ convStrideW }}u;
59
  const CONV_DILATION_H: u32 = {{ convDilationH }}u;
@@ -63,8 +75,13 @@ const CONV_PAD_LEFT: i32 = {{ convPadLeft }};
63
  const CONV_IN_H: i32 = {{ convInH }};
64
  const CONV_IN_W: i32 = {{ convInW }};
65
  const CONV_OUT_W: u32 = {{ convOutW }}u;
 
 
 
 
66
  // X is [batch, inChannels, inH, inW]; the tile gather indexes it directly.
67
  const B_BATCH_STRIDE: u32 = {{ convInChannels }}u * u32(CONV_IN_H) * u32(CONV_IN_W);
 
68
  {% else %}
69
  {% if padded %}
70
  const N_COLS: u32 = {{ nPadded }}u;
@@ -124,7 +141,12 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
124
  let col = c_idx * {{ bLoadWidth }}u;
125
  {% if implicit %}
126
  // One output position per tile column, so its window origin is loop-invariant.
 
 
 
 
127
  let ih0 = i32((b_col / CONV_OUT_W) * CONV_STRIDE_H) - CONV_PAD_TOP;
 
128
  let iw0 = i32((b_col % CONV_OUT_W) * CONV_STRIDE_W) - CONV_PAD_LEFT;
129
  let in_column = b_col < N;
130
  // k advances by exactly one per element, so the (in-channel, tap-row, tap-col)
@@ -135,8 +157,16 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
135
  var k = k_idx + col;
136
  var ic = k / CONV_KSIZE;
137
  var kq = k % CONV_KSIZE;
 
 
 
 
 
 
 
138
  var kh = kq / CONV_KERNEL_W;
139
  var kw = kq % CONV_KERNEL_W;
 
140
  var ih = ih0 + i32(kh * CONV_DILATION_H);
141
  var iw = iw0 + i32(kw * CONV_DILATION_W);
142
  // A tile column whose whole kernel window and K span are in bounds cannot
@@ -144,16 +174,30 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
144
  // the interior loop carries a plain address instead of recomputing coordinates.
145
  let interior = in_column
146
  && k_idx + col + {{ bLoadWidth }}u <= K
 
 
 
147
  && ih0 >= 0 && ih0 + i32((CONV_KERNEL_H - 1u) * CONV_DILATION_H) < CONV_IN_H
148
  && iw0 >= 0 && iw0 + i32((CONV_KERNEL_W - 1u) * CONV_DILATION_W) < CONV_IN_W;
149
  if (interior) {
150
  let colStep = i32(CONV_DILATION_W);
151
  let rowStep = i32(CONV_DILATION_H) * CONV_IN_W;
 
 
 
 
152
  let planeStep = CONV_IN_H * CONV_IN_W;
 
153
  // Advancing k already added one colStep, so each wrap rewinds that plus the taps it walked.
154
  let kwWrap = colStep * i32(CONV_KERNEL_W) - rowStep;
 
 
 
 
 
155
  let khWrap = rowStep * i32(CONV_KERNEL_H) - planeStep;
156
  var addr = i32(b_base) + i32(ic) * planeStep + ih * CONV_IN_W + iw;
 
157
  for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
158
  tile_B[row * TILE_K + col + i] = {{ operandScalar }}(xm[u32(addr)]);
159
  addr = addr + colStep;
@@ -165,6 +209,13 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
165
  if (kh == CONV_KERNEL_H) {
166
  kh = 0u;
167
  addr = addr - khWrap;
 
 
 
 
 
 
 
168
  }
169
  }
170
  }
@@ -174,8 +225,12 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
174
  var value = {{ operandScalar }}(0.0);
175
  // The K and N tails read as exact zeros, exactly as the materialized padded
176
  // cols buffer does, so they contribute nothing to the dot product.
177
- if (in_column && k < K && ih >= 0 && ih < CONV_IN_H && iw >= 0 && iw < CONV_IN_W) {
 
 
 
178
  value = {{ operandScalar }}(xm[b_base + (ic * u32(CONV_IN_H) + u32(ih)) * u32(CONV_IN_W) + u32(iw)]);
 
179
  }
180
  tile_B[row * TILE_K + col + i] = value;
181
  k = k + 1u;
@@ -189,7 +244,17 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
189
  if (kh == CONV_KERNEL_H) {
190
  kh = 0u;
191
  ih = ih0;
 
 
 
 
 
 
 
 
 
192
  ic = ic + 1u;
 
193
  }
194
  }
195
  }
@@ -318,14 +383,14 @@ fn main(
318
 
319
  for (var step = 0u; step < TILE_K; step = step + 8u) {
320
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
321
- var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
322
- var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
323
 
324
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
325
- var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
326
- var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
327
- var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
328
- var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
329
 
330
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
331
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -345,14 +410,14 @@ fn main(
345
  // cross-subgroup handoff into the narrow projection. All collective stores
346
  // remain subgroup-uniform, and every fragment has a disjoint destination.
347
  let fused_tile_offset = base_A * TILE_COLS + base_B;
348
- subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS + 0u, matC00, false, TILE_COLS);
349
- subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS + 8u, matC01, false, TILE_COLS);
350
- subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS + 16u, matC02, false, TILE_COLS);
351
- subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS + 24u, matC03, false, TILE_COLS);
352
- subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS + 0u, matC10, false, TILE_COLS);
353
- subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS + 8u, matC11, false, TILE_COLS);
354
- subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS + 16u, matC12, false, TILE_COLS);
355
- subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS + 24u, matC13, false, TILE_COLS);
356
  workgroupBarrier();
357
 
358
  // One invocation owns one spatial column and walks the producer channels in
@@ -384,14 +449,14 @@ fn main(
384
  // scratch is whole tiles in both axes, so this needs none of the column or
385
  // row guards the real output store carries.
386
  let partial_offset = c_base + (a_global_base + base_A) * PARTIAL_COLS + b_global_base + base_B;
387
- subgroupMatrixStore(&y, partial_offset + 0u * PARTIAL_COLS + 0u, matC00, false, PARTIAL_COLS);
388
- subgroupMatrixStore(&y, partial_offset + 0u * PARTIAL_COLS + 8u, matC01, false, PARTIAL_COLS);
389
- subgroupMatrixStore(&y, partial_offset + 0u * PARTIAL_COLS + 16u, matC02, false, PARTIAL_COLS);
390
- subgroupMatrixStore(&y, partial_offset + 0u * PARTIAL_COLS + 24u, matC03, false, PARTIAL_COLS);
391
- subgroupMatrixStore(&y, partial_offset + 8u * PARTIAL_COLS + 0u, matC10, false, PARTIAL_COLS);
392
- subgroupMatrixStore(&y, partial_offset + 8u * PARTIAL_COLS + 8u, matC11, false, PARTIAL_COLS);
393
- subgroupMatrixStore(&y, partial_offset + 8u * PARTIAL_COLS + 16u, matC12, false, PARTIAL_COLS);
394
- subgroupMatrixStore(&y, partial_offset + 8u * PARTIAL_COLS + 24u, matC13, false, PARTIAL_COLS);
395
  {% elif useDirectMatrixStore %}
396
  {% if padded %}
397
  // Padded N only needs guarded publication for the final partial-width tile.
@@ -404,14 +469,14 @@ fn main(
404
  // When an epilogue is present, publish the raw accumulators first, then apply
405
  // bias, residual Z, and activation cooperatively after one storage barrier.
406
  let matrix_c_offset = c_base + (a_global_base + base_A) * N + b_global_base + base_B;
407
- subgroupMatrixStore(&y, matrix_c_offset + 0u * N + 0u, matC00, false, N);
408
- subgroupMatrixStore(&y, matrix_c_offset + 0u * N + 8u, matC01, false, N);
409
- subgroupMatrixStore(&y, matrix_c_offset + 0u * N + 16u, matC02, false, N);
410
- subgroupMatrixStore(&y, matrix_c_offset + 0u * N + 24u, matC03, false, N);
411
- subgroupMatrixStore(&y, matrix_c_offset + 8u * N + 0u, matC10, false, N);
412
- subgroupMatrixStore(&y, matrix_c_offset + 8u * N + 8u, matC11, false, N);
413
- subgroupMatrixStore(&y, matrix_c_offset + 8u * N + 16u, matC12, false, N);
414
- subgroupMatrixStore(&y, matrix_c_offset + 8u * N + 24u, matC13, false, N);
415
  {% if hasBias or hasActivation or hasZ %}
416
  storageBarrier();
417
  for (var tile_idx = local_idx; tile_idx < TILE_ROWS * TILE_COLS; tile_idx += {{ workgroupThreadsValue }}u) {
@@ -442,10 +507,10 @@ fn main(
442
  // give the reads visibility of the store AND stop the second row-group's store from
443
  // clobbering the first's still-in-flight readback when a partial final M-tile
444
  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
445
- subgroupMatrixStore(&scratch[subtile_id][0], 0u, matC00, false, 8u);
446
- subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC01, false, 8u);
447
- subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC02, false, 8u);
448
- subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC03, false, 8u);
449
  workgroupBarrier();
450
  let row = sg_id / 4u;
451
  let col = (sg_id % 4u) * 2u;
@@ -455,10 +520,10 @@ fn main(
455
  storeOutput(matrix_c_offset, {% if hasBias or polyphase %}row_base, {% endif %}row, col, subtile_id, row_limit{% if padded %}, b_global_base + base_B{% endif %});
456
  workgroupBarrier();
457
 
458
- subgroupMatrixStore(&scratch[subtile_id][0], 0u, matC10, false, 8u);
459
- subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC11, false, 8u);
460
- subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC12, false, 8u);
461
- subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC13, false, 8u);
462
  workgroupBarrier();
463
  matrix_c_offset = matrix_c_offset + 8u * N;
464
  row_limit = i32(M) - i32(row_base + 8u);
 
13
  {% set fusedNarrowProjection = fusedNarrowProjection if fusedNarrowProjection is defined else false %}
14
  {% set polyphase = polyphaseConvTranspose is defined and polyphaseConvTranspose %}
15
  {% set implicit = implicitIm2col is defined and implicitIm2col %}
16
+ {% set implicit3d = implicit and (convSpatialDims is defined and convSpatialDims == 3) %}
17
  {% set splitKValue = splitK if splitK is defined else 1 %}
18
  {% set splitKPartial = splitKValue > 1 %}
19
  {% set kLoopVar = "K_LOOP" if padded else "K" %}
 
54
  {% if implicit %}
55
  const CONV_KERNEL_H: u32 = {{ convKernelH }}u;
56
  const CONV_KERNEL_W: u32 = {{ convKernelW }}u;
57
+ {% if implicit3d %}
58
+ const CONV_KERNEL_D: u32 = {{ convKernelD }}u;
59
+ const CONV_KSIZE_HW: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
60
+ const CONV_KSIZE: u32 = CONV_KERNEL_D * CONV_KSIZE_HW;
61
+ const CONV_STRIDE_D: u32 = {{ convStrideD }}u;
62
+ const CONV_DILATION_D: u32 = {{ convDilationD }}u;
63
+ const CONV_PAD_FRONT: i32 = {{ convPadFront }};
64
+ const CONV_IN_D: i32 = {{ convInD }};
65
+ const CONV_OUT_H: u32 = {{ convOutH }}u;
66
+ {% else %}
67
  const CONV_KSIZE: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
68
+ {% endif %}
69
  const CONV_STRIDE_H: u32 = {{ convStrideH }}u;
70
  const CONV_STRIDE_W: u32 = {{ convStrideW }}u;
71
  const CONV_DILATION_H: u32 = {{ convDilationH }}u;
 
75
  const CONV_IN_H: i32 = {{ convInH }};
76
  const CONV_IN_W: i32 = {{ convInW }};
77
  const CONV_OUT_W: u32 = {{ convOutW }}u;
78
+ {% if implicit3d %}
79
+ // X is [batch, inChannels, inD, inH, inW]; the tile gather indexes it directly.
80
+ const B_BATCH_STRIDE: u32 = {{ convInChannels }}u * u32(CONV_IN_D) * u32(CONV_IN_H) * u32(CONV_IN_W);
81
+ {% else %}
82
  // X is [batch, inChannels, inH, inW]; the tile gather indexes it directly.
83
  const B_BATCH_STRIDE: u32 = {{ convInChannels }}u * u32(CONV_IN_H) * u32(CONV_IN_W);
84
+ {% endif %}
85
  {% else %}
86
  {% if padded %}
87
  const N_COLS: u32 = {{ nPadded }}u;
 
141
  let col = c_idx * {{ bLoadWidth }}u;
142
  {% if implicit %}
143
  // One output position per tile column, so its window origin is loop-invariant.
144
+ {% if implicit3d %}
145
+ let id0 = i32((b_col / (CONV_OUT_W * CONV_OUT_H)) * CONV_STRIDE_D) - CONV_PAD_FRONT;
146
+ let ih0 = i32(((b_col / CONV_OUT_W) % CONV_OUT_H) * CONV_STRIDE_H) - CONV_PAD_TOP;
147
+ {% else %}
148
  let ih0 = i32((b_col / CONV_OUT_W) * CONV_STRIDE_H) - CONV_PAD_TOP;
149
+ {% endif %}
150
  let iw0 = i32((b_col % CONV_OUT_W) * CONV_STRIDE_W) - CONV_PAD_LEFT;
151
  let in_column = b_col < N;
152
  // k advances by exactly one per element, so the (in-channel, tap-row, tap-col)
 
157
  var k = k_idx + col;
158
  var ic = k / CONV_KSIZE;
159
  var kq = k % CONV_KSIZE;
160
+ {% if implicit3d %}
161
+ var kd = kq / CONV_KSIZE_HW;
162
+ var khw = kq % CONV_KSIZE_HW;
163
+ var kh = khw / CONV_KERNEL_W;
164
+ var kw = khw % CONV_KERNEL_W;
165
+ var id = id0 + i32(kd * CONV_DILATION_D);
166
+ {% else %}
167
  var kh = kq / CONV_KERNEL_W;
168
  var kw = kq % CONV_KERNEL_W;
169
+ {% endif %}
170
  var ih = ih0 + i32(kh * CONV_DILATION_H);
171
  var iw = iw0 + i32(kw * CONV_DILATION_W);
172
  // A tile column whose whole kernel window and K span are in bounds cannot
 
174
  // the interior loop carries a plain address instead of recomputing coordinates.
175
  let interior = in_column
176
  && k_idx + col + {{ bLoadWidth }}u <= K
177
+ {% if implicit3d %}
178
+ && id0 >= 0 && id0 + i32((CONV_KERNEL_D - 1u) * CONV_DILATION_D) < CONV_IN_D
179
+ {% endif %}
180
  && ih0 >= 0 && ih0 + i32((CONV_KERNEL_H - 1u) * CONV_DILATION_H) < CONV_IN_H
181
  && iw0 >= 0 && iw0 + i32((CONV_KERNEL_W - 1u) * CONV_DILATION_W) < CONV_IN_W;
182
  if (interior) {
183
  let colStep = i32(CONV_DILATION_W);
184
  let rowStep = i32(CONV_DILATION_H) * CONV_IN_W;
185
+ {% if implicit3d %}
186
+ let depthStep = i32(CONV_DILATION_D) * CONV_IN_H * CONV_IN_W;
187
+ let planeStep = CONV_IN_D * CONV_IN_H * CONV_IN_W;
188
+ {% else %}
189
  let planeStep = CONV_IN_H * CONV_IN_W;
190
+ {% endif %}
191
  // Advancing k already added one colStep, so each wrap rewinds that plus the taps it walked.
192
  let kwWrap = colStep * i32(CONV_KERNEL_W) - rowStep;
193
+ {% if implicit3d %}
194
+ let khWrap = rowStep * i32(CONV_KERNEL_H) - depthStep;
195
+ let kdWrap = depthStep * i32(CONV_KERNEL_D) - planeStep;
196
+ var addr = i32(b_base) + i32(ic) * planeStep + (id * CONV_IN_H + ih) * CONV_IN_W + iw;
197
+ {% else %}
198
  let khWrap = rowStep * i32(CONV_KERNEL_H) - planeStep;
199
  var addr = i32(b_base) + i32(ic) * planeStep + ih * CONV_IN_W + iw;
200
+ {% endif %}
201
  for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
202
  tile_B[row * TILE_K + col + i] = {{ operandScalar }}(xm[u32(addr)]);
203
  addr = addr + colStep;
 
209
  if (kh == CONV_KERNEL_H) {
210
  kh = 0u;
211
  addr = addr - khWrap;
212
+ {% if implicit3d %}
213
+ kd = kd + 1u;
214
+ if (kd == CONV_KERNEL_D) {
215
+ kd = 0u;
216
+ addr = addr - kdWrap;
217
+ }
218
+ {% endif %}
219
  }
220
  }
221
  }
 
225
  var value = {{ operandScalar }}(0.0);
226
  // The K and N tails read as exact zeros, exactly as the materialized padded
227
  // cols buffer does, so they contribute nothing to the dot product.
228
+ if (in_column && k < K{% if implicit3d %} && id >= 0 && id < CONV_IN_D{% endif %} && ih >= 0 && ih < CONV_IN_H && iw >= 0 && iw < CONV_IN_W) {
229
+ {% if implicit3d %}
230
+ value = {{ operandScalar }}(xm[b_base + ((ic * u32(CONV_IN_D) + u32(id)) * u32(CONV_IN_H) + u32(ih)) * u32(CONV_IN_W) + u32(iw)]);
231
+ {% else %}
232
  value = {{ operandScalar }}(xm[b_base + (ic * u32(CONV_IN_H) + u32(ih)) * u32(CONV_IN_W) + u32(iw)]);
233
+ {% endif %}
234
  }
235
  tile_B[row * TILE_K + col + i] = value;
236
  k = k + 1u;
 
244
  if (kh == CONV_KERNEL_H) {
245
  kh = 0u;
246
  ih = ih0;
247
+ {% if implicit3d %}
248
+ kd = kd + 1u;
249
+ id = id + i32(CONV_DILATION_D);
250
+ if (kd == CONV_KERNEL_D) {
251
+ kd = 0u;
252
+ id = id0;
253
+ ic = ic + 1u;
254
+ }
255
+ {% else %}
256
  ic = ic + 1u;
257
+ {% endif %}
258
  }
259
  }
260
  }
 
383
 
384
  for (var step = 0u; step < TILE_K; step = step + 8u) {
385
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
386
+ var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
387
+ var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
388
 
389
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
390
+ var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
391
+ var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
392
+ var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
393
+ var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
394
 
395
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
396
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
 
410
  // cross-subgroup handoff into the narrow projection. All collective stores
411
  // remain subgroup-uniform, and every fragment has a disjoint destination.
412
  let fused_tile_offset = base_A * TILE_COLS + base_B;
413
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 0u, matC00, TILE_COLS);
414
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 8u, matC01, TILE_COLS);
415
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 16u, matC02, TILE_COLS);
416
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 24u, matC03, TILE_COLS);
417
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 0u, matC10, TILE_COLS);
418
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 8u, matC11, TILE_COLS);
419
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 16u, matC12, TILE_COLS);
420
+ subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 24u, matC13, TILE_COLS);
421
  workgroupBarrier();
422
 
423
  // One invocation owns one spatial column and walks the producer channels in
 
449
  // scratch is whole tiles in both axes, so this needs none of the column or
450
  // row guards the real output store carries.
451
  let partial_offset = c_base + (a_global_base + base_A) * PARTIAL_COLS + b_global_base + base_B;
452
+ subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS + 0u, matC00, PARTIAL_COLS);
453
+ subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS + 8u, matC01, PARTIAL_COLS);
454
+ subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS + 16u, matC02, PARTIAL_COLS);
455
+ subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS + 24u, matC03, PARTIAL_COLS);
456
+ subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS + 0u, matC10, PARTIAL_COLS);
457
+ subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS + 8u, matC11, PARTIAL_COLS);
458
+ subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS + 16u, matC12, PARTIAL_COLS);
459
+ subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS + 24u, matC13, PARTIAL_COLS);
460
  {% elif useDirectMatrixStore %}
461
  {% if padded %}
462
  // Padded N only needs guarded publication for the final partial-width tile.
 
469
  // When an epilogue is present, publish the raw accumulators first, then apply
470
  // bias, residual Z, and activation cooperatively after one storage barrier.
471
  let matrix_c_offset = c_base + (a_global_base + base_A) * N + b_global_base + base_B;
472
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 0u * N + 0u, matC00, N);
473
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 0u * N + 8u, matC01, N);
474
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 0u * N + 16u, matC02, N);
475
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 0u * N + 24u, matC03, N);
476
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 8u * N + 0u, matC10, N);
477
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 8u * N + 8u, matC11, N);
478
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 8u * N + 16u, matC12, N);
479
+ subgroupMatrixStore<row_major>(&y, matrix_c_offset + 8u * N + 24u, matC13, N);
480
  {% if hasBias or hasActivation or hasZ %}
481
  storageBarrier();
482
  for (var tile_idx = local_idx; tile_idx < TILE_ROWS * TILE_COLS; tile_idx += {{ workgroupThreadsValue }}u) {
 
507
  // give the reads visibility of the store AND stop the second row-group's store from
508
  // clobbering the first's still-in-flight readback when a partial final M-tile
509
  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
510
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC00, 8u);
511
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC01, 8u);
512
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC02, 8u);
513
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC03, 8u);
514
  workgroupBarrier();
515
  let row = sg_id / 4u;
516
  let col = (sg_id % 4u) * 2u;
 
520
  storeOutput(matrix_c_offset, {% if hasBias or polyphase %}row_base, {% endif %}row, col, subtile_id, row_limit{% if padded %}, b_global_base + base_B{% endif %});
521
  workgroupBarrier();
522
 
523
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
524
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC11, 8u);
525
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC12, 8u);
526
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC13, 8u);
527
  workgroupBarrier();
528
  matrix_c_offset = matrix_c_offset + 8u * N;
529
  row_limit = i32(M) - i32(row_base + 8u);
build/webgpu/manifest.json CHANGED
@@ -249,7 +249,10 @@
249
  "tiledReduceWorkgroupSize": "min(256, deviceWorkgroupCap)",
250
  "pointwiseUnitStride": "kernelHeight == 1 and kernelWidth == 1 and strideH == 1 and strideW == 1 and dilationH == 1 and dilationW == 1 and effectivePadTop == 0 and effectivePadLeft == 0 and outputHeight == inputHeight and outputWidth == inputWidth",
251
  "channelReduceOk": "nchwRanksOk and attrs.group == 1 and weightInChannels == inChannels and pointwiseUnitStride and conv2dOutputOk and outChannelsOk and inChannels >= 1 and outChannels <= tunables.CHANNEL_REDUCE_MAX_OUT_CHANNELS and outputSpatial >= tunables.CHANNEL_REDUCE_MIN_SPATIAL",
252
- "channelReduceWorkgroupSize": "min(tunables.CHANNEL_REDUCE_WORKGROUP_SIZE, deviceWorkgroupCap)"
 
 
 
253
  },
254
  "bindingSets": {
255
  "sgmatDirect": [
@@ -1385,6 +1388,122 @@
1385
  }
1386
  ]
1387
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1388
  {
1389
  "id": "implicit_im2col_subgroup_matrix_splitk",
1390
  "description": "Partitions the implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums from scratch. The combine reassociates floating-point addition relative to the unsplit route.",
 
249
  "tiledReduceWorkgroupSize": "min(256, deviceWorkgroupCap)",
250
  "pointwiseUnitStride": "kernelHeight == 1 and kernelWidth == 1 and strideH == 1 and strideW == 1 and dilationH == 1 and dilationW == 1 and effectivePadTop == 0 and effectivePadLeft == 0 and outputHeight == inputHeight and outputWidth == inputWidth",
251
  "channelReduceOk": "nchwRanksOk and attrs.group == 1 and weightInChannels == inChannels and pointwiseUnitStride and conv2dOutputOk and outChannelsOk and inChannels >= 1 and outChannels <= tunables.CHANNEL_REDUCE_MAX_OUT_CHANNELS and outputSpatial >= tunables.CHANNEL_REDUCE_MIN_SPATIAL",
252
+ "channelReduceWorkgroupSize": "min(tunables.CHANNEL_REDUCE_WORKGROUP_SIZE, deviceWorkgroupCap)",
253
+ "implicitIm2col3dOk": "conv3dRanksOk and attrs.group == 1 and weightInChannels == inChannels and conv3dOutputOk",
254
+ "implicitIm2col3dWorthIt": "sgmatYWorkgroups <= tunables.IMPLICIT_MAX_M_TILES and outputVolume >= tunables.IMPLICIT_MIN_OUTPUT_SPATIAL",
255
+ "paddedKernelVolumeRows": "ceilDiv(kernelVolumeRows, 32) * 32"
256
  },
257
  "bindingSets": {
258
  "sgmatDirect": [
 
1388
  }
1389
  ]
1390
  },
1391
+ {
1392
+ "id": "implicit_im2col3d_subgroup_matrix",
1393
+ "description": "Gathers each logical 3-D im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth a volumetric materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.",
1394
+ "priority": 160,
1395
+ "requires": {
1396
+ "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
1397
+ "subgroupMatrixConfigs": [
1398
+ { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
1399
+ { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
1400
+ ]
1401
+ },
1402
+ "when": ["wave32Effective", "noBiasContract", "f16Ok(dtypes.T)", "implicitIm2col3dOk", "implicitIm2col3dWorthIt", "outChannelsOk", "outChannels >= tunables.SGMAT_MIN_M", "outputVolume >= tunables.SGMAT_MIN_ALIGNED_N", "kernelVolumeRows >= tunables.SGMAT_MIN_PADDED_K", "paddedKernelVolumeRows <= device.limits.maxComputeWorkgroupsPerDimension", "batchSize >= 1", "batchSize <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(outputVolume / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatYWorkgroups <= device.limits.maxComputeWorkgroupsPerDimension"],
1403
+ "constants": {
1404
+ "usesF16": "dtypes.T == \"f16\"",
1405
+ "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
1406
+ "hasBias": false,
1407
+ "directMatrixStore": "sgmatDirectMatrixStore",
1408
+ "tileRows": "sgmatTileRows",
1409
+ "subgroupRows": "sgmatSubgroupRows",
1410
+ "M": "outChannels",
1411
+ "K": "kernelVolumeRows",
1412
+ "N": "outputVolume",
1413
+ "padded": true,
1414
+ "kPadded": "paddedKernelVolumeRows",
1415
+ "implicitIm2col": true,
1416
+ "convKernelH": "kernelHeight",
1417
+ "convKernelW": "kernelWidth",
1418
+ "convStrideH": "strideH",
1419
+ "convStrideW": "strideW",
1420
+ "convDilationH": "dilationH",
1421
+ "convDilationW": "dilationW",
1422
+ "convPadTop": "effectivePadTop",
1423
+ "convPadLeft": "effectivePadLeft",
1424
+ "convInH": "inputHeight",
1425
+ "convInW": "inputWidth",
1426
+ "convOutW": "outputWidth",
1427
+ "convInChannels": "inChannels",
1428
+ "subgroupCols": 2,
1429
+ "workgroupThreads": "sgmatWorkgroupThreads",
1430
+ "convSpatialDims": 3,
1431
+ "convKernelD": "kernelDepth",
1432
+ "convStrideD": "strideD",
1433
+ "convDilationD": "dilationD",
1434
+ "convPadFront": "effectivePadFront",
1435
+ "convInD": "inputDepth",
1436
+ "convOutH": "outputHeight"
1437
+ },
1438
+ "bindings": "sgmatDirect",
1439
+ "passes": [
1440
+ {
1441
+ "id": "main",
1442
+ "name": "Conv.ImplicitIm2col3dSubgroupMatrix",
1443
+ "shader": "conv-1x1-subgroup-matrix.wgsl.jinja",
1444
+ "dispatch": { "x": "ceil(outputVolume / 64)", "y": "sgmatYWorkgroups", "z": "batchSize" },
1445
+ "profile": { "flops": "convFlops", "gemmK": "kernelVolumeRows" }
1446
+ }
1447
+ ]
1448
+ },
1449
+ {
1450
+ "id": "implicit_im2col3d_subgroup_matrix_bias",
1451
+ "description": "Gathers each logical 3-D im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth a volumetric materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.",
1452
+ "priority": 161,
1453
+ "requires": {
1454
+ "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
1455
+ "subgroupMatrixConfigs": [
1456
+ { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
1457
+ { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
1458
+ ]
1459
+ },
1460
+ "when": ["wave32Effective", "biasContract", "f16Ok(dtypes.T)", "implicitIm2col3dOk", "implicitIm2col3dWorthIt", "outChannelsOk", "outChannels >= tunables.SGMAT_MIN_M", "outputVolume >= tunables.SGMAT_MIN_ALIGNED_N", "kernelVolumeRows >= tunables.SGMAT_MIN_PADDED_K", "paddedKernelVolumeRows <= device.limits.maxComputeWorkgroupsPerDimension", "batchSize >= 1", "batchSize <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(outputVolume / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatYWorkgroups <= device.limits.maxComputeWorkgroupsPerDimension"],
1461
+ "constants": {
1462
+ "usesF16": "dtypes.T == \"f16\"",
1463
+ "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
1464
+ "hasBias": true,
1465
+ "directMatrixStore": "sgmatDirectMatrixStore",
1466
+ "tileRows": "sgmatTileRows",
1467
+ "subgroupRows": "sgmatSubgroupRows",
1468
+ "M": "outChannels",
1469
+ "K": "kernelVolumeRows",
1470
+ "N": "outputVolume",
1471
+ "padded": true,
1472
+ "kPadded": "paddedKernelVolumeRows",
1473
+ "implicitIm2col": true,
1474
+ "convKernelH": "kernelHeight",
1475
+ "convKernelW": "kernelWidth",
1476
+ "convStrideH": "strideH",
1477
+ "convStrideW": "strideW",
1478
+ "convDilationH": "dilationH",
1479
+ "convDilationW": "dilationW",
1480
+ "convPadTop": "effectivePadTop",
1481
+ "convPadLeft": "effectivePadLeft",
1482
+ "convInH": "inputHeight",
1483
+ "convInW": "inputWidth",
1484
+ "convOutW": "outputWidth",
1485
+ "convInChannels": "inChannels",
1486
+ "subgroupCols": 2,
1487
+ "workgroupThreads": "sgmatWorkgroupThreads",
1488
+ "convSpatialDims": 3,
1489
+ "convKernelD": "kernelDepth",
1490
+ "convStrideD": "strideD",
1491
+ "convDilationD": "dilationD",
1492
+ "convPadFront": "effectivePadFront",
1493
+ "convInD": "inputDepth",
1494
+ "convOutH": "outputHeight"
1495
+ },
1496
+ "bindings": "sgmatDirectBias",
1497
+ "passes": [
1498
+ {
1499
+ "id": "main",
1500
+ "name": "Conv.ImplicitIm2col3dSubgroupMatrixBias",
1501
+ "shader": "conv-1x1-subgroup-matrix.wgsl.jinja",
1502
+ "dispatch": { "x": "ceil(outputVolume / 64)", "y": "sgmatYWorkgroups", "z": "batchSize" },
1503
+ "profile": { "flops": "convFlops", "gemmK": "kernelVolumeRows" }
1504
+ }
1505
+ ]
1506
+ },
1507
  {
1508
  "id": "implicit_im2col_subgroup_matrix_splitk",
1509
  "description": "Partitions the implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums from scratch. The combine reassociates floating-point addition relative to the unsplit route.",
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Conv",
3
- "id": "_ai_onnx_conv_webgpu_1c7bb07",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -11,7 +11,7 @@
11
  "conv-1x1-channel-reduce.wgsl.jinja": "N0hmEYctrx1dFezCieFBMSZw1AJkB9k1jqWbqwY73yw=",
12
  "conv-1x1-gemm-tiled-reg.wgsl.jinja": "uhup4og6kzn9fGBicWDhp8u7tRMWOxV1FP4h4B+1QC4=",
13
  "conv-1x1-gemm-tiled.wgsl.jinja": "VeRKrzHV+9CvWm0Lrsgpt5C98Gkz5MubGVEnYApQhdk=",
14
- "conv-1x1-subgroup-matrix.wgsl.jinja": "xx4b+ferts28wrhST9PyLEEsmreDv/2qdNwU7o5Id1o=",
15
  "conv-direct-nd.wgsl.jinja": "ANLdecurm5wdc72OIpjU8oiYtrMhL5a5z8XhDjYJXGg=",
16
  "conv-direct-unrolled.wgsl.jinja": "ILQ1vB/EtaZvB65S5LwNliPq91fcT7LCzVoPDRoDolY=",
17
  "conv-im2col-ncdhw.wgsl.jinja": "AsmaNPOjoGmh+YHO91mT5EIc7RSlGqgzm7uKASMS3vw=",
@@ -19,10 +19,10 @@
19
  "conv-splitk-reduce.wgsl.jinja": "qAkIOXGzAj6HWYSscyL1cfNgM6A1Ym/1L/CjQ8ogCEA=",
20
  "conv1d-tiled-reg.wgsl.jinja": "8xjz6LBfrQDZDFGDSYo8aV1m+yeVDvHMngjSAbOAmk8=",
21
  "conv2d-grouped-large-w4.wgsl.jinja": "jt2Z54ENODhKgmScaR19Mq3Y3gAMybt2fP49zoZfwEA=",
22
- "manifest.json": "de+S+inva5VNYxBD0nNBRaf4n1133OuqnIuYFaQ6tRA=",
23
  "test.json": "ron7OZMUp/fWx5XDPPNom2Zg0q0lYcWgRzLyfgE4N5I="
24
  }
25
  },
26
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
27
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Conv" }
28
  }
 
1
  {
2
  "name": "ai.onnx.Conv",
3
+ "id": "_ai_onnx_conv_webgpu_4822108",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
11
  "conv-1x1-channel-reduce.wgsl.jinja": "N0hmEYctrx1dFezCieFBMSZw1AJkB9k1jqWbqwY73yw=",
12
  "conv-1x1-gemm-tiled-reg.wgsl.jinja": "uhup4og6kzn9fGBicWDhp8u7tRMWOxV1FP4h4B+1QC4=",
13
  "conv-1x1-gemm-tiled.wgsl.jinja": "VeRKrzHV+9CvWm0Lrsgpt5C98Gkz5MubGVEnYApQhdk=",
14
+ "conv-1x1-subgroup-matrix.wgsl.jinja": "CTPjuc+j4mBdgIKG/LMdUSNTO/qeyS0810Gvzyl9waU=",
15
  "conv-direct-nd.wgsl.jinja": "ANLdecurm5wdc72OIpjU8oiYtrMhL5a5z8XhDjYJXGg=",
16
  "conv-direct-unrolled.wgsl.jinja": "ILQ1vB/EtaZvB65S5LwNliPq91fcT7LCzVoPDRoDolY=",
17
  "conv-im2col-ncdhw.wgsl.jinja": "AsmaNPOjoGmh+YHO91mT5EIc7RSlGqgzm7uKASMS3vw=",
 
19
  "conv-splitk-reduce.wgsl.jinja": "qAkIOXGzAj6HWYSscyL1cfNgM6A1Ym/1L/CjQ8ogCEA=",
20
  "conv1d-tiled-reg.wgsl.jinja": "8xjz6LBfrQDZDFGDSYo8aV1m+yeVDvHMngjSAbOAmk8=",
21
  "conv2d-grouped-large-w4.wgsl.jinja": "jt2Z54ENODhKgmScaR19Mq3Y3gAMybt2fP49zoZfwEA=",
22
+ "manifest.json": "hVZagZKncZQkjXOpaqdBUrFx0HXtV5u0K7rEqSuus1Q=",
23
  "test.json": "ron7OZMUp/fWx5XDPPNom2Zg0q0lYcWgRzLyfgE4N5I="
24
  }
25
  },
26
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
27
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Conv" }
28
  }