sync 6fdf6301e2bb

Browse files

Files changed (15) hide show

README.md +2 -2
build/webgpu/bench.json +5 -11
build/webgpu/manifest.json +38 -49
build/webgpu/metadata.json +16 -16
build/webgpu/softmax-longrow-normalize.wgsl.jinja +0 -4
build/webgpu/softmax-longrow-stats.wgsl.jinja +7 -95
build/webgpu/softmax-normalize.wgsl.jinja +12 -14
build/webgpu/softmax-online-packed-rows.wgsl.jinja +1 -1
build/webgpu/softmax-online.wgsl.jinja +14 -41
build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja +15 -12
build/webgpu/softmax-row-stage.wgsl.jinja +6 -11
build/webgpu/softmax-strided-online-coop.wgsl.jinja +7 -13
build/webgpu/softmax-strided-packed4-tail.wgsl.jinja +13 -5
build/webgpu/softmax-subgroup-rows.wgsl.jinja +10 -11
build/webgpu/test.json +11 -13

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)
@@ -66,7 +66,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

@@ -86,10 +86,7 @@
     {
       "name": "logsoftmax-f32-strided-axis1-64x2048x130-packed4-control",
       "preset": "stress",
-      "provenance": {
-        "source": "strided-axis selector crossover",
-        "notes": "At 8,320 independent strided rows with an inner extent of 130, the selector uses the cooperative packed-four reduction path."
-      },
       "vars": { "dtype": "float32", "outer": 64, "axisDim": 2048, "inner": 130 },
       "attrs": { "axis": 1 },
       "inputs": { "x": { "shape": [64, 2048, 130], "dtype": "float32", "dist": "normal", "seed": 414, "scale": 2 } },
@@ -104,8 +101,7 @@
       "name": "logsoftmax-f32-strided-axis1-96x2048x130-crossover",
       "preset": "stress",
       "provenance": {
-        "source": "strided-axis selector crossover",
-        "notes": "A large-inner tail with 12,480 independent strided rows measures the selector boundary between the serial-lane capacity path and cooperative packed-four reduction."
       },
       "vars": { "dtype": "float32", "outer": 96, "axisDim": 2048, "inner": 130 },
       "attrs": { "axis": 1 },
@@ -136,9 +132,7 @@
     },
     {
       "name": "logsoftmax-f32-8192x512-midwidth",
-      "provenance": {
-        "notes": "Covers the 512-column midpoint between the 128- and 2048-column cases and probes the width term in online_wg_vec4's selector."
-      },
       "preset": "smoke",
       "vars": { "dtype": "float32", "rows": 8192, "cols": 512 },
       "attrs": { "axis": 1 },
@@ -150,7 +144,7 @@
       "name": "logsoftmax-f32-64x50257-vocab-odd-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "GPT-2 vocabulary width (50257 % 4 != 0): online_wg scalar because online_wg_vec4 needs cols % 4 == 0 (it is otherwise eligible here, cols >= 1024 with subgroups) and the subgroup-rows route is out of range (> 8 vec4 per lane). Control: logsoftmax-f32-64x50260-vocab-control."
       },
       "vars": { "dtype": "float32", "rows": 64, "cols": 50257 },
       "attrs": { "axis": 1 },
@@ -162,7 +156,7 @@
       "name": "logsoftmax-f32-64x50260-vocab-control",
       "preset": "stress",
       "provenance": {
-        "notes": "Nearest multiple-of-4 vocabulary width for logsoftmax-f32-64x50257-vocab-odd-pathology; same bytes within 0.01%, selects online_wg_vec4."
       },
       "vars": { "dtype": "float32", "rows": 64, "cols": 50260 },
       "attrs": { "axis": 1 },

     {
       "name": "logsoftmax-f32-strided-axis1-64x2048x130-packed4-control",
       "preset": "stress",
+      "provenance": { "notes": "Measures LogSoftmax over 8,320 independent strided rows with an inner extent of 130." },
       "vars": { "dtype": "float32", "outer": 64, "axisDim": 2048, "inner": 130 },
       "attrs": { "axis": 1 },
       "inputs": { "x": { "shape": [64, 2048, 130], "dtype": "float32", "dist": "normal", "seed": 414, "scale": 2 } },
       "name": "logsoftmax-f32-strided-axis1-96x2048x130-crossover",
       "preset": "stress",
       "provenance": {
+        "notes": "A large inner tail across 12,480 independent strided rows measures LogSoftmax on noncontiguous data."
       },
       "vars": { "dtype": "float32", "outer": 96, "axisDim": 2048, "inner": 130 },
       "attrs": { "axis": 1 },
     },
     {
       "name": "logsoftmax-f32-8192x512-midwidth",
+      "provenance": { "notes": "A 512-column case measures LogSoftmax between the 128- and 2048-column shapes." },
       "preset": "smoke",
       "vars": { "dtype": "float32", "rows": 8192, "cols": 512 },
       "attrs": { "axis": 1 },
       "name": "logsoftmax-f32-64x50257-vocab-odd-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures log-softmax over GPT-2's 50,257-token vocabulary width, which is not a multiple of four; its 50,260-wide sibling in this file is the nearest four-aligned control at nearly the same byte count."
       },
       "vars": { "dtype": "float32", "rows": 64, "cols": 50257 },
       "attrs": { "axis": 1 },
       "name": "logsoftmax-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 },

build/webgpu/manifest.json CHANGED Viewed

@@ -29,7 +29,6 @@
     "axisRows": "axisOuter * axisInner",
     "longrowBlocks": "ceilDiv(lastCols, tunables.BLOCK_COLS)",
     "lastRowScratchFits": "4 * lastRows <= storageBufferLimit",
-    "axisRowScratchFits": "4 * axisRows <= storageBufferLimit",
     "longrowBlockScratchFits": "4 * lastRows * longrowBlocks <= storageBufferLimit",
     "lastDispatchFits": "ceilDiv(lastRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "axisDispatchFits": "ceilDiv(axisRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
@@ -40,37 +39,33 @@
     "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 == \"f32\" else 1"
   },
   "bindings": {
     "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" },
-    "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" }
       ]
     },
-    "x_3": { "name": "x", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "y_3": { "name": "y", "buffer": "storage", "elementType": "$scalar" },
-    "params_5": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "lastRows" },
         { "name": "cols", "type": "u32", "value": "lastCols" }
       ]
-    },
-    "x_4": { "name": "x", "buffer": "read-only-storage", "elementType": "vec4<f32>" },
-    "params_7": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.x)" }]
     }
   },
   "variants": [
@@ -97,11 +92,11 @@
           "id": "block_stats",
           "name": "LogSoftmax.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": [
@@ -119,7 +114,7 @@
           "id": "row_stats",
           "name": "LogSoftmax.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" },
@@ -139,8 +134,8 @@
           "derive": { "op": "\"logsoftmax\"" },
           "bindings": [
             "x",
-            "rowMax_2",
-            "rowSum_2",
             "y",
             { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "lastCols" }] }
           ],
@@ -168,13 +163,12 @@
           "name": "LogSoftmax.OnlineSubgroupRowsVec4",
           "shader": "softmax-subgroup-rows.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"" },
-          "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"
         }
       ]
     },
@@ -193,9 +187,8 @@
           "id": "main",
           "name": "LogSoftmax.OnlineWorkgroupVec4",
           "shader": "softmax-online.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": { "op": "\"logsoftmax\"", "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)",
@@ -218,9 +211,8 @@
           "id": "main",
           "name": "LogSoftmax.OnlineWorkgroup",
           "shader": "softmax-online.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": { "op": "\"logsoftmax\"", "useVec4": "false" },
-          "bindings": ["x_3", "y_3", "params_5"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -250,7 +242,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,
@@ -275,7 +267,7 @@
             "lanesPerRow": "lanesPerRow",
             "rowsPerWorkgroup": "rowsPerWorkgroup"
           },
-          "bindings": ["x_3", "y_3", "params_5"],
           "dispatch": {
             "x": "min(ceilDiv(rows(shapes.x, ranks.x - 1), rowsPerWorkgroup), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -301,7 +293,7 @@
           "name": "LogSoftmax.OnlineLocalVec4NoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"", "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)",
@@ -322,7 +314,7 @@
           "name": "LogSoftmax.OnlineLocalNoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"", "useVec4": "false" },
-          "bindings": ["x_3", "y_3", "params_5"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -345,7 +337,7 @@
           "name": "LogSoftmax.ReduceMax",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"max\"", "strided": "false" },
-          "bindings": ["x", "rowMax", "params_5"],
           "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)",
@@ -357,7 +349,7 @@
           "name": "LogSoftmax.ExpSum",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "strided": "false" },
-          "bindings": ["x", "rowMax_2", "y", "rowSum", "params_5"],
           "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)",
@@ -371,8 +363,8 @@
           "derive": { "strided": "false", "op": "\"logsoftmax\"" },
           "bindings": [
             "x",
-            "rowMax_2",
-            "rowSum_2",
             "y",
             {
               "name": "params",
@@ -394,7 +386,6 @@
       "id": "strided_3pass_vec4",
       "priority": 14,
       "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 == 0", "axisDispatchFits", "axisRowScratchFits"],
-      "derive": { "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" },
       "intermediates": [
         { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" },
         { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" }
@@ -405,7 +396,7 @@
           "name": "LogSoftmax.StridedReduceMaxVec4",
           "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
           "derive": { "stage": "\"max\"", "xShape": "shapes.x", "axisSpec": "axis" },
-          "bindings": ["x_4", { "name": "rowMax", "buffer": "storage", "elementType": "vec4<f32>" }],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, axis) / 4, 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, axis) / 4, 65535)",
@@ -418,10 +409,10 @@
           "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "xShape": "shapes.x", "axisSpec": "axis" },
           "bindings": [
-            "x_4",
             { "name": "rowMax", "buffer": "read-only-storage", "elementType": "vec4<f32>" },
             { "arg": "y", "elementType": "vec4<f32>" },
-            { "name": "rowSum", "buffer": "storage", "elementType": "vec4<f32>" }
           ],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, axis) / 4, 65535)",
@@ -434,7 +425,7 @@
           "name": "LogSoftmax.StridedNormalize",
           "shader": "softmax-normalize.wgsl.jinja",
           "derive": { "strided": "true", "xShape": "shapes.x", "axisSpec": "axis", "op": "\"logsoftmax\"" },
-          "bindings": ["x", "rowMax_2", "rowSum_2", "y", "params_7"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -446,7 +437,7 @@
     {
       "id": "strided_online_coop",
       "priority": 19,
-      "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisThreadDispatchFits"],
       "passes": [
         {
           "id": "main",
@@ -455,9 +446,9 @@
           "derive": {
             "xShape": "shapes.x",
             "axisSpec": "axis",
-            "op": "\"logsoftmax\"",
             "tileCols": "stridedCoopTileCols",
-            "onlineStreams": "stridedOnlineStreams"
           },
           "bindings": ["x", "y"],
           "dispatch": {
@@ -472,7 +463,6 @@
       "id": "strided_online_packed4_tail",
       "priority": 18,
       "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 4", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 != 0", "dim(shapes.x, axis) < 2048 or axisRows < tunables.STRIDED_CAPACITY_MIN_ROWS", "axisDispatchFits"],
-      "derive": { "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" },
       "passes": [
         {
           "id": "main",
@@ -493,7 +483,6 @@
       "priority": 12,
       "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisDispatchFits", "axisRowScratchFits"],
       "demoteWhen": ["dtypes.T == \"f32\" and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 4 and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 != 0"],
-      "derive": { "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" },
       "intermediates": [
         { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" },
         { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" }
@@ -516,7 +505,7 @@
           "name": "LogSoftmax.StridedExpSum",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "strided": "true", "xShape": "shapes.x", "axisSpec": "axis" },
-          "bindings": ["x", "rowMax_2", "y", "rowSum"],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, axis), 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, axis), 65535)",
@@ -528,7 +517,7 @@
           "name": "LogSoftmax.StridedNormalize",
           "shader": "softmax-normalize.wgsl.jinja",
           "derive": { "strided": "true", "xShape": "shapes.x", "axisSpec": "axis", "op": "\"logsoftmax\"" },
-          "bindings": ["x", "rowMax_2", "rowSum_2", "y", "params_7"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",

     "axisRows": "axisOuter * axisInner",
     "longrowBlocks": "ceilDiv(lastCols, tunables.BLOCK_COLS)",
     "lastRowScratchFits": "4 * lastRows <= storageBufferLimit",
     "longrowBlockScratchFits": "4 * lastRows * longrowBlocks <= storageBufferLimit",
     "lastDispatchFits": "ceilDiv(lastRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "axisDispatchFits": "ceilDiv(axisRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "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)",
+    "axisRowScratchFits": "4 * axisRows <= storageBufferLimit",
     "stridedOnlineStreams": "min(4, pow2ceil(max(1, tunables.STRIDED_ONLINE_STREAMS)), max(1, ceilDiv(dim(shapes.x, axis), 256 / stridedCoopTileCols))) if dtypes.T == \"f32\" else 1"
   },
   "bindings": {
+    "x_scalar": { "name": "x", "elementType": "$scalar" },
+    "y_scalar": { "name": "y", "elementType": "$scalar" },
+    "x_exp_sum": { "name": "x", "elementType": "vec4<f32>" },
+    "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.x)" }] },
     "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" }
       ]
     }
   },
   "variants": [
           "id": "block_stats",
           "name": "LogSoftmax.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": "LogSoftmax.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": "\"logsoftmax\"" },
           "bindings": [
             "x",
+            "rowMax_f32",
+            "rowSum_f32",
             "y",
             { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "lastCols" }] }
           ],
           "name": "LogSoftmax.OnlineSubgroupRowsVec4",
           "shader": "softmax-subgroup-rows.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"" },
+          "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": "main",
           "name": "LogSoftmax.OnlineWorkgroupVec4",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"", "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": "LogSoftmax.OnlineWorkgroup",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"", "useVec4": "false" },
+          "bindings": ["x_scalar", "y_scalar", "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_scalar", "y_scalar", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv(rows(shapes.x, ranks.x - 1), rowsPerWorkgroup), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "LogSoftmax.OnlineLocalVec4NoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"", "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": "LogSoftmax.OnlineLocalNoSg",
           "shader": "softmax-online.wgsl.jinja",
           "derive": { "op": "\"logsoftmax\"", "useVec4": "false" },
+          "bindings": ["x_scalar", "y_scalar", "params_main"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
           "name": "LogSoftmax.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": "LogSoftmax.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)",
           "derive": { "strided": "false", "op": "\"logsoftmax\"" },
           "bindings": [
             "x",
+            "rowMax_f32",
+            "rowSum_f32",
             "y",
             {
               "name": "params",
       "id": "strided_3pass_vec4",
       "priority": 14,
       "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 == 0", "axisDispatchFits", "axisRowScratchFits"],
       "intermediates": [
         { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" },
         { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" }
           "name": "LogSoftmax.StridedReduceMaxVec4",
           "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
           "derive": { "stage": "\"max\"", "xShape": "shapes.x", "axisSpec": "axis" },
+          "bindings": ["x_exp_sum", { "name": "rowMax", "elementType": "vec4<f32>" }],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, axis) / 4, 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, axis) / 4, 65535)",
           "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "xShape": "shapes.x", "axisSpec": "axis" },
           "bindings": [
+            "x_exp_sum",
             { "name": "rowMax", "buffer": "read-only-storage", "elementType": "vec4<f32>" },
             { "arg": "y", "elementType": "vec4<f32>" },
+            { "name": "rowSum", "elementType": "vec4<f32>" }
           ],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, axis) / 4, 65535)",
           "name": "LogSoftmax.StridedNormalize",
           "shader": "softmax-normalize.wgsl.jinja",
           "derive": { "strided": "true", "xShape": "shapes.x", "axisSpec": "axis", "op": "\"logsoftmax\"" },
+          "bindings": ["x", "rowMax_f32", "rowSum_f32", "y", "params_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.x)), (tunables.WORKGROUP_SIZE)), 65535)",
     {
       "id": "strided_online_coop",
       "priority": 19,
+      "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "stridedCoopDispatchFits"],
       "passes": [
         {
           "id": "main",
           "derive": {
             "xShape": "shapes.x",
             "axisSpec": "axis",
             "tileCols": "stridedCoopTileCols",
+            "onlineStreams": "stridedOnlineStreams",
+            "op": "\"logsoftmax\""
           },
           "bindings": ["x", "y"],
           "dispatch": {
       "id": "strided_online_packed4_tail",
       "priority": 18,
       "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 4", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 != 0", "dim(shapes.x, axis) < 2048 or axisRows < tunables.STRIDED_CAPACITY_MIN_ROWS", "axisDispatchFits"],
       "passes": [
         {
           "id": "main",
       "priority": 12,
       "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisDispatchFits", "axisRowScratchFits"],
       "demoteWhen": ["dtypes.T == \"f32\" and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 4 and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 != 0"],
       "intermediates": [
         { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" },
         { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" }
           "name": "LogSoftmax.StridedExpSum",
           "shader": "softmax-row-stage.wgsl.jinja",
           "derive": { "stage": "\"exp_sum\"", "strided": "true", "xShape": "shapes.x", "axisSpec": "axis" },
+          "bindings": ["x", "rowMax_f32", "y", "rowSum"],
           "dispatch": {
             "x": "min(numel(shapes.x) / dim(shapes.x, axis), 65535)",
             "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, axis), 65535)",
           "name": "LogSoftmax.StridedNormalize",
           "shader": "softmax-normalize.wgsl.jinja",
           "derive": { "strided": "true", "xShape": "shapes.x", "axisSpec": "axis", "op": "\"logsoftmax\"" },
+          "bindings": ["x", "rowMax_f32", "rowSum_f32", "y", "params_count"],
           "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,30 +1,30 @@
 {
   "name": "ai.onnx.LogSoftmax",
-  "id": "_ai_onnx_logsoftmax_webgpu_eafe4e3",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "PypxQO0F40vAAACduE1EoHuJbXFSoUdf/1sxjfO3XYI=",
-      "manifest.json": "YEgGrgHNsUKssjpYJg/VI7JNOc7ROOzoE7TnwK8uD48=",
-      "softmax-longrow-normalize.wgsl.jinja": "WUMowCIMomWax9n1OZifCH/q1sfWCKgwU5OdAH6eOMo=",
-      "softmax-longrow-stats.wgsl.jinja": "prBVBse5ULq+Wl8maEVHnHtV24ispCW8XYS6ys3rzmA=",
-      "softmax-normalize.wgsl.jinja": "5HX4KpHm5VW4oAq0LfwDIv2h4ZcaW5F6DIw8hYHiwBA=",
-      "softmax-online-packed-rows.wgsl.jinja": "R05RjrEHTCi8zUwQKccDnrznkHAi3i4Xl1+VAXECEyQ=",
-      "softmax-online.wgsl.jinja": "t/yk5HbGLSgCjBNT2NbQ93rACba+TWWe0Hc/I4M3YhY=",
-      "softmax-row-stage-strided-vec4.wgsl.jinja": "bJgInlBtWpRzBoQLGZzAIWQ+crsQ9rt4Sokaq/fozFE=",
-      "softmax-row-stage.wgsl.jinja": "HyD4uC2ncynJAQwamSNH9orhDje96yEs5IecERo2ivI=",
-      "softmax-strided-online-coop.wgsl.jinja": "/OpczYeV5thgixotFxCXjWVCQvqO3kEKY0fmCgm0GPw=",
-      "softmax-strided-packed4-tail.wgsl.jinja": "CUut2x8eExSySbqp4R97Rt6F0GbFJ6b5LOHLKLpWG5Q=",
-      "softmax-subgroup-rows.wgsl.jinja": "jwyn5Nf1BBfBuaAIKIfJ5tMg16FcOeMW7llo2hy/ktw=",
-      "test.json": "pLfIl89rWpvYoFMWrF35r5mJTQQ1BkpZPmChvLUF2nQ="
     }
   },
-  "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.LogSoftmax",
+  "id": "_ai_onnx_logsoftmax_webgpu_4896bd3",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "SvuAkJFufs4pJCvvQl68ccv1olqep5tmHlqxCp4Dd/s=",
+      "manifest.json": "agl/2kS1pO7NsCz4/7uItvvoLvoK5flm7PTJ89i8bHw=",
+      "softmax-longrow-normalize.wgsl.jinja": "jy4t+YROOPczKJwk1qn7dgbA4TFpslDRX4CLql5Ychk=",
+      "softmax-longrow-stats.wgsl.jinja": "S8w3EVvJXQavrn/79+OYm6yMZWERwlyJiZhYF8HkAj0=",
+      "softmax-normalize.wgsl.jinja": "MRBOzfP0VMvC2H8uWhMJ3wm5RckrdGxMneOy88Y92b4=",
+      "softmax-online-packed-rows.wgsl.jinja": "8SJM2qyv5RnJK/nzoif+8RZBcIIovrtcCVq+tQQ6n0k=",
+      "softmax-online.wgsl.jinja": "kspDFq+qGAmhnhKKcOmOWTLOQ2vX6y54NHbUlPnxhHY=",
+      "softmax-row-stage-strided-vec4.wgsl.jinja": "zTL9bEcqInxFgFd5V2HamSyPqzPzZyEcuW87QlO4aU4=",
+      "softmax-row-stage.wgsl.jinja": "Uv2knZMUyaueKOSxAhn7Vko2N0QajW/vyOJWSNbyux4=",
+      "softmax-strided-online-coop.wgsl.jinja": "ZMjzbUYJpr/AI6VDYIoPvqIt2vOadAtKxkfc/eFcvNE=",
+      "softmax-strided-packed4-tail.wgsl.jinja": "16AesZJMCqOU7HANkDNZlJdjpkohzbOk/M7pQu/vMNQ=",
+      "softmax-subgroup-rows.wgsl.jinja": "l8/ApzifNWPXI4+nQlz+dJ4ZI+t7IeiDqLrRj/XmmQw=",
+      "test.json": "Purex4g7SQtnFKtBVaRg6FkZWPLvCWGk/5pQ1aol1iI="
     }
   },
+  "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;
@@ -34,7 +31,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;
   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,9 +1,13 @@
-{% if T == "f16" %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% if strided %}
 {% set stridedAxisUses = ["row_for_offset"] %}
 // Compiled indexing for row-wise reductions over a non-last axis. A row is one
 // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
@@ -42,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 {
@@ -60,7 +64,7 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif -%}
 {% endif %}
 // 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
@@ -88,15 +92,9 @@ fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @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;
-  }
   // Final element-parallel log-softmax pass. Each element rereads x and writes
   // (x - rowMax) - log(rowSum); strided mode changes only the row map.
 {% if strided %}

+{% 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 strided %}
 {% set stridedAxisUses = ["row_for_offset"] %}
 // Compiled indexing for row-wise reductions over a non-last axis. A row is one
 // slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
 {% 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 %}
 {% endif %}
 // 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
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
   // Final element-parallel log-softmax pass. Each element rereads x and writes
   // (x - rowMax) - log(rowSum); strided mode changes only the row map.
 {% if strided %}

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);
@@ -218,7 +191,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   let logD = log(finalD);
 {% if useVec4 %}
   for (var c = lid.x; c < params.vecCols; c = c + WG) {
-    let v = vec4<f32>(x[rowBase + c]);
     var o = vec4<f32>(
       shifted_value(v.x, finalM),
       shifted_value(v.y, finalM),
@@ -227,7 +200,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
     if (rowIsNan) {
       o = vec4<f32>(nanOut);
     }
-    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);
   let logD = log(finalD);
 {% if useVec4 %}
   for (var c = lid.x; c < params.vecCols; c = c + WG) {
+    let v = {{ read_vec4("c") }};
     var o = vec4<f32>(
       shifted_value(v.x, finalM),
       shifted_value(v.y, finalM),
     if (rowIsNan) {
       o = vec4<f32>(nanOut);
     }
+{{ store_vec4("o", "c") }}
   }
 {% else %}
   for (var c = lid.x; c < params.cols; c = c + WG) {

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

@@ -46,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 {
@@ -64,7 +64,8 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif -%}
 {% if stage == "max" %}
 // max(1u,..) so an empty strided tensor (INNER_SIZE 0) still compiles — the kernel
 // returns before any divide (VEC4_ROWS 0 -> guard), this only avoids a const `/0`.
@@ -84,17 +85,19 @@ var<workgroup> partial: array<vec4<f32>, WGS>;
 fn {{ reduce_fn }}(value: vec4<f32>, tid: u32) -> vec4<f32> {
   partial[tid] = value;
   workgroupBarrier();
-  var stride: u32 = WGS / 2u;
   loop {
-    if (stride == 0u) { break; }
-{% if stage == "max" %}
-    if (tid < stride) { partial[tid] = max(partial[tid], partial[tid + stride]); }
-{% else %}
-    if (tid < stride) { partial[tid] = partial[tid] + partial[tid + stride]; }
-{% endif %}
-    stride = stride / 2u;
     workgroupBarrier();
-  }
   return partial[0];
 }

 {% 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 stage == "max" %}
 // max(1u,..) so an empty strided tensor (INNER_SIZE 0) still compiles — the kernel
 // returns before any divide (VEC4_ROWS 0 -> guard), this only avoids a const `/0`.
 fn {{ reduce_fn }}(value: vec4<f32>, tid: u32) -> vec4<f32> {
   partial[tid] = value;
   workgroupBarrier();
+{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" or op == "min" %}
+{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
+  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
+    if ({{ svar }} == 0u) { break; }
+    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
+    {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
+{{ wgsl_tree_fold(["partial"], op=stage, idx="tid", wg="WGS", typed=true, form="head", breakInline=true, bodyInline=true) }}
   return partial[0];
 }

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

@@ -1,9 +1,6 @@
 // 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;
@@ -46,17 +43,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 {
@@ -64,10 +61,10 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif -%}
 {% endif %}
-{% 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
@@ -93,14 +90,13 @@ 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 %}
@@ -108,7 +104,6 @@ var<workgroup> partial: array<f32, 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>) {
 {% if strided %}

 // 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;
 {% 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 %}
+{% endif %}
+{% 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 %}
 {% 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>) {
 {% if strided %}

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

@@ -5,12 +5,7 @@
 // Each invocation maintains the manifest-selected number of independent
 // 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 +45,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,7 +63,7 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif -%}
 // 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
@@ -88,11 +83,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);
 }
-const TILE_COLS: u32 = {{ tileCols }}u;
-const AXIS_LANES: u32 = {{ axisLanes }}u;
-{% 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 }};
@@ -106,6 +97,9 @@ const AXIS_LANES: u32 = {{ axisLanes }}u;
       }
 {% endmacro %}
 var<workgroup> part_max: array<f32, TILE_COLS * AXIS_LANES>;
 var<workgroup> part_sum: array<f32, TILE_COLS * AXIS_LANES>;

 // Each invocation maintains the manifest-selected number of independent
 // 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 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 %}
 // 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
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
+{% set negInit = "-FLT_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;
 var<workgroup> part_max: array<f32, TILE_COLS * AXIS_LANES>;
 var<workgroup> part_sum: array<f32, TILE_COLS * AXIS_LANES>;

build/webgpu/softmax-strided-packed4-tail.wgsl.jinja CHANGED Viewed

@@ -44,17 +44,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 {
@@ -62,7 +62,7 @@ fn row_for_offset(offset: u32) -> u32 {
   let inner = offset % INNER_SIZE_SAFE;
   return outer * INNER_SIZE_SAFE + inner;
 }
-{%- endif -%}
 // 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
@@ -82,6 +82,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));
 }
@@ -98,6 +99,9 @@ const PACKED_ROWS: u32 = OUTER_SIZE * PACKED_INNER;
 var<workgroup> partial: array<vec4<f32>, WGS>;
 fn reduce_max4(value: vec4<f32>, tid: u32) -> vec4<f32> {
   partial[tid] = value;
   workgroupBarrier();
@@ -107,7 +111,9 @@ fn reduce_max4(value: vec4<f32>, tid: u32) -> vec4<f32> {
     }
     workgroupBarrier();
   }
-  return partial[0];
 }
 fn reduce_add4(value: vec4<f32>, tid: u32) -> vec4<f32> {
@@ -119,7 +125,9 @@ fn reduce_add4(value: vec4<f32>, tid: u32) -> vec4<f32> {
     }
     workgroupBarrier();
   }
-  return partial[0];
 }
 fn stable_exp(value: f32, maximum: f32) -> f32 {

 {% 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 %}
 // 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
   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));
 }
 var<workgroup> partial: array<vec4<f32>, WGS>;
+// Both reductions fold in `partial`, so each returns only after every lane has
+// loaded slot 0: without the trailing barrier, a lane that starts the next
+// reduction overwrites slot 0 before slower lanes read the previous result.
 fn reduce_max4(value: vec4<f32>, tid: u32) -> vec4<f32> {
   partial[tid] = value;
   workgroupBarrier();
     }
     workgroupBarrier();
   }
+  let result = partial[0];
+  workgroupBarrier();
+  return result;
 }
 fn reduce_add4(value: vec4<f32>, tid: u32) -> vec4<f32> {
     }
     workgroupBarrier();
   }
+  let result = partial[0];
+  workgroupBarrier();
+  return result;
 }
 fn stable_exp(value: f32, maximum: f32) -> f32 {

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

@@ -22,7 +22,7 @@
     {
       "name": "strided_capacity_large_axis_tail_boundary",
       "provenance": {
-        "notes": "At the large-inner crossover, axis length 2,048 and inner size 6 select the capacity path with a non-multiple-of-four tail."
       },
       "tunables": { "STRIDED_CAPACITY_MIN_ROWS": 1 },
       "attrs": { "axis": 1 },
@@ -73,7 +73,7 @@
     {
       "name": "packed_rows_tail_axis1_33x32",
       "provenance": {
-        "notes": "One row beyond a full 32-row packed workgroup locks the shared portable multirow kernel's partial-workgroup barrier and output guards."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -102,9 +102,7 @@
     },
     {
       "name": "many_rows_axis1_2048x32",
-      "provenance": {
-        "notes": "Many 32-column rows exercise the packed-rows path below the 64-column subgroup-per-row threshold."
-      },
       "attrs": { "axis": 1 },
       "inputs": {
         "x": {
@@ -712,7 +710,7 @@
     {
       "name": "strided_rowmax_scratch_over_128mib_f16_capacity_fallback",
       "provenance": {
-        "notes": "With 33,554,433 strided rows, a float32 row-maximum buffer would exceed 128 MiB. Float16 keeps the input within the minimum storage-binding limit and selects the single-pass online path without row scratch."
       },
       "attrs": { "axis": 0 },
       "inputs": {
@@ -790,7 +788,7 @@
     {
       "name": "online_stream_axis31_f32",
       "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": {
@@ -805,7 +803,7 @@
     {
       "name": "online_stream_axis32_f32",
       "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": {
@@ -820,7 +818,7 @@
     {
       "name": "online_stream_axis33_f32",
       "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": {
@@ -835,7 +833,7 @@
     {
       "name": "online_stream_axis64_f32",
       "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": {
@@ -850,7 +848,7 @@
     {
       "name": "online_stream_axis65_f32",
       "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": {
@@ -865,7 +863,7 @@
     {
       "name": "online_stream_columns_tail_f32",
       "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": {
@@ -906,7 +904,7 @@
     {
       "name": "online_stream_four_way_tail_f32",
       "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": {

     {
       "name": "strided_capacity_large_axis_tail_boundary",
       "provenance": {
+        "notes": "LogSoftmax axis length 2048 is strided by an inner size of 6, which is not a multiple of four, leaving a scalar tail; checks correctness on a non-contiguous reduction axis at this length."
       },
       "tunables": { "STRIDED_CAPACITY_MIN_ROWS": 1 },
       "attrs": { "axis": 1 },
     {
       "name": "packed_rows_tail_axis1_33x32",
       "provenance": {
+        "notes": "Thirty-three rows leave one row after a full 32-row group; every row must be normalized independently and written once."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     },
     {
       "name": "many_rows_axis1_2048x32",
+      "provenance": { "notes": "Many 32-column rows exercise short-row LogSoftmax." },
       "attrs": { "axis": 1 },
       "inputs": {
         "x": {
     {
       "name": "strided_rowmax_scratch_over_128mib_f16_capacity_fallback",
       "provenance": {
+        "notes": "33,554,433 independent strided float16 rows exceed 128 MiB of float32 per-row statistics; the output must still be correct within the minimum storage-binding limit."
       },
       "attrs": { "axis": 0 },
       "inputs": {
     {
       "name": "online_stream_axis31_f32",
       "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_f32",
       "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_f32",
       "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_f32",
       "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_f32",
       "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_f32",
       "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_f32",
       "provenance": {
+        "notes": "Exercises the per-lane stream-count boundary, reduction tail, or column tail. f16 comparison includes its storage rounding."
       },
       "attrs": { "axis": 1 },
       "inputs": {