sync 6fdf6301e2bb

Browse files

Files changed (14) hide show

README.md +2 -2
build/webgpu/bench.json +5 -5
build/webgpu/manifest.json +34 -38
build/webgpu/metadata.json +15 -15
build/webgpu/softmax-longrow-normalize.wgsl.jinja +1 -4
build/webgpu/softmax-longrow-stats.wgsl.jinja +7 -95
build/webgpu/softmax-normalize.wgsl.jinja +9 -21
build/webgpu/softmax-online-packed-rows.wgsl.jinja +1 -1
build/webgpu/softmax-online.wgsl.jinja +14 -41
build/webgpu/softmax-row-stage.wgsl.jinja +5 -10
build/webgpu/softmax-strided-online-coop.wgsl.jinja +9 -25
build/webgpu/softmax-strided-online-lane.wgsl.jinja +10 -8
build/webgpu/softmax-subgroup-rows.wgsl.jinja +10 -11
build/webgpu/test.json +15 -15

README.md CHANGED Viewed

@@ -51,7 +51,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`softmax-longrow-normalize.wgsl.jinja`](build/webgpu/softmax-longrow-normalize.wgsl.jinja)
 - [`softmax-longrow-stats.wgsl.jinja`](build/webgpu/softmax-longrow-stats.wgsl.jinja)
 - [`softmax-normalize.wgsl.jinja`](build/webgpu/softmax-normalize.wgsl.jinja)
@@ -65,7 +65,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`softmax-longrow-normalize.wgsl.jinja`](build/webgpu/softmax-longrow-normalize.wgsl.jinja)
 - [`softmax-longrow-stats.wgsl.jinja`](build/webgpu/softmax-longrow-stats.wgsl.jinja)
 - [`softmax-normalize.wgsl.jinja`](build/webgpu/softmax-normalize.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -110,7 +110,7 @@
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
     },
     {
-      "name": "softmax-f32-packed-rows-131072x32-healthy",
       "preset": "stress",
       "vars": { "dtype": "float32", "rows": 131072, "cols": 32 },
       "attrs": { "axis": 1 },
@@ -166,7 +166,7 @@
       "name": "softmax-f32-131072x33-online-row-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "One column past the packed-rows ceiling (<= 32 columns): 33 columns leave online_packed_rows_vec4_nosg for online_wg at one workgroup per row, about 32x the workgroup count of the same-bytes control softmax-f32-packed-rows-131072x32-healthy for a single added column."
       },
       "vars": { "dtype": "float32", "rows": 131072, "cols": 33 },
       "attrs": { "axis": 1 },
@@ -178,7 +178,7 @@
       "name": "softmax-f32-131072x64-subgroup-rows-control",
       "preset": "stress",
       "provenance": {
-        "notes": "Width-64 context row for the 33-column pathology: the first width the subgroup-per-row vec4 route admits (>= 64 columns, % 4 == 0) at the same row count. Moves twice the bytes of the 32/33-column rows, so compare per element rather than per case."
       },
       "vars": { "dtype": "float32", "rows": 131072, "cols": 64 },
       "attrs": { "axis": 1 },
@@ -202,7 +202,7 @@
       "name": "softmax-f32-64x50260-vocab-control",
       "preset": "stress",
       "provenance": {
-        "notes": "Nearest multiple-of-4 vocabulary width for softmax-f32-64x50257-vocab-odd-pathology; same bytes within 0.01%, selects online_wg_vec4."
       },
       "vars": { "dtype": "float32", "rows": 64, "cols": 50260 },
       "attrs": { "axis": 1 },
@@ -226,7 +226,7 @@
       "name": "softmax-f32-4096x200-vit-tokens-control",
       "preset": "stress",
       "provenance": {
-        "notes": "Nearest multiple-of-4 width for softmax-f32-4096x197-vit-tokens-pathology (1.5% more bytes); with subgroups it selects online_subgroup_rows_vec4, without them online_wg_vec4 is demoted below 4096 columns and online_wg runs on both sides."
       },
       "vars": { "dtype": "float32", "rows": 4096, "cols": 200 },
       "attrs": { "axis": 1 },

       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
     },
     {
+      "name": "softmax-f32-packed-rows-131072x32-control",
       "preset": "stress",
       "vars": { "dtype": "float32", "rows": 131072, "cols": 32 },
       "attrs": { "axis": 1 },
       "name": "softmax-f32-131072x33-online-row-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "One column past the packed-rows ceiling (<= 32 columns): 33 columns leave online_packed_rows_vec4_nosg for online_wg at one workgroup per row, about 32x the workgroup count of the same-bytes control softmax-f32-packed-rows-131072x32-control for a single added column."
       },
       "vars": { "dtype": "float32", "rows": 131072, "cols": 33 },
       "attrs": { "axis": 1 },
       "name": "softmax-f32-131072x64-subgroup-rows-control",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures softmax over 131,072 rows of 64 columns (a multiple of four); its 33-column sibling in this file has the same row count but moves half the bytes, so compare per element rather than per case."
       },
       "vars": { "dtype": "float32", "rows": 131072, "cols": 64 },
       "attrs": { "axis": 1 },
       "name": "softmax-f32-64x50260-vocab-control",
       "preset": "stress",
       "provenance": {
+        "notes": "The next multiple-of-four vocabulary width above the paired 50,257-column case, with nearly the same input bytes."
       },
       "vars": { "dtype": "float32", "rows": 64, "cols": 50260 },
       "attrs": { "axis": 1 },
       "name": "softmax-f32-4096x200-vit-tokens-control",
       "preset": "stress",
       "provenance": {
+        "notes": "A 200-column control for the paired 197-column ViT token case; the width changes input size by about 1.5%. Both subgroup and portable tiers must normalize every row."
       },
       "vars": { "dtype": "float32", "rows": 4096, "cols": 200 },
       "attrs": { "axis": 1 },

build/webgpu/manifest.json CHANGED Viewed

@@ -20,6 +20,7 @@
     "PACKED_ROWS_MAX_COLS": { "default": 256 },
     "PACKED_ROWS_MIN_ROW_GROUPS": { "default": 64 },
     "PACKED_ROWS_VEC4_MAX_COLS": { "default": 128 },
     "STRIDED_ONLINE_STREAMS": { "default": 2 }
   },
   "derive": {
@@ -29,7 +30,6 @@
     "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
     "lastCols": "dim(shapes.x, ranks.x - 1)",
     "lastRows": "rows(shapes.x, ranks.x - 1)",
-    "axisCols": "dim(shapes.x, axis)",
     "axisInner": "inner(shapes.x, axis)",
     "axisOuter": "outer(shapes.x, axis)",
     "axisRows": "axisOuter * axisInner",
@@ -45,37 +45,36 @@
     "rowWorkgroupSize": "min(tunables.WORKGROUP_SIZE, max(tunables.ROW_MIN_WORKGROUP, pow2ceil(ceilDiv(lastCols, tunables.ROW_COLS_PER_LANE))))",
     "subgroupRowsWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "subgroupRowsLaneVecs": "ceilDiv(lastCols / 4, device.adapterInfo.subgroupMinSize) if (device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize > 0) else 0",
     "stridedOnlineStreams": "min(4, pow2ceil(max(1, tunables.STRIDED_ONLINE_STREAMS)), max(1, ceilDiv(dim(shapes.x, axis), 256 / stridedCoopTileCols))) if dtypes.T == \"f16\" else 1"
   },
   "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$T" },
-    "rowMax": { "buffer": "storage", "elementType": "f32" },
-    "rowSum": { "buffer": "storage", "elementType": "f32" },
-    "rowMax_2": { "name": "rowMax", "buffer": "read-only-storage", "elementType": "f32" },
-    "rowSum_2": { "name": "rowSum", "buffer": "read-only-storage", "elementType": "f32" },
-    "y": { "buffer": "storage", "elementType": "$T" },
-    "params_4": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "lastRows" },
         { "name": "cols", "type": "u32", "value": "lastCols" }
       ]
     },
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "y_2": { "name": "y", "buffer": "storage", "elementType": "$vectorScalar" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "lastRows" },
         { "name": "vecCols", "type": "u32", "value": "lastCols / 4" }
       ]
     },
-    "params_6": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "count", "type": "u32", "value": "numel(shapes.x)" },
         { "name": "cols", "type": "u32", "value": "lastCols" }
       ]
     }
@@ -105,11 +104,11 @@
           "id": "block_stats",
           "name": "Softmax.LongRowBlockStats",
           "shader": "softmax-longrow-stats.wgsl.jinja",
-          "derive": { "stage": "\"block\"" },
           "bindings": [
             "x",
-            { "name": "blockMax", "buffer": "storage", "elementType": "f32" },
-            { "name": "blockSum", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
@@ -127,7 +126,7 @@
           "id": "row_stats",
           "name": "Softmax.LongRowStats",
           "shader": "softmax-longrow-stats.wgsl.jinja",
-          "derive": { "stage": "\"row\"" },
           "bindings": [
             { "name": "blockMax", "buffer": "read-only-storage", "elementType": "f32" },
             { "name": "blockSum", "buffer": "read-only-storage", "elementType": "f32" },
@@ -147,8 +146,8 @@
           "derive": { "op": "\"softmax\"" },
           "bindings": [
             "x",
-            "rowMax_2",
-            "rowSum_2",
             "y",
             { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "lastCols" }] }
           ],
@@ -176,13 +175,12 @@
           "name": "Softmax.OnlineSubgroupRowsVec4",
           "shader": "softmax-subgroup-rows.wgsl.jinja",
           "derive": { "op": "\"softmax\"" },
-          "bindings": ["x_2", "y_2", "params"],
           "dispatch": {
             "x": "min(ceilDiv(lastRows, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)",
             "y": "ceilDiv(ceilDiv(lastRows, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -190,7 +188,7 @@
       "id": "online_wg_vec4",
       "priority": 16,
       "when": ["dim(shapes.x, ranks.x - 1) % 4 == 0", "ceilDiv(rows(shapes.x, ranks.x - 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "lastAxisOk"],
-      "demoteWhen": ["dim(shapes.x, ranks.x - 1) < 4096"],
       "derive": {
         "combineSubgroups": "device.features.has(\"subgroups\") and not narrowSubgroupRange",
         "scalar": "dtypes.T",
@@ -202,9 +200,8 @@
           "id": "main",
           "name": "Softmax.OnlineWorkgroupVec4",
           "shader": "softmax-online.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": { "op": "\"softmax\"", "useVec4": "true" },
-          "bindings": ["x_2", "y_2", "params"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -227,9 +224,8 @@
           "id": "main",
           "name": "Softmax.OnlineWorkgroup",
           "shader": "softmax-online.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": { "op": "\"softmax\"", "useVec4": "false" },
-          "bindings": ["x", "y", "params_4"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -259,7 +255,7 @@
             "rowsPerWorkgroup": "rowsPerWorkgroup",
             "useVec4": "true"
           },
-          "bindings": ["x_2", "y_2", "params"],
           "dispatch": {
             "x": "min(ceilDiv(rows(shapes.x, ranks.x - 1), rowsPerWorkgroup), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -284,7 +280,7 @@
             "lanesPerRow": "lanesPerRow",
             "rowsPerWorkgroup": "rowsPerWorkgroup"
           },
-          "bindings": ["x", "y", "params_4"],
           "dispatch": {
             "x": "min(ceilDiv(rows(shapes.x, ranks.x - 1), rowsPerWorkgroup), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -310,7 +306,7 @@
           "name": "Softmax.OnlineLocalVec4NoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"softmax\"", "useVec4": "true" },
-          "bindings": ["x_2", "y_2", "params"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -331,7 +327,7 @@
           "name": "Softmax.OnlineLocalNoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"softmax\"", "useVec4": "false" },
-          "bindings": ["x", "y", "params_4"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -343,7 +339,7 @@
     {
       "id": "strided_online_coop",
       "priority": 13,
-      "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisDispatchFits"],
       "passes": [
         {
           "id": "main",
@@ -378,7 +374,7 @@
           "name": "Softmax.ReduceMax",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"max\"", "strided": "false" },
-          "bindings": ["x", "rowMax", "params_4"],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
@@ -390,7 +386,7 @@
           "name": "Softmax.ExpSum",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "strided": "false" },
-          "bindings": ["x", "rowMax_2", "y", "rowSum", "params_4"],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
@@ -402,7 +398,7 @@
           "name": "Softmax.Normalize",
           "shader": "softmax-normalize.wgsl.jinja",
           "derive": { "strided": "false" },
-          "bindings": ["y", "rowSum_2", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",

     "PACKED_ROWS_MAX_COLS": { "default": 256 },
     "PACKED_ROWS_MIN_ROW_GROUPS": { "default": 64 },
     "PACKED_ROWS_VEC4_MAX_COLS": { "default": 128 },
+    "WG_VEC4_MIN_COLS": { "default": 128 },
     "STRIDED_ONLINE_STREAMS": { "default": 2 }
   },
   "derive": {
     "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
     "lastCols": "dim(shapes.x, ranks.x - 1)",
     "lastRows": "rows(shapes.x, ranks.x - 1)",
     "axisInner": "inner(shapes.x, axis)",
     "axisOuter": "outer(shapes.x, axis)",
     "axisRows": "axisOuter * axisInner",
     "rowWorkgroupSize": "min(tunables.WORKGROUP_SIZE, max(tunables.ROW_MIN_WORKGROUP, pow2ceil(ceilDiv(lastCols, tunables.ROW_COLS_PER_LANE))))",
     "subgroupRowsWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "subgroupRowsLaneVecs": "ceilDiv(lastCols / 4, device.adapterInfo.subgroupMinSize) if (device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize > 0) else 0",
+    "stridedCoopDispatchFits": "min(ceilDiv(axisRows, stridedCoopTileCols), DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(axisRows, stridedCoopTileCols), DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "axisCols": "dim(shapes.x, axis)",
     "stridedOnlineStreams": "min(4, pow2ceil(max(1, tunables.STRIDED_ONLINE_STREAMS)), max(1, ceilDiv(dim(shapes.x, axis), 256 / stridedCoopTileCols))) if dtypes.T == \"f16\" else 1"
   },
   "bindings": {
+    "x": { "elementType": "$T" },
+    "y": { "elementType": "$T" },
+    "params__uniform": {
       "name": "params",
       "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.x)" },
         { "name": "cols", "type": "u32", "value": "lastCols" }
       ]
     },
+    "rowMax": { "buffer": "storage", "elementType": "f32" },
+    "rowSum": { "buffer": "storage", "elementType": "f32" },
+    "rowMax_f32": { "name": "rowMax", "buffer": "read-only-storage", "elementType": "f32" },
+    "rowSum_f32": { "name": "rowSum", "buffer": "read-only-storage", "elementType": "f32" },
+    "x_main": { "name": "x", "elementType": "$vectorScalar" },
+    "y_main": { "name": "y", "elementType": "$vectorScalar" },
     "params": {
       "struct": [
         { "name": "rows", "type": "u32", "value": "lastRows" },
         { "name": "vecCols", "type": "u32", "value": "lastCols / 4" }
       ]
     },
+    "params_main": {
       "name": "params",
       "struct": [
+        { "name": "rows", "type": "u32", "value": "lastRows" },
         { "name": "cols", "type": "u32", "value": "lastCols" }
       ]
     }
           "id": "block_stats",
           "name": "Softmax.LongRowBlockStats",
           "shader": "softmax-longrow-stats.wgsl.jinja",
+          "derive": { "stage": "\"block\"", "combineSubgroups": false },
           "bindings": [
             "x",
+            { "name": "blockMax", "elementType": "f32" },
+            { "name": "blockSum", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
           "id": "row_stats",
           "name": "Softmax.LongRowStats",
           "shader": "softmax-longrow-stats.wgsl.jinja",
+          "derive": { "stage": "\"row\"", "combineSubgroups": false },
           "bindings": [
             { "name": "blockMax", "buffer": "read-only-storage", "elementType": "f32" },
             { "name": "blockSum", "buffer": "read-only-storage", "elementType": "f32" },
           "derive": { "op": "\"softmax\"" },
           "bindings": [
             "x",
+            "rowMax_f32",
+            "rowSum_f32",
             "y",
             { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "lastCols" }] }
           ],
           "name": "Softmax.OnlineSubgroupRowsVec4",
           "shader": "softmax-subgroup-rows.wgsl.jinja",
           "derive": { "op": "\"softmax\"" },
+          "bindings": ["x_main", "y_main", "params"],
           "dispatch": {
             "x": "min(ceilDiv(lastRows, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)",
             "y": "ceilDiv(ceilDiv(lastRows, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)",
             "z": 1
+          }
         }
       ]
     },
       "id": "online_wg_vec4",
       "priority": 16,
       "when": ["dim(shapes.x, ranks.x - 1) % 4 == 0", "ceilDiv(rows(shapes.x, ranks.x - 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "lastAxisOk"],
+      "demoteWhen": ["dim(shapes.x, ranks.x - 1) < tunables.WG_VEC4_MIN_COLS"],
       "derive": {
         "combineSubgroups": "device.features.has(\"subgroups\") and not narrowSubgroupRange",
         "scalar": "dtypes.T",
           "id": "main",
           "name": "Softmax.OnlineWorkgroupVec4",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"softmax\"", "useVec4": "true" },
+          "bindings": ["x_main", "y_main", "params"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
           "id": "main",
           "name": "Softmax.OnlineWorkgroup",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"softmax\"", "useVec4": "false" },
+          "bindings": ["x", "y", "params_main"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "rowsPerWorkgroup": "rowsPerWorkgroup",
             "useVec4": "true"
           },
+          "bindings": ["x_main", "y_main", "params"],
           "dispatch": {
             "x": "min(ceilDiv(rows(shapes.x, ranks.x - 1), rowsPerWorkgroup), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
             "lanesPerRow": "lanesPerRow",
             "rowsPerWorkgroup": "rowsPerWorkgroup"
           },
+          "bindings": ["x", "y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv(rows(shapes.x, ranks.x - 1), rowsPerWorkgroup), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "Softmax.OnlineLocalVec4NoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"softmax\"", "useVec4": "true" },
+          "bindings": ["x_main", "y_main", "params"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
           "name": "Softmax.OnlineLocalNoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"softmax\"", "useVec4": "false" },
+          "bindings": ["x", "y", "params_main"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
     {
       "id": "strided_online_coop",
       "priority": 13,
+      "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "stridedCoopDispatchFits"],
       "passes": [
         {
           "id": "main",
           "name": "Softmax.ReduceMax",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"max\"", "strided": "false" },
+          "bindings": ["x", "rowMax", "params_main"],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
           "name": "Softmax.ExpSum",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "strided": "false" },
+          "bindings": ["x", "rowMax_f32", "y", "rowSum", "params_main"],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
           "name": "Softmax.Normalize",
           "shader": "softmax-normalize.wgsl.jinja",
           "derive": { "strided": "false" },
+          "bindings": ["y", "rowSum_f32", "params__uniform"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,29 +1,29 @@
 {
   "name": "ai.onnx.Softmax",
-  "id": "_ai_onnx_softmax_webgpu_3fa7f40",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "aU1fchujehBREF4OEI2lDx52FqWQ7KG71PLit0YBOJk=",
-      "manifest.json": "QkoJAAKNK3VdbvesprtvngDdr6fwrsUrZXOAzMepQHg=",
-      "softmax-longrow-normalize.wgsl.jinja": "ttaSzDzYfSH4XuQxo9EyAPOGcYiV0dkbUEug5e7seZg=",
-      "softmax-longrow-stats.wgsl.jinja": "prBVBse5ULq+Wl8maEVHnHtV24ispCW8XYS6ys3rzmA=",
-      "softmax-normalize.wgsl.jinja": "HMCOne5hCi5FFNWwfx/cDYOWSX9HbQbxGc64fNCYNeE=",
-      "softmax-online-packed-rows.wgsl.jinja": "+WeaU/Vi2axEcmzfW9HbLo5W0cLyYskd4QU0JCdZXPU=",
-      "softmax-online.wgsl.jinja": "HW9xDlnVtDtOuoQ9V0qdCldBRSPksX3OWjGnaAUSuts=",
-      "softmax-row-stage.wgsl.jinja": "QwJMAIkPQ6XppyMHf8uM5oBCHi/Dk2Nbur2cgTOuNGE=",
-      "softmax-strided-online-coop.wgsl.jinja": "Stf+8ldvQ37YZRZqK2eMwxFE5IRtSHI5CnDuX6cBaqg=",
-      "softmax-strided-online-lane.wgsl.jinja": "8g3/Y0koYxoMY2EwXlBYfwV/fxs8Qo9lOxJNcymmRzs=",
-      "softmax-subgroup-rows.wgsl.jinja": "wU5WO+CSSIcJt5KKD13affqBjFjrNbdvSNglO15sVDg=",
-      "test.json": "HYBGxPhprRUNZ7v+qibfypCv6gGWJpb30yH2IXo1WL8="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "longrow_split": ["softmax-longrow-normalize.wgsl.jinja", "softmax-longrow-stats.wgsl.jinja"],
       "online_subgroup_rows_vec4": ["softmax-subgroup-rows.wgsl.jinja"],

 {
   "name": "ai.onnx.Softmax",
+  "id": "_ai_onnx_softmax_webgpu_9c9600e",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "CnIUtxX0uN+sBSS6S7sXw4PhPgJN0quVT/7N7+m+kvc=",
+      "manifest.json": "bvTch1ywn3onsk2ac2P6UQN14EdCTr7IXYOQVJxkwhA=",
+      "softmax-longrow-normalize.wgsl.jinja": "GbW6zPjxGeiY5RQZTFwS5pUJ++ZTI/HQ6xtfmrssn18=",
+      "softmax-longrow-stats.wgsl.jinja": "S8w3EVvJXQavrn/79+OYm6yMZWERwlyJiZhYF8HkAj0=",
+      "softmax-normalize.wgsl.jinja": "WJcysxIi1A/NUtzd8CQXbAGTyVRH8J2D3fBJjH8oQlM=",
+      "softmax-online-packed-rows.wgsl.jinja": "uUzYkkedsLP0Jbu7C6puiDpvzvN+h5bSnsSs2lEY2xY=",
+      "softmax-online.wgsl.jinja": "/uJEBPE+2Yh0kTputbT4qMMyeB8IbkaWnRwlw+jQ8Qk=",
+      "softmax-row-stage.wgsl.jinja": "7+SK74287ITFa+SmGFdWhyjVYO0uKbQ8fN1X2eP/TbA=",
+      "softmax-strided-online-coop.wgsl.jinja": "80DwMJFFvyEWnTcPQXUmhFoFPhk1Uj9X/smbmCMn930=",
+      "softmax-strided-online-lane.wgsl.jinja": "Z/lssYWxtZYq1aNDCkA3+U4lJUwHce22W7gFaAKi/PQ=",
+      "softmax-subgroup-rows.wgsl.jinja": "aD8MbAhJiaMqEWPg64GAAg6af5rUzC2v3LRYJI3ItcI=",
+      "test.json": "hNvyxaj27rdwa+Th6fTNY8lBe80QVerINYUaKVO4NO4="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "longrow_split": ["softmax-longrow-normalize.wgsl.jinja", "softmax-longrow-stats.wgsl.jinja"],
       "online_subgroup_rows_vec4": ["softmax-subgroup-rows.wgsl.jinja"],

build/webgpu/softmax-longrow-normalize.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,6 @@
 // Final normalization pass for long-row split softmax. Each workgroup handles
 // one block of a row and re-reads the input against the completed row maximum
 // and denominator. The epilogue supports softmax and log-softmax.
-{% if T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
@@ -27,6 +24,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -37,7 +35,6 @@ fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {

 // Final normalization pass for long-row split softmax. Each workgroup handles
 // one block of a row and re-reads the input against the completed row maximum
 // and denominator. The epilogue supports softmax and log-softmax.
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {

build/webgpu/softmax-longrow-stats.wgsl.jinja CHANGED Viewed

@@ -3,9 +3,8 @@
 // The block stage assigns one workgroup to each (block, row), merges per-thread
 // online pairs, and writes one partial pair. The row stage assigns one
 // workgroup per row and folds those block partials into rowMax and rowSum.
-{% if stage == "block" and T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
@@ -31,6 +30,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -40,94 +40,9 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
-{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
-{% if combineSubgroups %}
-// Cross-subgroup merge that assumes nothing about which invocations share a
-// subgroup or how many subgroups there are: each subgroup's elected lane
-// publishes the subgroup pair in the slot at its OWN invocation index and sets
-// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
-// published slots, in ascending index order (the online (m, d) merge is not
-// float-associative, so the order is fixed), and clears the mask for the next
-// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
-// setup. Same three collectives as a single-subgroup reduce, two barriers.
-var<workgroup> partialM: array<f32, WG>;
-var<workgroup> partialD: array<f32, WG>;
-var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
-var<workgroup> combinedMD: vec2<f32>;
-// When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state). `subgroup_size` is the size of the current
-// subgroup and uniform, so the test is exact and may guard the barriers below.
-fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
-  let sgM = subgroupMax(m);
-  // A lane with no elements contributes d == 0 (exact identity). A +inf
-  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
-  // landed in d via exp(NaN); both survive the merge and are detected by the
-  // code after the reduction.
-  let sgD = subgroupAdd(d * exp_shift(m, sgM));
-  if (sgSize == WG) {
-    return vec2<f32>(sgM, sgD);
-  }
-  if (subgroupElect()) {
-    partialM[lidx] = sgM;
-    partialD[lidx] = sgD;
-    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
-  }
-  workgroupBarrier();
-  if (lidx == 0u) {
-    var accM = -FLT_MAX;
-    var accD = 0.0;
-    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
-      var bits = atomicExchange(&leaderMask[w], 0u);
-      while (bits != 0u) {
-        let slot = w * 32u + firstTrailingBit(bits);
-        bits = bits & (bits - 1u);
-        let mNew = max(accM, partialM[slot]);
-        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
-        accM = mNew;
-      }
-    }
-    combinedMD = vec2<f32>(accM, accD);
-  }
-  workgroupBarrier();
-  return combinedMD;
-}
-{% else %}
-{% set mdStreamed = mdStreams is defined %}
-{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
-{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
-{% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. Input partials occupy
-// partialM/partialD; stream s returns its merged pair in slot s * WG.
-fn combine_partials_streams(lidx: u32) {
-  workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (lidx < stride) {
-{% for s in range(mdStreams) %}
-      {
-        let slot = {{ s }}u * WG + lidx;
-        let m1 = partialM[slot];
-        let d1 = partialD[slot];
-        let m2 = partialM[slot + stride];
-        let d2 = partialD[slot + stride];
-        let mNew = max(m1, m2);
-        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
-        partialM[slot] = mNew;
-      }
-{% endfor %}
-    }
-    workgroupBarrier();
-    stride = stride / 2u;
-  }
-}
-{% else %}
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
@@ -156,9 +71,6 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   workgroupBarrier();
   return merged;
 }
-{% endif %}
-{% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
@@ -166,14 +78,14 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
 {% if stage == "block" %}
   let block = wg.x;
   let row = wg.y;
-  let rowBase = row * params.cols;
   let colStart = block * BLOCK_COLS;
-  let colEnd = min(params.cols, colStart + BLOCK_COLS);
   var m = -FLT_MAX;
   var d = 0.0;
   for (var col = colStart + lid.x; col < colEnd; col = col + WG) {
-    let v = f32(x[rowBase + col]);
     let mNew = max(m, v);
     d = d * exp_shift(m, mNew) + exp_shift(v, mNew);
     m = mNew;

 // The block stage assigns one workgroup to each (block, row), merges per-thread
 // online pairs, and writes one partial pair. The row stage assigns one
 // workgroup per row and folds those block partials into rowMax and rowSum.
+{% set biasRow = false %}
+{% set cols = "params.cols" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
+{% set mdExtent = "WG" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
   workgroupBarrier();
   return merged;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
 {% if stage == "block" %}
   let block = wg.x;
   let row = wg.y;
+  let rowBase = row * {{ cols }};
   let colStart = block * BLOCK_COLS;
+  let colEnd = min({{ cols }}, colStart + BLOCK_COLS);
   var m = -FLT_MAX;
   var d = 0.0;
   for (var col = colStart + lid.x; col < colEnd; col = col + WG) {
+    let v = {{ "f32(data[rowBase + col]) + f32(bias[biasBase + col])" if biasRow else "f32(x[rowBase + col])" }};
     let mNew = max(m, v);
     d = d * exp_shift(m, mNew) + exp_shift(v, mNew);
     m = mNew;

build/webgpu/softmax-normalize.wgsl.jinja CHANGED Viewed

@@ -1,17 +1,19 @@
-{% if T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% if op == "logsoftmax" %}
-{% set stableHelperUsage = "shifted" %}
-{% set stableUsage = stableHelperUsage if stableHelperUsage is defined else "all" %}
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
 // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
 // a zero final denominator according to their public semantics. Using -inf
 // here changes +inf-row behavior.
 const FLT_MAX: f32 = 3.4028234663852886e38;
-{% if stableUsage != "constant" %}
 fn is_finite_f32(value: f32) -> bool {
   return select(false, value <= FLT_MAX, value >= -FLT_MAX);
@@ -25,31 +27,17 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
-{%- endif %}
-{% if stableUsage == "all" %}
-fn exp_shift(value: f32, maxValue: f32) -> f32 {
-  return exp(shifted_value(value, maxValue));
-}
-{%- endif %}
 // Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
 // and a constant NaN bit pattern is rejected at WGSL const-evaluation.
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // per-axis dispatch fold width (outputs > 16.7M elements).
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (i >= params.count) {
-    return;
-  }
 {% if op == "logsoftmax" %}
   // Final element-parallel log-softmax pass. Each element rereads x and writes
   // (x - rowMax) - log(rowSum); strided mode changes only the row map.

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% if op == "logsoftmax" %}
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
 // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
 // a zero final denominator according to their public semantics. Using -inf
 // here changes +inf-row behavior.
 const FLT_MAX: f32 = 3.4028234663852886e38;
 fn is_finite_f32(value: f32) -> bool {
   return select(false, value <= FLT_MAX, value >= -FLT_MAX);
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 // Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
 // and a constant NaN bit pattern is rejected at WGSL const-evaluation.
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 {% if op == "logsoftmax" %}
   // Final element-parallel log-softmax pass. Each element rereads x and writes
   // (x - rowMax) - log(rowSum); strided mode changes only the row map.

build/webgpu/softmax-online-packed-rows.wgsl.jinja CHANGED Viewed

@@ -31,6 +31,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -41,7 +42,6 @@ fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 // Each logical row owns one contiguous LANES_PER_ROW slice. All reductions
 // are workgroup-portable; no subgroup width or subgroup feature is assumed.
 var<workgroup> partialM: array<f32, WG>;

   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 // Each logical row owns one contiguous LANES_PER_ROW slice. All reductions
 // are workgroup-portable; no subgroup width or subgroup feature is assumed.
 var<workgroup> partialM: array<f32, WG>;

build/webgpu/softmax-online.wgsl.jinja CHANGED Viewed

@@ -6,9 +6,17 @@ enable subgroups;
 // Workgroup-level online normalization over one contiguous last-axis row per
 // workgroup. Each thread keeps an online (maximum, denominator) pair over its
 // strided chunk. The workgroup merges those pairs, then a fused second sweep
-// writes the normalized result. Vectorized rows use vec4 accesses.
 //
 // `workgroupSize` fixes the merge-tree width for this specialization.
 const WG: u32 = {{ workgroupSize }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
@@ -29,6 +37,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -38,14 +47,12 @@ fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 // Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
 // mNew = max(m1, m2)
 // dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
-{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
 // Cross-subgroup merge that assumes nothing about which invocations share a
 // subgroup or how many subgroups there are: each subgroup's elected lane
@@ -98,41 +105,9 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   return combinedMD;
 }
 {% else %}
-{% set mdStreamed = mdStreams is defined %}
-{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
-{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
-{% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. Input partials occupy
-// partialM/partialD; stream s returns its merged pair in slot s * WG.
-fn combine_partials_streams(lidx: u32) {
-  workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (lidx < stride) {
-{% for s in range(mdStreams) %}
-      {
-        let slot = {{ s }}u * WG + lidx;
-        let m1 = partialM[slot];
-        let d1 = partialD[slot];
-        let m2 = partialM[slot + stride];
-        let d2 = partialD[slot + stride];
-        let mNew = max(m1, m2);
-        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
-        partialM[slot] = mNew;
-      }
-{% endfor %}
-    }
-    workgroupBarrier();
-    stride = stride / 2u;
-  }
-}
-{% else %}
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
@@ -162,8 +137,6 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
 @compute @workgroup_size({{ workgroupSize }})
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
@@ -186,7 +159,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   var d = 0.0;
 {% if useVec4 %}
   for (var c = lid.x; c < params.vecCols; c = c + WG) {
-    let v = vec4<f32>(x[rowBase + c]);
     let mNew = max(m, max(max(v.x, v.y), max(v.z, v.w)));
     d = d * exp_shift(m, mNew)
       + exp_shift(v.x, mNew) + exp_shift(v.y, mNew) + exp_shift(v.z, mNew) + exp_shift(v.w, mNew);
@@ -217,7 +190,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   // denominator from an all--infinity row produces an all-NaN row.
 {% if useVec4 %}
   for (var c = lid.x; c < params.vecCols; c = c + WG) {
-    let v = vec4<f32>(x[rowBase + c]);
     var o = vec4<f32>(nanOut);
     if (!rowIsNan) {
       o = vec4<f32>(
@@ -226,7 +199,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
         exp_shift(v.z, finalM),
         exp_shift(v.w, finalM)) / finalD;
     }
-    y[rowBase + c] = vec4<{{ scalar }}>(o);
   }
 {% else %}
   for (var c = lid.x; c < params.cols; c = c + WG) {

 // Workgroup-level online normalization over one contiguous last-axis row per
 // workgroup. Each thread keeps an online (maximum, denominator) pair over its
 // strided chunk. The workgroup merges those pairs, then a fused second sweep
+// writes the normalized result. Vectorized rows use native vec4 storage when
+// aligned, or guarded scalar storage at unaligned row strides. Both paths retain
+// f32 bias and reduction math.
 //
 // `workgroupSize` fixes the merge-tree width for this specialization.
+// `scalarVec4` selects scalar storage for fused bias rows only.
+{% macro read_vec4(c) %}
+vec4<f32>(x[rowBase + {{ c }}]){% endmacro %}
+{% macro store_vec4(value, c) %}
+    y[rowBase + {{ c }}] = vec4<{{ scalar }}>({{ value }});
+{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 // Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
 // mNew = max(m1, m2)
 // dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
 {% if combineSubgroups %}
 // Cross-subgroup merge that assumes nothing about which invocations share a
 // subgroup or how many subgroups there are: each subgroup's elected lane
   return combinedMD;
 }
 {% else %}
+{% set mdExtent = "WG" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
   return merged;
 }
 {% endif %}
 @compute @workgroup_size({{ workgroupSize }})
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
   var d = 0.0;
 {% if useVec4 %}
   for (var c = lid.x; c < params.vecCols; c = c + WG) {
+    let v = {{ read_vec4("c") }};
     let mNew = max(m, max(max(v.x, v.y), max(v.z, v.w)));
     d = d * exp_shift(m, mNew)
       + exp_shift(v.x, mNew) + exp_shift(v.y, mNew) + exp_shift(v.z, mNew) + exp_shift(v.w, mNew);
   // denominator from an all--infinity row produces an all-NaN row.
 {% if useVec4 %}
   for (var c = lid.x; c < params.vecCols; c = c + WG) {
+    let v = {{ read_vec4("c") }};
     var o = vec4<f32>(nanOut);
     if (!rowIsNan) {
       o = vec4<f32>(
         exp_shift(v.z, finalM),
         exp_shift(v.w, finalM)) / finalD;
     }
+{{ store_vec4("o", "c") }}
   }
 {% else %}
   for (var c = lid.x; c < params.cols; c = c + WG) {

build/webgpu/softmax-row-stage.wgsl.jinja CHANGED Viewed

@@ -1,14 +1,11 @@
 // First and second passes of the three-pass softmax, with one workgroup per
 // row. The max stage reduces the row maximum. The exp_sum stage writes
 // exp(x - rowMax) and reduces the row denominator.
-{% if T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
-{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
-fn {{ name }}(value: f32, tid: u32) -> f32 {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
@@ -34,22 +31,20 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
-{% if trailingBarrier %}
   workgroupBarrier();
-{% endif %}
   return reduced;
 }
 {% endmacro %}
 var<workgroup> partial: array<f32, WGS>;
 {% if stage == "max" %}
 {{ wgsl_tree_reduce_f32("reduce_max_f32", "max", "partial", "WGS") }}
 {% else %}
 {{ wgsl_tree_reduce_f32("reduce_add_f32", "add", "partial", "WGS") }}
 {% endif %}
-{% set count = "REDUCE_SIZE" if strided else "params.cols" %}
-{% set offset_of = "element_offset(row, k)" if strided else "row * params.cols + k" %}
 @compute @workgroup_size(WGS)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits once the row count passes the

 // First and second passes of the three-pass softmax, with one workgroup per
 // row. The max stage reduces the row maximum. The exp_sum stage writes
 // exp(x - rowMax) and reduces the row denominator.
 {{ env.wgsl.resourceDeclarations }}
 const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
   workgroupBarrier();
   return reduced;
 }
 {% endmacro %}
 var<workgroup> partial: array<f32, WGS>;
 {% if stage == "max" %}
 {{ wgsl_tree_reduce_f32("reduce_max_f32", "max", "partial", "WGS") }}
 {% else %}
 {{ wgsl_tree_reduce_f32("reduce_add_f32", "add", "partial", "WGS") }}
 {% endif %}
+{% set count = "params.cols" %}
+{% set offset_of = "row * params.cols + k" %}
 @compute @workgroup_size(WGS)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits once the row count passes the

build/webgpu/softmax-strided-online-coop.wgsl.jinja CHANGED Viewed

@@ -6,11 +6,7 @@
 // online streams. Their partials and the per-lane partials merge in workgroup memory with the same update form as the serial
 // walk, so NaN and +Inf rows reproduce the single-lane kernel's outputs.
 {% set operation = op if op is defined else "softmax" %}
-{% set tileCols = tileCols %}
 {% set axisLanes = (256 / tileCols) | int %}
-{% if T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
 // Compiled indexing for row-wise reductions over a non-last axis. A row is one
@@ -50,17 +46,17 @@ const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1
 {% if usesAxisStrideSafe %}
 const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
 {% endif %}
 {% if usesElementOffset %}
 // Flat offset of element k (along the reduce axis) of `row`.
 fn element_offset(row: u32, k: u32) -> u32 {
   let outer = row / INNER_SIZE_SAFE;
   let inner = row % INNER_SIZE_SAFE;
   return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
 }
 {% endif %}
 {% if usesRowForOffset %}
 // Row that owns the flat element `offset` (inverse of element_offset over
 // the non-axis coordinates).
 fn row_for_offset(offset: u32) -> u32 {
@@ -68,17 +64,14 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif -%}
 {% if operation == "logsoftmax" %}
-{% set stableHelperUsage = "shifted" %}
-{% set stableUsage = stableHelperUsage if stableHelperUsage is defined else "all" %}
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
 // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
 // a zero final denominator according to their public semantics. Using -inf
 // here changes +inf-row behavior.
 const FLT_MAX: f32 = 3.4028234663852886e38;
-{% if stableUsage != "constant" %}
 fn is_finite_f32(value: f32) -> bool {
   return select(false, value <= FLT_MAX, value >= -FLT_MAX);
@@ -92,23 +85,8 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
-{%- endif %}
-{% if stableUsage == "all" %}
-fn exp_shift(value: f32, maxValue: f32) -> f32 {
-  return exp(shifted_value(value, maxValue));
-}
-{%- endif %}
-{% endif %}
-const TILE_COLS: u32 = {{ tileCols }}u;
-const AXIS_LANES: u32 = {{ axisLanes }}u;
-{% if operation != "logsoftmax" %}
-const F32_MAX: f32 = 3.4028234663852886e38;
 {% endif %}
 {% set negInit = "-FLT_MAX" if operation == "logsoftmax" else "-F32_MAX" %}
 {% macro update_partial(maxName, sumName, valueName, sumValue) %}
       if ({{ valueName }} > {{ maxName }}) {
         {{ sumName }} = {{ sumName }} * exp({{ maxName }} - {{ valueName }}) + {{ sumValue }};
@@ -122,6 +100,12 @@ const F32_MAX: f32 = 3.4028234663852886e38;
       }
 {% endmacro %}
 var<workgroup> part_max: array<f32, TILE_COLS * AXIS_LANES>;
 var<workgroup> part_sum: array<f32, TILE_COLS * AXIS_LANES>;

 // online streams. Their partials and the per-lane partials merge in workgroup memory with the same update form as the serial
 // walk, so NaN and +Inf rows reproduce the single-lane kernel's outputs.
 {% set operation = op if op is defined else "softmax" %}
 {% set axisLanes = (256 / tileCols) | int %}
 {{ env.wgsl.resourceDeclarations }}
 {% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
 // Compiled indexing for row-wise reductions over a non-last axis. A row is one
 {% if usesAxisStrideSafe %}
 const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
 {% endif %}
 {% if usesElementOffset %}
 // Flat offset of element k (along the reduce axis) of `row`.
 fn element_offset(row: u32, k: u32) -> u32 {
   let outer = row / INNER_SIZE_SAFE;
   let inner = row % INNER_SIZE_SAFE;
   return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
 }
 {% endif %}
 {% if usesRowForOffset %}
 // Row that owns the flat element `offset` (inverse of element_offset over
 // the non-axis coordinates).
 fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
+{% endif %}
 {% if operation == "logsoftmax" %}
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
 // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
 // a zero final denominator according to their public semantics. Using -inf
 // here changes +inf-row behavior.
 const FLT_MAX: f32 = 3.4028234663852886e38;
 fn is_finite_f32(value: f32) -> bool {
   return select(false, value <= FLT_MAX, value >= -FLT_MAX);
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 {% endif %}
 {% set negInit = "-FLT_MAX" if operation == "logsoftmax" else "-F32_MAX" %}
 {% macro update_partial(maxName, sumName, valueName, sumValue) %}
       if ({{ valueName }} > {{ maxName }}) {
         {{ sumName }} = {{ sumName }} * exp({{ maxName }} - {{ valueName }}) + {{ sumValue }};
       }
 {% endmacro %}
+const TILE_COLS: u32 = {{ tileCols }}u;
+const AXIS_LANES: u32 = {{ axisLanes }}u;
+{% if operation != "logsoftmax" %}
+const F32_MAX: f32 = 3.4028234663852886e38;
+{% endif %}
 var<workgroup> part_max: array<f32, TILE_COLS * AXIS_LANES>;
 var<workgroup> part_sum: array<f32, TILE_COLS * AXIS_LANES>;

build/webgpu/softmax-strided-online-lane.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,12 @@
 // Scratch-free online softmax for a non-last reduction axis. One invocation
 // owns one (outer, inner) row. At each reduction step, neighboring invocations
 // access consecutive inner positions. Row statistics remain private.
-{% if T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
 // Compiled indexing for row-wise reductions over a non-last axis. A row is one
@@ -43,17 +46,17 @@ const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1
 {% if usesAxisStrideSafe %}
 const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
 {% endif %}
 {% if usesElementOffset %}
 // Flat offset of element k (along the reduce axis) of `row`.
 fn element_offset(row: u32, k: u32) -> u32 {
   let outer = row / INNER_SIZE_SAFE;
   let inner = row % INNER_SIZE_SAFE;
   return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
 }
 {% endif %}
 {% if usesRowForOffset %}
 // Row that owns the flat element `offset` (inverse of element_offset over
 // the non-axis coordinates).
 fn row_for_offset(offset: u32) -> u32 {
@@ -61,15 +64,14 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif %}
 const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 const F32_MAX: f32 = 3.4028234663852886e38;
 @compute @workgroup_size(WGS)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WGS;
-  if (row >= STRIDED_ROWS) { return; }
   var row_max = -F32_MAX;
   var row_sum = 0.0;

 // Scratch-free online softmax for a non-last reduction axis. One invocation
 // owns one (outer, inner) row. At each reduction step, neighboring invocations
 // access consecutive inner positions. Row statistics remain private.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
 // Compiled indexing for row-wise reductions over a non-last axis. A row is one
 {% if usesAxisStrideSafe %}
 const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
 {% endif %}
 {% if usesElementOffset %}
 // Flat offset of element k (along the reduce axis) of `row`.
 fn element_offset(row: u32, k: u32) -> u32 {
   let outer = row / INNER_SIZE_SAFE;
   let inner = row % INNER_SIZE_SAFE;
   return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
 }
 {% endif %}
 {% if usesRowForOffset %}
 // Row that owns the flat element `offset` (inverse of element_offset over
 // the non-axis coordinates).
 fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
+{% endif %}
 const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 const F32_MAX: f32 = 3.4028234663852886e38;
 @compute @workgroup_size(WGS)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WGS", "row", "STRIDED_ROWS", guardInline=true) }}
   var row_max = -F32_MAX;
   var row_sum = 0.0;

build/webgpu/softmax-subgroup-rows.wgsl.jinja CHANGED Viewed

@@ -28,6 +28,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -37,7 +38,6 @@ fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 var<workgroup> sgLaneClaims: atomic<u32>;
 struct SubgroupLogicalLanes {
@@ -47,28 +47,27 @@ struct SubgroupLogicalLanes {
   count: u32,
 }
-fn subgroup_logical_lanes() -> SubgroupLogicalLanes {
-  let rank = subgroupExclusiveAdd(1u);
-  let count = subgroupAdd(1u);
-  // Every lane performs the atomic so no collective follows a lane guard (a
-  // subgroup op after a closed `if (rank == 0u)` block is the reconvergence
-  // hazard the render gate flags): only the rank-0 lane adds its subgroup's
-  // claim, every other lane adds 0 and discards its snapshot. The rank-0 lane is
-  // the lowest active lane, which is the lane `subgroupBroadcastFirst` reads.
   let ticket = atomicAdd(&sgLaneClaims, select(0u, count | (1u << 16u), rank == 0u));
   let claim = subgroupBroadcastFirst(ticket);
   return SubgroupLogicalLanes((claim & 0xffffu) + rank, claim >> 16u, rank, count);
 }
 @compute @workgroup_size({{ workgroupSize }})
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(subgroup_size) sgSize: u32) {
   // Rows tile the folded workgroup grid; the lanes of one subgroup share a row.
   // Which subgroup owns which row, and which lane holds which columns, come from
   // the logical lane coordinates, so no assumption is made about which
   // invocations share a subgroup.
-  let L = subgroup_logical_lanes();
   let sgLane = L.rank;
   let rowsPerWorkgroup = WG / sgSize;
   let row = (wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u) * rowsPerWorkgroup + L.ord;

   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 var<workgroup> sgLaneClaims: atomic<u32>;
 struct SubgroupLogicalLanes {
   count: u32,
 }
+fn subgroup_logical_lanes(rank: u32, count: u32) -> SubgroupLogicalLanes {
+  // Every lane performs the atomic so no collective follows a lane guard: a
+  // subgroup op after a closed `if (rank == 0u)` block is a reconvergence
+  // hazard, since every subgroup operation must be reached in uniform control
+  // flow. Only the lane-0 lane adds its subgroup's claim, every other lane
+  // adds 0 and discards its snapshot. Lane 0 is the lowest lane, which is the
+  // lane `subgroupBroadcastFirst` reads.
   let ticket = atomicAdd(&sgLaneClaims, select(0u, count | (1u << 16u), rank == 0u));
   let claim = subgroupBroadcastFirst(ticket);
   return SubgroupLogicalLanes((claim & 0xffffu) + rank, claim >> 16u, rank, count);
 }
 @compute @workgroup_size({{ workgroupSize }})
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(subgroup_invocation_id) sgLaneId: u32,
         @builtin(subgroup_size) sgSize: u32) {
   // Rows tile the folded workgroup grid; the lanes of one subgroup share a row.
   // Which subgroup owns which row, and which lane holds which columns, come from
   // the logical lane coordinates, so no assumption is made about which
   // invocations share a subgroup.
+  let L = subgroup_logical_lanes(sgLaneId, sgSize);
   let sgLane = L.rank;
   let rowsPerWorkgroup = WG / sgSize;
   let row = (wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u) * rowsPerWorkgroup + L.ord;

build/webgpu/test.json CHANGED Viewed

@@ -17,7 +17,7 @@
       "outputs": { "y": { "dtype": "float32", "shape": [65540, 40], "tolerance": 0.0001, "relTolerance": 0.0001 } }
     },
     {
-      "name": "dispatch_cliff_stable_3pass_257x65535",
       "attrs": { "axis": -1 },
       "inputs": {
         "x": {
@@ -28,7 +28,7 @@
       },
       "outputs": { "y": { "dtype": "float32", "shape": [257, 65535], "tolerance": 1e-8, "relTolerance": 0.0001 } },
       "provenance": {
-        "notes": "A 65,535-column row forces the three-pass stable route past the single-dispatch cliff. The tolerance is scaled to the roughly 1/65,535 outputs so denominator rescaling or a missing partial is observable."
       }
     },
     {
@@ -80,7 +80,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/softmax_test.cc",
         "test": "SoftmaxOperator.LargeNumber",
-        "notes": "An 87.5-point logit gap leaves a valid positive subnormal probability tail in ORT CPU; stable softmax should not flush it to zero."
       },
       "attrs": { "axis": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, -87.5] } } },
@@ -517,7 +517,7 @@
       "outputs": { "y": { "dtype": "float32", "shape": [1, 1025], "tolerance": 0.000001 } }
     },
     {
-      "name": "ort_gh15949_rank4_axis1_singleton_spatial",
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/softmax_test.cc",
         "test": "SoftmaxOperator.GH15949_regression_test"
@@ -748,7 +748,7 @@
     {
       "name": "rank7_axis1_strided",
       "provenance": {
-        "notes": "Exercises strided Softmax over axis 1 of a rank-7 tensor, with expectations from ORT CPU."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -891,7 +891,7 @@
     {
       "name": "packed_rows_tail_axis1_33x32",
       "provenance": {
-        "notes": "One row beyond a full 32-row packed workgroup locks the grid-stride tail guard and requires every inactive logical row to keep participating in workgroup barriers."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -936,7 +936,7 @@
     {
       "name": "strided_scalar4_tail_f32_rank3_axis1_inner6",
       "provenance": {
-        "notes": "Compact correctness lock for coarsening four adjacent strided rows when the inner dimension has a two-lane tail."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -961,7 +961,7 @@
       "outputs": { "y": { "dtype": "float16", "shape": [2, 17, 6], "tolerance": 0.003, "relTolerance": 0.003 } }
     },
     {
-      "name": "strided_empty_inner_f16_forces_plain_3pass",
       "provenance": {
         "notes": "A float16 tensor with a zero-length inner dimension exercises the plain three-pass strided implementation, since scalar-four processing requires a nonempty inner dimension and vectorized processing requires float32."
       },
@@ -1080,7 +1080,7 @@
     {
       "name": "online_stream_axis31_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1095,7 +1095,7 @@
     {
       "name": "online_stream_axis32_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1110,7 +1110,7 @@
     {
       "name": "online_stream_axis33_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1125,7 +1125,7 @@
     {
       "name": "online_stream_axis64_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1140,7 +1140,7 @@
     {
       "name": "online_stream_axis65_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1155,7 +1155,7 @@
     {
       "name": "online_stream_columns_tail_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1190,7 +1190,7 @@
     {
       "name": "online_stream_four_way_tail_f16",
       "provenance": {
-        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding; existing fixtures and tolerances are unchanged."
       },
       "attrs": { "axis": 1 },
       "inputs": {

       "outputs": { "y": { "dtype": "float32", "shape": [65540, 40], "tolerance": 0.0001, "relTolerance": 0.0001 } }
     },
     {
+      "name": "wide_rows_257x65535",
       "attrs": { "axis": -1 },
       "inputs": {
         "x": {
       },
       "outputs": { "y": { "dtype": "float32", "shape": [257, 65535], "tolerance": 1e-8, "relTolerance": 0.0001 } },
       "provenance": {
+        "notes": "A 257-row softmax over 65,535 columns; the tolerance, scaled to the roughly 1/65,535 outputs, exposes a dropped column or an incorrectly rescaled denominator."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/softmax_test.cc",
         "test": "SoftmaxOperator.LargeNumber",
+        "notes": "An 87.5-point logit gap leaves a valid positive subnormal probability tail in ONNX Runtime's CPU provider; stable softmax should not flush it to zero."
       },
       "attrs": { "axis": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, -87.5] } } },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 1025], "tolerance": 0.000001 } }
     },
     {
+      "name": "ort_rank4_axis1_singleton_spatial",
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/softmax_test.cc",
         "test": "SoftmaxOperator.GH15949_regression_test"
     {
       "name": "rank7_axis1_strided",
       "provenance": {
+        "notes": "Exercises strided Softmax over axis 1 of a rank-7 tensor, with expectations from ONNX Runtime's CPU provider."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "packed_rows_tail_axis1_33x32",
       "provenance": {
+        "notes": "33 rows of 32 columns leave one row past a multiple of 32; that final row must normalize independently of the others."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "strided_scalar4_tail_f32_rank3_axis1_inner6",
       "provenance": {
+        "notes": "Four adjacent strided rows with a two-element inner tail check normalized output independently per row."
       },
       "attrs": { "axis": 1 },
       "inputs": {
       "outputs": { "y": { "dtype": "float16", "shape": [2, 17, 6], "tolerance": 0.003, "relTolerance": 0.003 } }
     },
     {
+      "name": "strided_empty_inner_f16_plain_3pass",
       "provenance": {
         "notes": "A float16 tensor with a zero-length inner dimension exercises the plain three-pass strided implementation, since scalar-four processing requires a nonempty inner dimension and vectorized processing requires float32."
       },
     {
       "name": "online_stream_axis31_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "online_stream_axis32_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "online_stream_axis33_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "online_stream_axis64_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "online_stream_axis65_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "online_stream_columns_tail_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "online_stream_four_way_tail_f16",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {