Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Softmax

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

sync c928d21e6cc1

Browse files
Files changed (9) hide show
  1. README.md +1 -3
  2. build/webgpu/manifest.json +26 -245
  3. build/webgpu/metadata.json +7 -9
  4. build/webgpu/softmax-normalize.wgsl.jinja +0 -64
  5. build/webgpu/softmax-online.wgsl.jinja +6 -2
  6. build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja +0 -138
  7. build/webgpu/softmax-row-stage.wgsl.jinja +0 -73
  8. build/webgpu/{softmax-strided-scalar4-reduce.wgsl.jinja → softmax-strided-online-coop.wgsl.jinja} +75 -64
  9. build/webgpu/softmax-strided-scalar4-normalize.wgsl.jinja +0 -80
README.md CHANGED
@@ -58,11 +58,9 @@ Some implementation variants require `subgroups`. These are route-specific capab
58
  - [`softmax-online-local.wgsl.jinja`](build/webgpu/softmax-online-local.wgsl.jinja)
59
  - [`softmax-online-packed-rows.wgsl.jinja`](build/webgpu/softmax-online-packed-rows.wgsl.jinja)
60
  - [`softmax-online.wgsl.jinja`](build/webgpu/softmax-online.wgsl.jinja)
61
- - [`softmax-row-stage-strided-vec4.wgsl.jinja`](build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja)
62
  - [`softmax-row-stage.wgsl.jinja`](build/webgpu/softmax-row-stage.wgsl.jinja)
 
63
  - [`softmax-strided-online-lane.wgsl.jinja`](build/webgpu/softmax-strided-online-lane.wgsl.jinja)
64
- - [`softmax-strided-scalar4-normalize.wgsl.jinja`](build/webgpu/softmax-strided-scalar4-normalize.wgsl.jinja)
65
- - [`softmax-strided-scalar4-reduce.wgsl.jinja`](build/webgpu/softmax-strided-scalar4-reduce.wgsl.jinja)
66
 
67
  ## Use with `@huggingface/kernels`
68
 
 
58
  - [`softmax-online-local.wgsl.jinja`](build/webgpu/softmax-online-local.wgsl.jinja)
59
  - [`softmax-online-packed-rows.wgsl.jinja`](build/webgpu/softmax-online-packed-rows.wgsl.jinja)
60
  - [`softmax-online.wgsl.jinja`](build/webgpu/softmax-online.wgsl.jinja)
 
61
  - [`softmax-row-stage.wgsl.jinja`](build/webgpu/softmax-row-stage.wgsl.jinja)
62
+ - [`softmax-strided-online-coop.wgsl.jinja`](build/webgpu/softmax-strided-online-coop.wgsl.jinja)
63
  - [`softmax-strided-online-lane.wgsl.jinja`](build/webgpu/softmax-strided-online-lane.wgsl.jinja)
 
 
64
 
65
  ## Use with `@huggingface/kernels`
66
 
build/webgpu/manifest.json CHANGED
@@ -24,6 +24,8 @@
24
  },
25
  "tunables": {
26
  "WORKGROUP_SIZE": 256,
 
 
27
  "BLOCK_COLS": 1024,
28
  "LONGROW_SPLIT_MIN_COLS": 65536,
29
  "LONGROW_SPLIT_MAX_ROWS": 32,
@@ -44,13 +46,14 @@
44
  "axisRows": "axisOuter * axisInner",
45
  "longrowBlocks": "ceilDiv(lastCols, tunables.BLOCK_COLS)",
46
  "lastRowScratchFits": "4 * lastRows <= storageBufferLimit",
47
- "axisRowScratchFits": "4 * axisRows <= storageBufferLimit",
48
  "longrowBlockScratchFits": "4 * lastRows * longrowBlocks <= storageBufferLimit",
49
  "lastDispatchFits": "ceilDiv(lastRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
50
  "axisDispatchFits": "ceilDiv(axisRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
51
  "axisThreadDispatchFits": "ceilDiv(ceilDiv(axisRows, tunables.WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
52
  "lastAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 1 and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and f16Ok(dtypes.T)",
53
- "stridedAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 2 and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)"
 
 
54
  },
55
  "bindingSets": {
56
  "longrowBlockStats": [
@@ -187,30 +190,9 @@
187
  }
188
  }
189
  ],
190
- "stridedScalarReduce": [
191
- { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
192
- { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" }
193
- ],
194
- "stridedScalarExpSum": [
195
- { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
196
- { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
197
- { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" },
198
- { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" }
199
- ],
200
  "stridedOnline": [
201
  { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
202
  { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" }
203
- ],
204
- "stridedNormalize": [
205
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
206
- { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
207
- {
208
- "name": "params",
209
- "semantic": "kernel.params",
210
- "buffer": { "type": "uniform" },
211
- "immediate": true,
212
- "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.x)" }] }
213
- }
214
  ]
215
  },
216
  "variants": [
@@ -287,7 +269,8 @@
287
  "combineSubgroups": "device.features.has(\"subgroups\") and not narrowSubgroupRange",
288
  "scalar": "dtypes.T",
289
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
290
- "usesF16": "dtypes.T == \"f16\""
 
291
  },
292
  "passes": [
293
  {
@@ -306,7 +289,8 @@
306
  "constants": {
307
  "combineSubgroups": "device.features.has(\"subgroups\") and not narrowSubgroupRange",
308
  "scalar": "dtypes.T",
309
- "usesF16": "dtypes.T == \"f16\""
 
310
  },
311
  "passes": [
312
  {
@@ -355,7 +339,8 @@
355
  "combineSubgroups": false,
356
  "scalar": "dtypes.T",
357
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
358
- "usesF16": "dtypes.T == \"f16\""
 
359
  },
360
  "passes": [
361
  {
@@ -372,7 +357,12 @@
372
  "priority": 13,
373
  "tunables": { "WORKGROUP_SIZE": 32 },
374
  "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
375
- "constants": { "combineSubgroups": false, "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
 
 
 
 
 
376
  "passes": [
377
  {
378
  "id": "main",
@@ -384,161 +374,19 @@
384
  ]
385
  },
386
  {
387
- "id": "strided_3pass_vec4",
388
- "priority": 8,
389
- "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 == 0", "axisDispatchFits", "axisRowScratchFits"],
390
- "intermediates": [
391
- {
392
- "id": "rowMax",
393
- "dtype": "float32",
394
- "shape": "[numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)]"
395
- },
396
- {
397
- "id": "rowSum",
398
- "dtype": "float32",
399
- "shape": "[numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)]"
400
- }
401
- ],
402
- "passes": [
403
- {
404
- "id": "reduce_max",
405
- "name": "Softmax.StridedReduceMaxVec4",
406
- "source": {
407
- "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
408
- "inputs": {
409
- "stage": "\"max\"",
410
- "xShape": "shapes.x",
411
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
412
- }
413
- },
414
- "bindings": [
415
- {
416
- "name": "x",
417
- "arg": "x",
418
- "semantic": "input",
419
- "buffer": { "type": "read-only-storage" },
420
- "elementType": "vec4<f32>"
421
- },
422
- { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }
423
- ],
424
- "dispatch": {
425
- "workgroups": "numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) / 4"
426
- }
427
- },
428
- {
429
- "id": "exp_sum",
430
- "name": "Softmax.StridedExpSumVec4",
431
- "source": {
432
- "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
433
- "inputs": {
434
- "stage": "\"exp_sum\"",
435
- "xShape": "shapes.x",
436
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
437
- }
438
- },
439
- "bindings": [
440
- {
441
- "name": "x",
442
- "arg": "x",
443
- "semantic": "input",
444
- "buffer": { "type": "read-only-storage" },
445
- "elementType": "vec4<f32>"
446
- },
447
- {
448
- "name": "rowMax",
449
- "semantic": "rowMax",
450
- "buffer": { "type": "read-only-storage" },
451
- "elementType": "vec4<f32>"
452
- },
453
- {
454
- "name": "y",
455
- "arg": "y",
456
- "semantic": "output",
457
- "buffer": { "type": "storage" },
458
- "elementType": "vec4<f32>"
459
- },
460
- { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }
461
- ],
462
- "dispatch": {
463
- "workgroups": "numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) / 4"
464
- }
465
- },
466
- {
467
- "id": "normalize",
468
- "name": "Softmax.StridedNormalize",
469
- "source": {
470
- "shader": "softmax-normalize.wgsl.jinja",
471
- "inputs": {
472
- "strided": "true",
473
- "xShape": "shapes.x",
474
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
475
- }
476
- },
477
- "bindings": "stridedNormalize",
478
- "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
479
- }
480
- ]
481
- },
482
- {
483
- "id": "strided_3pass_scalar4",
484
- "priority": 9,
485
- "tunables": { "WORKGROUP_SIZE": 128 },
486
- "demoteWhen": ["dtypes.T == \"f32\" and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 == 0"],
487
- "when": ["stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) > 0", "f16Ok(dtypes.T)", "axisDispatchFits", "axisRowScratchFits"],
488
- "intermediates": [
489
- {
490
- "id": "rowMax",
491
- "dtype": "float32",
492
- "shape": "[numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)]"
493
- },
494
- {
495
- "id": "rowSum",
496
- "dtype": "float32",
497
- "shape": "[numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)]"
498
- }
499
- ],
500
  "passes": [
501
  {
502
- "id": "reduce_max",
503
- "name": "Softmax.StridedReduceMaxScalar4",
504
- "source": {
505
- "shader": "softmax-strided-scalar4-reduce.wgsl.jinja",
506
- "inputs": {
507
- "stage": "\"max\"",
508
- "xShape": "shapes.x",
509
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
510
- }
511
- },
512
- "bindings": "stridedScalarReduce",
513
- "dispatch": {
514
- "workgroups": "outer(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) * ceil(inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) / 4)"
515
- }
516
- },
517
- {
518
- "id": "exp_sum",
519
- "name": "Softmax.StridedExpSumScalar4",
520
- "source": {
521
- "shader": "softmax-strided-scalar4-reduce.wgsl.jinja",
522
- "inputs": {
523
- "stage": "\"exp_sum\"",
524
- "xShape": "shapes.x",
525
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
526
- }
527
- },
528
- "bindings": "stridedScalarExpSum",
529
- "dispatch": {
530
- "workgroups": "outer(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) * ceil(inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) / 4)"
531
- }
532
- },
533
- {
534
- "id": "normalize",
535
- "name": "Softmax.StridedNormalizeScalar4",
536
  "source": {
537
- "shader": "softmax-strided-scalar4-normalize.wgsl.jinja",
538
- "inputs": { "xShape": "shapes.x", "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" }
539
  },
540
- "bindings": "stridedNormalize",
541
- "dispatch": { "threads": "ceil(numel(shapes.x) / 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
542
  }
543
  ]
544
  },
@@ -577,73 +425,6 @@
577
  }
578
  ]
579
  },
580
- {
581
- "id": "strided_3pass",
582
- "priority": 5,
583
- "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisDispatchFits", "axisRowScratchFits"],
584
- "intermediates": [
585
- {
586
- "id": "rowMax",
587
- "dtype": "float32",
588
- "shape": "[numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)]"
589
- },
590
- {
591
- "id": "rowSum",
592
- "dtype": "float32",
593
- "shape": "[numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)]"
594
- }
595
- ],
596
- "passes": [
597
- {
598
- "id": "reduce_max",
599
- "name": "Softmax.StridedReduceMax",
600
- "source": {
601
- "shader": "softmax-row-stage.wgsl.jinja",
602
- "inputs": {
603
- "stage": "\"max\"",
604
- "strided": "true",
605
- "xShape": "shapes.x",
606
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
607
- }
608
- },
609
- "bindings": "stridedScalarReduce",
610
- "dispatch": {
611
- "workgroups": "numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
612
- }
613
- },
614
- {
615
- "id": "exp_sum",
616
- "name": "Softmax.StridedExpSum",
617
- "source": {
618
- "shader": "softmax-row-stage.wgsl.jinja",
619
- "inputs": {
620
- "stage": "\"exp_sum\"",
621
- "strided": "true",
622
- "xShape": "shapes.x",
623
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
624
- }
625
- },
626
- "bindings": "stridedScalarExpSum",
627
- "dispatch": {
628
- "workgroups": "numel(shapes.x) / dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
629
- }
630
- },
631
- {
632
- "id": "normalize",
633
- "name": "Softmax.StridedNormalize",
634
- "source": {
635
- "shader": "softmax-normalize.wgsl.jinja",
636
- "inputs": {
637
- "strided": "true",
638
- "xShape": "shapes.x",
639
- "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x"
640
- }
641
- },
642
- "bindings": "stridedNormalize",
643
- "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
644
- }
645
- ]
646
- },
647
  {
648
  "id": "strided_online_lane_large_inner",
649
  "priority": 12,
 
24
  },
25
  "tunables": {
26
  "WORKGROUP_SIZE": 256,
27
+ "ROW_MIN_WORKGROUP": 64,
28
+ "ROW_COLS_PER_LANE": 8,
29
  "BLOCK_COLS": 1024,
30
  "LONGROW_SPLIT_MIN_COLS": 65536,
31
  "LONGROW_SPLIT_MAX_ROWS": 32,
 
46
  "axisRows": "axisOuter * axisInner",
47
  "longrowBlocks": "ceilDiv(lastCols, tunables.BLOCK_COLS)",
48
  "lastRowScratchFits": "4 * lastRows <= storageBufferLimit",
 
49
  "longrowBlockScratchFits": "4 * lastRows * longrowBlocks <= storageBufferLimit",
50
  "lastDispatchFits": "ceilDiv(lastRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
51
  "axisDispatchFits": "ceilDiv(axisRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
52
  "axisThreadDispatchFits": "ceilDiv(ceilDiv(axisRows, tunables.WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
53
  "lastAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 1 and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and f16Ok(dtypes.T)",
54
+ "stridedAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 2 and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)",
55
+ "stridedCoopTileCols": "min(128, max(8, pow2ceil(ceilDiv(axisOuter * axisInner, 256))))",
56
+ "rowWorkgroupSize": "min(tunables.WORKGROUP_SIZE, max(tunables.ROW_MIN_WORKGROUP, pow2ceil(ceilDiv(lastCols, tunables.ROW_COLS_PER_LANE))))"
57
  },
58
  "bindingSets": {
59
  "longrowBlockStats": [
 
190
  }
191
  }
192
  ],
 
 
 
 
 
 
 
 
 
 
193
  "stridedOnline": [
194
  { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
195
  { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" }
 
 
 
 
 
 
 
 
 
 
 
196
  ]
197
  },
198
  "variants": [
 
269
  "combineSubgroups": "device.features.has(\"subgroups\") and not narrowSubgroupRange",
270
  "scalar": "dtypes.T",
271
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
272
+ "usesF16": "dtypes.T == \"f16\"",
273
+ "workgroupSize": "rowWorkgroupSize"
274
  },
275
  "passes": [
276
  {
 
289
  "constants": {
290
  "combineSubgroups": "device.features.has(\"subgroups\") and not narrowSubgroupRange",
291
  "scalar": "dtypes.T",
292
+ "usesF16": "dtypes.T == \"f16\"",
293
+ "workgroupSize": "rowWorkgroupSize"
294
  },
295
  "passes": [
296
  {
 
339
  "combineSubgroups": false,
340
  "scalar": "dtypes.T",
341
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
342
+ "usesF16": "dtypes.T == \"f16\"",
343
+ "workgroupSize": "tunables.WORKGROUP_SIZE"
344
  },
345
  "passes": [
346
  {
 
357
  "priority": 13,
358
  "tunables": { "WORKGROUP_SIZE": 32 },
359
  "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
360
+ "constants": {
361
+ "combineSubgroups": false,
362
+ "scalar": "dtypes.T",
363
+ "usesF16": "dtypes.T == \"f16\"",
364
+ "workgroupSize": "tunables.WORKGROUP_SIZE"
365
+ },
366
  "passes": [
367
  {
368
  "id": "main",
 
374
  ]
375
  },
376
  {
377
+ "id": "strided_online_coop",
378
+ "priority": 13,
379
+ "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisDispatchFits"],
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
380
  "passes": [
381
  {
382
+ "id": "main",
383
+ "name": "Softmax.StridedOnlineCoop",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
384
  "source": {
385
+ "shader": "softmax-strided-online-coop.wgsl.jinja",
386
+ "inputs": { "xShape": "shapes.x", "axis": "axis", "tileCols": "stridedCoopTileCols" }
387
  },
388
+ "bindings": "stridedOnline",
389
+ "dispatch": { "workgroups": "ceilDiv(axisOuter * axisInner, stridedCoopTileCols)" }
390
  }
391
  ]
392
  },
 
425
  }
426
  ]
427
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
428
  {
429
  "id": "strided_online_lane_large_inner",
430
  "priority": 12,
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Softmax",
3
- "id": "_ai_onnx_softmax_webgpu_c404d21",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,21 +8,19 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "IWhtZGnfzrlXJApME4TReOxOt/WgFYva6gNiC0pSRXY=",
11
- "manifest.json": "AJJM7OMdQ7WDrMFllmIIj+pCqK0rV3wPvB5y/AE0pks=",
12
  "softmax-longrow-normalize.wgsl.jinja": "9G+cx3d86wJm/QrOlAxkKqAKC9B/Qkh+BxtE4utrfyA=",
13
  "softmax-longrow-stats.wgsl.jinja": "K0PI+9cmRgjJsOOcBRyugN0Dm7yxBy0xvyPamEpCgxo=",
14
- "softmax-normalize.wgsl.jinja": "JMUDWqv66mpTdzhasbI6mb8TUaAtF5zv2M8IPfHo+dU=",
15
  "softmax-online-local.wgsl.jinja": "P+bXw3Rk32WOdy9setNhL40emyF7xT/xKLKUW+ggzFE=",
16
  "softmax-online-packed-rows.wgsl.jinja": "qMLjill/PhoAXExDpUcwYKJzYyH9CX6hRqni8ifJqhI=",
17
- "softmax-online.wgsl.jinja": "NMj90A8eK97QK2RACDYQlA/9I0txsZdgAflmjGF+eFM=",
18
- "softmax-row-stage-strided-vec4.wgsl.jinja": "l1byT6xoLl6Q7kSPd/uws6X5NsS2PwlSv2TBMCioVZA=",
19
- "softmax-row-stage.wgsl.jinja": "dxlM6y02DwXeLeavHDiq70h0NCRZgz+bniCltyEgRX4=",
20
  "softmax-strided-online-lane.wgsl.jinja": "4zyPmcDECprAw/kqbi6eyp4msdKlUmslPxf7zzm4+kU=",
21
- "softmax-strided-scalar4-normalize.wgsl.jinja": "+Mxy5B3Yh91xnVGEeqdy/8jToAr+Tug2C7ee2z44mqM=",
22
- "softmax-strided-scalar4-reduce.wgsl.jinja": "BnX4GiA+GWD2pBute0sRLX4AaI842mcsId2Ufy4Ny4M=",
23
  "test.json": "g2KZicqIV4q33Ekn/Sj8ZGgRPLPRam59CzZmHesRLu8="
24
  }
25
  },
26
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
27
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Softmax" }
28
  }
 
1
  {
2
  "name": "ai.onnx.Softmax",
3
+ "id": "_ai_onnx_softmax_webgpu_f440486",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "IWhtZGnfzrlXJApME4TReOxOt/WgFYva6gNiC0pSRXY=",
11
+ "manifest.json": "iB6hH3a+gt1He/iVMiq85duryk5DPtH4O2YgOunmISY=",
12
  "softmax-longrow-normalize.wgsl.jinja": "9G+cx3d86wJm/QrOlAxkKqAKC9B/Qkh+BxtE4utrfyA=",
13
  "softmax-longrow-stats.wgsl.jinja": "K0PI+9cmRgjJsOOcBRyugN0Dm7yxBy0xvyPamEpCgxo=",
14
+ "softmax-normalize.wgsl.jinja": "+1V54QupLFjjPfEwDjElnmjnDNV8DhjErqsIeP8qPag=",
15
  "softmax-online-local.wgsl.jinja": "P+bXw3Rk32WOdy9setNhL40emyF7xT/xKLKUW+ggzFE=",
16
  "softmax-online-packed-rows.wgsl.jinja": "qMLjill/PhoAXExDpUcwYKJzYyH9CX6hRqni8ifJqhI=",
17
+ "softmax-online.wgsl.jinja": "Rzfdp0Ri4pPG1hnlwps9SwV+RWgL1a9aYqZKXYoeZ+g=",
18
+ "softmax-row-stage.wgsl.jinja": "Mc34LCJFgUyX85CMaxM8QpTkPeCTlrzkK5bRn7wEJjw=",
19
+ "softmax-strided-online-coop.wgsl.jinja": "rnwhlx2bR0eH2SdsmIUKLizepWJZJGZvBhpgiPJgn0s=",
20
  "softmax-strided-online-lane.wgsl.jinja": "4zyPmcDECprAw/kqbi6eyp4msdKlUmslPxf7zzm4+kU=",
 
 
21
  "test.json": "g2KZicqIV4q33Ekn/Sj8ZGgRPLPRam59CzZmHesRLu8="
22
  }
23
  },
24
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
25
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Softmax" }
26
  }
build/webgpu/softmax-normalize.wgsl.jinja CHANGED
@@ -7,66 +7,6 @@
7
  enable f16;
8
  {% endif %}
9
  {{ env.wgsl.resourceDeclarations }}
10
- {% if source.strided %}
11
-
12
- {% set stridedAxisUses = ["row_for_offset"] %}
13
- // Compiled indexing for row-wise reductions over a non-last axis. A row is one
14
- // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
15
- {% set reduce_size = source.xShape[source.axis] %}
16
- {% set inner_size = namespace(value=1) %}
17
- {% for i in range(source.axis + 1, source.xShape | length) %}
18
- {% set inner_size.value = inner_size.value * source.xShape[i] %}
19
- {% endfor %}
20
- {% set outer_size = namespace(value=1) %}
21
- {% for i in range(source.axis) %}
22
- {% set outer_size.value = outer_size.value * source.xShape[i] %}
23
- {% endfor %}
24
- {% set usesAll = stridedAxisUses is not defined %}
25
- {% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
26
- {% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
27
- {% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
28
- {% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
29
- {% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
30
- {% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
31
- {% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
32
- {% if usesReduceSize %}
33
- const REDUCE_SIZE: u32 = {{ reduce_size }}u;
34
- {% endif %}
35
- {% if usesInnerSize %}
36
- const INNER_SIZE: u32 = {{ inner_size.value }}u;
37
- {% endif %}
38
- {% if usesStridedRows %}
39
- // Total one-per-row workgroup count (product of every non-axis dim). Lets the
40
- // row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
41
- // the over-dispatched tail without a runtime `params` uniform.
42
- const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
43
- {% endif %}
44
- {% if usesInnerSizeSafe %}
45
- const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
46
- {% endif %}
47
- {% if usesAxisStrideSafe %}
48
- const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
49
- {% endif %}
50
-
51
- {% if usesElementOffset %}
52
- // Flat offset of element k (along the reduce axis) of `row`.
53
- fn element_offset(row: u32, k: u32) -> u32 {
54
- let outer = row / INNER_SIZE_SAFE;
55
- let inner = row % INNER_SIZE_SAFE;
56
- return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
57
- }
58
-
59
- {% endif %}
60
- {% if usesRowForOffset %}
61
- // Row that owns the flat element `offset` (inverse of element_offset over
62
- // the non-axis coordinates).
63
- fn row_for_offset(offset: u32) -> u32 {
64
- let outer = offset / AXIS_STRIDE_SAFE;
65
- let inner = offset % INNER_SIZE_SAFE;
66
- return outer * INNER_SIZE_SAFE + inner;
67
- }
68
- {%- endif -%}
69
- {% endif %}
70
 
71
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
72
  fn main(@builtin(global_invocation_id) gid : vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
@@ -76,11 +16,7 @@ fn main(@builtin(global_invocation_id) gid : vec3<u32>, @builtin(num_workgroups)
76
  if (i >= params.count) {
77
  return;
78
  }
79
- {% if source.strided %}
80
- let sum = rowSum[row_for_offset(i)];
81
- {% else %}
82
  let sum = rowSum[i / params.cols];
83
- {% endif %}
84
  // Bit-exact NaN test: `sum != sum` can be folded to false by fast-math.
85
  let sumBits = bitcast<u32>(sum);
86
  let sumIsNan = (sumBits & 0x7f800000u) == 0x7f800000u && (sumBits & 0x007fffffu) != 0u;
 
7
  enable f16;
8
  {% endif %}
9
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
 
11
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
12
  fn main(@builtin(global_invocation_id) gid : vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 
16
  if (i >= params.count) {
17
  return;
18
  }
 
 
 
19
  let sum = rowSum[i / params.cols];
 
20
  // Bit-exact NaN test: `sum != sum` can be folded to false by fast-math.
21
  let sumBits = bitcast<u32>(sum);
22
  let sumIsNan = (sumBits & 0x7f800000u) == 0x7f800000u && (sumBits & 0x007fffffu) != 0u;
build/webgpu/softmax-online.wgsl.jinja CHANGED
@@ -12,7 +12,11 @@ enable subgroups;
12
  // the shared online-softmax combiner) and a fused second sweep writes
13
  // the result. The epilogue is either softmax or log-softmax, and vec4 row access
14
  // is used when rows are vectorized.
15
- const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 
 
 
 
16
  // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
17
  // `m - m` finite so an empty lane / all--inf row contributes the exact
18
  // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
@@ -174,7 +178,7 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
174
  {% endif %}
175
 
176
 
177
- @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
178
  fn main(@builtin(workgroup_id) wid: vec3<u32>,
179
  @builtin(num_workgroups) nwg: vec3<u32>,
180
  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
 
12
  // the shared online-softmax combiner) and a fused second sweep writes
13
  // the result. The epilogue is either softmax or log-softmax, and vec4 row access
14
  // is used when rows are vectorized.
15
+ //
16
+ // `workgroupSize` is a variant constant, not the op-wide tunable: the merge tree
17
+ // costs a barrier per level whatever the row length, so a short row wants fewer,
18
+ // busier lanes and a long row wants the widest workgroup the device allows.
19
+ const WG: u32 = {{ workgroupSize }}u;
20
  // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
21
  // `m - m` finite so an empty lane / all--inf row contributes the exact
22
  // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
 
178
  {% endif %}
179
 
180
 
181
+ @compute @workgroup_size({{ workgroupSize }})
182
  fn main(@builtin(workgroup_id) wid: vec3<u32>,
183
  @builtin(num_workgroups) nwg: vec3<u32>,
184
  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja DELETED
@@ -1,138 +0,0 @@
1
- // Vec4-over-inner specialization for passes 1 and 2 of strided three-pass
2
- // softmax. A strided row reduces REDUCE_SIZE elements at INNER_SIZE stride, so
3
- // adjacent scalar threads would read non-coalesced locations. When INNER_SIZE
4
- // is divisible by four, four adjacent inner positions remain contiguous at
5
- // every axis coordinate. One workgroup therefore handles four rows packed as a
6
- // vec4: one coalesced 128-bit access per axis step through one reduction tree.
7
- // rowMax and rowSum use the same bytes the scalar normalize pass reads.
8
- {{ env.wgsl.resourceDeclarations }}
9
-
10
- const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
11
- {% set stridedAxisUses = ["REDUCE_SIZE", "INNER_SIZE", "STRIDED_ROWS"] %}
12
- // Compiled indexing for row-wise reductions over a non-last axis. A row is one
13
- // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
14
- {% set reduce_size = source.xShape[source.axis] %}
15
- {% set inner_size = namespace(value=1) %}
16
- {% for i in range(source.axis + 1, source.xShape | length) %}
17
- {% set inner_size.value = inner_size.value * source.xShape[i] %}
18
- {% endfor %}
19
- {% set outer_size = namespace(value=1) %}
20
- {% for i in range(source.axis) %}
21
- {% set outer_size.value = outer_size.value * source.xShape[i] %}
22
- {% endfor %}
23
- {% set usesAll = stridedAxisUses is not defined %}
24
- {% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
25
- {% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
26
- {% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
27
- {% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
28
- {% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
29
- {% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
30
- {% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
31
- {% if usesReduceSize %}
32
- const REDUCE_SIZE: u32 = {{ reduce_size }}u;
33
- {% endif %}
34
- {% if usesInnerSize %}
35
- const INNER_SIZE: u32 = {{ inner_size.value }}u;
36
- {% endif %}
37
- {% if usesStridedRows %}
38
- // Total one-per-row workgroup count (product of every non-axis dim). Lets the
39
- // row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
40
- // the over-dispatched tail without a runtime `params` uniform.
41
- const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
42
- {% endif %}
43
- {% if usesInnerSizeSafe %}
44
- const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
45
- {% endif %}
46
- {% if usesAxisStrideSafe %}
47
- const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
48
- {% endif %}
49
-
50
- {% if usesElementOffset %}
51
- // Flat offset of element k (along the reduce axis) of `row`.
52
- fn element_offset(row: u32, k: u32) -> u32 {
53
- let outer = row / INNER_SIZE_SAFE;
54
- let inner = row % INNER_SIZE_SAFE;
55
- return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
56
- }
57
-
58
- {% endif %}
59
- {% if usesRowForOffset %}
60
- // Row that owns the flat element `offset` (inverse of element_offset over
61
- // the non-axis coordinates).
62
- fn row_for_offset(offset: u32) -> u32 {
63
- let outer = offset / AXIS_STRIDE_SAFE;
64
- let inner = offset % INNER_SIZE_SAFE;
65
- return outer * INNER_SIZE_SAFE + inner;
66
- }
67
- {%- endif -%}
68
- {% if source.stage == "max" %}
69
- // max(1u,..) so an empty strided tensor (INNER_SIZE 0) still compiles — the kernel
70
- // returns before any divide (VEC4_ROWS 0 -> guard), this only avoids a const `/0`.
71
- const INNER4: u32 = max(1u, INNER_SIZE / 4u); // vec4 groups of inner positions per outer
72
- const STRIDE4: u32 = INNER_SIZE / 4u; // vec4 stride along the reduce axis
73
- const VEC4_ROWS: u32 = STRIDED_ROWS / 4u; // = outer_size * INNER4
74
- {% else %}
75
- // max(1u,..) so an empty strided tensor (INNER_SIZE 0) still compiles (kernel
76
- // returns before any divide; only avoids a const `/0`).
77
- const INNER4: u32 = max(1u, INNER_SIZE / 4u);
78
- const STRIDE4: u32 = INNER_SIZE / 4u;
79
- const VEC4_ROWS: u32 = STRIDED_ROWS / 4u;
80
- {% endif %}
81
-
82
- var<workgroup> partial: array<vec4<f32>, WGS>;
83
- {% set reduce_fn = "reduce_max_vec4" if source.stage == "max" else "reduce_add_vec4" %}
84
- fn {{ reduce_fn }}(value: vec4<f32>, tid: u32) -> vec4<f32> {
85
- partial[tid] = value;
86
- workgroupBarrier();
87
- var stride: u32 = WGS / 2u;
88
- loop {
89
- if (stride == 0u) { break; }
90
- {% if source.stage == "max" %}
91
- if (tid < stride) { partial[tid] = max(partial[tid], partial[tid + stride]); }
92
- {% else %}
93
- if (tid < stride) { partial[tid] = partial[tid] + partial[tid + stride]; }
94
- {% endif %}
95
- stride = stride / 2u;
96
- workgroupBarrier();
97
- }
98
- return partial[0];
99
- }
100
-
101
- @compute @workgroup_size(WGS)
102
- fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
103
- @builtin(num_workgroups) nwg: vec3<u32>) {
104
- let rg = wg.x + wg.y * nwg.x; // vec4-row-group index
105
- if (rg >= VEC4_ROWS) { return; }
106
- let outer = rg / INNER4;
107
- let inner_group = rg % INNER4;
108
- {% if source.stage == "max" %}
109
- // vec4 index of element k=0 for this group (offset is a multiple of 4 since
110
- // INNER_SIZE % 4 == 0 and inner_group*4 is too).
111
- {% endif %}
112
- let base4 = (outer * REDUCE_SIZE * INNER_SIZE + inner_group * 4u) / 4u;
113
-
114
- {% if source.stage == "max" %}
115
- var acc = vec4<f32>(-3.4028234663852886e38);
116
- for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
117
- acc = max(acc, x[base4 + k * STRIDE4]);
118
- }
119
- let rowM = reduce_max_vec4(acc, lid.x);
120
- if (lid.x == 0u) { rowMax[rg] = rowM; }
121
- {% else %}
122
- let m = rowMax[rg];
123
- // +Inf row max -> NaN row (ONNX). Derive the quiet NaN per-component from the
124
- // runtime bits (a constant NaN pattern is rejected at const-eval); m - m folds
125
- // to 0 for finite m so the exp arg stays correct on the non-Inf lanes.
126
- let nanv = bitcast<vec4<f32>>(bitcast<vec4<u32>>(m) | vec4<u32>(0x7fc00000u));
127
- let isInf = m > vec4<f32>(3.4028234663852886e38);
128
- var acc = vec4<f32>(0.0);
129
- for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
130
- let idx = base4 + k * STRIDE4;
131
- let e = select(exp(x[idx] - m), nanv, isInf);
132
- y[idx] = e;
133
- acc = acc + e;
134
- }
135
- let rowSumValue = reduce_add_vec4(acc, lid.x);
136
- if (lid.x == 0u) { rowSum[rg] = rowSumValue; }
137
- {% endif %}
138
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
build/webgpu/softmax-row-stage.wgsl.jinja CHANGED
@@ -7,65 +7,6 @@ enable f16;
7
  {{ env.wgsl.resourceDeclarations }}
8
 
9
  const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
10
- {% if source.strided %}
11
- {% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
12
- // Compiled indexing for row-wise reductions over a non-last axis. A row is one
13
- // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
14
- {% set reduce_size = source.xShape[source.axis] %}
15
- {% set inner_size = namespace(value=1) %}
16
- {% for i in range(source.axis + 1, source.xShape | length) %}
17
- {% set inner_size.value = inner_size.value * source.xShape[i] %}
18
- {% endfor %}
19
- {% set outer_size = namespace(value=1) %}
20
- {% for i in range(source.axis) %}
21
- {% set outer_size.value = outer_size.value * source.xShape[i] %}
22
- {% endfor %}
23
- {% set usesAll = stridedAxisUses is not defined %}
24
- {% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
25
- {% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
26
- {% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
27
- {% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
28
- {% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
29
- {% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
30
- {% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
31
- {% if usesReduceSize %}
32
- const REDUCE_SIZE: u32 = {{ reduce_size }}u;
33
- {% endif %}
34
- {% if usesInnerSize %}
35
- const INNER_SIZE: u32 = {{ inner_size.value }}u;
36
- {% endif %}
37
- {% if usesStridedRows %}
38
- // Total one-per-row workgroup count (product of every non-axis dim). Lets the
39
- // row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
40
- // the over-dispatched tail without a runtime `params` uniform.
41
- const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
42
- {% endif %}
43
- {% if usesInnerSizeSafe %}
44
- const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
45
- {% endif %}
46
- {% if usesAxisStrideSafe %}
47
- const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
48
- {% endif %}
49
-
50
- {% if usesElementOffset %}
51
- // Flat offset of element k (along the reduce axis) of `row`.
52
- fn element_offset(row: u32, k: u32) -> u32 {
53
- let outer = row / INNER_SIZE_SAFE;
54
- let inner = row % INNER_SIZE_SAFE;
55
- return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
56
- }
57
-
58
- {% endif %}
59
- {% if usesRowForOffset %}
60
- // Row that owns the flat element `offset` (inverse of element_offset over
61
- // the non-axis coordinates).
62
- fn row_for_offset(offset: u32) -> u32 {
63
- let outer = offset / AXIS_STRIDE_SAFE;
64
- let inner = offset % INNER_SIZE_SAFE;
65
- return outer * INNER_SIZE_SAFE + inner;
66
- }
67
- {%- endif -%}
68
- {% endif %}
69
  {% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
70
  fn {{ name }}(value: f32, tid: u32) -> f32 {
71
  {{ buffer }}[tid] = value;
@@ -112,19 +53,6 @@ var<workgroup> partial: array<f32, WGS>;
112
  @compute @workgroup_size(WGS)
113
  fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
114
  @builtin(num_workgroups) nwg: vec3<u32>) {
115
- {% if source.strided %}
116
- // Strided mode has no `params` uniform (offsets come from compile-time
117
- // constants), so it folds against the compile-time STRIDED_ROWS instead of
118
- // params.rows. wg.y carries the high bits of the row past the maxComputeWorkgroupsPerDimension limit;
119
- {% if source.stage == "max" %}
120
- // the guard drops the over-dispatched tail before any rowMax[row] store.
121
- {% else %}
122
- // the guard drops the over-dispatched tail before any rowSum[row]/y[offset]
123
- // store.
124
- {% endif %}
125
- let row = wg.x + wg.y * nwg.x;
126
- if (row >= STRIDED_ROWS) { return; }
127
- {% else %}
128
  // 2D-folded row index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
129
  {% if source.stage == "max" %}
130
  // workgroup-per-dimension dispatch limit (rows > maxComputeWorkgroupsPerDimension). The guard returns
@@ -138,7 +66,6 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
138
  // nwg.x are workgroup-uniform, so the return cannot diverge across a barrier.
139
  let row = wg.x + wg.y * nwg.x;
140
  if (row >= params.rows) { return; }
141
- {% endif %}
142
  {% if source.stage == "max" %}
143
  var acc = -3.4028234663852886e38;
144
  for (var k = lid.x; k < {{ count }}; k = k + WGS) {
 
7
  {{ env.wgsl.resourceDeclarations }}
8
 
9
  const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
  {% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
11
  fn {{ name }}(value: f32, tid: u32) -> f32 {
12
  {{ buffer }}[tid] = value;
 
53
  @compute @workgroup_size(WGS)
54
  fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
55
  @builtin(num_workgroups) nwg: vec3<u32>) {
 
 
 
 
 
 
 
 
 
 
 
 
 
56
  // 2D-folded row index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
57
  {% if source.stage == "max" %}
58
  // workgroup-per-dimension dispatch limit (rows > maxComputeWorkgroupsPerDimension). The guard returns
 
66
  // nwg.x are workgroup-uniform, so the return cannot diverge across a barrier.
67
  let row = wg.x + wg.y * nwg.x;
68
  if (row >= params.rows) { return; }
 
69
  {% if source.stage == "max" %}
70
  var acc = -3.4028234663852886e38;
71
  for (var k = lid.x; k < {{ count }}; k = k + WGS) {
build/webgpu/{softmax-strided-scalar4-reduce.wgsl.jinja → softmax-strided-online-coop.wgsl.jinja} RENAMED
@@ -1,13 +1,18 @@
1
- // Four adjacent inner rows share one workgroup. Scalar storage keeps odd inner
2
- // widths and f16 valid, while the four values at every reduced-axis step are
3
- // contiguous and accumulate together as vec4<f32>.
 
 
 
 
 
 
 
4
  {% if T == "f16" %}
5
  enable f16;
6
  {% endif %}
7
  {{ env.wgsl.resourceDeclarations }}
8
-
9
- const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
10
- {% set stridedAxisUses = ["REDUCE_SIZE", "INNER_SIZE", "STRIDED_ROWS"] %}
11
  // Compiled indexing for row-wise reductions over a non-last axis. A row is one
12
  // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
13
  {% set reduce_size = source.xShape[source.axis] %}
@@ -63,74 +68,80 @@ fn row_for_offset(offset: u32) -> u32 {
63
  let inner = offset % INNER_SIZE_SAFE;
64
  return outer * INNER_SIZE_SAFE + inner;
65
  }
66
- {%- endif -%}
67
- const GROUPS_PER_OUTER: u32 = (INNER_SIZE + 3u) / 4u;
68
- const OUTER_SIZE: u32 = STRIDED_ROWS / INNER_SIZE;
69
- const TOTAL_GROUPS: u32 = OUTER_SIZE * GROUPS_PER_OUTER;
70
 
71
- var<workgroup> partial: array<vec4<f32>, WGS>;
72
- fn reduce_vec4(value: vec4<f32>, tid: u32) -> vec4<f32> {
73
- partial[tid] = value;
74
- workgroupBarrier();
75
- var stride = WGS / 2u;
76
- loop {
77
- if (stride == 0u) { break; }
78
- if (tid < stride) {
79
- {% if source.stage == "max" %}
80
- partial[tid] = max(partial[tid], partial[tid + stride]);
81
- {% else %}
82
- partial[tid] += partial[tid + stride];
83
- {% endif %}
84
- }
85
- stride /= 2u;
86
- workgroupBarrier();
87
- }
88
- return partial[0];
89
- }
90
 
91
- @compute @workgroup_size(WGS)
92
  fn main(@builtin(workgroup_id) wg: vec3<u32>,
93
  @builtin(num_workgroups) nwg: vec3<u32>,
94
  @builtin(local_invocation_id) lid: vec3<u32>) {
95
- let group = wg.x + wg.y * nwg.x;
96
- if (group >= TOTAL_GROUPS) { return; }
97
- let outerIndex = group / GROUPS_PER_OUTER;
98
- let innerBase = (group % GROUPS_PER_OUTER) * 4u;
99
- let rowBase = outerIndex * INNER_SIZE + innerBase;
100
- {% if source.stage == "max" %}
101
- var acc = vec4<f32>(-3.4028234663852886e38);
102
- {% else %}
103
- var maxima = vec4<f32>(0.0);
104
- for (var c = 0u; c < 4u; c += 1u) {
105
- if (innerBase + c < INNER_SIZE) { maxima[c] = rowMax[rowBase + c]; }
 
 
 
 
 
106
  }
107
- var acc = vec4<f32>(0.0);
108
- {% endif %}
109
- for (var k = lid.x; k < REDUCE_SIZE; k += WGS) {
110
- let base = outerIndex * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + innerBase;
111
- for (var c = 0u; c < 4u; c += 1u) {
112
- if (innerBase + c < INNER_SIZE) {
113
- {% if source.stage == "max" %}
114
- acc[c] = max(acc[c], f32(x[base + c]));
115
- {% else %}
116
- let m = maxima[c];
117
- let e = select(exp(f32(x[base + c]) - m), bitcast<f32>(bitcast<u32>(m) | 0x7fc00000u), m > 3.4028234663852886e38);
118
- y[base + c] = {{ T }}(e);
119
- acc[c] += e;
120
- {% endif %}
 
 
 
 
 
121
  }
122
  }
 
 
123
  }
124
- let result = reduce_vec4(acc, lid.x);
125
- if (lid.x == 0u) {
126
- for (var c = 0u; c < 4u; c += 1u) {
127
- if (innerBase + c < INNER_SIZE) {
128
- {% if source.stage == "max" %}
129
- rowMax[rowBase + c] = result[c];
130
- {% else %}
131
- rowSum[rowBase + c] = result[c];
132
- {% endif %}
 
 
 
 
 
 
 
 
133
  }
 
134
  }
135
  }
136
  }
 
1
+ // Cooperative online softmax/log-softmax for a non-last reduction axis. A
2
+ // workgroup owns TILE_COLS consecutive (outer, inner) rows and splits the
3
+ // reduction walk across AXIS_LANES invocations per row, so short-row shapes
4
+ // still fill the machine and every load stays coalesced across the column
5
+ // lanes regardless of whether the inner extent divides four. Per-lane online
6
+ // partials merge in workgroup memory with the same update form as the serial
7
+ // walk, so NaN and +Inf rows reproduce the single-lane kernel's outputs.
8
+ {% set operation = source.op if source.op is defined else "softmax" %}
9
+ {% set tileCols = source.tileCols %}
10
+ {% set axisLanes = (256 / tileCols) | int %}
11
  {% if T == "f16" %}
12
  enable f16;
13
  {% endif %}
14
  {{ env.wgsl.resourceDeclarations }}
15
+ {% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
 
 
16
  // Compiled indexing for row-wise reductions over a non-last axis. A row is one
17
  // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
18
  {% set reduce_size = source.xShape[source.axis] %}
 
68
  let inner = offset % INNER_SIZE_SAFE;
69
  return outer * INNER_SIZE_SAFE + inner;
70
  }
71
+ {%- endif %}
 
 
 
72
 
73
+ const TILE_COLS: u32 = {{ tileCols }}u;
74
+ const AXIS_LANES: u32 = {{ axisLanes }}u;
75
+ const F32_MAX: f32 = 3.4028234663852886e38;
76
+ {% set negInit = "-FLT_MAX" if operation == "logsoftmax" else "-F32_MAX" %}
77
+
78
+ var<workgroup> part_max: array<f32, TILE_COLS * AXIS_LANES>;
79
+ var<workgroup> part_sum: array<f32, TILE_COLS * AXIS_LANES>;
 
 
 
 
 
 
 
 
 
 
 
 
80
 
81
+ @compute @workgroup_size(TILE_COLS, AXIS_LANES, 1)
82
  fn main(@builtin(workgroup_id) wg: vec3<u32>,
83
  @builtin(num_workgroups) nwg: vec3<u32>,
84
  @builtin(local_invocation_id) lid: vec3<u32>) {
85
+ let tile = wg.x + wg.y * nwg.x;
86
+ let row = tile * TILE_COLS + lid.x;
87
+ let row_live = row < STRIDED_ROWS;
88
+
89
+ var lane_max = {{ negInit }};
90
+ var lane_sum = 0.0;
91
+ if (row_live) {
92
+ for (var k = lid.y; k < REDUCE_SIZE; k += AXIS_LANES) {
93
+ let value = f32(x[element_offset(row, k)]);
94
+ if (value > lane_max) {
95
+ lane_sum = lane_sum * exp(lane_max - value) + 1.0;
96
+ lane_max = value;
97
+ } else {
98
+ lane_sum += exp(value - lane_max);
99
+ }
100
+ }
101
  }
102
+ part_max[lid.y * TILE_COLS + lid.x] = lane_max;
103
+ part_sum[lid.y * TILE_COLS + lid.x] = lane_sum;
104
+ workgroupBarrier();
105
+
106
+ // One lane per column merges the axis partials serially, in the same
107
+ // (max, rescale, add) form as the walk above: an empty partial
108
+ // ({{ negInit }}, 0) is the identity, a NaN partial poisons the sum the
109
+ // way a NaN element does, and +Inf carries into the merged max.
110
+ if (lid.y == 0u) {
111
+ var merged_max = part_max[lid.x];
112
+ var merged_sum = part_sum[lid.x];
113
+ for (var t = 1u; t < AXIS_LANES; t += 1u) {
114
+ let m2 = part_max[t * TILE_COLS + lid.x];
115
+ let s2 = part_sum[t * TILE_COLS + lid.x];
116
+ if (m2 > merged_max) {
117
+ merged_sum = merged_sum * exp(merged_max - m2) + s2;
118
+ merged_max = m2;
119
+ } else {
120
+ merged_sum += s2 * exp(m2 - merged_max);
121
  }
122
  }
123
+ part_max[lid.x] = merged_max;
124
+ part_sum[lid.x] = merged_sum;
125
  }
126
+ workgroupBarrier();
127
+
128
+ let row_max = part_max[lid.x];
129
+ let row_sum = part_sum[lid.x];
130
+ let sum_bits = bitcast<u32>(row_sum);
131
+ let sum_is_nan = (sum_bits & 0x7f800000u) == 0x7f800000u && (sum_bits & 0x007fffffu) != 0u;
132
+ if (row_live) {
133
+ for (var k = lid.y; k < REDUCE_SIZE; k += AXIS_LANES) {
134
+ let offset = element_offset(row, k);
135
+ var out = bitcast<f32>(bitcast<u32>(row_sum) | 0x7fc00000u);
136
+ if (row_max > F32_MAX) {
137
+ // Preserve the established +Inf-row behavior: every output is NaN.
138
+ out = bitcast<f32>(bitcast<u32>(row_max) | 0x7fc00000u);
139
+ } else if (sum_is_nan) {
140
+ out = row_sum;
141
+ } else if (row_sum != 0.0) {
142
+ out = exp(f32(x[offset]) - row_max) / row_sum;
143
  }
144
+ y[offset] = {{ T }}(out);
145
  }
146
  }
147
  }
build/webgpu/softmax-strided-scalar4-normalize.wgsl.jinja DELETED
@@ -1,80 +0,0 @@
1
- {% if T == "f16" %}
2
- enable f16;
3
- {% endif %}
4
- {{ env.wgsl.resourceDeclarations }}
5
- {% set stridedAxisUses = ["row_for_offset"] %}
6
- // Compiled indexing for row-wise reductions over a non-last axis. A row is one
7
- // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
8
- {% set reduce_size = source.xShape[source.axis] %}
9
- {% set inner_size = namespace(value=1) %}
10
- {% for i in range(source.axis + 1, source.xShape | length) %}
11
- {% set inner_size.value = inner_size.value * source.xShape[i] %}
12
- {% endfor %}
13
- {% set outer_size = namespace(value=1) %}
14
- {% for i in range(source.axis) %}
15
- {% set outer_size.value = outer_size.value * source.xShape[i] %}
16
- {% endfor %}
17
- {% set usesAll = stridedAxisUses is not defined %}
18
- {% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
19
- {% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
20
- {% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
21
- {% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
22
- {% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
23
- {% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
24
- {% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
25
- {% if usesReduceSize %}
26
- const REDUCE_SIZE: u32 = {{ reduce_size }}u;
27
- {% endif %}
28
- {% if usesInnerSize %}
29
- const INNER_SIZE: u32 = {{ inner_size.value }}u;
30
- {% endif %}
31
- {% if usesStridedRows %}
32
- // Total one-per-row workgroup count (product of every non-axis dim). Lets the
33
- // row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
34
- // the over-dispatched tail without a runtime `params` uniform.
35
- const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
36
- {% endif %}
37
- {% if usesInnerSizeSafe %}
38
- const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
39
- {% endif %}
40
- {% if usesAxisStrideSafe %}
41
- const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
42
- {% endif %}
43
-
44
- {% if usesElementOffset %}
45
- // Flat offset of element k (along the reduce axis) of `row`.
46
- fn element_offset(row: u32, k: u32) -> u32 {
47
- let outer = row / INNER_SIZE_SAFE;
48
- let inner = row % INNER_SIZE_SAFE;
49
- return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
50
- }
51
-
52
- {% endif %}
53
- {% if usesRowForOffset %}
54
- // Row that owns the flat element `offset` (inverse of element_offset over
55
- // the non-axis coordinates).
56
- fn row_for_offset(offset: u32) -> u32 {
57
- let outer = offset / AXIS_STRIDE_SAFE;
58
- let inner = offset % INNER_SIZE_SAFE;
59
- return outer * INNER_SIZE_SAFE + inner;
60
- }
61
- {%- endif -%}
62
- @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
63
- fn main(@builtin(global_invocation_id) gid: vec3<u32>,
64
- @builtin(num_workgroups) nwg: vec3<u32>) {
65
- let group = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
66
- let base = group * 4u;
67
- for (var c = 0u; c < 4u; c += 1u) {
68
- let i = base + c;
69
- if (i < params.count) {
70
- let sum = rowSum[row_for_offset(i)];
71
- let bits = bitcast<u32>(sum);
72
- let isNan = (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
73
- var out = bitcast<f32>(bits | 0x7fc00000u);
74
- if (!isNan && sum != 0.0) {
75
- out = f32(y[i]) / sum;
76
- }
77
- y[i] = {{ T }}(out);
78
- }
79
- }
80
- }