sync 91d990483a17

Browse files

Files changed (12) hide show

README.md +11 -7
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +198 -284
build/webgpu/metadata.json +27 -13
build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja +3 -4
build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja +9 -10
build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja +4 -4
build/webgpu/reduce-arg-axis-tiled.wgsl.jinja +9 -10
build/webgpu/reduce-arg-axis.wgsl.jinja +4 -5
build/webgpu/reduce-arg-row-split.wgsl.jinja +44 -25
build/webgpu/reduce-arg-row-subgroup.wgsl.jinja +55 -27
build/webgpu/test.json +35 -38

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `ArgMin` spec](https://onnx.ai/onnx/operators/onnx__ArgMin.html) f
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `data` | `x` | `T` | — | — | Values whose minimum index is selected along `axis`. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `reduced` | `y` | `I` | `uint32` | derived | derived; see description | Logical int64 indices of the minimum values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
 ## Attributes
@@ -51,7 +51,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -65,10 +65,14 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `data` | `T` | — | — | Values whose minimum index is selected along `axis`. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `reduced` | `I` | `uint32` | derived | derived | Logical int64 indices of the minimum values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
 ## Attributes
 ## Files
+- [`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
 ## 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.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.ArgMin",
   "cases": [
     {
       "name": "argmin-axis0-f32-1024x512",

 {
   "cases": [
     {
       "name": "argmin-axis0-f32-1024x512",

build/webgpu/manifest.json CHANGED Viewed

@@ -2,57 +2,47 @@
   "domain": "ai.onnx",
   "name": "ArgMin",
   "sinceVersion": 13,
-  "description": "Returns the index of the minimum value along an axis, choosing the first equal value unless `select_last_index` is enabled.",
-  "inputs": [{ "role": "data", "dtype": "T", "description": "Values whose minimum index is selected along `axis`." }],
-  "outputs": [
-    {
-      "role": "reduced",
       "dtype": "I",
-      "rank": "ranks.data if attrs.keepdims == 1 else ranks.data - 1",
-      "description": "Logical int64 indices of the minimum values along the reduced axis; WebGPU stores these bounded indices as uint32.",
-      "shape": "prefix(shapes.data, normalizedAxis) + ([1] if attrs.keepdims == 1 else []) + suffix(shapes.data, normalizedAxis + 1)"
     }
-  ],
-  "attributes": { "axis": 0, "keepdims": 1, "select_last_index": 0 },
-  "attributeDescriptions": {
-    "axis": "Axis to reduce; negative values count from the back.",
-    "keepdims": "Retain the reduced dimension with length one when non-zero.",
-    "select_last_index": "Choose the last equal minimum instead of the first when non-zero."
   },
   "attributeConstraints": { "keepdims": { "values": [0, 1] }, "select_last_index": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"], "I": ["int64"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "data", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "reduced", "role": "output", "dtype": "uint32" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "TILE_COLUMNS": 16,
-    "TILED_MIN_AXIS": 64,
-    "TILED_MIN_OUTPUTS": 16,
-    "SPLIT_MIN_AXIS": 8192,
-    "SPLIT_MAX_OUTPUTS": 4096,
-    "SPLIT_TARGET_AXIS": 256,
-    "SPLIT_TILE_COLUMNS": 8,
-    "MAX_SPLITS": 128,
-    "ROW_SPLIT_MIN_AXIS": 32768,
-    "ROW_SPLIT_MAX_OUTPUTS": 32,
-    "ROW_SPLIT_TARGET_AXIS": 2048,
-    "ROW_SPLIT_MAX_SPLITS": 64
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
-    "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
-    "normalizedAxis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data",
-    "axisInRange": "ranks.data >= 1 and normalizedAxis >= 0 and normalizedAxis < ranks.data",
-    "axisDim": "dim(shapes.data, normalizedAxis)",
-    "axisInner": "inner(shapes.data, normalizedAxis)",
-    "outputCount": "numel(shapes.reduced)",
-    "outputRankOk": "(attrs.keepdims == 1 and ranks.reduced == ranks.data) or (attrs.keepdims == 0 and ranks.reduced == ranks.data - 1)",
-    "outputShapeOk": "outputRankOk and outputCount * axisDim == numel(shapes.data)",
     "baseContract": "axisInRange and axisDim > 0 and outputShapeOk and f16Ok(dtypes.T)",
-    "lastAxis": "normalizedAxis == ranks.data - 1",
     "nativeParallelType": "dtypes.T == \"f32\" or dtypes.T == \"i32\" or dtypes.T == \"u32\"",
     "rowParallelType": "nativeParallelType or dtypes.T == \"f16\"",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
@@ -65,190 +55,84 @@
     "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.SPLIT_TARGET_AXIS)))",
     "splitScratchBytes": "splitCount * outputCount * 8",
     "splitScratchFits": "splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize",
-    "splitDispatchFits": "splitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(outputCount, workgroupSize), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
-    "splitTileDispatchFits": "splitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS) <= device.limits.maxComputeWorkgroupsPerDimension",
     "rowSplitCount": "min(tunables.ROW_SPLIT_MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.ROW_SPLIT_TARGET_AXIS)))",
     "rowSplitScratchBytes": "rowSplitCount * outputCount * 8",
     "rowSplitScratchFits": "rowSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and rowSplitScratchBytes <= device.limits.maxBufferSize",
-    "rowSplitDispatchFits": "outputCount <= device.limits.maxComputeWorkgroupsPerDimension and rowSplitCount <= device.limits.maxComputeWorkgroupsPerDimension",
     "tileDispatchFits": "ceilDiv(outputCount, tunables.TILE_COLUMNS) <= foldedDispatchCapacity"
   },
-  "bindingSets": {
-    "contiguousVec4": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputCount" },
-            { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(outputCount, device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
-        }
-      }
-    ],
-    "contiguousScalar": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputCount" },
-            { "name": "chunkCount", "type": "u32", "value": "axisDim" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(outputCount, device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
-        }
-      }
-    ],
-    "axisGeometry": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "axisDim", "type": "u32", "value": "axisDim" },
-            { "name": "outputCount", "type": "u32", "value": "outputCount" },
-            { "name": "innerSize", "type": "u32", "value": "axisInner" }
-          ]
-        }
-      }
-    ],
-    "rowSplitVec4": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "partials_val", "semantic": "partials_val", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "partials_idx", "semantic": "partials_idx", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputCount" },
-            { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" }
-          ]
-        }
-      }
-    ],
-    "rowSplitScalar": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "partials_val", "semantic": "partials_val", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "partials_idx", "semantic": "partials_idx", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputCount" },
-            { "name": "chunkCount", "type": "u32", "value": "axisDim" }
-          ]
         }
-      }
-    ],
-    "splitReduce": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "partials_val", "semantic": "partials_val", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "partials_idx", "semantic": "partials_idx", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "axisDim", "type": "u32", "value": "axisDim" },
-            { "name": "outputCount", "type": "u32", "value": "outputCount" },
-            { "name": "innerSize", "type": "u32", "value": "axisInner" }
-          ]
         }
-      }
-    ],
-    "splitCombine": [
-      {
-        "name": "partials_val",
-        "semantic": "partials_val",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "partials_idx",
-        "semantic": "partials_idx",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "outputCount", "type": "u32", "value": "outputCount" }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "last_axis_split_vec4",
       "priority": 41,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= tunables.ROW_SPLIT_MIN_AXIS", "axisDim % 4 == 0", "outputCount <= tunables.ROW_SPLIT_MAX_OUTPUTS", "rowSplitScratchFits", "rowSplitDispatchFits", "rowStorageFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "usesF16": "dtypes.T == \"f16\"",
@@ -264,19 +148,23 @@
         {
           "id": "split_reduce",
           "name": "ArgMin.LastAxisSplitVec4Reduce",
-          "source": {
-            "shader": "reduce-arg-row-split.wgsl.jinja",
-            "inputs": { "mode": "\"min\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" }
-          },
-          "bindings": "rowSplitVec4",
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.LastAxisSplitCombine",
-          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "splitCombine",
-          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
         }
       ]
     },
@@ -284,7 +172,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"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "selectLastIndex": "attrs.select_last_index != 0",
@@ -299,19 +187,23 @@
         {
           "id": "split_reduce",
           "name": "ArgMin.LastAxisSplitReduce",
-          "source": {
-            "shader": "reduce-arg-row-split.wgsl.jinja",
-            "inputs": { "mode": "\"min\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" }
-          },
-          "bindings": "rowSplitScalar",
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.LastAxisSplitCombine",
-          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "splitCombine",
-          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
         }
       ]
     },
@@ -319,7 +211,7 @@
       "id": "axis_split_tiled_narrow",
       "priority": 31,
       "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"],
-      "constants": {
         "scalar": "dtypes.T",
         "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "workgroupSize",
@@ -334,25 +226,31 @@
         {
           "id": "split_reduce",
           "name": "ArgMin.AxisSplitTiledReduce",
-          "source": { "shader": "reduce-arg-axis-split-tiled.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "splitReduce",
           "dispatch": { "x": "ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS)", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.AxisSplitCombine",
-          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "splitCombine",
-          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
         }
       ]
     },
     {
       "id": "subgroup_min_last_axis_vec4",
       "priority": 41,
-      "requires": { "features": ["subgroups"] },
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= 4", "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", "ceilDiv(axisDim, 4) <= device.adapterInfo.subgroupMinSize * 16", "outputCount >= device.adapterInfo.subgroupMinSize * 4", "rowDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "usesF16": "dtypes.T == \"f16\"",
@@ -363,21 +261,20 @@
         {
           "id": "main",
           "name": "ArgMin.SubgroupMinLastAxisVec4",
-          "source": {
-            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-            "inputs": { "mode": "\"min\"", "vec4": true, "singleSubgroup": true }
-          },
-          "bindings": "contiguousVec4",
-          "dispatch": { "workgroups": "outputCount" }
         }
       ]
     },
     {
       "id": "subgroup_min_last_axis",
       "priority": 41,
-      "requires": { "features": ["subgroups"] },
       "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"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "selectLastIndex": "attrs.select_last_index != 0",
@@ -387,12 +284,11 @@
         {
           "id": "main",
           "name": "ArgMin.SubgroupMinLastAxis",
-          "source": {
-            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-            "inputs": { "mode": "\"min\"", "vec4": false, "singleSubgroup": true }
-          },
-          "bindings": "contiguousScalar",
-          "dispatch": { "workgroups": "outputCount" }
         }
       ]
     },
@@ -400,7 +296,7 @@
       "id": "last_axis_vec4",
       "priority": 40,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= 4", "axisDim % 4 == 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "usesF16": "dtypes.T == \"f16\"",
@@ -411,12 +307,11 @@
         {
           "id": "main",
           "name": "ArgMin.LastAxisVec4",
-          "source": {
-            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-            "inputs": { "mode": "\"min\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" }
-          },
-          "bindings": "contiguousVec4",
-          "dispatch": { "workgroups": "outputCount" }
         }
       ]
     },
@@ -424,7 +319,7 @@
       "id": "last_axis",
       "priority": 39,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim % 4 != 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "selectLastIndex": "attrs.select_last_index != 0",
@@ -434,25 +329,24 @@
         {
           "id": "main",
           "name": "ArgMin.LastAxis",
-          "source": {
-            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "mode": "\"min\"",
-              "vec4": false,
-              "useSubgroups": "device.features.has(\"subgroups\") and not narrowOddRowPortablePreferred"
-            }
           },
-          "bindings": "contiguousScalar",
-          "dispatch": { "workgroups": "outputCount" }
         }
       ]
     },
     {
       "id": "axis_split",
       "priority": 30,
-      "demoteWhen": ["narrowLargeAxisTilePreferred"],
       "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= 1", "outputCount <= tunables.SPLIT_MAX_OUTPUTS", "splitScratchFits", "splitDispatchFits", "threadDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "workgroupSize",
@@ -466,16 +360,26 @@
         {
           "id": "split_reduce",
           "name": "ArgMin.AxisSplitReduce",
-          "source": { "shader": "reduce-arg-axis-split-reduce.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "splitReduce",
-          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.AxisSplitCombine",
-          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "splitCombine",
-          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
         }
       ]
     },
@@ -483,7 +387,7 @@
       "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"],
-      "constants": {
         "scalar": "dtypes.T",
         "selectLastIndex": "attrs.select_last_index != 0",
         "workgroupSize": "workgroupSize",
@@ -493,18 +397,23 @@
         {
           "id": "main",
           "name": "ArgMin.AxisTiled",
-          "source": { "shader": "reduce-arg-axis-tiled.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "axisGeometry",
-          "dispatch": { "workgroups": "ceilDiv((outputCount), (tunables.TILE_COLUMNS))" }
         }
       ]
     },
     {
       "id": "axis_serial",
       "priority": 0,
-      "supersededBy": ["axis_tiled"],
       "when": ["baseContract", "threadDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "selectLastIndex": "attrs.select_last_index != 0",
@@ -514,9 +423,14 @@
         {
           "id": "main",
           "name": "ArgMin.AxisSerial",
-          "source": { "shader": "reduce-arg-axis.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
-          "bindings": "axisGeometry",
-          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "ArgMin",
   "sinceVersion": 13,
+  "inputs": { "x": { "onnx": "data", "dtype": "T" } },
+  "outputs": {
+    "y": {
+      "onnx": "reduced",
       "dtype": "I",
+      "rank": "ranks.x if attrs.keepdims == 1 else ranks.x - 1",
+      "shape": "prefix(shapes.x, normalizedAxis) + ([1] if attrs.keepdims == 1 else []) + suffix(shapes.x, normalizedAxis + 1)",
+      "storage": "uint32"
     }
   },
+  "attributes": { "axis": { "default": 0 }, "keepdims": { "default": 1 }, "select_last_index": { "default": 0 } },
   "attributeConstraints": { "keepdims": { "values": [0, 1] }, "select_last_index": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"], "I": ["int64"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "TILE_COLUMNS": { "default": 16 },
+    "TILED_MIN_AXIS": { "default": 64 },
+    "TILED_MIN_OUTPUTS": { "default": 16 },
+    "SPLIT_MIN_AXIS": { "default": 8192 },
+    "SPLIT_MAX_OUTPUTS": { "default": 4096 },
+    "SPLIT_TARGET_AXIS": { "default": 256 },
+    "SPLIT_TILE_COLUMNS": { "default": 8 },
+    "MAX_SPLITS": { "default": 128 },
+    "ROW_SPLIT_MIN_AXIS": { "default": 32768 },
+    "ROW_SPLIT_MAX_OUTPUTS": { "default": 32 },
+    "ROW_SPLIT_TARGET_AXIS": { "default": 2048 },
+    "ROW_SPLIT_MAX_SPLITS": { "default": 64 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
+    "normalizedAxis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
+    "axisInRange": "ranks.x >= 1 and normalizedAxis >= 0 and normalizedAxis < ranks.x",
+    "axisDim": "dim(shapes.x, normalizedAxis)",
+    "axisInner": "inner(shapes.x, normalizedAxis)",
+    "outputCount": "numel(shapes.y)",
+    "outputRankOk": "(attrs.keepdims == 1 and ranks.y == ranks.x) or (attrs.keepdims == 0 and ranks.y == ranks.x - 1)",
+    "outputShapeOk": "outputRankOk and outputCount * axisDim == numel(shapes.x)",
     "baseContract": "axisInRange and axisDim > 0 and outputShapeOk and f16Ok(dtypes.T)",
+    "lastAxis": "normalizedAxis == ranks.x - 1",
     "nativeParallelType": "dtypes.T == \"f32\" or dtypes.T == \"i32\" or dtypes.T == \"u32\"",
     "rowParallelType": "nativeParallelType or dtypes.T == \"f16\"",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.SPLIT_TARGET_AXIS)))",
     "splitScratchBytes": "splitCount * outputCount * 8",
     "splitScratchFits": "splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize",
+    "splitDispatchFits": "splitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(outputCount, workgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "splitTileDispatchFits": "splitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "rowSplitCount": "min(tunables.ROW_SPLIT_MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.ROW_SPLIT_TARGET_AXIS)))",
     "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": [
     {
       "id": "last_axis_split_vec4",
       "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",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "usesF16": "dtypes.T == \"f16\"",
         {
           "id": "split_reduce",
           "name": "ArgMin.LastAxisSplitVec4Reduce",
+          "shader": "reduce-arg-row-split.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "modeSpec": "\"min\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["x", "partials_val", "partials_idx", "params"],
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.LastAxisSplitCombine",
+          "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "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",
         {
           "id": "split_reduce",
           "name": "ArgMin.LastAxisSplitReduce",
+          "shader": "reduce-arg-row-split.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "modeSpec": "\"min\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["x_2", "partials_val", "partials_idx", "params_3"],
           "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.LastAxisSplitCombine",
+          "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "axis_split_tiled_narrow",
       "priority": 31,
       "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",
         {
           "id": "split_reduce",
           "name": "ArgMin.AxisSplitTiledReduce",
+          "shader": "reduce-arg-axis-split-tiled.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["x_2", "partials_val", "partials_idx", "params_4"],
           "dispatch": { "x": "ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS)", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ArgMin.AxisSplitCombine",
+          "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "subgroup_min_last_axis_vec4",
       "priority": 41,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= 4", "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", "ceilDiv(axisDim, 4) <= device.adapterInfo.subgroupMinSize * 16", "outputCount >= device.adapterInfo.subgroupMinSize * 4", "rowDispatchFits"],
+      "requires": { "features": ["subgroups"] },
+      "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "usesF16": "dtypes.T == \"f16\"",
         {
           "id": "main",
           "name": "ArgMin.SubgroupMinLastAxisVec4",
+          "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "modeSpec": "\"min\"", "vec4": true, "singleSubgroup": true },
+          "bindings": ["x", "y", "params_5"],
+          "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "subgroup_min_last_axis",
       "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",
         {
           "id": "main",
           "name": "ArgMin.SubgroupMinLastAxis",
+          "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "modeSpec": "\"min\"", "vec4": false, "singleSubgroup": true },
+          "bindings": ["x_2", "y", "params_6"],
+          "dispatch": { "x": "min(outputCount, 65535)", "y": "ceilDiv(outputCount, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_vec4",
       "priority": 40,
       "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= 4", "axisDim % 4 == 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
+      "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "usesF16": "dtypes.T == \"f16\"",
         {
           "id": "main",
           "name": "ArgMin.LastAxisVec4",
+          "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "modeSpec": "\"min\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["x", "y", "params_5"],
+          "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",
         "usesF16": "dtypes.T == \"f16\"",
         "selectLastIndex": "attrs.select_last_index != 0",
         {
           "id": "main",
           "name": "ArgMin.LastAxis",
+          "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": {
+            "modeSpec": "\"min\"",
+            "vec4": false,
+            "useSubgroups": "device.features.has(\"subgroups\") and not narrowOddRowPortablePreferred"
           },
+          "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",
         {
           "id": "split_reduce",
           "name": "ArgMin.AxisSplitReduce",
+          "shader": "reduce-arg-axis-split-reduce.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["x_2", "partials_val", "partials_idx", "params_4"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ArgMin.AxisSplitCombine",
+          "shader": "reduce-arg-axis-split-combine.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["partials_val_2", "partials_idx_2", "y", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "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",
         {
           "id": "main",
           "name": "ArgMin.AxisTiled",
+          "shader": "reduce-arg-axis-tiled.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["x_2", "y", "params_4"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (tunables.TILE_COLUMNS)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (tunables.TILE_COLUMNS)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis_serial",
       "priority": 0,
       "when": ["baseContract", "threadDispatchFits"],
+      "supersededBy": ["axis_tiled"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "selectLastIndex": "attrs.select_last_index != 0",
         {
           "id": "main",
           "name": "ArgMin.AxisSerial",
+          "shader": "reduce-arg-axis.wgsl.jinja",
+          "derive": { "modeSpec": "\"min\"" },
+          "bindings": ["x_2", "y", "params_4"],
+          "dispatch": {
+            "x": "min(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputCount), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,24 +1,38 @@
 {
   "name": "ai.onnx.ArgMin",
-  "id": "_ai_onnx_argmin_webgpu_a0fbb5d",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "mm/x1LvCBXT2wlKgCEbxB5Oc506zsg2sfh5+8s0SzyI=",
-      "manifest.json": "6YtVbWUVmqO89IGSyNQYIvaRetZCne1BxSikonXTKaQ=",
-      "reduce-arg-axis-split-combine.wgsl.jinja": "rdvAqkyd5MkxXLWQtVzpTxqqKwANhk2L1h/SMW9P2qU=",
-      "reduce-arg-axis-split-reduce.wgsl.jinja": "VC4pBZyx0kVu5zmny2clKWdMU2D6HWpM/qiym+b664o=",
-      "reduce-arg-axis-split-tiled.wgsl.jinja": "/CtqlF3QH7OIQi9q0WzSOJC6K/KDUzcH4sDHdrcGTLg=",
-      "reduce-arg-axis-tiled.wgsl.jinja": "382jqrusVAtzqbRnn0b++oLtQK53ezOx203+9kyWs34=",
-      "reduce-arg-axis.wgsl.jinja": "LzoRz8NbCeL+4FA0A6K/5xmPw6jr1n3rqnq3h8lXXFs=",
-      "reduce-arg-row-split.wgsl.jinja": "t2dWKMUzJLdeIBMtllRHEKE+V5YQDSry6dxe/WmwiUg=",
-      "reduce-arg-row-subgroup.wgsl.jinja": "IkTnxOCG5uHuSYBWOr8vUW7MfqnJKlBY8ByKbA4+K60=",
-      "test.json": "stypKi/CuonXgtb+pFhpPyL6q5ymL6WWV0M9V27ANfc="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.ArgMin" }
 }

 {
   "name": "ai.onnx.ArgMin",
+  "id": "_ai_onnx_argmin_webgpu_e54f586",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "eUbzwtoY0FL3vKlPzo8M3cBq2mhnew0DQY8yNp6XVJE=",
+      "manifest.json": "b/mmulEG774GnpjcB/ZKqlAMwgzQSVLE/hPpLeDjJjA=",
+      "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": "YQ6LPhs3SyGX/XSXRk9g0lEW91scGJn9qE4OcHi9W/s=",
+      "reduce-arg-row-subgroup.wgsl.jinja": "dz9xdylOz8BK6INB0dcwdKa2UpJSwvkqQejLkllg7to=",
+      "test.json": "a/in/WTnKwukbD6HmiZkClzWll8kawFe503uUr81g6U="
     }
   },
+  "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"],
+      "axis_split_tiled_narrow": ["reduce-arg-axis-split-combine.wgsl.jinja", "reduce-arg-axis-split-tiled.wgsl.jinja"],
+      "subgroup_min_last_axis_vec4": ["reduce-arg-row-subgroup.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": ["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,7 +7,7 @@
 // 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 source.mode == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
@@ -17,9 +17,8 @@ const SPLIT: u32 = {{ split }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  let output_index = gid.x + gid.y * nwg.x * WG;
   if (output_index >= params.outputCount) {
     return;
   }

 // 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 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;
   }

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

@@ -1,4 +1,4 @@
-// Pass 1 of a rank-independent split-axis ArgMax/ArgMin reduction. The tensor
 // is flattened as (outer, axis, inner). Each thread scans one contiguous segment
 // of the reduced axis for one output position.
 //
@@ -6,7 +6,7 @@
 // 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 source.mode == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
@@ -20,22 +20,21 @@ const SENTINEL_IDX: u32 = 4294967295u;
 // is not folded into an unrepresentable infinity literal.
 fn lane_identity() -> {{ work }} {
 {% if scalar == "i32" %}
-  return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
 {% elif scalar == "u32" %}
-  return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
 {% else %}
-  var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded column index: wg.z carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit on the x dimension.
-  let output_index = (wg.x + wg.z * nwg.x) * WG + (gid.x % WG);
   let split_index = wg.y;
   if (output_index >= params.outputCount) { return; }
   let outer_index = output_index / params.innerSize;

+// Pass 1 of a rank-independent split-axis arg reduction. The tensor
 // is flattened as (outer, axis, inner). Each thread scans one contiguous segment
 // of the reduced axis for one output position.
 //
 // 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 }}
 // is not folded into an unrepresentable infinity literal.
 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 %}
 }
 @compute @workgroup_size(WG, 1, 1)
 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;

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

@@ -7,7 +7,7 @@
 //
 // 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 source.mode == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
@@ -20,11 +20,11 @@ const SENTINEL_IDX: u32 = 4294967295u;
 fn lane_identity() -> {{ work }} {
 {% if scalar == "i32" %}
-  return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
 {% elif scalar == "u32" %}
-  return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
 {% else %}
-  var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
 }

 //
 // 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 }}
 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 %}
 }

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

@@ -1,10 +1,10 @@
-// Tiled strided-axis ArgMax/ArgMin reduction. f32 compares directly; i32 and
 // u32 compare natively, so large-magnitude integers keep full precision. The
 // workgroup partials carry the native comparison type without a cross-pass
 // 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 source.mode == "max" else "<" %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set work = scalar if is_int else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
@@ -19,11 +19,11 @@ const SENTINEL_IDX: u32 = 4294967295u;
 // is not folded into an unrepresentable infinity literal.
 fn lane_identity() -> {{ work }} {
 {% if scalar == "i32" %}
-  return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
 {% elif scalar == "u32" %}
-  return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
 {% else %}
-  var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
 }
@@ -32,15 +32,14 @@ var<workgroup> partial_val: array<{{ work }}, WG>;
 var<workgroup> partial_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>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
   let col_lane = tid % TILE_COLS;
   let row_lane = tid / TILE_COLS;
-  // 2D-folded column-tile index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to wg.x when nwg.y == 1;
   // the bounds guard drops the over-dispatched tail.
-  let output_index = (wg.x + wg.y * nwg.x) * TILE_COLS + col_lane;
   let in_bounds = output_index < params.outputCount;
   let outer_index = output_index / params.innerSize;
   let inner_index = output_index % params.innerSize;

+// Tiled strided-axis arg reduction. f32 compares directly; i32 and
 // u32 compare natively, so large-magnitude integers keep full precision. The
 // workgroup partials carry the native comparison type without a cross-pass
 // 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 }}
 // is not folded into an unrepresentable infinity literal.
 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> partial_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>) {
   let tid = lid.x;
   let col_lane = tid % TILE_COLS;
   let row_lane = tid / TILE_COLS;
+  // 2D-folded column-tile index: wg.y carries the high bits past the per-axis dispatch fold width.
+  // Reduces to wg.x when the dispatch does not fold;
   // the bounds guard drops the over-dispatched tail.
+  let output_index = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * TILE_COLS + col_lane;
   let in_bounds = output_index < params.outputCount;
   let outer_index = output_index / params.innerSize;
   let inner_index = output_index % params.innerSize;

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

@@ -1,18 +1,17 @@
-// Rank-independent serial arg reduction for ArgMax and ArgMin. Flattening a
 // 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 source.mode == "max" else "<" %}
 {% if usesF16 %}
 enable f16;
 {%- endif %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ workgroupSize }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   // Fold the output grid into x/y when one dispatch dimension is insufficient.
-  let output_index = gid.x + gid.y * nwg.x * {{ workgroupSize }}u;
   if (output_index >= params.outputCount) {
     return;
   }

+// Rank-independent serial arg reduction. Flattening a
 // 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;
   }

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

@@ -1,4 +1,4 @@
-{% set useSubgroups = source.useSubgroups %}
 {% if usesF16 %}
 enable f16;
 {%- endif %}
@@ -17,15 +17,15 @@ const SENTINEL_IDX: u32 = 4294967295u;
 {% else %}
 {% set work = "f32" %}
 {% endif %}
-{% set cmp = ">" if source.mode == "max" else "<" %}
 fn lane_identity() -> {{ work }} {
 {%- if scalar == "i32" %}
-  return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
 {%- elif scalar == "u32" %}
-  return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
 {%- else %}
-  var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {%- endif %}
 }
@@ -36,7 +36,6 @@ 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_invocation_id) subgroup_lane: u32,
         @builtin(subgroup_size) subgroup_size: u32{% endif %}) {
   let row = wg.x;
   let split_index = wg.y;
@@ -49,7 +48,7 @@ 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 source.vec4 %}
     let value4 = x[base + chunk];
 {%- for component in ["x", "y", "z", "w"] %}
     {
@@ -93,31 +92,51 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   let subgroup_candidate = select(SENTINEL_IDX, best_idx, best_val == subgroup_val);
   let subgroup_idx = subgroupMin(subgroup_candidate);
 {%- endif %}
-  let safe_subgroup_size = max(subgroup_size, 1u);
-  // The final subgroup may be partial on devices whose native subgroup size
-  // does not divide WG. Merge that leader as well.
-  let subgroup_count = max(1u, (WG + safe_subgroup_size - 1u) / safe_subgroup_size);
-  if (subgroup_lane == 0u) {
-    let slot = min(tid / safe_subgroup_size, WG - 1u);
-    wg_val[slot] = subgroup_val;
-    wg_idx[slot] = subgroup_idx;
-  }
-  workgroupBarrier();
-  if (tid == 0u) {
-    var output_val = wg_val[0];
-    var output_idx = wg_idx[0];
-    for (var slot = 1u; slot < subgroup_count; slot++) {
       let candidate_val = wg_val[slot];
       let candidate_idx = wg_idx[slot];
 {%- if selectLastIndex %}
-      if (candidate_idx != SENTINEL_IDX && (output_idx == SENTINEL_IDX || candidate_val {{ cmp }} output_val || (candidate_val == output_val && candidate_idx > output_idx))) {
 {%- else %}
-      if (candidate_idx != SENTINEL_IDX && (output_idx == SENTINEL_IDX || candidate_val {{ cmp }} output_val || (candidate_val == output_val && candidate_idx < output_idx))) {
 {%- endif %}
-        output_val = candidate_val;
-        output_idx = candidate_idx;
       }
     }
     let scratch_index = split_index * params.rows + row;
 {%- if scalar == "u32" %}
     partials_val[scratch_index] = output_val;

+{% set useSubgroups = useSubgroups %}
 {% if usesF16 %}
 enable f16;
 {%- endif %}
 {% 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 %}
 }
 @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;
   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 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);
+    workgroupBarrier();
+    var fold_val = lane_identity();
+    var fold_idx = SENTINEL_IDX;
+    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 = subgroupMin(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;

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

@@ -1,16 +1,17 @@
-// Subgroup (value, index) row reduction for ArgMax and ArgMin over a contiguous
 // last axis, with one workgroup per output row. Each thread scans a strided
 // chunk, carrying a (best value, best index) pair with ONNX tie semantics.
 // Subgroup combination first reduces the value, then reduces the matching
-// lanes' indices. Leaders deposit per-subgroup pairs in shared memory for a
-// final fold. A one-subgroup workgroup writes its collective result directly,
-// avoiding shared memory and the workgroup barrier.
 //
 // Float inputs compare in f32. Sentinels are runtime-bitcast infinities or the
 // 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 falls back to index zero like the serial scan.
-{% if source.useSubgroups is defined %}{% set useSubgroups = source.useSubgroups %}{% else %}{% set useSubgroups = true %}{% endif %}
 {% if usesF16 %}
 enable f16;
 {%- endif %}
@@ -21,7 +22,7 @@ enable subgroups;
 const WG: u32 = {{ workgroupSize }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
-{% if source.mode == "min" and scalar != "i32" and scalar != "u32" %}
 fn pos_inf_f32() -> f32 {
   var bits = 0x7f800000u;
   return bitcast<f32>(bits);
@@ -35,16 +36,22 @@ fn pos_inf_f32() -> f32 {
 {% set cmpType = "f32" %}
 {% endif %}
 {% set cmp = "<" %}
-{%- if not (useSubgroups and source.singleSubgroup is defined and source.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 %},
-        @builtin(subgroup_invocation_id) sgLid: u32{% if not (source.singleSubgroup is defined and source.singleSubgroup) %},
-        @builtin(subgroup_size) sgSize: u32{% endif %}{% endif %}) {
   let row = wg.x + wg.y * params.rowStride;
   if (row >= params.rows) {
     return;
@@ -58,7 +65,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   var bestVal: f32 = pos_inf_f32();
 {%- endif %}
   var bestIdx: u32 = SENTINEL_IDX;
-{%- if source.vec4 %}
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v4 = x[base + c];
 {%- for comp in ["x", "y", "z", "w"] %}
@@ -120,23 +127,37 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
   let sgIdx = subgroupMin(cand);
 {%- endif %}
-{% if source.singleSubgroup is defined and source.singleSubgroup %}
-  if (sgLid == 0u) {
     y[row] = select(sgIdx, 0u, sgIdx == SENTINEL_IDX);
   }
 {%- else %}
-  let safeSg = max(sgSize, 1u);
-  let slotCount = max(1u, (WG + safeSg - 1u) / safeSg);
-  if (sgLid == 0u) {
-    let slot = min(tid / safeSg, WG - 1u);
-    wgVal[slot] = m;
-    wgIdx[slot] = sgIdx;
-  }
-  workgroupBarrier();
-  if (tid == 0u) {
-    var outVal = wgVal[0];
-    var outIdx = wgIdx[0];
-    for (var i = 1u; i < slotCount; i = i + 1u) {
       let v = wgVal[i];
       let vi = wgIdx[i];
 {%- if selectLastIndex %}
@@ -148,7 +169,14 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
         outIdx = vi;
       }
     }
-    y[row] = select(outIdx, 0u, outIdx == SENTINEL_IDX);
   }
 {%- endif %}
 {% else %}

+// Subgroup (value, index) arg reduction over a contiguous
 // last axis, with one workgroup per output row. Each thread scans a strided
 // chunk, carrying a (best value, best index) pair with ONNX tie semantics.
 // Subgroup combination first reduces the value, then reduces the matching
+// lanes' indices. Every invocation then publishes one shared-memory slot (its
+// subgroup's pair from the elected lane, the identity pair from every other
+// lane) and each subgroup folds all of them. A one-subgroup workgroup writes
+// its collective result directly, avoiding shared memory and the barrier.
 //
 // Float inputs compare in f32. Sentinels are runtime-bitcast infinities or the
 // 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 %}
 const WG: u32 = {{ workgroupSize }}u;
 const SENTINEL_IDX: u32 = 4294967295u;
+{% if modeSpec == "min" and scalar != "i32" and scalar != "u32" %}
 fn pos_inf_f32() -> f32 {
   var bits = 0x7f800000u;
   return bitcast<f32>(bits);
 {% set cmpType = "f32" %}
 {% endif %}
 {% set cmp = "<" %}
+{%- if cmpType == "i32" %}
+{% set identityVal = "2147483647i" %}
+{%- elif cmpType == "u32" %}
+{% set identityVal = "4294967295u" %}
+{%- else %}
+{% set identityVal = "pos_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;
   var bestVal: f32 = pos_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"] %}
   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);
+    workgroupBarrier();
+    var outVal = {{ identityVal }};
+    var outIdx = SENTINEL_IDX;
+    for (var i = rank; i < WG; i = i + count) {
       let v = wgVal[i];
       let vi = wgIdx[i];
 {%- if selectLastIndex %}
         outIdx = vi;
       }
     }
+    let foldVal = subgroupMin(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 %}

build/webgpu/test.json CHANGED Viewed

@@ -1,9 +1,8 @@
 {
-  "op": "ai.onnx.ArgMin",
   "fixtureArrays": {
     "onnx_backend_argmin_input_x": [0.9762700796127319, 4.3037872314453125, 2.055267572402954, 0.8976636528968811, -1.5269039869308472, 2.917882204055786, -1.248255729675293, 7.835460186004639, 9.273255348205566, -2.331169605255127, 5.834500789642334, 0.577898383140564, 1.3608912229537964, 8.511932373046875, -8.579278945922852, -8.257413864135742, -9.59563159942627, 6.6523966789245605, 5.563135147094727, 7.400242805480957, 9.572366714477539, 5.983171463012695, -0.7704127430915833, 5.610583305358887],
     "axis0_splitk_8192x32_ties_first_index_input_x": [2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2],
-    "subgroup_vec4_ties_first_index_input_x": [1, 5, 5, 2, 5, 0, 1, 0, 7, 7, 7, 7, 7, 7, 7, 7],
     "axis0_splitk_i32_8192x16_precision_min_input_x": [16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216]
   },
   "cases": [
@@ -26,7 +25,7 @@
       "provenance": {
         "source": "https://onnx.ai/onnx/operators/onnx__ArgMin.html",
         "test": "ArgMin axis default",
-        "notes": "Omitting axis must reduce the leading dimension (ONNX default 0), not axis 1. The rank-2 output shape and per-column indices distinguish the two semantics. ONNX int64 indices use the framework's representable uint32 slot."
       },
       "inputs": {
         "x": {
@@ -84,7 +83,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ArgMin",
-        "notes": "Axis-0 companion for finite subnormal ordering: a negative subnormal is strictly less than zero and must not be tie-broken as flushed zero."
       },
       "attrs": { "axis": 0, "keepdims": 0 },
       "inputs": {
@@ -95,9 +94,7 @@
     {
       "name": "dispatch_cliff_axis1_16777216x1",
       "attrs": { "axis": 1, "keepdims": 0 },
-      "inputs": {
-        "x": { "dtype": "float32", "shape": [16777216, 1], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } }
-      },
       "outputs": { "y": { "dtype": "uint32", "shape": [16777216], "tolerance": 0 } }
     },
     {
@@ -200,12 +197,12 @@
       "name": "ort_axis1_nan_first_incumbent_gpu_gap",
       "skipGpu": {
         "category": "todo",
-        "reason": "The parallel min/max reduction routes do not yet preserve the reference's first-element NaN incumbent semantics; explicit NaN and index tracking is implementable in WGSL."
       },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ArgMin",
-        "notes": "NaN extension: ORT seeds ArgMin from the first reduced element, so a leading NaN remains the selected incumbent while later NaNs are ignored after a finite incumbent."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -224,7 +221,7 @@
       "name": "ort_axis1_nan_select_last_index_gpu_gap",
       "skipGpu": {
         "category": "todo",
-        "reason": "The parallel min/max reduction routes do not yet preserve the reference's first-element NaN incumbent semantics; explicit NaN and index tracking is implementable in WGSL."
       },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
@@ -279,7 +276,7 @@
           "shape": [2, 4],
           "data": {
             "kind": "values",
-            "values": ["Infinity", "Infinity", "Infinity", "Infinity", "Infinity", 5.0, "Infinity", 5.0]
           }
         }
       },
@@ -537,7 +534,7 @@
       "name": "onnx_backend_argmin_default_axis_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_example",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "keepdims": 1, "axis": 0 },
       "inputs": {
@@ -549,7 +546,7 @@
       "name": "onnx_backend_argmin_default_axis_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_example_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "keepdims": 1, "select_last_index": 1, "axis": 0 },
       "inputs": {
@@ -561,7 +558,7 @@
       "name": "onnx_backend_argmin_default_axis_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_random",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "keepdims": 1, "axis": 0 },
       "inputs": {
@@ -577,7 +574,7 @@
       "name": "onnx_backend_argmin_default_axis_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_random_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "keepdims": 1, "select_last_index": 1, "axis": 0 },
       "inputs": {
@@ -593,7 +590,7 @@
       "name": "onnx_backend_argmin_keepdims_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_example",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 1 },
       "inputs": {
@@ -605,7 +602,7 @@
       "name": "onnx_backend_argmin_keepdims_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_example_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
@@ -617,7 +614,7 @@
       "name": "onnx_backend_argmin_keepdims_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 1 },
       "inputs": {
@@ -633,7 +630,7 @@
       "name": "onnx_backend_argmin_keepdims_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
@@ -649,7 +646,7 @@
       "name": "onnx_backend_argmin_negative_axis_keepdims_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_example",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": -1, "keepdims": 1 },
       "inputs": {
@@ -661,7 +658,7 @@
       "name": "onnx_backend_argmin_negative_axis_keepdims_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_example_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
@@ -673,7 +670,7 @@
       "name": "onnx_backend_argmin_negative_axis_keepdims_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_random",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": -1, "keepdims": 1 },
       "inputs": {
@@ -689,7 +686,7 @@
       "name": "onnx_backend_argmin_negative_axis_keepdims_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_random_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
@@ -705,7 +702,7 @@
       "name": "onnx_backend_argmin_no_keepdims_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_example",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -717,7 +714,7 @@
       "name": "onnx_backend_argmin_no_keepdims_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_example_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": {
@@ -729,7 +726,7 @@
       "name": "onnx_backend_argmin_no_keepdims_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_random",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -745,7 +742,7 @@
       "name": "onnx_backend_argmin_no_keepdims_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_random_select_last_index",
-        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
       },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": {
@@ -767,7 +764,7 @@
           "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
         }
       },
-      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [5, 0] } } }
     },
     {
       "name": "subgroup_vec4_ties_select_last_index",
@@ -779,7 +776,7 @@
           "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
         }
       },
-      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [7, 7] } } }
     },
     {
       "name": "subgroup_scalar_f16_2x5",
@@ -889,7 +886,7 @@
       "name": "rank4_axis0_batch_argmin_f32_keepdims",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random (extended to rank4 batch axis)",
-        "notes": "Same rank4 axis=0 selection hole on the default float32 dtype with keepdims=1, to prove the gap is dtype-independent (not an integer-only branch)."
       },
       "attrs": { "axis": 0, "keepdims": 1 },
       "inputs": {
@@ -909,7 +906,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ArgMin_int8 (rank4 NCHW channel argmin)",
-        "notes": "int8 channel-axis argmin on a rank4 NCHW tensor. rank4_axis1 gates dtype only via f16Ok(T) (true for int8) so this is a WORKING-but-untested integer dtype on the rank4 selection path."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
@@ -929,7 +926,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ArgMin (rank3 middle-axis f16 select_last)",
-        "notes": "f16 middle-axis argmin on rank3 with keepdims=1 and select_last_index=1. rank3_axis1 admits f16 via f16Ok(T) on an f16 device; no existing rank3/rank4 case is f16."
       },
       "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
@@ -1060,7 +1057,7 @@
     {
       "name": "axis0_narrow_splitk_16384x8_f32",
       "provenance": {
-        "notes": "Tall narrow axis-0 (16384 rows, 8 cols) routes to the split-and-tiled kernel: all 256 lanes cooperate across eight columns instead of leaving 248 lanes idle. Cross-checks the tiled partials and combine pass against the reference."
       },
       "attrs": { "axis": 0, "keepdims": 0 },
       "inputs": {
@@ -1088,7 +1085,7 @@
     {
       "name": "last_axis_split_scalar_logits_ties_last",
       "provenance": {
-        "notes": "The unaligned 32769-element companion route-locks the scalar split-row reduction and verifies select_last_index across split boundaries."
       },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 32769], "data": { "kind": "constant", "value": -5.0 } } },
@@ -1132,7 +1129,7 @@
     {
       "name": "last_axis_split_vec4_f16_rows",
       "provenance": {
-        "notes": "float16 on the vec4 split-row reduction: each vec4 component must be widened to f32 before comparison, and the shader must enable f16. Rank 3 (1, 2, 32768) reduces the last axis exactly as the rank-2 logits case does while keeping the reference on its hoisted-decode path. All seven cycle values are exactly representable in float16, so the answer does not depend on rounding. ONNX ArgMin admits float16 via T; first index on ties."
       },
       "attrs": { "axis": 2, "keepdims": 0 },
       "inputs": {
@@ -1149,7 +1146,7 @@
     {
       "name": "last_axis_split_scalar_f16_rows",
       "provenance": {
-        "notes": "The unaligned (32769-column) companion: float16 on the scalar split-row reduction, where the widening happens on the plain load rather than on a vec4 component. Rank 3 keeps the reference on its hoisted-decode path; all cycle values are exact in float16. ONNX ArgMin admits float16 via T; first index on ties."
       },
       "attrs": { "axis": 2, "keepdims": 0 },
       "inputs": {
@@ -1183,7 +1180,7 @@
     {
       "name": "rank5_middle_axis_generic_geometry",
       "provenance": {
-        "notes": "Proves the rank-independent (outer, axis, inner) fallback covers ranks beyond the former rank-4 manifest specializations."
       },
       "attrs": { "axis": 2, "keepdims": 0, "select_last_index": 1 },
       "inputs": {
@@ -1198,7 +1195,7 @@
     {
       "name": "axis0_splitk_select_last_index_ties_8192x16",
       "provenance": {
-        "notes": "select_last_index=1 on the split-tiled route, which no case had ever rendered: every existing case for that geometry leaves the attribute at its default. The 17-value cycle is coprime with the 16 columns, so each column sees every value about 482 times and the minimum is a 482-way tie whose winning index differs per column -- exactly the tie-break the two arms disagree about, and one a per-lane fold can get wrong even when its lane-to-index map is monotone."
       },
       "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
       "inputs": {

 {
   "fixtureArrays": {
     "onnx_backend_argmin_input_x": [0.9762700796127319, 4.3037872314453125, 2.055267572402954, 0.8976636528968811, -1.5269039869308472, 2.917882204055786, -1.248255729675293, 7.835460186004639, 9.273255348205566, -2.331169605255127, 5.834500789642334, 0.577898383140564, 1.3608912229537964, 8.511932373046875, -8.579278945922852, -8.257413864135742, -9.59563159942627, 6.6523966789245605, 5.563135147094727, 7.400242805480957, 9.572366714477539, 5.983171463012695, -0.7704127430915833, 5.610583305358887],
     "axis0_splitk_8192x32_ties_first_index_input_x": [2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2],
+    "subgroup_vec4_ties_first_index_input_x": [1, 5, 0, 2, 5, 0, 1, 4, 7, 7, 7, 7, 7, 7, 7, 7],
     "axis0_splitk_i32_8192x16_precision_min_input_x": [16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216]
   },
   "cases": [
       "provenance": {
         "source": "https://onnx.ai/onnx/operators/onnx__ArgMin.html",
         "test": "ArgMin axis default",
+        "notes": "Omitting axis must reduce the leading dimension (ONNX default 0), not axis 1. The rank-2 output shape and per-column indices distinguish the two semantics. Representable ONNX int64 indices are stored in uint32 slots."
       },
       "inputs": {
         "x": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ArgMin",
+        "notes": "Along axis 0, a negative subnormal is strictly less than zero and must not be tie-broken as though it were flushed to zero."
       },
       "attrs": { "axis": 0, "keepdims": 0 },
       "inputs": {
     {
       "name": "dispatch_cliff_axis1_16777216x1",
       "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [16777216, 1], "data": { "kind": "constant", "value": 1.0 } } },
       "outputs": { "y": { "dtype": "uint32", "shape": [16777216], "tolerance": 0 } }
     },
     {
       "name": "ort_axis1_nan_first_incumbent_gpu_gap",
       "skipGpu": {
         "category": "todo",
+        "reason": "The parallel reduction routes do not preserve first-element NaN incumbent semantics. Explicit NaN and index tracking is implementable in WGSL."
       },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ArgMin",
+        "notes": "NaN extension: ArgMin seeds the reduction from the first reduced element, so a leading NaN remains the selected incumbent while later NaNs are ignored after a finite incumbent."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
       "name": "ort_axis1_nan_select_last_index_gpu_gap",
       "skipGpu": {
         "category": "todo",
+        "reason": "The parallel reduction routes do not preserve first-element NaN incumbent semantics. Explicit NaN and index tracking is implementable in WGSL."
       },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
           "shape": [2, 4],
           "data": {
             "kind": "values",
+            "values": ["Infinity", "Infinity", "Infinity", "Infinity", 5.0, 5.0, "Infinity", "Infinity"]
           }
         }
       },
       "name": "onnx_backend_argmin_default_axis_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_example",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "keepdims": 1, "axis": 0 },
       "inputs": {
       "name": "onnx_backend_argmin_default_axis_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_example_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "keepdims": 1, "select_last_index": 1, "axis": 0 },
       "inputs": {
       "name": "onnx_backend_argmin_default_axis_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_random",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "keepdims": 1, "axis": 0 },
       "inputs": {
       "name": "onnx_backend_argmin_default_axis_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_random_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "keepdims": 1, "select_last_index": 1, "axis": 0 },
       "inputs": {
       "name": "onnx_backend_argmin_keepdims_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_example",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_keepdims_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_example_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_keepdims_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_keepdims_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_negative_axis_keepdims_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_example",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": -1, "keepdims": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_negative_axis_keepdims_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_example_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_negative_axis_keepdims_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_random",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": -1, "keepdims": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_negative_axis_keepdims_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_random_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_no_keepdims_example",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_example",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
       "name": "onnx_backend_argmin_no_keepdims_example_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_example_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": {
       "name": "onnx_backend_argmin_no_keepdims_random",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_random",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
       "name": "onnx_backend_argmin_no_keepdims_random_select_last_index",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_random_select_last_index",
+        "notes": "ONNX arg-reduction outputs use int64 indices; this WebGPU package stores representable indices in uint32 slots."
       },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": {
           "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
         }
       },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [2, 0] } } }
     },
     {
       "name": "subgroup_vec4_ties_select_last_index",
           "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
         }
       },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [5, 7] } } }
     },
     {
       "name": "subgroup_scalar_f16_2x5",
       "name": "rank4_axis0_batch_argmin_f32_keepdims",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random (extended to rank4 batch axis)",
+        "notes": "A rank-4 float32 axis-0 reduction with keepdims=1 exercises the generic batch-axis selection path."
       },
       "attrs": { "axis": 0, "keepdims": 1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ArgMin_int8 (rank4 NCHW channel argmin)",
+        "notes": "An int8 channel-axis ArgMin over a rank-4 NCHW tensor exercises integer comparison on the rank-4 axis-1 path."
       },
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ArgMin (rank3 middle-axis f16 select_last)",
+        "notes": "A float16 rank-3 middle-axis reduction with keepdims=1 and select_last_index=1 exercises the float16 strided route and last-index tie breaking."
       },
       "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
       "inputs": {
     {
       "name": "axis0_narrow_splitk_16384x8_f32",
       "provenance": {
+        "notes": "A tall, narrow axis-0 reduction with 16,384 rows and eight columns exercises split-and-tiled partial reductions and their combine pass across all columns."
       },
       "attrs": { "axis": 0, "keepdims": 0 },
       "inputs": {
     {
       "name": "last_axis_split_scalar_logits_ties_last",
       "provenance": {
+        "notes": "An unaligned 32,769-element row exercises scalar split-row reduction and verifies select_last_index across split boundaries."
       },
       "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 32769], "data": { "kind": "constant", "value": -5.0 } } },
     {
       "name": "last_axis_split_vec4_f16_rows",
       "provenance": {
+        "notes": "A rank-3 float16 tensor with 32,768-element rows exercises vec4 split-row reduction. Each component is widened before comparison, all seven cycle values are exactly representable, and ties select the first index."
       },
       "attrs": { "axis": 2, "keepdims": 0 },
       "inputs": {
     {
       "name": "last_axis_split_scalar_f16_rows",
       "provenance": {
+        "notes": "A rank-3 float16 tensor with 32,769-element rows exercises scalar split-row reduction. Loaded values are widened before comparison, all cycle values are exactly representable, and ties select the first index."
       },
       "attrs": { "axis": 2, "keepdims": 0 },
       "inputs": {
     {
       "name": "rank5_middle_axis_generic_geometry",
       "provenance": {
+        "notes": "A rank-5 middle-axis reduction exercises the rank-independent outer-axis-inner geometry."
       },
       "attrs": { "axis": 2, "keepdims": 0, "select_last_index": 1 },
       "inputs": {
     {
       "name": "axis0_splitk_select_last_index_ties_8192x16",
       "provenance": {
+        "notes": "On the split-tiled route, a 17-value cycle coprime with the 16 columns gives each column an approximately 482-way minimum tie. With select_last_index=1, each column must retain its final matching axis index across lane-local folds."
       },
       "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
       "inputs": {