Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Hardmax

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
df073ed
·
verified ·
1 Parent(s): 5c58519

sync 91d990483a17

Browse files
Files changed (9) hide show
  1. README.md +11 -7
  2. build/webgpu/bench.json +0 -1
  3. build/webgpu/hardmax-axis-tree.wgsl.jinja +1 -4
  4. build/webgpu/hardmax-last-axis-subgroup.wgsl.jinja +30 -17
  5. build/webgpu/hardmax-last-axis-vec4.wgsl.jinja +3 -6
  6. build/webgpu/hardmax.wgsl.jinja +4 -7
  7. build/webgpu/manifest.json +75 -132
  8. build/webgpu/metadata.json +18 -10
  9. build/webgpu/test.json +9 -10
README.md CHANGED
@@ -18,15 +18,15 @@ See the [ONNX `Hardmax` spec](https://onnx.ai/onnx/operators/onnx__Hardmax.html)
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
- | `input` | `x` | `T` | — | — | Input tensor with rank at least 1. | required |
24
 
25
  ## Outputs
26
 
27
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
28
  | --- | --- | --- | --- | --- | --- | --- |
29
- | `output` | `y` | `T` | same as `input` | same as `input` | The output tensor with the same shape as the input, containing hardmax values. | required |
30
 
31
  ## Attributes
32
 
@@ -44,7 +44,7 @@ Default values (overridable per request):
44
 
45
  ## Files
46
 
47
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
48
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
49
  - [`test.json`](build/webgpu/test.json) — correctness cases
50
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -55,10 +55,14 @@ Default values (overridable per request):
55
 
56
  ## Use with `@huggingface/kernels`
57
 
58
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
59
- It then allocates the result tensors automatically.
 
 
 
60
 
61
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
62
 
63
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
64
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
+ | `x` | `input` | `T` | — | — | Input tensor with rank at least 1. | required |
24
 
25
  ## Outputs
26
 
27
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
28
  | --- | --- | --- | --- | --- | --- | --- |
29
+ | `y` | `output` | `T` | same as `x` | same as `x` | The output tensor with the same shape as the input, containing hardmax values. | required |
30
 
31
  ## Attributes
32
 
 
44
 
45
  ## Files
46
 
47
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
48
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
49
  - [`test.json`](build/webgpu/test.json) — correctness cases
50
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
55
 
56
  ## Use with `@huggingface/kernels`
57
 
58
+ ```sh
59
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
60
+ ```
61
+
62
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
63
 
64
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
65
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
66
 
67
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
68
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "ai.onnx.Hardmax",
3
  "cases": [
4
  {
5
  "name": "rows_4096_cols_128",
 
1
  {
 
2
  "cases": [
3
  {
4
  "name": "rows_4096_cols_128",
build/webgpu/hardmax-axis-tree.wgsl.jinja CHANGED
@@ -1,10 +1,7 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {{ env.wgsl.resourceDeclarations }}
5
 
6
  const AXIS_DIM: u32 = {{ axisDim }}u;
7
- const INNER: u32 = {{ inner }}u;
8
  const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
9
  const NEG_INF: f32 = -3.4028234663852886e38;
10
  var<workgroup> values: array<f32, WG>;
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
 
3
  const AXIS_DIM: u32 = {{ axisDim }}u;
4
+ const INNER: u32 = {{ innerSize }}u;
5
  const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
6
  const NEG_INF: f32 = -3.4028234663852886e38;
7
  var<workgroup> values: array<f32, WG>;
build/webgpu/hardmax-last-axis-subgroup.wgsl.jinja CHANGED
@@ -1,6 +1,3 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {% if useSubgroups %}
5
  enable subgroups;
6
  {% endif %}
@@ -27,7 +24,6 @@ var<workgroup> rowBestIdx: u32;
27
  @compute @workgroup_size(WG, 1, 1)
28
  fn main(@builtin(workgroup_id) wg: vec3<u32>,
29
  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
30
- @builtin(subgroup_invocation_id) sgLid: u32,
31
  @builtin(subgroup_size) sgSize: u32{% endif %}) {
32
  let row = wg.x + wg.y * params.rowStride;
33
  if (row >= params.rows) {
@@ -53,19 +49,30 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
53
  let m = subgroupMax(bestVal);
54
  let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
55
  let sgIdx = subgroupMin(cand);
56
- let safeSg = max(sgSize, 1u);
57
- let slotCount = max(1u, WG / safeSg);
58
- if (sgLid == 0u) {
59
- let slot = min(tid / safeSg, WG - 1u);
60
- wgVal[slot] = m;
61
- wgIdx[slot] = sgIdx;
62
- }
63
- workgroupBarrier();
64
-
65
- if (tid == 0u) {
66
- var outVal = wgVal[0];
67
- var outIdx = wgIdx[0];
68
- for (var i = 1u; i < slotCount; i = i + 1u) {
 
 
 
 
 
 
 
 
 
 
 
69
  let v = wgVal[i];
70
  let vi = wgIdx[i];
71
  if (vi != SENTINEL_IDX && (v > outVal || (v == outVal && vi < outIdx))) {
@@ -73,6 +80,12 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
73
  outIdx = vi;
74
  }
75
  }
 
 
 
 
 
 
76
  // Match the scalar shader: if x[row, 0] is NaN no later comparison can
77
  // replace it, and an all-NaN row selects nothing. The operator writes 1 only
78
  // where the value equals the winner, so both cases emit an all-zero row; an index
 
 
 
 
1
  {% if useSubgroups %}
2
  enable subgroups;
3
  {% endif %}
 
24
  @compute @workgroup_size(WG, 1, 1)
25
  fn main(@builtin(workgroup_id) wg: vec3<u32>,
26
  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
 
27
  @builtin(subgroup_size) sgSize: u32{% endif %}) {
28
  let row = wg.x + wg.y * params.rowStride;
29
  if (row >= params.rows) {
 
49
  let m = subgroupMax(bestVal);
50
  let cand = select(SENTINEL_IDX, bestIdx, bestVal == m);
51
  let sgIdx = subgroupMin(cand);
52
+ // Cross-subgroup fold that assumes nothing about which invocations share a
53
+ // subgroup, how many subgroups there are, or which of a subgroup's lanes are
54
+ // active: every invocation owns the slot at its own index, the elected lane
55
+ // publishes its subgroup pair there and every other lane publishes the
56
+ // identity pair (the sentinel index, which the merge skips). Each subgroup
57
+ // then folds all WG slots — lane `rank`, its dense position among the active
58
+ // lanes, walks slots rank, rank + count, ... — and one more collective merges
59
+ // the lane partials, so every slot is merged exactly once at any legal width
60
+ // and partition (the fold can never miss a trailing partial subgroup). The
61
+ // merge (greatest value, then lowest index) is commutative and associative.
62
+ var winnerIdx = sgIdx;
63
+ // A one-subgroup workgroup is already fully reduced by the collectives above.
64
+ // The test reads the `subgroup_size` builtin, which is uniform; a collective's
65
+ // result is not uniform to WGSL's analysis and may not guard a barrier.
66
+ if (sgSize != WG) {
67
+ let rank = subgroupExclusiveAdd(1u);
68
+ let count = subgroupAdd(1u);
69
+ let leader = rank == 0u;
70
+ wgVal[tid] = select(neg_inf_f32(), m, leader);
71
+ wgIdx[tid] = select(SENTINEL_IDX, sgIdx, leader);
72
+ workgroupBarrier();
73
+ var outVal = neg_inf_f32();
74
+ var outIdx = SENTINEL_IDX;
75
+ for (var i = rank; i < WG; i = i + count) {
76
  let v = wgVal[i];
77
  let vi = wgIdx[i];
78
  if (vi != SENTINEL_IDX && (v > outVal || (v == outVal && vi < outIdx))) {
 
80
  outIdx = vi;
81
  }
82
  }
83
+ let foldVal = subgroupMax(outVal);
84
+ winnerIdx = subgroupMin(select(SENTINEL_IDX, outIdx, outVal == foldVal));
85
+ }
86
+
87
+ if (tid == 0u) {
88
+ var outIdx = winnerIdx;
89
  // Match the scalar shader: if x[row, 0] is NaN no later comparison can
90
  // replace it, and an all-NaN row selects nothing. The operator writes 1 only
91
  // where the value equals the winner, so both cases emit an all-zero row; an index
build/webgpu/hardmax-last-axis-vec4.wgsl.jinja CHANGED
@@ -1,6 +1,3 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {{ env.wgsl.resourceDeclarations }}
5
 
6
  const COLS_VEC: u32 = {{ colsVec }}u;
@@ -12,10 +9,10 @@ fn is_nan_f32(value: f32) -> bool {
12
 
13
 
14
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
15
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
16
  // 2D-folded flat index: gid.y carries the high bits past the
17
- // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
18
- let row = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
19
  if (row >= params.rows) {
20
  return;
21
  }
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
 
3
  const COLS_VEC: u32 = {{ colsVec }}u;
 
9
 
10
 
11
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
12
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
13
  // 2D-folded flat index: gid.y carries the high bits past the
14
+ // per-axis dispatch fold width (outputs > 16.7M elements).
15
+ let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
16
  if (row >= params.rows) {
17
  return;
18
  }
build/webgpu/hardmax.wgsl.jinja CHANGED
@@ -1,6 +1,3 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {{ env.wgsl.resourceDeclarations }}
5
  fn is_nan_f32(value: f32) -> bool {
6
  let bits = bitcast<u32>(value);
@@ -9,13 +6,13 @@ fn is_nan_f32(value: f32) -> bool {
9
 
10
 
11
  const AXIS_DIM: u32 = {{ axisDim }}u;
12
- const INNER: u32 = {{ inner }}u;
13
 
14
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
15
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
16
  // 2D-folded flat index: gid.y carries the high bits past the
17
- // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
18
- let row = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
19
  if (row >= params.rows) {
20
  return;
21
  }
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
  fn is_nan_f32(value: f32) -> bool {
3
  let bits = bitcast<u32>(value);
 
6
 
7
 
8
  const AXIS_DIM: u32 = {{ axisDim }}u;
9
+ const INNER: u32 = {{ innerSize }}u;
10
 
11
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
12
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
13
  // 2D-folded flat index: gid.y carries the high bits past the
14
+ // per-axis dispatch fold width (outputs > 16.7M elements).
15
+ let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
16
  if (row >= params.rows) {
17
  return;
18
  }
build/webgpu/manifest.json CHANGED
@@ -2,63 +2,19 @@
2
  "domain": "ai.onnx",
3
  "name": "Hardmax",
4
  "sinceVersion": 13,
5
- "description": "Computes the hardmax of the input along a single axis: sets the position of the first maximum value along `axis` to 1 and all other positions to 0. The output has the same shape as the input.",
6
- "inputs": [{ "role": "input", "dtype": "T", "description": "Input tensor with rank at least 1." }],
7
- "outputs": [
8
- {
9
- "role": "output",
10
- "dtype": "T",
11
- "rank": "ranks.input",
12
- "description": "The output tensor with the same shape as the input, containing hardmax values.",
13
- "shape": "shapes.input"
14
- }
15
- ],
16
- "attributes": { "axis": -1 },
17
- "attributeDescriptions": {
18
- "axis": "The dimension along which hardmax is computed. Negative values count from the back; accepted range is `[-r, r-1]` where `r` is the rank of the input."
19
- },
20
  "typeConstraints": { "T": ["float32", "float16"] },
21
- "args": {
22
- "x": { "kind": "tensor", "semantic": "input", "role": "input" },
23
- "y": { "kind": "tensor", "semantic": "output", "role": "output" }
24
- },
25
- "tunables": { "WORKGROUP_SIZE": 256 },
26
- "bindingSets": {
27
- "subgroupRow": [
28
- {
29
- "name": "x",
30
- "arg": "x",
31
- "semantic": "input",
32
- "buffer": { "type": "read-only-storage" },
33
- "elementType": "$scalar"
34
- },
35
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" },
36
- {
37
- "name": "params",
38
- "semantic": "kernel.params",
39
- "buffer": { "type": "uniform" },
40
- "struct": {
41
- "name": "Params",
42
- "fields": [
43
- { "name": "rows", "type": "u32", "value": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" },
44
- { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" },
45
- {
46
- "name": "rowStride",
47
- "type": "u32",
48
- "value": "max(1, min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension))"
49
- }
50
- ]
51
- }
52
- }
53
- ]
54
- },
55
- "constants": { "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
56
  "variants": [
57
  {
58
  "id": "last_axis_row",
59
  "priority": 20,
60
- "when": ["ranks.x >= 1", "ranks.y == ranks.x", "numel(shapes.x) == numel(shapes.y)", "(attrs.axis == -1 or attrs.axis == ranks.x - 1)", "dim(shapes.x, ranks.x - 1) >= 1024", "f16Ok(dtypes.T)"],
61
- "constants": {
62
  "workgroupSize": "min(tunables.WORKGROUP_SIZE, max(32, pow2ceil(dim(shapes.x, ranks.x - 1))))",
63
  "useSubgroups": "device.features.has(\"subgroups\")"
64
  },
@@ -67,49 +23,53 @@
67
  "id": "main",
68
  "name": "Hardmax.LastAxisRow",
69
  "shader": "hardmax-last-axis-subgroup.wgsl.jinja",
70
- "bindings": "subgroupRow",
71
- "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
  }
73
  ]
74
  },
75
  {
76
  "id": "last_axis_vec4",
77
  "priority": 30,
78
- "when": ["ranks.x >= 1", "ranks.y == ranks.x", "numel(shapes.x) == numel(shapes.y)", "(attrs.axis == -1 or attrs.axis == ranks.x - 1)", "dim(shapes.x, ranks.x - 1) >= 4", "dim(shapes.x, ranks.x - 1) % 4 == 0", "f16Ok(dtypes.T)"],
79
- "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "colsVec": "dim(shapes.x, ranks.x - 1) / 4" },
80
  "passes": [
81
  {
82
  "id": "main",
83
  "name": "Hardmax.LastAxisVec4",
84
  "shader": "hardmax-last-axis-vec4.wgsl.jinja",
85
  "bindings": [
86
- {
87
- "name": "x",
88
- "arg": "x",
89
- "semantic": "input",
90
- "buffer": { "type": "read-only-storage" },
91
- "elementType": "$vectorScalar"
92
- },
93
- {
94
- "name": "y",
95
- "arg": "y",
96
- "semantic": "output",
97
- "buffer": { "type": "storage" },
98
- "elementType": "$vectorScalar"
99
- },
100
  {
101
  "name": "params",
102
- "semantic": "kernel.params",
103
- "buffer": { "type": "uniform" },
104
- "struct": {
105
- "name": "Params",
106
- "fields": [{ "name": "rows", "type": "u32", "value": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" }]
107
- }
108
  }
109
  ],
110
  "dispatch": {
111
- "threads": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)",
112
- "workgroupSize": "tunables.WORKGROUP_SIZE"
 
113
  }
114
  }
115
  ]
@@ -117,11 +77,11 @@
117
  {
118
  "id": "axis_tree",
119
  "priority": 21,
 
120
  "demoteWhen": ["device.features.has(\"subgroups\") and (attrs.axis == -1 or attrs.axis == ranks.x - 1)"],
121
- "when": ["ranks.x >= 1", "ranks.y == ranks.x", "sameShape(shapes.y, shapes.x)", "attrs.axis + ranks.x >= 0", "attrs.axis < ranks.x", "dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 1024", "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension", "f16Ok(dtypes.T)"],
122
- "constants": {
123
  "axisDim": "dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)",
124
- "inner": "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
125
  },
126
  "passes": [
127
  {
@@ -129,46 +89,39 @@
129
  "name": "Hardmax.AxisTree",
130
  "shader": "hardmax-axis-tree.wgsl.jinja",
131
  "bindings": [
132
- {
133
- "name": "x",
134
- "arg": "x",
135
- "semantic": "input",
136
- "buffer": { "type": "read-only-storage" },
137
- "elementType": "$scalar"
138
- },
139
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" },
140
  {
141
  "name": "params",
142
- "semantic": "kernel.params",
143
- "buffer": { "type": "uniform" },
144
- "struct": {
145
- "name": "Params",
146
- "fields": [
147
- {
148
- "name": "rows",
149
- "type": "u32",
150
- "value": "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
151
- },
152
- {
153
- "name": "rowStride",
154
- "type": "u32",
155
- "value": "max(1, min(rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x), device.limits.maxComputeWorkgroupsPerDimension))"
156
- }
157
- ]
158
- }
159
  }
160
  ],
161
- "dispatch": { "workgroups": "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)" }
 
 
 
 
162
  }
163
  ]
164
  },
165
  {
166
  "id": "axis",
167
- "when": ["ranks.x >= 1", "ranks.y == ranks.x", "numel(shapes.x) == numel(shapes.y)", "attrs.axis + ranks.x >= 0", "attrs.axis < ranks.x", "f16Ok(dtypes.T)"],
168
- "constants": {
169
  "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
170
  "axisDim": "dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)",
171
- "inner": "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
172
  },
173
  "passes": [
174
  {
@@ -176,33 +129,23 @@
176
  "name": "Hardmax",
177
  "shader": "hardmax.wgsl.jinja",
178
  "bindings": [
179
- {
180
- "name": "x",
181
- "arg": "x",
182
- "semantic": "input",
183
- "buffer": { "type": "read-only-storage" },
184
- "elementType": "$scalar"
185
- },
186
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" },
187
  {
188
  "name": "params",
189
- "semantic": "kernel.params",
190
- "buffer": { "type": "uniform" },
191
- "struct": {
192
- "name": "Params",
193
- "fields": [
194
- {
195
- "name": "rows",
196
- "type": "u32",
197
- "value": "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
198
- }
199
- ]
200
- }
201
  }
202
  ],
203
  "dispatch": {
204
- "threads": "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)",
205
- "workgroupSize": "tunables.WORKGROUP_SIZE"
 
206
  }
207
  }
208
  ]
 
2
  "domain": "ai.onnx",
3
  "name": "Hardmax",
4
  "sinceVersion": 13,
5
+ "inputs": { "x": { "onnx": "input", "dtype": "T" } },
6
+ "outputs": { "y": { "onnx": "output", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
7
+ "attributes": { "axis": { "default": -1 } },
 
 
 
 
 
 
 
 
 
 
 
 
8
  "typeConstraints": { "T": ["float32", "float16"] },
9
+ "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
10
+ "derive": { "scalar": "dtypes.T" },
11
+ "when": ["ranks.x >= 1", "ranks.y == ranks.x", "f16Ok(dtypes.T)"],
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
12
  "variants": [
13
  {
14
  "id": "last_axis_row",
15
  "priority": 20,
16
+ "when": ["numel(shapes.x) == numel(shapes.y)", "(attrs.axis == -1 or attrs.axis == ranks.x - 1)", "dim(shapes.x, ranks.x - 1) >= 1024"],
17
+ "derive": {
18
  "workgroupSize": "min(tunables.WORKGROUP_SIZE, max(32, pow2ceil(dim(shapes.x, ranks.x - 1))))",
19
  "useSubgroups": "device.features.has(\"subgroups\")"
20
  },
 
23
  "id": "main",
24
  "name": "Hardmax.LastAxisRow",
25
  "shader": "hardmax-last-axis-subgroup.wgsl.jinja",
26
+ "subgroupCollectivesWidth": "portable",
27
+ "bindings": [
28
+ "x",
29
+ "y",
30
+ {
31
+ "name": "params",
32
+ "struct": [
33
+ { "name": "rows", "type": "u32", "value": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" },
34
+ { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" },
35
+ {
36
+ "name": "rowStride",
37
+ "type": "u32",
38
+ "value": "max(1, min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
39
+ }
40
+ ]
41
+ }
42
+ ],
43
+ "dispatch": {
44
+ "x": "min(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
45
+ "y": "ceilDiv(numel(shapes.x) / dim(shapes.x, ranks.x - 1), 65535)",
46
+ "z": 1
47
+ }
48
  }
49
  ]
50
  },
51
  {
52
  "id": "last_axis_vec4",
53
  "priority": 30,
54
+ "when": ["numel(shapes.x) == numel(shapes.y)", "(attrs.axis == -1 or attrs.axis == ranks.x - 1)", "dim(shapes.x, ranks.x - 1) >= 4", "dim(shapes.x, ranks.x - 1) % 4 == 0"],
55
+ "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "colsVec": "dim(shapes.x, ranks.x - 1) / 4" },
56
  "passes": [
57
  {
58
  "id": "main",
59
  "name": "Hardmax.LastAxisVec4",
60
  "shader": "hardmax-last-axis-vec4.wgsl.jinja",
61
  "bindings": [
62
+ { "arg": "x", "elementType": "$vectorScalar" },
63
+ { "arg": "y", "elementType": "$vectorScalar" },
 
 
 
 
 
 
 
 
 
 
 
 
64
  {
65
  "name": "params",
66
+ "struct": [{ "name": "rows", "type": "u32", "value": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" }]
 
 
 
 
 
67
  }
68
  ],
69
  "dispatch": {
70
+ "x": "min(ceilDiv((numel(shapes.x) / dim(shapes.x, ranks.x - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
71
+ "y": "ceilDiv(ceilDiv((numel(shapes.x) / dim(shapes.x, ranks.x - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
72
+ "z": 1
73
  }
74
  }
75
  ]
 
77
  {
78
  "id": "axis_tree",
79
  "priority": 21,
80
+ "when": ["sameShape(shapes.y, shapes.x)", "attrs.axis + ranks.x >= 0", "attrs.axis < ranks.x", "dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 1024", "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
81
  "demoteWhen": ["device.features.has(\"subgroups\") and (attrs.axis == -1 or attrs.axis == ranks.x - 1)"],
82
+ "derive": {
 
83
  "axisDim": "dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)",
84
+ "innerSize": "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
85
  },
86
  "passes": [
87
  {
 
89
  "name": "Hardmax.AxisTree",
90
  "shader": "hardmax-axis-tree.wgsl.jinja",
91
  "bindings": [
92
+ "x",
93
+ "y",
 
 
 
 
 
 
94
  {
95
  "name": "params",
96
+ "struct": [
97
+ {
98
+ "name": "rows",
99
+ "type": "u32",
100
+ "value": "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
101
+ },
102
+ {
103
+ "name": "rowStride",
104
+ "type": "u32",
105
+ "value": "max(1, min(rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
106
+ }
107
+ ]
 
 
 
 
 
108
  }
109
  ],
110
+ "dispatch": {
111
+ "x": "min(rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x), 65535)",
112
+ "y": "ceilDiv(rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x), 65535)",
113
+ "z": 1
114
+ }
115
  }
116
  ]
117
  },
118
  {
119
  "id": "axis",
120
+ "when": ["numel(shapes.x) == numel(shapes.y)", "attrs.axis + ranks.x >= 0", "attrs.axis < ranks.x"],
121
+ "derive": {
122
  "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
123
  "axisDim": "dim(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)",
124
+ "innerSize": "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
125
  },
126
  "passes": [
127
  {
 
129
  "name": "Hardmax",
130
  "shader": "hardmax.wgsl.jinja",
131
  "bindings": [
132
+ "x",
133
+ "y",
 
 
 
 
 
 
134
  {
135
  "name": "params",
136
+ "struct": [
137
+ {
138
+ "name": "rows",
139
+ "type": "u32",
140
+ "value": "rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)"
141
+ }
142
+ ]
 
 
 
 
 
143
  }
144
  ],
145
  "dispatch": {
146
+ "x": "min(ceilDiv((rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)), (tunables.WORKGROUP_SIZE)), 65535)",
147
+ "y": "ceilDiv(ceilDiv((rows(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x)), (tunables.WORKGROUP_SIZE)), 65535)",
148
+ "z": 1
149
  }
150
  }
151
  ]
build/webgpu/metadata.json CHANGED
@@ -1,21 +1,29 @@
1
  {
2
  "name": "ai.onnx.Hardmax",
3
- "id": "_ai_onnx_hardmax_webgpu_0d3207a",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "bZefItlvlTUgPlW5qCO0jHoXmOmOqQo8vZMsx7mgh9A=",
11
- "hardmax-axis-tree.wgsl.jinja": "1Qd9u20cZKcxIcoVUn8UpI27dye5BPfXCtyLSIpmNq8=",
12
- "hardmax-last-axis-subgroup.wgsl.jinja": "YSX63Lf95J690RQT40ZWuwbBXDLOWNZp87ngY9T0UWE=",
13
- "hardmax-last-axis-vec4.wgsl.jinja": "4S/FaIFjlrXr1LuWMV5X+3MgBQ6FA4JMZioHRoCjVfQ=",
14
- "hardmax.wgsl.jinja": "xRSmORDzxQisVCQ78ZE+txGsYp+E/TBp8et510dk1mI=",
15
- "manifest.json": "bbiRLG6GXQQCiUUgEJEgl943Ec7p4oPPgR5QiXLmT2c=",
16
- "test.json": "CZ8WfjN4jnNK9CE1kbRfFuJeCyntQ40ql17EsKp87JI="
17
  }
18
  },
19
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
20
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Hardmax" }
 
 
 
 
 
 
 
 
21
  }
 
1
  {
2
  "name": "ai.onnx.Hardmax",
3
+ "id": "_ai_onnx_hardmax_webgpu_a93aa35",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "NChImmRQqPSuIgSNjBv92ZGiB8f6qcUKZNEy+Ft1BRk=",
11
+ "hardmax-axis-tree.wgsl.jinja": "91Y2CYNvpEDkYX7s6LNeSdRqiXWG4B5My5AHa+Hs4pY=",
12
+ "hardmax-last-axis-subgroup.wgsl.jinja": "hHh93SwjisBeuDLQTjHM7OooCSA3SbRBCL8Oe1aoDOA=",
13
+ "hardmax-last-axis-vec4.wgsl.jinja": "WUOLTeV1s0vn3+uCSE0gFqg7+Koj0NbKEM9NBx/gkl8=",
14
+ "hardmax.wgsl.jinja": "Gc5xzkfNhBTwQsFGkJ1BZMit4M1G+3gZhvBQ5VkQscI=",
15
+ "manifest.json": "FDA4PNWGBQSzgNljirb5o9tFp9Kj0j2yLn/0nTRs5/U=",
16
+ "test.json": "Zk5G+e7+1Pz+l4hmebeMQfEZf+epETOhTqJzheUMID4="
17
  }
18
  },
19
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
20
+ "webgpu": {
21
+ "manifestSpec": "2.0",
22
+ "variants": {
23
+ "last_axis_row": ["hardmax-last-axis-subgroup.wgsl.jinja"],
24
+ "last_axis_vec4": ["hardmax-last-axis-vec4.wgsl.jinja"],
25
+ "axis_tree": ["hardmax-axis-tree.wgsl.jinja"],
26
+ "axis": ["hardmax.wgsl.jinja"]
27
+ }
28
+ }
29
  }
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "ai.onnx.Hardmax",
3
  "fixtureArrays": {
4
  "ort_three_dims_input_x": [1.0856307, 0.99734545, 0.2829785, 1.5062947, 0.5786002, 1.6514366, 2.4266791, 0.42891264, 1.2659363, 0.8667404, 0.6788862, 0.09470897, 1.4913896, 0.638902, 0.44398195, 0.43435127, 2.20593, 2.1867862, 1.004054, 0.3861864, 0.7373686, 1.4907321, 0.9358339, 1.175829, 1.2538806, 0.6377515, 0.9071052, 1.4286807, 0.14006872, 0.8617549, 0.25561938, 2.798589, 1.7715331, 0.69987726, 0.92746246, 0.17363568, 0.002845916, 0.6882227, 0.87953633, 0.28362733, 0.8053665, 1.7276695, 0.3908998, 0.57380587, 0.33858904, 0.011830495, 2.3923652, 0.41291216, 0.978736, 2.2381434, 1.2940853, 1.0387882, 1.7437122, 0.79806274, 0.02968323, 1.0693159, 0.8907064, 1.7548862, 1.4956441, 1.0693927],
5
  "onnx_backend_hardmax_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902],
@@ -32,7 +31,7 @@
32
  "provenance": {
33
  "source": "onnxruntime/test/providers/cpu/math/hardmax_test.cc",
34
  "test": "HardmaxOperator.Simple",
35
- "notes": "Axis-0 companion for finite subnormal ordering: a positive subnormal is strictly greater than zero and should receive the one-hot."
36
  },
37
  "attrs": { "axis": 0 },
38
  "inputs": {
@@ -49,7 +48,7 @@
49
  "provenance": {
50
  "source": "onnxruntime/test/providers/cpu/math/hardmax_test.cc",
51
  "test": "HardmaxOperator.Simple",
52
- "notes": "Vec4 last-axis companion: a positive subnormal is strictly greater than zero and should win over zero-valued lanes."
53
  },
54
  "attrs": { "axis": -1 },
55
  "inputs": {
@@ -73,7 +72,7 @@
73
  "provenance": {
74
  "source": "onnxruntime/test/providers/cpu/math/hardmax_test.cc",
75
  "test": "HardmaxOperator.LargeNumber",
76
- "notes": "Subgroup last-axis companion: width 1025 bypasses the vec4 specialization, and a positive subnormal must still beat zero-valued lanes."
77
  },
78
  "attrs": { "axis": -1 },
79
  "inputs": {
@@ -480,7 +479,7 @@
480
  "name": "rank7_last_axis",
481
  "provenance": {
482
  "source": "ONNX spec: Hardmax permits arbitrary rank; onnxruntime CPU coerces dims [0,axis) to rows and [axis,end) to cols with no rank cap.",
483
- "notes": "Rank-7 last-axis coverage for the generalized last_axis_vec4 path. ORT and WebGPU both compute the normal one-hot result."
484
  },
485
  "attrs": { "axis": -1 },
486
  "inputs": {
@@ -511,7 +510,7 @@
511
  {
512
  "name": "f16_last_axis_subgroup_1025",
513
  "provenance": {
514
- "source": "Clean coverage: the adaptive f16 last-axis row path (dim(last) >= 1024, f16Ok) was correctness-untested (only f16 vec4 cols=128 and f16 axis-0 rank2 existed).",
515
  "notes": "Last dim 1025 (>=1024, not %4==0) skips last_axis_vec4 and selects last_axis_row with f16. linspace 0..1025 gives exact distinct f16 integers (representable to 2048) so argmax is unambiguous at the final lane."
516
  },
517
  "attrs": { "axis": -1 },
@@ -523,7 +522,7 @@
523
  {
524
  "name": "f16_axis1_rank3",
525
  "provenance": {
526
- "source": "Clean coverage: f16 'axis' scalar fallback on a non-last axis (existing f16 fixtures only cover axis-0 rank2 and last-axis rank2).",
527
  "notes": "Rank-3 axis=1 f16 selects the generic 'axis' variant with usesF16. Distinct per-column values avoid ties."
528
  },
529
  "attrs": { "axis": 1 },
@@ -540,7 +539,7 @@
540
  "name": "vec4_lone_finite_among_neg_inf",
541
  "attrs": { "axis": -1 },
542
  "provenance": {
543
- "source": "Coverage gap: vec4 last-axis path with -Infinity seed. Existing vec4 tests use finite fillFloat32/values; a row of all -Infinity except one finite element is untested.",
544
  "notes": "cols=8 selects last_axis_vec4. Only index 6 is finite (3.0); every other lane is -Infinity, so the one-hot must land at index 6."
545
  },
546
  "inputs": {
@@ -566,7 +565,7 @@
566
  "name": "subgroup_cross_slot_lowest_index_tie_1025",
567
  "attrs": { "axis": -1 },
568
  "provenance": {
569
- "source": "Coverage gap: subgroup last-axis lowest-index tie reduction. Existing 1025-wide subgroup tests use unique-max linspace, so cross-subgroup/slot ties are untested.",
570
  "notes": "cols=1025 (>=1024, %4!=0) selects last_axis_row, whose capability-adaptive reduction uses subgroup or portable execution. Alternating 0/1 makes 1.0 the max at every odd index; first occurrence is index 1, so the one-hot must land at index 1 despite ties spanning multiple reduction slots."
571
  },
572
  "inputs": { "x": { "dtype": "float32", "shape": [1, 1025], "data": { "kind": "cycle", "values": [0.0, 1.0] } } },
@@ -576,7 +575,7 @@
576
  "name": "f16_subgroup_all_equal_tie_index0_1025",
577
  "attrs": { "axis": -1 },
578
  "provenance": {
579
- "source": "Coverage gap: the adaptive f16 last-axis row path under a full tie. The existing f16_last_axis_subgroup_1025 fixture uses distinct linspace values, so an all-equal f16 row is untested.",
580
  "notes": "cols=1025 (>=1024, %4!=0) with f16 selects the adaptive last_axis_row path. All values equal (1.0) => first max is index 0, so the one-hot must land at index 0."
581
  },
582
  "inputs": { "x": { "dtype": "float16", "shape": [1, 1025], "data": { "kind": "constant", "value": 1.0 } } },
 
1
  {
 
2
  "fixtureArrays": {
3
  "ort_three_dims_input_x": [1.0856307, 0.99734545, 0.2829785, 1.5062947, 0.5786002, 1.6514366, 2.4266791, 0.42891264, 1.2659363, 0.8667404, 0.6788862, 0.09470897, 1.4913896, 0.638902, 0.44398195, 0.43435127, 2.20593, 2.1867862, 1.004054, 0.3861864, 0.7373686, 1.4907321, 0.9358339, 1.175829, 1.2538806, 0.6377515, 0.9071052, 1.4286807, 0.14006872, 0.8617549, 0.25561938, 2.798589, 1.7715331, 0.69987726, 0.92746246, 0.17363568, 0.002845916, 0.6882227, 0.87953633, 0.28362733, 0.8053665, 1.7276695, 0.3908998, 0.57380587, 0.33858904, 0.011830495, 2.3923652, 0.41291216, 0.978736, 2.2381434, 1.2940853, 1.0387882, 1.7437122, 0.79806274, 0.02968323, 1.0693159, 0.8907064, 1.7548862, 1.4956441, 1.0693927],
4
  "onnx_backend_hardmax_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902],
 
31
  "provenance": {
32
  "source": "onnxruntime/test/providers/cpu/math/hardmax_test.cc",
33
  "test": "HardmaxOperator.Simple",
34
+ "notes": "Along axis 0, a positive subnormal is strictly greater than zero and must receive the one-hot value."
35
  },
36
  "attrs": { "axis": 0 },
37
  "inputs": {
 
48
  "provenance": {
49
  "source": "onnxruntime/test/providers/cpu/math/hardmax_test.cc",
50
  "test": "HardmaxOperator.Simple",
51
+ "notes": "On the vec4 last-axis path, a positive subnormal is strictly greater than every zero-valued lane."
52
  },
53
  "attrs": { "axis": -1 },
54
  "inputs": {
 
72
  "provenance": {
73
  "source": "onnxruntime/test/providers/cpu/math/hardmax_test.cc",
74
  "test": "HardmaxOperator.LargeNumber",
75
+ "notes": "A last-axis width of 1,025 selects the subgroup path, where a positive subnormal must remain greater than zero-valued lanes."
76
  },
77
  "attrs": { "axis": -1 },
78
  "inputs": {
 
479
  "name": "rank7_last_axis",
480
  "provenance": {
481
  "source": "ONNX spec: Hardmax permits arbitrary rank; onnxruntime CPU coerces dims [0,axis) to rows and [axis,end) to cols with no rank cap.",
482
+ "notes": "Rank-7 last-axis coverage for the generalized last_axis_vec4 path. The result is one-hot along the final axis."
483
  },
484
  "attrs": { "axis": -1 },
485
  "inputs": {
 
510
  {
511
  "name": "f16_last_axis_subgroup_1025",
512
  "provenance": {
513
+ "source": "constructed fixture",
514
  "notes": "Last dim 1025 (>=1024, not %4==0) skips last_axis_vec4 and selects last_axis_row with f16. linspace 0..1025 gives exact distinct f16 integers (representable to 2048) so argmax is unambiguous at the final lane."
515
  },
516
  "attrs": { "axis": -1 },
 
522
  {
523
  "name": "f16_axis1_rank3",
524
  "provenance": {
525
+ "source": "constructed fixture",
526
  "notes": "Rank-3 axis=1 f16 selects the generic 'axis' variant with usesF16. Distinct per-column values avoid ties."
527
  },
528
  "attrs": { "axis": 1 },
 
539
  "name": "vec4_lone_finite_among_neg_inf",
540
  "attrs": { "axis": -1 },
541
  "provenance": {
542
+ "source": "constructed fixture",
543
  "notes": "cols=8 selects last_axis_vec4. Only index 6 is finite (3.0); every other lane is -Infinity, so the one-hot must land at index 6."
544
  },
545
  "inputs": {
 
565
  "name": "subgroup_cross_slot_lowest_index_tie_1025",
566
  "attrs": { "axis": -1 },
567
  "provenance": {
568
+ "source": "constructed fixture",
569
  "notes": "cols=1025 (>=1024, %4!=0) selects last_axis_row, whose capability-adaptive reduction uses subgroup or portable execution. Alternating 0/1 makes 1.0 the max at every odd index; first occurrence is index 1, so the one-hot must land at index 1 despite ties spanning multiple reduction slots."
570
  },
571
  "inputs": { "x": { "dtype": "float32", "shape": [1, 1025], "data": { "kind": "cycle", "values": [0.0, 1.0] } } },
 
575
  "name": "f16_subgroup_all_equal_tie_index0_1025",
576
  "attrs": { "axis": -1 },
577
  "provenance": {
578
+ "source": "constructed fixture",
579
  "notes": "cols=1025 (>=1024, %4!=0) with f16 selects the adaptive last_axis_row path. All values equal (1.0) => first max is index 0, so the one-hot must land at index 0."
580
  },
581
  "inputs": { "x": { "dtype": "float16", "shape": [1, 1025], "data": { "kind": "constant", "value": 1.0 } } },