sync 2e7068faf55e

Browse files

Files changed (12) hide show

README.md +77 -0
build/webgpu/bench.json +121 -0
build/webgpu/manifest.json +519 -0
build/webgpu/metadata.json +24 -0
build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja +49 -0
build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja +69 -0
build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja +95 -0
build/webgpu/reduce-arg-axis-tiled.wgsl.jinja +85 -0
build/webgpu/reduce-arg-axis.wgsl.jinja +45 -0
build/webgpu/reduce-arg-row-split.wgsl.jinja +159 -0
build/webgpu/reduce-arg-row-subgroup.wgsl.jinja +185 -0
build/webgpu/test.json +1266 -0

README.md CHANGED Viewed

@@ -1,3 +1,80 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.ArgMax
+`ai.onnx`  ·  standard ONNX operator  ·  ONNX opset ≥ 13
+## Description
+Returns the index of the maximum value along an axis, choosing the first equal value unless `select_last_index` is enabled.
+See the [ONNX `ArgMax` spec](https://onnx.ai/onnx/operators/onnx__ArgMax.html) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `data` | `x` | `T` | — | — | Values whose maximum 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 maximum values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
+## Attributes
+Default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `axis` | `0` | Axis to reduce; negative values count from the back. |
+| `keepdims` | `1` | Retain the reduced dimension with length one when non-zero. |
+| `select_last_index` | `0` | Choose the last equal maximum instead of the first when non-zero. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16`, `int32`, `uint32`, `int16`, `int8`, `uint8` |
+| `I` | `int64` |
+## Device requirements
+Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## 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
+- [`reduce-arg-axis-split-combine.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja)
+- [`reduce-arg-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja)
+- [`reduce-arg-axis-split-tiled.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja)
+- [`reduce-arg-axis-tiled.wgsl.jinja`](build/webgpu/reduce-arg-axis-tiled.wgsl.jinja)
+- [`reduce-arg-axis.wgsl.jinja`](build/webgpu/reduce-arg-axis.wgsl.jinja)
+- [`reduce-arg-row-split.wgsl.jinja`](build/webgpu/reduce-arg-row-split.wgsl.jinja)
+- [`reduce-arg-row-subgroup.wgsl.jinja`](build/webgpu/reduce-arg-row-subgroup.wgsl.jinja)
+## 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.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/ai.onnx.ArgMax", { version: 1 });
+const { y } = await kernel({ x: { data: xData, shape: [2, 2] } });
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,121 @@

+{
+  "op": "ai.onnx.ArgMax",
+  "cases": [
+    {
+      "name": "argmax-axis0-f32-1024x512",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 1024, "cols": 512 },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": { "x": { "shape": [1024, 512], "dtype": "float32", "dist": "normal", "seed": 971, "scale": 1 } },
+      "outputs": { "y": { "shape": [512], "dtype": "uint32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "logits",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 32000] } },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1] } }
+    },
+    {
+      "name": "argmax-axis0-large-65536x512",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 65536, "cols": 512 },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": { "x": { "shape": [65536, 512], "dtype": "float32", "dist": "normal", "seed": 211, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [512], "dtype": "uint32" } },
+      "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
+    },
+    {
+      "name": "argmax-lastaxis-f32-4096x4096",
+      "preset": "smoke",
+      "vars": { "rows": 4096, "cols": 4096 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 117, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [4096], "dtype": "uint32" } },
+      "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
+    },
+    {
+      "name": "argmax-lastaxis-f32-4096x50257",
+      "preset": "stress",
+      "provenance": { "notes": "Stress-only capacity case: declared tensors occupy 785 MiB of GPU storage." },
+      "vars": { "dtype": "float32", "rows": 4096, "cols": 50257 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [4096, 50257], "dtype": "float32", "dist": "normal", "seed": 503, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [4096], "dtype": "uint32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "argmax-logits-f32-1x131072",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 1, "cols": 131072 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [1, 131072], "dtype": "float32", "dist": "normal", "seed": 733, "scale": 1 } },
+      "outputs": { "y": { "shape": [1], "dtype": "uint32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "argmax-vocab-unaligned-f32-1x50257",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 1, "cols": 50257 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [1, 50257], "dtype": "float32", "dist": "normal", "seed": 739, "scale": 1 } },
+      "outputs": { "y": { "shape": [1], "dtype": "uint32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "argmax-lastaxis-f16-4096x4096",
+      "preset": "smoke",
+      "vars": { "dtype": "float16", "rows": 4096, "cols": 4096 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [4096, 4096], "dtype": "float16", "dist": "normal", "seed": 617, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [4096], "dtype": "uint32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "argmax-axis0-int32-65536x512",
+      "preset": "smoke",
+      "vars": { "dtype": "int32", "rows": 65536, "cols": 512 },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": { "x": { "shape": [65536, 512], "dtype": "int32", "dist": "normal", "seed": 419, "scale": 1000 } },
+      "outputs": { "y": { "shape": [512], "dtype": "uint32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "argmax-lastaxis-int8-2048x4096-serial",
+      "preset": "stress",
+      "vars": { "dtype": "int8", "rows": 2048, "cols": 4096 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [2048, 4096], "dtype": "int8", "dist": "normal", "seed": 831, "scale": 50 } },
+      "outputs": { "y": { "shape": [2048], "dtype": "uint32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "argmax-axis0-narrow-starvation-131072x8",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 131072, "cols": 8 },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": { "x": { "shape": [131072, 8], "dtype": "float32", "dist": "normal", "seed": 823, "scale": 1 } },
+      "outputs": { "y": { "shape": [8], "dtype": "uint32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "argmax-temporal-rank3-axis1-starvation-1x65536x8",
+      "preset": "stress",
+      "vars": { "dtype": "float32", "frames": 65536, "classes": 8 },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "shape": [1, 65536, 8], "dtype": "float32", "dist": "normal", "seed": 829, "scale": 1 } },
+      "outputs": { "y": { "shape": [1, 8], "dtype": "uint32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.frames * args.classes * dtypeBytes(args.dtype)" }]
+      }
+    }
+  ],
+  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,519 @@

+{
+  "domain": "ai.onnx",
+  "name": "ArgMax",
+  "sinceVersion": 13,
+  "description": "Returns the index of the maximum value along an axis, choosing the first equal value unless `select_last_index` is enabled.",
+  "inputs": [{ "role": "data", "dtype": "T", "description": "Values whose maximum 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 maximum 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 maximum 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)",
+    "narrowLargeAxisTilePreferred": "narrowSubgroupRange and outputCount >= workgroupSize",
+    "threadDispatchFits": "ceilDiv(outputCount, workgroupSize) <= foldedDispatchCapacity",
+    "rowDispatchFits": "outputCount <= foldedDispatchCapacity",
+    "rowStorageFits": "workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
+    "parallelRowWorthwhile": "axisDim >= 256 or outputCount < 32768",
+    "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\"",
+        "selectLastIndex": "attrs.select_last_index != 0",
+        "workgroupSize": "workgroupSize",
+        "split": "rowSplitCount"
+      },
+      "intermediates": [
+        { "id": "partials_val", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" },
+        { "id": "partials_idx", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" }
+      ],
+      "passes": [
+        {
+          "id": "split_reduce",
+          "name": "ArgMax.LastAxisSplitVec4Reduce",
+          "source": {
+            "shader": "reduce-arg-row-split.wgsl.jinja",
+            "inputs": { "mode": "\"max\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" }
+          },
+          "bindings": "rowSplitVec4",
+          "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
+        },
+        {
+          "id": "combine",
+          "name": "ArgMax.LastAxisSplitCombine",
+          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "bindings": "splitCombine",
+          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
+        }
+      ]
+    },
+    {
+      "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",
+        "workgroupSize": "workgroupSize",
+        "split": "rowSplitCount"
+      },
+      "intermediates": [
+        { "id": "partials_val", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" },
+        { "id": "partials_idx", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" }
+      ],
+      "passes": [
+        {
+          "id": "split_reduce",
+          "name": "ArgMax.LastAxisSplitReduce",
+          "source": {
+            "shader": "reduce-arg-row-split.wgsl.jinja",
+            "inputs": { "mode": "\"max\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" }
+          },
+          "bindings": "rowSplitScalar",
+          "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
+        },
+        {
+          "id": "combine",
+          "name": "ArgMax.LastAxisSplitCombine",
+          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "bindings": "splitCombine",
+          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
+        }
+      ]
+    },
+    {
+      "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",
+        "split": "splitCount",
+        "tileCols": "tunables.SPLIT_TILE_COLUMNS"
+      },
+      "intermediates": [
+        { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
+        { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
+      ],
+      "passes": [
+        {
+          "id": "split_reduce",
+          "name": "ArgMax.AxisSplitTiledReduce",
+          "source": { "shader": "reduce-arg-axis-split-tiled.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "bindings": "splitReduce",
+          "dispatch": { "x": "ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS)", "y": "splitCount" }
+        },
+        {
+          "id": "combine",
+          "name": "ArgMax.AxisSplitCombine",
+          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "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\"",
+        "selectLastIndex": "attrs.select_last_index != 0",
+        "workgroupSize": "device.adapterInfo.subgroupMinSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ArgMax.SubgroupMinLastAxisVec4",
+          "source": {
+            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+            "inputs": { "mode": "\"max\"", "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",
+        "workgroupSize": "device.adapterInfo.subgroupMinSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ArgMax.SubgroupMinLastAxis",
+          "source": {
+            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+            "inputs": { "mode": "\"max\"", "vec4": false, "singleSubgroup": true }
+          },
+          "bindings": "contiguousScalar",
+          "dispatch": { "workgroups": "outputCount" }
+        }
+      ]
+    },
+    {
+      "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\"",
+        "selectLastIndex": "attrs.select_last_index != 0",
+        "workgroupSize": "min(workgroupSize, max(32, pow2ceil(ceilDiv(axisDim, 4))))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ArgMax.LastAxisVec4",
+          "source": {
+            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+            "inputs": { "mode": "\"max\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" }
+          },
+          "bindings": "contiguousVec4",
+          "dispatch": { "workgroups": "outputCount" }
+        }
+      ]
+    },
+    {
+      "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",
+        "workgroupSize": "min(workgroupSize, max(32, pow2ceil(axisDim)))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ArgMax.LastAxis",
+          "source": {
+            "shader": "reduce-arg-row-subgroup.wgsl.jinja",
+            "inputs": { "mode": "\"max\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" }
+          },
+          "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",
+        "split": "splitCount"
+      },
+      "intermediates": [
+        { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
+        { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
+      ],
+      "passes": [
+        {
+          "id": "split_reduce",
+          "name": "ArgMax.AxisSplitReduce",
+          "source": { "shader": "reduce-arg-axis-split-reduce.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "bindings": "splitReduce",
+          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize", "y": "splitCount" }
+        },
+        {
+          "id": "combine",
+          "name": "ArgMax.AxisSplitCombine",
+          "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "bindings": "splitCombine",
+          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
+        }
+      ]
+    },
+    {
+      "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",
+        "tileCols": "tunables.TILE_COLUMNS"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ArgMax.AxisTiled",
+          "source": { "shader": "reduce-arg-axis-tiled.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "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",
+        "workgroupSize": "workgroupSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ArgMax.AxisSerial",
+          "source": { "shader": "reduce-arg-axis.wgsl.jinja", "inputs": { "mode": "\"max\"" } },
+          "bindings": "axisGeometry",
+          "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,24 @@

+{
+  "name": "ai.onnx.ArgMax",
+  "id": "_ai_onnx_argmax_webgpu_97b7564",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "eUX0SrelUEHbrkOOtyZUmXYTcd70fLBI6oiQOQj3Uyg=",
+      "manifest.json": "csrYFYXbC3J/cmgTiTUK5e78NJOJceeiFXvyUe2ZAHI=",
+      "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": "w5ycl8KmqAXgOJSClyPwdHpvD9scRtewwYVn9g++A7o=",
+      "reduce-arg-row-subgroup.wgsl.jinja": "0OjqIIUwZiSWMJ/Uu6uJem+mFZrECLPJctBSQZR+IKc=",
+      "test.json": "gcD6r3sktg3/lFyaKAnQC4krQaZJDHYlXhQ4WHiw3EU="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.ArgMax" }
+}

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

	@@ -0,0 +1,49 @@

+// Pass 2 of the split-axis arg reduction folds SPLIT (best value, best index)
+// partials and emits the winning reduced-axis index. One thread handles each
+// flattened output. The reduce pass bit-stored each partial value in u32
+// scratch, so this pass restores the comparison type: floats compare in f32,
+// while i32 and u32 compare natively.
+//
+// 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 }}
+const WG: u32 = {{ workgroupSize }}u;
+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;
+  }
+{% if scalar == "u32" %}
+  var out_val = partials_val[output_index];
+{% else %}
+  var out_val = bitcast<{{ work }}>(partials_val[output_index]);
+{% endif %}
+  var out_idx = partials_idx[output_index];
+  for (var split_index = 1u; split_index < SPLIT; split_index++) {
+{% if scalar == "u32" %}
+    let candidate_val = partials_val[split_index * params.outputCount + output_index];
+{% else %}
+    let candidate_val = bitcast<{{ work }}>(partials_val[split_index * params.outputCount + output_index]);
+{% endif %}
+    let candidate_idx = partials_idx[split_index * params.outputCount + output_index];
+{% if selectLastIndex %}
+    if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
+{% else %}
+    if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
+{% endif %}
+      out_val = candidate_val;
+      out_idx = candidate_idx;
+    }
+  }
+  y[output_index] = select(out_idx, 0u, out_idx == SENTINEL_IDX);
+}

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

	@@ -0,0 +1,69 @@

+// 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.
+//
+// Partial values are bit-stored in u32 scratch, so f32, i32, and u32 share one
+// 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 }}
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+const SENTINEL_IDX: u32 = 4294967295u;
+// The scan identity: the most-losing value for the mode, in the compare type.
+// The float branch bitcasts through a runtime `var` so the +/-inf bit pattern
+// 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;
+  let inner_index = output_index % params.innerSize;
+  let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
+  // Even split of [0, rows) across SPLIT segments (last segment may be shorter).
+  let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
+  let axis_begin = split_index * chunk;
+  var axis_end = min(axis_begin + chunk, params.axisDim);
+  var best_val = lane_identity();
+  var best_idx = SENTINEL_IDX;
+  for (var axis_index = axis_begin; axis_index < axis_end; axis_index++) {
+    let v = x[input_base + axis_index * params.innerSize];
+{% if selectLastIndex %}
+    if (v {{ cmp }} best_val || (v == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
+{% else %}
+    if (v {{ cmp }} best_val || (v == best_val && axis_index < best_idx)) {
+{% endif %}
+      best_val = v;
+      best_idx = axis_index;
+    }
+  }
+{% if scalar == "u32" %}
+  partials_val[split_index * params.outputCount + output_index] = best_val;
+{% else %}
+  partials_val[split_index * params.outputCount + output_index] = bitcast<u32>(best_val);
+{% endif %}
+  partials_idx[split_index * params.outputCount + output_index] = best_idx;
+}

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

	@@ -0,0 +1,95 @@

+// Split-axis arg reduction for a narrow set of output columns. The ordinary
+// split pass assigns one invocation to each output column, leaving most lanes
+// idle when outputCount is small. This route tiles TILE_COLS columns across the
+// workgroup and uses ROW_LANES lanes to cooperatively scan each column. One
+// workgroup handles each (split, column tile), preserving coalesced reads for
+// flattened (outer, axis, inner) tensors.
+//
+// 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 }}
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+const TILE_COLS: u32 = {{ tileCols }}u;
+const ROW_LANES: u32 = WG / TILE_COLS;
+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 %}
+}
+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>
+) {
+  let tid = lid.x;
+  let col_lane = tid % TILE_COLS;
+  let row_lane = tid / TILE_COLS;
+  let output_index = wg.x * TILE_COLS + col_lane;
+  let split_index = wg.y;
+  let in_bounds = output_index < params.outputCount;
+  let outer_index = output_index / params.innerSize;
+  let inner_index = output_index % params.innerSize;
+  let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
+  let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
+  let axis_begin = split_index * chunk;
+  let axis_end = min(axis_begin + chunk, params.axisDim);
+  var best_val = lane_identity();
+  var best_idx = SENTINEL_IDX;
+  if (in_bounds) {
+    for (var axis_index = axis_begin + row_lane; axis_index < axis_end; axis_index += ROW_LANES) {
+      let value = x[input_base + axis_index * params.innerSize];
+{% if selectLastIndex %}
+      if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
+{% else %}
+      if (value {{ cmp }} best_val || (value == best_val && axis_index < best_idx)) {
+{% endif %}
+        best_val = value;
+        best_idx = axis_index;
+      }
+    }
+  }
+  partial_val[tid] = best_val;
+  partial_idx[tid] = best_idx;
+  workgroupBarrier();
+  if (row_lane == 0u && in_bounds) {
+    var out_val = partial_val[col_lane];
+    var out_idx = partial_idx[col_lane];
+    for (var lane = 1u; lane < ROW_LANES; lane++) {
+      let candidate_val = partial_val[lane * TILE_COLS + col_lane];
+      let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
+{% if selectLastIndex %}
+      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
+{% else %}
+      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
+{% endif %}
+        out_val = candidate_val;
+        out_idx = candidate_idx;
+      }
+    }
+{% if scalar == "u32" %}
+    partials_val[split_index * params.outputCount + output_index] = out_val;
+{% else %}
+    partials_val[split_index * params.outputCount + output_index] = bitcast<u32>(out_val);
+{% endif %}
+    partials_idx[split_index * params.outputCount + output_index] = out_idx;
+  }
+}

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

	@@ -0,0 +1,85 @@

+// 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 }}
+const WG: u32 = {{ workgroupSize }}u;
+const TILE_COLS: u32 = {{ tileCols }}u;
+const ROW_LANES: u32 = WG / TILE_COLS;
+const SENTINEL_IDX: u32 = 4294967295u;
+// The scan identity: the most-losing value for the mode, in the compare type.
+// The float branch bitcasts through a runtime `var` so the +/-inf bit pattern
+// 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 %}
+}
+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;
+  let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
+  var best_val = lane_identity();
+  var best_idx = SENTINEL_IDX;
+  if (in_bounds) {
+    for (var axis_index = row_lane; axis_index < params.axisDim; axis_index = axis_index + ROW_LANES) {
+      let v = x[input_base + axis_index * params.innerSize];
+{% if selectLastIndex %}
+      if (v {{ cmp }} best_val || (v == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
+{% else %}
+      if (v {{ cmp }} best_val || (v == best_val && axis_index < best_idx)) {
+{% endif %}
+        best_val = v;
+        best_idx = axis_index;
+      }
+    }
+  }
+  partial_val[tid] = best_val;
+  partial_idx[tid] = best_idx;
+  workgroupBarrier();
+  if (row_lane == 0u && in_bounds) {
+    var out_val = partial_val[col_lane];
+    var out_idx = partial_idx[col_lane];
+    for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
+      let candidate_val = partial_val[lane * TILE_COLS + col_lane];
+      let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
+{% if selectLastIndex %}
+      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
+{% else %}
+      if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
+{% endif %}
+        out_val = candidate_val;
+        out_idx = candidate_idx;
+      }
+    }
+    y[output_index] = select(out_idx, 0u, out_idx == SENTINEL_IDX);
+  }
+}

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

	@@ -0,0 +1,45 @@

+// 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;
+  }
+  let outer_index = output_index / params.innerSize;
+  let inner_index = output_index % params.innerSize;
+  let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
+  var best_index = 0u;
+{%- if scalar == "i32" or scalar == "u32" %}
+  var best = x[input_base];
+  for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
+    let candidate = x[input_base + axis_index * params.innerSize];
+    if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
+      best = candidate;
+      best_index = axis_index;
+    }
+  }
+{%- else %}
+  var best = f32(x[input_base]);
+  for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
+    let candidate = f32(x[input_base + axis_index * params.innerSize]);
+    if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
+      best = candidate;
+      best_index = axis_index;
+    }
+  }
+{%- endif %}
+  y[output_index] = best_index;
+}

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

	@@ -0,0 +1,159 @@

+{% set useSubgroups = source.useSubgroups %}
+{% if usesF16 %}
+enable f16;
+{%- endif %}
+{% if useSubgroups %}
+enable subgroups;
+{%- endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+const SENTINEL_IDX: u32 = 4294967295u;
+{%- if scalar == "i32" %}
+{% set work = "i32" %}
+{% elif scalar == "u32" %}
+{% set work = "u32" %}
+{% else %}
+{% set work = "f32" %}
+{% endif %}
+{% set cmp = ">" if 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 %}
+}
+var<workgroup> wg_val: array<{{ work }}, WG>;
+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;
+  let tid = lid.x;
+  let segment_size = (params.chunkCount + SPLIT - 1u) / SPLIT;
+  let segment_begin = split_index * segment_size;
+  let segment_end = min(segment_begin + segment_size, params.chunkCount);
+  let base = row * params.chunkCount;
+  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"] %}
+    {
+{%- if scalar == "f16" %}
+      let value = f32(value4.{{ component }});
+{%- else %}
+      let value = value4.{{ component }};
+{%- endif %}
+      let index = chunk * 4u + {{ loop.index0 }}u;
+{%- if selectLastIndex %}
+      if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || index > best_idx))) {
+{%- else %}
+      if (value {{ cmp }} best_val || (value == best_val && index < best_idx)) {
+{%- endif %}
+        best_val = value;
+        best_idx = index;
+      }
+    }
+{%- endfor %}
+{% else %}
+{% if scalar == "f16" %}
+    let value = f32(x[base + chunk]);
+{%- else %}
+    let value = x[base + chunk];
+{%- endif %}
+{% if selectLastIndex %}
+    if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || chunk > best_idx))) {
+{%- else %}
+    if (value {{ cmp }} best_val || (value == best_val && chunk < best_idx)) {
+{%- endif %}
+      best_val = value;
+      best_idx = chunk;
+    }
+{%- endif %}
+  }
+{%- if useSubgroups %}
+  let subgroup_val = subgroupMax(best_val);{% if selectLastIndex %}
+  let subgroup_candidate = select(0u, best_idx, best_val == subgroup_val && best_idx != SENTINEL_IDX);
+  let subgroup_idx = subgroupMax(subgroup_candidate);
+{%- else %}
+  let subgroup_candidate = select(SENTINEL_IDX, best_idx, best_val == subgroup_val);
+  let subgroup_idx = subgroupMin(subgroup_candidate);
+{%- endif %}
+  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;
+{%- else %}
+    partials_val[scratch_index] = bitcast<u32>(output_val);
+{%- endif %}
+    partials_idx[scratch_index] = output_idx;
+  }
+{%- else %}
+  wg_val[tid] = best_val;
+  wg_idx[tid] = best_idx;
+  workgroupBarrier();
+  for (var stride = WG / 2u; stride > 0u; stride >>= 1u) {
+    if (tid < stride) {
+      let candidate_val = wg_val[tid + stride];
+      let candidate_idx = wg_idx[tid + stride];
+      let current_idx = wg_idx[tid];
+{%- if selectLastIndex %}
+      if (candidate_idx != SENTINEL_IDX && (current_idx == SENTINEL_IDX || candidate_val {{ cmp }} wg_val[tid] || (candidate_val == wg_val[tid] && candidate_idx > current_idx))) {
+{%- else %}
+      if (candidate_idx != SENTINEL_IDX && (current_idx == SENTINEL_IDX || candidate_val {{ cmp }} wg_val[tid] || (candidate_val == wg_val[tid] && candidate_idx < current_idx))) {
+{%- endif %}
+        wg_val[tid] = candidate_val;
+        wg_idx[tid] = candidate_idx;
+      }
+    }
+    workgroupBarrier();
+  }
+  if (tid == 0u) {
+    let scratch_index = split_index * params.rows + row;
+{%- if scalar == "u32" %}
+    partials_val[scratch_index] = wg_val[0];
+{%- else %}
+    partials_val[scratch_index] = bitcast<u32>(wg_val[0]);
+{%- endif %}
+    partials_idx[scratch_index] = wg_idx[0];
+  }
+{%- endif %}
+}

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

	@@ -0,0 +1,185 @@

+// 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 %}
+{% if useSubgroups %}
+enable subgroups;
+{%- endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+const SENTINEL_IDX: u32 = 4294967295u;
+{% if source.mode == "max" and scalar != "i32" and scalar != "u32" %}
+fn neg_inf_f32() -> f32 {
+  var bits = 0xff800000u;
+  return bitcast<f32>(bits);
+}
+{% endif %}
+{%- if scalar == "i32" %}
+{% set cmpType = "i32" %}
+{% elif scalar == "u32" %}
+{% set cmpType = "u32" %}
+{% else %}
+{% set cmpType = "f32" %}
+{% endif %}
+{% set cmp = ">" %}
+{%- if 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;
+  }
+  let tid = lid.x;
+  let base = row * params.chunkCount;{% if cmpType == "i32" %}
+  var bestVal: i32 = -2147483647i - 1i;
+{%- elif cmpType == "u32" %}
+  var bestVal: u32 = 0u;
+{%- else %}
+  var bestVal: f32 = neg_inf_f32();
+{%- endif %}
+  var bestIdx: u32 = SENTINEL_IDX;
+{%- if source.vec4 %}
+  for (var c = tid; c < params.chunkCount; c = c + WG) {
+    let v4 = x[base + c];
+{%- for comp in ["x", "y", "z", "w"] %}
+    {
+{%- if scalar == "f16" %}
+      let v = f32(v4.{{ comp }});
+{%- else %}
+      let v = v4.{{ comp }};
+{%- endif %}
+      let idx = c * 4u + {{ loop.index0 }}u;
+{%- if selectLastIndex %}
+      if (v {{ cmp }} bestVal || (v == bestVal && (bestIdx == SENTINEL_IDX || idx > bestIdx))) {
+{%- else %}
+      if (v {{ cmp }} bestVal || (v == bestVal && idx < bestIdx)) {
+{%- endif %}
+        bestVal = v;
+        bestIdx = idx;
+      }
+    }
+{%- endfor %}
+  }
+{%- else %}
+  // Four independent, lane-coalesced loads per loop iteration. Keeping each
+  // unrolled plane WG elements apart preserves contiguous accesses across the
+  // workgroup while exposing enough ILP for long odd-width rows (where a vec4
+  // storage binding cannot represent row boundaries).
+  for (var c0 = tid; c0 < params.chunkCount; c0 = c0 + 4u * WG) {
+{%- for plane in [0, 1, 2, 3] %}
+    {
+      let c = c0 + {{ plane }}u * WG;
+      if (c < params.chunkCount) {
+{%- if scalar == "f16" %}
+        let v = f32(x[base + c]);
+{%- else %}
+        let v = x[base + c];
+{%- endif %}
+{% if selectLastIndex %}
+        if (v {{ cmp }} bestVal || (v == bestVal && (bestIdx == SENTINEL_IDX || c > bestIdx))) {
+{%- else %}
+        if (v {{ cmp }} bestVal || (v == bestVal && c < bestIdx)) {
+{%- endif %}
+          bestVal = v;
+          bestIdx = c;
+        }
+      }
+    }
+{%- endfor %}
+  }
+{%- endif %}
+{%- if useSubgroups %}
+  let m = subgroupMax(bestVal);{% if selectLastIndex %}
+  // Lanes that match the subgroup best contribute their (real) index; the
+  // rest contribute 0 so subgroupMax picks the highest matching index.
+  let cand = select(0u, bestIdx, bestVal == m && bestIdx != SENTINEL_IDX);
+  let sgIdx = subgroupMax(cand);
+{%- else %}
+  // Lanes that match the subgroup best contribute their index; the rest
+  // contribute the sentinel so subgroupMin picks the lowest matching index.
+  let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
+  let sgIdx = subgroupMin(cand);
+{%- endif %}
+{% if 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 %}
+      if (vi != SENTINEL_IDX && (v {{ cmp }} outVal || (v == outVal && (outIdx == SENTINEL_IDX || vi > outIdx)))) {
+{%- else %}
+      if (vi != SENTINEL_IDX && (v {{ cmp }} outVal || (v == outVal && vi < outIdx))) {
+{%- endif %}
+        outVal = v;
+        outIdx = vi;
+      }
+    }
+    y[row] = select(outIdx, 0u, outIdx == SENTINEL_IDX);
+  }
+{%- endif %}
+{% else %}
+  // No-subgroup tier: workgroup barrier tree-reduction over the per-thread
+  // (bestVal, bestIdx) pairs. Same ONNX tie rule (lowest index by default,
+  // highest with select_last_index); SENTINEL_IDX lanes (short rows, WG >
+  // chunkCount) never win. WG is a power of two,
+  // so the halving tree visits every lane.
+  wgVal[tid] = bestVal;
+  wgIdx[tid] = bestIdx;
+  workgroupBarrier();
+  for (var stride = WG / 2u; stride > 0u; stride = stride >> 1u) {
+    if (tid < stride) {
+      let v = wgVal[tid + stride];
+      let vi = wgIdx[tid + stride];
+      let curIdx = wgIdx[tid];
+{%- if selectLastIndex %}
+      if (vi != SENTINEL_IDX && (curIdx == SENTINEL_IDX || v {{ cmp }} wgVal[tid] || (v == wgVal[tid] && vi > curIdx))) {
+{%- else %}
+      if (vi != SENTINEL_IDX && (curIdx == SENTINEL_IDX || v {{ cmp }} wgVal[tid] || (v == wgVal[tid] && vi < curIdx))) {
+{%- endif %}
+        wgVal[tid] = v;
+        wgIdx[tid] = vi;
+      }
+    }
+    workgroupBarrier();
+  }
+  if (tid == 0u) {
+    y[row] = select(wgIdx[0], 0u, wgIdx[0] == SENTINEL_IDX);
+  }
+{%- endif %}
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,1266 @@

+{
+  "op": "ai.onnx.ArgMax",
+  "fixtureArrays": {
+    "onnx_backend_argmax_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": [1, 3, 3, 2, 1, 3, 3, 2, 1, 3, 3, 2, 1, 3, 3, 2, 1, 3, 3, 2, 1, 3, 3, 2, 1, 3, 3, 2, 1, 3, 3, 2, 1],
+    "subgroup_vec4_ties_first_index_input_x": [1, 5, 5, 2, 5, 0, 1, 0, 7, 7, 7, 7, 7, 7, 7, 7],
+    "rank4_axis0_f32_no_keepdims_input_x": [5, 1, 9, 2, 0, 7, 3, 4, 1, 6, 9, 8, 2, 7, 3, 1],
+    "axis0_splitk_i32_8192x16_precision_max_input_x": [16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217]
+  },
+  "cases": [
+    {
+      "name": "int16_last_axis_vec4_boundaries",
+      "attrs": { "axis": -1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int16",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [-32768, -1, 32767, 0, 32767, -32768, 32767, 1] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [2], "tolerance": 0, "data": { "kind": "values", "values": [2, 0] } }
+      }
+    },
+    {
+      "name": "onnx_default_axis0_omitted",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__ArgMax.html",
+        "test": "ArgMax 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": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [1.0, 9.0, 3.0, 4.0, 2.0, 8.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1, 3], "tolerance": 0, "data": { "kind": "values", "values": [1, 0, 1] } }
+      }
+    },
+    {
+      "name": "f32_positive_subnormal_beats_zero_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; bit-exact subnormal preservation is unattainable on GPU."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ArgMax",
+        "notes": "A positive subnormal is strictly greater than zero; index selection must not treat it as a tie with zero."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [-1e-40, 0.0, 1e-40] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "f32_positive_subnormal_beats_zero_select_last_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; bit-exact subnormal preservation is unattainable on GPU."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ArgMax2D_select_last",
+        "notes": "select_last_index must only apply to true ties; a positive subnormal at index 0 is greater than following zeros."
+      },
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [1e-40, 0.0, 0.0] } }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1], "tolerance": 0, "data": { "kind": "values", "values": [0] } }
+      }
+    },
+    {
+      "name": "f32_positive_subnormal_beats_zero_axis0_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; bit-exact subnormal preservation is unattainable on GPU."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ArgMax",
+        "notes": "Axis-0 companion for finite subnormal ordering: a positive subnormal is strictly greater than zero and must not be tie-broken as flushed zero."
+      },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [3, 1], "data": { "kind": "values", "values": [-1e-40, 0.0, 1e-40] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "dispatch_cliff_axis1_dim1_over_16M",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [16776961, 1], "data": { "kind": "constant", "value": 1.0 } } },
+      "outputs": { "y": { "dtype": "uint32", "shape": [16776961], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [1.0, 9.0, 3.0, 4.0, -1.0, 2.0, 7.0, 8.0, 0.0, 5.0, -3.0, 6.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [4] } }
+    },
+    {
+      "name": "axis0_splitk_8192x32_ties_first_index",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8192, 32],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_splitk_8192x32_ties_select_last_index",
+      "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8192, 32],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_splitk_8192x48_keepdims",
+      "attrs": { "axis": 0, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8192, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.2 }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 48], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_tiled_64x32_ties_first_index",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_tiled_64x32_ties_select_last_index",
+      "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "axis1_tie_first_index",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [1.0, 3.0, 3.0, 2.0, 3.0, -1.0, -1.0, -2.0, -1.0, -3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
+    },
+    {
+      "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.ArgMax",
+        "notes": "NaN extension: ORT seeds ArgMax 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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 4],
+          "data": {
+            "kind": "values",
+            "values": ["NaN", 1.0, 2.0, 0.0, 1.0, "NaN", 2.0, 0.0, 1.0, 2.0, "NaN", 0.0, "NaN", "NaN", "NaN", "NaN"]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "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",
+        "test": "ReductionOpTest.ArgMax_int32_last_index_dups",
+        "notes": "NaN extension: select_last_index only changes equal finite ties. A leading NaN remains the incumbent because finite values are not greater than NaN."
+      },
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 4],
+          "data": {
+            "kind": "values",
+            "values": ["NaN", 1.0, 2.0, 0.0, 1.0, "NaN", 3.0, 3.0, -1.0, -2.0, "NaN", -3.0, "NaN", "NaN", "NaN", "NaN"]
+          }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [4], "data": { "kind": "values", "values": [0, 3, 0, 0] }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "axis1_f16",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [1.0, 9.0, 3.0, 4.0, -1.0, 2.0, 7.0, 8.0, 0.0, 5.0, -3.0, 6.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3] } }
+    },
+    {
+      "name": "axis1_select_last_index_ties",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [1.0, 3.0, 3.0, 2.0, 3.0, -1.0, -1.0, -2.0, -1.0, -3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
+    },
+    {
+      "name": "axis1_select_last_all_negative_infinity",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": {
+            "kind": "values",
+            "values": ["-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", -5.0, "-Infinity", -5.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "axis_minus_one",
+      "attrs": { "axis": -1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [1.0, 2.0, 9.0, 4.0, 5.0, 8.0, 7.0, 6.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
+    },
+    {
+      "name": "axis1_keepdims",
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [1.0, 7.0, 7.0, 3.0, -2.0, -2.0, -5.0, -2.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1] } }
+    },
+    {
+      "name": "rank3_axis2_keepdims",
+      "attrs": { "axis": 2, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [1.0, 5.0, 5.0, 2.0, -1.0, -2.0, -3.0, -4.0, 0.0, 3.0, 2.0, 1.0, 9.0, 8.0, 9.0, 7.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "rank4_axis1_channel_select_last_ties",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 5.0, 3.0, 5.0, 3.0, 4.0, -1.0, 7.0, 3.0, 7.0, 0.0, 7.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "rank4_axis2_spatial_select_last_ties",
+      "attrs": { "axis": 2, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 4.0, 3.0, 4.0, 3.0, 2.0, -1.0, 5.0, 0.0, 5.0, 0.0, 3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "rank4_last_axis_nhwc_keepdims_select_last",
+      "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 2, 3],
+          "data": { "kind": "values", "values": [1.0, 3.0, 3.0, 5.0, 4.0, 7.0, 3.0, 3.0, 0.0, 5.0, 7.0, 7.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "rank1_axis0_scalar_output_select_last",
+      "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [5], "data": { "kind": "values", "values": [-1.0, 7.0, 3.0, 7.0, 2.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_int8_axis1_keepdims",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_int8"
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int8",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_uint8_axis1_keepdims",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_uint8"
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_int32_axis1_keepdims",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_int32",
+        "notes": "Native int32 ArgMax fixture from ORT; output indices use this framework's uint32 convention."
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_int32_axis1_keepdims_select_last_nodups",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_int32_last_index_nodups",
+        "notes": "ORT select_last_index path where no ties are present; should match the regular max positions."
+      },
+      "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [2, 2, 3, 4, 5, 6, 7, 8, 10, 10, 11, 12] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_float_first_index_infinity_deterministic",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_float_first_index_random",
+        "notes": "Deterministic compact projection of ORT's random infinity-sentinel first-index test."
+      },
+      "attrs": { "axis": 0, "keepdims": 1, "select_last_index": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [16],
+          "data": {
+            "kind": "values",
+            "values": [0.0, 1.0, "Infinity", 5.0, "Infinity", -3.0, 2.0, "Infinity", 4.0, -1.0, 3.0, "Infinity", 0.0, 2.0, 1.0, 0.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_int32_axis1_keepdims_select_last_dups",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_int32_last_index_dups"
+      },
+      "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [2, 4, 3, 4, 8, 6, 7, 8, 9, 13, 11, 12] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis1_no_keepdims_rank3_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_do_not_keepdims"
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis0_no_keepdims_rank1_scalar_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_do_not_keepdims_2"
+      },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1.0, 2.0, 3.0] } } },
+      "outputs": { "y": { "dtype": "uint32", "shape": [], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_negative_axis_int32_keepdims",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax_int32_neg_axis"
+      },
+      "attrs": { "axis": -2, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis1_keepdims_rank3_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax"
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis1_keepdims_rank2_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax2D"
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 6.0, 5.0, 9.0, 10.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis1_keepdims_rank2_select_last_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax2D_select_last"
+      },
+      "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 6.0, 5.0, 9.0, 10.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis1_dim1_rank2_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.ArgMax2D_dim1"
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [3, 1], "data": { "kind": "values", "values": [1.0, 6.0, 9.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "int32_axis1_exact_above_float24",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [16777216, 16777217, 5, -1, 123456788, 123456789, 0, -7] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_default_axis_example",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_default_axis_example",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "keepdims": 1, "axis": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_default_axis_example_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_default_axis_random",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_default_axis_random",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "keepdims": 1, "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 3, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_default_axis_random_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 3, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_keepdims_example",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_keepdims_example",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_keepdims_example_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_keepdims_random",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_keepdims_random",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "axis": 1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_keepdims_random_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_negative_axis_keepdims_example",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_negative_axis_keepdims_example",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "axis": -1, "keepdims": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_negative_axis_keepdims_example_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_negative_axis_keepdims_random",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_negative_axis_keepdims_random",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "axis": -1, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 3, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_negative_axis_keepdims_random_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 3, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_no_keepdims_example",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_no_keepdims_example",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_no_keepdims_example_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_no_keepdims_random",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_no_keepdims_random",
+        "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "onnx_backend_argmax_no_keepdims_random_select_last_index",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmax_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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "subgroup_vec4_ties_first_index",
+      "attrs": { "axis": -1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [1, 0] } } }
+    },
+    {
+      "name": "subgroup_vec4_ties_select_last_index",
+      "attrs": { "axis": -1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [4, 7] } } }
+    },
+    {
+      "name": "subgroup_scalar_f16_2x5",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [1.5, 2.5, 2.5, 0.5, 1.0, -1.0, -2.0, -0.5, -0.5, -3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [1, 2] } } }
+    },
+    {
+      "name": "subgroup_vec4_uint8_full_range",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "uint8", "shape": [2, 4], "data": { "kind": "values", "values": [255, 0, 254, 1, 2, 2, 1, 0] } }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [0, 0] } } }
+    },
+    {
+      "name": "subgroup_vec4_long_row_2x256",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
+    },
+    {
+      "name": "subgroup_min_vec4_rows128x256_ties_first",
+      "provenance": {
+        "notes": "Route-locks the one-minimum-subgroup vec4 path on variable-width adapters and verifies first-index tie ordering without a workgroup-memory combine."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [128, 256], "data": { "kind": "constant", "value": 7.0 } } },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [128], "data": { "kind": "constant", "value": 0 }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "subgroup_min_scalar_rows128x257_ties_last",
+      "provenance": {
+        "notes": "Route-locks the unaligned one-minimum-subgroup path and verifies select_last_index ordering across all lane-local chunks."
+      },
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [128, 257], "data": { "kind": "constant", "value": 7.0 } } },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [128], "data": { "kind": "constant", "value": 256 }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "subgroup_scalar_long_row_2x65",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.11 }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
+    },
+    {
+      "name": "empty_zero_dim",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [0, 4], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "uint32", "shape": [0], "tolerance": 0 } }
+    },
+    {
+      "name": "empty_zero_dim_f16",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "float16", "shape": [0, 4], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "uint32", "shape": [0], "tolerance": 0 } }
+    },
+    {
+      "name": "empty_zero_dim_int32",
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": { "x": { "dtype": "int32", "shape": [0, 4], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "uint32", "shape": [0], "tolerance": 0 } }
+    },
+    {
+      "name": "rank4_axis0_f32_no_keepdims",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/rank4_axis0_f32_no_keepdims_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "rank4_axis0_int32_keepdims",
+      "attrs": { "axis": 0, "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 2, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/rank4_axis0_f32_no_keepdims_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "rank4_neg_axis0_select_last",
+      "attrs": { "axis": -4, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/rank4_axis0_f32_no_keepdims_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "axis1_serial_fallback_rows32768_cols5_select_last",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32768, 5],
+          "data": { "kind": "cycle", "values": [1.0, 3.0, 3.0, 2.0, 3.0, 1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32768], "tolerance": 0 } }
+    },
+    {
+      "name": "axis1_serial_fallback_rows32768_cols5_first_index",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32768, 5],
+          "data": { "kind": "cycle", "values": [1.0, 3.0, 3.0, 2.0, 3.0, 1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [32768], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_fallback_rows4_int32_select_last_ties",
+      "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [4, 4],
+          "data": { "kind": "values", "values": [5, 1, 3, 2, 5, 7, 3, 2, 4, 7, 9, 2, 4, 1, 9, 8] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [4], "tolerance": 0, "data": { "kind": "values", "values": [1, 2, 3, 3] } }
+      }
+    },
+    {
+      "name": "axis1_int8_serial_fallback_select_last",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int8",
+          "shape": [2, 8],
+          "data": {
+            "kind": "values",
+            "values": [-5, -5, 127, 127, 3, 3, -128, -128, 0, -128, 0, 127, 100, 100, 100, 50]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "axis1_uint8_serial_fallback_boundary_values",
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [3, 8],
+          "data": {
+            "kind": "values",
+            "values": [255, 0, 255, 128, 1, 254, 255, 200, 0, 0, 0, 0, 0, 0, 0, 0, 128, 128, 255, 255, 100, 100, 255, 255]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [3], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_splitk_i32_8192x16_precision_max",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [8192, 16],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_max_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_splitk_u32_8192x16_precision_max",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [8192, 16],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_max_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_tilecols_i32_64x16_precision_max",
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [64, 16],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_max_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_tilecols_i32_64x16_precision_select_last_max",
+      "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [64, 16],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_max_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
+    },
+    {
+      "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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [16384, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [8], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_narrow_splitk_8192x4_int32",
+      "provenance": {
+        "notes": "Tall narrow axis-0 (16384 rows, 4 cols): cols<16 now routes to axis0_splitk (split-K over rows) instead of the few-lane axis0 serial path (which gives one thread per column -> only 4 lanes). Cross-checks split-K against the single-lane axis0 kernel bit-for-bit; reference is the oracle. int32 exercises the native-integer split-K compare path."
+      },
+      "attrs": { "axis": 0, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [8192, 4],
+          "data": { "kind": "cycle", "values": [-5000, 1200, -300, 4800, 700] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "rank3_axis1_splitk_temporal_select_last",
+      "provenance": {
+        "notes": "A long temporal middle axis with only four output classes used to launch four serial reduction lanes. The strided split-K path must map (batch, time, class) correctly and merge equal maxima to the last temporal index."
+      },
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 8192, 4], "data": { "kind": "constant", "value": 7.0 } } },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1, 4], "data": { "kind": "constant", "value": 8191 }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "last_axis_split_vec4_logits_ties_first",
+      "provenance": {
+        "notes": "A single 32768-element row route-locks the cooperative split-row vec4 reduction. Equal maxima span many split segments and must merge to the first global index."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 32768], "data": { "kind": "cycle", "values": [7.0, 3.0, 7.0, 1.0] } }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1], "data": { "kind": "constant", "value": 0 }, "tolerance": 0 }
+      }
+    },
+    {
+      "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 } } },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1], "data": { "kind": "constant", "value": 32768 }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "last_axis_split_vec4_i32_negative_rows",
+      "provenance": {
+        "notes": "int32 on the cooperative split-row vec4 reduction (32768 columns, two rows). Every value in the row is negative, so a lane identity of 0 instead of INT_MIN would beat the whole row and the merge would return the sentinel; this directly exercises the i32 identity and native signed compare. The period-7 cycle does not divide 32768, so the two rows are phase-shifted and their answers differ — a kernel that reduced row 0 for both rows cannot pass. ONNX ArgMax admits int32 via T and breaks ties on the first index."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 32768],
+          "data": { "kind": "cycle", "values": [-9, -5, -3, -1, -2, -4, -6] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [3, 2] }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "last_axis_split_scalar_u32_above_int32_max",
+      "provenance": {
+        "notes": "uint32 on the unaligned (32769-column) scalar split-row reduction. Six of the seven cycle values exceed 2^31, so a signed compare would rank them negative and answer [0, 5] instead of [3, 1]. This directly exercises the unsigned identity 0u and the partial store that must write the value through without bitcasting it. ONNX ArgMax admits uint32 via T; first index on ties."
+      },
+      "attrs": { "axis": 1, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [2, 32769],
+          "data": {
+            "kind": "cycle",
+            "values": [2147483000, 4200000000, 3000000000, 4294967290, 2500000000, 3500000000, 4000000000]
+          }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [3, 1] }, "tolerance": 0 }
+      }
+    },
+    {
+      "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 ArgMax admits float16 via T; first index on ties."
+      },
+      "attrs": { "axis": 2, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 2, 32768],
+          "data": { "kind": "cycle", "values": [1.5, -2.0, 0.5, 3.5, -1.0, 2.25, -3.75] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1, 2], "data": { "kind": "values", "values": [3, 2] }, "tolerance": 0 }
+      }
+    },
+    {
+      "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 ArgMax admits float16 via T; first index on ties."
+      },
+      "attrs": { "axis": 2, "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 2, 32769],
+          "data": { "kind": "cycle", "values": [1.5, -2.0, 0.5, 3.5, -1.0, 2.25, -3.75] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [1, 2], "data": { "kind": "values", "values": [3, 1] }, "tolerance": 0 }
+      }
+    },
+    {
+      "name": "last_axis_split_vec4_ties_select_last_index",
+      "provenance": {
+        "notes": "select_last_index on the vec4 split-row route checks the last-index tie rule in both the segment reduction and combine. The maximum 7.0 recurs at cycle offsets 0, 2 and 5, so every segment reports a tie; the two phase-shifted rows must return 32767 and 32766."
+      },
+      "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 32768],
+          "data": { "kind": "cycle", "values": [7.0, 3.0, 7.0, 1.0, 5.0, 7.0, 2.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [32767, 32766] }, "tolerance": 0 }
+      }
+    },
+    {
+      "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": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 3, 2, 2],
+          "data": { "kind": "cycle", "values": [-3.0, 5.0, 1.0, 5.0, 2.0, -1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 2, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "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 maximum 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": {
+        "x": {
+          "dtype": "int32",
+          "shape": [8192, 16],
+          "data": { "kind": "cycle", "values": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
+    }
+  ]
+}