Xenova HF Staff commited on
Commit
31ec3cf
·
verified ·
1 Parent(s): eefaea5

sync 2e7068faf55e

Browse files
README.md CHANGED
@@ -1,3 +1,89 @@
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.ReduceLogSum
10
+
11
+ `ai.onnx` · standard ONNX operator · ONNX opset ≥ 18
12
+
13
+ ## Description
14
+
15
+ Computes `log(sum(elements))` of the input tensor along the specified axes. The output rank matches the input when `keepdims` is 1; reduced dimensions are pruned when `keepdims` is 0. Reducing an empty set of values yields negative infinity.
16
+
17
+ See the [ONNX `ReduceLogSum` spec](https://onnx.ai/onnx/operators/onnx__ReduceLogSum.html) for the reference semantics.
18
+
19
+ ## Inputs
20
+
21
+ | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
+ | --- | --- | --- | --- | --- | --- | --- |
23
+ | `data` | `x` | `T` | — | — | Input tensor to reduce. | required |
24
+
25
+ ## Outputs
26
+
27
+ | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
28
+ | --- | --- | --- | --- | --- | --- | --- |
29
+ | `reduced` | `y` | `T` | derived | — | Reduced output tensor containing the log of the summed elements. | required |
30
+
31
+ ## Attributes
32
+
33
+ Default values (overridable per request):
34
+
35
+ | Attribute | Default | Description |
36
+ | --- | --- | --- |
37
+ | `keepdims` | `1` | If 1 (default in spec), the reduced dimension is retained with size 1; if 0, it is removed from the output shape. |
38
+ | `noop_with_empty_axes` | `0` | When 1 and axes is empty, acts as a no-op applying only the non-reduction step (log of input); when 0 (default), reduces over all axes. |
39
+ | `axes` | `[]` | Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`. |
40
+
41
+ ## Type constraints
42
+
43
+ | Variable | Allowed dtypes |
44
+ | --- | --- |
45
+ | `T` | `float32`, `float16` |
46
+
47
+ ## Device requirements
48
+
49
+ Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
50
+
51
+ ## Files
52
+
53
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
54
+ - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
55
+ - [`test.json`](build/webgpu/test.json) — correctness cases
56
+ - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
57
+ - [`reduce-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-axis-split-reduce.wgsl.jinja)
58
+ - [`reduce-axis0-splitk-combine.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja)
59
+ - [`reduce-axis0-splitk-reduce.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja)
60
+ - [`reduce-axis0-tilecols.wgsl.jinja`](build/webgpu/reduce-axis0-tilecols.wgsl.jinja)
61
+ - [`reduce-flat-partial.wgsl.jinja`](build/webgpu/reduce-flat-partial.wgsl.jinja)
62
+ - [`reduce-noop-empty-axes.wgsl.jinja`](build/webgpu/reduce-noop-empty-axes.wgsl.jinja)
63
+ - [`reduce-row-subgroup.wgsl.jinja`](build/webgpu/reduce-row-subgroup.wgsl.jinja)
64
+ - [`reduce-row-tree.wgsl.jinja`](build/webgpu/reduce-row-tree.wgsl.jinja)
65
+ - [`reduce-serial-axis.wgsl.jinja`](build/webgpu/reduce-serial-axis.wgsl.jinja)
66
+
67
+ ## Use with `@huggingface/kernels`
68
+
69
+ The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
70
+
71
+ The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
72
+
73
+ - `y`
74
+
75
+ Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
76
+
77
+ The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
78
+
79
+ Replace each `*Data` placeholder with a typed array containing the corresponding input data.
80
+
81
+ ```js
82
+ import { getKernel } from "@huggingface/kernels";
83
+
84
+ const kernel = await getKernel("webgpu-kernels/ai.onnx.ReduceLogSum", { version: 1 });
85
+ // Explicit destinations request optional results or supply metadata that cannot be inferred.
86
+ const { y } = await kernel({ x: { data: xData, shape: [] } }, {
87
+ outputs: { y: { shape: [], dtype: "float32" } },
88
+ });
89
+ ```
build/webgpu/bench.json ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "op": "ai.onnx.ReduceLogSum",
3
+ "cases": [
4
+ {
5
+ "name": "1024x1024_axis1",
6
+ "attrs": { "axes": [1], "keepdims": 0 },
7
+ "inputs": { "x": { "dtype": "float32", "shape": [1024, 1024], "data": { "kind": "constant", "value": 0.5 } } },
8
+ "outputs": { "y": { "dtype": "float32", "shape": [1024] } }
9
+ },
10
+ {
11
+ "name": "reducelogsum-axis0-f32-1024x512",
12
+ "preset": "smoke",
13
+ "vars": { "rows": 1024, "cols": 512 },
14
+ "attrs": { "axes": [0], "keepdims": 0 },
15
+ "inputs": {
16
+ "x": { "shape": [1024, 512], "dtype": "float32", "dist": "normal", "seed": 118, "scale": 0.2, "offset": 1 }
17
+ },
18
+ "outputs": { "y": { "shape": [512], "dtype": "float32" } },
19
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
20
+ },
21
+ {
22
+ "name": "reducelogsum-lastaxis-f32-4096x4096",
23
+ "preset": "smoke",
24
+ "vars": { "rows": 4096, "cols": 4096 },
25
+ "attrs": { "axes": [1], "keepdims": 0 },
26
+ "inputs": { "x": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 117, "scale": 0.2 } },
27
+ "outputs": { "y": { "shape": [4096], "dtype": "float32" } },
28
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
29
+ },
30
+ {
31
+ "name": "reducelogsum-fullreduce-r3-256x256x256-keepdims",
32
+ "preset": "smoke",
33
+ "vars": { "rows": 65536, "cols": 256 },
34
+ "attrs": { "keepdims": 1 },
35
+ "inputs": {
36
+ "x": { "shape": [256, 256, 256], "dtype": "float32", "dist": "uniform", "seed": 711, "scale": 2, "offset": 0.1 }
37
+ },
38
+ "outputs": { "y": { "shape": [1, 1, 1], "dtype": "float32" } },
39
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
40
+ },
41
+ {
42
+ "name": "reducelogsum-fullreduce-serial-numel-not-mul4-r3",
43
+ "preset": "stress",
44
+ "attrs": { "keepdims": 0 },
45
+ "inputs": {
46
+ "x": { "shape": [1001, 1001, 1], "dtype": "float32", "dist": "uniform", "seed": 501, "scale": 2, "offset": 0.1 }
47
+ },
48
+ "outputs": { "y": { "shape": [], "dtype": "float32" } },
49
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "1001 * 1001 * 1 * 4" }] }
50
+ },
51
+ {
52
+ "name": "reducelogsum-rank3-spatial-axes12-f32-128x256x256-pathology",
53
+ "preset": "stress",
54
+ "provenance": {
55
+ "source": "authored for branch coverage",
56
+ "notes": "Rank-3 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane. Positive inputs keep the post-reduction logarithm finite."
57
+ },
58
+ "attrs": { "axes": [1, 2], "keepdims": 1 },
59
+ "inputs": {
60
+ "x": {
61
+ "shape": [128, 256, 256],
62
+ "dtype": "float32",
63
+ "dist": "uniform",
64
+ "seed": 305,
65
+ "scale": 1,
66
+ "offset": 0.01
67
+ }
68
+ },
69
+ "outputs": { "y": { "shape": [128, 1, 1], "dtype": "float32", "dist": "empty" } },
70
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "128 * 256 * 256 * 4" }] }
71
+ },
72
+ {
73
+ "name": "reducelogsum-spatial-axes23-f32-2x64x256x256-pathology",
74
+ "preset": "stress",
75
+ "provenance": {
76
+ "source": "authored for branch coverage",
77
+ "notes": "Rank-4 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane. Positive inputs keep the post-reduction logarithm finite."
78
+ },
79
+ "attrs": { "axes": [2, 3], "keepdims": 1 },
80
+ "inputs": {
81
+ "x": {
82
+ "shape": [2, 64, 256, 256],
83
+ "dtype": "float32",
84
+ "dist": "uniform",
85
+ "seed": 304,
86
+ "scale": 1,
87
+ "offset": 0.01
88
+ }
89
+ },
90
+ "outputs": { "y": { "shape": [2, 64, 1, 1], "dtype": "float32", "dist": "empty" } },
91
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "2 * 64 * 256 * 256 * 4" }] }
92
+ },
93
+ {
94
+ "name": "reducelogsum-axis0-f32-131072x15-narrow-splitk-boundary",
95
+ "preset": "stress",
96
+ "vars": { "rows": 131072, "cols": 15 },
97
+ "attrs": { "axes": [0], "keepdims": 0 },
98
+ "inputs": {
99
+ "x": { "shape": [131072, 15], "dtype": "float32", "dist": "uniform", "seed": 502, "scale": 0.1, "offset": 1 }
100
+ },
101
+ "outputs": { "y": { "shape": [15], "dtype": "float32", "dist": "empty" } },
102
+ "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * 4" }] }
103
+ }
104
+ ],
105
+ "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }
106
+ }
build/webgpu/manifest.json ADDED
@@ -0,0 +1,1004 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "domain": "ai.onnx",
3
+ "name": "ReduceLogSum",
4
+ "sinceVersion": 18,
5
+ "description": "Computes `log(sum(elements))` of the input tensor along the specified axes. The output rank matches the input when `keepdims` is 1; reduced dimensions are pruned when `keepdims` is 0. Reducing an empty set of values yields negative infinity.",
6
+ "inputs": [{ "role": "data", "dtype": "T", "description": "Input tensor to reduce." }],
7
+ "outputs": [
8
+ {
9
+ "role": "reduced",
10
+ "dtype": "T",
11
+ "rank": "ranks.data if attrs.keepdims == 1 or ((attrs.axes | length) == 0 and attrs.noop_with_empty_axes == 1) else (ranks.data - (attrs.axes | length) if (attrs.axes | length) > 0 else 0)",
12
+ "description": "Reduced output tensor containing the log of the summed elements."
13
+ }
14
+ ],
15
+ "attributes": { "keepdims": 1, "noop_with_empty_axes": 0, "axes": [] },
16
+ "attributeDescriptions": {
17
+ "keepdims": "If 1 (default in spec), the reduced dimension is retained with size 1; if 0, it is removed from the output shape.",
18
+ "noop_with_empty_axes": "When 1 and axes is empty, acts as a no-op applying only the non-reduction step (log of input); when 0 (default), reduces over all axes.",
19
+ "axes": "Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`."
20
+ },
21
+ "attributeConstraints": { "keepdims": { "values": [0, 1] }, "noop_with_empty_axes": { "values": [0, 1] } },
22
+ "typeConstraints": { "T": ["float32", "float16"] },
23
+ "args": {
24
+ "x": { "kind": "tensor", "semantic": "data", "role": "input" },
25
+ "y": { "kind": "tensor", "semantic": "reduced", "role": "output" }
26
+ },
27
+ "derive": {
28
+ "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
29
+ "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
30
+ "treeWorkgroupOk": "reduceWorkgroupSize > 0 and pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize and reduceWorkgroupSize * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize",
31
+ "subgroupWorkgroupFloor": "min(reduceWorkgroupSize, max(1, device.adapterInfo.subgroupMaxSize))",
32
+ "lastAxisRows": "rows(shapes.data, ranks.data - 1) if ranks.data > 0 else 1",
33
+ "lastAxisCols": "dim(shapes.data, ranks.data - 1) if ranks.data > 0 else 1",
34
+ "rowSerialPreferred": "lastAxisRows >= tunables.ROW_SERIAL_MIN_ROWS and lastAxisCols <= tunables.ROW_SERIAL_MAX_COLS",
35
+ "axis0Rows": "dim(shapes.data, 0) if ranks.data >= 2 else 0",
36
+ "axis0Cols": "dim(shapes.data, 1) if ranks.data >= 2 else 0",
37
+ "axis0SplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axis0Rows, tunables.AXIS0_SPLIT_TARGET_ROWS)))",
38
+ "axis0SplitScratchBytes": "axis0SplitCount * axis0Cols * dtypeBytes(\"float32\")",
39
+ "axis0SplitPathFits": "axis0SplitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(axis0Cols, reduceWorkgroupSize), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension and axis0SplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axis0SplitScratchBytes <= device.limits.maxBufferSize",
40
+ "reduceAxis": "(attrs.axes[0] + ranks.data if attrs.axes[0] < 0 else attrs.axes[0]) if ((attrs.axes | length) == 1 and isUniqueIntList(attrs.axes, 0 - ranks.data, ranks.data, 1)) else ranks.data",
41
+ "axisSplitDim": "dim(shapes.data, reduceAxis) if ranks.data >= 2 and reduceAxis < ranks.data else 0",
42
+ "axisSplitInner": "inner(shapes.data, reduceAxis) if ranks.data >= 2 and reduceAxis < ranks.data else 1",
43
+ "axisSplitOutputs": "numel(shapes.reduced)",
44
+ "axisSplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axisSplitDim, tunables.AXIS0_SPLIT_TARGET_ROWS)))",
45
+ "axisSplitScratchBytes": "axisSplitCount * axisSplitOutputs * 4",
46
+ "axisSplitPathFits": "axisSplitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(axisSplitOutputs, reduceWorkgroupSize), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension and axisSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axisSplitScratchBytes <= device.limits.maxBufferSize",
47
+ "axis0TilePathFits": "treeWorkgroupOk and tunables.AXIS0_TILE_COLS > 0 and tunables.AXIS0_TILE_COLS <= reduceWorkgroupSize and reduceWorkgroupSize % tunables.AXIS0_TILE_COLS == 0",
48
+ "flatItems": "floor(numel(shapes.data) / tunables.VECTOR_WIDTH)",
49
+ "flatSplitCount": "max(1, min(tunables.FULL_REDUCE_MAX_SPLITS, ceilDiv(flatItems, reduceWorkgroupSize)))",
50
+ "flatScratchBytes": "flatSplitCount * dtypeBytes(\"float32\")",
51
+ "flatPathFits": "treeWorkgroupOk and flatSplitCount <= device.limits.maxComputeWorkgroupsPerDimension and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize",
52
+ "flatParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.reduced) == 1 and numel(shapes.data) >= tunables.FULL_REDUCE_MIN_ELEMENTS and flatPathFits",
53
+ "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.reduced) > 0 and numel(shapes.data) % numel(shapes.reduced) == 0 and numel(shapes.data) / numel(shapes.reduced) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.data == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.reduced) == dim(shapes.data, 0)) or (ranks.data == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) == false and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.reduced) == dim(shapes.data, 0) * dim(shapes.data, 1)) or (ranks.data == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.reduced) == dim(shapes.data, 0)))"
54
+ },
55
+ "tunables": {
56
+ "WORKGROUP_SIZE": 256,
57
+ "VECTOR_WIDTH": 4,
58
+ "ROW_PARALLEL_MIN_COLS": 64,
59
+ "SUBGROUP_MIN_COLS": 256,
60
+ "SUBGROUP_SMALL_ROW_LIMIT": 32768,
61
+ "AXIS0_SPLIT_MIN_ROWS": 8192,
62
+ "AXIS0_SPLIT_TARGET_ROWS": 256,
63
+ "AXIS0_MAX_SPLITS": 128,
64
+ "AXIS0_TILE_MIN_ROWS": 64,
65
+ "AXIS0_TILE_MIN_COLS": 16,
66
+ "AXIS0_TILE_COLS": 16,
67
+ "AXIS_SPLIT_TILE_COLS": 8,
68
+ "FULL_REDUCE_MIN_ELEMENTS": 8192,
69
+ "FULL_REDUCE_MAX_SPLITS": 256,
70
+ "CONTIGUOUS_SUFFIX_MIN_COLS": 256,
71
+ "ROW_SERIAL_MIN_ROWS": 8192,
72
+ "ROW_SERIAL_MAX_COLS": 1024
73
+ },
74
+ "bindingSets": {
75
+ "elementwise": [
76
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
77
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
78
+ {
79
+ "name": "params",
80
+ "semantic": "kernel.params",
81
+ "buffer": { "type": "uniform" },
82
+ "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.reduced)" }] }
83
+ }
84
+ ],
85
+ "lastAxisVec4": [
86
+ {
87
+ "name": "x",
88
+ "arg": "x",
89
+ "semantic": "data",
90
+ "buffer": { "type": "read-only-storage" },
91
+ "elementType": "$vectorScalar"
92
+ },
93
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
94
+ {
95
+ "name": "params",
96
+ "semantic": "kernel.params",
97
+ "buffer": { "type": "uniform" },
98
+ "struct": {
99
+ "name": "Params",
100
+ "fields": [
101
+ { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" },
102
+ { "name": "chunkCount", "type": "u32", "value": "dim(shapes.data, ranks.data - 1) / tunables.VECTOR_WIDTH" }
103
+ ]
104
+ }
105
+ }
106
+ ],
107
+ "lastAxisScalar": [
108
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
109
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
110
+ {
111
+ "name": "params",
112
+ "semantic": "kernel.params",
113
+ "buffer": { "type": "uniform" },
114
+ "struct": {
115
+ "name": "Params",
116
+ "fields": [
117
+ { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" },
118
+ { "name": "cols", "type": "u32", "value": "dim(shapes.data, ranks.data - 1)" }
119
+ ]
120
+ }
121
+ }
122
+ ],
123
+ "lastAxisScalarSubgroup": [
124
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
125
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
126
+ {
127
+ "name": "params",
128
+ "semantic": "kernel.params",
129
+ "buffer": { "type": "uniform" },
130
+ "struct": {
131
+ "name": "Params",
132
+ "fields": [
133
+ { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" },
134
+ { "name": "chunkCount", "type": "u32", "value": "dim(shapes.data, ranks.data - 1)" }
135
+ ]
136
+ }
137
+ }
138
+ ],
139
+ "scalar": [
140
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
141
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
142
+ {
143
+ "name": "params",
144
+ "semantic": "kernel.params",
145
+ "buffer": { "type": "uniform" },
146
+ "struct": {
147
+ "name": "Params",
148
+ "fields": [
149
+ { "name": "rows", "type": "u32", "value": "1" },
150
+ { "name": "cols", "type": "u32", "value": "1" },
151
+ { "name": "outCount", "type": "u32", "value": "1" }
152
+ ]
153
+ }
154
+ }
155
+ ],
156
+ "rank1Axis0": [
157
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
158
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
159
+ {
160
+ "name": "params",
161
+ "semantic": "kernel.params",
162
+ "buffer": { "type": "uniform" },
163
+ "struct": {
164
+ "name": "Params",
165
+ "fields": [
166
+ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
167
+ { "name": "cols", "type": "u32", "value": "1" },
168
+ { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
169
+ ]
170
+ }
171
+ }
172
+ ],
173
+ "rank2Serial": [
174
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
175
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
176
+ {
177
+ "name": "params",
178
+ "semantic": "kernel.params",
179
+ "buffer": { "type": "uniform" },
180
+ "struct": {
181
+ "name": "Params",
182
+ "fields": [
183
+ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
184
+ { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" },
185
+ { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
186
+ ]
187
+ }
188
+ }
189
+ ],
190
+ "rank2SerialAxis1": [
191
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
192
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
193
+ {
194
+ "name": "params",
195
+ "semantic": "kernel.params",
196
+ "buffer": { "type": "uniform" },
197
+ "struct": {
198
+ "name": "Params",
199
+ "fields": [
200
+ { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" },
201
+ { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
202
+ ]
203
+ }
204
+ }
205
+ ],
206
+ "axis0Parallel": [
207
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
208
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
209
+ {
210
+ "name": "params",
211
+ "semantic": "kernel.params",
212
+ "buffer": { "type": "uniform" },
213
+ "struct": {
214
+ "name": "Params",
215
+ "fields": [
216
+ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
217
+ { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }
218
+ ]
219
+ }
220
+ }
221
+ ],
222
+ "fullReduceSerial": [
223
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
224
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
225
+ {
226
+ "name": "params",
227
+ "semantic": "kernel.params",
228
+ "buffer": { "type": "uniform" },
229
+ "struct": {
230
+ "name": "Params",
231
+ "fields": [
232
+ { "name": "rows", "type": "u32", "value": "numel(shapes.data)" },
233
+ { "name": "cols", "type": "u32", "value": "1" },
234
+ { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
235
+ ]
236
+ }
237
+ }
238
+ ],
239
+ "axisSplitReduce": [
240
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
241
+ { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" },
242
+ {
243
+ "name": "params",
244
+ "semantic": "kernel.params",
245
+ "buffer": { "type": "uniform" },
246
+ "struct": {
247
+ "name": "Params",
248
+ "fields": [
249
+ { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
250
+ { "name": "inner", "type": "u32", "value": "axisSplitInner" },
251
+ { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" }
252
+ ]
253
+ }
254
+ }
255
+ ],
256
+ "axisSplitCombine": [
257
+ {
258
+ "name": "partials",
259
+ "semantic": "partials",
260
+ "buffer": { "type": "read-only-storage" },
261
+ "elementType": "$partialElement"
262
+ },
263
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
264
+ {
265
+ "name": "params",
266
+ "semantic": "kernel.params",
267
+ "buffer": { "type": "uniform" },
268
+ "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }] }
269
+ }
270
+ ],
271
+ "axis0SplitReduce": [
272
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
273
+ { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" },
274
+ {
275
+ "name": "params",
276
+ "semantic": "kernel.params",
277
+ "buffer": { "type": "uniform" },
278
+ "struct": {
279
+ "name": "Params",
280
+ "fields": [
281
+ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
282
+ { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }
283
+ ]
284
+ }
285
+ }
286
+ ],
287
+ "axis0SplitCombine": [
288
+ {
289
+ "name": "partials",
290
+ "semantic": "partials",
291
+ "buffer": { "type": "read-only-storage" },
292
+ "elementType": "$partialElement"
293
+ },
294
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
295
+ {
296
+ "name": "params",
297
+ "semantic": "kernel.params",
298
+ "buffer": { "type": "uniform" },
299
+ "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }] }
300
+ }
301
+ ],
302
+ "rankNAxis": [
303
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
304
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
305
+ {
306
+ "name": "params",
307
+ "semantic": "kernel.params",
308
+ "buffer": { "type": "uniform" },
309
+ "struct": {
310
+ "name": "Params",
311
+ "fields": [
312
+ { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
313
+ { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
314
+ ]
315
+ }
316
+ }
317
+ ],
318
+ "flatPartialF32": [
319
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
320
+ { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" },
321
+ {
322
+ "name": "params",
323
+ "semantic": "kernel.params",
324
+ "buffer": { "type": "uniform" },
325
+ "struct": {
326
+ "name": "Params",
327
+ "fields": [
328
+ { "name": "count4", "type": "u32", "value": "floor(numel(shapes.data) / tunables.VECTOR_WIDTH)" },
329
+ { "name": "numel", "type": "u32", "value": "numel(shapes.data)" }
330
+ ]
331
+ }
332
+ }
333
+ ],
334
+ "flatCombineF32": [
335
+ { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
336
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
337
+ {
338
+ "name": "params",
339
+ "semantic": "kernel.params",
340
+ "buffer": { "type": "uniform" },
341
+ "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "1" }] }
342
+ }
343
+ ],
344
+ "suffixVec4": [
345
+ {
346
+ "name": "x",
347
+ "arg": "x",
348
+ "semantic": "data",
349
+ "buffer": { "type": "read-only-storage" },
350
+ "elementType": "$vectorScalar"
351
+ },
352
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
353
+ {
354
+ "name": "params",
355
+ "semantic": "kernel.params",
356
+ "buffer": { "type": "uniform" },
357
+ "struct": {
358
+ "name": "Params",
359
+ "fields": [
360
+ { "name": "rows", "type": "u32", "value": "numel(shapes.reduced)" },
361
+ {
362
+ "name": "chunkCount",
363
+ "type": "u32",
364
+ "value": "numel(shapes.data) / numel(shapes.reduced) / tunables.VECTOR_WIDTH"
365
+ }
366
+ ]
367
+ }
368
+ }
369
+ ],
370
+ "suffixScalar": [
371
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
372
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
373
+ {
374
+ "name": "params",
375
+ "semantic": "kernel.params",
376
+ "buffer": { "type": "uniform" },
377
+ "struct": {
378
+ "name": "Params",
379
+ "fields": [
380
+ { "name": "rows", "type": "u32", "value": "numel(shapes.reduced)" },
381
+ { "name": "cols", "type": "u32", "value": "numel(shapes.data) / numel(shapes.reduced)" }
382
+ ]
383
+ }
384
+ }
385
+ ],
386
+ "multiAxis": [
387
+ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
388
+ { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
389
+ {
390
+ "name": "params",
391
+ "semantic": "kernel.params",
392
+ "buffer": { "type": "uniform" },
393
+ "struct": {
394
+ "name": "Params",
395
+ "fields": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }]
396
+ }
397
+ }
398
+ ]
399
+ },
400
+ "variants": [
401
+ {
402
+ "id": "contiguous_suffix_subgroup_vec4",
403
+ "priority": 30,
404
+ "requires": { "features": ["subgroups"] },
405
+ "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.data) / numel(shapes.reduced)) % tunables.VECTOR_WIDTH == 0"],
406
+ "constants": {
407
+ "scalar": "dtypes.T",
408
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
409
+ "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(numel(shapes.data) / numel(shapes.reduced), tunables.VECTOR_WIDTH))))"
410
+ },
411
+ "passes": [
412
+ {
413
+ "id": "main",
414
+ "name": "ReduceLogSum.ContiguousSuffixSubgroupVec4",
415
+ "source": {
416
+ "shader": "reduce-row-subgroup.wgsl.jinja",
417
+ "inputs": {
418
+ "op": "\"logsum\"",
419
+ "vec4": true,
420
+ "castF32": "dtypes.T == \"f16\"",
421
+ "usesF16": "dtypes.T == \"f16\""
422
+ }
423
+ },
424
+ "subgroupCollectivesWidth": "portable",
425
+ "bindings": "suffixVec4",
426
+ "dispatch": { "workgroups": "numel(shapes.reduced)" }
427
+ }
428
+ ]
429
+ },
430
+ {
431
+ "id": "contiguous_suffix_tree_vec4",
432
+ "priority": 22,
433
+ "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.data) / numel(shapes.reduced)) % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"],
434
+ "constants": {
435
+ "scalar": "dtypes.T",
436
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
437
+ "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(numel(shapes.data) / numel(shapes.reduced), tunables.VECTOR_WIDTH)))"
438
+ },
439
+ "passes": [
440
+ {
441
+ "id": "main",
442
+ "name": "ReduceLogSum.ContiguousSuffixTreeVec4",
443
+ "source": {
444
+ "shader": "reduce-row-tree.wgsl.jinja",
445
+ "inputs": {
446
+ "op": "\"logsum\"",
447
+ "vec4": true,
448
+ "castF32": "dtypes.T == \"f16\"",
449
+ "usesF16": "dtypes.T == \"f16\""
450
+ }
451
+ },
452
+ "bindings": "suffixVec4",
453
+ "dispatch": { "workgroups": "numel(shapes.reduced)" }
454
+ }
455
+ ]
456
+ },
457
+ {
458
+ "id": "contiguous_suffix_tree",
459
+ "priority": 21,
460
+ "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "treeWorkgroupOk"],
461
+ "constants": {
462
+ "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(numel(shapes.data) / numel(shapes.reduced)))",
463
+ "scalar": "dtypes.T"
464
+ },
465
+ "passes": [
466
+ {
467
+ "id": "main",
468
+ "name": "ReduceLogSum.ContiguousSuffixTree",
469
+ "source": {
470
+ "shader": "reduce-row-tree.wgsl.jinja",
471
+ "inputs": { "op": "\"logsum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
472
+ },
473
+ "bindings": "suffixScalar",
474
+ "dispatch": { "workgroups": "numel(shapes.reduced)" }
475
+ }
476
+ ]
477
+ },
478
+ {
479
+ "id": "multi_axis_rank3",
480
+ "priority": 8,
481
+ "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 3", "(attrs.keepdims == 1 and ranks.reduced == 3 and (dim(shapes.reduced, 0) == 1 if hasAxis(attrs.axes, 0, 3) else dim(shapes.reduced, 0) == dim(shapes.data, 0)) and (dim(shapes.reduced, 1) == 1 if hasAxis(attrs.axes, 1, 3) else dim(shapes.reduced, 1) == dim(shapes.data, 1)) and (dim(shapes.reduced, 2) == 1 if hasAxis(attrs.axes, 2, 3) else dim(shapes.reduced, 2) == dim(shapes.data, 2))) or (attrs.keepdims == 0 and ranks.reduced == 1)"],
482
+ "passes": [
483
+ {
484
+ "id": "main",
485
+ "name": "ReduceLogSum.MultiAxisRank3",
486
+ "source": {
487
+ "shader": "reduce-serial-axis.wgsl.jinja",
488
+ "inputs": {
489
+ "op": "\"logsum\"",
490
+ "indexing": "\"multiaxis\"",
491
+ "castF32": "dtypes.T == \"f16\"",
492
+ "usesF16": "dtypes.T == \"f16\"",
493
+ "rank": 3,
494
+ "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
495
+ "dataShape": "shapes.data",
496
+ "outputShape": "shapes.reduced",
497
+ "outputRank": "ranks.reduced",
498
+ "keepDims": "attrs.keepdims != 0",
499
+ "logicalBool": "tensorDtypes.data == \"bool\""
500
+ }
501
+ },
502
+ "bindings": "multiAxis",
503
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
504
+ }
505
+ ]
506
+ },
507
+ {
508
+ "id": "multi_axis_rank4",
509
+ "priority": 8,
510
+ "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 4", "attrs.noop_with_empty_axes == 0", "numel(shapes.reduced) == (1 if hasAxis(attrs.axes, 0, 4) else dim(shapes.data, 0)) * (1 if hasAxis(attrs.axes, 1, 4) else dim(shapes.data, 1)) * (1 if hasAxis(attrs.axes, 2, 4) else dim(shapes.data, 2)) * (1 if hasAxis(attrs.axes, 3, 4) else dim(shapes.data, 3))", "((attrs.keepdims == 1 and ranks.reduced == 4) or (attrs.keepdims == 0 and ranks.reduced < 4))"],
511
+ "passes": [
512
+ {
513
+ "id": "main",
514
+ "name": "ReduceLogSum.MultiAxisRank4",
515
+ "source": {
516
+ "shader": "reduce-serial-axis.wgsl.jinja",
517
+ "inputs": {
518
+ "op": "\"logsum\"",
519
+ "indexing": "\"multiaxis\"",
520
+ "castF32": "dtypes.T == \"f16\"",
521
+ "usesF16": "dtypes.T == \"f16\"",
522
+ "rank": 4,
523
+ "reduce": ["hasAxis(attrs.axes, 0, 4)", "hasAxis(attrs.axes, 1, 4)", "hasAxis(attrs.axes, 2, 4)", "hasAxis(attrs.axes, 3, 4)"],
524
+ "dataShape": "shapes.data",
525
+ "outputShape": "shapes.reduced",
526
+ "outputRank": "ranks.reduced",
527
+ "keepDims": "attrs.keepdims != 0",
528
+ "logicalBool": "tensorDtypes.data == \"bool\""
529
+ }
530
+ },
531
+ "bindings": "multiAxis",
532
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
533
+ }
534
+ ]
535
+ },
536
+ {
537
+ "id": "noop_empty_axes",
538
+ "priority": 40,
539
+ "when": ["dtypes.T == \"f32\"", "attrs.noop_with_empty_axes == 1", "(attrs.axes | length) == 0", "sameShape(shapes.data, shapes.reduced)"],
540
+ "passes": [
541
+ {
542
+ "id": "main",
543
+ "name": "ReduceLogSum.NoopEmptyAxes",
544
+ "source": { "shader": "reduce-noop-empty-axes.wgsl.jinja", "inputs": { "op": "\"log\"" } },
545
+ "bindings": "elementwise",
546
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
547
+ }
548
+ ]
549
+ },
550
+ {
551
+ "id": "tree_last_axis_vec4",
552
+ "priority": 23,
553
+ "demoteWhen": ["rowSerialPreferred"],
554
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 1", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "attrs.noop_with_empty_axes == 0", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "lastAxisCols % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"],
555
+ "constants": {
556
+ "scalar": "dtypes.T",
557
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
558
+ "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH)))"
559
+ },
560
+ "passes": [
561
+ {
562
+ "id": "main",
563
+ "name": "ReduceLogSum.TreeRowVec4",
564
+ "source": {
565
+ "shader": "reduce-row-tree.wgsl.jinja",
566
+ "inputs": {
567
+ "op": "\"logsum\"",
568
+ "vec4": true,
569
+ "castF32": "dtypes.T == \"f16\"",
570
+ "usesF16": "dtypes.T == \"f16\""
571
+ }
572
+ },
573
+ "bindings": "lastAxisVec4",
574
+ "dispatch": { "workgroups": "lastAxisRows" }
575
+ }
576
+ ]
577
+ },
578
+ {
579
+ "id": "rank0_scalar",
580
+ "priority": 40,
581
+ "constants": { "axis": 0 },
582
+ "when": ["f16Ok(dtypes.T)", "ranks.data == 0", "ranks.reduced == 0"],
583
+ "passes": [
584
+ {
585
+ "id": "main",
586
+ "name": "ReduceLogSum.Rank0Scalar",
587
+ "source": {
588
+ "shader": "reduce-serial-axis.wgsl.jinja",
589
+ "inputs": {
590
+ "op": "\"logsum\"",
591
+ "indexing": "\"axis2d\"",
592
+ "castF32": "dtypes.T == \"f16\"",
593
+ "usesF16": "dtypes.T == \"f16\"",
594
+ "logicalBool": "tensorDtypes.data == \"bool\""
595
+ }
596
+ },
597
+ "bindings": "scalar",
598
+ "dispatch": { "x": 1 }
599
+ }
600
+ ]
601
+ },
602
+ {
603
+ "id": "rank1_axis0",
604
+ "constants": { "axis": 0 },
605
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 1", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.reduced == 0) or (attrs.keepdims == 1 and ranks.reduced == 1 and dim(shapes.reduced, 0) == 1))"],
606
+ "passes": [
607
+ {
608
+ "id": "main",
609
+ "name": "ReduceLogSum.Rank1Axis0",
610
+ "source": {
611
+ "shader": "reduce-serial-axis.wgsl.jinja",
612
+ "inputs": {
613
+ "op": "\"logsum\"",
614
+ "indexing": "\"axis2d\"",
615
+ "castF32": "dtypes.T == \"f16\"",
616
+ "usesF16": "dtypes.T == \"f16\"",
617
+ "logicalBool": "tensorDtypes.data == \"bool\""
618
+ }
619
+ },
620
+ "bindings": "rank1Axis0",
621
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
622
+ }
623
+ ]
624
+ },
625
+ {
626
+ "id": "axis1_parallel",
627
+ "priority": 20,
628
+ "demoteWhen": ["rowSerialPreferred"],
629
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 2", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "treeWorkgroupOk"],
630
+ "constants": { "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.data, ranks.data - 1)))" },
631
+ "passes": [
632
+ {
633
+ "id": "main",
634
+ "name": "ReduceLogSum.Axis1Parallel",
635
+ "source": {
636
+ "shader": "reduce-row-tree.wgsl.jinja",
637
+ "inputs": { "op": "\"logsum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
638
+ },
639
+ "bindings": "lastAxisScalar",
640
+ "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
641
+ }
642
+ ]
643
+ },
644
+ {
645
+ "id": "axis_split",
646
+ "priority": 24,
647
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "attrs.noop_with_empty_axes == 0", "ranks.data >= 2", "reduceAxis < ranks.data - 1", "not (ranks.data == 2 and reduceAxis == 0)", "axisSplitDim >= tunables.AXIS0_SPLIT_MIN_ROWS", "axisSplitOutputs >= 1", "axisSplitOutputs <= 4096", "axisSplitOutputs == rows(shapes.data, reduceAxis)", "axisSplitPathFits"],
648
+ "derive": { "splitCount": "axisSplitCount" },
649
+ "constants": { "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" },
650
+ "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * axisSplitOutputs]" }],
651
+ "passes": [
652
+ {
653
+ "id": "split_reduce",
654
+ "name": "ReduceLogSum.AxisSplitReduce",
655
+ "source": {
656
+ "shader": "reduce-axis-split-reduce.wgsl.jinja",
657
+ "inputs": {
658
+ "op": "\"logsum\"",
659
+ "workgroupSize": "reduceWorkgroupSize",
660
+ "split": "splitCount",
661
+ "castF32": "dtypes.T == \"f16\"",
662
+ "usesF16": "dtypes.T == \"f16\""
663
+ }
664
+ },
665
+ "bindings": "axisSplitReduce",
666
+ "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
667
+ },
668
+ {
669
+ "id": "combine",
670
+ "name": "ReduceLogSum.AxisSplitCombine",
671
+ "source": {
672
+ "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
673
+ "inputs": {
674
+ "op": "\"logsum\"",
675
+ "workgroupSize": "reduceWorkgroupSize",
676
+ "split": "splitCount",
677
+ "outputF16": "dtypes.T == \"f16\""
678
+ }
679
+ },
680
+ "bindings": "axisSplitCombine",
681
+ "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize" }
682
+ }
683
+ ]
684
+ },
685
+ {
686
+ "id": "axis_split_tiled_narrow",
687
+ "priority": 25,
688
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "attrs.noop_with_empty_axes == 0", "ranks.data >= 2", "reduceAxis < ranks.data - 1", "axisSplitDim >= tunables.AXIS0_SPLIT_MIN_ROWS", "axisSplitOutputs >= 1", "axisSplitOutputs <= 2 * tunables.AXIS_SPLIT_TILE_COLS", "reduceWorkgroupSize % tunables.AXIS_SPLIT_TILE_COLS == 0", "axisSplitOutputs == rows(shapes.data, reduceAxis)", "axisSplitPathFits"],
689
+ "derive": { "splitCount": "axisSplitCount" },
690
+ "constants": {
691
+ "partialElement": "\"f32\"",
692
+ "scalar": "dtypes.T",
693
+ "workgroupSize": "reduceWorkgroupSize",
694
+ "split": "splitCount",
695
+ "tileCols": "tunables.AXIS_SPLIT_TILE_COLS"
696
+ },
697
+ "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * axisSplitOutputs]" }],
698
+ "passes": [
699
+ {
700
+ "id": "split_reduce",
701
+ "name": "ReduceLogSum.AxisSplitTiledReduce",
702
+ "source": {
703
+ "shader": "reduce-axis0-tilecols.wgsl.jinja",
704
+ "inputs": {
705
+ "op": "\"logsum\"",
706
+ "workgroupSize": "reduceWorkgroupSize",
707
+ "split": "splitCount",
708
+ "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
709
+ "castF32": "dtypes.T == \"f16\"",
710
+ "usesF16": "dtypes.T == \"f16\""
711
+ }
712
+ },
713
+ "bindings": "axisSplitReduce",
714
+ "dispatch": { "workgroups": "ceilDiv((axisSplitOutputs), (constants.tileCols))", "y": "splitCount" }
715
+ },
716
+ {
717
+ "id": "combine",
718
+ "name": "ReduceLogSum.AxisSplitCombine",
719
+ "source": {
720
+ "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
721
+ "inputs": {
722
+ "op": "\"logsum\"",
723
+ "workgroupSize": "reduceWorkgroupSize",
724
+ "split": "splitCount",
725
+ "outputF16": "dtypes.T == \"f16\""
726
+ }
727
+ },
728
+ "bindings": "axisSplitCombine",
729
+ "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize" }
730
+ }
731
+ ]
732
+ },
733
+ {
734
+ "id": "axis0_splitk",
735
+ "priority": 22,
736
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_SPLIT_MIN_ROWS", "dim(shapes.data, 1) > 0", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == dim(shapes.data, 1)))", "axis0SplitPathFits"],
737
+ "derive": { "splitCount": "axis0SplitCount" },
738
+ "constants": { "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" },
739
+ "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * dim(shapes.data, 1)]" }],
740
+ "passes": [
741
+ {
742
+ "id": "split_reduce",
743
+ "name": "ReduceLogSum.Axis0SplitKReduce",
744
+ "source": {
745
+ "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
746
+ "inputs": {
747
+ "op": "\"logsum\"",
748
+ "workgroupSize": "reduceWorkgroupSize",
749
+ "split": "splitCount",
750
+ "castF32": "dtypes.T == \"f16\"",
751
+ "usesF16": "dtypes.T == \"f16\""
752
+ }
753
+ },
754
+ "bindings": "axis0SplitReduce",
755
+ "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
756
+ },
757
+ {
758
+ "id": "combine",
759
+ "name": "ReduceLogSum.Axis0SplitKCombine",
760
+ "source": {
761
+ "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
762
+ "inputs": {
763
+ "op": "\"logsum\"",
764
+ "workgroupSize": "reduceWorkgroupSize",
765
+ "split": "splitCount",
766
+ "outputF16": "dtypes.T == \"f16\""
767
+ }
768
+ },
769
+ "bindings": "axis0SplitCombine",
770
+ "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize" }
771
+ }
772
+ ]
773
+ },
774
+ {
775
+ "id": "axis0_tilecols",
776
+ "priority": 20,
777
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_TILE_MIN_ROWS", "axis0Cols >= tunables.AXIS0_TILE_MIN_COLS", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == dim(shapes.data, 1)))", "axis0TilePathFits"],
778
+ "constants": { "workgroupSize": "reduceWorkgroupSize", "tileCols": "tunables.AXIS0_TILE_COLS" },
779
+ "passes": [
780
+ {
781
+ "id": "main",
782
+ "name": "ReduceLogSum.Axis0TileCols",
783
+ "source": {
784
+ "shader": "reduce-axis0-tilecols.wgsl.jinja",
785
+ "inputs": { "op": "\"logsum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
786
+ },
787
+ "bindings": "axis0Parallel",
788
+ "dispatch": { "workgroups": "ceilDiv((dim(shapes.data, 1)), (constants.tileCols))" }
789
+ }
790
+ ]
791
+ },
792
+ {
793
+ "id": "all_axes_flat",
794
+ "priority": 31,
795
+ "constants": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
796
+ "when": ["flatParallelCovered"],
797
+ "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
798
+ "passes": [
799
+ {
800
+ "id": "flat_partial",
801
+ "name": "ReduceLogSum.AllAxesFlatPartial",
802
+ "source": {
803
+ "shader": "reduce-flat-partial.wgsl.jinja",
804
+ "inputs": { "op": "\"logsum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
805
+ },
806
+ "bindings": "flatPartialF32",
807
+ "dispatch": { "x": "flatSplitCount" }
808
+ },
809
+ {
810
+ "id": "combine",
811
+ "name": "ReduceLogSum.AllAxesFlatCombine",
812
+ "source": {
813
+ "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
814
+ "inputs": { "op": "\"logsum\"", "outputF16": "dtypes.T == \"f16\"" }
815
+ },
816
+ "bindings": "flatCombineF32",
817
+ "dispatch": { "x": 1 }
818
+ }
819
+ ]
820
+ },
821
+ {
822
+ "id": "rankn_single_axis_generic",
823
+ "priority": 12,
824
+ "supersededBy": ["axis_split_tiled_narrow", "axis_split", "subgroup_last_axis_vec4", "subgroup_last_axis", "tree_last_axis_vec4"],
825
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.noop_with_empty_axes == 0", "reduceAxis < ranks.data", "numel(shapes.reduced) == rows(shapes.data, reduceAxis)", "((attrs.keepdims == 0 and ranks.reduced == ranks.data - 1) or (attrs.keepdims == 1 and ranks.reduced == ranks.data and dim(shapes.reduced, reduceAxis) == 1))"],
826
+ "passes": [
827
+ {
828
+ "id": "main",
829
+ "name": "ReduceLogSum.RankNSingleAxisGeneric",
830
+ "source": {
831
+ "shader": "reduce-serial-axis.wgsl.jinja",
832
+ "inputs": {
833
+ "op": "\"logsum\"",
834
+ "indexing": "\"rankn\"",
835
+ "castF32": "dtypes.T == \"f16\"",
836
+ "usesF16": "dtypes.T == \"f16\"",
837
+ "rank": "ranks.data",
838
+ "axis": "reduceAxis",
839
+ "dataShape": "shapes.data",
840
+ "outputShape": "shapes.reduced",
841
+ "outputRank": "ranks.reduced",
842
+ "keepDims": "attrs.keepdims != 0",
843
+ "logicalBool": "tensorDtypes.data == \"bool\""
844
+ }
845
+ },
846
+ "bindings": "rankNAxis",
847
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
848
+ }
849
+ ]
850
+ },
851
+ {
852
+ "id": "subgroup_last_axis_vec4",
853
+ "priority": 25,
854
+ "requires": { "features": ["subgroups"] },
855
+ "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 1", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "dim(shapes.data, ranks.data - 1) >= 4", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH == 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "not rowSerialPreferred"],
856
+ "constants": {
857
+ "scalar": "dtypes.T",
858
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
859
+ "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH))))"
860
+ },
861
+ "passes": [
862
+ {
863
+ "id": "main",
864
+ "name": "ReduceLogSum.SubgroupRowVec4",
865
+ "source": {
866
+ "shader": "reduce-row-subgroup.wgsl.jinja",
867
+ "inputs": {
868
+ "op": "\"logsum\"",
869
+ "vec4": true,
870
+ "castF32": "dtypes.T == \"f16\"",
871
+ "usesF16": "dtypes.T == \"f16\""
872
+ }
873
+ },
874
+ "subgroupCollectivesWidth": "portable",
875
+ "bindings": "lastAxisVec4",
876
+ "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
877
+ }
878
+ ]
879
+ },
880
+ {
881
+ "id": "subgroup_last_axis",
882
+ "priority": 24,
883
+ "requires": { "features": ["subgroups"] },
884
+ "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 1", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "dim(shapes.data, ranks.data - 1) > 0", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH != 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "not rowSerialPreferred"],
885
+ "constants": {
886
+ "scalar": "dtypes.T",
887
+ "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))"
888
+ },
889
+ "passes": [
890
+ {
891
+ "id": "main",
892
+ "name": "ReduceLogSum.SubgroupRow",
893
+ "source": {
894
+ "shader": "reduce-row-subgroup.wgsl.jinja",
895
+ "inputs": {
896
+ "op": "\"logsum\"",
897
+ "vec4": false,
898
+ "castF32": "dtypes.T == \"f16\"",
899
+ "usesF16": "dtypes.T == \"f16\""
900
+ }
901
+ },
902
+ "subgroupCollectivesWidth": "portable",
903
+ "bindings": "lastAxisScalarSubgroup",
904
+ "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
905
+ }
906
+ ]
907
+ },
908
+ {
909
+ "id": "axis0",
910
+ "priority": 0,
911
+ "supersededBy": ["axis_split_tiled_narrow", "axis0_splitk", "axis0_tilecols"],
912
+ "constants": { "axis": 0 },
913
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == dim(shapes.data, 1)))"],
914
+ "passes": [
915
+ {
916
+ "id": "main",
917
+ "name": "axis0",
918
+ "source": {
919
+ "shader": "reduce-serial-axis.wgsl.jinja",
920
+ "inputs": {
921
+ "op": "\"logsum\"",
922
+ "indexing": "\"axis2d\"",
923
+ "castF32": "dtypes.T == \"f16\"",
924
+ "usesF16": "dtypes.T == \"f16\""
925
+ }
926
+ },
927
+ "bindings": "rank2Serial",
928
+ "constants": { "axis": 0 },
929
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
930
+ }
931
+ ]
932
+ },
933
+ {
934
+ "id": "axis1",
935
+ "priority": 0,
936
+ "constants": { "axis": 1 },
937
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 0)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == 1))"],
938
+ "passes": [
939
+ {
940
+ "id": "main",
941
+ "name": "axis1",
942
+ "source": {
943
+ "shader": "reduce-serial-axis.wgsl.jinja",
944
+ "inputs": {
945
+ "op": "\"logsum\"",
946
+ "indexing": "\"axis2d\"",
947
+ "castF32": "dtypes.T == \"f16\"",
948
+ "usesF16": "dtypes.T == \"f16\""
949
+ }
950
+ },
951
+ "bindings": "rank2SerialAxis1",
952
+ "constants": { "axis": 1 },
953
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
954
+ }
955
+ ]
956
+ },
957
+ {
958
+ "id": "all_axes_keepdims",
959
+ "priority": 30,
960
+ "constants": { "axis": 0 },
961
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.keepdims == 1", "ranks.reduced == ranks.data", "numel(shapes.reduced) == 1"],
962
+ "passes": [
963
+ {
964
+ "id": "main",
965
+ "name": "ReduceLogSum.Rank3AllAxesKeepdims",
966
+ "source": {
967
+ "shader": "reduce-serial-axis.wgsl.jinja",
968
+ "inputs": {
969
+ "op": "\"logsum\"",
970
+ "indexing": "\"axis2d\"",
971
+ "castF32": "dtypes.T == \"f16\"",
972
+ "usesF16": "dtypes.T == \"f16\""
973
+ }
974
+ },
975
+ "bindings": "fullReduceSerial",
976
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
977
+ }
978
+ ]
979
+ },
980
+ {
981
+ "id": "all_axes_no_keepdims",
982
+ "priority": 30,
983
+ "constants": { "axis": 0 },
984
+ "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.keepdims == 0", "attrs.noop_with_empty_axes == 0", "ranks.reduced == 0"],
985
+ "passes": [
986
+ {
987
+ "id": "main",
988
+ "name": "ReduceLogSum.Rank3AllAxesNoKeepdims",
989
+ "source": {
990
+ "shader": "reduce-serial-axis.wgsl.jinja",
991
+ "inputs": {
992
+ "op": "\"logsum\"",
993
+ "indexing": "\"axis2d\"",
994
+ "castF32": "dtypes.T == \"f16\"",
995
+ "usesF16": "dtypes.T == \"f16\""
996
+ }
997
+ },
998
+ "bindings": "fullReduceSerial",
999
+ "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
1000
+ }
1001
+ ]
1002
+ }
1003
+ ]
1004
+ }
build/webgpu/metadata.json ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "ai.onnx.ReduceLogSum",
3
+ "id": "_ai_onnx_reducelogsum_webgpu_0cd4469",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "backend": { "type": "webgpu" },
7
+ "digest": {
8
+ "algorithm": "sha256",
9
+ "files": {
10
+ "bench.json": "tExpDbPGHls8UweHAWP2+FY+piqkNai/ffADNFmbuq4=",
11
+ "manifest.json": "/pyMW/xHAyNmrprpHqurMtEhzXppdfFwqyI9wd4JVoM=",
12
+ "reduce-axis-split-reduce.wgsl.jinja": "Q50AhcbpYvCADMPrbGtPXKbmfT8f6LmQY+6+/CEX224=",
13
+ "reduce-axis0-splitk-combine.wgsl.jinja": "ul0TxUqohcvm6CDpKIFB0WoPiPqqpuh1AvsW1+6+3Rk=",
14
+ "reduce-axis0-splitk-reduce.wgsl.jinja": "jc7OBcuREabFLa+kwdZs2oEI3W/BWfHh8nyeUQnIVyY=",
15
+ "reduce-axis0-tilecols.wgsl.jinja": "6PKq2qx5+gRoZjG0LEOrwi1PbWT9b7QGzMjEekpvqbQ=",
16
+ "reduce-flat-partial.wgsl.jinja": "Jz5ZMpVLRJ702gjunfgrUEwqOase9mEBr5SRHlwA0zU=",
17
+ "reduce-noop-empty-axes.wgsl.jinja": "77n6kZMxBe/1e8ye2C8S3HMvXWicW/c3pzHEnQS42Gk=",
18
+ "reduce-row-subgroup.wgsl.jinja": "oJiY7cuPhJ++GEdGcV/RauZy5PUHrVD4YHx0sZyh6Tk=",
19
+ "reduce-row-tree.wgsl.jinja": "YobwxXxVZkfKiUnfWXcLIlz6y4J7jurRcnNC5UjPmDc=",
20
+ "reduce-serial-axis.wgsl.jinja": "JzfMwQJ1gzBxszuL1+IC7d/kWKP9dArAHl63PkAfdmA=",
21
+ "test.json": "R7TN083sH8d7kSwSt1fUUftGCxWFgrsezAmGa+hTAKw="
22
+ }
23
+ },
24
+ "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
25
+ "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.ReduceLogSum" }
26
+ }
build/webgpu/reduce-axis-split-reduce.wgsl.jinja ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 1 of a split-K reduction for an (outer, axis, inner) flattening. Each
2
+ // (output, segment) invocation reduces one axis slice at stride inner and writes
3
+ // partials[segment * outputs + output]. Adjacent output threads read adjacent
4
+ // inner-axis elements. The combine pass folds the segments and finalizes the
5
+ // selected reduction.
6
+ //
7
+ // logsumexp writes three partial planes per output and segment: the segment
8
+ // maximum, the sum of exp(x - maximum), and a packed NaN marker.
9
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
10
+ {% set xa = "f32(" if castF32 else "" %}
11
+ {% set ax = ")" if castF32 else "" %}
12
+ {% if source.usesF16 is defined and source.usesF16 %}
13
+ enable f16;
14
+ {% endif %}
15
+ {{ env.wgsl.resourceDeclarations }}
16
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
17
+ * evaluation, so f32 identities are constructed at runtime from their
18
+ * IEEE-754 bit patterns. */
19
+
20
+
21
+ const WG: u32 = {{ workgroupSize }}u;
22
+ const SPLIT: u32 = {{ split }}u;
23
+
24
+ @compute @workgroup_size(WG, 1, 1)
25
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
26
+ @builtin(workgroup_id) wg: vec3<u32>,
27
+ @builtin(num_workgroups) nwg: vec3<u32>) {
28
+ // 2D-folded output index: wg.z carries the high bits past the
29
+ // per-dimension dispatch limit on the x dimension.
30
+ let output_index = (wg.x + wg.z * nwg.x) * WG + (gid.x % WG);
31
+ let seg = wg.y;
32
+ if (output_index >= params.outputs) { return; }
33
+
34
+ let outer_index = output_index / params.inner;
35
+ let inner_index = output_index % params.inner;
36
+ let input_base = outer_index * params.axisDim * params.inner + inner_index;
37
+
38
+ // Even split of [0, axisDim) across SPLIT segments (last may be shorter).
39
+ let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
40
+ let a0 = seg * chunk;
41
+ var a1 = a0 + chunk;
42
+ if (a1 > params.axisDim) { a1 = params.axisDim; }
43
+
44
+ var acc = 0.0;
45
+ for (var axis_index = a0; axis_index < a1; axis_index = axis_index + 1u) {
46
+ acc = acc + {{ xa }}x[input_base + axis_index * params.inner]{{ ax }};
47
+ }
48
+ partials[seg * params.outputs + output_index] = acc;
49
+ }
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 2 of the split-K column-wise reduction. One invocation per output column
2
+ // folds the segment partials and applies the selected reduction's final step.
3
+ // Segments are folded in ascending order for deterministic results. This order
4
+ // differs from the single-pass reduction but remains within the f32 tolerance.
5
+ {% set yv = "f16(" if source.outputF16 else "" %}
6
+ {% set vy = ")" if source.outputF16 else "" %}
7
+ {% if source.outputF16 %}
8
+ enable f16;
9
+ {% endif %}
10
+ {{ env.wgsl.resourceDeclarations }}
11
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
12
+ * evaluation, so f32 identities are constructed at runtime from their
13
+ * IEEE-754 bit patterns. */
14
+
15
+
16
+ const WG: u32 = {{ workgroupSize }}u;
17
+ const SPLIT: u32 = {{ split }}u;
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 stride = nwg.x * WG;
23
+ let start = (gid.y * nwg.x * WG) + gid.x;
24
+ for (var col = start; col < params.cols; col = col + stride) {
25
+ var total = 0.0;
26
+ for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
27
+ let p = partials[seg * params.cols + col];
28
+ total = total + p;
29
+ }
30
+ y[col] = {{ yv }}log(total){{ vy }};
31
+ }
32
+ }
build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 1 of a split-K column-wise reduction. Splitting rows across workgroup
2
+ // segments increases residency for tall matrices. Each (column, segment)
3
+ // invocation reduces one row slice and writes partials[segment * columns +
4
+ // column]. Adjacent column threads keep row reads coalesced.
5
+ //
6
+ // logsumexp writes the segment maximum and sum of exp(x - maximum) as separate
7
+ // partial planes; the combine pass merges them stably and handles NaN and +Inf.
8
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
9
+ {% set xa = "f32(" if castF32 else "" %}
10
+ {% set ax = ")" if castF32 else "" %}
11
+ {% if source.usesF16 is defined and source.usesF16 %}
12
+ enable f16;
13
+ {% endif %}
14
+ {{ env.wgsl.resourceDeclarations }}
15
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
16
+ * evaluation, so f32 identities are constructed at runtime from their
17
+ * IEEE-754 bit patterns. */
18
+
19
+
20
+ const WG: u32 = {{ workgroupSize }}u;
21
+ const SPLIT: u32 = {{ split }}u;
22
+
23
+ @compute @workgroup_size(WG, 1, 1)
24
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
25
+ @builtin(workgroup_id) wg: vec3<u32>,
26
+ @builtin(num_workgroups) nwg: vec3<u32>) {
27
+ // 2D-folded column index: wg.z carries the high bits past the maxComputeWorkgroupsPerDimension
28
+ // workgroup-per-dimension dispatch limit on the x dimension.
29
+ let col = (wg.x + wg.z * nwg.x) * WG + (gid.x % WG);
30
+ let seg = wg.y;
31
+ if (col >= params.cols) { return; }
32
+
33
+ // Even split of [0, rows) across SPLIT segments (last segment may be shorter).
34
+ let chunk = (params.rows + SPLIT - 1u) / SPLIT;
35
+ let r0 = seg * chunk;
36
+ var r1 = r0 + chunk;
37
+ if (r1 > params.rows) { r1 = params.rows; }
38
+
39
+ var acc = 0.0;
40
+ for (var row = r0; row < r1; row = row + 1u) {
41
+ acc = acc + {{ xa }}x[row * params.cols + col]{{ ax }};
42
+ }
43
+ partials[seg * params.cols + col] = acc;
44
+ }
build/webgpu/reduce-axis0-tilecols.wgsl.jinja ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Tiled column-wise reduction. Each workgroup owns TILE_COLS columns;
2
+ // ROW_LANES threads stride the rows of one column, then lane 0 folds their
3
+ // partials. In split mode, workgroup y selects an axis segment and finalization
4
+ // is deferred to the combine pass. Cooperative row lanes preserve occupancy
5
+ // when the flattened output has only a few elements.
6
+ {% set splitMode = source.split is defined %}
7
+ {% if splitMode %}
8
+ {% set rowBegin = "row_begin + row_lane" %}
9
+ {% set rowEnd = "row_end" %}
10
+ {% set elem = "x[input_base + row * params.inner]" %}
11
+ {% else %}
12
+ {% set rowBegin = "row_lane" %}
13
+ {% set rowEnd = "params.rows" %}
14
+ {% set elem = "x[inputBase + row * params.cols + col]" %}
15
+ {% endif %}
16
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
17
+ {% set intMode = source.intMode is defined and source.intMode %}
18
+ {% set scalar = "f32" if castF32 else scalar %}
19
+ {% if castF32 %}
20
+ {% set elem = "f32(" ~ elem ~ ")" %}
21
+ {% endif %}
22
+ {% set yv = "f16(" if castF32 else "" %}
23
+ {% set vy = ")" if castF32 else "" %}
24
+ {% if source.usesF16 is defined and source.usesF16 %}
25
+ enable f16;
26
+ {% endif %}
27
+ {{ env.wgsl.resourceDeclarations }}
28
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
29
+ * evaluation, so f32 identities are constructed at runtime from their
30
+ * IEEE-754 bit patterns. */
31
+
32
+ {% if not splitMode and not intMode and (source.op == "logsum" or source.op == "logsumexp") %}
33
+ fn negative_infinity() -> f32 {
34
+ var bits = 0xff800000u;
35
+ return bitcast<f32>(bits);
36
+ }
37
+
38
+ {% endif %}
39
+
40
+ const WG: u32 = {{ workgroupSize }}u;
41
+ const TILE_COLS: u32 = {{ tileCols }}u;
42
+ const ROW_LANES: u32 = WG / TILE_COLS;
43
+ {% if splitMode %}
44
+ const SPLIT: u32 = {{ source.split }}u;
45
+ {% endif %}
46
+
47
+ var<workgroup> partial: array<{{ scalar if (source.op == "max" or source.op == "min" or intMode) else "f32" }}, WG>;
48
+ @compute @workgroup_size(WG, 1, 1)
49
+ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if not splitMode %},
50
+ @builtin(num_workgroups) nwg: vec3<u32>{% endif %}) {
51
+ let tid = lid.x;
52
+ let col_lane = tid % TILE_COLS;
53
+ let row_lane = tid / TILE_COLS;
54
+ {% if splitMode %}
55
+ // Narrow outputs: wg.x covers every column tile, wg.y is the axis segment.
56
+ let col = wg.x * TILE_COLS + col_lane;
57
+ let outputIndex = col;
58
+ let in_bounds = col < params.outputs;
59
+ let seg = wg.y;
60
+ let outer_index = col / params.inner;
61
+ let inner_index = col % params.inner;
62
+ let input_base = outer_index * params.axisDim * params.inner + inner_index;
63
+ // Even split of [0, axisDim) across SPLIT segments (last may be shorter).
64
+ let chunk = (params.axisDim + SPLIT - 1u) / SPLIT;
65
+ let row_begin = seg * chunk;
66
+ let row_end = min(row_begin + chunk, params.axisDim);
67
+ {% else %}
68
+ // 2D-folded tile index: wg.y carries the high bits past the dispatch limit.
69
+ // The batched form reuses this same coalesced axis-0 reduction for a middle
70
+ // axis by assigning consecutive tiles to each outer slice.
71
+ let tile = wg.x + wg.y * nwg.x;
72
+ let col = tile * TILE_COLS + col_lane;
73
+ let inputBase = 0u;
74
+ let outputIndex = col;
75
+ let in_bounds = col < params.cols;
76
+ {% endif %}
77
+ {% if source.op == "logsum" and not splitMode %}
78
+
79
+ if (params.rows == 0u) {
80
+ if (row_lane == 0u && in_bounds) {
81
+ y[outputIndex] = {{ yv }}negative_infinity(){{ vy }};
82
+ }
83
+ return;
84
+ }
85
+ {% endif %}
86
+
87
+ var acc = 0.0;
88
+ if (in_bounds) {
89
+ for (var row = {{ rowBegin }}; row < {{ rowEnd }}; row = row + ROW_LANES) {
90
+ acc = acc + {{ elem }};
91
+ }
92
+ }
93
+ partial[tid] = acc;
94
+ workgroupBarrier();
95
+
96
+ if (row_lane == 0u && in_bounds) {
97
+ var total = partial[col_lane];
98
+ for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
99
+ total = total + partial[lane * TILE_COLS + col_lane];
100
+ }
101
+ {% if splitMode %}
102
+ partials[seg * params.outputs + outputIndex] = total;
103
+ {% else %}
104
+ y[outputIndex] = {{ yv }}log(total){{ vy }};
105
+ {% endif %}
106
+ }
107
+ }
build/webgpu/reduce-flat-partial.wgsl.jinja ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Pass 1 of a flat reduction to one scalar. Instead of assigning the whole
2
+ // tensor to one invocation, SPLIT workgroups traverse the flat input, accumulate
3
+ // in registers, reduce within each workgroup, and write one partial each. A
4
+ // following combine pass folds the partials and applies the operation finalizer.
5
+ //
6
+ // Scalar f32 bindings keep arbitrary element counts legal. The grid-stride loop
7
+ // manually assembles full vec4 groups from contiguous scalars, and one global
8
+ // thread folds the final zero-to-three scalar elements exactly once.
9
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
10
+ {% set xa = "f32(" if castF32 else "" %}
11
+ {% set ax = ")" if castF32 else "" %}
12
+ {% if source.usesF16 is defined and source.usesF16 %}
13
+ enable f16;
14
+ {% endif %}
15
+ {{ env.wgsl.resourceDeclarations }}
16
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
17
+ * evaluation, so f32 identities are constructed at runtime from their
18
+ * IEEE-754 bit patterns. */
19
+
20
+
21
+ const WG: u32 = {{ workgroupSize }}u;
22
+ var<workgroup> red: array<{{ "i32" if source.intMode else "f32" }}, WG>;
23
+ @compute @workgroup_size(WG)
24
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
25
+ @builtin(local_invocation_id) lid: vec3<u32>,
26
+ @builtin(workgroup_id) wg: vec3<u32>,
27
+ @builtin(num_workgroups) nwg: vec3<u32>) {
28
+ let tid = lid.x;
29
+ let gstride = nwg.x * WG;
30
+ var acc = 0.0;
31
+ // Grid-stride over the flat vec4 groups (params.count4 = numel / 4, floored).
32
+ for (var i = gid.x; i < params.count4; i = i + gstride) {
33
+ let b = 4u * i;
34
+ let v = vec4<{{ "i32" if source.intMode else "f32" }}>({{ xa }}x[b]{{ ax }}, {{ xa }}x[b + 1u]{{ ax }}, {{ xa }}x[b + 2u]{{ ax }}, {{ xa }}x[b + 3u]{{ ax }});
35
+ acc = acc + v.x + v.y + v.z + v.w;
36
+ }
37
+ // Scalar tail (the 0..3 elements past the last full vec4). One global thread
38
+ // folds it so it is counted exactly once; the count is tiny so serializing it
39
+ // is negligible. numel % 4 == 0 makes this loop empty.
40
+ if (gid.x == 0u) {
41
+ for (var i = 4u * params.count4; i < params.numel; i = i + 1u) {
42
+ let s = {{ xa }}x[i]{{ ax }};
43
+ acc = acc + s;
44
+ }
45
+ }
46
+ red[tid] = acc;
47
+ workgroupBarrier();
48
+ var stride: u32 = WG / 2u;
49
+ loop {
50
+ if (stride == 0u) { break; }
51
+ if (tid < stride) {
52
+ red[tid] = red[tid] + red[tid + stride];
53
+ }
54
+ stride = stride / 2u;
55
+ workgroupBarrier();
56
+ }
57
+ if (tid == 0u) {
58
+ partials[wg.x] = red[0];
59
+ }
60
+ }
build/webgpu/reduce-noop-empty-axes.wgsl.jinja ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {{ env.wgsl.resourceDeclarations }}
2
+
3
+ @compute @workgroup_size({{ reduceWorkgroupSize }})
4
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
5
+ // 2D-folded flat index: gid.y carries the high bits past the
6
+ // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
7
+ let i = gid.x + gid.y * nwg.x * {{ reduceWorkgroupSize }}u;
8
+ if (i >= params.count) {
9
+ return;
10
+ }
11
+ let v = x[i];
12
+ y[i] = log(v);
13
+ }
build/webgpu/reduce-row-subgroup.wgsl.jinja ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Subgroup row reduction for a contiguous last axis. One workgroup owns each
2
+ // output row. Threads reduce strided chunks, subgroup leaders deposit their
3
+ // partials in workgroup memory, and thread 0 folds those slots and finalizes.
4
+ //
5
+ // Accumulators are f32 except for integer max/min. Mean, l2, and logsum apply
6
+ // division, square root, or logarithm only at finalization. logsumexp uses
7
+ // max-subtraction, bit-exact NaN detection, NaN propagation, and a +Inf
8
+ // short-circuit. Runtime bitcasts provide max/min infinity identities because
9
+ // WGSL constant evaluation rejects infinite constants. f16 storage is widened
10
+ // before accumulation and narrowed only for the final store.
11
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
12
+ {% set scalar = "f32" if castF32 else scalar %}
13
+ {% set xv = ("vec4<f32>(" if source.vec4 else "f32(") if castF32 else "" %}
14
+ {% set vx = ")" if castF32 else "" %}
15
+ {% set yv = "f16(" if castF32 else "" %}
16
+ {% set vy = ")" if castF32 else "" %}
17
+ enable subgroups;
18
+ {% if source.usesF16 is defined and source.usesF16 %}
19
+ enable f16;
20
+ {% endif %}
21
+ {{ env.wgsl.resourceDeclarations }}
22
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
23
+ * evaluation, so f32 identities are constructed at runtime from their
24
+ * IEEE-754 bit patterns. */
25
+
26
+
27
+ const WG: u32 = {{ workgroupSize }}u;
28
+ var<workgroup> wgPartial: array<{{ scalar }}, WG>;
29
+
30
+ {% macro emit_reduce(name, collective, combine) %}
31
+ fn {{ name }}(value: {{ scalar }}, sgLid: u32, sgId: u32, numSg: u32) -> {{ scalar }} {
32
+ let sgVal = {{ collective }}(value);
33
+ if (numSg == 1u) {
34
+ return sgVal;
35
+ }
36
+ if (sgLid == 0u) {
37
+ wgPartial[sgId] = sgVal;
38
+ }
39
+ workgroupBarrier();
40
+ var total = wgPartial[0];
41
+ for (var i = 1u; i < numSg; i = i + 1u) {
42
+ {{ combine }}
43
+ }
44
+ workgroupBarrier();
45
+ return total;
46
+ }
47
+ {%- endmacro %}{{ emit_reduce("reduce_row", "subgroupAdd", "total = total + wgPartial[i];") }}
48
+ @compute @workgroup_size(WG, 1, 1)
49
+ fn main(@builtin(workgroup_id) wg: vec3<u32>,
50
+ @builtin(num_workgroups) nwg: vec3<u32>,
51
+ @builtin(local_invocation_id) lid: vec3<u32>,
52
+ @builtin(subgroup_invocation_id) sgLid: u32,
53
+ @builtin(subgroup_id) sgId: u32,
54
+ @builtin(num_subgroups) numSg: u32) {
55
+ let row = wg.x + wg.y * nwg.x;
56
+ if (row >= params.rows) {
57
+ return;
58
+ }
59
+ let tid = lid.x;
60
+ let base = row * params.chunkCount; let INIT: f32 = 0.0;{% if source.vec4 %}
61
+ var acc4 = vec4<{{ scalar }}>(INIT);
62
+ for (var c = tid; c < params.chunkCount; c = c + WG) {
63
+ let v = {{ xv }}x[base + c]{{ vx }}; acc4 = acc4 + v; } let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);{% else %}
64
+ var acc = INIT;
65
+ for (var c = tid; c < params.chunkCount; c = c + WG) {
66
+ let v = {{ xv }}x[base + c]{{ vx }}; acc = acc + v; }
67
+ {%- endif %}
68
+ let total = reduce_row(acc, sgLid, sgId, numSg);
69
+ if (tid == 0u) { y[row] = {{ yv }}log(total){{ vy }}; }}
build/webgpu/reduce-row-tree.wgsl.jinja ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Portable one-workgroup-per-row reduction for the Reduce value family.
2
+ // Threads stride a contiguous row, accumulate locally, and fold their values
3
+ // through a shared-memory tree without relying on subgroups.
4
+ //
5
+ // Max, min, and product use the native f32/i32/u32 value type. Additive
6
+ // transforms accumulate float inputs in f32 but retain integer accumulation
7
+ // for sum, L1, and sum-of-squares. Mean, L2, log-sum, and log-sum-exp are
8
+ // f32-only because their finalizers divide, take a square root, or take a log.
9
+ // f16 storage widens through f32 for both accumulation and the shared tree,
10
+ // then narrows only at the final store.
11
+ {% set isVec4 = source.vec4 is defined and source.vec4 %}
12
+ {% set rowIsEmpty = "params.chunkCount == 0u" if isVec4 else "params.cols == 0u" %}
13
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
14
+ {% set scalar = "f32" if castF32 else scalar %}
15
+ {% set xv = ("vec4<f32>(" if source.vec4 else "f32(") if castF32 else "" %}
16
+ {% set vx = ")" if castF32 else "" %}
17
+ {% set yv = "f16(" if castF32 else "" %}
18
+ {% set vy = ")" if castF32 else "" %}
19
+ {% if source.usesF16 is defined and source.usesF16 %}
20
+ enable f16;
21
+ {% endif %}
22
+ {{ env.wgsl.resourceDeclarations }}
23
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
24
+ * evaluation, so f32 identities are constructed at runtime from their
25
+ * IEEE-754 bit patterns. */
26
+
27
+ {% if scalar != "i32" and scalar != "u32" and (source.op == "logsum" or source.op == "logsumexp") %}
28
+ fn negative_infinity() -> f32 {
29
+ var bits = 0xff800000u;
30
+ return bitcast<f32>(bits);
31
+ }
32
+
33
+ {% endif %}
34
+
35
+ const WG: u32 = {{ workgroupSize }}u;
36
+
37
+ {% set is_int = scalar == "i32" or scalar == "u32" %}
38
+ {% set intAdditive = is_int and (source.op == "sum" or source.op == "l1" or source.op == "sumsquare") %}
39
+ {% set accType = scalar if source.op == "max" or source.op == "min" or source.op == "prod" or intAdditive else "f32" %}
40
+ fn identity() -> {{ accType }} { return {{ accType }}(0); }
41
+
42
+ fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
43
+ return a + b;
44
+ }
45
+
46
+ var<workgroup> partial: array<{{ accType }}, WG>;
47
+
48
+ @compute @workgroup_size(WG, 1, 1)
49
+ fn main(@builtin(workgroup_id) wg: vec3<u32>,
50
+ @builtin(num_workgroups) nwg: vec3<u32>,
51
+ @builtin(local_invocation_id) lid: vec3<u32>) {
52
+ let row = wg.x + wg.y * nwg.x;
53
+ if (row >= params.rows) {
54
+ return;
55
+ }
56
+ let tid = lid.x;
57
+ {% if source.vec4 %}
58
+ let base = row * params.chunkCount;
59
+ {% else %}
60
+ let base = row * params.cols;
61
+ {% endif %}
62
+
63
+ if ({{ rowIsEmpty }}) {
64
+ if (tid == 0u) { y[row] = {{ yv }}negative_infinity(){{ vy }}; }
65
+ return;
66
+ }
67
+
68
+ {% if source.vec4 %}
69
+ var acc4 = vec4<{{ accType }}>(identity());
70
+ for (var col = tid; col < params.chunkCount; col = col + WG) {
71
+ let value = {{ xv }}x[base + col]{{ vx }};
72
+ acc4 = acc4 + value;
73
+ }
74
+ let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
75
+ {% else %}
76
+ var acc = identity();
77
+ for (var col = tid; col < params.cols; col = col + WG) {
78
+ let value = {{ xv }}x[base + col]{{ vx }};
79
+ acc = combine(acc, value);
80
+ }
81
+ {% endif %}
82
+ partial[tid] = acc;
83
+ workgroupBarrier();
84
+
85
+ for (var step = WG >> 1u; step > 0u; step = step >> 1u) {
86
+ if (tid < step) {
87
+ partial[tid] = combine(partial[tid], partial[tid + step]);
88
+ }
89
+ workgroupBarrier();
90
+ }
91
+
92
+ if (tid == 0u) {
93
+ let total = partial[0];
94
+ y[row] = {{ yv }}log(total){{ vy }};
95
+ }
96
+ }
build/webgpu/reduce-serial-axis.wgsl.jinja ADDED
@@ -0,0 +1,144 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // Serial one-thread-per-output reduction for the no-feature tier. f16 storage
2
+ // is widened before every accumulation and narrowed only for the final store.
3
+ {% set castF32 = source.castF32 is defined and source.castF32 %}
4
+ {% set yv = "f16(" if castF32 else "" %}
5
+ {% set vy = ")" if castF32 else "" %}
6
+ {% if source.usesF16 is defined and source.usesF16 %}
7
+ enable f16;
8
+ {% endif %}
9
+ {{ env.wgsl.resourceDeclarations }}
10
+ /* Exact max/min reduction identities. WGSL rejects infinity during constant
11
+ * evaluation, so f32 identities are constructed at runtime from their
12
+ * IEEE-754 bit patterns. */
13
+
14
+ fn negative_infinity() -> f32 {
15
+ var bits = 0xff800000u;
16
+ return bitcast<f32>(bits);
17
+ }
18
+
19
+ {% if source.indexing == "rankn" %}
20
+
21
+ fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
22
+ var rem = out_index;
23
+ {% for axis in range(source.outputRank) %}
24
+ {% set out_stride = namespace(value=1) %}
25
+ {% for j in range(axis + 1, source.outputRank) %}
26
+ {% set out_stride.value = out_stride.value * source.outputShape[j] %}
27
+ {% endfor %}
28
+ {% set safe_out_stride = 1 if out_stride.value == 0 else out_stride.value %}
29
+ {% if not source.keepDims or axis != source.axis %}
30
+ let out_coord{{ axis }} = rem / {{ safe_out_stride }}u;
31
+ {% endif %}
32
+ rem = rem % {{ safe_out_stride }}u;
33
+ {% endfor %}
34
+ {% for axis in range(source.rank) %}
35
+ {% if axis == source.axis %}
36
+ let coord{{ axis }} = reduce_index;
37
+ {% elif source.keepDims %}
38
+ let coord{{ axis }} = out_coord{{ axis }};
39
+ {% elif axis < source.axis %}
40
+ let coord{{ axis }} = out_coord{{ axis }};
41
+ {% else %}
42
+ let coord{{ axis }} = out_coord{{ axis - 1 }};
43
+ {% endif %}
44
+ {% endfor %}
45
+ {% set src = namespace(value="coord0") %}
46
+ {% for axis in range(1, source.rank) %}
47
+ {% set src.value = "(" ~ src.value ~ " * " ~ source.dataShape[axis] ~ "u + coord" ~ axis ~ ")" %}
48
+ {% endfor %}
49
+ return {{ src.value }};
50
+ }
51
+ {% endif %}
52
+ {% if source.indexing == "multiaxis" %}
53
+ {% set hasReducedAxis = namespace(value=false) %}
54
+ {% for a in range(source.rank) %}{% if source.reduce[a] %}{% set hasReducedAxis.value = true %}{% endif %}{% endfor %}
55
+
56
+ // One thread per output element walks the Cartesian product of the reduced axes,
57
+ // linearized as reduce_linear. Specialized shapes make every input offset a sum
58
+ // of coordinate-times-constant terms.
59
+ fn input_offset(out_index: u32{% if hasReducedAxis.value %}, reduce_linear: u32{% endif %}) -> u32 {
60
+ var rem = out_index;
61
+ {% for oaxis in range(source.outputRank) %}
62
+ {% set ostride = namespace(value=1) %}
63
+ {% for j in range(oaxis + 1, source.outputRank) %}
64
+ {% set ostride.value = ostride.value * source.outputShape[j] %}
65
+ {% endfor %}
66
+ {% set osafe = 1 if ostride.value == 0 else ostride.value %}
67
+ {% if not source.keepDims or not source.reduce[oaxis] %}
68
+ let out_coord{{ oaxis }} = rem / {{ osafe }}u;
69
+ {% endif %}
70
+ rem = rem % {{ osafe }}u;
71
+ {% endfor %}
72
+ {% if hasReducedAxis.value %}
73
+ var rrem = reduce_linear;
74
+ {% endif %}
75
+ {% for a in range(source.rank) if source.reduce[a] %}
76
+ {% set rstride = namespace(value=1) %}
77
+ {% for b in range(a + 1, source.rank) if source.reduce[b] %}
78
+ {% set rstride.value = rstride.value * source.dataShape[b] %}
79
+ {% endfor %}
80
+ {% set rsafe = 1 if rstride.value == 0 else rstride.value %}
81
+ let red_coord{{ a }} = rrem / {{ rsafe }}u;
82
+ rrem = rrem % {{ rsafe }}u;
83
+ {% endfor %}
84
+ {% set oc = namespace(i=0) %}
85
+ {% for a in range(source.rank) %}
86
+ {% if source.reduce[a] %}
87
+ let coord{{ a }} = red_coord{{ a }};
88
+ {% elif source.keepDims %}
89
+ let coord{{ a }} = out_coord{{ a }};
90
+ {% else %}
91
+ let coord{{ a }} = out_coord{{ oc.i }};
92
+ {% set oc.i = oc.i + 1 %}
93
+ {% endif %}
94
+ {% endfor %}
95
+ {% set src = namespace(value="coord0") %}
96
+ {% for a in range(1, source.rank) %}
97
+ {% set src.value = "(" ~ src.value ~ " * " ~ source.dataShape[a] ~ "u + coord" ~ a ~ ")" %}
98
+ {% endfor %}
99
+ return {{ src.value }};
100
+ }
101
+ {% endif %}
102
+ {% if source.indexing == "multiaxis" %}
103
+ {% set mcount = namespace(value=1) %}
104
+ {% for a in range(source.rank) if source.reduce[a] %}
105
+ {% set mcount.value = mcount.value * source.dataShape[a] %}
106
+ {% endfor %}
107
+ {% set count = mcount.value ~ "u" %}
108
+ {% if hasReducedAxis.value %}
109
+ {% set at = "x[input_offset(i, r)]" %}
110
+ {% else %}
111
+ {% set at = "x[input_offset(i)]" %}
112
+ {% endif %}
113
+ {% elif source.indexing == "rankn" %}
114
+ {% set count = "params.axisDim" %}
115
+ {% set at = "x[input_offset(i, r)]" %}
116
+ {% elif axis == 0 %}
117
+ {% set count = "params.rows" %}
118
+ {% set at = "x[r * params.cols + i]" %}
119
+ {% else %}
120
+ {% set count = "params.cols" %}
121
+ {% set at = "x[i * params.cols + r]" %}
122
+ {% endif %}
123
+ {% if castF32 %}
124
+ {% set at = "f32(" ~ at ~ ")" %}
125
+ {% endif %}
126
+
127
+ @compute @workgroup_size({{ reduceWorkgroupSize }})
128
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
129
+ // 2D-folded flat index: gid.y carries the high bits past the
130
+ // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
131
+ let i = gid.x + gid.y * nwg.x * {{ reduceWorkgroupSize }}u;
132
+ if (i >= params.outCount) {
133
+ return;
134
+ }
135
+ if ({{ count }} == 0u) {
136
+ y[i] = {{ yv }}negative_infinity(){{ vy }};
137
+ return;
138
+ }
139
+ var acc = {{ "1.0" if source.op == "prod" else "0.0" }};
140
+ for (var r = 0u; r < {{ count }}; r = r + 1u) {
141
+ acc = acc + {{ at }};
142
+ }
143
+ y[i] = {{ yv }}log(acc){{ vy }};
144
+ }
build/webgpu/test.json ADDED
@@ -0,0 +1,746 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "op": "ai.onnx.ReduceLogSum",
3
+ "fixtureArrays": {
4
+ "onnx_backend_reduce_log_sum_input_x": [0.54881352186203, 0.7151893377304077, 0.6027633547782898, 0.5448831915855408, 0.42365479469299316, 0.6458941102027893, 0.4375872015953064, 0.891772985458374, 0.9636627435684204, 0.3834415078163147, 0.7917250394821167, 0.5288949012756348, 0.5680445432662964, 0.9255966544151306, 0.07103605568408966, 0.08712930232286453, 0.020218396559357643, 0.832619845867157, 0.7781567573547363, 0.8700121641159058, 0.978618323802948, 0.7991585731506348, 0.4614793658256531, 0.7805292010307312, 0.11827442795038223, 0.6399210095405579, 0.14335328340530396, 0.9446688890457153, 0.5218483209609985, 0.4146619439125061, 0.26455560326576233, 0.7742336988449097, 0.4561503231525421, 0.568433940410614, 0.018789799883961678, 0.6176354885101318, 0.6120957136154175, 0.6169340014457703, 0.9437480568885803, 0.681820273399353, 0.35950788855552673, 0.43703195452690125, 0.6976311802864075, 0.0602254718542099, 0.6667667031288147, 0.670637845993042, 0.21038256585597992, 0.12892629206180573, 0.31542834639549255, 0.36371076107025146, 0.5701967477798462, 0.4386015236377716, 0.9883738160133362, 0.10204481333494186, 0.20887675881385803, 0.16130951046943665, 0.6531082987785339, 0.25329160690307617, 0.4663107693195343, 0.24442559480667114]
5
+ },
6
+ "cases": [
7
+ {
8
+ "name": "all_axes_flat_rank1_boundary_8192",
9
+ "provenance": {
10
+ "notes": "The parallel full-reduction threshold must supersede the rank1 serial and row-reduction fallbacks."
11
+ },
12
+ "attrs": { "axes": [0], "keepdims": 0 },
13
+ "inputs": { "x": { "dtype": "float32", "shape": [8192], "data": { "kind": "constant", "value": 1.0 } } },
14
+ "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.0001 } }
15
+ },
16
+ {
17
+ "name": "all_axes_flat_fullreduce_32x32x32_keepdims",
18
+ "attrs": { "keepdims": 1 },
19
+ "inputs": {
20
+ "x": { "dtype": "float32", "shape": [32, 32, 32], "data": { "kind": "linspace", "start": 0.1, "end": 2.0 } }
21
+ },
22
+ "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 1], "tolerance": 0.0005, "relTolerance": 0.0001 } }
23
+ },
24
+ {
25
+ "name": "dispatch_cliff_axis1_rank2",
26
+ "attrs": { "axes": [1], "keepdims": 0 },
27
+ "inputs": {
28
+ "x": {
29
+ "dtype": "float32",
30
+ "shape": [16776961, 1],
31
+ "data": { "kind": "cycle", "values": [1.0, 2.0, 0.5, 3.25, 1.5, 2.75, 0.75, 4.0] }
32
+ }
33
+ },
34
+ "outputs": { "y": { "dtype": "float32", "shape": [16776961], "tolerance": 0.00001 } }
35
+ },
36
+ {
37
+ "name": "axis0",
38
+ "attrs": { "axes": [0], "keepdims": 0 },
39
+ "inputs": {
40
+ "x": {
41
+ "dtype": "float32",
42
+ "shape": [2, 3],
43
+ "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
44
+ }
45
+ },
46
+ "outputs": { "y": { "dtype": "float32", "shape": [3], "tolerance": 0.000001 } }
47
+ },
48
+ {
49
+ "name": "axis0_tiled_64x32",
50
+ "attrs": { "axes": [0], "keepdims": 0 },
51
+ "inputs": {
52
+ "x": {
53
+ "dtype": "float32",
54
+ "shape": [64, 32],
55
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
56
+ }
57
+ },
58
+ "outputs": { "y": { "dtype": "float32", "shape": [32], "tolerance": 0.00002 } }
59
+ },
60
+ {
61
+ "name": "axis1",
62
+ "attrs": { "axes": [1], "keepdims": 0 },
63
+ "inputs": {
64
+ "x": {
65
+ "dtype": "float32",
66
+ "shape": [2, 3],
67
+ "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
68
+ }
69
+ },
70
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.000001 } }
71
+ },
72
+ {
73
+ "name": "f32_axis1_parallel_cancellation_finite_logsum_gpu_gap",
74
+ "skipGpu": {
75
+ "category": "todo",
76
+ "reason": "The current parallel reduction changes the fixture's required sequential evaluation order, so f32 rounding is not bit-exact. An order-preserving reduction route can implement this behavior."
77
+ },
78
+ "provenance": {
79
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
80
+ "test": "ReductionOpTest.ReduceLogSum",
81
+ "notes": "Serial float32 row summation of repeated [1e20, 1, -1e20, 1] blocks keeps only the final trailing 1 before log; a parallel tree can preserve one small term per block and return log(256)."
82
+ },
83
+ "attrs": { "axes": [1], "keepdims": 0 },
84
+ "inputs": {
85
+ "x": {
86
+ "dtype": "float32",
87
+ "shape": [1, 1024],
88
+ "data": { "kind": "cycle", "values": [100000000000000000000.0, 1.0, -100000000000000000000.0, 1.0] }
89
+ }
90
+ },
91
+ "outputs": { "y": { "dtype": "float32", "shape": [1], "tolerance": 0 } }
92
+ },
93
+ {
94
+ "name": "f32_subnormal_axis1_logsum_gpu_gap",
95
+ "skipGpu": {
96
+ "category": "permanent",
97
+ "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."
98
+ },
99
+ "provenance": {
100
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
101
+ "test": "ReductionOpTest.ReduceLogSum",
102
+ "notes": "The row sum is finite subnormal, so ReduceLogSum should produce a large finite negative log rather than -Infinity."
103
+ },
104
+ "attrs": { "axes": [1], "keepdims": 0 },
105
+ "inputs": {
106
+ "x": {
107
+ "dtype": "float32",
108
+ "shape": [2, 3],
109
+ "data": { "kind": "values", "values": [1e-40, 1e-40, 1e-40, 1e-39, 1e-39, 1e-39] }
110
+ }
111
+ },
112
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.000001 } }
113
+ },
114
+ {
115
+ "name": "f32_subnormal_axis0_logsum_gpu_gap",
116
+ "skipGpu": {
117
+ "category": "permanent",
118
+ "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."
119
+ },
120
+ "provenance": {
121
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
122
+ "test": "ReductionOpTest.ReduceLogSum",
123
+ "notes": "Axis-0 companion: the reduced sums are finite subnormal, so ReduceLogSum should stay finite rather than returning -Infinity."
124
+ },
125
+ "attrs": { "axes": [0], "keepdims": 0 },
126
+ "inputs": {
127
+ "x": {
128
+ "dtype": "float32",
129
+ "shape": [2, 3],
130
+ "data": { "kind": "values", "values": [1e-40, 1e-39, 1e-38, 1e-40, 2e-39, 2e-38] }
131
+ }
132
+ },
133
+ "outputs": { "y": { "dtype": "float32", "shape": [3], "tolerance": 0.00001 } }
134
+ },
135
+ {
136
+ "name": "f32_subnormal_axis0_tilecols_logsum_gpu_gap",
137
+ "skipGpu": {
138
+ "category": "permanent",
139
+ "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."
140
+ },
141
+ "provenance": {
142
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
143
+ "test": "ReductionOpTest.ReduceLogSum",
144
+ "notes": "Tile-column axis-0 path: many positive subnormal inputs sum to a finite subnormal value whose log is finite."
145
+ },
146
+ "attrs": { "axes": [0], "keepdims": 0 },
147
+ "inputs": { "x": { "dtype": "float32", "shape": [64, 16], "data": { "kind": "constant", "value": 1e-40 } } },
148
+ "outputs": { "y": { "dtype": "float32", "shape": [16], "tolerance": 0.00001 } }
149
+ },
150
+ {
151
+ "name": "f32_subnormal_last_axis_vec4_logsum_gpu_gap",
152
+ "skipGpu": {
153
+ "category": "permanent",
154
+ "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."
155
+ },
156
+ "provenance": {
157
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
158
+ "test": "ReductionOpTest.ReduceLogSum",
159
+ "notes": "Vec4 last-axis path: finite subnormal row sums should produce finite logs."
160
+ },
161
+ "attrs": { "axes": [-1], "keepdims": 0 },
162
+ "inputs": {
163
+ "x": {
164
+ "dtype": "float32",
165
+ "shape": [2, 4],
166
+ "data": { "kind": "values", "values": [1e-40, 1e-40, 1e-40, 1e-40, 1e-39, 1e-39, 1e-39, 1e-39] }
167
+ }
168
+ },
169
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.00001 } }
170
+ },
171
+ {
172
+ "name": "f32_subnormal_rank3_axis1_logsum_gpu_gap",
173
+ "skipGpu": {
174
+ "category": "permanent",
175
+ "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."
176
+ },
177
+ "provenance": {
178
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
179
+ "test": "ReductionOpTest.ReduceLogSum",
180
+ "notes": "Rank-3 non-last-axis path: finite subnormal sums along axis 1 should produce finite logs."
181
+ },
182
+ "attrs": { "axes": [1], "keepdims": 0 },
183
+ "inputs": {
184
+ "x": {
185
+ "dtype": "float32",
186
+ "shape": [2, 2, 2],
187
+ "data": { "kind": "values", "values": [1e-40, 2e-40, 1e-40, 2e-40, 1e-39, 2e-39, 1e-39, 2e-39] }
188
+ }
189
+ },
190
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.00001 } }
191
+ },
192
+ {
193
+ "name": "axis1_empty_cols_negative_infinity",
194
+ "attrs": { "axes": [1], "keepdims": 0 },
195
+ "inputs": { "x": { "dtype": "float32", "shape": [2, 0], "data": { "kind": "values", "values": [] } } },
196
+ "outputs": {
197
+ "y": {
198
+ "dtype": "float32",
199
+ "shape": [2],
200
+ "tolerance": 0,
201
+ "data": { "kind": "values", "values": ["-Infinity", "-Infinity"] }
202
+ }
203
+ }
204
+ },
205
+ {
206
+ "name": "axis0_empty_rows_negative_infinity",
207
+ "attrs": { "axes": [0], "keepdims": 0 },
208
+ "inputs": { "x": { "dtype": "float32", "shape": [0, 3], "data": { "kind": "values", "values": [] } } },
209
+ "outputs": {
210
+ "y": {
211
+ "dtype": "float32",
212
+ "shape": [3],
213
+ "tolerance": 0,
214
+ "data": { "kind": "values", "values": ["-Infinity", "-Infinity", "-Infinity"] }
215
+ }
216
+ }
217
+ },
218
+ {
219
+ "name": "f16_empty_reduction_negative_infinity",
220
+ "attrs": { "axes": [1], "keepdims": 0 },
221
+ "inputs": { "x": { "dtype": "float16", "shape": [1, 0], "data": { "kind": "values", "values": [] } } },
222
+ "outputs": {
223
+ "y": { "dtype": "float16", "shape": [1], "tolerance": 0, "data": { "kind": "values", "values": ["-Infinity"] } }
224
+ }
225
+ },
226
+ {
227
+ "name": "axis1_zero_rows_noop",
228
+ "attrs": { "axes": [1], "keepdims": 0 },
229
+ "inputs": { "x": { "dtype": "float32", "shape": [0, 3], "data": { "kind": "values", "values": [] } } },
230
+ "outputs": { "y": { "dtype": "float32", "shape": [0], "tolerance": 0 } }
231
+ },
232
+ {
233
+ "name": "axis_minus_one",
234
+ "attrs": { "axes": [-1], "keepdims": 0 },
235
+ "inputs": {
236
+ "x": {
237
+ "dtype": "float32",
238
+ "shape": [2, 3],
239
+ "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
240
+ }
241
+ },
242
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.000001 } }
243
+ },
244
+ {
245
+ "name": "rank3_axis_minus_one_keepdims",
246
+ "attrs": { "axes": [-1], "keepdims": 1 },
247
+ "inputs": {
248
+ "x": {
249
+ "dtype": "float32",
250
+ "shape": [2, 3, 4],
251
+ "data": {
252
+ "kind": "values",
253
+ "values": [1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5, 5.0, 6.0, 7.0, 8.0, 2.0, 4.0, 6.0, 8.0, 1.25, 2.25, 3.25, 4.25, 0.75, 1.75, 2.75, 3.75]
254
+ }
255
+ }
256
+ },
257
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 1], "tolerance": 0.000001 } }
258
+ },
259
+ {
260
+ "name": "rank1_axis0_scalar_output",
261
+ "attrs": { "axes": [0], "keepdims": 0 },
262
+ "inputs": {
263
+ "x": { "dtype": "float32", "shape": [5], "data": { "kind": "values", "values": [1.0, 2.0, 0.5, 4.0, 8.0] } }
264
+ },
265
+ "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.000001 } }
266
+ },
267
+ {
268
+ "name": "ort_axis1_rank3_keepdims",
269
+ "provenance": {
270
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
271
+ "test": "ReductionOpTest.ReduceLogSum"
272
+ },
273
+ "attrs": { "axes": [1], "keepdims": 1 },
274
+ "inputs": {
275
+ "x": {
276
+ "dtype": "float32",
277
+ "shape": [3, 2, 2],
278
+ "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] }
279
+ }
280
+ },
281
+ "outputs": { "y": { "dtype": "float32", "shape": [3, 1, 2], "tolerance": 0.00001 } }
282
+ },
283
+ {
284
+ "name": "ort_axis2_singleton_keepdims_noop",
285
+ "provenance": {
286
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
287
+ "test": "ReductionOpTest.ReduceLogSum_samesize"
288
+ },
289
+ "attrs": { "axes": [2], "keepdims": 1 },
290
+ "inputs": {
291
+ "x": {
292
+ "dtype": "float32",
293
+ "shape": [3, 2, 1],
294
+ "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
295
+ }
296
+ },
297
+ "outputs": { "y": { "dtype": "float32", "shape": [3, 2, 1], "tolerance": 0.00001 } }
298
+ },
299
+ {
300
+ "name": "ort_axis0_rank1_scalar",
301
+ "provenance": {
302
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
303
+ "test": "ReductionOpTest.ReduceLogSum_do_not_keepdims_2"
304
+ },
305
+ "attrs": { "axes": [0], "keepdims": 0 },
306
+ "inputs": { "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1.0, 2.0, 3.0] } } },
307
+ "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.00001 } }
308
+ },
309
+ {
310
+ "name": "ort_rank0_scalar",
311
+ "provenance": {
312
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
313
+ "test": "ReductionOpTest.ReduceLogSum0DTensor"
314
+ },
315
+ "inputs": { "x": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [2.0] } } },
316
+ "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.000001 } }
317
+ },
318
+ {
319
+ "name": "ort_mixed_infinities_axis1_nan_rows",
320
+ "provenance": {
321
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
322
+ "test": "ReductionOpTest.ReduceInfLogSum"
323
+ },
324
+ "attrs": { "axes": [1], "keepdims": 0 },
325
+ "inputs": {
326
+ "x": {
327
+ "dtype": "float32",
328
+ "shape": [6, 2],
329
+ "data": {
330
+ "kind": "values",
331
+ "values": [1.0, "Infinity", "Infinity", 1.0, "Infinity", "-Infinity", "-Infinity", "Infinity", 1.0, "-Infinity", "-Infinity", 1.0]
332
+ }
333
+ }
334
+ },
335
+ "outputs": { "y": { "dtype": "float32", "shape": [6], "tolerance": 0, "allowNaN": true } }
336
+ },
337
+ {
338
+ "name": "onnx_backend_reduce_log_sum_empty_set",
339
+ "attrs": { "keepdims": 1, "axes": [1] },
340
+ "inputs": { "x": { "dtype": "float32", "shape": [2, 0, 4], "data": { "kind": "values", "values": [] } } },
341
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 1, 4] } },
342
+ "provenance": {
343
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_reduce_log_sum_empty_set",
344
+ "notes": "The ONNX int64 axes input is materialized as this compile-time axes list."
345
+ }
346
+ },
347
+ {
348
+ "name": "onnx_backend_reduce_log_sum_negative_axes",
349
+ "attrs": { "axes": [-2], "keepdims": 1 },
350
+ "inputs": {
351
+ "x": {
352
+ "dtype": "float32",
353
+ "shape": [3, 4, 5],
354
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_reduce_log_sum_input_x" } }
355
+ }
356
+ },
357
+ "outputs": { "y": { "dtype": "float32", "shape": [3, 1, 5] } },
358
+ "provenance": {
359
+ "source": "cmake/external/onnx/onnx/backend/test/data/node/test_reduce_log_sum_negative_axes",
360
+ "notes": "The ONNX int64 axes input is materialized as this compile-time axes list."
361
+ }
362
+ },
363
+ {
364
+ "name": "default_axes_rank3_no_keepdims_scalar",
365
+ "attrs": { "keepdims": 0 },
366
+ "inputs": {
367
+ "x": {
368
+ "dtype": "float32",
369
+ "shape": [3, 2, 2],
370
+ "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] }
371
+ }
372
+ },
373
+ "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.000001 } }
374
+ },
375
+ {
376
+ "name": "onnx_backend_reduce_log_sum_default_axes_keepdims_random",
377
+ "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_reduce_log_sum_default" },
378
+ "attrs": { "keepdims": 1 },
379
+ "inputs": {
380
+ "x": {
381
+ "dtype": "float32",
382
+ "shape": [3, 4, 5],
383
+ "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_reduce_log_sum_input_x" } }
384
+ }
385
+ },
386
+ "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 1], "tolerance": 0.000001 } }
387
+ },
388
+ {
389
+ "name": "subgroup_vec4_last_axis_2x256",
390
+ "attrs": { "axes": [-1], "keepdims": 0 },
391
+ "inputs": {
392
+ "x": {
393
+ "dtype": "float32",
394
+ "shape": [2, 256],
395
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "offset": 2.0 }
396
+ }
397
+ },
398
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.0002, "relTolerance": 0.0001 } }
399
+ },
400
+ {
401
+ "name": "subgroup_scalar_last_axis_2x65",
402
+ "attrs": { "axes": [1], "keepdims": 0 },
403
+ "inputs": {
404
+ "x": {
405
+ "dtype": "float32",
406
+ "shape": [2, 65],
407
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.11, "offset": 2.0 }
408
+ }
409
+ },
410
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.0002, "relTolerance": 0.0001 } }
411
+ },
412
+ {
413
+ "name": "ort_noop_empty_axes_2d_elementwise_log",
414
+ "provenance": {
415
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
416
+ "test": "ReductionOpTest.ReduceLogSum_NoopWithEmptyAxes_2D_ElementwiseLog"
417
+ },
418
+ "attrs": { "noop_with_empty_axes": 1 },
419
+ "inputs": {
420
+ "x": {
421
+ "dtype": "float32",
422
+ "shape": [2, 2],
423
+ "data": { "kind": "values", "values": [2.7182817, 7.389056, 1.6487213, 20.085537] }
424
+ }
425
+ },
426
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.000001 } }
427
+ },
428
+ {
429
+ "name": "ort_noop_empty_axes_scalar_log",
430
+ "provenance": {
431
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
432
+ "test": "ReductionOpTest.ReduceLogSum_NoopWithEmptyAxes_Scalar"
433
+ },
434
+ "attrs": { "noop_with_empty_axes": 1 },
435
+ "inputs": { "x": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [2.7182817] } } },
436
+ "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.000001 } }
437
+ },
438
+ {
439
+ "name": "ort_noop_empty_axes_3d_elementwise_log",
440
+ "provenance": {
441
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
442
+ "test": "ReductionOpTest.ReduceLogSum_NoopWithEmptyAxes_3D_ElementwiseLog"
443
+ },
444
+ "attrs": { "noop_with_empty_axes": 1 },
445
+ "inputs": {
446
+ "x": {
447
+ "dtype": "float32",
448
+ "shape": [2, 1, 3],
449
+ "data": { "kind": "values", "values": [2.7182817, 7.389056, 1.6487213, 20.085537, 54.59815, 148.41316] }
450
+ }
451
+ },
452
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 1, 3], "tolerance": 0.000001 } }
453
+ },
454
+ {
455
+ "name": "ort_float_multi_axis_no_keepdims",
456
+ "provenance": {
457
+ "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
458
+ "test": "ReductionOpTest.ReduceLogSumAxes01",
459
+ "notes": "Compact positive tensor covering the same multi-axis ReduceLogSum surface as the ORT case."
460
+ },
461
+ "attrs": { "axes": [0, 1], "keepdims": 0 },
462
+ "inputs": {
463
+ "x": {
464
+ "dtype": "float32",
465
+ "shape": [3, 2, 2],
466
+ "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] }
467
+ }
468
+ },
469
+ "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0.000001 } }
470
+ },
471
+ {
472
+ "name": "rank3_lastaxis_cols1024_tree_nosubgroup",
473
+ "attrs": { "axes": [2], "keepdims": 0 },
474
+ "inputs": {
475
+ "x": {
476
+ "dtype": "float32",
477
+ "shape": [2, 2, 1024],
478
+ "data": {
479
+ "kind": "cycle",
480
+ "values": [1.0, 2.0, 0.5, 3.25, 1.5, 2.0, 0.75, 4.0, 3.5, 1.25, 0.25, 2.25, 5.0, 4.0, 2.75, 1.0]
481
+ }
482
+ }
483
+ },
484
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.00001 } }
485
+ },
486
+ {
487
+ "name": "axis0_splitk_8192x32",
488
+ "attrs": { "axes": [0], "keepdims": 0 },
489
+ "inputs": {
490
+ "x": {
491
+ "dtype": "float32",
492
+ "shape": [8192, 32],
493
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
494
+ }
495
+ },
496
+ "outputs": { "y": { "dtype": "float32", "shape": [32], "tolerance": 0.0001 } }
497
+ },
498
+ {
499
+ "name": "axis0_splitk_8192x48_keepdims",
500
+ "attrs": { "axes": [0], "keepdims": 1 },
501
+ "inputs": {
502
+ "x": {
503
+ "dtype": "float32",
504
+ "shape": [8192, 48],
505
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
506
+ }
507
+ },
508
+ "outputs": { "y": { "dtype": "float32", "shape": [1, 48], "tolerance": 0.0001 } }
509
+ },
510
+ {
511
+ "name": "rank4_multi_axis_12_keepdims",
512
+ "attrs": { "axes": [1, 2], "keepdims": 1 },
513
+ "inputs": {
514
+ "x": {
515
+ "dtype": "float32",
516
+ "shape": [2, 3, 2, 2],
517
+ "data": {
518
+ "kind": "values",
519
+ "values": [1.0, 2.0, 0.5, 1.5, 3.0, 0.25, 2.5, 1.25, 0.75, 2.0, 1.0, 0.5, 1.5, 2.5, 0.25, 1.75, 3.0, 0.5, 2.0, 1.0, 0.75, 1.25, 1.5, 2.0]
520
+ }
521
+ }
522
+ },
523
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 1, 1, 2], "tolerance": 0.0001 } }
524
+ },
525
+ {
526
+ "name": "row_with_negative_sum_produces_nan",
527
+ "attrs": { "axes": [1], "keepdims": 0 },
528
+ "inputs": {
529
+ "x": {
530
+ "dtype": "float32",
531
+ "shape": [3, 3],
532
+ "data": { "kind": "values", "values": [1.0, 2.0, 3.0, -4.0, -5.0, -6.0, 1.0, -2.0, 0.5] }
533
+ }
534
+ },
535
+ "outputs": { "y": { "dtype": "float32", "shape": [3], "allowNaN": true, "tolerance": 0 } },
536
+ "provenance": {
537
+ "notes": "Row 0 sums to 6.0 (finite positive, log(6)=1.7917...), row 1 sums to -15.0 (log(-15)=NaN), row 2 sums to -0.5 (log(-0.5)=NaN). Only row 0 is finite; rows 1 and 2 are NaN. Evaluated against the trusted TS reference."
538
+ }
539
+ },
540
+ {
541
+ "name": "row_summing_to_zero_produces_neg_inf",
542
+ "attrs": { "axes": [1], "keepdims": 0 },
543
+ "inputs": {
544
+ "x": {
545
+ "dtype": "float32",
546
+ "shape": [2, 4],
547
+ "data": { "kind": "values", "values": [1.0, -1.0, 2.0, -2.0, 0.0, 0.0, 0.0, 0.0] }
548
+ }
549
+ },
550
+ "outputs": {
551
+ "y": {
552
+ "dtype": "float32",
553
+ "shape": [2],
554
+ "data": { "kind": "values", "values": ["-Infinity", "-Infinity"] },
555
+ "tolerance": 0
556
+ }
557
+ }
558
+ },
559
+ {
560
+ "name": "all_axes_flat_log_applied_after_combine",
561
+ "attrs": { "keepdims": 1 },
562
+ "inputs": { "x": { "dtype": "float32", "shape": [1, 8192], "data": { "kind": "constant", "value": 1.0 } } },
563
+ "outputs": { "y": { "dtype": "float32", "shape": [1, 1], "tolerance": 0.0001 } }
564
+ },
565
+ {
566
+ "name": "axis0_splitk_log_applied_after_combine",
567
+ "attrs": { "axes": [0], "keepdims": 0 },
568
+ "inputs": { "x": { "dtype": "float32", "shape": [8192, 16], "data": { "kind": "constant", "value": 1.0 } } },
569
+ "outputs": { "y": { "dtype": "float32", "shape": [16], "tolerance": 0.0001 } }
570
+ },
571
+ {
572
+ "name": "axis0_narrow_f32_8192x3_splitk_guard_lock",
573
+ "provenance": {
574
+ "notes": "Compact lock below the historical 16-column split-K guard. Constant ones verify log is applied once after combining all partial sums."
575
+ },
576
+ "attrs": { "axes": [0], "keepdims": 0 },
577
+ "inputs": { "x": { "dtype": "float32", "shape": [8192, 3], "data": { "kind": "constant", "value": 1.0 } } },
578
+ "outputs": { "y": { "dtype": "float32", "shape": [3], "tolerance": 0.0001, "relTolerance": 0.0001 } }
579
+ },
580
+ {
581
+ "name": "contiguous_suffix_axes12_parallel",
582
+ "provenance": {
583
+ "notes": "Contiguous axes {1,2} exercise the shared cooperative suffix reduction instead of one serial lane per output."
584
+ },
585
+ "attrs": { "axes": [1, 2], "keepdims": 1 },
586
+ "inputs": {
587
+ "x": {
588
+ "dtype": "float32",
589
+ "shape": [3, 16, 16],
590
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.07, "offset": 1.0, "scale": 0.001 }
591
+ }
592
+ },
593
+ "outputs": { "y": { "dtype": "float32", "shape": [3, 1, 1], "tolerance": 0.00001 } }
594
+ },
595
+ {
596
+ "name": "axis_split_rank3_axis1_2x8192x4",
597
+ "attrs": { "axes": [1], "keepdims": 0 },
598
+ "inputs": {
599
+ "x": {
600
+ "dtype": "float32",
601
+ "shape": [2, 8192, 4],
602
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
603
+ }
604
+ },
605
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.0001 } }
606
+ },
607
+ {
608
+ "name": "f16_axis_split_tiled_narrow_2x8192x4",
609
+ "attrs": { "axes": [1], "keepdims": 0 },
610
+ "inputs": {
611
+ "x": {
612
+ "dtype": "float16",
613
+ "shape": [2, 8192, 4],
614
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
615
+ }
616
+ },
617
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 4], "tolerance": 0.05, "relTolerance": 0.002 } }
618
+ },
619
+ {
620
+ "name": "axis_split_rank3_axis1_wide_2x8192x32",
621
+ "attrs": { "axes": [1], "keepdims": 0 },
622
+ "inputs": {
623
+ "x": {
624
+ "dtype": "float32",
625
+ "shape": [2, 8192, 32],
626
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
627
+ }
628
+ },
629
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 32], "tolerance": 0.0001 } }
630
+ },
631
+ {
632
+ "name": "f16_axis_split_wide_2x8192x32",
633
+ "attrs": { "axes": [1], "keepdims": 0 },
634
+ "inputs": {
635
+ "x": {
636
+ "dtype": "float16",
637
+ "shape": [2, 8192, 32],
638
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
639
+ }
640
+ },
641
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 32], "tolerance": 0.05, "relTolerance": 0.002 } }
642
+ },
643
+ {
644
+ "name": "f16_rank3_axis1_serial",
645
+ "attrs": { "axes": [1], "keepdims": 1 },
646
+ "inputs": {
647
+ "x": {
648
+ "dtype": "float16",
649
+ "shape": [3, 2, 2],
650
+ "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] }
651
+ }
652
+ },
653
+ "outputs": { "y": { "dtype": "float16", "shape": [3, 1, 2], "tolerance": 0.02 } }
654
+ },
655
+ {
656
+ "name": "f16_all_axes",
657
+ "attrs": { "axes": [0], "keepdims": 0 },
658
+ "inputs": { "x": { "dtype": "float16", "shape": [8192], "data": { "kind": "constant", "value": 1.0 } } },
659
+ "outputs": { "y": { "dtype": "float16", "shape": [], "tolerance": 0.02 } }
660
+ },
661
+ {
662
+ "name": "f16_axis0_splitk_8192x8",
663
+ "attrs": { "axes": [0], "keepdims": 0 },
664
+ "inputs": {
665
+ "x": {
666
+ "dtype": "float16",
667
+ "shape": [8192, 8],
668
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
669
+ }
670
+ },
671
+ "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.002 } }
672
+ },
673
+ {
674
+ "name": "f16_last_axis_vec4_8x1024",
675
+ "attrs": { "axes": [1], "keepdims": 0 },
676
+ "inputs": {
677
+ "x": {
678
+ "dtype": "float16",
679
+ "shape": [8, 1024],
680
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
681
+ }
682
+ },
683
+ "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.002 } }
684
+ },
685
+ {
686
+ "name": "f16_last_axis_scalar_8x1023",
687
+ "attrs": { "axes": [1], "keepdims": 0 },
688
+ "inputs": {
689
+ "x": {
690
+ "dtype": "float16",
691
+ "shape": [8, 1023],
692
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
693
+ }
694
+ },
695
+ "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.002 } }
696
+ },
697
+ {
698
+ "name": "f16_all_axes_flat_65543",
699
+ "attrs": { "keepdims": 0 },
700
+ "inputs": {
701
+ "x": {
702
+ "dtype": "float16",
703
+ "shape": [65543],
704
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
705
+ }
706
+ },
707
+ "outputs": { "y": { "dtype": "float16", "shape": [], "tolerance": 0.05, "relTolerance": 0.002 } }
708
+ },
709
+ {
710
+ "name": "f16_suffix_vec4_4x8x128",
711
+ "attrs": { "axes": [1, 2], "keepdims": 0 },
712
+ "inputs": {
713
+ "x": {
714
+ "dtype": "float16",
715
+ "shape": [4, 8, 128],
716
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
717
+ }
718
+ },
719
+ "outputs": { "y": { "dtype": "float16", "shape": [4], "tolerance": 0.05, "relTolerance": 0.002 } }
720
+ },
721
+ {
722
+ "name": "f16_suffix_scalar_4x7x37",
723
+ "attrs": { "axes": [1, 2], "keepdims": 0 },
724
+ "inputs": {
725
+ "x": {
726
+ "dtype": "float16",
727
+ "shape": [4, 7, 37],
728
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
729
+ }
730
+ },
731
+ "outputs": { "y": { "dtype": "float16", "shape": [4], "tolerance": 0.05, "relTolerance": 0.002 } }
732
+ },
733
+ {
734
+ "name": "f16_axis0_tilecols_4096x64",
735
+ "attrs": { "axes": [0], "keepdims": 0 },
736
+ "inputs": {
737
+ "x": {
738
+ "dtype": "float16",
739
+ "shape": [4096, 64],
740
+ "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 2.0 }
741
+ }
742
+ },
743
+ "outputs": { "y": { "dtype": "float16", "shape": [64], "tolerance": 0.05, "relTolerance": 0.002 } }
744
+ }
745
+ ]
746
+ }