Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.NonZero

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
0f561dd
·
verified ·
1 Parent(s): 47e90a2

sync 2e7068faf55e

Browse files
Files changed (9) hide show
  1. README.md +68 -0
  2. build/webgpu/bench.json +145 -0
  3. build/webgpu/manifest.json +175 -0
  4. build/webgpu/metadata.json +21 -0
  5. build/webgpu/nonzero-flag-block-scan-vec4.wgsl.jinja +69 -0
  6. build/webgpu/nonzero-scatter-vec4.wgsl.jinja +56 -0
  7. build/webgpu/nonzero.wgsl.jinja +51 -0
  8. build/webgpu/scan-block-prefix-u32.wgsl.jinja +44 -0
  9. build/webgpu/test.json +470 -0
README.md CHANGED
@@ -1,3 +1,71 @@
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.NonZero
10
+
11
+ `ai.onnx` · standard ONNX operator · ONNX opset ≥ 9
12
+
13
+ ## Description
14
+
15
+ Returns the indices of all non-zero elements in `X`, in row-major order, as a 2-D tensor of shape `[rank(X), nnz]` where each column is an N-dimensional index. This follows ONNX `NonZero`, including scalar output shape `[0, N]` and logical int64 output type. Because every emitted coordinate is bounded by a WebGPU-addressable input dimension, the backend stores this logical int64 tensor losslessly as uint32. The caller must provide the exact data-dependent output shape. String, float64, 64-bit input, and other ONNX input types not listed below are not yet implemented.
16
+
17
+ See the [ONNX `NonZero` spec](https://onnx.ai/onnx/operators/onnx__NonZero.html) for the reference semantics.
18
+
19
+ ## Inputs
20
+
21
+ | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
+ | --- | --- | --- | --- | --- | --- | --- |
23
+ | `X` | `x` | `T` | — | — | Input tensor of any shape whose non-zero element indices are to be found. | required |
24
+
25
+ ## Outputs
26
+
27
+ | Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
28
+ | --- | --- | --- | --- | --- | --- | --- | --- |
29
+ | `Y` | `y` | `I` | `uint32` | `2` | — | Logical int64 tensor of shape `[rank(X), nnz]` containing the multi-dimensional indices of non-zero elements, one index per column. The WebGPU storage representation is uint32 because every coordinate is within an addressable input dimension. | required |
30
+
31
+ ## Type constraints
32
+
33
+ | Variable | Allowed dtypes |
34
+ | --- | --- |
35
+ | `T` | `float32`, `float16`, `uint32`, `int32`, `int16`, `uint8`, `int8`, `bool` |
36
+ | `I` | `int64` |
37
+
38
+ ## Files
39
+
40
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
41
+ - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
42
+ - [`test.json`](build/webgpu/test.json) — correctness cases
43
+ - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
44
+ - [`nonzero-flag-block-scan-vec4.wgsl.jinja`](build/webgpu/nonzero-flag-block-scan-vec4.wgsl.jinja)
45
+ - [`nonzero-scatter-vec4.wgsl.jinja`](build/webgpu/nonzero-scatter-vec4.wgsl.jinja)
46
+ - [`nonzero.wgsl.jinja`](build/webgpu/nonzero.wgsl.jinja)
47
+ - [`scan-block-prefix-u32.wgsl.jinja`](build/webgpu/scan-block-prefix-u32.wgsl.jinja)
48
+
49
+ ## Use with `@huggingface/kernels`
50
+
51
+ The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
52
+
53
+ The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
54
+
55
+ - `y`
56
+
57
+ Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
58
+
59
+ The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
60
+
61
+ Replace each `*Data` placeholder with a typed array containing the corresponding input data.
62
+
63
+ ```js
64
+ import { getKernel } from "@huggingface/kernels";
65
+
66
+ const kernel = await getKernel("webgpu-kernels/ai.onnx.NonZero", { version: 1 });
67
+ // Explicit destinations request optional results or supply metadata that cannot be inferred.
68
+ const { y } = await kernel({ x: { data: xData, shape: [4] } }, {
69
+ outputs: { y: { shape: [1, 2], dtype: "int64" } },
70
+ });
71
+ ```
build/webgpu/bench.json ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "op": "ai.onnx.NonZero",
3
+ "tunableSpace": {},
4
+ "cases": [
5
+ {
6
+ "name": "f32_1024_capacity",
7
+ "inputs": {
8
+ "x": {
9
+ "dtype": "float32",
10
+ "shape": [1024],
11
+ "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.02 }
12
+ }
13
+ },
14
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1024] } }
15
+ },
16
+ {
17
+ "name": "nonzero-rank2-u32-1m",
18
+ "preset": "smoke",
19
+ "inputs": { "x": { "dtype": "uint32", "shape": [1024, 1024], "dist": "linearMod", "mod": 4 } },
20
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 786432] } },
21
+ "bench": {
22
+ "primary": true,
23
+ "metrics": [
24
+ {
25
+ "type": "bandwidth",
26
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
27
+ }
28
+ ]
29
+ }
30
+ },
31
+ {
32
+ "name": "nonzero-exact-capacity-1024-f32",
33
+ "preset": "smoke",
34
+ "provenance": {
35
+ "notes": "Legacy case name retained for longitudinal comparisons. At the measured crossover, 1024 elements now select the parallel scan. The traffic metric uses algorithm-independent logical I/O so an all-variant comparison weighs serial and parallel kernels fairly."
36
+ },
37
+ "inputs": { "x": { "dtype": "float32", "shape": [1024], "dist": "constant", "value": 1 } },
38
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1024] } },
39
+ "bench": {
40
+ "metrics": [
41
+ {
42
+ "type": "bandwidth",
43
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
44
+ }
45
+ ]
46
+ }
47
+ },
48
+ {
49
+ "name": "nonzero-parallel-1280-f32-healthy",
50
+ "preset": "smoke",
51
+ "inputs": { "x": { "dtype": "float32", "shape": [1280], "dist": "constant", "value": 1 } },
52
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1280] } },
53
+ "bench": {
54
+ "metrics": [
55
+ {
56
+ "type": "bandwidth",
57
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
58
+ }
59
+ ]
60
+ }
61
+ },
62
+ {
63
+ "name": "nonzero-rank2-f32-1m-dense-100pct",
64
+ "preset": "smoke",
65
+ "inputs": { "x": { "dtype": "float32", "shape": [1024, 1024], "dist": "constant", "value": 1 } },
66
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1048576] } },
67
+ "bench": {
68
+ "metrics": [
69
+ {
70
+ "type": "bandwidth",
71
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
72
+ }
73
+ ]
74
+ }
75
+ },
76
+ {
77
+ "name": "nonzero-rank2-f16-1m-dense",
78
+ "preset": "smoke",
79
+ "inputs": { "x": { "dtype": "float16", "shape": [1024, 1024], "dist": "constant", "value": 1 } },
80
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 1048576] } },
81
+ "bench": {
82
+ "metrics": [
83
+ {
84
+ "type": "bandwidth",
85
+ "value": "2 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
86
+ }
87
+ ]
88
+ }
89
+ },
90
+ {
91
+ "name": "nonzero-dispatch-cliff-16m-f32-sparse",
92
+ "preset": "smoke",
93
+ "inputs": { "x": { "dtype": "float32", "shape": [4096, 4096], "dist": "constant", "value": 0 } },
94
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 0] } },
95
+ "bench": {
96
+ "metrics": [{ "type": "bandwidth", "value": "4 * numel(shapes.x) + 4 * (2 * ceilDiv(numel(shapes.x), 4096))" }]
97
+ }
98
+ },
99
+ {
100
+ "name": "nonzero-2dfold-boundary-16m-int32-exact",
101
+ "preset": "stress",
102
+ "inputs": { "x": { "dtype": "int32", "shape": [4096, 4097], "dist": "normal", "seed": 7, "scale": 2 } },
103
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 16781312] } },
104
+ "bench": {
105
+ "metrics": [
106
+ {
107
+ "type": "bandwidth",
108
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
109
+ }
110
+ ]
111
+ }
112
+ },
113
+ {
114
+ "name": "nonzero-rank4-u32-1x32x512x512-exact-activation",
115
+ "preset": "stress",
116
+ "inputs": { "x": { "dtype": "uint32", "shape": [1, 32, 512, 512], "dist": "linearMod", "mod": 16 } },
117
+ "outputs": { "y": { "dtype": "uint32", "shape": [4, 7864320] } },
118
+ "bench": {
119
+ "metrics": [
120
+ {
121
+ "type": "bandwidth",
122
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
123
+ }
124
+ ]
125
+ }
126
+ },
127
+ {
128
+ "name": "nonzero-rank3-u32-32x512x512-exact-activation",
129
+ "preset": "stress",
130
+ "provenance": {
131
+ "notes": "Rank-3 twin of the rank-4 activation case, same element count and density. The scatter writes one coordinate stream per rank, so it exists to place the items-per-thread threshold on measured rank-2/3/4 behaviour instead of extrapolating from two points."
132
+ },
133
+ "inputs": { "x": { "dtype": "uint32", "shape": [32, 512, 512], "dist": "linearMod", "mod": 16 } },
134
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 7864320] } },
135
+ "bench": {
136
+ "metrics": [
137
+ {
138
+ "type": "bandwidth",
139
+ "value": "4 * numel(shapes.x) + 4 * numel(shapes.y) + 4 * (4 * ceilDiv(numel(shapes.x), 4) + 5 * ceilDiv(numel(shapes.x), 1024))"
140
+ }
141
+ ]
142
+ }
143
+ }
144
+ ]
145
+ }
build/webgpu/manifest.json ADDED
@@ -0,0 +1,175 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "domain": "ai.onnx",
3
+ "name": "NonZero",
4
+ "sinceVersion": 9,
5
+ "description": "Returns the indices of all non-zero elements in `X`, in row-major order, as a 2-D tensor of shape `[rank(X), nnz]` where each column is an N-dimensional index. This follows ONNX `NonZero`, including scalar output shape `[0, N]` and logical int64 output type. Because every emitted coordinate is bounded by a WebGPU-addressable input dimension, the backend stores this logical int64 tensor losslessly as uint32. The caller must provide the exact data-dependent output shape. String, float64, 64-bit input, and other ONNX input types not listed below are not yet implemented.",
6
+ "inputs": [
7
+ {
8
+ "role": "X",
9
+ "dtype": "T",
10
+ "description": "Input tensor of any shape whose non-zero element indices are to be found."
11
+ }
12
+ ],
13
+ "outputs": [
14
+ {
15
+ "role": "Y",
16
+ "dtype": "I",
17
+ "rank": 2,
18
+ "description": "Logical int64 tensor of shape `[rank(X), nnz]` containing the multi-dimensional indices of non-zero elements, one index per column. The WebGPU storage representation is uint32 because every coordinate is within an addressable input dimension."
19
+ }
20
+ ],
21
+ "typeConstraints": {
22
+ "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"],
23
+ "I": ["int64"]
24
+ },
25
+ "args": {
26
+ "x": { "kind": "tensor", "semantic": "X", "role": "input" },
27
+ "y": { "kind": "tensor", "semantic": "Y", "role": "output", "dtype": "uint32" }
28
+ },
29
+ "tunables": {
30
+ "SCAN_ITEMS_PER_THREAD": 16,
31
+ "SCAN_ITEMS_PER_THREAD_WIDE": 32,
32
+ "SCAN_ITEMS_PER_THREAD_WIDE_3COORD": 8,
33
+ "SCAN_ITEMS_PER_THREAD_WIDE_4COORD": 4,
34
+ "SCAN_MIN_BLOCKS": 256
35
+ },
36
+ "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
37
+ "variants": [
38
+ {
39
+ "id": "parallel_scan",
40
+ "priority": 15,
41
+ "when": ["ranks.X >= 1", "ranks.Y == 2", "dim(shapes.Y, 0) == ranks.X", "dim(shapes.Y, 1) <= numel(shapes.X)", "tensorDtypes.Y == \"uint32\"", "numel(shapes.X) >= 1024", "numel(shapes.X) <= 1073741823", "f16Ok(tensorDtypes.X)"],
42
+ "derive": {
43
+ "scanN": "numel(shapes.X)",
44
+ "scanItemsWide": "tunables.SCAN_ITEMS_PER_THREAD_WIDE_4COORD if ranks.X >= 4 else (tunables.SCAN_ITEMS_PER_THREAD_WIDE_3COORD if ranks.X == 3 else tunables.SCAN_ITEMS_PER_THREAD_WIDE)",
45
+ "scanItems": "scanItemsWide if ceilDiv(numel(shapes.X), 256 * scanItemsWide) >= tunables.SCAN_MIN_BLOCKS else tunables.SCAN_ITEMS_PER_THREAD",
46
+ "scanThreads": "ceilDiv(scanN, scanItems)",
47
+ "scanBlocks": "ceilDiv(scanN, 256 * scanItems)"
48
+ },
49
+ "constants": { "rank": "ranks.X", "shape": "shapes.X", "workgroupSize": 256 },
50
+ "intermediates": [
51
+ { "id": "offsets", "dtype": "uint32", "shape": "[scanThreads]" },
52
+ { "id": "flags", "dtype": "uint32", "shape": "[scanThreads]" },
53
+ { "id": "blockSums", "dtype": "uint32", "shape": "[scanBlocks]" },
54
+ { "id": "blockPrefix", "dtype": "uint32", "shape": "[scanBlocks]" }
55
+ ],
56
+ "passes": [
57
+ {
58
+ "id": "flag_scan",
59
+ "name": "NonZero.FlagBlockScan",
60
+ "source": { "shader": "nonzero-flag-block-scan-vec4.wgsl.jinja", "inputs": { "itemsPerThread": "scanItems" } },
61
+ "bindings": [
62
+ {
63
+ "name": "src",
64
+ "arg": "x",
65
+ "semantic": "X",
66
+ "buffer": { "type": "read-only-storage" },
67
+ "elementType": "$scalar"
68
+ },
69
+ { "name": "offsets", "semantic": "offsets", "buffer": { "type": "storage" }, "elementType": "u32" },
70
+ { "name": "flags", "semantic": "flags", "buffer": { "type": "storage" }, "elementType": "u32" },
71
+ { "name": "blockSums", "semantic": "blockSums", "buffer": { "type": "storage" }, "elementType": "u32" },
72
+ {
73
+ "name": "params",
74
+ "semantic": "kernel.params",
75
+ "buffer": { "type": "uniform" },
76
+ "struct": { "name": "Params", "fields": [{ "name": "n", "type": "u32", "value": "scanN" }] }
77
+ }
78
+ ],
79
+ "dispatch": { "workgroups": "scanBlocks" }
80
+ },
81
+ {
82
+ "id": "block_prefix",
83
+ "name": "NonZero.BlockPrefixScan",
84
+ "shader": "scan-block-prefix-u32.wgsl.jinja",
85
+ "bindings": [
86
+ {
87
+ "name": "blockSums",
88
+ "semantic": "blockSums",
89
+ "buffer": { "type": "read-only-storage" },
90
+ "elementType": "u32"
91
+ },
92
+ { "name": "blockPrefix", "semantic": "blockPrefix", "buffer": { "type": "storage" }, "elementType": "u32" },
93
+ {
94
+ "name": "params",
95
+ "semantic": "kernel.params",
96
+ "buffer": { "type": "uniform" },
97
+ "struct": { "name": "Params", "fields": [{ "name": "numBlocks", "type": "u32", "value": "scanBlocks" }] }
98
+ }
99
+ ],
100
+ "dispatch": { "x": 1 }
101
+ },
102
+ {
103
+ "id": "scatter",
104
+ "name": "NonZero.CoordinateScatter",
105
+ "source": { "shader": "nonzero-scatter-vec4.wgsl.jinja", "inputs": { "itemsPerThread": "scanItems" } },
106
+ "bindings": [
107
+ {
108
+ "name": "offsets",
109
+ "semantic": "offsets",
110
+ "buffer": { "type": "read-only-storage" },
111
+ "elementType": "u32"
112
+ },
113
+ { "name": "flags", "semantic": "flags", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
114
+ {
115
+ "name": "blockPrefix",
116
+ "semantic": "blockPrefix",
117
+ "buffer": { "type": "read-only-storage" },
118
+ "elementType": "u32"
119
+ },
120
+ { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "u32" },
121
+ {
122
+ "name": "params",
123
+ "semantic": "kernel.params",
124
+ "buffer": { "type": "uniform" },
125
+ "struct": {
126
+ "name": "Params",
127
+ "fields": [
128
+ { "name": "inputCount", "type": "u32", "value": "scanN" },
129
+ { "name": "capacity", "type": "u32", "value": "dim(shapes.Y, 1)" },
130
+ { "name": "numBlocks", "type": "u32", "value": "scanBlocks" }
131
+ ]
132
+ }
133
+ }
134
+ ],
135
+ "dispatch": { "workgroups": "scanBlocks" }
136
+ }
137
+ ]
138
+ },
139
+ {
140
+ "id": "serial_rankn",
141
+ "when": ["ranks.X >= 0", "ranks.Y == 2", "dim(shapes.Y, 0) == ranks.X", "dim(shapes.Y, 1) <= numel(shapes.X)", "tensorDtypes.Y == \"uint32\"", "f16Ok(tensorDtypes.X)"],
142
+ "constants": { "rank": "ranks.X", "shape": "shapes.X" },
143
+ "passes": [
144
+ {
145
+ "id": "main",
146
+ "name": "NonZero",
147
+ "shader": "nonzero.wgsl.jinja",
148
+ "bindings": [
149
+ {
150
+ "name": "x",
151
+ "arg": "x",
152
+ "semantic": "X",
153
+ "buffer": { "type": "read-only-storage" },
154
+ "elementType": "$scalar"
155
+ },
156
+ { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "u32" },
157
+ {
158
+ "name": "params",
159
+ "semantic": "kernel.params",
160
+ "buffer": { "type": "uniform" },
161
+ "struct": {
162
+ "name": "Params",
163
+ "fields": [
164
+ { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
165
+ { "name": "capacity", "type": "u32", "value": "dim(shapes.Y, 1)" }
166
+ ]
167
+ }
168
+ }
169
+ ],
170
+ "dispatch": { "x": 1 }
171
+ }
172
+ ]
173
+ }
174
+ ]
175
+ }
build/webgpu/metadata.json ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "ai.onnx.NonZero",
3
+ "id": "_ai_onnx_nonzero_webgpu_3ec09bb",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "backend": { "type": "webgpu" },
7
+ "digest": {
8
+ "algorithm": "sha256",
9
+ "files": {
10
+ "bench.json": "M3F+WpZnjPCDWEpO1dO4qDLcYExiCtbL15AykZVxep0=",
11
+ "manifest.json": "RtKjb+++Sh/aUmIAT4P7uGoEBp0DvBS7qMpDaTQ2+qA=",
12
+ "nonzero-flag-block-scan-vec4.wgsl.jinja": "ep5md2gJ0PBS1CUVQKzAB4jSKLjRvQRSIjzi0Q4pO24=",
13
+ "nonzero-scatter-vec4.wgsl.jinja": "0jNME7EBPGS5DujgBi2FGHSbl+l7j3BM/rOiiXy5K/I=",
14
+ "nonzero.wgsl.jinja": "3E27IFWj7Prj1aOYH54NhyQkFS+FKxJwBEz9KuNz7r8=",
15
+ "scan-block-prefix-u32.wgsl.jinja": "EDY0dOyl7TX0GjGdfkGRqho+8VaqNTv69XIUAGWwW7w=",
16
+ "test.json": "1xLlAXOvht1/Q9/I55mRO2vGHBtJ7O3pu0SNjQp4TWU="
17
+ }
18
+ },
19
+ "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
20
+ "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.NonZero" }
21
+ }
build/webgpu/nonzero-flag-block-scan-vec4.wgsl.jinja ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {% if usesF16 %}
2
+ enable f16;
3
+ {% endif %}
4
+ {{ env.wgsl.resourceDeclarations }}
5
+
6
+ const WG: u32 = 256u;
7
+ const ITEMS_PER_THREAD: u32 = {{ source.itemsPerThread }}u;
8
+ const BLOCK_ITEMS: u32 = WG * ITEMS_PER_THREAD;
9
+
10
+ var<workgroup> wgScan: array<u32, WG>;
11
+
12
+ fn is_nonzero(index: u32) -> bool {
13
+ {% if scalar == "f32" %}
14
+ return (bitcast<u32>(src[index]) & 0x7fffffffu) != 0u;
15
+ {% elif scalar == "f16" %}
16
+ return f32(src[index]) != 0.0;
17
+ {% elif scalar == "u32" %}
18
+ return src[index] != 0u;
19
+ {% else %}
20
+ return src[index] != 0i;
21
+ {% endif %}
22
+ }
23
+
24
+ @compute @workgroup_size(WG, 1, 1)
25
+ fn main(
26
+ @builtin(workgroup_id) workgroup: vec3<u32>,
27
+ @builtin(num_workgroups) nwg: vec3<u32>,
28
+ @builtin(local_invocation_id) local: vec3<u32>
29
+ ) {
30
+ let block = workgroup.x + workgroup.y * nwg.x;
31
+ let tid = local.x;
32
+ if (block * BLOCK_ITEMS >= params.n) {
33
+ return;
34
+ }
35
+ let base = block * BLOCK_ITEMS + tid * ITEMS_PER_THREAD;
36
+ var packed_flags = 0u;
37
+ var thread_sum = 0u;
38
+ for (var item = 0u; item < ITEMS_PER_THREAD; item = item + 1u) {
39
+ let index = base + item;
40
+ var flag = 0u;
41
+ if (index < params.n && is_nonzero(index)) {
42
+ flag = 1u;
43
+ }
44
+ packed_flags = packed_flags | (flag << item);
45
+ thread_sum = thread_sum + flag;
46
+ }
47
+
48
+ // Scan one ITEMS_PER_THREAD-element aggregate per lane. This preserves
49
+ // flattened order while shrinking the block table by the same factor, which
50
+ // is what bounds the single-workgroup block-prefix pass that follows.
51
+ wgScan[tid] = thread_sum;
52
+ for (var step = 1u; step < WG; step = step << 1u) {
53
+ workgroupBarrier();
54
+ var previous = 0u;
55
+ if (tid >= step) {
56
+ previous = wgScan[tid - step];
57
+ }
58
+ workgroupBarrier();
59
+ wgScan[tid] = wgScan[tid] + previous;
60
+ }
61
+ let thread_index = block * WG + tid;
62
+ if (base < params.n) {
63
+ offsets[thread_index] = wgScan[tid] - thread_sum;
64
+ flags[thread_index] = packed_flags;
65
+ }
66
+ if (tid == WG - 1u) {
67
+ blockSums[block] = wgScan[WG - 1u];
68
+ }
69
+ }
build/webgpu/nonzero-scatter-vec4.wgsl.jinja ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {{ env.wgsl.resourceDeclarations }}
2
+
3
+ const WG: u32 = 256u;
4
+ const ITEMS_PER_THREAD: u32 = {{ source.itemsPerThread }}u;
5
+ const BLOCK_ITEMS: u32 = WG * ITEMS_PER_THREAD;
6
+
7
+ fn write_coords(flat_index: u32, out_col: u32) {
8
+ var rem = flat_index;
9
+ {% for axis in range(rank) %}
10
+ {% set stride = namespace(value=1) %}
11
+ {% for j in range(axis + 1, rank) %}
12
+ {% set stride.value = stride.value * shape[j] %}
13
+ {% endfor %}
14
+ let coord{{ axis }} = rem / {{ stride.value }}u;
15
+ rem = rem % {{ stride.value }}u;
16
+ y[{{ axis }}u * params.capacity + out_col] = coord{{ axis }};
17
+ {% endfor %}
18
+ }
19
+
20
+ @compute @workgroup_size(WG, 1, 1)
21
+ fn main(
22
+ @builtin(workgroup_id) workgroup: vec3<u32>,
23
+ @builtin(num_workgroups) nwg: vec3<u32>,
24
+ @builtin(local_invocation_id) local: vec3<u32>,
25
+ ) {
26
+ let block = workgroup.x + workgroup.y * nwg.x;
27
+ let tid = local.x;
28
+ if (block >= params.numBlocks) {
29
+ return;
30
+ }
31
+ let base = block * BLOCK_ITEMS + tid * ITEMS_PER_THREAD;
32
+ let block_base = blockPrefix[block];
33
+
34
+ if (base >= params.inputCount) {
35
+ return;
36
+ }
37
+
38
+ // One packed flag mask and one exclusive thread offset replace ITEMS_PER_THREAD
39
+ // element offsets, and scatter never rereads X.
40
+ let thread_index = block * WG + tid;
41
+ let packed_flags = flags[thread_index];
42
+ var offset = offsets[thread_index];
43
+ for (var item = 0u; item < ITEMS_PER_THREAD; item = item + 1u) {
44
+ let index = base + item;
45
+ if (index >= params.inputCount) {
46
+ continue;
47
+ }
48
+ if ((packed_flags & (1u << item)) != 0u) {
49
+ let position = block_base + offset;
50
+ if (position < params.capacity) {
51
+ write_coords(index, position);
52
+ }
53
+ offset = offset + 1u;
54
+ }
55
+ }
56
+ }
build/webgpu/nonzero.wgsl.jinja ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {% if usesF16 %}
2
+ enable f16;
3
+ {% endif %}
4
+ {{ env.wgsl.resourceDeclarations }}
5
+
6
+ fn is_nonzero(index: u32) -> bool {
7
+ {% if scalar == "f32" %}
8
+ // Bit-test because float ALUs may flush subnormals to zero. Only +0 and -0
9
+ // have a zero magnitude field; subnormals and NaNs remain nonzero.
10
+ let bits = bitcast<u32>(x[index]);
11
+ return (bits & 0x7fffffffu) != 0u;
12
+ {% elif scalar == "f16" %}
13
+ return f32(x[index]) != 0.0;
14
+ {% elif scalar == "u32" %}
15
+ return x[index] != 0u;
16
+ {% else %}
17
+ return x[index] != 0i;
18
+ {% endif %}
19
+ }
20
+
21
+ {% if rank > 0 %}
22
+ fn write_coords(flat_index: u32, out_col: u32) {
23
+ var rem = flat_index;
24
+ {% for axis in range(rank) %}
25
+ {% set stride = namespace(value=1) %}
26
+ {% for j in range(axis + 1, rank) %}
27
+ {% set stride.value = stride.value * shape[j] %}
28
+ {% endfor %}
29
+ {% set divisor = 1 if stride.value == 0 else stride.value %}
30
+ let coord{{ axis }} = rem / {{ divisor }}u;
31
+ rem = rem % {{ divisor }}u;
32
+ y[{{ axis }}u * params.capacity + out_col] = coord{{ axis }};
33
+ {% endfor %}
34
+ }
35
+ {% endif %}
36
+
37
+ @compute @workgroup_size(1)
38
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
39
+ if (gid.x != 0u) { return; }
40
+ var total = 0u;
41
+ for (var i = 0u; i < params.inputCount; i = i + 1u) {
42
+ if (is_nonzero(i)) {
43
+ if (total < params.capacity) {
44
+ {% if rank > 0 %}
45
+ write_coords(i, total);
46
+ {% endif %}
47
+ }
48
+ total = total + 1u;
49
+ }
50
+ }
51
+ }
build/webgpu/scan-block-prefix-u32.wgsl.jinja ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Single-workgroup exclusive scan over per-block sums for stream compaction.
2
+ // The dispatch is (1, 1, 1). It walks blockSums in workgroup-sized chunks,
3
+ // scans each chunk with Hillis-Steele or subgroup collectives, and links chunks
4
+ // with a running carry. blockPrefix[b] is therefore the number of set flags in
5
+ // all blocks before b. The total count is the last prefix plus the last sum.
6
+ {{ env.wgsl.resourceDeclarations }}
7
+
8
+ const WG: u32 = {{ workgroupSize }}u;
9
+
10
+ var<workgroup> wgScan: array<u32, WG>;
11
+
12
+ @compute @workgroup_size(WG, 1, 1)
13
+ fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
14
+ let tid = lid.x;
15
+ var carry = 0u;
16
+ let chunks = (params.numBlocks + WG - 1u) / WG;
17
+ for (var c = 0u; c < chunks; c = c + 1u) {
18
+ let j = c * WG + tid;
19
+ var v = 0u;
20
+ if (j < params.numBlocks) {
21
+ v = blockSums[j];
22
+ }
23
+ workgroupBarrier();
24
+ wgScan[tid] = v;
25
+ for (var step = 1u; step < WG; step = step << 1u) {
26
+ workgroupBarrier();
27
+ var prev = 0u;
28
+ if (tid >= step) {
29
+ prev = wgScan[tid - step];
30
+ }
31
+ workgroupBarrier();
32
+ wgScan[tid] = wgScan[tid] + prev;
33
+ }
34
+ workgroupBarrier();
35
+ if (j < params.numBlocks) {
36
+ var value = carry;
37
+ if (tid > 0u) {
38
+ value = carry + wgScan[tid - 1u];
39
+ }
40
+ blockPrefix[j] = value;
41
+ }
42
+ carry = carry + wgScan[WG - 1u];
43
+ }
44
+ }
build/webgpu/test.json ADDED
@@ -0,0 +1,470 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "op": "ai.onnx.NonZero",
3
+ "cases": [
4
+ {
5
+ "name": "int16_scalar_boundaries",
6
+ "inputs": {
7
+ "x": { "dtype": "int16", "shape": [4], "data": { "kind": "values", "values": [-32768, 0, 32767, -1] } }
8
+ },
9
+ "outputs": {
10
+ "y": { "dtype": "uint32", "shape": [1, 3], "tolerance": 0, "data": { "kind": "values", "values": [0, 2, 3] } }
11
+ }
12
+ },
13
+ {
14
+ "name": "rank2_exact_capacity_f32",
15
+ "inputs": {
16
+ "x": {
17
+ "dtype": "float32",
18
+ "shape": [2, 3],
19
+ "data": { "kind": "values", "values": [0.0, 1.0, 0.0, -2.0, 3.0, 0.0] }
20
+ }
21
+ },
22
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 3] } }
23
+ },
24
+ {
25
+ "name": "f32_subnormal_nonzero_coordinates",
26
+ "provenance": {
27
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
28
+ "test": "NonZeroOpTest.BasicNumeric",
29
+ "notes": "Subnormal finite values are nonzero under ONNX equality semantics; flushing them to zero changes the emitted coordinates."
30
+ },
31
+ "inputs": {
32
+ "x": {
33
+ "dtype": "float32",
34
+ "shape": [2, 3],
35
+ "data": { "kind": "values", "values": [0.0, 1e-40, -1e-40, 0.0, 1e-39, 0.0] }
36
+ }
37
+ },
38
+ "outputs": {
39
+ "y": {
40
+ "dtype": "uint32",
41
+ "shape": [2, 3],
42
+ "tolerance": 0,
43
+ "data": { "kind": "values", "values": [0, 0, 1, 1, 2, 1] }
44
+ }
45
+ }
46
+ },
47
+ {
48
+ "name": "f32_scalar_subnormal_nonzero_coordinates",
49
+ "provenance": {
50
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
51
+ "test": "NonZeroOpTest.Scalar",
52
+ "notes": "A finite subnormal rank-0 input is still nonzero. ONNX represents its one hit with shape [0, 1], containing no coordinate values."
53
+ },
54
+ "inputs": { "x": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [1e-40] } } },
55
+ "outputs": {
56
+ "y": { "dtype": "uint32", "shape": [0, 1], "tolerance": 0, "data": { "kind": "values", "values": [] } }
57
+ }
58
+ },
59
+ {
60
+ "name": "f32_parallel_subnormal_nonzero_coordinates",
61
+ "provenance": {
62
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
63
+ "test": "NonZeroOpTest.BasicNumeric",
64
+ "notes": "Large enough to select the parallel-scan path; every even column is a finite subnormal and must produce an output coordinate."
65
+ },
66
+ "inputs": { "x": { "dtype": "float32", "shape": [1, 1026], "data": { "kind": "cycle", "values": [1e-40, 0.0] } } },
67
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 513], "tolerance": 0 } }
68
+ },
69
+ {
70
+ "name": "rank3_exact_capacity_uint8",
71
+ "inputs": {
72
+ "x": { "dtype": "uint8", "shape": [2, 2, 2], "data": { "kind": "values", "values": [1, 0, 2, 3, 0, 4, 5, 6] } }
73
+ },
74
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 6] } }
75
+ },
76
+ {
77
+ "name": "rank1_no_hits_int8",
78
+ "inputs": { "x": { "dtype": "int8", "shape": [4], "data": { "kind": "values", "values": [0, 0, 0, 0] } } },
79
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 0] } }
80
+ },
81
+ {
82
+ "name": "rank1_nonzero_int8",
83
+ "inputs": { "x": { "dtype": "int8", "shape": [4], "data": { "kind": "values", "values": [-2, 0, 7, 0] } } },
84
+ "outputs": {
85
+ "y": { "dtype": "uint32", "shape": [1, 2], "tolerance": 0, "data": { "kind": "values", "values": [0, 2] } }
86
+ },
87
+ "provenance": {
88
+ "notes": "Synthetic compact non-empty int8 witness for the standard numeric NonZero route; the existing int8 zero-hit case does not dispatch the dtype-specific shader path."
89
+ }
90
+ },
91
+ {
92
+ "name": "rank4_exact_capacity_f32",
93
+ "inputs": {
94
+ "x": {
95
+ "dtype": "float32",
96
+ "shape": [1, 2, 2, 3],
97
+ "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 0.0, 0.0, -3.0, 4.0, 0.0, 5.0, 6.0, 0.0, 7.0] }
98
+ }
99
+ },
100
+ "outputs": { "y": { "dtype": "uint32", "shape": [4, 7] } }
101
+ },
102
+ {
103
+ "name": "ort_basic_numeric_rank3_exact_capacity_int32",
104
+ "provenance": {
105
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
106
+ "test": "NonZeroOpTest.BasicNumeric",
107
+ "notes": "The runtime represents ONNX int64 coordinate metadata as uint32."
108
+ },
109
+ "inputs": {
110
+ "x": { "dtype": "int32", "shape": [1, 2, 3], "data": { "kind": "values", "values": [0, 1, 2, 0, 3, 4] } }
111
+ },
112
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 4], "tolerance": 0 } }
113
+ },
114
+ {
115
+ "name": "ort_basic_numeric_rank3_exact_capacity_float",
116
+ "provenance": {
117
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
118
+ "test": "NonZeroOpTest.BasicNumeric",
119
+ "notes": "Float variant from ORT’s templated BasicNumeric case; the runtime represents ONNX int64 coordinate metadata as uint32."
120
+ },
121
+ "inputs": {
122
+ "x": {
123
+ "dtype": "float32",
124
+ "shape": [1, 2, 3],
125
+ "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 0.0, 3.0, 4.0] }
126
+ }
127
+ },
128
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 4], "tolerance": 0 } }
129
+ },
130
+ {
131
+ "name": "ort_three_dims_rank3_exact_capacity_int32",
132
+ "provenance": {
133
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
134
+ "test": "NonZeroOpTest.ThreeDims",
135
+ "notes": "Integer-width adapted to int32; the runtime represents ONNX int64 coordinate metadata as uint32."
136
+ },
137
+ "inputs": {
138
+ "x": { "dtype": "int32", "shape": [2, 2, 2], "data": { "kind": "values", "values": [0, 1, 1, 0, 1, 0, 1, 0] } }
139
+ },
140
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 4], "tolerance": 0 } }
141
+ },
142
+ {
143
+ "name": "ort_empty_input_rank3_zero_hits",
144
+ "provenance": {
145
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
146
+ "test": "NonZeroOpTest.EmptyInput",
147
+ "notes": "Empty-input coverage with the exact ONNX output shape [rank(X), 0]."
148
+ },
149
+ "inputs": { "x": { "dtype": "int32", "shape": [1, 0, 2], "data": { "kind": "values", "values": [] } } },
150
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 0], "tolerance": 0 } }
151
+ },
152
+ {
153
+ "name": "ort_basic_bool_adapted_int32",
154
+ "provenance": {
155
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
156
+ "test": "NonZeroOpTest.BasicBool",
157
+ "notes": "Integer twin of the boolean pattern: a [2,3] int32 0/1 tensor pins row-major coordinates for its two nonzero elements and the uint32 metadata output."
158
+ },
159
+ "inputs": {
160
+ "x": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [1, 0, 0, 0, 0, 1] } }
161
+ },
162
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 2], "tolerance": 0 } }
163
+ },
164
+ {
165
+ "name": "ort_basic_bool_exact",
166
+ "provenance": {
167
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
168
+ "test": "NonZeroOpTest.BasicBool"
169
+ },
170
+ "inputs": {
171
+ "x": { "dtype": "bool", "shape": [2, 3], "data": { "kind": "values", "values": [1, 0, 0, 0, 0, 1] } }
172
+ },
173
+ "outputs": {
174
+ "y": {
175
+ "dtype": "uint32",
176
+ "shape": [2, 2],
177
+ "tolerance": 0,
178
+ "data": { "kind": "values", "values": [0, 1, 0, 2] }
179
+ }
180
+ }
181
+ },
182
+ {
183
+ "name": "ort_scalar_zero_int32",
184
+ "provenance": {
185
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
186
+ "test": "NonZeroOpTest.Scalar",
187
+ "notes": "Pins the ONNX scalar-zero output shape [0, 0]; the runtime represents ONNX int64 coordinate metadata as uint32."
188
+ },
189
+ "inputs": { "x": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [0] } } },
190
+ "outputs": { "y": { "dtype": "uint32", "shape": [0, 0], "tolerance": 0 } }
191
+ },
192
+ {
193
+ "name": "ort_scalar_nonzero_int32",
194
+ "provenance": {
195
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
196
+ "test": "NonZeroOpTest.Scalar",
197
+ "notes": "Pins the ONNX scalar-nonzero output shape [0, 1]; the runtime represents ONNX int64 coordinate metadata as uint32."
198
+ },
199
+ "inputs": { "x": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [1] } } },
200
+ "outputs": { "y": { "dtype": "uint32", "shape": [0, 1], "tolerance": 0 } }
201
+ },
202
+ {
203
+ "name": "bool_scalar_nonzero",
204
+ "provenance": {
205
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
206
+ "test": "NonZeroOpTest.Scalar",
207
+ "notes": "Scalar NonZero shape convention from ORT's scalar case, using ONNX-valid bool input coverage from NonZeroOpTest.BasicBool."
208
+ },
209
+ "inputs": { "x": { "dtype": "bool", "shape": [], "data": { "kind": "values", "values": [1] } } },
210
+ "outputs": {
211
+ "y": { "dtype": "uint32", "shape": [0, 1], "tolerance": 0, "data": { "kind": "values", "values": [] } }
212
+ }
213
+ },
214
+ {
215
+ "name": "onnx_backend_nonzero_example",
216
+ "inputs": { "x": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [1, 0, 1, 1] } } },
217
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 3] } },
218
+ "provenance": {
219
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_nonzero_example",
220
+ "notes": "Official ONNX example with its exact dynamic output shape. The runtime represents ONNX int64 coordinate metadata as uint32."
221
+ }
222
+ },
223
+ {
224
+ "name": "float32_signed_zero_nan_and_infinity",
225
+ "provenance": {
226
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
227
+ "test": "NonZeroOpTest.BasicNumeric",
228
+ "notes": "Additional numeric edge: +0 and -0 are zero, while NaN and infinities compare nonzero."
229
+ },
230
+ "inputs": {
231
+ "x": {
232
+ "dtype": "float32",
233
+ "shape": [2, 3],
234
+ "data": { "kind": "values", "values": [0.0, 0.0, "NaN", "Infinity", "-Infinity", 1.0] }
235
+ }
236
+ },
237
+ "outputs": {
238
+ "y": {
239
+ "dtype": "uint32",
240
+ "shape": [2, 4],
241
+ "tolerance": 0,
242
+ "data": { "kind": "values", "values": [0, 1, 1, 1, 2, 0, 1, 2] }
243
+ }
244
+ }
245
+ },
246
+ {
247
+ "name": "float16_signed_zero_nan_exact_capacity",
248
+ "provenance": {
249
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
250
+ "test": "NonZeroOpTest.BasicNumeric",
251
+ "notes": "Additional numeric edge: -0 is zero, while NaN and the finite nonzero value produce the two exact output coordinates."
252
+ },
253
+ "inputs": {
254
+ "x": { "dtype": "float16", "shape": [4], "data": { "kind": "values", "values": [0.0, "NaN", 0.0, -2.0] } }
255
+ },
256
+ "outputs": {
257
+ "y": { "dtype": "uint32", "shape": [1, 2], "tolerance": 0, "data": { "kind": "values", "values": [1, 3] } }
258
+ }
259
+ },
260
+ {
261
+ "name": "parallel_rank2_exact_capacity_f32",
262
+ "provenance": {
263
+ "notes": "65536-element rank-2 input with 25% nonzeros at exact capacity; exercises the parallel_scan 3-pass coordinate-emission variant."
264
+ },
265
+ "inputs": {
266
+ "x": {
267
+ "dtype": "float32",
268
+ "shape": [256, 256],
269
+ "data": { "kind": "cycle", "values": [0.0, 0.0, 2.5, 0.0, -1.25, 0.0, 0.0, 0.0] }
270
+ }
271
+ },
272
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 16384], "tolerance": 0 } }
273
+ },
274
+ {
275
+ "name": "parallel_rank1_f32_subnormal_nonzero",
276
+ "provenance": {
277
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
278
+ "test": "NonZeroOpTest.BasicNumeric",
279
+ "notes": "Parallel-scan companion to the serial subnormal case: finite subnormal float32 values are nonzero and must survive the shared predicate scan."
280
+ },
281
+ "inputs": {
282
+ "x": { "dtype": "float32", "shape": [4096], "data": { "kind": "cycle", "values": [0.0, 1e-40, 0.0, -1e-40] } }
283
+ },
284
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 2048], "tolerance": 0 } }
285
+ },
286
+ {
287
+ "name": "parallel_rank1_exact_capacity_int32",
288
+ "provenance": {
289
+ "notes": "Exactly 10000 nonzeros; the parallel scatter must emit every coordinate into the caller-provided exact output shape."
290
+ },
291
+ "inputs": { "x": { "dtype": "int32", "shape": [50000], "data": { "kind": "cycle", "values": [0, 3, 0, 0, 0] } } },
292
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 10000], "tolerance": 0 } }
293
+ },
294
+ {
295
+ "name": "parallel_rank3_u32",
296
+ "provenance": {
297
+ "notes": "Rank-3 uint32 input; render-time stride unrolling must emit three coordinate rows at the scanned columns."
298
+ },
299
+ "inputs": {
300
+ "x": { "dtype": "uint32", "shape": [32, 60, 8], "data": { "kind": "cycle", "values": [7, 0, 0, 0, 1, 0] } }
301
+ },
302
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 5120], "tolerance": 0 } }
303
+ },
304
+ {
305
+ "name": "parallel_rank1_f16",
306
+ "provenance": {
307
+ "notes": "f16 input through the parallel predicate scan (f32-widened zero test, same as the serial kernel)."
308
+ },
309
+ "inputs": {
310
+ "x": { "dtype": "float16", "shape": [4096], "data": { "kind": "cycle", "values": [0.0, 1.5, 0.0, 0.0] } }
311
+ },
312
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1024], "tolerance": 0 } }
313
+ },
314
+ {
315
+ "name": "f16_subnormal_nonzero_coordinates",
316
+ "provenance": {
317
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
318
+ "test": "NonZeroOpTest.BasicNumeric",
319
+ "notes": "The f16 predicate widens values to f32 before comparing with zero. The three finite f16 subnormals must therefore produce coordinates (0,1), (0,2), and (1,1), proving that widening preserves their nonzero values."
320
+ },
321
+ "inputs": {
322
+ "x": {
323
+ "dtype": "float16",
324
+ "shape": [2, 3],
325
+ "data": { "kind": "values", "values": [0.0, 6e-8, -6e-8, 0.0, 0.00006, 0.0] }
326
+ }
327
+ },
328
+ "outputs": {
329
+ "y": {
330
+ "dtype": "uint32",
331
+ "shape": [2, 3],
332
+ "tolerance": 0,
333
+ "data": { "kind": "values", "values": [0, 0, 1, 1, 2, 1] }
334
+ }
335
+ }
336
+ },
337
+ {
338
+ "name": "parallel_f16_subnormal_nonzero_coordinates",
339
+ "provenance": {
340
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
341
+ "test": "NonZeroOpTest.BasicNumeric",
342
+ "notes": "Parallel companion to the serial f16-subnormal case. Every even element rounds to the finite f16 subnormal 0x0001 and must survive the shared widened-f32 predicate, producing exactly 1024 coordinate columns."
343
+ },
344
+ "inputs": { "x": { "dtype": "float16", "shape": [2048], "data": { "kind": "cycle", "values": [6e-8, 0.0] } } },
345
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1024], "tolerance": 0 } }
346
+ },
347
+ {
348
+ "name": "parallel_rank1_bool_mask",
349
+ "provenance": {
350
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
351
+ "test": "NonZeroOpTest.BasicBool",
352
+ "notes": "Clean coverage: a sequence-scale boolean mask (the canonical real-model NonZero producer) through the parallel_scan path (numel>1024). bool maps to the u32 scalar with predicate `src[i] != 0u`; only the serial variant carried bool coverage before. ORT CPU NonZero(bool) and the TS reference agree on 1024 hits at the even columns."
353
+ },
354
+ "inputs": { "x": { "dtype": "bool", "shape": [2048], "data": { "kind": "cycle", "values": [1, 0] } } },
355
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1024], "tolerance": 0 } }
356
+ },
357
+ {
358
+ "name": "parallel_rank2_uint8",
359
+ "provenance": {
360
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
361
+ "test": "NonZeroOpTest.BasicNumeric",
362
+ "notes": "Clean coverage: rank-2 uint8 input through the parallel_scan path (numel>1024); uint8 maps to the u32 scalar with predicate `src[i] != 0u`. uint8 NonZero was only exercised by the tiny serial case before. ORT CPU NonZero(uint8) and the TS reference agree (every third element nonzero)."
363
+ },
364
+ "inputs": { "x": { "dtype": "uint8", "shape": [30, 50], "data": { "kind": "cycle", "values": [7, 0, 0] } } },
365
+ "outputs": { "y": { "dtype": "uint32", "shape": [2, 500], "tolerance": 0 } }
366
+ },
367
+ {
368
+ "name": "parallel_f32_signed_zero_nan_infinity",
369
+ "provenance": {
370
+ "source": "onnxruntime/test/providers/cpu/tensor/nonzero_op_test.cc",
371
+ "test": "NonZeroOpTest.BasicNumeric",
372
+ "notes": "Parallel companion to float32_signed_zero_nan_and_infinity. The flag scan bit-tests `(bitcast<u32>(src[i]) & 0x7fffffffu) != 0u`, so +0/-0 are zero while NaN and +/-Inf are nonzero; the six-value cycle produces exactly 1024 coordinate columns."
373
+ },
374
+ "inputs": {
375
+ "x": {
376
+ "dtype": "float32",
377
+ "shape": [1536],
378
+ "data": { "kind": "cycle", "values": [0.0, 0.0, "NaN", "Infinity", "-Infinity", 1.0] }
379
+ }
380
+ },
381
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 1024], "tolerance": 0 } }
382
+ },
383
+ {
384
+ "name": "parallel_rank6_max_rank_stride_unroll",
385
+ "provenance": {
386
+ "notes": "Six-dimensional int32 input through the parallel path. This exercises six compile-time coordinate strides and emits six coordinate rows; the [3,0,0,0] cycle produces 512 deterministic columns in row-major order."
387
+ },
388
+ "inputs": {
389
+ "x": { "dtype": "int32", "shape": [2, 2, 2, 2, 2, 64], "data": { "kind": "cycle", "values": [3, 0, 0, 0] } }
390
+ },
391
+ "outputs": { "y": { "dtype": "uint32", "shape": [6, 512], "tolerance": 0 } }
392
+ },
393
+ {
394
+ "name": "parallel_rank4_vec4_tail_scan_scatter",
395
+ "provenance": {
396
+ "notes": "Compact route lock for the rank-4 activation-map benchmark. 1536 elements cross one 1024-element vectorized block boundary; the 50% mask validates stable flattened order, rank-4 coordinate decomposition, and the partial second block."
397
+ },
398
+ "inputs": { "x": { "dtype": "uint32", "shape": [2, 3, 8, 32], "data": { "kind": "cycle", "values": [0, 7] } } },
399
+ "outputs": { "y": { "dtype": "uint32", "shape": [4, 768], "tolerance": 0 } }
400
+ },
401
+ {
402
+ "name": "parallel_rank1_exact_crossover_1024",
403
+ "provenance": {
404
+ "notes": "Pins the first profitable parallel-scan size. Exactly 1024 elements form one complete vectorized scan block; the alternating mask validates stable coordinates and an exact 512-column output at the serial/parallel selector boundary."
405
+ },
406
+ "inputs": { "x": { "dtype": "float32", "shape": [1024], "data": { "kind": "cycle", "values": [0.0, 3.0] } } },
407
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 512], "tolerance": 0 } }
408
+ },
409
+ {
410
+ "name": "rank7_exact_capacity_coordinates",
411
+ "inputs": {
412
+ "x": {
413
+ "dtype": "float32",
414
+ "shape": [1, 2, 1, 2, 1, 2, 3],
415
+ "data": { "kind": "cycle", "values": [0.0, 1.0, 0.0, -2.0] }
416
+ }
417
+ },
418
+ "outputs": { "y": { "dtype": "uint32", "shape": [7, 12], "tolerance": 0 } }
419
+ },
420
+ {
421
+ "name": "parallel_sparse_exact_capacity_2048",
422
+ "provenance": {
423
+ "notes": "Sparse parallel-route fixture with 512 deterministic hits and an exact [1, 512] output shape."
424
+ },
425
+ "inputs": {
426
+ "x": { "dtype": "float32", "shape": [2048], "data": { "kind": "cycle", "values": [0.0, 3.0, 0.0, 0.0] } }
427
+ },
428
+ "outputs": { "y": { "dtype": "uint32", "shape": [1, 512], "tolerance": 0 } }
429
+ },
430
+ {
431
+ "name": "rank8_exact_capacity_coordinates",
432
+ "inputs": {
433
+ "x": {
434
+ "dtype": "float32",
435
+ "shape": [1, 2, 1, 2, 1, 2, 2, 3],
436
+ "data": { "kind": "cycle", "values": [0.0, 1.0, 0.0, -2.0] }
437
+ }
438
+ },
439
+ "outputs": { "y": { "dtype": "uint32", "shape": [8, 24], "tolerance": 0 } }
440
+ },
441
+ {
442
+ "name": "parallel_rank4_wide_items_scan",
443
+ "provenance": {
444
+ "notes": "Large enough (524288 elements, 512 scan blocks) to take the wide items-per-thread branch, which for rank >= 4 resolves to 4. Every other parallel fixture is small enough to fall to the narrow 16-item branch, so without this the 4-item scan and scatter were exercised only by benchmarks, which do not check output."
445
+ },
446
+ "inputs": {
447
+ "x": {
448
+ "dtype": "uint32",
449
+ "shape": [4, 8, 128, 128],
450
+ "data": { "kind": "cycle", "values": [0, 0, 0, 0, 0, 0, 0, 7] }
451
+ }
452
+ },
453
+ "outputs": { "y": { "dtype": "uint32", "shape": [4, 65536], "tolerance": 0 } }
454
+ },
455
+ {
456
+ "name": "parallel_rank3_wide_items_scan",
457
+ "provenance": {
458
+ "notes": "Rank-3 twin of parallel_rank4_wide_items_scan: 1048576 elements is 512 scan blocks at the rank-3 wide value of 8 items per thread, the other items-per-thread value no smaller fixture reaches."
459
+ },
460
+ "inputs": {
461
+ "x": {
462
+ "dtype": "uint32",
463
+ "shape": [8, 128, 1024],
464
+ "data": { "kind": "cycle", "values": [0, 0, 0, 0, 0, 0, 0, 7] }
465
+ }
466
+ },
467
+ "outputs": { "y": { "dtype": "uint32", "shape": [3, 131072], "tolerance": 0 } }
468
+ }
469
+ ]
470
+ }