sync 6fdf6301e2bb

Browse files

Files changed (12) hide show

README.md +8 -2
build/webgpu/bench.json +1 -1
build/webgpu/manifest.json +129 -143
build/webgpu/metadata.json +14 -13
build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja +14 -20
build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja +28 -14
build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja +5 -16
build/webgpu/reduce-arg-axis-tiled.wgsl.jinja +4 -11
build/webgpu/reduce-arg-axis.wgsl.jinja +13 -22
build/webgpu/reduce-arg-row-split.wgsl.jinja +49 -79
build/webgpu/reduce-arg-row-subgroup.wgsl.jinja +75 -95
build/webgpu/test.json +122 -10

README.md CHANGED Viewed

@@ -45,6 +45,12 @@ Default values (overridable per request):
 | `T` | `float32`, `float16`, `int32`, `uint32`, `int16`, `int8`, `uint8` |
 | `I` | `int64` |
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
@@ -54,7 +60,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
 - [`reduce-arg-axis-split-combine.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja)
 - [`reduce-arg-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja)
 - [`reduce-arg-axis-split-tiled.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja)
@@ -66,7 +72,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.

 | `T` | `float32`, `float16`, `int32`, `uint32`, `int16`, `int8`, `uint8` |
 | `I` | `int64` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `axis_split_packed` — Scans four adjacent float32 inner-axis columns per invocation while preserving each column’s axis order and the scalar partial/combine layout. Small output grids retain the scalar route.
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 - [`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
 - [`reduce-arg-axis-split-combine.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja)
 - [`reduce-arg-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja)
 - [`reduce-arg-axis-split-tiled.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-tiled.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

@@ -137,7 +137,7 @@
       "name": "argmax-lastaxis-f32-64x50260-vocab-control",
       "preset": "stress",
       "provenance": {
-        "notes": "Nearest multiple-of-4 vocabulary width for argmax-lastaxis-f32-64x50257-vocab-odd-pathology; same bytes within 0.01%, selects last_axis_vec4."
       },
       "vars": { "dtype": "float32", "rows": 64, "cols": 50260 },
       "attrs": { "axis": 1 },

       "name": "argmax-lastaxis-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

@@ -60,71 +60,9 @@
     "rowSplitScratchBytes": "rowSplitCount * outputCount * 8",
     "rowSplitScratchFits": "rowSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and rowSplitScratchBytes <= device.limits.maxBufferSize",
     "rowSplitDispatchFits": "outputCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and rowSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "tileDispatchFits": "ceilDiv(outputCount, tunables.TILE_COLUMNS) <= foldedDispatchCapacity"
-  },
-  "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "partials_val": { "buffer": "storage", "elementType": "u32" },
-    "partials_idx": { "buffer": "storage", "elementType": "u32" },
-    "params": {
-      "buffer": "uniform",
-      "struct": [
-        { "name": "rows", "type": "u32", "value": "outputCount" },
-        { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" }
-      ]
-    },
-    "partials_val_2": { "name": "partials_val", "buffer": "read-only-storage", "elementType": "u32" },
-    "partials_idx_2": { "name": "partials_idx", "buffer": "read-only-storage", "elementType": "u32" },
-    "y": { "buffer": "storage", "elementType": "u32" },
-    "params_2": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "outputCount", "type": "u32", "value": "outputCount" }]
-    },
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "params_3": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [
-        { "name": "rows", "type": "u32", "value": "outputCount" },
-        { "name": "chunkCount", "type": "u32", "value": "axisDim" }
-      ]
-    },
-    "params_4": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [
-        { "name": "axisDim", "type": "u32", "value": "axisDim" },
-        { "name": "outputCount", "type": "u32", "value": "outputCount" },
-        { "name": "innerSize", "type": "u32", "value": "axisInner" }
-      ]
-    },
-    "params_5": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [
-        { "name": "rows", "type": "u32", "value": "outputCount" },
-        { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" },
-        {
-          "name": "rowStride",
-          "type": "u32",
-          "value": "max(1, min(outputCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
-        }
-      ]
-    },
-    "params_6": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [
-        { "name": "rows", "type": "u32", "value": "outputCount" },
-        { "name": "chunkCount", "type": "u32", "value": "axisDim" },
-        {
-          "name": "rowStride",
-          "type": "u32",
-          "value": "max(1, min(outputCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
-        }
-      ]
-    }
   },
   "variants": [
     {
@@ -134,8 +72,6 @@
       "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "workgroupSize",
         "split": "rowSplitCount"
       },
@@ -148,8 +84,7 @@
           "id": "split_reduce",
           "name": "ArgMax.LastAxisSplitVec4Reduce",
           "shader": "reduce-arg-row-split.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": { "modeSpec": "\"max\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" },
           "bindings": ["x", "partials_val", "partials_idx", "params"],
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
@@ -157,8 +92,7 @@
           "id": "combine",
           "name": "ArgMax.LastAxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
@@ -171,13 +105,7 @@
       "id": "last_axis_split",
       "priority": 41,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= tunables.ROW_SPLIT_MIN_AXIS", "axisDim % 4 != 0", "outputCount <= tunables.ROW_SPLIT_MAX_OUTPUTS", "rowSplitScratchFits", "rowSplitDispatchFits", "rowStorageFits"],
-      "derive": {
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
-        "workgroupSize": "workgroupSize",
-        "split": "rowSplitCount"
-      },
       "intermediates": [
         { "id": "partials_val", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" },
         { "id": "partials_idx", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" }
@@ -187,17 +115,15 @@
           "id": "split_reduce",
           "name": "ArgMax.LastAxisSplitReduce",
           "shader": "reduce-arg-row-split.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": { "modeSpec": "\"max\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" },
-          "bindings": ["x_2", "partials_val", "partials_idx", "params_3"],
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMax.LastAxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
@@ -212,7 +138,6 @@
       "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= tunables.SPLIT_TILE_COLUMNS", "outputCount <= 2 * tunables.SPLIT_TILE_COLUMNS", "workgroupSize % tunables.SPLIT_TILE_COLUMNS == 0", "splitScratchFits", "splitTileDispatchFits"],
       "derive": {
         "scalar": "dtypes.T",
-        "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "workgroupSize",
         "split": "splitCount",
         "tileCols": "tunables.SPLIT_TILE_COLUMNS"
@@ -226,16 +151,14 @@
           "id": "split_reduce",
           "name": "ArgMax.AxisSplitTiledReduce",
           "shader": "reduce-arg-axis-split-tiled.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["x_2", "partials_val", "partials_idx", "params_4"],
           "dispatch": { "x": "ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS)", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMax.AxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
@@ -252,8 +175,6 @@
       "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "device.adapterInfo.subgroupMinSize"
       },
       "passes": [
@@ -261,9 +182,8 @@
           "id": "main",
           "name": "ArgMax.SubgroupMinLastAxisVec4",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": { "modeSpec": "\"max\"", "vec4": true, "singleSubgroup": true },
-          "bindings": ["x", "y", "params_5"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
@@ -273,20 +193,14 @@
       "priority": 41,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim % 4 != 0", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize >= 4", "device.adapterInfo.subgroupMinSize <= device.adapterInfo.subgroupMaxSize", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeInvocationsPerWorkgroup", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeWorkgroupSizeX", "axisDim >= device.adapterInfo.subgroupMinSize * 4", "axisDim <= device.adapterInfo.subgroupMinSize * 64", "outputCount >= device.adapterInfo.subgroupMinSize * 4", "rowDispatchFits"],
       "requires": { "features": ["subgroups"] },
-      "derive": {
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
-        "workgroupSize": "device.adapterInfo.subgroupMinSize"
-      },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.SubgroupMinLastAxis",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": { "modeSpec": "\"max\"", "vec4": false, "singleSubgroup": true },
-          "bindings": ["x_2", "y", "params_6"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
@@ -298,8 +212,6 @@
       "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "min(workgroupSize, max(32, pow2ceil(ceilDiv(axisDim, 4))))"
       },
       "passes": [
@@ -307,9 +219,8 @@
           "id": "main",
           "name": "ArgMax.LastAxisVec4",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": { "modeSpec": "\"max\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" },
-          "bindings": ["x", "y", "params_5"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
@@ -318,46 +229,79 @@
       "id": "last_axis",
       "priority": 39,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim % 4 != 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
-      "derive": {
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
-        "workgroupSize": "min(workgroupSize, max(32, pow2ceil(axisDim)))"
-      },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.LastAxis",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": { "modeSpec": "\"max\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" },
-          "bindings": ["x_2", "y", "params_6"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
     },
     {
-      "id": "axis_split",
-      "priority": 30,
-      "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= 1", "outputCount <= tunables.SPLIT_MAX_OUTPUTS", "splitScratchFits", "splitDispatchFits", "threadDispatchFits"],
-      "demoteWhen": ["narrowLargeAxisTilePreferred"],
       "derive": {
         "scalar": "dtypes.T",
-        "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "workgroupSize",
-        "split": "splitCount"
       },
       "intermediates": [
         { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
         { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
       ],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ArgMax.AxisSplitReduce",
           "shader": "reduce-arg-axis-split-reduce.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["x_2", "partials_val", "partials_idx", "params_4"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -368,8 +312,7 @@
           "id": "combine",
           "name": "ArgMax.AxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
@@ -382,19 +325,13 @@
       "id": "axis_tiled",
       "priority": 20,
       "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.TILED_MIN_AXIS", "outputCount >= tunables.TILED_MIN_OUTPUTS", "workgroupSize % tunables.TILE_COLUMNS == 0", "rowStorageFits", "tileDispatchFits"],
-      "derive": {
-        "scalar": "dtypes.T",
-        "selectLastIndex": "attrs.select_last_index != 0",
-        "workgroupSize": "workgroupSize",
-        "tileCols": "tunables.TILE_COLUMNS"
-      },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.AxisTiled",
           "shader": "reduce-arg-axis-tiled.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["x_2", "y", "params_4"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (tunables.TILE_COLUMNS)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (tunables.TILE_COLUMNS)), 65535)",
@@ -408,19 +345,13 @@
       "priority": 0,
       "when": ["baseContract", "threadDispatchFits"],
       "supersededBy": ["axis_tiled"],
-      "derive": {
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "selectLastIndex": "attrs.select_last_index != 0",
-        "workgroupSize": "workgroupSize"
-      },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.AxisSerial",
           "shader": "reduce-arg-axis.wgsl.jinja",
-          "derive": { "modeSpec": "\"max\"" },
-          "bindings": ["x_2", "y", "params_4"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
@@ -429,5 +360,60 @@
         }
       ]
     }
-  ]
 }

     "rowSplitScratchBytes": "rowSplitCount * outputCount * 8",
     "rowSplitScratchFits": "rowSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and rowSplitScratchBytes <= device.limits.maxBufferSize",
     "rowSplitDispatchFits": "outputCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and rowSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "tileDispatchFits": "ceilDiv(outputCount, tunables.TILE_COLUMNS) <= foldedDispatchCapacity",
+    "selectLastIndex": "attrs.select_last_index != 0",
+    "modeSpec": "\"max\""
   },
   "variants": [
     {
       "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "workgroupSize",
         "split": "rowSplitCount"
       },
           "id": "split_reduce",
           "name": "ArgMax.LastAxisSplitVec4Reduce",
           "shader": "reduce-arg-row-split.wgsl.jinja",
+          "derive": { "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" },
           "bindings": ["x", "partials_val", "partials_idx", "params"],
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
           "id": "combine",
           "name": "ArgMax.LastAxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "bindings": ["partials_val_u32", "partials_idx_u32", "y", "params_combine"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
       "id": "last_axis_split",
       "priority": 41,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= tunables.ROW_SPLIT_MIN_AXIS", "axisDim % 4 != 0", "outputCount <= tunables.ROW_SPLIT_MAX_OUTPUTS", "rowSplitScratchFits", "rowSplitDispatchFits", "rowStorageFits"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "workgroupSize", "split": "rowSplitCount" },
       "intermediates": [
         { "id": "partials_val", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" },
         { "id": "partials_idx", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" }
           "id": "split_reduce",
           "name": "ArgMax.LastAxisSplitReduce",
           "shader": "reduce-arg-row-split.wgsl.jinja",
+          "derive": { "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["x_main", "partials_val", "partials_idx", "params__uniform"],
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMax.LastAxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "bindings": ["partials_val_u32", "partials_idx_u32", "y", "params_combine"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
       "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= tunables.SPLIT_TILE_COLUMNS", "outputCount <= 2 * tunables.SPLIT_TILE_COLUMNS", "workgroupSize % tunables.SPLIT_TILE_COLUMNS == 0", "splitScratchFits", "splitTileDispatchFits"],
       "derive": {
         "scalar": "dtypes.T",
         "workgroupSize": "workgroupSize",
         "split": "splitCount",
         "tileCols": "tunables.SPLIT_TILE_COLUMNS"
           "id": "split_reduce",
           "name": "ArgMax.AxisSplitTiledReduce",
           "shader": "reduce-arg-axis-split-tiled.wgsl.jinja",
+          "bindings": ["x_main", "partials_val", "partials_idx", "params_main"],
           "dispatch": { "x": "ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS)", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMax.AxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "bindings": ["partials_val_u32", "partials_idx_u32", "y", "params_combine"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
       "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "device.adapterInfo.subgroupMinSize"
       },
       "passes": [
           "id": "main",
           "name": "ArgMax.SubgroupMinLastAxisVec4",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "derive": { "vec4": true, "singleSubgroup": true },
+          "bindings": ["x", "y", "params_last_axis_vec4"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
       "priority": 41,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim % 4 != 0", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize >= 4", "device.adapterInfo.subgroupMinSize <= device.adapterInfo.subgroupMaxSize", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeInvocationsPerWorkgroup", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeWorkgroupSizeX", "axisDim >= device.adapterInfo.subgroupMinSize * 4", "axisDim <= device.adapterInfo.subgroupMinSize * 64", "outputCount >= device.adapterInfo.subgroupMinSize * 4", "rowDispatchFits"],
       "requires": { "features": ["subgroups"] },
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "device.adapterInfo.subgroupMinSize" },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.SubgroupMinLastAxis",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "derive": { "vec4": false, "singleSubgroup": true },
+          "bindings": ["x_main", "y", "params_last_axis"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
       "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "min(workgroupSize, max(32, pow2ceil(ceilDiv(axisDim, 4))))"
       },
       "passes": [
           "id": "main",
           "name": "ArgMax.LastAxisVec4",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "derive": { "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["x", "y", "params_last_axis_vec4"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
       "id": "last_axis",
       "priority": 39,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim % 4 != 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "min(workgroupSize, max(32, pow2ceil(axisDim)))" },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.LastAxis",
           "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "derive": {
+            "vec4": false,
+            "useSubgroups": "device.features.has(\"subgroups\")",
+            "adjacentScalarLoads": "false"
+          },
+          "bindings": ["x_main", "y", "params_last_axis"],
           "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
     },
     {
+      "id": "axis_split_packed",
+      "priority": 32,
+      "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= 1", "outputCount <= tunables.SPLIT_MAX_OUTPUTS", "splitScratchFits", "splitDispatchFits", "threadDispatchFits", "dtypes.T == \"f32\"", "axisInner % 4 == 0"],
+      "demoteWhen": ["narrowLargeAxisTilePreferred", "outputCount < workgroupSize * 2"],
       "derive": {
         "scalar": "dtypes.T",
         "workgroupSize": "workgroupSize",
+        "split": "splitCount",
+        "vectorScalar": "\"vec4<f32>\""
       },
       "intermediates": [
         { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
         { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
       ],
+      "passes": [
+        {
+          "id": "split_reduce",
+          "name": "ArgMax.AxisSplitPackedReduce",
+          "shader": "reduce-arg-axis-split-reduce.wgsl.jinja",
+          "derive": { "columnVectorWidth": 4 },
+          "bindings": ["x", "partials_val", "partials_idx", "params_main"],
+          "dispatch": {
+            "x": "min(ceilDiv(outputCount, workgroupSize * 4), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv(outputCount, workgroupSize * 4), DISPATCH_FOLD_WIDTH)"
+          }
+        },
+        {
+          "id": "combine",
+          "name": "ArgMax.AxisSplitCombine",
+          "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "bindings": ["partials_val_u32", "partials_idx_u32", "y", "params_combine"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "axis_split",
+      "priority": 30,
+      "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= 1", "outputCount <= tunables.SPLIT_MAX_OUTPUTS", "splitScratchFits", "splitDispatchFits", "threadDispatchFits"],
+      "demoteWhen": ["narrowLargeAxisTilePreferred"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "workgroupSize", "split": "splitCount" },
+      "intermediates": [
+        { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
+        { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
+      ],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ArgMax.AxisSplitReduce",
           "shader": "reduce-arg-axis-split-reduce.wgsl.jinja",
+          "bindings": ["x_main", "partials_val", "partials_idx", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
           "id": "combine",
           "name": "ArgMax.AxisSplitCombine",
           "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "bindings": ["partials_val_u32", "partials_idx_u32", "y", "params_combine"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
       "id": "axis_tiled",
       "priority": 20,
       "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.TILED_MIN_AXIS", "outputCount >= tunables.TILED_MIN_OUTPUTS", "workgroupSize % tunables.TILE_COLUMNS == 0", "rowStorageFits", "tileDispatchFits"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "workgroupSize", "tileCols": "tunables.TILE_COLUMNS" },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.AxisTiled",
           "shader": "reduce-arg-axis-tiled.wgsl.jinja",
+          "bindings": ["x_main", "y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (tunables.TILE_COLUMNS)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (tunables.TILE_COLUMNS)), 65535)",
       "priority": 0,
       "when": ["baseContract", "threadDispatchFits"],
       "supersededBy": ["axis_tiled"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "workgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ArgMax.AxisSerial",
           "shader": "reduce-arg-axis.wgsl.jinja",
+          "bindings": ["x_main", "y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
         }
       ]
     }
+  ],
+  "bindings": {
+    "x": { "elementType": "$vectorScalar" },
+    "partials_val": { "buffer": "storage", "elementType": "u32" },
+    "partials_idx": { "buffer": "storage", "elementType": "u32" },
+    "params": {
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputCount" },
+        { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" }
+      ]
+    },
+    "partials_val_u32": { "name": "partials_val", "buffer": "read-only-storage", "elementType": "u32" },
+    "partials_idx_u32": { "name": "partials_idx", "buffer": "read-only-storage", "elementType": "u32" },
+    "y": { "elementType": "u32" },
+    "params_combine": { "name": "params", "struct": [{ "name": "outputCount", "type": "u32", "value": "outputCount" }] },
+    "x_main": { "name": "x", "elementType": "$scalar" },
+    "params__uniform": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputCount" },
+        { "name": "chunkCount", "type": "u32", "value": "axisDim" }
+      ]
+    },
+    "params_main": {
+      "name": "params",
+      "struct": [
+        { "name": "axisDim", "type": "u32", "value": "axisDim" },
+        { "name": "outputCount", "type": "u32", "value": "outputCount" },
+        { "name": "innerSize", "type": "u32", "value": "axisInner" }
+      ]
+    },
+    "params_last_axis_vec4": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputCount" },
+        { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" },
+        {
+          "name": "rowStride",
+          "type": "u32",
+          "value": "max(1, min(outputCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+        }
+      ]
+    },
+    "params_last_axis": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputCount" },
+        { "name": "chunkCount", "type": "u32", "value": "axisDim" },
+        {
+          "name": "rowStride",
+          "type": "u32",
+          "value": "max(1, min(outputCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+        }
+      ]
+    }
+  }
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,27 +1,27 @@
 {
   "name": "ai.onnx.ArgMax",
-  "id": "_ai_onnx_argmax_webgpu_7a0b0be",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "Pmvyno397Q/iYsA6Ig9/ltAyWvlxIqkGycwLnpwWjUg=",
-      "manifest.json": "b/CUpSEy7gFwzPXKc2bK0XEQQDCO0MiKBs/Ynk6Jals=",
-      "reduce-arg-axis-split-combine.wgsl.jinja": "vNqRyvgMuCTgGo/a3v+7SKnhIWNU/SHjGEGOYSIeo6o=",
-      "reduce-arg-axis-split-reduce.wgsl.jinja": "3GqAQKBSBkxqVaKQg2ycI9Vq3xVteDHQ3PJBMCJqulM=",
-      "reduce-arg-axis-split-tiled.wgsl.jinja": "Z3JRqGsy97pIswTbs/9DR3QgTJblOetszM3RiAhhvb8=",
-      "reduce-arg-axis-tiled.wgsl.jinja": "ndfXWO62yzI5tgZUJeFg9tFmQYm9bRw3QRq0EwgVxTo=",
-      "reduce-arg-axis.wgsl.jinja": "RUQfKNZwkYusxT68lh2FcrAxsx4x/ENaaQ7MXWaV2Ow=",
-      "reduce-arg-row-split.wgsl.jinja": "U8wriN+k8hXANjPPj7ZRRD++kmWr1wh3oCfOcac3qVg=",
-      "reduce-arg-row-subgroup.wgsl.jinja": "QVlX6iupq59hl4a3lLGuLsm62X4VGmwDVvNU+J9/Pwo=",
-      "test.json": "bOm2DRCrjR5ovA0Q7f5/UeTwJ6e1dqlOwSwWilMPCJo="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "last_axis_split_vec4": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-row-split.wgsl.jinja"],
       "last_axis_split": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-row-split.wgsl.jinja"],
@@ -30,6 +30,7 @@
       "subgroup_min_last_axis": ["reduce-arg-row-subgroup.wgsl.jinja"],
       "last_axis_vec4": ["reduce-arg-row-subgroup.wgsl.jinja"],
       "last_axis": ["reduce-arg-row-subgroup.wgsl.jinja"],
       "axis_split": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-axis-split-reduce.wgsl.jinja"],
       "axis_tiled": ["reduce-arg-axis-tiled.wgsl.jinja"],
       "axis_serial": ["reduce-arg-axis.wgsl.jinja"]

 {
   "name": "ai.onnx.ArgMax",
+  "id": "_ai_onnx_argmax_webgpu_1e3a265",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "iuzwJJYAdNsUw+DglCvHX2WAtQEjRzfBOZt6MvlMa1U=",
+      "manifest.json": "TiCizYs12kW8ZsFx/7wB+8OysF5fLvrk+gdZ0EoFqNk=",
+      "reduce-arg-axis-split-combine.wgsl.jinja": "Zgdr9BmWms05/uDEY+YMSJE5jj6soIVfmxLsIfzt6M4=",
+      "reduce-arg-axis-split-reduce.wgsl.jinja": "iYeNFJ7YX9FMCNdlFlhUIW21KpKmU4ouODT2u9MjjjE=",
+      "reduce-arg-axis-split-tiled.wgsl.jinja": "MvEKY6B3FsRB7SfFPVWK7rWWCUsdhTSYOUmbQxtzP2w=",
+      "reduce-arg-axis-tiled.wgsl.jinja": "sWYpTFWzc2ATUo2VGrKw+zLAyWFVD6O9og+WpY9f28g=",
+      "reduce-arg-axis.wgsl.jinja": "9wjP1QeFI3ydGt9T537YVC3O8tPYSK4qSvw+urw1ruY=",
+      "reduce-arg-row-split.wgsl.jinja": "yfXWnH5jaZHa5PW28zyTSF0O7A1yGKOIJUCZdZliFS0=",
+      "reduce-arg-row-subgroup.wgsl.jinja": "A1sX+NMUilPMxLDzHhdKFlwdHIDgFTe9a6igb7pOX7k=",
+      "test.json": "6RGFz3E0tdm0K3gRicwYZ3hU54MqthGV/tnBeAxFxHs="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "last_axis_split_vec4": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-row-split.wgsl.jinja"],
       "last_axis_split": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-row-split.wgsl.jinja"],
       "subgroup_min_last_axis": ["reduce-arg-row-subgroup.wgsl.jinja"],
       "last_axis_vec4": ["reduce-arg-row-subgroup.wgsl.jinja"],
       "last_axis": ["reduce-arg-row-subgroup.wgsl.jinja"],
+      "axis_split_packed": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-axis-split-reduce.wgsl.jinja"],
       "axis_split": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-axis-split-reduce.wgsl.jinja"],
       "axis_tiled": ["reduce-arg-axis-tiled.wgsl.jinja"],
       "axis_serial": ["reduce-arg-axis.wgsl.jinja"]

build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja CHANGED Viewed

@@ -7,9 +7,17 @@
 // The fold compares (value, index), never segment order. Equal values keep the
 // smaller index for first-index mode and the larger index for last-index mode.
 // Empty segments carry SENTINEL_IDX and are skipped.
-{% set cmp = ">" if modeSpec == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -18,28 +26,14 @@ const SENTINEL_IDX: u32 = 4294967295u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let output_index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (output_index >= params.outputCount) {
-    return;
-  }
-{% if scalar == "u32" %}
-  var out_val = partials_val[output_index];
-{% else %}
-  var out_val = bitcast<{{ work }}>(partials_val[output_index]);
-{% endif %}
   var out_idx = partials_idx[output_index];
   for (var split_index = 1u; split_index < SPLIT; split_index++) {
-{% if scalar == "u32" %}
-    let candidate_val = partials_val[split_index * params.outputCount + output_index];
-{% else %}
-    let candidate_val = bitcast<{{ work }}>(partials_val[split_index * params.outputCount + output_index]);
-{% endif %}
     let candidate_idx = partials_idx[split_index * params.outputCount + output_index];
-{% if selectLastIndex %}
-    if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
-{% else %}
-    if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
-{% endif %}
       out_val = candidate_val;
       out_idx = candidate_idx;
     }

 // The fold compares (value, index), never segment order. Equal values keep the
 // smaller index for first-index mode and the larger index for last-index mode.
 // Empty segments carry SENTINEL_IDX and are skipped.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
+{% macro arg_merge_wins(v, i, best, bestIdx) %}{{ i }} != SENTINEL_IDX && ({{ bestIdx }} == SENTINEL_IDX || {{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ i }} {{ ">" if selectLastIndex else "<" }} {{ bestIdx }})){% endmacro %}
+{% 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 }}
 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "output_index", "params.outputCount") }}
+  var out_val = {{ "partials_val[output_index]" if scalar == "u32" else "bitcast<" ~ work ~ ">(partials_val[output_index])" }};
   var out_idx = partials_idx[output_index];
   for (var split_index = 1u; split_index < SPLIT; split_index++) {
+{% set candidate = "partials_val[split_index * params.outputCount + output_index]" %}
+    let candidate_val = {{ candidate if scalar == "u32" else "bitcast<" ~ work ~ ">(" ~ candidate ~ ")" }};
     let candidate_idx = partials_idx[split_index * params.outputCount + output_index];
+    if ({{ arg_merge_wins("candidate_val", "candidate_idx", "out_val", "out_idx") }}) {
       out_val = candidate_val;
       out_idx = candidate_idx;
     }

build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja CHANGED Viewed

@@ -6,9 +6,10 @@
 // layout without losing integer precision. The combine pass restores the native
 // comparison type. Equal values keep the lowest index by default or the highest
 // index when select_last_index is enabled.
-{% set cmp = ">" if modeSpec == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -34,35 +35,48 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(workgroup_id) wg: vec3<u32>) {
   // 2D-folded column index: wg.z carries the high bits past the per-axis dispatch
   // fold width on the x dimension.
-  let output_index = (wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u) * WG + (gid.x % WG);
   let split_index = wg.y;
   if (output_index >= params.outputCount) { return; }
   let outer_index = output_index / params.innerSize;
   let inner_index = output_index % params.innerSize;
-  let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
   // Even split of [0, rows) across SPLIT segments (last segment may be shorter).
   let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
   let axis_begin = split_index * chunk;
   var axis_end = min(axis_begin + chunk, params.axisDim);
-  var best_val = lane_identity();
-  var best_idx = SENTINEL_IDX;
   for (var axis_index = axis_begin; axis_index < axis_end; axis_index++) {
-    let v = x[input_base + axis_index * params.innerSize];
-{% if selectLastIndex %}
-    if (v {{ cmp }} best_val || (v == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
 {% else %}
-    if (v {{ cmp }} best_val || (v == best_val && axis_index < best_idx)) {
-{% endif %}
       best_val = v;
       best_idx = axis_index;
     }
   }
-{% if scalar == "u32" %}
-  partials_val[split_index * params.outputCount + output_index] = best_val;
 {% else %}
-  partials_val[split_index * params.outputCount + output_index] = bitcast<u32>(best_val);
-{% endif %}
   partials_idx[split_index * params.outputCount + output_index] = best_idx;
 }

 // layout without losing integer precision. The combine pass restores the native
 // comparison type. Equal values keep the lowest index by default or the highest
 // index when select_last_index is enabled.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
+{% set columnVectorWidth = columnVectorWidth | default(1) %}
+{% macro arg_wins(v, i, best, bestIdx) %}{{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ "(" ~ bestIdx ~ " == SENTINEL_IDX || " ~ i ~ " > " ~ bestIdx ~ ")" if selectLastIndex else i ~ " < " ~ bestIdx }}){% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
         @builtin(workgroup_id) wg: vec3<u32>) {
   // 2D-folded column index: wg.z carries the high bits past the per-axis dispatch
   // fold width on the x dimension.
+  let output_index = {% if columnVectorWidth > 1 %}({% endif %}(wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u) * WG + (gid.x % WG){% if columnVectorWidth > 1 %}) * {{ columnVectorWidth }}u{% endif %};
   let split_index = wg.y;
   if (output_index >= params.outputCount) { return; }
   let outer_index = output_index / params.innerSize;
   let inner_index = output_index % params.innerSize;
+  let input_base = {% if columnVectorWidth > 1 %}({% endif %}outer_index * params.axisDim * params.innerSize + inner_index{% if columnVectorWidth > 1 %}) / {{ columnVectorWidth }}u{% endif %};
+{% if columnVectorWidth > 1 %}
+  let inner_vectors = params.innerSize / {{ columnVectorWidth }}u;
+{% endif %}
   // Even split of [0, rows) across SPLIT segments (last segment may be shorter).
   let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
   let axis_begin = split_index * chunk;
   var axis_end = min(axis_begin + chunk, params.axisDim);
+  var best_val = {% if columnVectorWidth > 1 %}vec{{ columnVectorWidth }}<{{ work }}>({% endif %}lane_identity(){% if columnVectorWidth > 1 %}){% endif %};
+  var best_idx = {% if columnVectorWidth > 1 %}vec{{ columnVectorWidth }}<u32>({% endif %}SENTINEL_IDX{% if columnVectorWidth > 1 %}){% endif %};
   for (var axis_index = axis_begin; axis_index < axis_end; axis_index++) {
+{% if columnVectorWidth > 1 %}
+    let v = x[input_base + axis_index * inner_vectors];
+    let take = vec{{ columnVectorWidth }}<bool>(
+{% for c in ['x', 'y', 'z', 'w'][:columnVectorWidth] %}
+      {{ arg_wins("v." ~ c, "axis_index", "best_val." ~ c, "best_idx." ~ c) }}{{ ',' if not loop.last else '' }}
+{% endfor %}
+    );
+    best_val = select(best_val, v, take);
+    best_idx = select(best_idx, vec{{ columnVectorWidth }}<u32>(axis_index), take);
 {% else %}
+    let v = x[input_base + axis_index * params.innerSize];
+    if ({{ arg_wins("v", "axis_index", "best_val", "best_idx") }}) {
       best_val = v;
       best_idx = axis_index;
     }
+{% endif %}
   }
+{% if columnVectorWidth > 1 %}
+{% for component in ['x', 'y', 'z', 'w'][:columnVectorWidth] %}
+  partials_val[split_index * params.outputCount + output_index + {{ loop.index0 }}u] = bitcast<u32>(best_val.{{ component }});
+  partials_idx[split_index * params.outputCount + output_index + {{ loop.index0 }}u] = best_idx.{{ component }};
+{% endfor %}
 {% else %}
+  partials_val[split_index * params.outputCount + output_index] = {{ "best_val" if scalar == "u32" else "bitcast<u32>(best_val)" }};
   partials_idx[split_index * params.outputCount + output_index] = best_idx;
+{% endif %}
 }

build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja CHANGED Viewed

@@ -7,9 +7,10 @@
 //
 // The final fold is intentionally subgroup-free: it is small, deterministic,
 // and preserves identical first/last-index tie semantics on every feature tier.
-{% set cmp = ">" if modeSpec == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -56,11 +57,7 @@ fn main(
   if (in_bounds) {
     for (var axis_index = axis_begin + row_lane; axis_index < axis_end; axis_index += ROW_LANES) {
       let value = x[input_base + axis_index * params.innerSize];
-{% if selectLastIndex %}
-      if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
-{% else %}
-      if (value {{ cmp }} best_val || (value == best_val && axis_index < best_idx)) {
-{% endif %}
         best_val = value;
         best_idx = axis_index;
       }
@@ -76,20 +73,12 @@ fn main(
     for (var lane = 1u; lane < ROW_LANES; lane++) {
       let candidate_val = partial_val[lane * TILE_COLS + col_lane];
       let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
-{% if selectLastIndex %}
-      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
-{% else %}
-      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
-{% endif %}
         out_val = candidate_val;
         out_idx = candidate_idx;
       }
     }
-{% if scalar == "u32" %}
-    partials_val[split_index * params.outputCount + output_index] = out_val;
-{% else %}
-    partials_val[split_index * params.outputCount + output_index] = bitcast<u32>(out_val);
-{% endif %}
     partials_idx[split_index * params.outputCount + output_index] = out_idx;
   }
 }

 //
 // The final fold is intentionally subgroup-free: it is small, deterministic,
 // and preserves identical first/last-index tie semantics on every feature tier.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
+{% macro arg_wins(v, i, best, bestIdx) %}{{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ "(" ~ bestIdx ~ " == SENTINEL_IDX || " ~ i ~ " > " ~ bestIdx ~ ")" if selectLastIndex else i ~ " < " ~ bestIdx }}){% endmacro %}
+{% macro arg_merge_wins(v, i, best, bestIdx) %}{{ i }} != SENTINEL_IDX && ({{ bestIdx }} == SENTINEL_IDX || {{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ i }} {{ ">" if selectLastIndex else "<" }} {{ bestIdx }})){% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
   if (in_bounds) {
     for (var axis_index = axis_begin + row_lane; axis_index < axis_end; axis_index += ROW_LANES) {
       let value = x[input_base + axis_index * params.innerSize];
+      if ({{ arg_wins("value", "axis_index", "best_val", "best_idx") }}) {
         best_val = value;
         best_idx = axis_index;
       }
     for (var lane = 1u; lane < ROW_LANES; lane++) {
       let candidate_val = partial_val[lane * TILE_COLS + col_lane];
       let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
+      if ({{ arg_merge_wins("candidate_val", "candidate_idx", "out_val", "out_idx") }}) {
         out_val = candidate_val;
         out_idx = candidate_idx;
       }
     }
+    partials_val[split_index * params.outputCount + output_index] = {{ "out_val" if scalar == "u32" else "bitcast<u32>(out_val)" }};
     partials_idx[split_index * params.outputCount + output_index] = out_idx;
   }
 }

build/webgpu/reduce-arg-axis-tiled.wgsl.jinja CHANGED Viewed

@@ -4,9 +4,10 @@
 // bitcast round trip. Flattening the output as (outer, inner) lets one coalesced
 // column tile handle every non-contiguous axis. f16 uses the serial or
 // contiguous routes because this kernel has no widened-f16 scratch path.
-{% set cmp = ">" if modeSpec == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -50,11 +51,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   if (in_bounds) {
     for (var axis_index = row_lane; axis_index < params.axisDim; axis_index = axis_index + ROW_LANES) {
       let v = x[input_base + axis_index * params.innerSize];
-{% if selectLastIndex %}
-      if (v {{ cmp }} best_val || (v == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
-{% else %}
-      if (v {{ cmp }} best_val || (v == best_val && axis_index < best_idx)) {
-{% endif %}
         best_val = v;
         best_idx = axis_index;
       }
@@ -70,11 +67,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
       let candidate_val = partial_val[lane * TILE_COLS + col_lane];
       let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
-{% if selectLastIndex %}
-      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
-{% else %}
-      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
-{% endif %}
         out_val = candidate_val;
         out_idx = candidate_idx;
       }

 // bitcast round trip. Flattening the output as (outer, inner) lets one coalesced
 // column tile handle every non-contiguous axis. f16 uses the serial or
 // contiguous routes because this kernel has no widened-f16 scratch path.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
+{% macro arg_wins(v, i, best, bestIdx) %}{{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ "(" ~ bestIdx ~ " == SENTINEL_IDX || " ~ i ~ " > " ~ bestIdx ~ ")" if selectLastIndex else i ~ " < " ~ bestIdx }}){% endmacro %}
+{% macro arg_merge_wins(v, i, best, bestIdx) %}{{ i }} != SENTINEL_IDX && ({{ bestIdx }} == SENTINEL_IDX || {{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ i }} {{ ">" if selectLastIndex else "<" }} {{ bestIdx }})){% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
   if (in_bounds) {
     for (var axis_index = row_lane; axis_index < params.axisDim; axis_index = axis_index + ROW_LANES) {
       let v = x[input_base + axis_index * params.innerSize];
+      if ({{ arg_wins("v", "axis_index", "best_val", "best_idx") }}) {
         best_val = v;
         best_idx = axis_index;
       }
     for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
       let candidate_val = partial_val[lane * TILE_COLS + col_lane];
       let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
+      if ({{ arg_merge_wins("candidate_val", "candidate_idx", "out_val", "out_idx") }}) {
         out_val = candidate_val;
         out_idx = candidate_idx;
       }

build/webgpu/reduce-arg-axis.wgsl.jinja CHANGED Viewed

@@ -2,43 +2,34 @@
 // tensor around the reduced axis into (outer, axis, inner) gives one
 // rank-independent coordinate calculation. One thread owns one output position
 // and walks the reduced axis at innerSize stride.
-{% set cmp = ">" if modeSpec == "max" else "<" %}
-{% if usesF16 %}
-enable f16;
-{%- endif %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ workgroupSize }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // Fold the output grid into x/y when one dispatch dimension is insufficient.
-  let output_index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ workgroupSize }}u;
-  if (output_index >= params.outputCount) {
-    return;
-  }
   let outer_index = output_index / params.innerSize;
   let inner_index = output_index % params.innerSize;
   let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
   var best_index = 0u;
-{%- if scalar == "i32" or scalar == "u32" %}
-  var best = x[input_base];
-  for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
-    let candidate = x[input_base + axis_index * params.innerSize];
-    if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
-      best = candidate;
-      best_index = axis_index;
-    }
-  }
-{%- else %}
-  var best = f32(x[input_base]);
   for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
-    let candidate = f32(x[input_base + axis_index * params.innerSize]);
     if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
       best = candidate;
       best_index = axis_index;
     }
   }
-{%- endif %}
   y[output_index] = best_index;
 }

 // tensor around the reduced axis into (outer, axis, inner) gives one
 // rank-independent coordinate calculation. One thread owns one output position
 // and walks the reduced axis at innerSize stride.
+{% set cmp = ">" %}
+{% set widen = scalar != "i32" and scalar != "u32" %}
+{% 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 }}
 @compute @workgroup_size({{ workgroupSize }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(workgroupSize, "output_index", "params.outputCount") }}
   let outer_index = output_index / params.innerSize;
   let inner_index = output_index % params.innerSize;
   let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
   var best_index = 0u;
+  var best = {{ "f32(x[input_base])" if widen else "x[input_base]" }};
   for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
+    let candidate = {{ "f32(x[input_base + axis_index * params.innerSize])" if widen else "x[input_base + axis_index * params.innerSize]" }};
     if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
       best = candidate;
       best_index = axis_index;
     }
   }
   y[output_index] = best_index;
 }

build/webgpu/reduce-arg-row-split.wgsl.jinja CHANGED Viewed

@@ -1,33 +1,26 @@
-{% set useSubgroups = useSubgroups %}
-{% if usesF16 %}
-enable f16;
-{%- endif %}
 {% if useSubgroups %}
 enable subgroups;
-{%- endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
-{%- if scalar == "i32" %}
-{% set work = "i32" %}
-{% elif scalar == "u32" %}
-{% set work = "u32" %}
-{% else %}
-{% set work = "f32" %}
-{% endif %}
-{% set cmp = ">" if modeSpec == "max" else "<" %}
 fn lane_identity() -> {{ work }} {
-{%- if scalar == "i32" %}
   return bitcast<i32>({{ "0x80000000u" if modeSpec == "max" else "0x7fffffffu" }});
-{%- elif scalar == "u32" %}
   return {{ "0u" if modeSpec == "max" else "0xffffffffu" }};
-{%- else %}
   var bits = {{ "0xff800000u" if modeSpec == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
-{%- endif %}
 }
 var<workgroup> wg_val: array<{{ work }}, WG>;
@@ -36,6 +29,7 @@ var<workgroup> wg_idx: array<u32, WG>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
         @builtin(subgroup_size) subgroup_size: u32{% endif %}) {
   let row = wg.x;
   let split_index = wg.y;
@@ -48,68 +42,59 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   var best_val = lane_identity();
   var best_idx = SENTINEL_IDX;
   for (var chunk = segment_begin + tid; chunk < segment_end; chunk += WG) {
-{%- if vec4 %}
     let value4 = x[base + chunk];
-{%- for component in ["x", "y", "z", "w"] %}
     {
-{%- if scalar == "f16" %}
-      let value = f32(value4.{{ component }});
-{%- else %}
-      let value = value4.{{ component }};
-{%- endif %}
       let index = chunk * 4u + {{ loop.index0 }}u;
-{%- if selectLastIndex %}
-      if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || index > best_idx))) {
-{%- else %}
-      if (value {{ cmp }} best_val || (value == best_val && index < best_idx)) {
-{%- endif %}
         best_val = value;
         best_idx = index;
       }
     }
-{%- endfor %}
 {% else %}
-{% if scalar == "f16" %}
-    let value = f32(x[base + chunk]);
-{%- else %}
-    let value = x[base + chunk];
-{%- endif %}
-{% if selectLastIndex %}
-    if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || chunk > best_idx))) {
-{%- else %}
-    if (value {{ cmp }} best_val || (value == best_val && chunk < best_idx)) {
-{%- endif %}
       best_val = value;
       best_idx = chunk;
     }
-{%- endif %}
   }
-{%- if useSubgroups %}
-  let subgroup_val = subgroupMax(best_val);{% if selectLastIndex %}
   let subgroup_candidate = select(0u, best_idx, best_val == subgroup_val && best_idx != SENTINEL_IDX);
   let subgroup_idx = subgroupMax(subgroup_candidate);
-{%- else %}
   let subgroup_candidate = select(SENTINEL_IDX, best_idx, best_val == subgroup_val);
   let subgroup_idx = subgroupMin(subgroup_candidate);
-{%- endif %}
   // The cross-subgroup fold assumes nothing about which invocations share a
-  // subgroup, how many subgroups there are, or which of a subgroup's lanes are
-  // active: every invocation owns the slot at its own index, the elected lane
   // publishes its subgroup pair there and every other lane publishes the
   // identity pair (the sentinel index, which the merge skips). Each subgroup
-  // then folds all WG slots — lane `rank`, its dense position among the active
-  // lanes, walks slots rank, rank + count, ... — and one more collective merges
-  // the lane partials, so every slot is merged exactly once at any legal width
-  // and partition. The merge (best value, then the tie-break index) is
-  // commutative and associative, so the fold order does not change the result.
   var output_val = subgroup_val;
   var output_idx = subgroup_idx;
   // A one-subgroup workgroup is already fully reduced by the collectives above.
   // The test reads the `subgroup_size` builtin, which is uniform; a collective's
   // result is not uniform to WGSL's analysis and may not guard a barrier.
   if (subgroup_size != WG) {
-    let rank = subgroupExclusiveAdd(1u);
-    let count = subgroupAdd(1u);
     let leader = rank == 0u;
     wg_val[tid] = select(lane_identity(), subgroup_val, leader);
     wg_idx[tid] = select(SENTINEL_IDX, subgroup_idx, leader);
@@ -119,33 +104,26 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     for (var slot = rank; slot < WG; slot = slot + count) {
       let candidate_val = wg_val[slot];
       let candidate_idx = wg_idx[slot];
-{%- if selectLastIndex %}
-      if (candidate_idx != SENTINEL_IDX && (fold_idx == SENTINEL_IDX || candidate_val {{ cmp }} fold_val || (candidate_val == fold_val && candidate_idx > fold_idx))) {
-{%- else %}
-      if (candidate_idx != SENTINEL_IDX && (fold_idx == SENTINEL_IDX || candidate_val {{ cmp }} fold_val || (candidate_val == fold_val && candidate_idx < fold_idx))) {
-{%- endif %}
         fold_val = candidate_val;
         fold_idx = candidate_idx;
       }
     }
-    output_val = subgroupMax(fold_val);{% if selectLastIndex %}
     let fold_candidate = select(0u, fold_idx, fold_val == output_val && fold_idx != SENTINEL_IDX);
     output_idx = subgroupMax(fold_candidate);
-{%- else %}
     let fold_candidate = select(SENTINEL_IDX, fold_idx, fold_val == output_val);
     output_idx = subgroupMin(fold_candidate);
-{%- endif %}
   }
   if (tid == 0u) {
     let scratch_index = split_index * params.rows + row;
-{%- if scalar == "u32" %}
-    partials_val[scratch_index] = output_val;
-{%- else %}
-    partials_val[scratch_index] = bitcast<u32>(output_val);
-{%- endif %}
     partials_idx[scratch_index] = output_idx;
   }
-{%- else %}
   wg_val[tid] = best_val;
   wg_idx[tid] = best_idx;
   workgroupBarrier();
@@ -154,11 +132,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
       let candidate_val = wg_val[tid + stride];
       let candidate_idx = wg_idx[tid + stride];
       let current_idx = wg_idx[tid];
-{%- if selectLastIndex %}
-      if (candidate_idx != SENTINEL_IDX && (current_idx == SENTINEL_IDX || candidate_val {{ cmp }} wg_val[tid] || (candidate_val == wg_val[tid] && candidate_idx > current_idx))) {
-{%- else %}
-      if (candidate_idx != SENTINEL_IDX && (current_idx == SENTINEL_IDX || candidate_val {{ cmp }} wg_val[tid] || (candidate_val == wg_val[tid] && candidate_idx < current_idx))) {
-{%- endif %}
         wg_val[tid] = candidate_val;
         wg_idx[tid] = candidate_idx;
       }
@@ -167,12 +141,8 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   }
   if (tid == 0u) {
     let scratch_index = split_index * params.rows + row;
-{%- if scalar == "u32" %}
-    partials_val[scratch_index] = wg_val[0];
-{%- else %}
-    partials_val[scratch_index] = bitcast<u32>(wg_val[0]);
-{%- endif %}
     partials_idx[scratch_index] = wg_idx[0];
   }
-{%- endif %}
 }

+{% set is_int = scalar == "i32" or scalar == "u32" %}
+{% set work = scalar if is_int else "f32" %}
+{% set sgReduce = "subgroupMax" %}
+{% macro arg_wins(v, i, best, bestIdx) %}{{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ "(" ~ bestIdx ~ " == SENTINEL_IDX || " ~ i ~ " > " ~ bestIdx ~ ")" if selectLastIndex else i ~ " < " ~ bestIdx }}){% endmacro %}
+{% macro arg_merge_wins(v, i, best, bestIdx) %}{{ i }} != SENTINEL_IDX && ({{ bestIdx }} == SENTINEL_IDX || {{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ i }} {{ ">" if selectLastIndex else "<" }} {{ bestIdx }})){% endmacro %}
 {% if useSubgroups %}
 enable subgroups;
+{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
 fn lane_identity() -> {{ work }} {
+{% if scalar == "i32" %}
   return bitcast<i32>({{ "0x80000000u" if modeSpec == "max" else "0x7fffffffu" }});
+{% elif scalar == "u32" %}
   return {{ "0u" if modeSpec == "max" else "0xffffffffu" }};
+{% else %}
   var bits = {{ "0xff800000u" if modeSpec == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+{% endif %}
 }
 var<workgroup> wg_val: array<{{ work }}, WG>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+        @builtin(subgroup_invocation_id) subgroup_lane: u32,
         @builtin(subgroup_size) subgroup_size: u32{% endif %}) {
   let row = wg.x;
   let split_index = wg.y;
   var best_val = lane_identity();
   var best_idx = SENTINEL_IDX;
   for (var chunk = segment_begin + tid; chunk < segment_end; chunk += WG) {
+{% if vec4 %}
     let value4 = x[base + chunk];
+{% for component in ["x", "y", "z", "w"] %}
     {
+      let value = {{ "f32(value4." ~ component ~ ")" if scalar == "f16" else "value4." ~ component }};
       let index = chunk * 4u + {{ loop.index0 }}u;
+      if ({{ arg_wins("value", "index", "best_val", "best_idx") }}) {
         best_val = value;
         best_idx = index;
       }
     }
+{% endfor %}
 {% else %}
+    let value = {{ "f32(x[base + chunk])" if scalar == "f16" else "x[base + chunk]" }};
+    if ({{ arg_wins("value", "chunk", "best_val", "best_idx") }}) {
       best_val = value;
       best_idx = chunk;
     }
+{% endif %}
   }
+{% if useSubgroups %}
+  let subgroup_val = {{ sgReduce }}(best_val);
+{% if selectLastIndex %}
   let subgroup_candidate = select(0u, best_idx, best_val == subgroup_val && best_idx != SENTINEL_IDX);
   let subgroup_idx = subgroupMax(subgroup_candidate);
+{% else %}
   let subgroup_candidate = select(SENTINEL_IDX, best_idx, best_val == subgroup_val);
   let subgroup_idx = subgroupMin(subgroup_candidate);
+{% endif %}
   // The cross-subgroup fold assumes nothing about which invocations share a
+  // subgroup or how many subgroups there are (it does require the uniform
+  // control flow this entry point already has): every invocation owns the slot
+  // at its own index, the elected lane
   // publishes its subgroup pair there and every other lane publishes the
   // identity pair (the sentinel index, which the merge skips). Each subgroup
+  // then folds all WG slots — lane `subgroup_lane` walks slots subgroup_lane,
+  // subgroup_lane + subgroup_size, ... — and one more collective merges the lane
+  // partials, so every slot is merged exactly once at any legal width. The merge
+  // (best value, then the tie-break index) is commutative and associative, so
+  // the fold order does not change the result.
   var output_val = subgroup_val;
   var output_idx = subgroup_idx;
   // A one-subgroup workgroup is already fully reduced by the collectives above.
   // The test reads the `subgroup_size` builtin, which is uniform; a collective's
   // result is not uniform to WGSL's analysis and may not guard a barrier.
   if (subgroup_size != WG) {
+    // The coordinates are the BUILTINS, never `subgroupExclusiveAdd(1u)` /
+    // `subgroupAdd(1u)`: those agree with them under this contract, but a driver
+    // in the wild answers a claim derived from them with zero and leaves most of
+    // the workgroup's slots unclaimed. Every lane of this entry point is active
+    // here, so the full subgroup width is the right stride.
+    let rank = subgroup_lane;
+    let count = subgroup_size;
     let leader = rank == 0u;
     wg_val[tid] = select(lane_identity(), subgroup_val, leader);
     wg_idx[tid] = select(SENTINEL_IDX, subgroup_idx, leader);
     for (var slot = rank; slot < WG; slot = slot + count) {
       let candidate_val = wg_val[slot];
       let candidate_idx = wg_idx[slot];
+      if ({{ arg_merge_wins("candidate_val", "candidate_idx", "fold_val", "fold_idx") }}) {
         fold_val = candidate_val;
         fold_idx = candidate_idx;
       }
     }
+    output_val = {{ sgReduce }}(fold_val);
+{% if selectLastIndex %}
     let fold_candidate = select(0u, fold_idx, fold_val == output_val && fold_idx != SENTINEL_IDX);
     output_idx = subgroupMax(fold_candidate);
+{% else %}
     let fold_candidate = select(SENTINEL_IDX, fold_idx, fold_val == output_val);
     output_idx = subgroupMin(fold_candidate);
+{% endif %}
   }
   if (tid == 0u) {
     let scratch_index = split_index * params.rows + row;
+    partials_val[scratch_index] = {{ "output_val" if scalar == "u32" else "bitcast<u32>(output_val)" }};
     partials_idx[scratch_index] = output_idx;
   }
+{% else %}
   wg_val[tid] = best_val;
   wg_idx[tid] = best_idx;
   workgroupBarrier();
       let candidate_val = wg_val[tid + stride];
       let candidate_idx = wg_idx[tid + stride];
       let current_idx = wg_idx[tid];
+      if ({{ arg_merge_wins("candidate_val", "candidate_idx", "wg_val[tid]", "current_idx") }}) {
         wg_val[tid] = candidate_val;
         wg_idx[tid] = candidate_idx;
       }
   }
   if (tid == 0u) {
     let scratch_index = split_index * params.rows + row;
+    partials_val[scratch_index] = {{ "wg_val[0]" if scalar == "u32" else "bitcast<u32>(wg_val[0])" }};
     partials_idx[scratch_index] = wg_idx[0];
   }
+{% endif %}
 }

build/webgpu/reduce-arg-row-subgroup.wgsl.jinja CHANGED Viewed

@@ -11,148 +11,135 @@
 // integer bounds; ties against the sentinel use the same index rule, so a row
 // filled with the sentinel still produces the correct index. NaNs are never
 // accepted, and an all-NaN row returns index zero.
-{% if useSubgroups is defined %}{% set useSubgroups = useSubgroups %}{% else %}{% set useSubgroups = true %}{% endif %}
-{% if usesF16 %}
-enable f16;
-{%- endif %}
 {% if useSubgroups %}
 enable subgroups;
-{%- endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
-{% if modeSpec == "max" and scalar != "i32" and scalar != "u32" %}
-fn neg_inf_f32() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 {% endif %}
-{%- if scalar == "i32" %}
-{% set cmpType = "i32" %}
-{% elif scalar == "u32" %}
-{% set cmpType = "u32" %}
-{% else %}
-{% set cmpType = "f32" %}
-{% endif %}
-{% set cmp = ">" %}
-{%- if cmpType == "i32" %}
-{% set identityVal = "-2147483647i - 1i" %}
-{%- elif cmpType == "u32" %}
-{% set identityVal = "0u" %}
-{%- else %}
-{% set identityVal = "neg_inf_f32()" %}
-{%- endif %}
-{%- if not (useSubgroups and singleSubgroup is defined and singleSubgroup) %}
 var<workgroup> wgVal: array<{{ cmpType }}, WG>;
 var<workgroup> wgIdx: array<u32, WG>;
-{%- endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups and not (singleSubgroup is defined and singleSubgroup) %},
         @builtin(subgroup_size) sgSize: u32{% endif %}) {
   let row = wg.x + wg.y * params.rowStride;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
-  let base = row * params.chunkCount;{% if cmpType == "i32" %}
-  var bestVal: i32 = -2147483647i - 1i;
-{%- elif cmpType == "u32" %}
-  var bestVal: u32 = 0u;
-{%- else %}
-  var bestVal: f32 = neg_inf_f32();
-{%- endif %}
   var bestIdx: u32 = SENTINEL_IDX;
-{%- if vec4 %}
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v4 = x[base + c];
-{%- for comp in ["x", "y", "z", "w"] %}
     {
-{%- if scalar == "f16" %}
-      let v = f32(v4.{{ comp }});
-{%- else %}
-      let v = v4.{{ comp }};
-{%- endif %}
       let idx = c * 4u + {{ loop.index0 }}u;
-{%- if selectLastIndex %}
-      if (v {{ cmp }} bestVal || (v == bestVal && (bestIdx == SENTINEL_IDX || idx > bestIdx))) {
-{%- else %}
-      if (v {{ cmp }} bestVal || (v == bestVal && idx < bestIdx)) {
-{%- endif %}
         bestVal = v;
         bestIdx = idx;
       }
     }
-{%- endfor %}
   }
-{%- else %}
-  // Four independent, lane-coalesced loads per loop iteration. Keeping each
-  // unrolled plane WG elements apart preserves contiguous accesses across the
-  // workgroup while exposing enough ILP for long odd-width rows (where a vec4
-  // storage binding cannot represent row boundaries).
-  for (var c0 = tid; c0 < params.chunkCount; c0 = c0 + 4u * WG) {
-{%- for plane in [0, 1, 2, 3] %}
     {
-      let c = c0 + {{ plane }}u * WG;
       if (c < params.chunkCount) {
-{%- if scalar == "f16" %}
-        let v = f32(x[base + c]);
-{%- else %}
-        let v = x[base + c];
-{%- endif %}
-{% if selectLastIndex %}
-        if (v {{ cmp }} bestVal || (v == bestVal && (bestIdx == SENTINEL_IDX || c > bestIdx))) {
-{%- else %}
-        if (v {{ cmp }} bestVal || (v == bestVal && c < bestIdx)) {
-{%- endif %}
           bestVal = v;
           bestIdx = c;
         }
       }
     }
-{%- endfor %}
   }
-{%- endif %}
-{%- if useSubgroups %}
-  let m = subgroupMax(bestVal);{% if selectLastIndex %}
   // Lanes that match the subgroup best contribute their (real) index; the
   // rest contribute 0 so subgroupMax picks the highest matching index.
   let cand = select(0u, bestIdx, bestVal == m && bestIdx != SENTINEL_IDX);
   let sgIdx = subgroupMax(cand);
-{%- else %}
   // Lanes that match the subgroup best contribute their index; the rest
   // contribute the sentinel so subgroupMin picks the lowest matching index.
   let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
   let sgIdx = subgroupMin(cand);
-{%- endif %}
-{% if singleSubgroup is defined and singleSubgroup %}
   // The workgroup is one (possibly partial) subgroup: elect the writer rather
   // than assuming lane 0 is active.
   if (subgroupElect()) {
     y[row] = select(sgIdx, 0u, sgIdx == SENTINEL_IDX);
   }
-{%- else %}
   // Cross-subgroup fold that assumes nothing about which invocations share a
-  // subgroup, how many subgroups there are, or which of a subgroup's lanes are
-  // active: every invocation owns the slot at its own index, the elected lane
   // publishes its subgroup pair there and every other lane publishes the
   // identity pair (the sentinel index, which the merge skips). Each subgroup
-  // then folds all WG slots — lane `rank`, its dense position among the active
-  // lanes, walks slots rank, rank + count, ... — and one more collective merges
-  // the lane partials, so every slot is merged exactly once at any legal width
-  // and partition. The merge is commutative and associative (best value, then
-  // the tie-break index), so the fold order does not change the result.
   var winnerIdx = sgIdx;
   // A one-subgroup workgroup is already fully reduced by the collectives above.
   // The test reads the `subgroup_size` builtin, which is uniform; a collective's
   // result is not uniform to WGSL's analysis and may not guard a barrier.
   if (sgSize != WG) {
-    let rank = subgroupExclusiveAdd(1u);
-    let count = subgroupAdd(1u);
     let leader = rank == 0u;
     wgVal[tid] = select({{ identityVal }}, m, leader);
     wgIdx[tid] = select(SENTINEL_IDX, sgIdx, leader);
@@ -162,25 +149,22 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     for (var i = rank; i < WG; i = i + count) {
       let v = wgVal[i];
       let vi = wgIdx[i];
-{%- if selectLastIndex %}
-      if (vi != SENTINEL_IDX && (v {{ cmp }} outVal || (v == outVal && (outIdx == SENTINEL_IDX || vi > outIdx)))) {
-{%- else %}
-      if (vi != SENTINEL_IDX && (v {{ cmp }} outVal || (v == outVal && vi < outIdx))) {
-{%- endif %}
         outVal = v;
         outIdx = vi;
       }
     }
-    let foldVal = subgroupMax(outVal);{% if selectLastIndex %}
     winnerIdx = subgroupMax(select(0u, outIdx, outVal == foldVal && outIdx != SENTINEL_IDX));
-{%- else %}
     winnerIdx = subgroupMin(select(SENTINEL_IDX, outIdx, outVal == foldVal));
-{%- endif %}
   }
   if (tid == 0u) {
     y[row] = select(winnerIdx, 0u, winnerIdx == SENTINEL_IDX);
   }
-{%- endif %}
 {% else %}
   // No-subgroup tier: workgroup barrier tree-reduction over the per-thread
   // (bestVal, bestIdx) pairs. Same ONNX tie rule (lowest index by default,
@@ -195,11 +179,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
       let v = wgVal[tid + stride];
       let vi = wgIdx[tid + stride];
       let curIdx = wgIdx[tid];
-{%- if selectLastIndex %}
-      if (vi != SENTINEL_IDX && (curIdx == SENTINEL_IDX || v {{ cmp }} wgVal[tid] || (v == wgVal[tid] && vi > curIdx))) {
-{%- else %}
-      if (vi != SENTINEL_IDX && (curIdx == SENTINEL_IDX || v {{ cmp }} wgVal[tid] || (v == wgVal[tid] && vi < curIdx))) {
-{%- endif %}
         wgVal[tid] = v;
         wgIdx[tid] = vi;
       }
@@ -209,5 +189,5 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   if (tid == 0u) {
     y[row] = select(wgIdx[0], 0u, wgIdx[0] == SENTINEL_IDX);
   }
-{%- endif %}
 }

 // integer bounds; ties against the sentinel use the same index rule, so a row
 // filled with the sentinel still produces the correct index. NaNs are never
 // accepted, and an all-NaN row returns index zero.
+{% set useSubgroups = useSubgroups | default(true) %}
+{% set singleSubgroup = useSubgroups and singleSubgroup | default(false) %}
+{% set adjacent = adjacentScalarLoads | default(false) %}
+{% set cmpType = scalar if scalar == "i32" or scalar == "u32" else "f32" %}
+{% set sgReduce = "subgroupMax" %}
+{% if cmpType == "i32" %}
+{% set identityVal = "-2147483647i - 1i" %}
+{% elif cmpType == "u32" %}
+{% set identityVal = "0u" %}
+{% else %}
+{% set identityVal = "negative_infinity()" %}
+{% endif %}
+{% macro arg_wins(v, i, best, bestIdx) %}{{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ "(" ~ bestIdx ~ " == SENTINEL_IDX || " ~ i ~ " > " ~ bestIdx ~ ")" if selectLastIndex else i ~ " < " ~ bestIdx }}){% endmacro %}
+{% macro arg_merge_wins(v, i, best, bestIdx) %}{{ i }} != SENTINEL_IDX && ({{ bestIdx }} == SENTINEL_IDX || {{ v }} {{ ">" if modeSpec == "max" else "<" }} {{ best }} || ({{ v }} == {{ best }} && {{ i }} {{ ">" if selectLastIndex else "<" }} {{ bestIdx }})){% endmacro %}
 {% if useSubgroups %}
 enable subgroups;
+{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
+{% if cmpType == "f32" and modeSpec == "max" %}
+fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
+{% elif cmpType == "f32" %}
+fn pos_inf_f32() -> f32 {
+  var bits = 0x7f800000u;
+  return bitcast<f32>(bits);
+}
 {% endif %}
+{% if not singleSubgroup %}
 var<workgroup> wgVal: array<{{ cmpType }}, WG>;
 var<workgroup> wgIdx: array<u32, WG>;
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups and not singleSubgroup %},
+        @builtin(subgroup_invocation_id) sgLane: u32,
         @builtin(subgroup_size) sgSize: u32{% endif %}) {
   let row = wg.x + wg.y * params.rowStride;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
+  let base = row * params.chunkCount;
+  var bestVal: {{ cmpType }} = {{ identityVal }};
   var bestIdx: u32 = SENTINEL_IDX;
+{% if vec4 %}
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v4 = x[base + c];
+{% for comp in ["x", "y", "z", "w"] %}
     {
+      let v = {{ "f32(v4." ~ comp ~ ")" if scalar == "f16" else "v4." ~ comp }};
       let idx = c * 4u + {{ loop.index0 }}u;
+      if ({{ arg_wins("v", "idx", "bestVal", "bestIdx") }}) {
         bestVal = v;
         bestIdx = idx;
       }
     }
+{% endfor %}
   }
+{% else %}
+  // Scan four bounded scalar values per lane. For long rows the manifest may
+  // assign adjacent values to a lane; otherwise the four planes stay WG apart.
+  // Both layouts cover each input exactly once and use the same tie comparison.
+  for (var c0 = {{ "4u * tid" if adjacent else "tid" }}; c0 < params.chunkCount; c0 = c0 + 4u * WG) {
+{% for plane in [0, 1, 2, 3] %}
     {
+      let c = c0 + {{ plane }}u{{ "" if adjacent else " * WG" }};
       if (c < params.chunkCount) {
+        let v = {{ "f32(x[base + c])" if scalar == "f16" else "x[base + c]" }};
+        if ({{ arg_wins("v", "c", "bestVal", "bestIdx") }}) {
           bestVal = v;
           bestIdx = c;
         }
       }
     }
+{% endfor %}
   }
+{% endif %}
+{% if useSubgroups %}
+  let m = {{ sgReduce }}(bestVal);
+{% if selectLastIndex %}
   // Lanes that match the subgroup best contribute their (real) index; the
   // rest contribute 0 so subgroupMax picks the highest matching index.
   let cand = select(0u, bestIdx, bestVal == m && bestIdx != SENTINEL_IDX);
   let sgIdx = subgroupMax(cand);
+{% else %}
   // Lanes that match the subgroup best contribute their index; the rest
   // contribute the sentinel so subgroupMin picks the lowest matching index.
   let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
   let sgIdx = subgroupMin(cand);
+{% endif %}
+{% if singleSubgroup %}
   // The workgroup is one (possibly partial) subgroup: elect the writer rather
   // than assuming lane 0 is active.
   if (subgroupElect()) {
     y[row] = select(sgIdx, 0u, sgIdx == SENTINEL_IDX);
   }
+{% else %}
   // Cross-subgroup fold that assumes nothing about which invocations share a
+  // subgroup or how many subgroups there are (it does require the uniform
+  // control flow this entry point already has): every invocation owns the slot
+  // at its own index, the elected lane
   // publishes its subgroup pair there and every other lane publishes the
   // identity pair (the sentinel index, which the merge skips). Each subgroup
+  // then folds all WG slots — lane `sgLane` walks slots sgLane, sgLane + sgSize,
+  // ... — and one more collective merges the lane partials, so every slot is
+  // merged exactly once at any legal width. The merge is commutative and
+  // associative (best value, then the tie-break index), so the fold order does
+  // not change the result.
   var winnerIdx = sgIdx;
   // A one-subgroup workgroup is already fully reduced by the collectives above.
   // The test reads the `subgroup_size` builtin, which is uniform; a collective's
   // result is not uniform to WGSL's analysis and may not guard a barrier.
   if (sgSize != WG) {
+    // The coordinates are the BUILTINS, never `subgroupExclusiveAdd(1u)` /
+    // `subgroupAdd(1u)`: those agree with them under this contract, but a driver
+    // in the wild answers a claim derived from them with zero and leaves most of
+    // the workgroup's slots unclaimed. Every lane of this entry point is active
+    // here, so the full subgroup width is the right stride.
+    let rank = sgLane;
+    let count = sgSize;
     let leader = rank == 0u;
     wgVal[tid] = select({{ identityVal }}, m, leader);
     wgIdx[tid] = select(SENTINEL_IDX, sgIdx, leader);
     for (var i = rank; i < WG; i = i + count) {
       let v = wgVal[i];
       let vi = wgIdx[i];
+      if (vi != SENTINEL_IDX && ({{ arg_wins("v", "vi", "outVal", "outIdx") }})) {
         outVal = v;
         outIdx = vi;
       }
     }
+    let foldVal = {{ sgReduce }}(outVal);
+{% if selectLastIndex %}
     winnerIdx = subgroupMax(select(0u, outIdx, outVal == foldVal && outIdx != SENTINEL_IDX));
+{% else %}
     winnerIdx = subgroupMin(select(SENTINEL_IDX, outIdx, outVal == foldVal));
+{% endif %}
   }
   if (tid == 0u) {
     y[row] = select(winnerIdx, 0u, winnerIdx == SENTINEL_IDX);
   }
+{% endif %}
 {% else %}
   // No-subgroup tier: workgroup barrier tree-reduction over the per-thread
   // (bestVal, bestIdx) pairs. Same ONNX tie rule (lowest index by default,
       let v = wgVal[tid + stride];
       let vi = wgIdx[tid + stride];
       let curIdx = wgIdx[tid];
+      if ({{ arg_merge_wins("v", "vi", "wgVal[tid]", "curIdx") }}) {
         wgVal[tid] = v;
         wgIdx[tid] = vi;
       }
   if (tid == 0u) {
     y[row] = select(wgIdx[0], 0u, wgIdx[0] == SENTINEL_IDX);
   }
+{% endif %}
 }

build/webgpu/test.json CHANGED Viewed

@@ -864,9 +864,7 @@
     },
     {
       "name": "subgroup_min_vec4_rows128x256_ties_first",
-      "provenance": {
-        "notes": "Route-locks the one-minimum-subgroup vec4 path on variable-width adapters and verifies first-index tie ordering without a workgroup-memory combine."
-      },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [128, 256], "data": { "kind": "constant", "value": 7.0 } } },
       "outputs": {
@@ -875,9 +873,7 @@
     },
     {
       "name": "subgroup_min_scalar_rows128x257_ties_last",
-      "provenance": {
-        "notes": "Route-locks the unaligned one-minimum-subgroup path and verifies select_last_index ordering across all lane-local chunks."
-      },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [128, 257], "data": { "kind": "constant", "value": 7.0 } } },
       "outputs": {
@@ -1122,7 +1118,7 @@
     {
       "name": "last_axis_split_vec4_logits_ties_first",
       "provenance": {
-        "notes": "A single 32768-element row route-locks the cooperative split-row vec4 reduction. Equal maxima span many split segments and must merge to the first global index."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -1146,7 +1142,7 @@
     {
       "name": "last_axis_split_vec4_i32_negative_rows",
       "provenance": {
-        "notes": "int32 on the cooperative split-row vec4 reduction (32768 columns, two rows). Every value in the row is negative, so a lane identity of 0 instead of INT_MIN would beat the whole row and the merge would return the sentinel; this directly exercises the i32 identity and native signed compare. The period-7 cycle does not divide 32768, so the two rows are phase-shifted and their answers differ — a kernel that reduced row 0 for both rows cannot pass. ONNX ArgMax admits int32 via T and breaks ties on the first index."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -1163,7 +1159,7 @@
     {
       "name": "last_axis_split_scalar_u32_above_int32_max",
       "provenance": {
-        "notes": "uint32 on the unaligned (32769-column) scalar split-row reduction. Six of the seven cycle values exceed 2^31, so a signed compare would rank them negative and answer [0, 5] instead of [3, 1]. This directly exercises the unsigned identity 0u and the partial store that must write the value through without bitcasting it. ONNX ArgMax admits uint32 via T; first index on ties."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -1264,7 +1260,7 @@
     {
       "name": "subgroup_min_scalar_keepdims_128x197_vit_tokens",
       "provenance": {
-        "notes": "With 128 rows, the 197-element last axis satisfies the subgroup row floor but is not vec4-aligned, selecting the scalar subgroup-minimum reduction with default keepdims."
       },
       "attrs": { "axis": -1 },
       "inputs": {
@@ -1275,6 +1271,122 @@
         }
       },
       "outputs": { "y": { "dtype": "uint32", "shape": [128, 1], "tolerance": 0 } }
     }
   ]
 }

     },
     {
       "name": "subgroup_min_vec4_rows128x256_ties_first",
+      "provenance": { "notes": "Aligned 256-element rows with equal maxima verify first-index tie ordering." },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [128, 256], "data": { "kind": "constant", "value": 7.0 } } },
       "outputs": {
     },
     {
       "name": "subgroup_min_scalar_rows128x257_ties_last",
+      "provenance": { "notes": "Unaligned 257-element rows with equal maxima verify last-index tie ordering." },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [128, 257], "data": { "kind": "constant", "value": 7.0 } } },
       "outputs": {
     {
       "name": "last_axis_split_vec4_logits_ties_first",
       "provenance": {
+        "notes": "A 32,768-element row has equal maxima across distant segments; the result must be the first global index."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
     {
       "name": "last_axis_split_vec4_i32_negative_rows",
       "provenance": {
+        "notes": "A [2, 32768] int32 tensor's last axis cycles seven negative values (-9, -5, -3, -1, -2, -4, -6), so a zero-initialized accumulator would misreport the maximum. Since 32768 is not a multiple of 7, row 1 begins one step into the cycle, putting its first-occurrence maximum at column 2 versus row 0's column 3."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
     {
       "name": "last_axis_split_scalar_u32_above_int32_max",
       "provenance": {
+        "notes": "A [2, 32769] uint32 tensor's last axis cycles seven values; six of seven exceed 2^31 (up to 4,294,967,290, near UINT32_MAX), so a signed comparison would misorder them. Since 32769 is not a multiple of 7, row 1 begins two steps into the cycle, moving its maximum's column from 3 (row 0) to 1 (row 1)."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
     {
       "name": "subgroup_min_scalar_keepdims_128x197_vit_tokens",
       "provenance": {
+        "notes": "A 197-element last axis across 128 rows checks an unaligned reduction and default keepdims."
       },
       "attrs": { "axis": -1 },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "uint32", "shape": [128, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_wide_last_axis_2x202048_keepdims_negative_axis",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_float_wide_last_axis_keepdims",
+        "notes": "Vocabulary-scale last axis with keepdims and a negative axis; upstream added it so the CUDA multi-block scratch path is covered end to end. 2 x 202048 float32 = 1.6 MB."
+      },
+      "attrs": { "axis": -1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 202048],
+          "data": {
+            "kind": "cycle",
+            "values": [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5, 12.0, 12.5, 13.0, 13.5, 14.0, 14.5, 15.0, 15.5, 16.0, 16.5, 17.0, 17.5, 18.0, 18.5, 19.0, 19.5, 20.0, 20.5, 21.0, 21.5, 22.0, 22.5, 23.0, 23.5, 24.0, 24.5, 25.0, 25.5, 26.0, 26.5, 27.0, 27.5, 28.0, 28.5, 29.0, 29.5, 30.0, 30.5, 31.0, 31.5, 32.0, 32.5, 33.0, 33.5, 34.0, 34.5, 35.0, 35.5, 36.0, 36.5, 37.0, 37.5, 38.0, 38.5, 39.0, 39.5, 40.0, 40.5, 41.0, 41.5, 42.0, 42.5, 43.0, 43.5, 44.0, 44.5, 45.0, 45.5, 46.0, 46.5, 47.0, 47.5, 48.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "packed_axis_columns_8192x32_ties_first",
+      "provenance": {
+        "notes": "A [8192, 32] float32 tensor's axis=-2 cycles a 13-value pattern whose maximum, 8.0, repeats twice per period; every column sees that tie across its 8192 rows, and select_last_index=0 resolves each to the earliest matching row."
+      },
+      "tunables": { "WORKGROUP_SIZE": 16 },
+      "attrs": { "axis": -2, "keepdims": 0, "select_last_index": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8192, 32],
+          "data": { "kind": "cycle", "values": [-3.0, 7.0, 7.0, -9.0, 2.0, 8.0, -9.0, 8.0, -1.0, 0.0, 2.0, 7.0, -9.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "packed_axis_columns_8192x32_ties_last",
+      "provenance": {
+        "notes": "A [8192, 32] float32 tensor's axis=-2 cycles a 13-value pattern whose maximum, 8.0, repeats twice per period; every column sees that tie across its 8192 rows, and select_last_index=1 resolves each to the latest matching row."
+      },
+      "tunables": { "WORKGROUP_SIZE": 16 },
+      "attrs": { "axis": -2, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8192, 32],
+          "data": { "kind": "cycle", "values": [-3.0, 7.0, 7.0, -9.0, 2.0, 8.0, -9.0, 8.0, -1.0, 0.0, 2.0, 7.0, -9.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "packed_axis_columns_8193x28_ties_first",
+      "provenance": {
+        "notes": "A [8193, 28] float32 tensor's axis=-2 (one row more than the 8192-row case) cycles the same 13-value pattern, whose maximum 8.0 repeats twice per period; keepdims=1 keeps that axis as size 1 in a [1, 28] output, and select_last_index=0 resolves each column's tie to its earliest row."
+      },
+      "tunables": { "WORKGROUP_SIZE": 8 },
+      "attrs": { "axis": -2, "keepdims": 1, "select_last_index": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8193, 28],
+          "data": { "kind": "cycle", "values": [-3.0, 7.0, 7.0, -9.0, 2.0, 8.0, -9.0, 8.0, -1.0, 0.0, 2.0, 7.0, -9.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 28], "tolerance": 0 } }
+    },
+    {
+      "name": "packed_axis_columns_8193x28_ties_last",
+      "provenance": {
+        "notes": "A [8193, 28] float32 tensor's axis=-2 (one row more than the 8192-row case) cycles the same 13-value pattern, whose maximum 8.0 repeats twice per period; keepdims=1 keeps that axis as size 1 in a [1, 28] output, and select_last_index=1 resolves each column's tie to its latest row."
+      },
+      "tunables": { "WORKGROUP_SIZE": 8 },
+      "attrs": { "axis": -2, "keepdims": 1, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8193, 28],
+          "data": { "kind": "cycle", "values": [-3.0, 7.0, 7.0, -9.0, 2.0, 8.0, -9.0, 8.0, -1.0, 0.0, 2.0, 7.0, -9.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 28], "tolerance": 0 } }
+    },
+    {
+      "name": "packed_axis_columns_2x8193x12_ties_first",
+      "provenance": {
+        "notes": "A [2, 8193, 12] float32 tensor's axis=-2 (the 8193-length middle axis) cycles the same 13-value pattern for each of the 2 batches' 12 columns; the maximum 8.0 repeats twice per period, and select_last_index=0 resolves each tie to its earliest row, giving a [2, 12] output."
+      },
+      "tunables": { "WORKGROUP_SIZE": 8 },
+      "attrs": { "axis": -2, "keepdims": 0, "select_last_index": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8193, 12],
+          "data": { "kind": "cycle", "values": [-3.0, 7.0, 7.0, -9.0, 2.0, 8.0, -9.0, 8.0, -1.0, 0.0, 2.0, 7.0, -9.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 12], "tolerance": 0 } }
+    },
+    {
+      "name": "packed_axis_columns_2x8193x12_ties_last",
+      "provenance": {
+        "notes": "A [2, 8193, 12] float32 tensor's axis=-2 (the 8193-length middle axis) cycles the same 13-value pattern for each of the 2 batches' 12 columns; the maximum 8.0 repeats twice per period, and select_last_index=1 resolves each tie to its latest row, giving a [2, 12] output."
+      },
+      "tunables": { "WORKGROUP_SIZE": 8 },
+      "attrs": { "axis": -2, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8193, 12],
+          "data": { "kind": "cycle", "values": [-3.0, 7.0, 7.0, -9.0, 2.0, 8.0, -9.0, 8.0, -1.0, 0.0, 2.0, 7.0, -9.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 12], "tolerance": 0 } }
     }
   ]
 }