Xenova HF Staff commited on
Commit
7c6615c
·
verified ·
1 Parent(s): 8d1ebb5

sync 2e7068faf55e

Browse files
README.md CHANGED
@@ -1,3 +1,80 @@
1
  ---
 
2
  license: apache-2.0
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ library_name: kernels
3
  license: apache-2.0
4
+ tags:
5
+ - kernel
6
+ - webgpu
7
+ - wgsl
8
  ---
9
+ # ai.onnx.ArgMin
10
+
11
+ `ai.onnx` · standard ONNX operator · ONNX opset ≥ 13
12
+
13
+ ## Description
14
+
15
+ Returns the index of the minimum value along an axis, choosing the first equal value unless `select_last_index` is enabled.
16
+
17
+ See the [ONNX `ArgMin` spec](https://onnx.ai/onnx/operators/onnx__ArgMin.html) for the reference semantics.
18
+
19
+ ## Inputs
20
+
21
+ | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
+ | --- | --- | --- | --- | --- | --- | --- |
23
+ | `data` | `x` | `T` | — | — | Values whose minimum index is selected along `axis`. | required |
24
+
25
+ ## Outputs
26
+
27
+ | Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
28
+ | --- | --- | --- | --- | --- | --- | --- | --- |
29
+ | `reduced` | `y` | `I` | `uint32` | derived | derived; see description | Logical int64 indices of the minimum values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
30
+
31
+ ## Attributes
32
+
33
+ Default values (overridable per request):
34
+
35
+ | Attribute | Default | Description |
36
+ | --- | --- | --- |
37
+ | `axis` | `0` | Axis to reduce; negative values count from the back. |
38
+ | `keepdims` | `1` | Retain the reduced dimension with length one when non-zero. |
39
+ | `select_last_index` | `0` | Choose the last equal minimum instead of the first when non-zero. |
40
+
41
+ ## Type constraints
42
+
43
+ | Variable | Allowed dtypes |
44
+ | --- | --- |
45
+ | `T` | `float32`, `float16`, `int32`, `uint32`, `int16`, `int8`, `uint8` |
46
+ | `I` | `int64` |
47
+
48
+ ## Device requirements
49
+
50
+ Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
51
+
52
+ ## Files
53
+
54
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
55
+ - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
56
+ - [`test.json`](build/webgpu/test.json) — correctness cases
57
+ - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
58
+ - [`reduce-arg-axis-split-combine.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja)
59
+ - [`reduce-arg-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja)
60
+ - [`reduce-arg-axis-split-tiled.wgsl.jinja`](build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja)
61
+ - [`reduce-arg-axis-tiled.wgsl.jinja`](build/webgpu/reduce-arg-axis-tiled.wgsl.jinja)
62
+ - [`reduce-arg-axis.wgsl.jinja`](build/webgpu/reduce-arg-axis.wgsl.jinja)
63
+ - [`reduce-arg-row-split.wgsl.jinja`](build/webgpu/reduce-arg-row-split.wgsl.jinja)
64
+ - [`reduce-arg-row-subgroup.wgsl.jinja`](build/webgpu/reduce-arg-row-subgroup.wgsl.jinja)
65
+
66
+ ## Use with `@huggingface/kernels`
67
+
68
+ The loader derives every required output's shape and logical dtype from the manifest contract and this call.
69
+ It then allocates the result tensors automatically.
70
+
71
+ The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
72
+
73
+ Replace each `*Data` placeholder with a typed array containing the corresponding input data.
74
+
75
+ ```js
76
+ import { getKernel } from "@huggingface/kernels";
77
+
78
+ const kernel = await getKernel("webgpu-kernels/ai.onnx.ArgMin", { version: 1 });
79
+ const { y } = await kernel({ x: { data: xData, shape: [2, 2] } });
80
+ ```
build/webgpu/bench.json ADDED
@@ -0,0 +1,101 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "op": "ai.onnx.ArgMin",
3
+ "cases": [
4
+ {
5
+ "name": "argmin-axis0-f32-1024x512",
6
+ "preset": "smoke",
7
+ "vars": { "dtype": "float32", "rows": 1024, "cols": 512 },
8
+ "attrs": { "axis": 0, "keepdims": 0 },
9
+ "inputs": { "x": { "shape": [1024, 512], "dtype": "float32", "dist": "normal", "seed": 977, "scale": 1 } },
10
+ "outputs": { "y": { "shape": [512], "dtype": "uint32" } },
11
+ "bench": {
12
+ "primary": true,
13
+ "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }]
14
+ }
15
+ },
16
+ {
17
+ "name": "scores",
18
+ "attrs": { "axis": 1, "keepdims": 0 },
19
+ "inputs": { "x": { "dtype": "float32", "shape": [1024, 768] } },
20
+ "outputs": { "y": { "dtype": "uint32", "shape": [1024] } }
21
+ },
22
+ {
23
+ "name": "argmin-axis0-large-65536x512",
24
+ "preset": "smoke",
25
+ "vars": { "dtype": "float32", "rows": 65536, "cols": 512 },
26
+ "attrs": { "axis": 0, "keepdims": 0 },
27
+ "inputs": { "x": { "shape": [65536, 512], "dtype": "float32", "dist": "normal", "seed": 211, "scale": 0.2 } },
28
+ "outputs": { "y": { "shape": [512], "dtype": "uint32" } },
29
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
30
+ },
31
+ {
32
+ "name": "argmin-lastaxis-f32-4096x4096",
33
+ "preset": "smoke",
34
+ "vars": { "rows": 4096, "cols": 4096 },
35
+ "attrs": { "axis": 1, "keepdims": 0 },
36
+ "inputs": { "x": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 117, "scale": 0.2 } },
37
+ "outputs": { "y": { "shape": [4096], "dtype": "uint32" } },
38
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
39
+ },
40
+ {
41
+ "name": "argmin-lastaxis-f32-4096x4095",
42
+ "preset": "smoke",
43
+ "vars": { "dtype": "float32", "rows": 4096, "cols": 4095 },
44
+ "attrs": { "axis": 1, "keepdims": 0 },
45
+ "inputs": { "x": { "shape": [4096, 4095], "dtype": "float32", "dist": "normal", "seed": 419, "scale": 0.2 } },
46
+ "outputs": { "y": { "shape": [4096], "dtype": "uint32" } },
47
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
48
+ },
49
+ {
50
+ "name": "argmin-vocab-unaligned-f32-1x50257",
51
+ "preset": "smoke",
52
+ "vars": { "dtype": "float32", "rows": 1, "cols": 50257 },
53
+ "attrs": { "axis": 1, "keepdims": 0 },
54
+ "inputs": { "x": { "shape": [1, 50257], "dtype": "float32", "dist": "normal", "seed": 743, "scale": 1 } },
55
+ "outputs": { "y": { "shape": [1], "dtype": "uint32" } },
56
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
57
+ },
58
+ {
59
+ "name": "argmin-lastaxis-f16-1024x768",
60
+ "preset": "smoke",
61
+ "vars": { "dtype": "float16", "rows": 1024, "cols": 768 },
62
+ "attrs": { "axis": 1, "keepdims": 0 },
63
+ "inputs": { "x": { "shape": [1024, 768], "dtype": "float16", "dist": "normal", "seed": 523, "scale": 0.2 } },
64
+ "outputs": { "y": { "shape": [1024], "dtype": "uint32" } },
65
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
66
+ },
67
+ {
68
+ "name": "argmin-axis0-f32-32x512",
69
+ "preset": "smoke",
70
+ "vars": { "dtype": "float32", "rows": 32, "cols": 512 },
71
+ "attrs": { "axis": 0, "keepdims": 0 },
72
+ "inputs": { "x": { "shape": [32, 512], "dtype": "float32", "dist": "normal", "seed": 631, "scale": 0.2 } },
73
+ "outputs": { "y": { "shape": [512], "dtype": "uint32" } },
74
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
75
+ },
76
+ {
77
+ "name": "argmin-lastaxis-uint8-2048x4096-serial",
78
+ "preset": "stress",
79
+ "vars": { "dtype": "uint8", "rows": 2048, "cols": 4096 },
80
+ "attrs": { "axis": 1, "keepdims": 0 },
81
+ "inputs": {
82
+ "x": { "shape": [2048, 4096], "dtype": "uint8", "dist": "uniform", "seed": 743, "min": 0, "max": 255 }
83
+ },
84
+ "outputs": { "y": { "shape": [2048], "dtype": "uint32" } },
85
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
86
+ },
87
+ {
88
+ "name": "argmin-axis0-narrow-starvation-131072x8",
89
+ "preset": "smoke",
90
+ "vars": { "dtype": "float32", "rows": 131072, "cols": 8 },
91
+ "attrs": { "axis": 0, "keepdims": 0 },
92
+ "inputs": { "x": { "shape": [131072, 8], "dtype": "float32", "dist": "normal", "seed": 827, "scale": 1 } },
93
+ "outputs": { "y": { "shape": [8], "dtype": "uint32" } },
94
+ "bench": {
95
+ "primary": true,
96
+ "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }]
97
+ }
98
+ }
99
+ ],
100
+ "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }
101
+ }
build/webgpu/manifest.json ADDED
@@ -0,0 +1,524 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "domain": "ai.onnx",
3
+ "name": "ArgMin",
4
+ "sinceVersion": 13,
5
+ "description": "Returns the index of the minimum value along an axis, choosing the first equal value unless `select_last_index` is enabled.",
6
+ "inputs": [{ "role": "data", "dtype": "T", "description": "Values whose minimum index is selected along `axis`." }],
7
+ "outputs": [
8
+ {
9
+ "role": "reduced",
10
+ "dtype": "I",
11
+ "rank": "ranks.data if attrs.keepdims == 1 else ranks.data - 1",
12
+ "description": "Logical int64 indices of the minimum values along the reduced axis; WebGPU stores these bounded indices as uint32.",
13
+ "shape": "prefix(shapes.data, normalizedAxis) + ([1] if attrs.keepdims == 1 else []) + suffix(shapes.data, normalizedAxis + 1)"
14
+ }
15
+ ],
16
+ "attributes": { "axis": 0, "keepdims": 1, "select_last_index": 0 },
17
+ "attributeDescriptions": {
18
+ "axis": "Axis to reduce; negative values count from the back.",
19
+ "keepdims": "Retain the reduced dimension with length one when non-zero.",
20
+ "select_last_index": "Choose the last equal minimum instead of the first when non-zero."
21
+ },
22
+ "attributeConstraints": { "keepdims": { "values": [0, 1] }, "select_last_index": { "values": [0, 1] } },
23
+ "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"], "I": ["int64"] },
24
+ "args": {
25
+ "x": { "kind": "tensor", "semantic": "data", "role": "input" },
26
+ "y": { "kind": "tensor", "semantic": "reduced", "role": "output", "dtype": "uint32" }
27
+ },
28
+ "tunables": {
29
+ "WORKGROUP_SIZE": 256,
30
+ "TILE_COLUMNS": 16,
31
+ "TILED_MIN_AXIS": 64,
32
+ "TILED_MIN_OUTPUTS": 16,
33
+ "SPLIT_MIN_AXIS": 8192,
34
+ "SPLIT_MAX_OUTPUTS": 4096,
35
+ "SPLIT_TARGET_AXIS": 256,
36
+ "SPLIT_TILE_COLUMNS": 8,
37
+ "MAX_SPLITS": 128,
38
+ "ROW_SPLIT_MIN_AXIS": 32768,
39
+ "ROW_SPLIT_MAX_OUTPUTS": 32,
40
+ "ROW_SPLIT_TARGET_AXIS": 2048,
41
+ "ROW_SPLIT_MAX_SPLITS": 64
42
+ },
43
+ "derive": {
44
+ "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
45
+ "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
46
+ "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
47
+ "normalizedAxis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data",
48
+ "axisInRange": "ranks.data >= 1 and normalizedAxis >= 0 and normalizedAxis < ranks.data",
49
+ "axisDim": "dim(shapes.data, normalizedAxis)",
50
+ "axisInner": "inner(shapes.data, normalizedAxis)",
51
+ "outputCount": "numel(shapes.reduced)",
52
+ "outputRankOk": "(attrs.keepdims == 1 and ranks.reduced == ranks.data) or (attrs.keepdims == 0 and ranks.reduced == ranks.data - 1)",
53
+ "outputShapeOk": "outputRankOk and outputCount * axisDim == numel(shapes.data)",
54
+ "baseContract": "axisInRange and axisDim > 0 and outputShapeOk and f16Ok(dtypes.T)",
55
+ "lastAxis": "normalizedAxis == ranks.data - 1",
56
+ "nativeParallelType": "dtypes.T == \"f32\" or dtypes.T == \"i32\" or dtypes.T == \"u32\"",
57
+ "rowParallelType": "nativeParallelType or dtypes.T == \"f16\"",
58
+ "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
59
+ "narrowLargeAxisTilePreferred": "narrowSubgroupRange and outputCount >= workgroupSize",
60
+ "narrowOddRowPortablePreferred": "narrowSubgroupRange and axisDim >= workgroupSize * 8 and outputCount >= workgroupSize",
61
+ "threadDispatchFits": "ceilDiv(outputCount, workgroupSize) <= foldedDispatchCapacity",
62
+ "rowDispatchFits": "outputCount <= foldedDispatchCapacity",
63
+ "rowStorageFits": "workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
64
+ "parallelRowWorthwhile": "axisDim >= 256 or outputCount < 32768",
65
+ "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.SPLIT_TARGET_AXIS)))",
66
+ "splitScratchBytes": "splitCount * outputCount * 8",
67
+ "splitScratchFits": "splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize",
68
+ "splitDispatchFits": "splitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(outputCount, workgroupSize), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
69
+ "splitTileDispatchFits": "splitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS) <= device.limits.maxComputeWorkgroupsPerDimension",
70
+ "rowSplitCount": "min(tunables.ROW_SPLIT_MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.ROW_SPLIT_TARGET_AXIS)))",
71
+ "rowSplitScratchBytes": "rowSplitCount * outputCount * 8",
72
+ "rowSplitScratchFits": "rowSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and rowSplitScratchBytes <= device.limits.maxBufferSize",
73
+ "rowSplitDispatchFits": "outputCount <= device.limits.maxComputeWorkgroupsPerDimension and rowSplitCount <= device.limits.maxComputeWorkgroupsPerDimension",
74
+ "tileDispatchFits": "ceilDiv(outputCount, tunables.TILE_COLUMNS) <= foldedDispatchCapacity"
75
+ },
76
+ "bindingSets": {
77
+ "contiguousVec4": [
78
+ {
79
+ "name": "x",
80
+ "arg": "x",
81
+ "semantic": "data",
82
+ "buffer": { "type": "read-only-storage" },
83
+ "elementType": "$vectorScalar"
84
+ },
85
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
86
+ {
87
+ "name": "params",
88
+ "semantic": "kernel.params",
89
+ "buffer": { "type": "uniform" },
90
+ "struct": {
91
+ "name": "Params",
92
+ "fields": [
93
+ { "name": "rows", "type": "u32", "value": "outputCount" },
94
+ { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" },
95
+ {
96
+ "name": "rowStride",
97
+ "type": "u32",
98
+ "value": "max(1, min(outputCount, device.limits.maxComputeWorkgroupsPerDimension))"
99
+ }
100
+ ]
101
+ }
102
+ }
103
+ ],
104
+ "contiguousScalar": [
105
+ {
106
+ "name": "x",
107
+ "arg": "x",
108
+ "semantic": "data",
109
+ "buffer": { "type": "read-only-storage" },
110
+ "elementType": "$scalar"
111
+ },
112
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
113
+ {
114
+ "name": "params",
115
+ "semantic": "kernel.params",
116
+ "buffer": { "type": "uniform" },
117
+ "struct": {
118
+ "name": "Params",
119
+ "fields": [
120
+ { "name": "rows", "type": "u32", "value": "outputCount" },
121
+ { "name": "chunkCount", "type": "u32", "value": "axisDim" },
122
+ {
123
+ "name": "rowStride",
124
+ "type": "u32",
125
+ "value": "max(1, min(outputCount, device.limits.maxComputeWorkgroupsPerDimension))"
126
+ }
127
+ ]
128
+ }
129
+ }
130
+ ],
131
+ "axisGeometry": [
132
+ {
133
+ "name": "x",
134
+ "arg": "x",
135
+ "semantic": "data",
136
+ "buffer": { "type": "read-only-storage" },
137
+ "elementType": "$scalar"
138
+ },
139
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
140
+ {
141
+ "name": "params",
142
+ "semantic": "kernel.params",
143
+ "buffer": { "type": "uniform" },
144
+ "struct": {
145
+ "name": "Params",
146
+ "fields": [
147
+ { "name": "axisDim", "type": "u32", "value": "axisDim" },
148
+ { "name": "outputCount", "type": "u32", "value": "outputCount" },
149
+ { "name": "innerSize", "type": "u32", "value": "axisInner" }
150
+ ]
151
+ }
152
+ }
153
+ ],
154
+ "rowSplitVec4": [
155
+ {
156
+ "name": "x",
157
+ "arg": "x",
158
+ "semantic": "data",
159
+ "buffer": { "type": "read-only-storage" },
160
+ "elementType": "$vectorScalar"
161
+ },
162
+ { "name": "partials_val", "semantic": "partials_val", "buffer": { "type": "storage" }, "elementType": "u32" },
163
+ { "name": "partials_idx", "semantic": "partials_idx", "buffer": { "type": "storage" }, "elementType": "u32" },
164
+ {
165
+ "name": "params",
166
+ "semantic": "kernel.params",
167
+ "buffer": { "type": "uniform" },
168
+ "struct": {
169
+ "name": "Params",
170
+ "fields": [
171
+ { "name": "rows", "type": "u32", "value": "outputCount" },
172
+ { "name": "chunkCount", "type": "u32", "value": "axisDim / 4" }
173
+ ]
174
+ }
175
+ }
176
+ ],
177
+ "rowSplitScalar": [
178
+ {
179
+ "name": "x",
180
+ "arg": "x",
181
+ "semantic": "data",
182
+ "buffer": { "type": "read-only-storage" },
183
+ "elementType": "$scalar"
184
+ },
185
+ { "name": "partials_val", "semantic": "partials_val", "buffer": { "type": "storage" }, "elementType": "u32" },
186
+ { "name": "partials_idx", "semantic": "partials_idx", "buffer": { "type": "storage" }, "elementType": "u32" },
187
+ {
188
+ "name": "params",
189
+ "semantic": "kernel.params",
190
+ "buffer": { "type": "uniform" },
191
+ "struct": {
192
+ "name": "Params",
193
+ "fields": [
194
+ { "name": "rows", "type": "u32", "value": "outputCount" },
195
+ { "name": "chunkCount", "type": "u32", "value": "axisDim" }
196
+ ]
197
+ }
198
+ }
199
+ ],
200
+ "splitReduce": [
201
+ {
202
+ "name": "x",
203
+ "arg": "x",
204
+ "semantic": "data",
205
+ "buffer": { "type": "read-only-storage" },
206
+ "elementType": "$scalar"
207
+ },
208
+ { "name": "partials_val", "semantic": "partials_val", "buffer": { "type": "storage" }, "elementType": "u32" },
209
+ { "name": "partials_idx", "semantic": "partials_idx", "buffer": { "type": "storage" }, "elementType": "u32" },
210
+ {
211
+ "name": "params",
212
+ "semantic": "kernel.params",
213
+ "buffer": { "type": "uniform" },
214
+ "struct": {
215
+ "name": "Params",
216
+ "fields": [
217
+ { "name": "axisDim", "type": "u32", "value": "axisDim" },
218
+ { "name": "outputCount", "type": "u32", "value": "outputCount" },
219
+ { "name": "innerSize", "type": "u32", "value": "axisInner" }
220
+ ]
221
+ }
222
+ }
223
+ ],
224
+ "splitCombine": [
225
+ {
226
+ "name": "partials_val",
227
+ "semantic": "partials_val",
228
+ "buffer": { "type": "read-only-storage" },
229
+ "elementType": "u32"
230
+ },
231
+ {
232
+ "name": "partials_idx",
233
+ "semantic": "partials_idx",
234
+ "buffer": { "type": "read-only-storage" },
235
+ "elementType": "u32"
236
+ },
237
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "u32" },
238
+ {
239
+ "name": "params",
240
+ "semantic": "kernel.params",
241
+ "buffer": { "type": "uniform" },
242
+ "struct": { "name": "Params", "fields": [{ "name": "outputCount", "type": "u32", "value": "outputCount" }] }
243
+ }
244
+ ]
245
+ },
246
+ "variants": [
247
+ {
248
+ "id": "last_axis_split_vec4",
249
+ "priority": 41,
250
+ "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= tunables.ROW_SPLIT_MIN_AXIS", "axisDim % 4 == 0", "outputCount <= tunables.ROW_SPLIT_MAX_OUTPUTS", "rowSplitScratchFits", "rowSplitDispatchFits", "rowStorageFits"],
251
+ "constants": {
252
+ "scalar": "dtypes.T",
253
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
254
+ "usesF16": "dtypes.T == \"f16\"",
255
+ "selectLastIndex": "attrs.select_last_index != 0",
256
+ "workgroupSize": "workgroupSize",
257
+ "split": "rowSplitCount"
258
+ },
259
+ "intermediates": [
260
+ { "id": "partials_val", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" },
261
+ { "id": "partials_idx", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" }
262
+ ],
263
+ "passes": [
264
+ {
265
+ "id": "split_reduce",
266
+ "name": "ArgMin.LastAxisSplitVec4Reduce",
267
+ "source": {
268
+ "shader": "reduce-arg-row-split.wgsl.jinja",
269
+ "inputs": { "mode": "\"min\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" }
270
+ },
271
+ "bindings": "rowSplitVec4",
272
+ "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
273
+ },
274
+ {
275
+ "id": "combine",
276
+ "name": "ArgMin.LastAxisSplitCombine",
277
+ "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
278
+ "bindings": "splitCombine",
279
+ "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
280
+ }
281
+ ]
282
+ },
283
+ {
284
+ "id": "last_axis_split",
285
+ "priority": 41,
286
+ "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= tunables.ROW_SPLIT_MIN_AXIS", "axisDim % 4 != 0", "outputCount <= tunables.ROW_SPLIT_MAX_OUTPUTS", "rowSplitScratchFits", "rowSplitDispatchFits", "rowStorageFits"],
287
+ "constants": {
288
+ "scalar": "dtypes.T",
289
+ "usesF16": "dtypes.T == \"f16\"",
290
+ "selectLastIndex": "attrs.select_last_index != 0",
291
+ "workgroupSize": "workgroupSize",
292
+ "split": "rowSplitCount"
293
+ },
294
+ "intermediates": [
295
+ { "id": "partials_val", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" },
296
+ { "id": "partials_idx", "dtype": "uint32", "shape": "[rowSplitCount * outputCount]" }
297
+ ],
298
+ "passes": [
299
+ {
300
+ "id": "split_reduce",
301
+ "name": "ArgMin.LastAxisSplitReduce",
302
+ "source": {
303
+ "shader": "reduce-arg-row-split.wgsl.jinja",
304
+ "inputs": { "mode": "\"min\"", "vec4": false, "useSubgroups": "device.features.has(\"subgroups\")" }
305
+ },
306
+ "bindings": "rowSplitScalar",
307
+ "dispatch": { "x": "outputCount", "y": "rowSplitCount" }
308
+ },
309
+ {
310
+ "id": "combine",
311
+ "name": "ArgMin.LastAxisSplitCombine",
312
+ "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
313
+ "bindings": "splitCombine",
314
+ "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
315
+ }
316
+ ]
317
+ },
318
+ {
319
+ "id": "axis_split_tiled_narrow",
320
+ "priority": 31,
321
+ "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"],
322
+ "constants": {
323
+ "scalar": "dtypes.T",
324
+ "selectLastIndex": "attrs.select_last_index != 0",
325
+ "workgroupSize": "workgroupSize",
326
+ "split": "splitCount",
327
+ "tileCols": "tunables.SPLIT_TILE_COLUMNS"
328
+ },
329
+ "intermediates": [
330
+ { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
331
+ { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
332
+ ],
333
+ "passes": [
334
+ {
335
+ "id": "split_reduce",
336
+ "name": "ArgMin.AxisSplitTiledReduce",
337
+ "source": { "shader": "reduce-arg-axis-split-tiled.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
338
+ "bindings": "splitReduce",
339
+ "dispatch": { "x": "ceilDiv(outputCount, tunables.SPLIT_TILE_COLUMNS)", "y": "splitCount" }
340
+ },
341
+ {
342
+ "id": "combine",
343
+ "name": "ArgMin.AxisSplitCombine",
344
+ "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
345
+ "bindings": "splitCombine",
346
+ "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
347
+ }
348
+ ]
349
+ },
350
+ {
351
+ "id": "subgroup_min_last_axis_vec4",
352
+ "priority": 41,
353
+ "requires": { "features": ["subgroups"] },
354
+ "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"],
355
+ "constants": {
356
+ "scalar": "dtypes.T",
357
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
358
+ "usesF16": "dtypes.T == \"f16\"",
359
+ "selectLastIndex": "attrs.select_last_index != 0",
360
+ "workgroupSize": "device.adapterInfo.subgroupMinSize"
361
+ },
362
+ "passes": [
363
+ {
364
+ "id": "main",
365
+ "name": "ArgMin.SubgroupMinLastAxisVec4",
366
+ "source": {
367
+ "shader": "reduce-arg-row-subgroup.wgsl.jinja",
368
+ "inputs": { "mode": "\"min\"", "vec4": true, "singleSubgroup": true }
369
+ },
370
+ "bindings": "contiguousVec4",
371
+ "dispatch": { "workgroups": "outputCount" }
372
+ }
373
+ ]
374
+ },
375
+ {
376
+ "id": "subgroup_min_last_axis",
377
+ "priority": 41,
378
+ "requires": { "features": ["subgroups"] },
379
+ "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"],
380
+ "constants": {
381
+ "scalar": "dtypes.T",
382
+ "usesF16": "dtypes.T == \"f16\"",
383
+ "selectLastIndex": "attrs.select_last_index != 0",
384
+ "workgroupSize": "device.adapterInfo.subgroupMinSize"
385
+ },
386
+ "passes": [
387
+ {
388
+ "id": "main",
389
+ "name": "ArgMin.SubgroupMinLastAxis",
390
+ "source": {
391
+ "shader": "reduce-arg-row-subgroup.wgsl.jinja",
392
+ "inputs": { "mode": "\"min\"", "vec4": false, "singleSubgroup": true }
393
+ },
394
+ "bindings": "contiguousScalar",
395
+ "dispatch": { "workgroups": "outputCount" }
396
+ }
397
+ ]
398
+ },
399
+ {
400
+ "id": "last_axis_vec4",
401
+ "priority": 40,
402
+ "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim >= 4", "axisDim % 4 == 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
403
+ "constants": {
404
+ "scalar": "dtypes.T",
405
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
406
+ "usesF16": "dtypes.T == \"f16\"",
407
+ "selectLastIndex": "attrs.select_last_index != 0",
408
+ "workgroupSize": "min(workgroupSize, max(32, pow2ceil(ceilDiv(axisDim, 4))))"
409
+ },
410
+ "passes": [
411
+ {
412
+ "id": "main",
413
+ "name": "ArgMin.LastAxisVec4",
414
+ "source": {
415
+ "shader": "reduce-arg-row-subgroup.wgsl.jinja",
416
+ "inputs": { "mode": "\"min\"", "vec4": true, "useSubgroups": "device.features.has(\"subgroups\")" }
417
+ },
418
+ "bindings": "contiguousVec4",
419
+ "dispatch": { "workgroups": "outputCount" }
420
+ }
421
+ ]
422
+ },
423
+ {
424
+ "id": "last_axis",
425
+ "priority": 39,
426
+ "when": ["baseContract", "rowParallelType", "lastAxis", "axisDim % 4 != 0", "parallelRowWorthwhile", "rowDispatchFits", "rowStorageFits"],
427
+ "constants": {
428
+ "scalar": "dtypes.T",
429
+ "usesF16": "dtypes.T == \"f16\"",
430
+ "selectLastIndex": "attrs.select_last_index != 0",
431
+ "workgroupSize": "min(workgroupSize, max(32, pow2ceil(axisDim)))"
432
+ },
433
+ "passes": [
434
+ {
435
+ "id": "main",
436
+ "name": "ArgMin.LastAxis",
437
+ "source": {
438
+ "shader": "reduce-arg-row-subgroup.wgsl.jinja",
439
+ "inputs": {
440
+ "mode": "\"min\"",
441
+ "vec4": false,
442
+ "useSubgroups": "device.features.has(\"subgroups\") and not narrowOddRowPortablePreferred"
443
+ }
444
+ },
445
+ "bindings": "contiguousScalar",
446
+ "dispatch": { "workgroups": "outputCount" }
447
+ }
448
+ ]
449
+ },
450
+ {
451
+ "id": "axis_split",
452
+ "priority": 30,
453
+ "demoteWhen": ["narrowLargeAxisTilePreferred"],
454
+ "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.SPLIT_MIN_AXIS", "outputCount >= 1", "outputCount <= tunables.SPLIT_MAX_OUTPUTS", "splitScratchFits", "splitDispatchFits", "threadDispatchFits"],
455
+ "constants": {
456
+ "scalar": "dtypes.T",
457
+ "selectLastIndex": "attrs.select_last_index != 0",
458
+ "workgroupSize": "workgroupSize",
459
+ "split": "splitCount"
460
+ },
461
+ "intermediates": [
462
+ { "id": "partials_val", "dtype": "uint32", "shape": "[splitCount * outputCount]" },
463
+ { "id": "partials_idx", "dtype": "uint32", "shape": "[splitCount * outputCount]" }
464
+ ],
465
+ "passes": [
466
+ {
467
+ "id": "split_reduce",
468
+ "name": "ArgMin.AxisSplitReduce",
469
+ "source": { "shader": "reduce-arg-axis-split-reduce.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
470
+ "bindings": "splitReduce",
471
+ "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize", "y": "splitCount" }
472
+ },
473
+ {
474
+ "id": "combine",
475
+ "name": "ArgMin.AxisSplitCombine",
476
+ "source": { "shader": "reduce-arg-axis-split-combine.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
477
+ "bindings": "splitCombine",
478
+ "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
479
+ }
480
+ ]
481
+ },
482
+ {
483
+ "id": "axis_tiled",
484
+ "priority": 20,
485
+ "when": ["baseContract", "nativeParallelType", "not lastAxis", "axisDim >= tunables.TILED_MIN_AXIS", "outputCount >= tunables.TILED_MIN_OUTPUTS", "workgroupSize % tunables.TILE_COLUMNS == 0", "rowStorageFits", "tileDispatchFits"],
486
+ "constants": {
487
+ "scalar": "dtypes.T",
488
+ "selectLastIndex": "attrs.select_last_index != 0",
489
+ "workgroupSize": "workgroupSize",
490
+ "tileCols": "tunables.TILE_COLUMNS"
491
+ },
492
+ "passes": [
493
+ {
494
+ "id": "main",
495
+ "name": "ArgMin.AxisTiled",
496
+ "source": { "shader": "reduce-arg-axis-tiled.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
497
+ "bindings": "axisGeometry",
498
+ "dispatch": { "workgroups": "ceilDiv((outputCount), (tunables.TILE_COLUMNS))" }
499
+ }
500
+ ]
501
+ },
502
+ {
503
+ "id": "axis_serial",
504
+ "priority": 0,
505
+ "supersededBy": ["axis_tiled"],
506
+ "when": ["baseContract", "threadDispatchFits"],
507
+ "constants": {
508
+ "scalar": "dtypes.T",
509
+ "usesF16": "dtypes.T == \"f16\"",
510
+ "selectLastIndex": "attrs.select_last_index != 0",
511
+ "workgroupSize": "workgroupSize"
512
+ },
513
+ "passes": [
514
+ {
515
+ "id": "main",
516
+ "name": "ArgMin.AxisSerial",
517
+ "source": { "shader": "reduce-arg-axis.wgsl.jinja", "inputs": { "mode": "\"min\"" } },
518
+ "bindings": "axisGeometry",
519
+ "dispatch": { "threads": "outputCount", "workgroupSize": "workgroupSize" }
520
+ }
521
+ ]
522
+ }
523
+ ]
524
+ }
build/webgpu/metadata.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "ai.onnx.ArgMin",
3
+ "id": "_ai_onnx_argmin_webgpu_a0fbb5d",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "backend": { "type": "webgpu" },
7
+ "digest": {
8
+ "algorithm": "sha256",
9
+ "files": {
10
+ "bench.json": "mm/x1LvCBXT2wlKgCEbxB5Oc506zsg2sfh5+8s0SzyI=",
11
+ "manifest.json": "6YtVbWUVmqO89IGSyNQYIvaRetZCne1BxSikonXTKaQ=",
12
+ "reduce-arg-axis-split-combine.wgsl.jinja": "rdvAqkyd5MkxXLWQtVzpTxqqKwANhk2L1h/SMW9P2qU=",
13
+ "reduce-arg-axis-split-reduce.wgsl.jinja": "VC4pBZyx0kVu5zmny2clKWdMU2D6HWpM/qiym+b664o=",
14
+ "reduce-arg-axis-split-tiled.wgsl.jinja": "/CtqlF3QH7OIQi9q0WzSOJC6K/KDUzcH4sDHdrcGTLg=",
15
+ "reduce-arg-axis-tiled.wgsl.jinja": "382jqrusVAtzqbRnn0b++oLtQK53ezOx203+9kyWs34=",
16
+ "reduce-arg-axis.wgsl.jinja": "LzoRz8NbCeL+4FA0A6K/5xmPw6jr1n3rqnq3h8lXXFs=",
17
+ "reduce-arg-row-split.wgsl.jinja": "t2dWKMUzJLdeIBMtllRHEKE+V5YQDSry6dxe/WmwiUg=",
18
+ "reduce-arg-row-subgroup.wgsl.jinja": "IkTnxOCG5uHuSYBWOr8vUW7MfqnJKlBY8ByKbA4+K60=",
19
+ "test.json": "stypKi/CuonXgtb+pFhpPyL6q5ymL6WWV0M9V27ANfc="
20
+ }
21
+ },
22
+ "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
23
+ "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.ArgMin" }
24
+ }
build/webgpu/reduce-arg-axis-split-combine.wgsl.jinja ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 2 of the split-axis arg reduction folds SPLIT (best value, best index)
2
+ // partials and emits the winning reduced-axis index. One thread handles each
3
+ // flattened output. The reduce pass bit-stored each partial value in u32
4
+ // scratch, so this pass restores the comparison type: floats compare in f32,
5
+ // while i32 and u32 compare natively.
6
+ //
7
+ // The fold compares (value, index), never segment order. Equal values keep the
8
+ // smaller index for first-index mode and the larger index for last-index mode.
9
+ // Empty segments carry SENTINEL_IDX and are skipped.
10
+ {% set cmp = ">" if source.mode == "max" else "<" %}
11
+ {% set is_int = scalar == "i32" or scalar == "u32" %}
12
+ {% set work = scalar if is_int else "f32" %}
13
+ {{ env.wgsl.resourceDeclarations }}
14
+
15
+ const WG: u32 = {{ workgroupSize }}u;
16
+ const SPLIT: u32 = {{ split }}u;
17
+ const SENTINEL_IDX: u32 = 4294967295u;
18
+
19
+ @compute @workgroup_size(WG, 1, 1)
20
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
21
+ @builtin(num_workgroups) nwg: vec3<u32>) {
22
+ let output_index = gid.x + gid.y * nwg.x * WG;
23
+ if (output_index >= params.outputCount) {
24
+ return;
25
+ }
26
+ {% if scalar == "u32" %}
27
+ var out_val = partials_val[output_index];
28
+ {% else %}
29
+ var out_val = bitcast<{{ work }}>(partials_val[output_index]);
30
+ {% endif %}
31
+ var out_idx = partials_idx[output_index];
32
+ for (var split_index = 1u; split_index < SPLIT; split_index++) {
33
+ {% if scalar == "u32" %}
34
+ let candidate_val = partials_val[split_index * params.outputCount + output_index];
35
+ {% else %}
36
+ let candidate_val = bitcast<{{ work }}>(partials_val[split_index * params.outputCount + output_index]);
37
+ {% endif %}
38
+ let candidate_idx = partials_idx[split_index * params.outputCount + output_index];
39
+ {% if selectLastIndex %}
40
+ if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
41
+ {% else %}
42
+ if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
43
+ {% endif %}
44
+ out_val = candidate_val;
45
+ out_idx = candidate_idx;
46
+ }
47
+ }
48
+ y[output_index] = select(out_idx, 0u, out_idx == SENTINEL_IDX);
49
+ }
build/webgpu/reduce-arg-axis-split-reduce.wgsl.jinja ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 1 of a rank-independent split-axis ArgMax/ArgMin reduction. The tensor
2
+ // is flattened as (outer, axis, inner). Each thread scans one contiguous segment
3
+ // of the reduced axis for one output position.
4
+ //
5
+ // Partial values are bit-stored in u32 scratch, so f32, i32, and u32 share one
6
+ // layout without losing integer precision. The combine pass restores the native
7
+ // comparison type. Equal values keep the lowest index by default or the highest
8
+ // index when select_last_index is enabled.
9
+ {% set cmp = ">" if source.mode == "max" else "<" %}
10
+ {% set is_int = scalar == "i32" or scalar == "u32" %}
11
+ {% set work = scalar if is_int else "f32" %}
12
+ {{ env.wgsl.resourceDeclarations }}
13
+
14
+ const WG: u32 = {{ workgroupSize }}u;
15
+ const SPLIT: u32 = {{ split }}u;
16
+ const SENTINEL_IDX: u32 = 4294967295u;
17
+
18
+ // The scan identity: the most-losing value for the mode, in the compare type.
19
+ // The float branch bitcasts through a runtime `var` so the +/-inf bit pattern
20
+ // is not folded into an unrepresentable infinity literal.
21
+ fn lane_identity() -> {{ work }} {
22
+ {% if scalar == "i32" %}
23
+ return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
24
+ {% elif scalar == "u32" %}
25
+ return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
26
+ {% else %}
27
+ var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
28
+ return bitcast<f32>(bits);
29
+ {% endif %}
30
+ }
31
+
32
+ @compute @workgroup_size(WG, 1, 1)
33
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
34
+ @builtin(workgroup_id) wg: vec3<u32>,
35
+ @builtin(num_workgroups) nwg: vec3<u32>) {
36
+ // 2D-folded column index: wg.z carries the high bits past the maxComputeWorkgroupsPerDimension
37
+ // workgroup-per-dimension dispatch limit on the x dimension.
38
+ let output_index = (wg.x + wg.z * nwg.x) * WG + (gid.x % WG);
39
+ let split_index = wg.y;
40
+ if (output_index >= params.outputCount) { return; }
41
+ let outer_index = output_index / params.innerSize;
42
+ let inner_index = output_index % params.innerSize;
43
+ let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
44
+
45
+ // Even split of [0, rows) across SPLIT segments (last segment may be shorter).
46
+ let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
47
+ let axis_begin = split_index * chunk;
48
+ var axis_end = min(axis_begin + chunk, params.axisDim);
49
+
50
+ var best_val = lane_identity();
51
+ var best_idx = SENTINEL_IDX;
52
+ for (var axis_index = axis_begin; axis_index < axis_end; axis_index++) {
53
+ let v = x[input_base + axis_index * params.innerSize];
54
+ {% if selectLastIndex %}
55
+ if (v {{ cmp }} best_val || (v == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
56
+ {% else %}
57
+ if (v {{ cmp }} best_val || (v == best_val && axis_index < best_idx)) {
58
+ {% endif %}
59
+ best_val = v;
60
+ best_idx = axis_index;
61
+ }
62
+ }
63
+ {% if scalar == "u32" %}
64
+ partials_val[split_index * params.outputCount + output_index] = best_val;
65
+ {% else %}
66
+ partials_val[split_index * params.outputCount + output_index] = bitcast<u32>(best_val);
67
+ {% endif %}
68
+ partials_idx[split_index * params.outputCount + output_index] = best_idx;
69
+ }
build/webgpu/reduce-arg-axis-split-tiled.wgsl.jinja ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Split-axis arg reduction for a narrow set of output columns. The ordinary
2
+ // split pass assigns one invocation to each output column, leaving most lanes
3
+ // idle when outputCount is small. This route tiles TILE_COLS columns across the
4
+ // workgroup and uses ROW_LANES lanes to cooperatively scan each column. One
5
+ // workgroup handles each (split, column tile), preserving coalesced reads for
6
+ // flattened (outer, axis, inner) tensors.
7
+ //
8
+ // The final fold is intentionally subgroup-free: it is small, deterministic,
9
+ // and preserves identical first/last-index tie semantics on every feature tier.
10
+ {% set cmp = ">" if source.mode == "max" else "<" %}
11
+ {% set is_int = scalar == "i32" or scalar == "u32" %}
12
+ {% set work = scalar if is_int else "f32" %}
13
+ {{ env.wgsl.resourceDeclarations }}
14
+
15
+ const WG: u32 = {{ workgroupSize }}u;
16
+ const SPLIT: u32 = {{ split }}u;
17
+ const TILE_COLS: u32 = {{ tileCols }}u;
18
+ const ROW_LANES: u32 = WG / TILE_COLS;
19
+ const SENTINEL_IDX: u32 = 4294967295u;
20
+
21
+ fn lane_identity() -> {{ work }} {
22
+ {% if scalar == "i32" %}
23
+ return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
24
+ {% elif scalar == "u32" %}
25
+ return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
26
+ {% else %}
27
+ var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
28
+ return bitcast<f32>(bits);
29
+ {% endif %}
30
+ }
31
+
32
+ var<workgroup> partial_val: array<{{ work }}, WG>;
33
+ var<workgroup> partial_idx: array<u32, WG>;
34
+
35
+ @compute @workgroup_size(WG, 1, 1)
36
+ fn main(
37
+ @builtin(workgroup_id) wg: vec3<u32>,
38
+ @builtin(local_invocation_id) lid: vec3<u32>
39
+ ) {
40
+ let tid = lid.x;
41
+ let col_lane = tid % TILE_COLS;
42
+ let row_lane = tid / TILE_COLS;
43
+ let output_index = wg.x * TILE_COLS + col_lane;
44
+ let split_index = wg.y;
45
+ let in_bounds = output_index < params.outputCount;
46
+
47
+ let outer_index = output_index / params.innerSize;
48
+ let inner_index = output_index % params.innerSize;
49
+ let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
50
+ let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
51
+ let axis_begin = split_index * chunk;
52
+ let axis_end = min(axis_begin + chunk, params.axisDim);
53
+
54
+ var best_val = lane_identity();
55
+ var best_idx = SENTINEL_IDX;
56
+ if (in_bounds) {
57
+ for (var axis_index = axis_begin + row_lane; axis_index < axis_end; axis_index += ROW_LANES) {
58
+ let value = x[input_base + axis_index * params.innerSize];
59
+ {% if selectLastIndex %}
60
+ if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
61
+ {% else %}
62
+ if (value {{ cmp }} best_val || (value == best_val && axis_index < best_idx)) {
63
+ {% endif %}
64
+ best_val = value;
65
+ best_idx = axis_index;
66
+ }
67
+ }
68
+ }
69
+ partial_val[tid] = best_val;
70
+ partial_idx[tid] = best_idx;
71
+ workgroupBarrier();
72
+
73
+ if (row_lane == 0u && in_bounds) {
74
+ var out_val = partial_val[col_lane];
75
+ var out_idx = partial_idx[col_lane];
76
+ for (var lane = 1u; lane < ROW_LANES; lane++) {
77
+ let candidate_val = partial_val[lane * TILE_COLS + col_lane];
78
+ let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
79
+ {% if selectLastIndex %}
80
+ if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
81
+ {% else %}
82
+ if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
83
+ {% endif %}
84
+ out_val = candidate_val;
85
+ out_idx = candidate_idx;
86
+ }
87
+ }
88
+ {% if scalar == "u32" %}
89
+ partials_val[split_index * params.outputCount + output_index] = out_val;
90
+ {% else %}
91
+ partials_val[split_index * params.outputCount + output_index] = bitcast<u32>(out_val);
92
+ {% endif %}
93
+ partials_idx[split_index * params.outputCount + output_index] = out_idx;
94
+ }
95
+ }
build/webgpu/reduce-arg-axis-tiled.wgsl.jinja ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Tiled strided-axis ArgMax/ArgMin reduction. f32 compares directly; i32 and
2
+ // u32 compare natively, so large-magnitude integers keep full precision. The
3
+ // workgroup partials carry the native comparison type without a cross-pass
4
+ // bitcast round trip. Flattening the output as (outer, inner) lets one coalesced
5
+ // column tile handle every non-contiguous axis. f16 uses the serial or
6
+ // contiguous routes because this kernel has no widened-f16 scratch path.
7
+ {% set cmp = ">" if source.mode == "max" else "<" %}
8
+ {% set is_int = scalar == "i32" or scalar == "u32" %}
9
+ {% set work = scalar if is_int else "f32" %}
10
+ {{ env.wgsl.resourceDeclarations }}
11
+
12
+ const WG: u32 = {{ workgroupSize }}u;
13
+ const TILE_COLS: u32 = {{ tileCols }}u;
14
+ const ROW_LANES: u32 = WG / TILE_COLS;
15
+ const SENTINEL_IDX: u32 = 4294967295u;
16
+
17
+ // The scan identity: the most-losing value for the mode, in the compare type.
18
+ // The float branch bitcasts through a runtime `var` so the +/-inf bit pattern
19
+ // is not folded into an unrepresentable infinity literal.
20
+ fn lane_identity() -> {{ work }} {
21
+ {% if scalar == "i32" %}
22
+ return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
23
+ {% elif scalar == "u32" %}
24
+ return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
25
+ {% else %}
26
+ var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
27
+ return bitcast<f32>(bits);
28
+ {% endif %}
29
+ }
30
+
31
+ var<workgroup> partial_val: array<{{ work }}, WG>;
32
+ var<workgroup> partial_idx: array<u32, WG>;
33
+
34
+ @compute @workgroup_size(WG, 1, 1)
35
+ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
36
+ @builtin(num_workgroups) nwg: vec3<u32>) {
37
+ let tid = lid.x;
38
+ let col_lane = tid % TILE_COLS;
39
+ let row_lane = tid / TILE_COLS;
40
+ // 2D-folded column-tile index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
41
+ // workgroup-per-dimension dispatch limit. Reduces to wg.x when nwg.y == 1;
42
+ // the bounds guard drops the over-dispatched tail.
43
+ let output_index = (wg.x + wg.y * nwg.x) * TILE_COLS + col_lane;
44
+ let in_bounds = output_index < params.outputCount;
45
+ let outer_index = output_index / params.innerSize;
46
+ let inner_index = output_index % params.innerSize;
47
+ let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
48
+
49
+ var best_val = lane_identity();
50
+ var best_idx = SENTINEL_IDX;
51
+ if (in_bounds) {
52
+ for (var axis_index = row_lane; axis_index < params.axisDim; axis_index = axis_index + ROW_LANES) {
53
+ let v = x[input_base + axis_index * params.innerSize];
54
+ {% if selectLastIndex %}
55
+ if (v {{ cmp }} best_val || (v == best_val && (best_idx == SENTINEL_IDX || axis_index > best_idx))) {
56
+ {% else %}
57
+ if (v {{ cmp }} best_val || (v == best_val && axis_index < best_idx)) {
58
+ {% endif %}
59
+ best_val = v;
60
+ best_idx = axis_index;
61
+ }
62
+ }
63
+ }
64
+ partial_val[tid] = best_val;
65
+ partial_idx[tid] = best_idx;
66
+ workgroupBarrier();
67
+
68
+ if (row_lane == 0u && in_bounds) {
69
+ var out_val = partial_val[col_lane];
70
+ var out_idx = partial_idx[col_lane];
71
+ for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
72
+ let candidate_val = partial_val[lane * TILE_COLS + col_lane];
73
+ let candidate_idx = partial_idx[lane * TILE_COLS + col_lane];
74
+ {% if selectLastIndex %}
75
+ if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx > out_idx))) {
76
+ {% else %}
77
+ if (candidate_idx != SENTINEL_IDX && (out_idx == SENTINEL_IDX || candidate_val {{ cmp }} out_val || (candidate_val == out_val && candidate_idx < out_idx))) {
78
+ {% endif %}
79
+ out_val = candidate_val;
80
+ out_idx = candidate_idx;
81
+ }
82
+ }
83
+ y[output_index] = select(out_idx, 0u, out_idx == SENTINEL_IDX);
84
+ }
85
+ }
build/webgpu/reduce-arg-axis.wgsl.jinja ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Rank-independent serial arg reduction for ArgMax and ArgMin. Flattening a
2
+ // tensor around the reduced axis into (outer, axis, inner) gives one
3
+ // rank-independent coordinate calculation. One thread owns one output position
4
+ // and walks the reduced axis at innerSize stride.
5
+ {% set cmp = ">" if source.mode == "max" else "<" %}
6
+ {% if usesF16 %}
7
+ enable f16;
8
+ {%- endif %}
9
+ {{ env.wgsl.resourceDeclarations }}
10
+
11
+ @compute @workgroup_size({{ workgroupSize }})
12
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
13
+ @builtin(num_workgroups) nwg: vec3<u32>) {
14
+ // Fold the output grid into x/y when one dispatch dimension is insufficient.
15
+ let output_index = gid.x + gid.y * nwg.x * {{ workgroupSize }}u;
16
+ if (output_index >= params.outputCount) {
17
+ return;
18
+ }
19
+
20
+ let outer_index = output_index / params.innerSize;
21
+ let inner_index = output_index % params.innerSize;
22
+ let input_base = outer_index * params.axisDim * params.innerSize + inner_index;
23
+
24
+ var best_index = 0u;
25
+ {%- if scalar == "i32" or scalar == "u32" %}
26
+ var best = x[input_base];
27
+ for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
28
+ let candidate = x[input_base + axis_index * params.innerSize];
29
+ if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
30
+ best = candidate;
31
+ best_index = axis_index;
32
+ }
33
+ }
34
+ {%- else %}
35
+ var best = f32(x[input_base]);
36
+ for (var axis_index = 1u; axis_index < params.axisDim; axis_index++) {
37
+ let candidate = f32(x[input_base + axis_index * params.innerSize]);
38
+ if (candidate {{ cmp }} best{% if selectLastIndex %} || candidate == best{% endif %}) {
39
+ best = candidate;
40
+ best_index = axis_index;
41
+ }
42
+ }
43
+ {%- endif %}
44
+ y[output_index] = best_index;
45
+ }
build/webgpu/reduce-arg-row-split.wgsl.jinja ADDED
@@ -0,0 +1,159 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {% set useSubgroups = source.useSubgroups %}
2
+ {% if usesF16 %}
3
+ enable f16;
4
+ {%- endif %}
5
+ {% if useSubgroups %}
6
+ enable subgroups;
7
+ {%- endif %}
8
+ {{ env.wgsl.resourceDeclarations }}
9
+
10
+ const WG: u32 = {{ workgroupSize }}u;
11
+ const SPLIT: u32 = {{ split }}u;
12
+ const SENTINEL_IDX: u32 = 4294967295u;
13
+ {%- if scalar == "i32" %}
14
+ {% set work = "i32" %}
15
+ {% elif scalar == "u32" %}
16
+ {% set work = "u32" %}
17
+ {% else %}
18
+ {% set work = "f32" %}
19
+ {% endif %}
20
+ {% set cmp = ">" if source.mode == "max" else "<" %}
21
+
22
+ fn lane_identity() -> {{ work }} {
23
+ {%- if scalar == "i32" %}
24
+ return bitcast<i32>({{ "0x80000000u" if source.mode == "max" else "0x7fffffffu" }});
25
+ {%- elif scalar == "u32" %}
26
+ return {{ "0u" if source.mode == "max" else "0xffffffffu" }};
27
+ {%- else %}
28
+ var bits = {{ "0xff800000u" if source.mode == "max" else "0x7f800000u" }};
29
+ return bitcast<f32>(bits);
30
+ {%- endif %}
31
+ }
32
+
33
+ var<workgroup> wg_val: array<{{ work }}, WG>;
34
+ var<workgroup> wg_idx: array<u32, WG>;
35
+
36
+ @compute @workgroup_size(WG, 1, 1)
37
+ fn main(@builtin(workgroup_id) wg: vec3<u32>,
38
+ @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
39
+ @builtin(subgroup_invocation_id) subgroup_lane: u32,
40
+ @builtin(subgroup_size) subgroup_size: u32{% endif %}) {
41
+ let row = wg.x;
42
+ let split_index = wg.y;
43
+ let tid = lid.x;
44
+ let segment_size = (params.chunkCount + SPLIT - 1u) / SPLIT;
45
+ let segment_begin = split_index * segment_size;
46
+ let segment_end = min(segment_begin + segment_size, params.chunkCount);
47
+ let base = row * params.chunkCount;
48
+
49
+ var best_val = lane_identity();
50
+ var best_idx = SENTINEL_IDX;
51
+ for (var chunk = segment_begin + tid; chunk < segment_end; chunk += WG) {
52
+ {%- if source.vec4 %}
53
+ let value4 = x[base + chunk];
54
+ {%- for component in ["x", "y", "z", "w"] %}
55
+ {
56
+ {%- if scalar == "f16" %}
57
+ let value = f32(value4.{{ component }});
58
+ {%- else %}
59
+ let value = value4.{{ component }};
60
+ {%- endif %}
61
+ let index = chunk * 4u + {{ loop.index0 }}u;
62
+ {%- if selectLastIndex %}
63
+ if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || index > best_idx))) {
64
+ {%- else %}
65
+ if (value {{ cmp }} best_val || (value == best_val && index < best_idx)) {
66
+ {%- endif %}
67
+ best_val = value;
68
+ best_idx = index;
69
+ }
70
+ }
71
+ {%- endfor %}
72
+ {% else %}
73
+ {% if scalar == "f16" %}
74
+ let value = f32(x[base + chunk]);
75
+ {%- else %}
76
+ let value = x[base + chunk];
77
+ {%- endif %}
78
+ {% if selectLastIndex %}
79
+ if (value {{ cmp }} best_val || (value == best_val && (best_idx == SENTINEL_IDX || chunk > best_idx))) {
80
+ {%- else %}
81
+ if (value {{ cmp }} best_val || (value == best_val && chunk < best_idx)) {
82
+ {%- endif %}
83
+ best_val = value;
84
+ best_idx = chunk;
85
+ }
86
+ {%- endif %}
87
+ }
88
+ {%- if useSubgroups %}
89
+ let subgroup_val = subgroupMin(best_val);{% if selectLastIndex %}
90
+ let subgroup_candidate = select(0u, best_idx, best_val == subgroup_val && best_idx != SENTINEL_IDX);
91
+ let subgroup_idx = subgroupMax(subgroup_candidate);
92
+ {%- else %}
93
+ let subgroup_candidate = select(SENTINEL_IDX, best_idx, best_val == subgroup_val);
94
+ let subgroup_idx = subgroupMin(subgroup_candidate);
95
+ {%- endif %}
96
+ let safe_subgroup_size = max(subgroup_size, 1u);
97
+ // The final subgroup may be partial on devices whose native subgroup size
98
+ // does not divide WG. Merge that leader as well.
99
+ let subgroup_count = max(1u, (WG + safe_subgroup_size - 1u) / safe_subgroup_size);
100
+ if (subgroup_lane == 0u) {
101
+ let slot = min(tid / safe_subgroup_size, WG - 1u);
102
+ wg_val[slot] = subgroup_val;
103
+ wg_idx[slot] = subgroup_idx;
104
+ }
105
+ workgroupBarrier();
106
+ if (tid == 0u) {
107
+ var output_val = wg_val[0];
108
+ var output_idx = wg_idx[0];
109
+ for (var slot = 1u; slot < subgroup_count; slot++) {
110
+ let candidate_val = wg_val[slot];
111
+ let candidate_idx = wg_idx[slot];
112
+ {%- if selectLastIndex %}
113
+ if (candidate_idx != SENTINEL_IDX && (output_idx == SENTINEL_IDX || candidate_val {{ cmp }} output_val || (candidate_val == output_val && candidate_idx > output_idx))) {
114
+ {%- else %}
115
+ if (candidate_idx != SENTINEL_IDX && (output_idx == SENTINEL_IDX || candidate_val {{ cmp }} output_val || (candidate_val == output_val && candidate_idx < output_idx))) {
116
+ {%- endif %}
117
+ output_val = candidate_val;
118
+ output_idx = candidate_idx;
119
+ }
120
+ }
121
+ let scratch_index = split_index * params.rows + row;
122
+ {%- if scalar == "u32" %}
123
+ partials_val[scratch_index] = output_val;
124
+ {%- else %}
125
+ partials_val[scratch_index] = bitcast<u32>(output_val);
126
+ {%- endif %}
127
+ partials_idx[scratch_index] = output_idx;
128
+ }
129
+ {%- else %}
130
+ wg_val[tid] = best_val;
131
+ wg_idx[tid] = best_idx;
132
+ workgroupBarrier();
133
+ for (var stride = WG / 2u; stride > 0u; stride >>= 1u) {
134
+ if (tid < stride) {
135
+ let candidate_val = wg_val[tid + stride];
136
+ let candidate_idx = wg_idx[tid + stride];
137
+ let current_idx = wg_idx[tid];
138
+ {%- if selectLastIndex %}
139
+ if (candidate_idx != SENTINEL_IDX && (current_idx == SENTINEL_IDX || candidate_val {{ cmp }} wg_val[tid] || (candidate_val == wg_val[tid] && candidate_idx > current_idx))) {
140
+ {%- else %}
141
+ if (candidate_idx != SENTINEL_IDX && (current_idx == SENTINEL_IDX || candidate_val {{ cmp }} wg_val[tid] || (candidate_val == wg_val[tid] && candidate_idx < current_idx))) {
142
+ {%- endif %}
143
+ wg_val[tid] = candidate_val;
144
+ wg_idx[tid] = candidate_idx;
145
+ }
146
+ }
147
+ workgroupBarrier();
148
+ }
149
+ if (tid == 0u) {
150
+ let scratch_index = split_index * params.rows + row;
151
+ {%- if scalar == "u32" %}
152
+ partials_val[scratch_index] = wg_val[0];
153
+ {%- else %}
154
+ partials_val[scratch_index] = bitcast<u32>(wg_val[0]);
155
+ {%- endif %}
156
+ partials_idx[scratch_index] = wg_idx[0];
157
+ }
158
+ {%- endif %}
159
+ }
build/webgpu/reduce-arg-row-subgroup.wgsl.jinja ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Subgroup (value, index) row reduction for ArgMax and ArgMin over a contiguous
2
+ // last axis, with one workgroup per output row. Each thread scans a strided
3
+ // chunk, carrying a (best value, best index) pair with ONNX tie semantics.
4
+ // Subgroup combination first reduces the value, then reduces the matching
5
+ // lanes' indices. Leaders deposit per-subgroup pairs in shared memory for a
6
+ // final fold. A one-subgroup workgroup writes its collective result directly,
7
+ // avoiding shared memory and the workgroup barrier.
8
+ //
9
+ // Float inputs compare in f32. Sentinels are runtime-bitcast infinities or the
10
+ // integer bounds; ties against the sentinel use the same index rule, so a row
11
+ // filled with the sentinel still produces the correct index. NaNs are never
12
+ // accepted, and an all-NaN row falls back to index zero like the serial scan.
13
+ {% if source.useSubgroups is defined %}{% set useSubgroups = source.useSubgroups %}{% else %}{% set useSubgroups = true %}{% endif %}
14
+ {% if usesF16 %}
15
+ enable f16;
16
+ {%- endif %}
17
+ {% if useSubgroups %}
18
+ enable subgroups;
19
+ {%- endif %}
20
+ {{ env.wgsl.resourceDeclarations }}
21
+
22
+ const WG: u32 = {{ workgroupSize }}u;
23
+ const SENTINEL_IDX: u32 = 4294967295u;
24
+ {% if source.mode == "min" and scalar != "i32" and scalar != "u32" %}
25
+ fn pos_inf_f32() -> f32 {
26
+ var bits = 0x7f800000u;
27
+ return bitcast<f32>(bits);
28
+ }
29
+ {% endif %}
30
+ {%- if scalar == "i32" %}
31
+ {% set cmpType = "i32" %}
32
+ {% elif scalar == "u32" %}
33
+ {% set cmpType = "u32" %}
34
+ {% else %}
35
+ {% set cmpType = "f32" %}
36
+ {% endif %}
37
+ {% set cmp = "<" %}
38
+ {%- if not (useSubgroups and source.singleSubgroup is defined and source.singleSubgroup) %}
39
+ var<workgroup> wgVal: array<{{ cmpType }}, WG>;
40
+ var<workgroup> wgIdx: array<u32, WG>;
41
+ {%- endif %}
42
+
43
+ @compute @workgroup_size(WG, 1, 1)
44
+ fn main(@builtin(workgroup_id) wg: vec3<u32>,
45
+ @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
46
+ @builtin(subgroup_invocation_id) sgLid: u32{% if not (source.singleSubgroup is defined and source.singleSubgroup) %},
47
+ @builtin(subgroup_size) sgSize: u32{% endif %}{% endif %}) {
48
+ let row = wg.x + wg.y * params.rowStride;
49
+ if (row >= params.rows) {
50
+ return;
51
+ }
52
+ let tid = lid.x;
53
+ let base = row * params.chunkCount;{% if cmpType == "i32" %}
54
+ var bestVal: i32 = 2147483647i;
55
+ {%- elif cmpType == "u32" %}
56
+ var bestVal: u32 = 4294967295u;
57
+ {%- else %}
58
+ var bestVal: f32 = pos_inf_f32();
59
+ {%- endif %}
60
+ var bestIdx: u32 = SENTINEL_IDX;
61
+ {%- if source.vec4 %}
62
+ for (var c = tid; c < params.chunkCount; c = c + WG) {
63
+ let v4 = x[base + c];
64
+ {%- for comp in ["x", "y", "z", "w"] %}
65
+ {
66
+ {%- if scalar == "f16" %}
67
+ let v = f32(v4.{{ comp }});
68
+ {%- else %}
69
+ let v = v4.{{ comp }};
70
+ {%- endif %}
71
+ let idx = c * 4u + {{ loop.index0 }}u;
72
+ {%- if selectLastIndex %}
73
+ if (v {{ cmp }} bestVal || (v == bestVal && (bestIdx == SENTINEL_IDX || idx > bestIdx))) {
74
+ {%- else %}
75
+ if (v {{ cmp }} bestVal || (v == bestVal && idx < bestIdx)) {
76
+ {%- endif %}
77
+ bestVal = v;
78
+ bestIdx = idx;
79
+ }
80
+ }
81
+ {%- endfor %}
82
+ }
83
+ {%- else %}
84
+ // Four independent, lane-coalesced loads per loop iteration. Keeping each
85
+ // unrolled plane WG elements apart preserves contiguous accesses across the
86
+ // workgroup while exposing enough ILP for long odd-width rows (where a vec4
87
+ // storage binding cannot represent row boundaries).
88
+ for (var c0 = tid; c0 < params.chunkCount; c0 = c0 + 4u * WG) {
89
+ {%- for plane in [0, 1, 2, 3] %}
90
+ {
91
+ let c = c0 + {{ plane }}u * WG;
92
+ if (c < params.chunkCount) {
93
+ {%- if scalar == "f16" %}
94
+ let v = f32(x[base + c]);
95
+ {%- else %}
96
+ let v = x[base + c];
97
+ {%- endif %}
98
+ {% if selectLastIndex %}
99
+ if (v {{ cmp }} bestVal || (v == bestVal && (bestIdx == SENTINEL_IDX || c > bestIdx))) {
100
+ {%- else %}
101
+ if (v {{ cmp }} bestVal || (v == bestVal && c < bestIdx)) {
102
+ {%- endif %}
103
+ bestVal = v;
104
+ bestIdx = c;
105
+ }
106
+ }
107
+ }
108
+ {%- endfor %}
109
+ }
110
+ {%- endif %}
111
+ {%- if useSubgroups %}
112
+ let m = subgroupMin(bestVal);{% if selectLastIndex %}
113
+ // Lanes that match the subgroup best contribute their (real) index; the
114
+ // rest contribute 0 so subgroupMax picks the highest matching index.
115
+ let cand = select(0u, bestIdx, bestVal == m && bestIdx != SENTINEL_IDX);
116
+ let sgIdx = subgroupMax(cand);
117
+ {%- else %}
118
+ // Lanes that match the subgroup best contribute their index; the rest
119
+ // contribute the sentinel so subgroupMin picks the lowest matching index.
120
+ let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
121
+ let sgIdx = subgroupMin(cand);
122
+ {%- endif %}
123
+ {% if source.singleSubgroup is defined and source.singleSubgroup %}
124
+ if (sgLid == 0u) {
125
+ y[row] = select(sgIdx, 0u, sgIdx == SENTINEL_IDX);
126
+ }
127
+ {%- else %}
128
+ let safeSg = max(sgSize, 1u);
129
+ let slotCount = max(1u, (WG + safeSg - 1u) / safeSg);
130
+ if (sgLid == 0u) {
131
+ let slot = min(tid / safeSg, WG - 1u);
132
+ wgVal[slot] = m;
133
+ wgIdx[slot] = sgIdx;
134
+ }
135
+ workgroupBarrier();
136
+ if (tid == 0u) {
137
+ var outVal = wgVal[0];
138
+ var outIdx = wgIdx[0];
139
+ for (var i = 1u; i < slotCount; i = i + 1u) {
140
+ let v = wgVal[i];
141
+ let vi = wgIdx[i];
142
+ {%- if selectLastIndex %}
143
+ if (vi != SENTINEL_IDX && (v {{ cmp }} outVal || (v == outVal && (outIdx == SENTINEL_IDX || vi > outIdx)))) {
144
+ {%- else %}
145
+ if (vi != SENTINEL_IDX && (v {{ cmp }} outVal || (v == outVal && vi < outIdx))) {
146
+ {%- endif %}
147
+ outVal = v;
148
+ outIdx = vi;
149
+ }
150
+ }
151
+ y[row] = select(outIdx, 0u, outIdx == SENTINEL_IDX);
152
+ }
153
+ {%- endif %}
154
+ {% else %}
155
+ // No-subgroup tier: workgroup barrier tree-reduction over the per-thread
156
+ // (bestVal, bestIdx) pairs. Same ONNX tie rule (lowest index by default,
157
+ // highest with select_last_index); SENTINEL_IDX lanes (short rows, WG >
158
+ // chunkCount) never win. WG is a power of two,
159
+ // so the halving tree visits every lane.
160
+ wgVal[tid] = bestVal;
161
+ wgIdx[tid] = bestIdx;
162
+ workgroupBarrier();
163
+ for (var stride = WG / 2u; stride > 0u; stride = stride >> 1u) {
164
+ if (tid < stride) {
165
+ let v = wgVal[tid + stride];
166
+ let vi = wgIdx[tid + stride];
167
+ let curIdx = wgIdx[tid];
168
+ {%- if selectLastIndex %}
169
+ if (vi != SENTINEL_IDX && (curIdx == SENTINEL_IDX || v {{ cmp }} wgVal[tid] || (v == wgVal[tid] && vi > curIdx))) {
170
+ {%- else %}
171
+ if (vi != SENTINEL_IDX && (curIdx == SENTINEL_IDX || v {{ cmp }} wgVal[tid] || (v == wgVal[tid] && vi < curIdx))) {
172
+ {%- endif %}
173
+ wgVal[tid] = v;
174
+ wgIdx[tid] = vi;
175
+ }
176
+ }
177
+ workgroupBarrier();
178
+ }
179
+ if (tid == 0u) {
180
+ y[row] = select(wgIdx[0], 0u, wgIdx[0] == SENTINEL_IDX);
181
+ }
182
+ {%- endif %}
183
+ }
build/webgpu/test.json ADDED
@@ -0,0 +1,1214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "op": "ai.onnx.ArgMin",
3
+ "fixtureArrays": {
4
+ "onnx_backend_argmin_input_x": [0.9762700796127319, 4.3037872314453125, 2.055267572402954, 0.8976636528968811, -1.5269039869308472, 2.917882204055786, -1.248255729675293, 7.835460186004639, 9.273255348205566, -2.331169605255127, 5.834500789642334, 0.577898383140564, 1.3608912229537964, 8.511932373046875, -8.579278945922852, -8.257413864135742, -9.59563159942627, 6.6523966789245605, 5.563135147094727, 7.400242805480957, 9.572366714477539, 5.983171463012695, -0.7704127430915833, 5.610583305358887],
5
+ "axis0_splitk_8192x32_ties_first_index_input_x": [2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2, -3, -3, 1, 2],
6
+ "subgroup_vec4_ties_first_index_input_x": [1, 5, 5, 2, 5, 0, 1, 0, 7, 7, 7, 7, 7, 7, 7, 7],
7
+ "axis0_splitk_i32_8192x16_precision_min_input_x": [16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777217, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216, 16777216]
8
+ },
9
+ "cases": [
10
+ {
11
+ "name": "int16_last_axis_vec4_boundaries",
12
+ "attrs": { "axis": -1, "keepdims": 0 },
13
+ "inputs": {
14
+ "x": {
15
+ "dtype": "int16",
16
+ "shape": [2, 4],
17
+ "data": { "kind": "values", "values": [-32768, -1, 32767, 0, 32767, -32768, 32767, 1] }
18
+ }
19
+ },
20
+ "outputs": {
21
+ "y": { "dtype": "uint32", "shape": [2], "tolerance": 0, "data": { "kind": "values", "values": [0, 1] } }
22
+ }
23
+ },
24
+ {
25
+ "name": "onnx_default_axis0_omitted",
26
+ "provenance": {
27
+ "source": "https://onnx.ai/onnx/operators/onnx__ArgMin.html",
28
+ "test": "ArgMin axis default",
29
+ "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."
30
+ },
31
+ "inputs": {
32
+ "x": {
33
+ "dtype": "float32",
34
+ "shape": [2, 3],
35
+ "data": { "kind": "values", "values": [1.0, 9.0, 3.0, 4.0, 2.0, 8.0] }
36
+ }
37
+ },
38
+ "outputs": {
39
+ "y": { "dtype": "uint32", "shape": [1, 3], "tolerance": 0, "data": { "kind": "values", "values": [0, 1, 0] } }
40
+ }
41
+ },
42
+ {
43
+ "name": "f32_negative_subnormal_beats_zero_gpu_gap",
44
+ "skipGpu": {
45
+ "category": "permanent",
46
+ "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."
47
+ },
48
+ "provenance": {
49
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
50
+ "test": "ArgMin",
51
+ "notes": "A negative subnormal is strictly less than zero; index selection must not treat it as a tie with zero."
52
+ },
53
+ "attrs": { "axis": 1, "keepdims": 0 },
54
+ "inputs": {
55
+ "x": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [1e-40, 0.0, -1e-40] } }
56
+ },
57
+ "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
58
+ },
59
+ {
60
+ "name": "f32_negative_subnormal_beats_zero_select_last_gpu_gap",
61
+ "skipGpu": {
62
+ "category": "permanent",
63
+ "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."
64
+ },
65
+ "provenance": {
66
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
67
+ "test": "ArgMin_do_not_keepdims_2_select_last",
68
+ "notes": "select_last_index must only apply to true ties; a negative subnormal at index 0 is less than following zeros."
69
+ },
70
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
71
+ "inputs": {
72
+ "x": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [-1e-40, 0.0, 0.0] } }
73
+ },
74
+ "outputs": {
75
+ "y": { "dtype": "uint32", "shape": [1], "tolerance": 0, "data": { "kind": "values", "values": [0] } }
76
+ }
77
+ },
78
+ {
79
+ "name": "f32_negative_subnormal_beats_zero_axis0_gpu_gap",
80
+ "skipGpu": {
81
+ "category": "permanent",
82
+ "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."
83
+ },
84
+ "provenance": {
85
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
86
+ "test": "ArgMin",
87
+ "notes": "Axis-0 companion for finite subnormal ordering: a negative subnormal is strictly less than zero and must not be tie-broken as flushed zero."
88
+ },
89
+ "attrs": { "axis": 0, "keepdims": 0 },
90
+ "inputs": {
91
+ "x": { "dtype": "float32", "shape": [3, 1], "data": { "kind": "values", "values": [1e-40, 0.0, -1e-40] } }
92
+ },
93
+ "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
94
+ },
95
+ {
96
+ "name": "dispatch_cliff_axis1_16777216x1",
97
+ "attrs": { "axis": 1, "keepdims": 0 },
98
+ "inputs": {
99
+ "x": { "dtype": "float32", "shape": [16777216, 1], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } }
100
+ },
101
+ "outputs": { "y": { "dtype": "uint32", "shape": [16777216], "tolerance": 0 } }
102
+ },
103
+ {
104
+ "name": "axis0",
105
+ "attrs": { "axis": 0, "keepdims": 0 },
106
+ "inputs": {
107
+ "x": {
108
+ "dtype": "float32",
109
+ "shape": [3, 4],
110
+ "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] }
111
+ }
112
+ },
113
+ "outputs": { "y": { "dtype": "uint32", "shape": [4] } }
114
+ },
115
+ {
116
+ "name": "axis0_splitk_8192x32_ties_first_index",
117
+ "attrs": { "axis": 0, "keepdims": 0 },
118
+ "inputs": {
119
+ "x": {
120
+ "dtype": "float32",
121
+ "shape": [8192, 32],
122
+ "data": {
123
+ "kind": "cycle",
124
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
125
+ }
126
+ }
127
+ },
128
+ "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
129
+ },
130
+ {
131
+ "name": "axis0_splitk_8192x32_ties_select_last_index",
132
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
133
+ "inputs": {
134
+ "x": {
135
+ "dtype": "float32",
136
+ "shape": [8192, 32],
137
+ "data": {
138
+ "kind": "cycle",
139
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
140
+ }
141
+ }
142
+ },
143
+ "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
144
+ },
145
+ {
146
+ "name": "axis0_splitk_8192x48_keepdims",
147
+ "attrs": { "axis": 0, "keepdims": 1 },
148
+ "inputs": {
149
+ "x": {
150
+ "dtype": "float32",
151
+ "shape": [8192, 48],
152
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.2 }
153
+ }
154
+ },
155
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 48], "tolerance": 0 } }
156
+ },
157
+ {
158
+ "name": "axis0_tiled_64x32_ties_first_index",
159
+ "attrs": { "axis": 0, "keepdims": 0 },
160
+ "inputs": {
161
+ "x": {
162
+ "dtype": "float32",
163
+ "shape": [64, 32],
164
+ "data": {
165
+ "kind": "cycle",
166
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
167
+ }
168
+ }
169
+ },
170
+ "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
171
+ },
172
+ {
173
+ "name": "axis0_tiled_64x32_ties_select_last_index",
174
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
175
+ "inputs": {
176
+ "x": {
177
+ "dtype": "float32",
178
+ "shape": [64, 32],
179
+ "data": {
180
+ "kind": "cycle",
181
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_8192x32_ties_first_index_input_x" }
182
+ }
183
+ }
184
+ },
185
+ "outputs": { "y": { "dtype": "uint32", "shape": [32], "tolerance": 0 } }
186
+ },
187
+ {
188
+ "name": "axis1_tie_first_index",
189
+ "attrs": { "axis": 1, "keepdims": 0 },
190
+ "inputs": {
191
+ "x": {
192
+ "dtype": "float32",
193
+ "shape": [2, 5],
194
+ "data": { "kind": "values", "values": [1.0, -3.0, -3.0, 2.0, -3.0, 1.0, 1.0, 2.0, 1.0, 3.0] }
195
+ }
196
+ },
197
+ "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
198
+ },
199
+ {
200
+ "name": "ort_axis1_nan_first_incumbent_gpu_gap",
201
+ "skipGpu": {
202
+ "category": "todo",
203
+ "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."
204
+ },
205
+ "provenance": {
206
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
207
+ "test": "ReductionOpTest.ArgMin",
208
+ "notes": "NaN extension: ORT seeds ArgMin from the first reduced element, so a leading NaN remains the selected incumbent while later NaNs are ignored after a finite incumbent."
209
+ },
210
+ "attrs": { "axis": 1, "keepdims": 0 },
211
+ "inputs": {
212
+ "x": {
213
+ "dtype": "float32",
214
+ "shape": [4, 4],
215
+ "data": {
216
+ "kind": "values",
217
+ "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"]
218
+ }
219
+ }
220
+ },
221
+ "outputs": { "y": { "dtype": "uint32", "shape": [4], "tolerance": 0 } }
222
+ },
223
+ {
224
+ "name": "ort_axis1_nan_select_last_index_gpu_gap",
225
+ "skipGpu": {
226
+ "category": "todo",
227
+ "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."
228
+ },
229
+ "provenance": {
230
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
231
+ "test": "ReductionOpTest.ArgMin_int32_last_index_dups",
232
+ "notes": "NaN extension: select_last_index only changes equal finite ties. A leading NaN remains the incumbent because finite values are not less than NaN."
233
+ },
234
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
235
+ "inputs": {
236
+ "x": {
237
+ "dtype": "float32",
238
+ "shape": [4, 4],
239
+ "data": {
240
+ "kind": "values",
241
+ "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"]
242
+ }
243
+ }
244
+ },
245
+ "outputs": {
246
+ "y": { "dtype": "uint32", "shape": [4], "data": { "kind": "values", "values": [0, 3, 0, 0] }, "tolerance": 0 }
247
+ }
248
+ },
249
+ {
250
+ "name": "axis1_f16",
251
+ "attrs": { "axis": 1, "keepdims": 0 },
252
+ "inputs": {
253
+ "x": {
254
+ "dtype": "float16",
255
+ "shape": [3, 4],
256
+ "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] }
257
+ }
258
+ },
259
+ "outputs": { "y": { "dtype": "uint32", "shape": [3] } }
260
+ },
261
+ {
262
+ "name": "axis1_select_last_index_ties",
263
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
264
+ "inputs": {
265
+ "x": {
266
+ "dtype": "float32",
267
+ "shape": [2, 5],
268
+ "data": { "kind": "values", "values": [1.0, -3.0, -3.0, 2.0, -3.0, 1.0, 1.0, 2.0, 1.0, 3.0] }
269
+ }
270
+ },
271
+ "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
272
+ },
273
+ {
274
+ "name": "axis1_select_last_all_positive_infinity",
275
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
276
+ "inputs": {
277
+ "x": {
278
+ "dtype": "float32",
279
+ "shape": [2, 4],
280
+ "data": {
281
+ "kind": "values",
282
+ "values": ["Infinity", "Infinity", "Infinity", "Infinity", "Infinity", 5.0, "Infinity", 5.0]
283
+ }
284
+ }
285
+ },
286
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
287
+ },
288
+ {
289
+ "name": "axis_minus_one",
290
+ "attrs": { "axis": -1, "keepdims": 0 },
291
+ "inputs": {
292
+ "x": {
293
+ "dtype": "float32",
294
+ "shape": [2, 4],
295
+ "data": { "kind": "values", "values": [1.0, -2.0, 9.0, 4.0, 5.0, 8.0, -7.0, 6.0] }
296
+ }
297
+ },
298
+ "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
299
+ },
300
+ {
301
+ "name": "axis1_keepdims",
302
+ "attrs": { "axis": 1, "keepdims": 1 },
303
+ "inputs": {
304
+ "x": {
305
+ "dtype": "float32",
306
+ "shape": [2, 4],
307
+ "data": { "kind": "values", "values": [4.0, -3.0, -3.0, 1.0, 2.0, 2.0, -5.0, -5.0] }
308
+ }
309
+ },
310
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1] } }
311
+ },
312
+ {
313
+ "name": "rank3_axis0",
314
+ "attrs": { "axis": 0, "keepdims": 0 },
315
+ "inputs": {
316
+ "x": {
317
+ "dtype": "float32",
318
+ "shape": [2, 2, 3],
319
+ "data": { "kind": "values", "values": [4.0, -3.0, 5.0, 1.0, 2.0, -5.0, 3.0, -4.0, 6.0, 0.0, 3.0, -6.0] }
320
+ }
321
+ },
322
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 3], "tolerance": 0 } }
323
+ },
324
+ {
325
+ "name": "rank4_axis1_channel_select_last_ties",
326
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
327
+ "inputs": {
328
+ "x": {
329
+ "dtype": "float32",
330
+ "shape": [1, 3, 2, 2],
331
+ "data": { "kind": "values", "values": [1.0, -5.0, 3.0, -5.0, -2.0, -5.0, 3.0, -7.0, -2.0, -4.0, 0.0, -7.0] }
332
+ }
333
+ },
334
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2], "tolerance": 0 } }
335
+ },
336
+ {
337
+ "name": "rank4_axis2_spatial_keepdims_select_last_ties",
338
+ "attrs": { "axis": 2, "keepdims": 1, "select_last_index": 1 },
339
+ "inputs": {
340
+ "x": {
341
+ "dtype": "float32",
342
+ "shape": [1, 2, 3, 2],
343
+ "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] }
344
+ }
345
+ },
346
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 1, 2], "tolerance": 0 } }
347
+ },
348
+ {
349
+ "name": "rank4_last_axis_nhwc_keepdims_select_last",
350
+ "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
351
+ "inputs": {
352
+ "x": {
353
+ "dtype": "float32",
354
+ "shape": [1, 2, 2, 3],
355
+ "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] }
356
+ }
357
+ },
358
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2, 1], "tolerance": 0 } }
359
+ },
360
+ {
361
+ "name": "rank1_axis0_keepdims_select_last",
362
+ "attrs": { "axis": 0, "keepdims": 1, "select_last_index": 1 },
363
+ "inputs": {
364
+ "x": {
365
+ "dtype": "float32",
366
+ "shape": [6],
367
+ "data": { "kind": "values", "values": [4.0, -8.0, 2.0, -8.0, 0.0, -1.0] }
368
+ }
369
+ },
370
+ "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
371
+ },
372
+ {
373
+ "name": "ort_uint8_axis0_no_keepdims",
374
+ "provenance": {
375
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
376
+ "test": "ReductionOpTest.ArgMin_uint8"
377
+ },
378
+ "attrs": { "axis": 0, "keepdims": 0 },
379
+ "inputs": {
380
+ "x": {
381
+ "dtype": "uint8",
382
+ "shape": [3, 2, 2],
383
+ "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
384
+ }
385
+ },
386
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
387
+ },
388
+ {
389
+ "name": "ort_int8_axis0_no_keepdims",
390
+ "provenance": {
391
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
392
+ "test": "ReductionOpTest.ArgMin_int8"
393
+ },
394
+ "attrs": { "axis": 0, "keepdims": 0 },
395
+ "inputs": {
396
+ "x": {
397
+ "dtype": "int8",
398
+ "shape": [3, 2, 2],
399
+ "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
400
+ }
401
+ },
402
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
403
+ },
404
+ {
405
+ "name": "ort_int32_axis0_select_last_dups",
406
+ "provenance": {
407
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
408
+ "test": "ReductionOpTest.ArgMin_int32_select_last"
409
+ },
410
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
411
+ "inputs": {
412
+ "x": {
413
+ "dtype": "int32",
414
+ "shape": [3, 2, 2],
415
+ "data": { "kind": "values", "values": [1, 2, 3, 4, 1, 6, 7, 8, 9, 10, 11, 12] }
416
+ }
417
+ },
418
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
419
+ },
420
+ {
421
+ "name": "ort_int32_axis0_no_keepdims",
422
+ "provenance": {
423
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
424
+ "test": "ReductionOpTest.ArgMin_int32"
425
+ },
426
+ "attrs": { "axis": 0, "keepdims": 0 },
427
+ "inputs": {
428
+ "x": {
429
+ "dtype": "int32",
430
+ "shape": [3, 2, 2],
431
+ "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
432
+ }
433
+ },
434
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
435
+ },
436
+ {
437
+ "name": "ort_axis0_keepdims_rank3_f32",
438
+ "provenance": {
439
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
440
+ "test": "ReductionOpTest.ArgMin"
441
+ },
442
+ "attrs": { "axis": 0, "keepdims": 1 },
443
+ "inputs": {
444
+ "x": {
445
+ "dtype": "float32",
446
+ "shape": [3, 2, 2],
447
+ "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] }
448
+ }
449
+ },
450
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2], "tolerance": 0 } }
451
+ },
452
+ {
453
+ "name": "ort_axis0_no_keepdims_rank3_f32",
454
+ "provenance": {
455
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
456
+ "test": "ReductionOpTest.ArgMin_do_not_keepdims"
457
+ },
458
+ "attrs": { "axis": 0, "keepdims": 0 },
459
+ "inputs": {
460
+ "x": {
461
+ "dtype": "float32",
462
+ "shape": [3, 2, 2],
463
+ "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] }
464
+ }
465
+ },
466
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
467
+ },
468
+ {
469
+ "name": "ort_axis0_no_keepdims_rank1_scalar_f32",
470
+ "provenance": {
471
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
472
+ "test": "ReductionOpTest.ArgMin_do_not_keepdims_2"
473
+ },
474
+ "attrs": { "axis": 0, "keepdims": 0 },
475
+ "inputs": { "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1.0, 2.0, 3.0] } } },
476
+ "outputs": { "y": { "dtype": "uint32", "shape": [], "tolerance": 0 } }
477
+ },
478
+ {
479
+ "name": "ort_axis0_no_keepdims_rank1_scalar_select_last_f32",
480
+ "provenance": {
481
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
482
+ "test": "ReductionOpTest.ArgMin_do_not_keepdims_2_select_last"
483
+ },
484
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
485
+ "inputs": { "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1.0, 2.0, 3.0] } } },
486
+ "outputs": { "y": { "dtype": "uint32", "shape": [], "tolerance": 0 } }
487
+ },
488
+ {
489
+ "name": "ort_float_first_index_negative_infinity_deterministic",
490
+ "provenance": {
491
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
492
+ "test": "ReductionOpTest.ArgMin_float_first_index_random",
493
+ "notes": "Deterministic compact projection of ORT's random negative-infinity first-index test."
494
+ },
495
+ "attrs": { "axis": 0, "keepdims": 1, "select_last_index": 0 },
496
+ "inputs": {
497
+ "x": {
498
+ "dtype": "float32",
499
+ "shape": [16],
500
+ "data": {
501
+ "kind": "values",
502
+ "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]
503
+ }
504
+ }
505
+ },
506
+ "outputs": { "y": { "dtype": "uint32", "shape": [1], "tolerance": 0 } }
507
+ },
508
+ {
509
+ "name": "ort_negative_axis_int32_no_keepdims",
510
+ "provenance": {
511
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
512
+ "test": "ReductionOpTest.ArgMin_int32_neg_axis"
513
+ },
514
+ "attrs": { "axis": -3, "keepdims": 0 },
515
+ "inputs": {
516
+ "x": {
517
+ "dtype": "int32",
518
+ "shape": [3, 2, 2],
519
+ "data": { "kind": "values", "values": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] }
520
+ }
521
+ },
522
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
523
+ },
524
+ {
525
+ "name": "int32_axis1_exact_above_float24",
526
+ "attrs": { "axis": 1, "keepdims": 0 },
527
+ "inputs": {
528
+ "x": {
529
+ "dtype": "int32",
530
+ "shape": [2, 4],
531
+ "data": { "kind": "values", "values": [16777217, 16777216, 5, 9, -16777216, -16777217, 0, 3] }
532
+ }
533
+ },
534
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
535
+ },
536
+ {
537
+ "name": "onnx_backend_argmin_default_axis_example",
538
+ "provenance": {
539
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_example",
540
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
541
+ },
542
+ "attrs": { "keepdims": 1, "axis": 0 },
543
+ "inputs": {
544
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 1.0, 3.0, 10.0] } }
545
+ },
546
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2], "tolerance": 0 } }
547
+ },
548
+ {
549
+ "name": "onnx_backend_argmin_default_axis_example_select_last_index",
550
+ "provenance": {
551
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_example_select_last_index",
552
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
553
+ },
554
+ "attrs": { "keepdims": 1, "select_last_index": 1, "axis": 0 },
555
+ "inputs": {
556
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
557
+ },
558
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2], "tolerance": 0 } }
559
+ },
560
+ {
561
+ "name": "onnx_backend_argmin_default_axis_random",
562
+ "provenance": {
563
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_random",
564
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
565
+ },
566
+ "attrs": { "keepdims": 1, "axis": 0 },
567
+ "inputs": {
568
+ "x": {
569
+ "dtype": "float32",
570
+ "shape": [2, 3, 4],
571
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
572
+ }
573
+ },
574
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 3, 4], "tolerance": 0 } }
575
+ },
576
+ {
577
+ "name": "onnx_backend_argmin_default_axis_random_select_last_index",
578
+ "provenance": {
579
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_default_axis_random_select_last_index",
580
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
581
+ },
582
+ "attrs": { "keepdims": 1, "select_last_index": 1, "axis": 0 },
583
+ "inputs": {
584
+ "x": {
585
+ "dtype": "float32",
586
+ "shape": [2, 3, 4],
587
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
588
+ }
589
+ },
590
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 3, 4], "tolerance": 0 } }
591
+ },
592
+ {
593
+ "name": "onnx_backend_argmin_keepdims_example",
594
+ "provenance": {
595
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_example",
596
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
597
+ },
598
+ "attrs": { "axis": 1, "keepdims": 1 },
599
+ "inputs": {
600
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 1.0, 3.0, 10.0] } }
601
+ },
602
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
603
+ },
604
+ {
605
+ "name": "onnx_backend_argmin_keepdims_example_select_last_index",
606
+ "provenance": {
607
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_example_select_last_index",
608
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
609
+ },
610
+ "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
611
+ "inputs": {
612
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
613
+ },
614
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
615
+ },
616
+ {
617
+ "name": "onnx_backend_argmin_keepdims_random",
618
+ "provenance": {
619
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random",
620
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
621
+ },
622
+ "attrs": { "axis": 1, "keepdims": 1 },
623
+ "inputs": {
624
+ "x": {
625
+ "dtype": "float32",
626
+ "shape": [2, 3, 4],
627
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
628
+ }
629
+ },
630
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1, 4], "tolerance": 0 } }
631
+ },
632
+ {
633
+ "name": "onnx_backend_argmin_keepdims_random_select_last_index",
634
+ "provenance": {
635
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random_select_last_index",
636
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
637
+ },
638
+ "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
639
+ "inputs": {
640
+ "x": {
641
+ "dtype": "float32",
642
+ "shape": [2, 3, 4],
643
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
644
+ }
645
+ },
646
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1, 4], "tolerance": 0 } }
647
+ },
648
+ {
649
+ "name": "onnx_backend_argmin_negative_axis_keepdims_example",
650
+ "provenance": {
651
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_example",
652
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
653
+ },
654
+ "attrs": { "axis": -1, "keepdims": 1 },
655
+ "inputs": {
656
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 1.0, 3.0, 10.0] } }
657
+ },
658
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
659
+ },
660
+ {
661
+ "name": "onnx_backend_argmin_negative_axis_keepdims_example_select_last_index",
662
+ "provenance": {
663
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_example_select_last_index",
664
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
665
+ },
666
+ "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
667
+ "inputs": {
668
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
669
+ },
670
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1], "tolerance": 0 } }
671
+ },
672
+ {
673
+ "name": "onnx_backend_argmin_negative_axis_keepdims_random",
674
+ "provenance": {
675
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_random",
676
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
677
+ },
678
+ "attrs": { "axis": -1, "keepdims": 1 },
679
+ "inputs": {
680
+ "x": {
681
+ "dtype": "float32",
682
+ "shape": [2, 3, 4],
683
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
684
+ }
685
+ },
686
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 3, 1], "tolerance": 0 } }
687
+ },
688
+ {
689
+ "name": "onnx_backend_argmin_negative_axis_keepdims_random_select_last_index",
690
+ "provenance": {
691
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_negative_axis_keepdims_random_select_last_index",
692
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
693
+ },
694
+ "attrs": { "axis": -1, "keepdims": 1, "select_last_index": 1 },
695
+ "inputs": {
696
+ "x": {
697
+ "dtype": "float32",
698
+ "shape": [2, 3, 4],
699
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
700
+ }
701
+ },
702
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 3, 1], "tolerance": 0 } }
703
+ },
704
+ {
705
+ "name": "onnx_backend_argmin_no_keepdims_example",
706
+ "provenance": {
707
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_example",
708
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
709
+ },
710
+ "attrs": { "axis": 1, "keepdims": 0 },
711
+ "inputs": {
712
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 1.0, 3.0, 10.0] } }
713
+ },
714
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
715
+ },
716
+ {
717
+ "name": "onnx_backend_argmin_no_keepdims_example_select_last_index",
718
+ "provenance": {
719
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_example_select_last_index",
720
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
721
+ },
722
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
723
+ "inputs": {
724
+ "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [2.0, 2.0, 3.0, 10.0] } }
725
+ },
726
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
727
+ },
728
+ {
729
+ "name": "onnx_backend_argmin_no_keepdims_random",
730
+ "provenance": {
731
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_random",
732
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
733
+ },
734
+ "attrs": { "axis": 1, "keepdims": 0 },
735
+ "inputs": {
736
+ "x": {
737
+ "dtype": "float32",
738
+ "shape": [2, 3, 4],
739
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
740
+ }
741
+ },
742
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 4], "tolerance": 0 } }
743
+ },
744
+ {
745
+ "name": "onnx_backend_argmin_no_keepdims_random_select_last_index",
746
+ "provenance": {
747
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_no_keepdims_random_select_last_index",
748
+ "notes": "ONNX arg-reduction outputs are int64; this framework uses uint32 index storage."
749
+ },
750
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
751
+ "inputs": {
752
+ "x": {
753
+ "dtype": "float32",
754
+ "shape": [2, 3, 4],
755
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_argmin_input_x" } }
756
+ }
757
+ },
758
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 4], "tolerance": 0 } }
759
+ },
760
+ {
761
+ "name": "subgroup_vec4_ties_first_index",
762
+ "attrs": { "axis": -1, "keepdims": 0 },
763
+ "inputs": {
764
+ "x": {
765
+ "dtype": "float32",
766
+ "shape": [2, 8],
767
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
768
+ }
769
+ },
770
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [5, 0] } } }
771
+ },
772
+ {
773
+ "name": "subgroup_vec4_ties_select_last_index",
774
+ "attrs": { "axis": -1, "keepdims": 0, "select_last_index": 1 },
775
+ "inputs": {
776
+ "x": {
777
+ "dtype": "float32",
778
+ "shape": [2, 8],
779
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/subgroup_vec4_ties_first_index_input_x" } }
780
+ }
781
+ },
782
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [7, 7] } } }
783
+ },
784
+ {
785
+ "name": "subgroup_scalar_f16_2x5",
786
+ "attrs": { "axis": 1, "keepdims": 0 },
787
+ "inputs": {
788
+ "x": {
789
+ "dtype": "float16",
790
+ "shape": [2, 5],
791
+ "data": { "kind": "values", "values": [1.5, 2.5, 2.5, 0.5, 1.0, -1.0, -2.0, -0.5, -0.5, -3.0] }
792
+ }
793
+ },
794
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [3, 4] } } }
795
+ },
796
+ {
797
+ "name": "subgroup_vec4_uint8_full_range",
798
+ "attrs": { "axis": 1, "keepdims": 0 },
799
+ "inputs": {
800
+ "x": { "dtype": "uint8", "shape": [2, 4], "data": { "kind": "values", "values": [255, 0, 254, 1, 2, 2, 1, 0] } }
801
+ },
802
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [1, 3] } } }
803
+ },
804
+ {
805
+ "name": "subgroup_vec4_long_row_2x256",
806
+ "attrs": { "axis": 1, "keepdims": 0 },
807
+ "inputs": {
808
+ "x": {
809
+ "dtype": "float32",
810
+ "shape": [2, 256],
811
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
812
+ }
813
+ },
814
+ "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
815
+ },
816
+ {
817
+ "name": "subgroup_min_vec4_rows128x256_ties_last",
818
+ "provenance": {
819
+ "notes": "Route-locks the one-minimum-subgroup vec4 path on variable-width adapters and verifies last-index tie ordering without a workgroup-memory combine."
820
+ },
821
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
822
+ "inputs": { "x": { "dtype": "float32", "shape": [128, 256], "data": { "kind": "constant", "value": -7.0 } } },
823
+ "outputs": {
824
+ "y": { "dtype": "uint32", "shape": [128], "data": { "kind": "constant", "value": 255 }, "tolerance": 0 }
825
+ }
826
+ },
827
+ {
828
+ "name": "subgroup_min_scalar_rows128x257_ties_first",
829
+ "provenance": {
830
+ "notes": "Route-locks the unaligned one-minimum-subgroup path and verifies first-index ordering across all lane-local chunks."
831
+ },
832
+ "attrs": { "axis": 1, "keepdims": 0 },
833
+ "inputs": { "x": { "dtype": "float32", "shape": [128, 257], "data": { "kind": "constant", "value": -7.0 } } },
834
+ "outputs": {
835
+ "y": { "dtype": "uint32", "shape": [128], "data": { "kind": "constant", "value": 0 }, "tolerance": 0 }
836
+ }
837
+ },
838
+ {
839
+ "name": "subgroup_scalar_long_row_2x65",
840
+ "attrs": { "axis": 1, "keepdims": 0 },
841
+ "inputs": {
842
+ "x": {
843
+ "dtype": "float32",
844
+ "shape": [2, 65],
845
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.11 }
846
+ }
847
+ },
848
+ "outputs": { "y": { "dtype": "uint32", "shape": [2] } }
849
+ },
850
+ {
851
+ "name": "empty_zero_dim",
852
+ "attrs": { "axis": 1, "keepdims": 0 },
853
+ "inputs": { "x": { "dtype": "float32", "shape": [0, 4], "data": { "kind": "values", "values": [] } } },
854
+ "outputs": { "y": { "dtype": "uint32", "shape": [0], "tolerance": 0 } }
855
+ },
856
+ {
857
+ "name": "empty_zero_dim_f16",
858
+ "attrs": { "axis": 1, "keepdims": 0 },
859
+ "inputs": { "x": { "dtype": "float16", "shape": [0, 4], "data": { "kind": "values", "values": [] } } },
860
+ "outputs": { "y": { "dtype": "uint32", "shape": [0], "tolerance": 0 } }
861
+ },
862
+ {
863
+ "name": "empty_zero_dim_int32",
864
+ "attrs": { "axis": 1, "keepdims": 0 },
865
+ "inputs": { "x": { "dtype": "int32", "shape": [0, 4], "data": { "kind": "values", "values": [] } } },
866
+ "outputs": { "y": { "dtype": "uint32", "shape": [0], "tolerance": 0 } }
867
+ },
868
+ {
869
+ "name": "rank4_axis0_batch_argmin_int32",
870
+ "provenance": {
871
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
872
+ "test": "ReductionOpTest.ArgMin_int32 (extended to rank4 batch axis)",
873
+ "notes": "Reduces the leading batch axis of a rank-4 NCHW int32 tensor."
874
+ },
875
+ "attrs": { "axis": 0, "keepdims": 0 },
876
+ "inputs": {
877
+ "x": {
878
+ "dtype": "int32",
879
+ "shape": [3, 2, 2, 2],
880
+ "data": {
881
+ "kind": "values",
882
+ "values": [5, 9, 4, 8, 7, 2, 6, 3, 3, 1, 9, 2, 0, 8, 5, 7, 8, 4, 1, 6, 9, 3, 2, 1]
883
+ }
884
+ }
885
+ },
886
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2, 2], "tolerance": 0 } }
887
+ },
888
+ {
889
+ "name": "rank4_axis0_batch_argmin_f32_keepdims",
890
+ "provenance": {
891
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_argmin_keepdims_random (extended to rank4 batch axis)",
892
+ "notes": "Same rank4 axis=0 selection hole on the default float32 dtype with keepdims=1, to prove the gap is dtype-independent (not an integer-only branch)."
893
+ },
894
+ "attrs": { "axis": 0, "keepdims": 1 },
895
+ "inputs": {
896
+ "x": {
897
+ "dtype": "float32",
898
+ "shape": [3, 2, 2, 2],
899
+ "data": {
900
+ "kind": "values",
901
+ "values": [4.0, -3.0, 5.0, 1.0, 2.0, -5.0, 3.0, -4.0, 6.0, 0.0, 3.0, -6.0, -1.0, 7.0, 2.0, 8.0, -9.0, 4.0, 5.0, 1.0, -2.0, 3.0, 0.0, 6.0]
902
+ }
903
+ }
904
+ },
905
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2, 2, 2], "tolerance": 0 } }
906
+ },
907
+ {
908
+ "name": "rank4_axis1_channel_int8_nchw",
909
+ "provenance": {
910
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
911
+ "test": "ReductionOpTest.ArgMin_int8 (rank4 NCHW channel argmin)",
912
+ "notes": "int8 channel-axis argmin on a rank4 NCHW tensor. rank4_axis1 gates dtype only via f16Ok(T) (true for int8) so this is a WORKING-but-untested integer dtype on the rank4 selection path."
913
+ },
914
+ "attrs": { "axis": 1, "keepdims": 0 },
915
+ "inputs": {
916
+ "x": {
917
+ "dtype": "int8",
918
+ "shape": [2, 4, 3, 3],
919
+ "data": {
920
+ "kind": "values",
921
+ "values": [3, 10, 17, 24, 31, 38, 45, 52, 59, 66, 73, 80, 87, 94, 101, 108, 115, 122, 127, -120, -113, -106, -99, -92, -85, -78, -71, -64, -57, -50, -43, -36, -29, -22, -15, -8, -1, 6, 13, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90, 97, 104, 111, 118, 125, -124, -117, -110, -103, -96, -89, -82, -75, -68, -61, -54, -47, -40, -33, -26, -19, -12]
922
+ }
923
+ }
924
+ },
925
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 3, 3], "tolerance": 0 } }
926
+ },
927
+ {
928
+ "name": "rank3_axis1_f16_keepdims_select_last",
929
+ "provenance": {
930
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
931
+ "test": "ReductionOpTest.ArgMin (rank3 middle-axis f16 select_last)",
932
+ "notes": "f16 middle-axis argmin on rank3 with keepdims=1 and select_last_index=1. rank3_axis1 admits f16 via f16Ok(T) on an f16 device; no existing rank3/rank4 case is f16."
933
+ },
934
+ "attrs": { "axis": 1, "keepdims": 1, "select_last_index": 1 },
935
+ "inputs": {
936
+ "x": {
937
+ "dtype": "float16",
938
+ "shape": [2, 5, 4],
939
+ "data": {
940
+ "kind": "values",
941
+ "values": [-1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, -1.5, -1.0, -0.5, 0.0, 0.5]
942
+ }
943
+ }
944
+ },
945
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1, 4], "tolerance": 0 } }
946
+ },
947
+ {
948
+ "name": "axis1_serial_fallback_rows32768_cols5_select_last",
949
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
950
+ "inputs": {
951
+ "x": {
952
+ "dtype": "float32",
953
+ "shape": [32768, 5],
954
+ "data": { "kind": "cycle", "values": [2.0, -3.0, -3.0, 1.0, -3.0, 2.0] }
955
+ }
956
+ },
957
+ "outputs": { "y": { "dtype": "uint32", "shape": [32768], "tolerance": 0 } }
958
+ },
959
+ {
960
+ "name": "axis1_serial_fallback_rows32768_cols5_first_index",
961
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 0 },
962
+ "inputs": {
963
+ "x": {
964
+ "dtype": "float32",
965
+ "shape": [32768, 5],
966
+ "data": { "kind": "cycle", "values": [2.0, -3.0, -3.0, 1.0, -3.0, 2.0] }
967
+ }
968
+ },
969
+ "outputs": { "y": { "dtype": "uint32", "shape": [32768], "tolerance": 0 } }
970
+ },
971
+ {
972
+ "name": "axis0_fallback_rows4_int32_select_last_ties",
973
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
974
+ "inputs": {
975
+ "x": {
976
+ "dtype": "int32",
977
+ "shape": [4, 4],
978
+ "data": { "kind": "values", "values": [5, 1, 3, 2, 5, 7, 3, 2, 4, 7, 9, 2, 4, 1, 9, 8] }
979
+ }
980
+ },
981
+ "outputs": {
982
+ "y": { "dtype": "uint32", "shape": [4], "tolerance": 0, "data": { "kind": "values", "values": [3, 3, 1, 2] } }
983
+ }
984
+ },
985
+ {
986
+ "name": "axis1_int8_serial_fallback_negative_min_select_last",
987
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
988
+ "inputs": {
989
+ "x": {
990
+ "dtype": "int8",
991
+ "shape": [2, 8],
992
+ "data": {
993
+ "kind": "values",
994
+ "values": [-5, -5, -128, -128, 3, 3, 127, 127, 0, -128, 0, -127, -100, -100, -100, -50]
995
+ }
996
+ }
997
+ },
998
+ "outputs": { "y": { "dtype": "uint32", "shape": [2], "tolerance": 0 } }
999
+ },
1000
+ {
1001
+ "name": "axis0_splitk_i32_8192x16_precision_min",
1002
+ "attrs": { "axis": 0, "keepdims": 0 },
1003
+ "inputs": {
1004
+ "x": {
1005
+ "dtype": "int32",
1006
+ "shape": [8192, 16],
1007
+ "data": {
1008
+ "kind": "cycle",
1009
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_min_input_x" }
1010
+ }
1011
+ }
1012
+ },
1013
+ "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
1014
+ },
1015
+ {
1016
+ "name": "axis0_splitk_u32_8192x16_precision_min",
1017
+ "attrs": { "axis": 0, "keepdims": 0 },
1018
+ "inputs": {
1019
+ "x": {
1020
+ "dtype": "uint32",
1021
+ "shape": [8192, 16],
1022
+ "data": {
1023
+ "kind": "cycle",
1024
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_min_input_x" }
1025
+ }
1026
+ }
1027
+ },
1028
+ "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
1029
+ },
1030
+ {
1031
+ "name": "axis0_tilecols_i32_64x16_precision_min",
1032
+ "attrs": { "axis": 0, "keepdims": 0 },
1033
+ "inputs": {
1034
+ "x": {
1035
+ "dtype": "int32",
1036
+ "shape": [64, 16],
1037
+ "data": {
1038
+ "kind": "cycle",
1039
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_min_input_x" }
1040
+ }
1041
+ }
1042
+ },
1043
+ "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
1044
+ },
1045
+ {
1046
+ "name": "axis0_tilecols_i32_64x16_precision_select_last_min",
1047
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
1048
+ "inputs": {
1049
+ "x": {
1050
+ "dtype": "int32",
1051
+ "shape": [64, 16],
1052
+ "data": {
1053
+ "kind": "cycle",
1054
+ "values": { "$ref": "#/fixtureArrays/axis0_splitk_i32_8192x16_precision_min_input_x" }
1055
+ }
1056
+ }
1057
+ },
1058
+ "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
1059
+ },
1060
+ {
1061
+ "name": "axis0_narrow_splitk_16384x8_f32",
1062
+ "provenance": {
1063
+ "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."
1064
+ },
1065
+ "attrs": { "axis": 0, "keepdims": 0 },
1066
+ "inputs": {
1067
+ "x": {
1068
+ "dtype": "float32",
1069
+ "shape": [16384, 8],
1070
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
1071
+ }
1072
+ },
1073
+ "outputs": { "y": { "dtype": "uint32", "shape": [8], "tolerance": 0 } }
1074
+ },
1075
+ {
1076
+ "name": "last_axis_split_vec4_logits_ties_first",
1077
+ "provenance": {
1078
+ "notes": "A single 32768-element row route-locks the cooperative split-row vec4 reduction. Equal minima span many split segments and must merge to the first global index."
1079
+ },
1080
+ "attrs": { "axis": 1, "keepdims": 0 },
1081
+ "inputs": {
1082
+ "x": { "dtype": "float32", "shape": [1, 32768], "data": { "kind": "cycle", "values": [-7.0, 3.0, -7.0, 1.0] } }
1083
+ },
1084
+ "outputs": {
1085
+ "y": { "dtype": "uint32", "shape": [1], "data": { "kind": "constant", "value": 0 }, "tolerance": 0 }
1086
+ }
1087
+ },
1088
+ {
1089
+ "name": "last_axis_split_scalar_logits_ties_last",
1090
+ "provenance": {
1091
+ "notes": "The unaligned 32769-element companion route-locks the scalar split-row reduction and verifies select_last_index across split boundaries."
1092
+ },
1093
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
1094
+ "inputs": { "x": { "dtype": "float32", "shape": [1, 32769], "data": { "kind": "constant", "value": -5.0 } } },
1095
+ "outputs": {
1096
+ "y": { "dtype": "uint32", "shape": [1], "data": { "kind": "constant", "value": 32768 }, "tolerance": 0 }
1097
+ }
1098
+ },
1099
+ {
1100
+ "name": "last_axis_split_vec4_i32_positive_rows",
1101
+ "provenance": {
1102
+ "notes": "int32 on the cooperative split-row vec4 reduction (32768 columns, two rows). Every value in the row is positive, so a lane identity of 0 instead of INT_MAX 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 ArgMin admits int32 via T and breaks ties on the first index."
1103
+ },
1104
+ "attrs": { "axis": 1, "keepdims": 0 },
1105
+ "inputs": {
1106
+ "x": { "dtype": "int32", "shape": [2, 32768], "data": { "kind": "cycle", "values": [9, 5, 3, 1, 2, 4, 6] } }
1107
+ },
1108
+ "outputs": {
1109
+ "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [3, 2] }, "tolerance": 0 }
1110
+ }
1111
+ },
1112
+ {
1113
+ "name": "last_axis_split_scalar_u32_above_int32_max",
1114
+ "provenance": {
1115
+ "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 [4, 2] instead of [0, 5]. This directly exercises the unsigned identity 0xffffffffu and the partial store that must write the value through without bitcasting it. ONNX ArgMin admits uint32 via T; first index on ties."
1116
+ },
1117
+ "attrs": { "axis": 1, "keepdims": 0 },
1118
+ "inputs": {
1119
+ "x": {
1120
+ "dtype": "uint32",
1121
+ "shape": [2, 32769],
1122
+ "data": {
1123
+ "kind": "cycle",
1124
+ "values": [2147483000, 4200000000, 3000000000, 4294967290, 2500000000, 3500000000, 4000000000]
1125
+ }
1126
+ }
1127
+ },
1128
+ "outputs": {
1129
+ "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [0, 5] }, "tolerance": 0 }
1130
+ }
1131
+ },
1132
+ {
1133
+ "name": "last_axis_split_vec4_f16_rows",
1134
+ "provenance": {
1135
+ "notes": "float16 on the vec4 split-row reduction: each vec4 component must be widened to f32 before comparison, and the shader must enable f16. Rank 3 (1, 2, 32768) reduces the last axis exactly as the rank-2 logits case does while keeping the reference on its hoisted-decode path. All seven cycle values are exactly representable in float16, so the answer does not depend on rounding. ONNX ArgMin admits float16 via T; first index on ties."
1136
+ },
1137
+ "attrs": { "axis": 2, "keepdims": 0 },
1138
+ "inputs": {
1139
+ "x": {
1140
+ "dtype": "float16",
1141
+ "shape": [1, 2, 32768],
1142
+ "data": { "kind": "cycle", "values": [1.5, -2.0, 0.5, 3.5, -1.0, 2.25, -3.75] }
1143
+ }
1144
+ },
1145
+ "outputs": {
1146
+ "y": { "dtype": "uint32", "shape": [1, 2], "data": { "kind": "values", "values": [6, 5] }, "tolerance": 0 }
1147
+ }
1148
+ },
1149
+ {
1150
+ "name": "last_axis_split_scalar_f16_rows",
1151
+ "provenance": {
1152
+ "notes": "The unaligned (32769-column) companion: float16 on the scalar split-row reduction, where the widening happens on the plain load rather than on a vec4 component. Rank 3 keeps the reference on its hoisted-decode path; all cycle values are exact in float16. ONNX ArgMin admits float16 via T; first index on ties."
1153
+ },
1154
+ "attrs": { "axis": 2, "keepdims": 0 },
1155
+ "inputs": {
1156
+ "x": {
1157
+ "dtype": "float16",
1158
+ "shape": [1, 2, 32769],
1159
+ "data": { "kind": "cycle", "values": [1.5, -2.0, 0.5, 3.5, -1.0, 2.25, -3.75] }
1160
+ }
1161
+ },
1162
+ "outputs": {
1163
+ "y": { "dtype": "uint32", "shape": [1, 2], "data": { "kind": "values", "values": [6, 4] }, "tolerance": 0 }
1164
+ }
1165
+ },
1166
+ {
1167
+ "name": "last_axis_split_vec4_ties_select_last_index",
1168
+ "provenance": {
1169
+ "notes": "select_last_index on the vec4 split-row route checks the last-index tie rule in both the segment reduction and combine. The minimum -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."
1170
+ },
1171
+ "attrs": { "axis": 1, "keepdims": 0, "select_last_index": 1 },
1172
+ "inputs": {
1173
+ "x": {
1174
+ "dtype": "float32",
1175
+ "shape": [2, 32768],
1176
+ "data": { "kind": "cycle", "values": [-7.0, -3.0, -7.0, -1.0, -5.0, -7.0, -2.0] }
1177
+ }
1178
+ },
1179
+ "outputs": {
1180
+ "y": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [32767, 32766] }, "tolerance": 0 }
1181
+ }
1182
+ },
1183
+ {
1184
+ "name": "rank5_middle_axis_generic_geometry",
1185
+ "provenance": {
1186
+ "notes": "Proves the rank-independent (outer, axis, inner) fallback covers ranks beyond the former rank-4 manifest specializations."
1187
+ },
1188
+ "attrs": { "axis": 2, "keepdims": 0, "select_last_index": 1 },
1189
+ "inputs": {
1190
+ "x": {
1191
+ "dtype": "float32",
1192
+ "shape": [2, 2, 3, 2, 2],
1193
+ "data": { "kind": "cycle", "values": [3.0, -5.0, -1.0, -5.0, -2.0, 1.0] }
1194
+ }
1195
+ },
1196
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2, 2, 2], "tolerance": 0 } }
1197
+ },
1198
+ {
1199
+ "name": "axis0_splitk_select_last_index_ties_8192x16",
1200
+ "provenance": {
1201
+ "notes": "select_last_index=1 on the split-tiled route, which no case had ever rendered: every existing case for that geometry leaves the attribute at its default. The 17-value cycle is coprime with the 16 columns, so each column sees every value about 482 times and the minimum is a 482-way tie whose winning index differs per column -- exactly the tie-break the two arms disagree about, and one a per-lane fold can get wrong even when its lane-to-index map is monotone."
1202
+ },
1203
+ "attrs": { "axis": 0, "keepdims": 0, "select_last_index": 1 },
1204
+ "inputs": {
1205
+ "x": {
1206
+ "dtype": "int32",
1207
+ "shape": [8192, 16],
1208
+ "data": { "kind": "cycle", "values": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] }
1209
+ }
1210
+ },
1211
+ "outputs": { "y": { "dtype": "uint32", "shape": [16], "tolerance": 0 } }
1212
+ }
1213
+ ]
1214
+ }