sync 91d990483a17

Browse files

Files changed (15) hide show

README.md +23 -10
build/webgpu/bench.json +41 -4
build/webgpu/manifest.json +720 -818
build/webgpu/metadata.json +50 -14
build/webgpu/reduce-axis-split-reduce.wgsl.jinja +76 -10
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +91 -7
build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja +81 -14
build/webgpu/reduce-axis0-tilecols.wgsl.jinja +183 -14
build/webgpu/reduce-flat-partial.wgsl.jinja +67 -8
build/webgpu/reduce-multi-axis-coop.wgsl.jinja +130 -0
build/webgpu/reduce-row-subgroup-rows.wgsl.jinja +167 -0
build/webgpu/reduce-row-subgroup.wgsl.jinja +169 -21
build/webgpu/reduce-row-tree.wgsl.jinja +228 -19
build/webgpu/reduce-serial-axis.wgsl.jinja +218 -73
build/webgpu/test.json +192 -20

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `ReduceSum` spec](https://onnx.ai/onnx/operators/onnx__ReduceSum.h
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `data` | `x` | `T` | — | — | The input tensor to reduce. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `reduced` | `y` | `T` | derived | — | The summed output tensor, with reduced dimensions either kept as size 1 or removed. | required |
 ## Attributes
@@ -34,9 +34,9 @@ Default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `keepdims` | `1` | If 1 (default), retains reduced dimensions with size 1; if 0, removes them from the output shape. |
 | `noop_with_empty_axes` | `0` | When axes is empty, if 0 (default) reduce over all axes; if 1, treat as a no-op identity and return the input unchanged. |
-| `axes` | `[]` | Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`. |
 ## Type constraints
@@ -44,13 +44,21 @@ Default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16`, `int32` |
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -60,21 +68,26 @@ Some implementation variants require `subgroups`. These are route-specific capab
 - [`reduce-axis0-splitk-reduce.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja)
 - [`reduce-axis0-tilecols.wgsl.jinja`](build/webgpu/reduce-axis0-tilecols.wgsl.jinja)
 - [`reduce-flat-partial.wgsl.jinja`](build/webgpu/reduce-flat-partial.wgsl.jinja)
 - [`reduce-row-subgroup.wgsl.jinja`](build/webgpu/reduce-row-subgroup.wgsl.jinja)
 - [`reduce-row-tree.wgsl.jinja`](build/webgpu/reduce-row-tree.wgsl.jinja)
 - [`reduce-serial-axis.wgsl.jinja`](build/webgpu/reduce-serial-axis.wgsl.jinja)
 ## Use with `@huggingface/kernels`
-The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
-The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
 - `y`
-Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `data` | `T` | — | — | The input tensor to reduce. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `reduced` | `T` | derived | — | The summed output tensor, with reduced dimensions either kept as size 1 or removed. | required |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
+| `axes` | `[]` | Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`. |
 | `keepdims` | `1` | If 1 (default), retains reduced dimensions with size 1; if 0, removes them from the output shape. |
 | `noop_with_empty_axes` | `0` | When axes is empty, if 0 (default) reduce over all axes; if 1, treat as a no-op identity and return the input unchanged. |
 ## Type constraints
 | --- | --- |
 | `T` | `float32`, `float16`, `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `axis0_splitk_i32` — Partitions a long rank-two axis-zero integer reduction across workgroups and combines exact int32 partials. It applies when the reduced row dimension is too large for one pass to expose enough parallelism.
+- `subgroup_last_axis_vec4` — Reduces each contiguous last-axis row with subgroup collectives and vec4-packed reads.
+- `subgroup_last_axis` — Reduces each contiguous last-axis row with subgroup collectives and scalar reads for an unaligned row width.
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`reduce-axis0-splitk-reduce.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja)
 - [`reduce-axis0-tilecols.wgsl.jinja`](build/webgpu/reduce-axis0-tilecols.wgsl.jinja)
 - [`reduce-flat-partial.wgsl.jinja`](build/webgpu/reduce-flat-partial.wgsl.jinja)
+- [`reduce-multi-axis-coop.wgsl.jinja`](build/webgpu/reduce-multi-axis-coop.wgsl.jinja)
+- [`reduce-row-subgroup-rows.wgsl.jinja`](build/webgpu/reduce-row-subgroup-rows.wgsl.jinja)
 - [`reduce-row-subgroup.wgsl.jinja`](build/webgpu/reduce-row-subgroup.wgsl.jinja)
 - [`reduce-row-tree.wgsl.jinja`](build/webgpu/reduce-row-tree.wgsl.jinja)
 - [`reduce-serial-axis.wgsl.jinja`](build/webgpu/reduce-serial-axis.wgsl.jinja)
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
+This example supplies explicit metadata for:
 - `y`
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.ReduceSum",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
@@ -143,7 +142,7 @@
       "name": "reducesum-rank3-spatial-axes12-f32-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
-        "source": "authored for branch coverage",
         "notes": "Rank-3 multi-axis f32 reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
@@ -155,7 +154,7 @@
       "name": "reducesum-rank3-spatial-axes12-f16-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
-        "source": "authored for branch coverage",
         "notes": "Rank-3 multi-axis f16 reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
@@ -167,7 +166,7 @@
       "name": "reducesum-rank4-spatial-axes23-f32-2x64x256x256-pathology",
       "preset": "stress",
       "provenance": {
-        "source": "authored for branch coverage",
         "notes": "Rank-4 multi-axis f32 reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [2, 3], "keepdims": 1 },
@@ -201,6 +200,44 @@
       "inputs": { "x": { "shape": [8, 65536, 64], "dtype": "float32", "dist": "normal", "seed": 118, "scale": 0.2 } },
       "outputs": { "y": { "shape": [8, 64], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.numel * 4" }] }
     }
   ]
 }

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
       "name": "reducesum-rank3-spatial-axes12-f32-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "repository-authored",
         "notes": "Rank-3 multi-axis f32 reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
       "name": "reducesum-rank3-spatial-axes12-f16-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "repository-authored",
         "notes": "Rank-3 multi-axis f16 reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
       "name": "reducesum-rank4-spatial-axes23-f32-2x64x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "repository-authored",
         "notes": "Rank-4 multi-axis f32 reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [2, 3], "keepdims": 1 },
       "inputs": { "x": { "shape": [8, 65536, 64], "dtype": "float32", "dist": "normal", "seed": 118, "scale": 0.2 } },
       "outputs": { "y": { "shape": [8, 64], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.numel * 4" }] }
+    },
+    {
+      "name": "reducesum-rank4-axes023-f32-8x256x56x56-noncontiguous-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Rank-4 reduction over axes {0,2,3} (per-channel statistics over batch and space with the channel axis second): the kept axis sits in the middle, so the reduced elements are not a contiguous suffix and the multi-axis rank-4 kernel launches one serial lane per output, 256 lanes each walking 25088 strided inputs. Its control reduces the same values with the channel axis first."
+      },
+      "vars": { "dtype": "float32", "batch": 8, "channels": 256, "height": 56, "width": 56 },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": { "x": { "shape": [8, 256, 56, 56], "dtype": "float32", "dist": "normal", "seed": 930, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [256], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.channels * args.height * args.width * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "reducesum-rank4-axes123-f32-256x8x56x56-contiguous-suffix-control",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Contiguous-suffix control for the axes {0,2,3} reduction: the same 256 reductions of 25088 elements with the kept axis first, which the cooperative suffix reduction covers with one workgroup per output."
+      },
+      "vars": { "dtype": "float32", "batch": 256, "channels": 8, "height": 56, "width": 56 },
+      "attrs": { "axes": [1, 2, 3], "keepdims": 0 },
+      "inputs": { "x": { "shape": [256, 8, 56, 56], "dtype": "float32", "dist": "normal", "seed": 931, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [256], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.channels * args.height * args.width * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,469 +2,246 @@
   "domain": "ai.onnx",
   "name": "ReduceSum",
   "sinceVersion": 13,
-  "description": "Computes the sum of elements along specified axes of the input tensor. The output rank matches the input rank when `keepdims` is 1; otherwise reduced dimensions are pruned. When no axes are provided, behavior is controlled by `noop_with_empty_axes`: reduce over all axes (default) or act as identity.",
-  "inputs": [{ "role": "data", "dtype": "T", "description": "The input tensor to reduce." }],
-  "outputs": [
-    {
-      "role": "reduced",
       "dtype": "T",
-      "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)",
-      "description": "The summed output tensor, with reduced dimensions either kept as size 1 or removed."
     }
-  ],
-  "attributes": { "keepdims": 1, "noop_with_empty_axes": 0, "axes": [] },
-  "attributeDescriptions": {
-    "keepdims": "If 1 (default), retains reduced dimensions with size 1; if 0, removes them from the output shape.",
-    "noop_with_empty_axes": "When axes is empty, if 0 (default) reduce over all axes; if 1, treat as a no-op identity and return the input unchanged.",
-    "axes": "Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`."
   },
   "attributeConstraints": { "keepdims": { "values": [0, 1] }, "noop_with_empty_axes": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int32"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "data", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "reduced", "role": "output" }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "treeWorkgroupOk": "reduceWorkgroupSize > 0 and pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize and reduceWorkgroupSize * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize",
     "subgroupWorkgroupFloor": "min(reduceWorkgroupSize, max(1, device.adapterInfo.subgroupMaxSize))",
-    "lastAxisRows": "rows(shapes.data, ranks.data - 1) if ranks.data > 0 else 1",
-    "lastAxisCols": "dim(shapes.data, ranks.data - 1) if ranks.data > 0 else 1",
     "rowSerialPreferred": "lastAxisRows >= tunables.ROW_SERIAL_MIN_ROWS and lastAxisCols <= tunables.ROW_SERIAL_MAX_COLS",
-    "axis0Rows": "dim(shapes.data, 0) if ranks.data >= 2 else 0",
-    "axis0Cols": "dim(shapes.data, 1) if ranks.data >= 2 else 0",
     "axis0SplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axis0Rows, tunables.AXIS0_SPLIT_TARGET_ROWS)))",
     "axis0SplitScratchBytes": "axis0SplitCount * axis0Cols * dtypeBytes(\"float32\")",
-    "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",
-    "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",
-    "axisSplitDim": "dim(shapes.data, reduceAxis) if ranks.data >= 2 and reduceAxis < ranks.data else 0",
-    "axisSplitInner": "inner(shapes.data, reduceAxis) if ranks.data >= 2 and reduceAxis < ranks.data else 1",
-    "axisSplitOutputs": "numel(shapes.reduced)",
     "axisSplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axisSplitDim, tunables.AXIS0_SPLIT_TARGET_ROWS)))",
     "axisSplitScratchBytes": "axisSplitCount * axisSplitOutputs * 4",
-    "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",
     "axis0TilePathFits": "treeWorkgroupOk and tunables.AXIS0_TILE_COLS > 0 and tunables.AXIS0_TILE_COLS <= reduceWorkgroupSize and reduceWorkgroupSize % tunables.AXIS0_TILE_COLS == 0",
-    "flatItems": "floor(numel(shapes.data) / tunables.VECTOR_WIDTH)",
     "flatSplitCount": "max(1, min(tunables.FULL_REDUCE_MAX_SPLITS, ceilDiv(flatItems, reduceWorkgroupSize)))",
     "flatScratchBytes": "flatSplitCount * dtypeBytes(\"float32\")",
-    "flatPathFits": "treeWorkgroupOk and flatSplitCount <= device.limits.maxComputeWorkgroupsPerDimension and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize",
-    "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",
-    "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)))"
-  },
-  "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "VECTOR_WIDTH": 4,
-    "ROW_PARALLEL_MIN_COLS": 64,
-    "SUBGROUP_MIN_COLS": 256,
-    "SUBGROUP_SMALL_ROW_LIMIT": 32768,
-    "AXIS0_SPLIT_MIN_ROWS": 8192,
-    "AXIS0_SPLIT_TARGET_ROWS": 256,
-    "AXIS0_MAX_SPLITS": 128,
-    "AXIS0_TILE_MIN_ROWS": 64,
-    "AXIS0_TILE_MIN_COLS": 16,
-    "AXIS0_TILE_COLS": 16,
-    "AXIS_SPLIT_TILE_COLS": 8,
-    "FULL_REDUCE_MIN_ELEMENTS": 8192,
-    "FULL_REDUCE_MAX_SPLITS": 256,
-    "CONTIGUOUS_SUFFIX_MIN_COLS": 256,
-    "ROW_SERIAL_MIN_ROWS": 8192,
-    "ROW_SERIAL_MAX_COLS": 1024
-  },
-  "tunableDescriptions": {
-    "SUBGROUP_SMALL_ROW_LIMIT": "Sets the row-count threshold that demotes subgroup last-axis reduction when each row is shorter than `SUBGROUP_MIN_COLS`.",
-    "AXIS0_SPLIT_MIN_ROWS": "Sets the minimum reduced-row count admitted to axis-zero split reduction.",
-    "AXIS0_SPLIT_TARGET_ROWS": "Sets the target reduced-row count per partial; the selector rounds the resulting split count up to a power of two.",
-    "AXIS0_MAX_SPLITS": "Caps the partial count for axis-zero and general axis-split reductions, bounding dispatch and scratch.",
-    "FULL_REDUCE_MAX_SPLITS": "Caps the partial count emitted by the flat all-axes reduction.",
-    "ROW_SERIAL_MIN_ROWS": "Sets the minimum row count for the many-short-rows region that prefers one serial reducer per row.",
-    "ROW_SERIAL_MAX_COLS": "Sets the maximum row width for the many-short-rows region that prefers one serial reducer per row."
   },
-  "bindingSets": {
-    "lastAxisVec4": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" },
-            { "name": "chunkCount", "type": "u32", "value": "dim(shapes.data, ranks.data - 1) / tunables.VECTOR_WIDTH" }
-          ]
-        }
-      }
-    ],
-    "lastAxisScalar": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" },
-            { "name": "cols", "type": "u32", "value": "dim(shapes.data, ranks.data - 1)" }
-          ]
-        }
-      }
-    ],
-    "lastAxisScalarSubgroup": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" },
-            { "name": "chunkCount", "type": "u32", "value": "dim(shapes.data, ranks.data - 1)" }
-          ]
-        }
-      }
-    ],
-    "scalar": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "1" },
-            { "name": "cols", "type": "u32", "value": "1" },
-            { "name": "outCount", "type": "u32", "value": "1" }
-          ]
-        }
-      }
-    ],
-    "rank1Axis0": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
-            { "name": "cols", "type": "u32", "value": "1" },
-            { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
-          ]
-        }
-      }
-    ],
-    "rank2Serial": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
-            { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" },
-            { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
-          ]
-        }
-      }
-    ],
-    "rank2SerialAxis1": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" },
-            { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
-          ]
-        }
-      }
-    ],
-    "axis0Parallel": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
-            { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }
-          ]
-        }
-      }
-    ],
-    "axisSplitReduce": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
-            { "name": "inner", "type": "u32", "value": "axisSplitInner" },
-            { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" }
-          ]
-        }
-      }
-    ],
-    "axisSplitCombine": [
-      {
-        "name": "partials",
-        "semantic": "partials",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$partialElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }] }
-      }
-    ],
-    "axis0SplitReduce": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" },
-            { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }
-          ]
-        }
-      }
-    ],
-    "axis0SplitCombine": [
-      {
-        "name": "partials",
-        "semantic": "partials",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$partialElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }] }
-      }
-    ],
-    "elementwise": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.reduced)" }] }
-      }
-    ],
-    "fullReduceSerial": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "numel(shapes.data)" },
-            { "name": "cols", "type": "u32", "value": "1" },
-            { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
-          ]
-        }
-      }
-    ],
-    "rankNAxis": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
-            { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }
-          ]
-        }
-      }
-    ],
-    "flatPartial": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count4", "type": "u32", "value": "floor(numel(shapes.data) / tunables.VECTOR_WIDTH)" },
-            { "name": "numel", "type": "u32", "value": "numel(shapes.data)" }
-          ]
-        }
-      }
-    ],
-    "flatCombine": [
-      {
-        "name": "partials",
-        "semantic": "partials",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$partialElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "1" }] }
-      }
-    ],
-    "suffixVec4": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "numel(shapes.reduced)" },
-            {
-              "name": "chunkCount",
-              "type": "u32",
-              "value": "numel(shapes.data) / numel(shapes.reduced) / tunables.VECTOR_WIDTH"
-            }
-          ]
-        }
-      }
-    ],
-    "suffixScalar": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "numel(shapes.reduced)" },
-            { "name": "cols", "type": "u32", "value": "numel(shapes.data) / numel(shapes.reduced)" }
-          ]
-        }
-      }
-    ],
-    "multiAxis": [
-      { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "contiguous_suffix_subgroup_vec4",
       "priority": 30,
       "requires": { "features": ["subgroups"] },
-      "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.data) / numel(shapes.reduced)) % tunables.VECTOR_WIDTH == 0"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(numel(shapes.data) / numel(shapes.reduced), tunables.VECTOR_WIDTH))))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.ContiguousSuffixSubgroupVec4",
-          "source": {
-            "shader": "reduce-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "vec4": true,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "suffixVec4",
-          "dispatch": { "workgroups": "numel(shapes.reduced)" }
         }
       ]
     },
     {
       "id": "contiguous_suffix_tree_vec4",
       "priority": 22,
-      "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.data) / numel(shapes.reduced)) % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(numel(shapes.data) / numel(shapes.reduced), tunables.VECTOR_WIDTH)))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.ContiguousSuffixTreeVec4",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "vec4": true,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "suffixVec4",
-          "dispatch": { "workgroups": "numel(shapes.reduced)" }
         }
       ]
     },
@@ -472,89 +249,174 @@
       "id": "contiguous_suffix_tree",
       "priority": 21,
       "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "treeWorkgroupOk"],
-      "constants": {
-        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(numel(shapes.data) / numel(shapes.reduced)))",
         "scalar": "dtypes.T"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.ContiguousSuffixTree",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
           },
-          "bindings": "suffixScalar",
-          "dispatch": { "workgroups": "numel(shapes.reduced)" }
         }
       ]
     },
     {
       "id": "multi_axis_rank3",
       "priority": 8,
-      "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)"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.MultiAxisRank3",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"multiaxis\"",
-              "rank": 3,
-              "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
-              "dataShape": "shapes.data",
-              "outputShape": "shapes.reduced",
-              "outputRank": "ranks.reduced",
-              "keepDims": "attrs.keepdims != 0",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "multiAxis",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
-      ],
-      "constants": { "scalar": "dtypes.T" }
     },
     {
       "id": "multi_axis_rank4",
       "priority": 8,
-      "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))"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.MultiAxisRank4",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"multiaxis\"",
-              "rank": 4,
-              "reduce": ["hasAxis(attrs.axes, 0, 4)", "hasAxis(attrs.axes, 1, 4)", "hasAxis(attrs.axes, 2, 4)", "hasAxis(attrs.axes, 3, 4)"],
-              "dataShape": "shapes.data",
-              "outputShape": "shapes.reduced",
-              "outputRank": "ranks.reduced",
-              "keepDims": "attrs.keepdims != 0",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "multiAxis",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
-      ],
-      "constants": { "scalar": "dtypes.T" }
     },
     {
       "id": "tree_last_axis_vec4",
       "priority": 23,
       "demoteWhen": ["rowSerialPreferred"],
-      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\" or dtypes.T == \"i32\")", "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"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH)))"
@@ -563,138 +425,140 @@
         {
           "id": "main",
           "name": "ReduceSum.TreeRowVec4",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "vec4": true,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "lastAxisVec4",
-          "dispatch": { "workgroups": "lastAxisRows" }
         }
       ]
     },
     {
       "id": "rank0_scalar",
       "priority": 40,
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "when": ["f16Ok(dtypes.T)", "ranks.data == 0", "ranks.reduced == 0"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank0Scalar",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"axis2d\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\"",
-              "logicalBool": "tensorDtypes.data == \"bool\""
-            }
           },
-          "bindings": "scalar",
           "dispatch": { "x": 1 }
         }
       ]
     },
     {
       "id": "rank1_axis0",
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "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))"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank1Axis0",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"axis2d\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\"",
-              "logicalBool": "tensorDtypes.data == \"bool\""
-            }
           },
-          "bindings": "rank1Axis0",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis1_parallel",
       "priority": 20,
       "demoteWhen": ["rowSerialPreferred"],
-      "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"],
-      "constants": {
         "scalar": "dtypes.T",
-        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.data, ranks.data - 1)))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Axis1Parallel",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "lastAxisScalar",
-          "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
         }
       ]
     },
     {
       "id": "axis_split",
       "priority": 24,
-      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "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"],
-      "derive": { "splitCount": "axisSplitCount" },
-      "constants": { "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * axisSplitOutputs]" }],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ReduceSum.AxisSplitReduce",
-          "source": {
-            "shader": "reduce-axis-split-reduce.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "axisSplitReduce",
-          "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AxisSplitCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "outputF16": "dtypes.T == \"f16\""
-            }
-          },
-          "bindings": "axisSplitCombine",
-          "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis_split_tiled_narrow",
       "priority": 25,
-      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "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"],
-      "derive": { "splitCount": "axisSplitCount" },
-      "constants": {
         "partialElement": "\"f32\"",
         "scalar": "dtypes.T",
         "workgroupSize": "reduceWorkgroupSize",
@@ -706,190 +570,189 @@
         {
           "id": "split_reduce",
           "name": "ReduceSum.AxisSplitTiledReduce",
-          "source": {
-            "shader": "reduce-axis0-tilecols.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "axisSplitReduce",
-          "dispatch": { "workgroups": "ceilDiv((axisSplitOutputs), (constants.tileCols))", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AxisSplitCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "outputF16": "dtypes.T == \"f16\""
-            }
-          },
-          "bindings": "axisSplitCombine",
-          "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis0_splitk",
       "priority": 22,
-      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "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"],
-      "derive": { "splitCount": "axis0SplitCount" },
-      "constants": { "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" },
-      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * dim(shapes.data, 1)]" }],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ReduceSum.Axis0SplitKReduce",
-          "source": {
-            "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "axis0SplitReduce",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.Axis0SplitKCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "outputF16": "dtypes.T == \"f16\""
-            }
-          },
-          "bindings": "axis0SplitCombine",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis0_splitk_i32",
-      "description": "Partitions a long rank-two axis-zero integer reduction across workgroups and combines exact int32 partials. It applies when the reduced row dimension is too large for one pass to expose enough parallelism.",
       "priority": 23,
-      "when": ["dtypes.T == \"i32\"", "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"],
-      "derive": { "splitCount": "axis0SplitCount" },
-      "constants": { "partialElement": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" },
-      "intermediates": [{ "id": "partials", "dtype": "int32", "shape": "[splitCount * dim(shapes.data, 1)]" }],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ReduceSum.Axis0SplitKReduceI32",
-          "source": {
-            "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "intMode": true,
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount"
-            }
-          },
-          "bindings": "axis0SplitReduce",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.Axis0SplitKCombineI32",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "intMode": true,
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount"
-            }
-          },
-          "bindings": "axis0SplitCombine",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis0_tilecols",
       "priority": 20,
-      "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"],
-      "constants": {
-        "scalar": "dtypes.T",
-        "workgroupSize": "reduceWorkgroupSize",
-        "tileCols": "tunables.AXIS0_TILE_COLS"
-      },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Axis0TileCols",
-          "source": {
-            "shader": "reduce-axis0-tilecols.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "axis0Parallel",
-          "dispatch": { "workgroups": "ceilDiv((dim(shapes.data, 1)), (constants.tileCols))" }
         }
       ]
     },
     {
       "id": "noop_empty_axes",
       "priority": 40,
-      "when": ["dtypes.T == \"f32\"", "attrs.noop_with_empty_axes == 1", "(attrs.axes | length) == 0", "sameShape(shapes.data, shapes.reduced)"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.NoopEmptyAxes",
           "shader": "datamove-elementwise-copy.wgsl.jinja",
-          "bindings": "elementwise",
-          "dispatch": { "gridStride": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "all_axes_flat",
       "priority": 31,
-      "constants": {
         "partialElement": "\"f32\"",
         "scalar": "dtypes.T",
         "workgroupSize": "reduceWorkgroupSize",
         "split": "flatSplitCount"
       },
-      "when": ["flatParallelCovered"],
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceSum.AllAxesFlatPartial",
-          "source": {
-            "shader": "reduce-flat-partial.wgsl.jinja",
-            "inputs": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "flatPartial",
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AllAxesFlatCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "flatCombine",
           "dispatch": { "x": 1 }
         }
       ]
@@ -897,25 +760,24 @@
     {
       "id": "all_axes_flat_i32",
       "priority": 32,
-      "constants": { "partialElement": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
-      "when": ["dtypes.T == \"i32\"", "numel(shapes.reduced) == 1", "numel(shapes.data) >= tunables.FULL_REDUCE_MIN_ELEMENTS", "flatPathFits"],
       "intermediates": [{ "id": "partials", "dtype": "int32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceSum.AllAxesFlatPartialI32",
-          "source": { "shader": "reduce-flat-partial.wgsl.jinja", "inputs": { "op": "\"sum\"", "intMode": true } },
-          "bindings": "flatPartial",
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AllAxesFlatCombineI32",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": { "op": "\"sum\"", "intMode": true, "workgroupSize": 1, "split": "flatSplitCount" }
-          },
-          "bindings": "flatCombine",
           "dispatch": { "x": 1 }
         }
       ]
@@ -923,42 +785,43 @@
     {
       "id": "rankn_single_axis_generic",
       "priority": 12,
       "supersededBy": ["axis_split_tiled_narrow", "axis_split", "subgroup_last_axis_vec4", "subgroup_last_axis", "tree_last_axis_vec4", "rank3_axis1", "rank3_axis2", "rank4_axis1"],
-      "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))"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.RankNSingleAxisGeneric",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"rankn\"",
-              "rank": "ranks.data",
-              "axis": "reduceAxis",
-              "dataShape": "shapes.data",
-              "outputShape": "shapes.reduced",
-              "outputRank": "ranks.reduced",
-              "keepDims": "attrs.keepdims != 0",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rankNAxis",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
-      ],
-      "constants": { "scalar": "dtypes.T" }
     },
     {
       "id": "subgroup_last_axis_vec4",
-      "description": "Reduces each contiguous last-axis row with subgroup collectives, using packed reads when the row width is vec4-aligned. It yields to the serial-row route when there are many short rows.",
       "priority": 25,
-      "requires": { "features": ["subgroups"] },
       "demoteWhen": ["rowSerialPreferred", "lastAxisCols < tunables.SUBGROUP_MIN_COLS and lastAxisRows >= tunables.SUBGROUP_SMALL_ROW_LIMIT"],
-      "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "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", "dim(shapes.data, ranks.data - 1) >= 4", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH == 0"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH))))"
@@ -967,29 +830,30 @@
         {
           "id": "main",
           "name": "ReduceSum.SubgroupRowVec4",
-          "source": {
-            "shader": "reduce-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "vec4": true,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "lastAxisVec4",
-          "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
         }
       ]
     },
     {
       "id": "subgroup_last_axis",
-      "description": "Reduces each contiguous last-axis row with subgroup collectives, using packed reads when the row width is vec4-aligned. It yields to the serial-row route when there are many short rows.",
       "priority": 24,
-      "requires": { "features": ["subgroups"] },
       "demoteWhen": ["rowSerialPreferred", "lastAxisCols < tunables.SUBGROUP_MIN_COLS and lastAxisRows >= tunables.SUBGROUP_SMALL_ROW_LIMIT"],
-      "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "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", "dim(shapes.data, ranks.data - 1) > 0", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH != 0"],
-      "constants": {
         "scalar": "dtypes.T",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))"
       },
@@ -997,210 +861,248 @@
         {
           "id": "main",
           "name": "ReduceSum.SubgroupRow",
-          "source": {
-            "shader": "reduce-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "vec4": false,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "lastAxisScalarSubgroup",
-          "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
         }
       ]
     },
     {
       "id": "rank3_axis1",
       "priority": 15,
       "supersededBy": ["axis_split"],
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 3", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.reduced == 2 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == dim(shapes.data, 2)) or (attrs.keepdims == 1 and ranks.reduced == 3 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == 1 and dim(shapes.reduced, 2) == dim(shapes.data, 2)))"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3Axis1",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"rankn\"",
-              "rank": 3,
-              "axis": 1,
-              "dataShape": "shapes.data",
-              "outputShape": "shapes.reduced",
-              "outputRank": "ranks.reduced",
-              "keepDims": "attrs.keepdims != 0",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rankNAxis",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
-      ],
-      "constants": { "scalar": "dtypes.T" }
     },
     {
       "id": "rank3_axis2",
       "priority": 15,
       "supersededBy": ["axis_split"],
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 3", "reduceAxis == 2", "((attrs.keepdims == 0 and ranks.reduced == 2 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 3 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == dim(shapes.data, 1) and dim(shapes.reduced, 2) == 1))"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3Axis2",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"rankn\"",
-              "rank": 3,
-              "axis": 2,
-              "dataShape": "shapes.data",
-              "outputShape": "shapes.reduced",
-              "outputRank": "ranks.reduced",
-              "keepDims": "attrs.keepdims != 0",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rankNAxis",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
-      ],
-      "constants": { "scalar": "dtypes.T" }
     },
     {
       "id": "rank4_axis1",
       "priority": 15,
       "supersededBy": ["axis_split"],
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 4", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.reduced == 3 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == dim(shapes.data, 2) and dim(shapes.reduced, 2) == dim(shapes.data, 3)) or (attrs.keepdims == 1 and ranks.reduced == 4 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == 1 and dim(shapes.reduced, 2) == dim(shapes.data, 2) and dim(shapes.reduced, 3) == dim(shapes.data, 3)))"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank4Axis1",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"rankn\"",
-              "rank": 4,
-              "axis": 1,
-              "dataShape": "shapes.data",
-              "outputShape": "shapes.reduced",
-              "outputRank": "ranks.reduced",
-              "keepDims": "attrs.keepdims != 0",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rankNAxis",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
-      ],
-      "constants": { "scalar": "dtypes.T" }
     },
     {
       "id": "axis0",
       "priority": 0,
       "supersededBy": ["axis_split_tiled_narrow", "axis0_splitk_i32", "axis0_splitk", "axis0_tilecols"],
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "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)))"],
       "passes": [
         {
           "id": "main",
           "name": "axis0",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"axis2d\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rank2Serial",
-          "constants": { "axis": 0 },
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis1",
       "priority": 0,
-      "constants": { "axis": 1, "scalar": "dtypes.T" },
-      "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))"],
       "passes": [
         {
           "id": "main",
           "name": "axis1",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"axis2d\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "rank2SerialAxis1",
-          "constants": { "axis": 1 },
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "all_axes_no_keepdims",
       "priority": 30,
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.keepdims == 0", "attrs.noop_with_empty_axes == 0", "ranks.reduced == 0"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3AllAxesNoKeepdims",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"axis2d\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "fullReduceSerial",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "all_axes_keepdims",
       "priority": 30,
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.keepdims == 1", "ranks.reduced == ranks.data", "numel(shapes.reduced) == 1"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3AllAxesKeepdims",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"sum\"",
-              "indexing": "\"axis2d\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "fullReduceSerial",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "ReduceSum",
   "sinceVersion": 13,
+  "inputs": { "x": { "onnx": "data", "dtype": "T" } },
+  "outputs": {
+    "y": {
+      "onnx": "reduced",
       "dtype": "T",
+      "rank": "ranks.x if attrs.keepdims == 1 or ((attrs.axes | length) == 0 and attrs.noop_with_empty_axes == 1) else (ranks.x - (attrs.axes | length) if (attrs.axes | length) > 0 else 0)"
     }
   },
+  "attributes": { "keepdims": { "default": 1 }, "noop_with_empty_axes": { "default": 0 }, "axes": { "default": [] } },
   "attributeConstraints": { "keepdims": { "values": [0, 1] }, "noop_with_empty_axes": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int32"] },
+  "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "VECTOR_WIDTH": { "default": 4 },
+    "ROW_PARALLEL_MIN_COLS": { "default": 64 },
+    "SUBGROUP_MIN_COLS": { "default": 256 },
+    "SUBGROUP_SMALL_ROW_LIMIT": { "default": 32768 },
+    "AXIS0_SPLIT_MIN_ROWS": { "default": 8192 },
+    "AXIS0_SPLIT_TARGET_ROWS": { "default": 256 },
+    "AXIS0_MAX_SPLITS": { "default": 128 },
+    "AXIS0_TILE_MIN_ROWS": { "default": 64 },
+    "AXIS0_TILE_MIN_COLS": { "default": 16 },
+    "AXIS0_TILE_COLS": { "default": 16 },
+    "AXIS_SPLIT_TILE_COLS": { "default": 8 },
+    "FULL_REDUCE_MIN_ELEMENTS": { "default": 8192 },
+    "FULL_REDUCE_MAX_SPLITS": { "default": 256 },
+    "CONTIGUOUS_SUFFIX_MIN_COLS": { "default": 256 },
+    "ROW_SERIAL_MIN_ROWS": { "default": 8192 },
+    "ROW_SERIAL_MAX_COLS": { "default": 1024 },
+    "MULTI_AXIS_COOP_MAX_OUTPUTS": { "default": 8192 },
+    "MULTI_AXIS_COOP_MIN_REDUCED": { "default": 256 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "treeWorkgroupOk": "reduceWorkgroupSize > 0 and pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize and reduceWorkgroupSize * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize",
     "subgroupWorkgroupFloor": "min(reduceWorkgroupSize, max(1, device.adapterInfo.subgroupMaxSize))",
+    "lastAxisRows": "rows(shapes.x, ranks.x - 1) if ranks.x > 0 else 1",
+    "lastAxisCols": "dim(shapes.x, ranks.x - 1) if ranks.x > 0 else 1",
     "rowSerialPreferred": "lastAxisRows >= tunables.ROW_SERIAL_MIN_ROWS and lastAxisCols <= tunables.ROW_SERIAL_MAX_COLS",
+    "axis0Rows": "dim(shapes.x, 0) if ranks.x >= 2 else 0",
+    "axis0Cols": "dim(shapes.x, 1) if ranks.x >= 2 else 0",
     "axis0SplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axis0Rows, tunables.AXIS0_SPLIT_TARGET_ROWS)))",
     "axis0SplitScratchBytes": "axis0SplitCount * axis0Cols * dtypeBytes(\"float32\")",
+    "axis0SplitPathFits": "axis0SplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(axis0Cols, reduceWorkgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and axis0SplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axis0SplitScratchBytes <= device.limits.maxBufferSize",
+    "reduceAxis": "(attrs.axes[0] + ranks.x if attrs.axes[0] < 0 else attrs.axes[0]) if ((attrs.axes | length) == 1 and isUniqueIntList(attrs.axes, 0 - ranks.x, ranks.x, 1)) else ranks.x",
+    "axisSplitDim": "dim(shapes.x, reduceAxis) if ranks.x >= 2 and reduceAxis < ranks.x else 0",
+    "axisSplitInner": "inner(shapes.x, reduceAxis) if ranks.x >= 2 and reduceAxis < ranks.x else 1",
+    "axisSplitOutputs": "numel(shapes.y)",
     "axisSplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axisSplitDim, tunables.AXIS0_SPLIT_TARGET_ROWS)))",
     "axisSplitScratchBytes": "axisSplitCount * axisSplitOutputs * 4",
+    "axisSplitPathFits": "axisSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(axisSplitOutputs, reduceWorkgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and axisSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axisSplitScratchBytes <= device.limits.maxBufferSize",
     "axis0TilePathFits": "treeWorkgroupOk and tunables.AXIS0_TILE_COLS > 0 and tunables.AXIS0_TILE_COLS <= reduceWorkgroupSize and reduceWorkgroupSize % tunables.AXIS0_TILE_COLS == 0",
+    "flatItems": "floor(numel(shapes.x) / tunables.VECTOR_WIDTH)",
     "flatSplitCount": "max(1, min(tunables.FULL_REDUCE_MAX_SPLITS, ceilDiv(flatItems, reduceWorkgroupSize)))",
     "flatScratchBytes": "flatSplitCount * dtypeBytes(\"float32\")",
+    "flatPathFits": "treeWorkgroupOk and flatSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize",
+    "flatParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) == 1 and numel(shapes.x) >= tunables.FULL_REDUCE_MIN_ELEMENTS and flatPathFits",
+    "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) > 0 and numel(shapes.x) % numel(shapes.y) == 0 and numel(shapes.x) / numel(shapes.y) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.x == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.y) == dim(shapes.x, 0)) or (ranks.x == 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.y) == dim(shapes.x, 0) * dim(shapes.x, 1)) or (ranks.x == 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.y) == dim(shapes.x, 0)))"
   },
+  "bindings": {
+    "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "y": { "buffer": "storage", "elementType": "$T" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
+        { "name": "chunkCount", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y) / tunables.VECTOR_WIDTH" }
+      ]
+    },
+    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$T" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
+        { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" }
+      ]
+    },
+    "params_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
+    },
+    "params_4": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
+        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1) / tunables.VECTOR_WIDTH" }
+      ]
+    },
+    "params_5": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "1" },
+        { "name": "cols", "type": "u32", "value": "1" },
+        { "name": "outCount", "type": "u32", "value": "1" }
+      ]
+    },
+    "params_6": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
+        { "name": "cols", "type": "u32", "value": "1" },
+        { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+      ]
+    },
+    "params_7": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
+      ]
+    },
+    "partials": { "buffer": "storage", "elementType": "$partialElement" },
+    "params_8": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
+        { "name": "inner", "type": "u32", "value": "axisSplitInner" },
+        { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" }
+      ]
+    },
+    "partials_2": { "name": "partials", "buffer": "read-only-storage", "elementType": "$partialElement" },
+    "params_9": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }]
+    },
+    "params_10": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
+      ]
+    },
+    "params_11": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }]
+    },
+    "params_13": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count4", "type": "u32", "value": "floor(numel(shapes.x) / tunables.VECTOR_WIDTH)" },
+        { "name": "numel", "type": "u32", "value": "numel(shapes.x)" }
+      ]
+    },
+    "params_14": { "name": "params", "buffer": "uniform", "struct": [{ "name": "cols", "type": "u32", "value": "1" }] },
+    "params_15": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
+        { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+      ]
+    },
+    "params_16": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
+        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
+      ]
+    },
+    "params_17": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+      ]
+    },
+    "params_18": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "contiguous_suffix_subgroup_vec4",
       "priority": 30,
+      "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.x) / numel(shapes.y)) % tunables.VECTOR_WIDTH == 0"],
       "requires": { "features": ["subgroups"] },
+      "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(numel(shapes.x) / numel(shapes.y), tunables.VECTOR_WIDTH))))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.ContiguousSuffixSubgroupVec4",
+          "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "vec4": true,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x", "y", "params"],
+          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "contiguous_suffix_tree_vec4",
       "priority": 22,
+      "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.x) / numel(shapes.y)) % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"],
+      "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(numel(shapes.x) / numel(shapes.y), tunables.VECTOR_WIDTH)))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.ContiguousSuffixTreeVec4",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "vec4": true,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x", "y", "params"],
+          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
     },
       "id": "contiguous_suffix_tree",
       "priority": 21,
       "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "treeWorkgroupOk"],
+      "derive": {
+        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(numel(shapes.x) / numel(shapes.y)))",
         "scalar": "dtypes.T"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.ContiguousSuffixTree",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": ["x_2", "y", "params_2"],
+          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "multi_axis_rank3_coop",
+      "priority": 9,
+      "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 3", "(attrs.keepdims == 1 and ranks.y == 3 and (dim(shapes.y, 0) == 1 if hasAxis(attrs.axes, 0, 3) else dim(shapes.y, 0) == dim(shapes.x, 0)) and (dim(shapes.y, 1) == 1 if hasAxis(attrs.axes, 1, 3) else dim(shapes.y, 1) == dim(shapes.x, 1)) and (dim(shapes.y, 2) == 1 if hasAxis(attrs.axes, 2, 3) else dim(shapes.y, 2) == dim(shapes.x, 2))) or (attrs.keepdims == 0 and ranks.y == 1)", "numel(shapes.y) >= 1", "numel(shapes.y) <= tunables.MULTI_AXIS_COOP_MAX_OUTPUTS", "numel(shapes.x) >= numel(shapes.y) * tunables.MULTI_AXIS_COOP_MIN_REDUCED", "pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize", "ceilDiv(numel(shapes.y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ReduceSum.MultiAxisRank3Coop",
+          "shader": "reduce-multi-axis-coop.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"multiaxis\"",
+            "rank": 3,
+            "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y"],
+          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "multi_axis_rank4_coop",
+      "priority": 9,
+      "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 4", "attrs.noop_with_empty_axes == 0", "numel(shapes.y) == (1 if hasAxis(attrs.axes, 0, 4) else dim(shapes.x, 0)) * (1 if hasAxis(attrs.axes, 1, 4) else dim(shapes.x, 1)) * (1 if hasAxis(attrs.axes, 2, 4) else dim(shapes.x, 2)) * (1 if hasAxis(attrs.axes, 3, 4) else dim(shapes.x, 3))", "((attrs.keepdims == 1 and ranks.y == 4) or (attrs.keepdims == 0 and ranks.y < 4))", "numel(shapes.y) >= 1", "numel(shapes.y) <= tunables.MULTI_AXIS_COOP_MAX_OUTPUTS", "numel(shapes.x) >= numel(shapes.y) * tunables.MULTI_AXIS_COOP_MIN_REDUCED", "pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize", "ceilDiv(numel(shapes.y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ReduceSum.MultiAxisRank4Coop",
+          "shader": "reduce-multi-axis-coop.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"multiaxis\"",
+            "rank": 4,
+            "reduce": ["hasAxis(attrs.axes, 0, 4)", "hasAxis(attrs.axes, 1, 4)", "hasAxis(attrs.axes, 2, 4)", "hasAxis(attrs.axes, 3, 4)"],
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
+          },
+          "bindings": ["x_2", "y"],
+          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "multi_axis_rank3",
       "priority": 8,
+      "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 3", "(attrs.keepdims == 1 and ranks.y == 3 and (dim(shapes.y, 0) == 1 if hasAxis(attrs.axes, 0, 3) else dim(shapes.y, 0) == dim(shapes.x, 0)) and (dim(shapes.y, 1) == 1 if hasAxis(attrs.axes, 1, 3) else dim(shapes.y, 1) == dim(shapes.x, 1)) and (dim(shapes.y, 2) == 1 if hasAxis(attrs.axes, 2, 3) else dim(shapes.y, 2) == dim(shapes.x, 2))) or (attrs.keepdims == 0 and ranks.y == 1)"],
+      "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.MultiAxisRank3",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"multiaxis\"",
+            "rank": 3,
+            "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_3"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
     },
     {
       "id": "multi_axis_rank4",
       "priority": 8,
+      "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 4", "attrs.noop_with_empty_axes == 0", "numel(shapes.y) == (1 if hasAxis(attrs.axes, 0, 4) else dim(shapes.x, 0)) * (1 if hasAxis(attrs.axes, 1, 4) else dim(shapes.x, 1)) * (1 if hasAxis(attrs.axes, 2, 4) else dim(shapes.x, 2)) * (1 if hasAxis(attrs.axes, 3, 4) else dim(shapes.x, 3))", "((attrs.keepdims == 1 and ranks.y == 4) or (attrs.keepdims == 0 and ranks.y < 4))"],
+      "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.MultiAxisRank4",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"multiaxis\"",
+            "rank": 4,
+            "reduce": ["hasAxis(attrs.axes, 0, 4)", "hasAxis(attrs.axes, 1, 4)", "hasAxis(attrs.axes, 2, 4)", "hasAxis(attrs.axes, 3, 4)"],
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_3"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
+    },
+    {
+      "id": "subgroup_rows_last_axis_vec4",
+      "priority": 26,
+      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\" or dtypes.T == \"i32\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "attrs.noop_with_empty_axes == 0", "lastAxisCols % tunables.VECTOR_WIDTH == 0", "lastAxisCols >= tunables.VECTOR_WIDTH", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize >= 16", "ceilDiv(lastAxisCols / tunables.VECTOR_WIDTH, device.adapterInfo.subgroupMinSize) <= 8", "lastAxisRows >= 64", "treeWorkgroupOk", "device.adapterInfo.subgroupMaxSize <= reduceWorkgroupSize", "(not rowSerialPreferred or lastAxisCols >= tunables.SUBGROUP_MIN_COLS)"],
+      "requires": { "features": ["subgroups"] },
+      "derive": {
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "workgroupSize": "reduceWorkgroupSize",
+        "vecsPerLane": "ceilDiv(lastAxisCols / tunables.VECTOR_WIDTH, device.adapterInfo.subgroupMinSize)"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "ReduceSum.SubgroupRowsVec4",
+          "shader": "reduce-row-subgroup-rows.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": ["x", "y", "params_4"],
+          "dispatch": {
+            "x": "min(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
+            "y": "ceilDiv(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
+            "z": 1
+          },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
     },
     {
       "id": "tree_last_axis_vec4",
       "priority": 23,
+      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\" or dtypes.T == \"i32\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "attrs.noop_with_empty_axes == 0", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "lastAxisCols % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"],
       "demoteWhen": ["rowSerialPreferred"],
+      "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH)))"
         {
           "id": "main",
           "name": "ReduceSum.TreeRowVec4",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "vec4": true,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x", "y", "params_4"],
+          "dispatch": { "x": "min(lastAxisRows, 65535)", "y": "ceilDiv(lastAxisRows, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "rank0_scalar",
       "priority": 40,
+      "when": ["f16Ok(dtypes.T)", "ranks.x == 0", "ranks.y == 0"],
+      "derive": { "axis": 0, "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank0Scalar",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_2", "y", "params_5"],
           "dispatch": { "x": 1 }
         }
       ]
     },
     {
       "id": "rank1_axis0",
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 1", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.y == 0) or (attrs.keepdims == 1 and ranks.y == 1 and dim(shapes.y, 0) == 1))"],
+      "derive": { "axis": 0, "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank1Axis0",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_2", "y", "params_6"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis1_parallel",
       "priority": 20,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 2", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "treeWorkgroupOk"],
       "demoteWhen": ["rowSerialPreferred"],
+      "derive": {
         "scalar": "dtypes.T",
+        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.x, ranks.x - 1)))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Axis1Parallel",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": ["x_2", "y", "params_7"],
+          "dispatch": {
+            "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
+            "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis_split",
       "priority": 24,
+      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "attrs.noop_with_empty_axes == 0", "ranks.x >= 2", "reduceAxis < ranks.x - 1", "not (ranks.x == 2 and reduceAxis == 0)", "axisSplitDim >= tunables.AXIS0_SPLIT_MIN_ROWS", "axisSplitOutputs >= 1", "axisSplitOutputs <= 4096", "axisSplitOutputs == rows(shapes.x, reduceAxis)", "axisSplitPathFits"],
+      "derive": {
+        "splitCount": "axisSplitCount",
+        "partialElement": "\"f32\"",
+        "workgroupSize": "reduceWorkgroupSize",
+        "split": "splitCount"
+      },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * axisSplitOutputs]" }],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ReduceSum.AxisSplitReduce",
+          "shader": "reduce-axis-split-reduce.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "splitSpec": "splitCount",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "partials", "params_8"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AxisSplitCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_9"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis_split_tiled_narrow",
       "priority": 25,
+      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "attrs.noop_with_empty_axes == 0", "ranks.x >= 2", "reduceAxis < ranks.x - 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.x, reduceAxis)", "axisSplitPathFits"],
+      "derive": {
+        "splitCount": "axisSplitCount",
         "partialElement": "\"f32\"",
         "scalar": "dtypes.T",
         "workgroupSize": "reduceWorkgroupSize",
         {
           "id": "split_reduce",
           "name": "ReduceSum.AxisSplitTiledReduce",
+          "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "splitSpec": "splitCount",
+            "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "partials", "params_8"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AxisSplitCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_9"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis0_splitk",
       "priority": 22,
+      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_SPLIT_MIN_ROWS", "dim(shapes.x, 1) > 0", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))", "axis0SplitPathFits"],
+      "derive": {
+        "splitCount": "axis0SplitCount",
+        "partialElement": "\"f32\"",
+        "workgroupSize": "reduceWorkgroupSize",
+        "split": "splitCount"
+      },
+      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitCount * dim(shapes.x, 1)]" }],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ReduceSum.Axis0SplitKReduce",
+          "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "splitSpec": "splitCount",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "partials", "params_10"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ReduceSum.Axis0SplitKCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_11"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis0_splitk_i32",
       "priority": 23,
+      "when": ["dtypes.T == \"i32\"", "ranks.x == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_SPLIT_MIN_ROWS", "dim(shapes.x, 1) > 0", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))", "axis0SplitPathFits"],
+      "derive": {
+        "splitCount": "axis0SplitCount",
+        "partialElement": "dtypes.T",
+        "workgroupSize": "reduceWorkgroupSize",
+        "split": "splitCount"
+      },
+      "intermediates": [{ "id": "partials", "dtype": "int32", "shape": "[splitCount * dim(shapes.x, 1)]" }],
       "passes": [
         {
           "id": "split_reduce",
           "name": "ReduceSum.Axis0SplitKReduceI32",
+          "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "intMode": true, "splitSpec": "splitCount" },
+          "bindings": ["x_2", "partials", "params_10"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ReduceSum.Axis0SplitKCombineI32",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "intMode": true, "splitSpec": "splitCount" },
+          "bindings": ["partials_2", "y", "params_11"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis0_tilecols",
       "priority": 20,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_TILE_MIN_ROWS", "axis0Cols >= tunables.AXIS0_TILE_MIN_COLS", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))", "axis0TilePathFits"],
+      "derive": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "tileCols": "tunables.AXIS0_TILE_COLS" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Axis0TileCols",
+          "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_10"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noop_empty_axes",
       "priority": 40,
+      "when": ["dtypes.T == \"f32\"", "attrs.noop_with_empty_axes == 1", "(attrs.axes | length) == 0", "sameShape(shapes.x, shapes.y)"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.NoopEmptyAxes",
           "shader": "datamove-elementwise-copy.wgsl.jinja",
+          "bindings": [
+            "x_2",
+            "y",
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "all_axes_flat",
       "priority": 31,
+      "when": ["flatParallelCovered"],
+      "derive": {
         "partialElement": "\"f32\"",
         "scalar": "dtypes.T",
         "workgroupSize": "reduceWorkgroupSize",
         "split": "flatSplitCount"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceSum.AllAxesFlatPartial",
+          "shader": "reduce-flat-partial.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": ["x_2", "partials", "params_13"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AllAxesFlatCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_14"],
           "dispatch": { "x": 1 }
         }
       ]
     {
       "id": "all_axes_flat_i32",
       "priority": 32,
+      "when": ["dtypes.T == \"i32\"", "numel(shapes.y) == 1", "numel(shapes.x) >= tunables.FULL_REDUCE_MIN_ELEMENTS", "flatPathFits"],
+      "derive": { "partialElement": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
       "intermediates": [{ "id": "partials", "dtype": "int32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceSum.AllAxesFlatPartialI32",
+          "shader": "reduce-flat-partial.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "intMode": true },
+          "bindings": ["x_2", "partials", "params_13"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceSum.AllAxesFlatCombineI32",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "intMode": true, "splitSpec": "flatSplitCount" },
+          "bindings": ["partials_2", "y", "params_14"],
           "dispatch": { "x": 1 }
         }
       ]
     {
       "id": "rankn_single_axis_generic",
       "priority": 12,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.noop_with_empty_axes == 0", "reduceAxis < ranks.x", "numel(shapes.y) == rows(shapes.x, reduceAxis)", "((attrs.keepdims == 0 and ranks.y == ranks.x - 1) or (attrs.keepdims == 1 and ranks.y == ranks.x and dim(shapes.y, reduceAxis) == 1))"],
       "supersededBy": ["axis_split_tiled_narrow", "axis_split", "subgroup_last_axis_vec4", "subgroup_last_axis", "tree_last_axis_vec4", "rank3_axis1", "rank3_axis2", "rank4_axis1"],
+      "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.RankNSingleAxisGeneric",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"rankn\"",
+            "rank": "ranks.x",
+            "axisSpec": "reduceAxis",
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_15"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
     },
     {
       "id": "subgroup_last_axis_vec4",
       "priority": 25,
+      "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "attrs.noop_with_empty_axes == 0", "dim(shapes.x, ranks.x - 1) >= 4", "dim(shapes.x, ranks.x - 1) % tunables.VECTOR_WIDTH == 0"],
       "demoteWhen": ["rowSerialPreferred", "lastAxisCols < tunables.SUBGROUP_MIN_COLS and lastAxisRows >= tunables.SUBGROUP_SMALL_ROW_LIMIT"],
+      "requires": { "features": ["subgroups"] },
+      "derive": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH))))"
         {
           "id": "main",
           "name": "ReduceSum.SubgroupRowVec4",
+          "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "vec4": true,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
+          },
+          "bindings": ["x", "y", "params_4"],
+          "dispatch": {
+            "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
+            "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
+            "z": 1
           },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "subgroup_last_axis",
       "priority": 24,
+      "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "attrs.noop_with_empty_axes == 0", "dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) % tunables.VECTOR_WIDTH != 0"],
       "demoteWhen": ["rowSerialPreferred", "lastAxisCols < tunables.SUBGROUP_MIN_COLS and lastAxisRows >= tunables.SUBGROUP_SMALL_ROW_LIMIT"],
+      "requires": { "features": ["subgroups"] },
+      "derive": {
         "scalar": "dtypes.T",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))"
       },
         {
           "id": "main",
           "name": "ReduceSum.SubgroupRow",
+          "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "vec4": false,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_16"],
+          "dispatch": {
+            "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
+            "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
+            "z": 1
+          },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "rank3_axis1",
       "priority": 15,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 3", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.y == 2 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 2)) or (attrs.keepdims == 1 and ranks.y == 3 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == 1 and dim(shapes.y, 2) == dim(shapes.x, 2)))"],
       "supersededBy": ["axis_split"],
+      "derive": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3Axis1",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"rankn\"",
+            "rank": 3,
+            "axisSpec": 1,
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_15"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
     },
     {
       "id": "rank3_axis2",
       "priority": 15,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 3", "reduceAxis == 2", "((attrs.keepdims == 0 and ranks.y == 2 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 3 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1) and dim(shapes.y, 2) == 1))"],
       "supersededBy": ["axis_split"],
+      "derive": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3Axis2",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"rankn\"",
+            "rank": 3,
+            "axisSpec": 2,
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_15"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
     },
     {
       "id": "rank4_axis1",
       "priority": 15,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 4", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.y == 3 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 2) and dim(shapes.y, 2) == dim(shapes.x, 3)) or (attrs.keepdims == 1 and ranks.y == 4 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == 1 and dim(shapes.y, 2) == dim(shapes.x, 2) and dim(shapes.y, 3) == dim(shapes.x, 3)))"],
       "supersededBy": ["axis_split"],
+      "derive": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank4Axis1",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"rankn\"",
+            "rank": 4,
+            "axisSpec": 1,
+            "dataShape": "shapes.x",
+            "outputShape": "shapes.y",
+            "outputRank": "ranks.y",
+            "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_15"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
     },
     {
       "id": "axis0",
       "priority": 0,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))"],
       "supersededBy": ["axis_split_tiled_narrow", "axis0_splitk_i32", "axis0_splitk", "axis0_tilecols"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "axis0",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "axis": 0,
+            "op": "\"sum\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_17"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis1",
       "priority": 0,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 0)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == 1))"],
+      "derive": { "axis": 1, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "axis1",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "axis": 1,
+            "op": "\"sum\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_18"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "all_axes_no_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 0", "attrs.noop_with_empty_axes == 0", "ranks.y == 0"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3AllAxesNoKeepdims",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": [
+            "x_2",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
+                { "name": "cols", "type": "u32", "value": "1" },
+                { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "all_axes_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 1", "ranks.y == ranks.x", "numel(shapes.y) == 1"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceSum.Rank3AllAxesKeepdims",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"sum\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": [
+            "x_2",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
+                { "name": "cols", "type": "u32", "value": "1" },
+                { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,26 +1,62 @@
 {
   "name": "ai.onnx.ReduceSum",
-  "id": "_ai_onnx_reducesum_webgpu_af45eb9",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "l6QAUKLa8vIxVVA0rZdVOMUrh/1wu9YyXFkaAvC3Pf0=",
       "datamove-elementwise-copy.wgsl.jinja": "Q1WuZCqcDf9b6cbT4rLilpeWj6327t/bSFmKr6IVBoQ=",
-      "manifest.json": "a28ZCFG58JYKI9Fz/xskSTbtmTzgBjjg1hnj11Qkhhw=",
-      "reduce-axis-split-reduce.wgsl.jinja": "Q50AhcbpYvCADMPrbGtPXKbmfT8f6LmQY+6+/CEX224=",
-      "reduce-axis0-splitk-combine.wgsl.jinja": "uBBEUUxLWHWFF9sdlIgcdZ3y5SlI46mwaa0qtq4oP6E=",
-      "reduce-axis0-splitk-reduce.wgsl.jinja": "Dsb/fe0TUftbM5lc45iDV/3aku1avSvAU1cYjHXAYrY=",
-      "reduce-axis0-tilecols.wgsl.jinja": "FaAQTXiES4TyKHlXhGBz2uc00pBlObzW5v8CxtnPTQo=",
-      "reduce-flat-partial.wgsl.jinja": "8Isi9lEY0j5IW8JpC0bjE4q+FxgdoJYMb+mSIg4mpTo=",
-      "reduce-row-subgroup.wgsl.jinja": "lRRp3W2sqq1oJk92iUiXcHzGsA/jZK3uxGfw+WUt+GE=",
-      "reduce-row-tree.wgsl.jinja": "eBfbOUoLED/RUxjWBKVEUfFXDlo1Nl3l30oqu0ljujc=",
-      "reduce-serial-axis.wgsl.jinja": "k7HGCftlMjDpZju1Yb57Ook5NJICOCtK2kCSsFoWXkc=",
-      "test.json": "/y+5OIa1acHzNtp0IHM424n+1/ptM4LM/9ZpFaEhZe0="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.ReduceSum" }
 }

 {
   "name": "ai.onnx.ReduceSum",
+  "id": "_ai_onnx_reducesum_webgpu_7169ef2",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "BI0nPERvR8p4G7FuvF0DHPAoslAI4pxFyIhmFGRy49A=",
       "datamove-elementwise-copy.wgsl.jinja": "Q1WuZCqcDf9b6cbT4rLilpeWj6327t/bSFmKr6IVBoQ=",
+      "manifest.json": "5i7cYgcHv/CRPtZy+xfWNEkFink8TlfQwSb8cKmvyaw=",
+      "reduce-axis-split-reduce.wgsl.jinja": "4+ep9xH4pHZOfaZ8abJhA8SW5M4mUhDC+y1CtDz6vjY=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "Yz1hjK55R/kndUrw3ugPqgPaOBwKmYupqZoVdJTHO5Q=",
+      "reduce-axis0-splitk-reduce.wgsl.jinja": "jB2h58emn6rfKhd8ALzqylSDsmBbrrkEOhtIoNhcQM0=",
+      "reduce-axis0-tilecols.wgsl.jinja": "PjYkEUQJBeG70td3W2xxmexH9x6XNIfeSzXBY47HbaU=",
+      "reduce-flat-partial.wgsl.jinja": "+qToL+wFi9QOxvY887aBAEwZK6Xu/eLQkucKvJYNwSk=",
+      "reduce-multi-axis-coop.wgsl.jinja": "ypqU1bkYaYU8YlNy9k4htstN7eLzpsRT2iJ7ErlExYQ=",
+      "reduce-row-subgroup-rows.wgsl.jinja": "76u7rAvFoZZKrFDs2A2jk0vkB0uNrPL6twdOBE9b+v8=",
+      "reduce-row-subgroup.wgsl.jinja": "2mu9LEsk8HfaLvucBCfcB1/ENpXkD6ELiCtRt+5UqiU=",
+      "reduce-row-tree.wgsl.jinja": "Bwa5xcI0bTmKXb4r9Cc1bfVbM5rNqqpQVrWWVqcb8xA=",
+      "reduce-serial-axis.wgsl.jinja": "fvUV9htqzKzt4Pg05pYtRmup/5QIHYaUGQHJZnthXKo=",
+      "test.json": "F7Iaw2W2jjR+gGYmEDa+6DCBeNv5zF6FSfLEmWfYx2A="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "contiguous_suffix_subgroup_vec4": ["reduce-row-subgroup.wgsl.jinja"],
+      "contiguous_suffix_tree_vec4": ["reduce-row-tree.wgsl.jinja"],
+      "contiguous_suffix_tree": ["reduce-row-tree.wgsl.jinja"],
+      "multi_axis_rank3_coop": ["reduce-multi-axis-coop.wgsl.jinja"],
+      "multi_axis_rank4_coop": ["reduce-multi-axis-coop.wgsl.jinja"],
+      "multi_axis_rank3": ["reduce-serial-axis.wgsl.jinja"],
+      "multi_axis_rank4": ["reduce-serial-axis.wgsl.jinja"],
+      "subgroup_rows_last_axis_vec4": ["reduce-row-subgroup-rows.wgsl.jinja"],
+      "tree_last_axis_vec4": ["reduce-row-tree.wgsl.jinja"],
+      "rank0_scalar": ["reduce-serial-axis.wgsl.jinja"],
+      "rank1_axis0": ["reduce-serial-axis.wgsl.jinja"],
+      "axis1_parallel": ["reduce-row-tree.wgsl.jinja"],
+      "axis_split": ["reduce-axis-split-reduce.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "axis_split_tiled_narrow": ["reduce-axis0-splitk-combine.wgsl.jinja", "reduce-axis0-tilecols.wgsl.jinja"],
+      "axis0_splitk": ["reduce-axis0-splitk-combine.wgsl.jinja", "reduce-axis0-splitk-reduce.wgsl.jinja"],
+      "axis0_splitk_i32": ["reduce-axis0-splitk-combine.wgsl.jinja", "reduce-axis0-splitk-reduce.wgsl.jinja"],
+      "axis0_tilecols": ["reduce-axis0-tilecols.wgsl.jinja"],
+      "noop_empty_axes": ["datamove-elementwise-copy.wgsl.jinja"],
+      "all_axes_flat": ["reduce-axis0-splitk-combine.wgsl.jinja", "reduce-flat-partial.wgsl.jinja"],
+      "all_axes_flat_i32": ["reduce-axis0-splitk-combine.wgsl.jinja", "reduce-flat-partial.wgsl.jinja"],
+      "rankn_single_axis_generic": ["reduce-serial-axis.wgsl.jinja"],
+      "subgroup_last_axis_vec4": ["reduce-row-subgroup.wgsl.jinja"],
+      "subgroup_last_axis": ["reduce-row-subgroup.wgsl.jinja"],
+      "rank3_axis1": ["reduce-serial-axis.wgsl.jinja"],
+      "rank3_axis2": ["reduce-serial-axis.wgsl.jinja"],
+      "rank4_axis1": ["reduce-serial-axis.wgsl.jinja"],
+      "axis0": ["reduce-serial-axis.wgsl.jinja"],
+      "axis1": ["reduce-serial-axis.wgsl.jinja"],
+      "all_axes_no_keepdims": ["reduce-serial-axis.wgsl.jinja"],
+      "all_axes_keepdims": ["reduce-serial-axis.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/reduce-axis-split-reduce.wgsl.jinja CHANGED Viewed

@@ -4,30 +4,53 @@
 // inner-axis elements. The combine pass folds the segments and finalizes the
 // selected reduction.
 //
-// logsumexp writes three partial planes per output and segment: the segment
-// maximum, the sum of exp(x - maximum), and a packed NaN marker.
-{% set castF32 = source.castF32 is defined and source.castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
-{% if source.usesF16 is defined and source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded output index: wg.z carries the high bits past the
   // per-dimension dispatch limit on the x dimension.
-  let output_index = (wg.x + wg.z * nwg.x) * WG + (gid.x % WG);
   let seg = wg.y;
   if (output_index >= params.outputs) { return; }
@@ -41,9 +64,52 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   var a1 = a0 + chunk;
   if (a1 > params.axisDim) { a1 = params.axisDim; }
   var acc = 0.0;
   for (var axis_index = a0; axis_index < a1; axis_index = axis_index + 1u) {
     acc = acc + {{ xa }}x[input_base + axis_index * params.inner]{{ ax }};
   }
   partials[seg * params.outputs + output_index] = acc;
 }

 // inner-axis elements. The combine pass folds the segments and finalizes the
 // selected reduction.
 //
+{% if op == "logsumexp" %}
+// Each output segment writes three partial planes: its maximum, the sum of
+// exp(x - maximum), and a packed NaN marker.
+{% endif %}
+{% set castF32 = castF32 is defined and castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% elif op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(workgroup_id) wg: vec3<u32>) {
   // 2D-folded output index: wg.z carries the high bits past the
   // per-dimension dispatch limit on the x dimension.
+  let output_index = (wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u) * WG + (gid.x % WG);
   let seg = wg.y;
   if (output_index >= params.outputs) { return; }
   var a1 = a0 + chunk;
   if (a1 > params.axisDim) { a1 = params.axisDim; }
+{% if op == "logsumexp" %}
+  var local_max = F32_MIN;
+  var local_nan_count = 0.0;
+  var local_nan_value = 0.0;
+  for (var axis_index = a0; axis_index < a1; axis_index = axis_index + 1u) {
+    let value = {{ xa }}x[input_base + axis_index * params.inner]{{ ax }};
+    if (is_nan_f32(value)) {
+      local_nan_count = local_nan_count + 1.0;
+      local_nan_value = value;
+    } else {
+      local_max = max(local_max, value);
+    }
+  }
+  var acc = 0.0;
+  if (local_nan_count == 0.0) {
+    for (var axis_index = a0; axis_index < a1; axis_index = axis_index + 1u) {
+      acc = acc + exp({{ xa }}x[input_base + axis_index * params.inner]{{ ax }} - local_max);
+    }
+  }
+  partials[seg * params.outputs + output_index] = local_max;
+  partials[(SPLIT + seg) * params.outputs + output_index] = acc;
+  partials[(2u * SPLIT + seg) * params.outputs + output_index] = select(0.0, local_nan_value, local_nan_count > 0.0);
+{% else %}
+{% if op == "max" %}
+  var acc = reduction_identity();
+{% elif op == "min" %}
+  var acc = reduction_identity();
+{% elif op == "prod" %}
+  var acc = 1.0;
+{% else %}
   var acc = 0.0;
+{% endif %}
   for (var axis_index = a0; axis_index < a1; axis_index = axis_index + 1u) {
+{% if op == "max" or op == "min" %}
+    acc = {{ op }}(acc, {{ xa }}x[input_base + axis_index * params.inner]{{ ax }});
+{% elif op == "prod" %}
+    acc = acc * {{ xa }}x[input_base + axis_index * params.inner]{{ ax }};
+{% elif op == "l1" %}
+    acc = acc + abs({{ xa }}x[input_base + axis_index * params.inner]{{ ax }});
+{% elif op == "l2" or op == "sumsquare" %}
+    let value = {{ xa }}x[input_base + axis_index * params.inner]{{ ax }};
+    acc = acc + value * value;
+{% else %}
     acc = acc + {{ xa }}x[input_base + axis_index * params.inner]{{ ax }};
+{% endif %}
   }
   partials[seg * params.outputs + output_index] = acc;
+{% endif %}
 }

build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja CHANGED Viewed

@@ -2,37 +2,121 @@
 // folds the segment partials and applies the selected reduction's final step.
 // Segments are folded in ascending order for deterministic results. This order
 // differs from the single-pass reduction but remains within the f32 tolerance.
-{% if source.outputF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let stride = nwg.x * WG;
-  let start = (gid.y * nwg.x * WG) + gid.x;
   for (var col = start; col < params.cols; col = col + stride) {
-{% if source.intMode %}
     var total = 0i;
 {% else %}
     var total = 0.0;
 {% endif %}
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];
       total = total + p;
     }
-{% if source.outputF16 %}
     y[col] = f16(total);
 {% else %}
     y[col] = total;
 {% endif %}
   }
 }

 // folds the segment partials and applies the selected reduction's final step.
 // Segments are folded in ascending order for deterministic results. This order
 // differs from the single-pass reduction but remains within the f32 tolerance.
+{% set addBias = addBias is defined and addBias %}
+{% set biasCols = biasCols | default(0) %}
+{% set intMode = intMode is defined and intMode %}
+{% set yv = "f16(" if outputF16 else "" %}
+{% set vy = ")" if outputF16 else "" %}
+{% if outputF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
+{% if addBias %}
+const BIAS_COLS: u32 = {{ biasCols }}u;
+{% endif %}
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% elif op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let stride = nwg.x * WG;
+  let start = (gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG) + gid.x;
   for (var col = start; col < params.cols; col = col + stride) {
+{% if op == "logsumexp" %}
+    // Merge SPLIT (segMax, segSumExp) pairs stably; carry NaN / +Inf markers.
+    var nan_value = 0.0;
+    var has_nan = false;
+    var global_max = F32_MIN;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let nv = partials[(2u * SPLIT + seg) * params.cols + col];
+      if (nv != 0.0 || is_nan_f32(nv)) {
+        has_nan = true;
+        nan_value = nv;
+      }
+      global_max = max(global_max, partials[seg * params.cols + col]);
+    }
+    var sum = 0.0;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let seg_max = partials[seg * params.cols + col];
+      let seg_sum = partials[(SPLIT + seg) * params.cols + col];
+      sum = sum + seg_sum * exp(seg_max - global_max);
+    }
+    let has_positive_inf = global_max > F32_MAX;
+    let finite_or_inf = select(global_max + log(sum), global_max, has_positive_inf);
+    y[col] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }};
+{% else %}
+{% if intMode %}
+{% if op == "prod" %}
+    var total = 1i;
+{% else %}
     var total = 0i;
+{% endif %}
+{% else %}
+{% if op == "max" %}
+    var total = reduction_identity();
+{% elif op == "min" %}
+    var total = reduction_identity();
+{% elif op == "prod" %}
+    var total = 1.0;
 {% else %}
     var total = 0.0;
+{% endif %}
 {% endif %}
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];
+{% if op == "max" or op == "min" %}
+      total = {{ op }}(total, p);
+{% elif op == "prod" %}
+      total = total * p;
+{% else %}
       total = total + p;
+{% endif %}
     }
+{% if addBias %}
+    total = total + f32(bias[col % BIAS_COLS]);
+{% endif %}
+{% if op == "l2" %}
+    y[col] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
+    y[col] = {{ yv }}log(total){{ vy }};
+{% elif op == "mean" %}
+    y[col] = {{ yv }}total / f32(params.rows){{ vy }};
+{% else %}
+{% if outputF16 %}
     y[col] = f16(total);
 {% else %}
     y[col] = total;
+{% endif %}
+{% endif %}
 {% endif %}
   }
 }

build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja CHANGED Viewed

@@ -2,31 +2,49 @@
 // segments increases residency for tall matrices. Each (column, segment)
 // invocation reduces one row slice and writes partials[segment * columns +
 // column]. Adjacent column threads keep row reads coalesced.
-//
-// logsumexp writes the segment maximum and sum of exp(x - maximum) as separate
-// partial planes; the combine pass merges them stably and handles NaN and +Inf.
-{% set castF32 = source.castF32 is defined and source.castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
-{% if source.usesF16 is defined and source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded column index: wg.z carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit on the x dimension.
-  let col = (wg.x + wg.z * nwg.x) * WG + (gid.x % WG);
   let seg = wg.y;
   if (col >= params.cols) { return; }
@@ -36,13 +54,62 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   var r1 = r0 + chunk;
   if (r1 > params.rows) { r1 = params.rows; }
-{% if source.intMode %}
   var acc = 0i;
 {% else %}
   var acc = 0.0;
 {% endif %}
   for (var row = r0; row < r1; row = row + 1u) {
     acc = acc + {{ xa }}x[row * params.cols + col]{{ ax }};
   }
   partials[seg * params.cols + col] = acc;
 }

 // segments increases residency for tall matrices. Each (column, segment)
 // invocation reduces one row slice and writes partials[segment * columns +
 // column]. Adjacent column threads keep row reads coalesced.
+{% set castF32 = castF32 is defined and castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% elif op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(workgroup_id) wg: vec3<u32>) {
+  // 2D-folded column index: wg.z carries the high bits past the per-axis dispatch
+  // fold width on the x dimension.
+  let col = (wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u) * WG + (gid.x % WG);
   let seg = wg.y;
   if (col >= params.cols) { return; }
   var r1 = r0 + chunk;
   if (r1 > params.rows) { r1 = params.rows; }
+{% if op == "logsumexp" %}
+  var local_max = F32_MIN;
+  var local_nan_count = 0.0;
+  var local_nan_value = 0.0;
+  for (var row = r0; row < r1; row = row + 1u) {
+    let value = {{ xa }}x[row * params.cols + col]{{ ax }};
+    if (is_nan_f32(value)) {
+      local_nan_count = local_nan_count + 1.0;
+      local_nan_value = value;
+    } else {
+      local_max = max(local_max, value);
+    }
+  }
+  var acc = 0.0;
+  if (local_nan_count == 0.0) {
+    for (var row = r0; row < r1; row = row + 1u) {
+      acc = acc + exp({{ xa }}x[row * params.cols + col]{{ ax }} - local_max);
+    }
+  }
+  // Three partial planes: seg max, seg sumexp(x - max), and a packed nan
+  // marker (count in the low slot, the nan bit-pattern smuggled as f32).
+  partials[seg * params.cols + col] = local_max;
+  partials[(SPLIT + seg) * params.cols + col] = acc;
+  partials[(2u * SPLIT + seg) * params.cols + col] = select(0.0, local_nan_value, local_nan_count > 0.0);
+{% else %}
+{% if intMode | default(false) %}
+{% if op == "prod" %}
+  var acc = 1i;
+{% else %}
   var acc = 0i;
+{% endif %}
+{% else %}
+{% if op == "max" %}
+  var acc = reduction_identity();
+{% elif op == "min" %}
+  var acc = reduction_identity();
+{% elif op == "prod" %}
+  var acc = 1.0;
 {% else %}
   var acc = 0.0;
+{% endif %}
 {% endif %}
   for (var row = r0; row < r1; row = row + 1u) {
+{% if op == "max" or op == "min" %}
+    acc = {{ op }}(acc, {{ xa }}x[row * params.cols + col]{{ ax }});
+{% elif op == "prod" %}
+    acc = acc * {{ xa }}x[row * params.cols + col]{{ ax }};
+{% elif op == "l1" %}
+    acc = acc + abs({{ xa }}x[row * params.cols + col]{{ ax }});
+{% elif op == "l2" or op == "sumsquare" %}
+    let value = {{ xa }}x[row * params.cols + col]{{ ax }};
+    acc = acc + value * value;
+{% else %}
     acc = acc + {{ xa }}x[row * params.cols + col]{{ ax }};
+{% endif %}
   }
   partials[seg * params.cols + col] = acc;
+{% endif %}
 }

build/webgpu/reduce-axis0-tilecols.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,9 @@
 // Tiled column-wise reduction. Each workgroup owns TILE_COLS columns;
 // ROW_LANES threads stride the rows of one column, then lane 0 folds their
 // partials. In split mode, workgroup y selects an axis segment and finalization
-// is deferred to the combine pass. Cooperative row lanes preserve occupancy
-// when the flattened output has only a few elements.
-{% set splitMode = source.split is defined %}
 {% if splitMode %}
 {% set rowBegin = "row_begin + row_lane" %}
 {% set rowEnd = "row_end" %}
@@ -13,34 +13,66 @@
 {% set rowEnd = "params.rows" %}
 {% set elem = "x[inputBase + row * params.cols + col]" %}
 {% endif %}
-{% set castF32 = source.castF32 is defined and source.castF32 %}
-{% set intMode = source.intMode is defined and source.intMode %}
 {% set scalar = "f32" if castF32 else scalar %}
 {% if castF32 %}
 {% set elem = "f32(" ~ elem ~ ")" %}
 {% endif %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-{% if source.usesF16 is defined and source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 const TILE_COLS: u32 = {{ tileCols }}u;
 const ROW_LANES: u32 = WG / TILE_COLS;
 {% if splitMode %}
-const SPLIT: u32 = {{ source.split }}u;
 {% endif %}
-var<workgroup> partial: array<{{ scalar if (source.op == "max" or source.op == "min" or intMode) else "f32" }}, WG>;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if not splitMode %},
-        @builtin(num_workgroups) nwg: vec3<u32>{% endif %}) {
   let tid = lid.x;
   let col_lane = tid % TILE_COLS;
   let row_lane = tid / TILE_COLS;
@@ -61,17 +93,134 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   // 2D-folded tile index: wg.y carries the high bits past the dispatch limit.
   // The batched form reuses this same coalesced axis-0 reduction for a middle
   // axis by assigning consecutive tiles to each outer slice.
-  let tile = wg.x + wg.y * nwg.x;
   let col = tile * TILE_COLS + col_lane;
   let inputBase = 0u;
   let outputIndex = col;
   let in_bounds = col < params.cols;
 {% endif %}
   var acc = {% if intMode %}{{ scalar }}(0){% else %}0.0{% endif %};
   if (in_bounds) {
     for (var row = {{ rowBegin }}; row < {{ rowEnd }}; row = row + ROW_LANES) {
       acc = acc + {{ elem }};
     }
   }
   partial[tid] = acc;
@@ -80,12 +229,32 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   if (row_lane == 0u && in_bounds) {
     var total = partial[col_lane];
     for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
       total = total + partial[lane * TILE_COLS + col_lane];
     }
 {% if splitMode %}
     partials[seg * params.outputs + outputIndex] = total;
 {% else %}
     y[outputIndex] = {{ yv }}total{{ vy }};
 {% endif %}
   }
 }

 // Tiled column-wise reduction. Each workgroup owns TILE_COLS columns;
 // ROW_LANES threads stride the rows of one column, then lane 0 folds their
 // partials. In split mode, workgroup y selects an axis segment and finalization
+// is deferred to the combine pass. Cooperative row lanes expose independent
+// reduction work when the flattened output has only a few elements.
+{% set splitMode = splitSpec is defined %}
 {% if splitMode %}
 {% set rowBegin = "row_begin + row_lane" %}
 {% set rowEnd = "row_end" %}
 {% set rowEnd = "params.rows" %}
 {% set elem = "x[inputBase + row * params.cols + col]" %}
 {% endif %}
+{% set castF32 = castF32 is defined and castF32 %}
+{% set intMode = intMode is defined and intMode %}
 {% set scalar = "f32" if castF32 else scalar %}
 {% if castF32 %}
 {% set elem = "f32(" ~ elem ~ ")" %}
 {% endif %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
+{% if not splitMode and not intMode and (op == "logsum" or op == "logsumexp") %}
+fn negative_infinity() -> f32 {
+  var bits = 0xff800000u;
+  return bitcast<f32>(bits);
+}
+{% endif %}
 const WG: u32 = {{ workgroupSize }}u;
 const TILE_COLS: u32 = {{ tileCols }}u;
 const ROW_LANES: u32 = WG / TILE_COLS;
 {% if splitMode %}
+const SPLIT: u32 = {{ splitSpec }}u;
+{% endif %}
+{% if op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("lane_identity", op, scalar) }}
+{% elif op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
 {% endif %}
+var<workgroup> partial: array<{{ scalar if (op == "max" or op == "min" or intMode) else "f32" }}, WG>;
+{% if op == "logsumexp" %}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let tid = lid.x;
   let col_lane = tid % TILE_COLS;
   let row_lane = tid / TILE_COLS;
   // 2D-folded tile index: wg.y carries the high bits past the dispatch limit.
   // The batched form reuses this same coalesced axis-0 reduction for a middle
   // axis by assigning consecutive tiles to each outer slice.
+  let tile = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let col = tile * TILE_COLS + col_lane;
   let inputBase = 0u;
   let outputIndex = col;
   let in_bounds = col < params.cols;
 {% endif %}
+{% if op == "logsumexp" %}
+{% if not splitMode %}
+  if (params.rows == 0u) {
+    if (row_lane == 0u && in_bounds) {
+      y[outputIndex] = {{ yv }}negative_infinity(){{ vy }};
+    }
+    return;
+  }
+{% endif %}
+  var local_max = F32_MIN;
+  var local_nan_count = 0.0;
+  var local_nan_value = 0.0;
+  if (in_bounds) {
+    for (var row = {{ rowBegin }}; row < {{ rowEnd }}; row = row + ROW_LANES) {
+      let value = {{ elem }};
+      if (is_nan_f32(value)) {
+        local_nan_count = local_nan_count + 1.0;
+        local_nan_value = value;
+      } else {
+        local_max = max(local_max, value);
+      }
+    }
+  }
+  partial[tid] = local_max;
+  workgroupBarrier();
+  if (row_lane == 0u && in_bounds) {
+    var max_lanes = partial[col_lane];
+    for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
+      max_lanes = max(max_lanes, partial[lane * TILE_COLS + col_lane]);
+    }
+    partial[col_lane] = max_lanes;
+  }
+  workgroupBarrier();
+  let max_value = partial[col_lane];
+  workgroupBarrier();
+  partial[tid] = local_nan_count;
+  workgroupBarrier();
+  if (row_lane == 0u && in_bounds) {
+    var nan_lanes = partial[col_lane];
+    for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
+      nan_lanes = nan_lanes + partial[lane * TILE_COLS + col_lane];
+    }
+    partial[col_lane] = nan_lanes;
+  }
+  workgroupBarrier();
+  let nan_count = partial[col_lane];
+  workgroupBarrier();
+  partial[tid] = local_nan_value;
+  workgroupBarrier();
+  if (row_lane == 0u && in_bounds) {
+    var nan_value_lanes = partial[col_lane];
+    for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
+      nan_value_lanes = nan_value_lanes + partial[lane * TILE_COLS + col_lane];
+    }
+    partial[col_lane] = nan_value_lanes;
+  }
+  workgroupBarrier();
+  let nan_value = partial[col_lane];
+  let has_nan = nan_count > 0.0;
+  let has_positive_inf = max_value > F32_MAX;
+  workgroupBarrier();
+  var acc = 0.0;
+  if (in_bounds && !has_nan && !has_positive_inf) {
+    for (var row = {{ rowBegin }}; row < {{ rowEnd }}; row = row + ROW_LANES) {
+      acc = acc + exp({{ elem }} - max_value);
+    }
+  }
+  partial[tid] = acc;
+  workgroupBarrier();
+  if (row_lane == 0u && in_bounds) {
+    var sum = partial[col_lane];
+    for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
+      sum = sum + partial[lane * TILE_COLS + col_lane];
+    }
+{% if splitMode %}
+    partials[seg * params.outputs + outputIndex] = max_value;
+    partials[(SPLIT + seg) * params.outputs + outputIndex] = sum;
+    partials[(2u * SPLIT + seg) * params.outputs + outputIndex] = select(0.0, nan_value, has_nan);
+{% else %}
+    let finite_or_inf = select(max_value + log(sum), max_value, has_positive_inf);
+    y[outputIndex] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }};
+{% endif %}
+  }
+{% else %}
+{% if op == "logsum" and not splitMode %}
+  if (params.rows == 0u) {
+    if (row_lane == 0u && in_bounds) {
+      y[outputIndex] = {{ yv }}negative_infinity(){{ vy }};
+    }
+    return;
+  }
+{% endif %}
+{% if op == "max" or op == "min" %}
+  var acc = lane_identity();
+{% elif op == "prod" %}
+  var acc = {% if intMode %}{{ scalar }}(1){% else %}1.0{% endif %};
+{% else %}
   var acc = {% if intMode %}{{ scalar }}(0){% else %}0.0{% endif %};
+{% endif %}
   if (in_bounds) {
     for (var row = {{ rowBegin }}; row < {{ rowEnd }}; row = row + ROW_LANES) {
+{% if op == "max" or op == "min" %}
+      acc = {{ op }}(acc, {{ elem }});
+{% elif op == "prod" %}
+      acc = acc * {{ elem }};
+{% elif op == "l1" %}
+      acc = acc + abs({{ elem }});
+{% elif op == "l2" or op == "sumsquare" %}
+      let value = {{ elem }};
+      acc = acc + value * value;
+{% else %}
       acc = acc + {{ elem }};
+{% endif %}
     }
   }
   partial[tid] = acc;
   if (row_lane == 0u && in_bounds) {
     var total = partial[col_lane];
     for (var lane = 1u; lane < ROW_LANES; lane = lane + 1u) {
+{% if op == "max" or op == "min" %}
+      total = {{ op }}(total, partial[lane * TILE_COLS + col_lane]);
+{% elif op == "prod" %}
+      total = total * partial[lane * TILE_COLS + col_lane];
+{% else %}
       total = total + partial[lane * TILE_COLS + col_lane];
+{% endif %}
     }
 {% if splitMode %}
     partials[seg * params.outputs + outputIndex] = total;
+{% else %}
+{% if op == "l2" and intMode %}
+    y[outputIndex] = {{ scalar }}(sqrt(f32(total)));
+{% elif op == "l2" %}
+    y[outputIndex] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
+    y[outputIndex] = {{ yv }}log(total){{ vy }};
+{% elif op == "mean" and intMode %}
+    // Integer division truncates toward zero.
+    y[outputIndex] = total / {{ scalar }}(params.rows);
+{% elif op == "mean" %}
+    y[outputIndex] = {{ yv }}total / f32(params.rows){{ vy }};
 {% else %}
     y[outputIndex] = {{ yv }}total{{ vy }};
+{% endif %}
 {% endif %}
   }
+{% endif %}
 }

build/webgpu/reduce-flat-partial.wgsl.jinja CHANGED Viewed

@@ -6,20 +6,37 @@
 // Scalar f32 bindings keep arbitrary element counts legal. The grid-stride loop
 // manually assembles full vec4 groups from contiguous scalars, and one global
 // thread folds the final zero-to-three scalar elements exactly once.
-{% set castF32 = source.castF32 is defined and source.castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
-{% if source.usesF16 is defined and source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
-var<workgroup> red: array<{{ "i32" if source.intMode else "f32" }}, WG>;
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
@@ -27,16 +44,38 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
   let gstride = nwg.x * WG;
-{% if source.intMode %}
   var acc = 0i;
 {% else %}
   var acc = 0.0;
 {% endif %}
   // Grid-stride over the flat vec4 groups (params.count4 = numel / 4, floored).
   for (var i = gid.x; i < params.count4; i = i + gstride) {
     let b = 4u * i;
-    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 }});
     acc = acc + v.x + v.y + v.z + v.w;
   }
   // Scalar tail (the 0..3 elements past the last full vec4). One global thread
   // folds it so it is counted exactly once; the count is tiny so serializing it
@@ -44,7 +83,19 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   if (gid.x == 0u) {
     for (var i = 4u * params.count4; i < params.numel; i = i + 1u) {
       let s = {{ xa }}x[i]{{ ax }};
       acc = acc + s;
     }
   }
   red[tid] = acc;
@@ -53,7 +104,15 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   loop {
     if (stride == 0u) { break; }
     if (tid < stride) {
       red[tid] = red[tid] + red[tid + stride];
     }
     stride = stride / 2u;
     workgroupBarrier();

 // Scalar f32 bindings keep arbitrary element counts legal. The grid-stride loop
 // manually assembles full vec4 groups from contiguous scalars, and one global
 // thread folds the final zero-to-three scalar elements exactly once.
+{% set intMode = intMode is defined and intMode %}
+{% set castF32 = castF32 is defined and castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
+var<workgroup> red: array<{{ "i32" if intMode else "f32" }}, WG>;
+{% if op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
   let gstride = nwg.x * WG;
+{% if intMode %}
+{% if op == "prod" %}
+  var acc = 1i;
+{% else %}
   var acc = 0i;
+{% endif %}
+{% elif op == "max" %}
+  var acc = reduction_identity();
+{% elif op == "min" %}
+  var acc = reduction_identity();
+{% elif op == "prod" %}
+  var acc = 1.0;
 {% else %}
   var acc = 0.0;
 {% endif %}
   // Grid-stride over the flat vec4 groups (params.count4 = numel / 4, floored).
   for (var i = gid.x; i < params.count4; i = i + gstride) {
     let b = 4u * i;
+    let v = vec4<{{ "i32" if intMode else "f32" }}>({{ xa }}x[b]{{ ax }}, {{ xa }}x[b + 1u]{{ ax }}, {{ xa }}x[b + 2u]{{ ax }}, {{ xa }}x[b + 3u]{{ ax }});
+{% if op == "max" %}
+    acc = max(acc, max(max(v.x, v.y), max(v.z, v.w)));
+{% elif op == "min" %}
+    acc = min(acc, min(min(v.x, v.y), min(v.z, v.w)));
+{% elif op == "prod" %}
+    acc = acc * v.x * v.y * v.z * v.w;
+{% elif op == "l1" %}
+    acc = acc + abs(v.x) + abs(v.y) + abs(v.z) + abs(v.w);
+{% elif op == "l2" or op == "sumsquare" %}
+    acc = acc + dot(v, v);
+{% else %}
     acc = acc + v.x + v.y + v.z + v.w;
+{% endif %}
   }
   // Scalar tail (the 0..3 elements past the last full vec4). One global thread
   // folds it so it is counted exactly once; the count is tiny so serializing it
   if (gid.x == 0u) {
     for (var i = 4u * params.count4; i < params.numel; i = i + 1u) {
       let s = {{ xa }}x[i]{{ ax }};
+{% if op == "max" %}
+      acc = max(acc, s);
+{% elif op == "min" %}
+      acc = min(acc, s);
+{% elif op == "prod" %}
+      acc = acc * s;
+{% elif op == "l1" %}
+      acc = acc + abs(s);
+{% elif op == "l2" or op == "sumsquare" %}
+      acc = acc + s * s;
+{% else %}
       acc = acc + s;
+{% endif %}
     }
   }
   red[tid] = acc;
   loop {
     if (stride == 0u) { break; }
     if (tid < stride) {
+{% if op == "max" %}
+      red[tid] = max(red[tid], red[tid + stride]);
+{% elif op == "min" %}
+      red[tid] = min(red[tid], red[tid + stride]);
+{% elif op == "prod" %}
+      red[tid] = red[tid] * red[tid + stride];
+{% else %}
       red[tid] = red[tid] + red[tid + stride];
+{% endif %}
     }
     stride = stride / 2u;
     workgroupBarrier();

build/webgpu/reduce-multi-axis-coop.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,130 @@

+{% macro reduce_multi_axis_offset(hasReduced) %}
+// One thread per output element walks the Cartesian product of the reduced axes,
+// linearized as reduce_linear. Specialized shapes make every input offset a sum
+// of coordinate-times-constant terms.
+fn input_offset(out_index: u32{% if hasReduced %}, reduce_linear: u32{% endif %}) -> u32 {
+  var rem = out_index;
+{% for oaxis in range(outputRank) %}
+{% set ostride = namespace(value=1) %}
+{% for j in range(oaxis + 1, outputRank) %}
+{% set ostride.value = ostride.value * outputShape[j] %}
+{% endfor %}
+{% set osafe = 1 if ostride.value == 0 else ostride.value %}
+{% if not keepDims or not reduce[oaxis] %}
+  let out_coord{{ oaxis }} = rem / {{ osafe }}u;
+{% endif %}
+  rem = rem % {{ osafe }}u;
+{% endfor %}
+{% if hasReduced %}
+  var rrem = reduce_linear;
+{% endif %}
+{% for a in range(rank) if reduce[a] %}
+{% set rstride = namespace(value=1) %}
+{% for b in range(a + 1, rank) if reduce[b] %}
+{% set rstride.value = rstride.value * dataShape[b] %}
+{% endfor %}
+{% set rsafe = 1 if rstride.value == 0 else rstride.value %}
+  let red_coord{{ a }} = rrem / {{ rsafe }}u;
+  rrem = rrem % {{ rsafe }}u;
+{% endfor %}
+{% set oc = namespace(i=0) %}
+{% for a in range(rank) %}
+{% if reduce[a] %}
+  let coord{{ a }} = red_coord{{ a }};
+{% elif keepDims %}
+  let coord{{ a }} = out_coord{{ a }};
+{% else %}
+  let coord{{ a }} = out_coord{{ oc.i }};
+{% set oc.i = oc.i + 1 %}
+{% endif %}
+{% endfor %}
+{% set src = namespace(value="coord0") %}
+{% for a in range(1, rank) %}
+{% set src.value = "(" ~ src.value ~ " * " ~ dataShape[a] ~ "u + coord" ~ a ~ ")" %}
+{% endfor %}
+  return {{ src.value }};
+}
+{%- endmacro %}
+// Cooperative multi-axis reduction with one workgroup per output element.
+// Threads take strided shares of the reduced index range, then combine their
+// partials through a workgroup tree. The shared offset helper maps each reduced
+// coordinate to the corresponding input element.
+//
+// The tree changes f32 association relative to a left-to-right fold, so the two
+// accumulation orders need not be bit-identical.
+{% set castF32 = castF32 is defined and castF32 %}
+{% set intMode = intMode is defined and intMode %}
+{% set yv = "f16(" if castF32 else "" %}
+{% set vy = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasReducedAxis = namespace(value=false) %}
+{% for a in range(rank) %}{% if reduce[a] %}{% set hasReducedAxis.value = true %}{% endif %}{% endfor %}
+{{- reduce_multi_axis_offset(hasReducedAxis.value) }}
+{% set ACC = scalar if intMode else "f32" %}
+const WG: u32 = {{ workgroupSize }}u;
+{% set mcount = namespace(value=1) %}
+{% for a in range(rank) if reduce[a] %}
+{% set mcount.value = mcount.value * dataShape[a] %}
+{% endfor %}
+const REDUCED: u32 = {{ mcount.value }}u;
+{% set ocount = namespace(value=1) %}
+{% for a in range(outputRank) %}
+{% set ocount.value = ocount.value * outputShape[a] %}
+{% endfor %}
+const OUT_COUNT: u32 = {{ ocount.value }}u;
+var<workgroup> partials: array<{{ ACC }}, WG>;
+// The reduction's combine, used identically by the per-thread accumulation and by
+// the workgroup fold, so the two cannot disagree about the operator.
+fn combine(a: {{ ACC }}, b: {{ ACC }}) -> {{ ACC }} {
+  return a + b;
+}
+@compute @workgroup_size({{ workgroupSize }})
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid3: vec3<u32>
+) {
+  // 2D-folded output index: wg.y carries the high bits past the per-axis
+  // dispatch fold width.
+  let i = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  if (i >= OUT_COUNT) {
+    return;
+  }
+  let lid = lid3.x;
+  var acc = {{ ACC }}(0);
+  for (var r = lid; r < REDUCED; r = r + WG) {
+{% set at = "x[input_offset(i, r)]" %}
+{% if castF32 %}
+{% set at = "f32(" ~ at ~ ")" %}
+{% endif %}
+    acc = combine(acc, {{ at }});
+  }
+  // Fold the per-thread accumulators. WORKGROUP_SIZE is a power of two, and every
+  // lane holds the reduction identity when its strided share was empty, so the
+  // tree needs no separate active-lane count.
+  partials[lid] = acc;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) { break; }
+    if (lid < stride) {
+      partials[lid] = combine(partials[lid], partials[lid + stride]);
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  if (lid != 0u) {
+    return;
+  }
+  let total = partials[0];
+  y[i] = {{ yv }}total{{ vy }};
+}

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

	@@ -0,0 +1,167 @@

+// Subgroup-per-row reduction for the selected operation on a contiguous last axis.
+// One SUBGROUP owns each output row and a workgroup carries as many rows as it
+// has subgroups, so the row fold is a single subgroup collective with no
+// workgroup memory or barriers. Every lane
+// reads vec4 words strided by the subgroup width, keeps them in registers, and
+// the row is folded by the selected collective.
+{% if op == "logsumexp" %}
+// The row remains in registers across the maximum, NaN census, and exponential
+// sum. Max subtraction and a positive-infinity branch define the stable result.
+{% elif op == "mean" %}
+// The finalizer divides the accumulated sum by the row width.
+{% elif op == "l2" %}
+// The finalizer takes the square root of the sum of squares.
+{% elif op == "logsum" %}
+// The finalizer takes the logarithm of the sum.
+{% endif %}
+// f16 storage widens before accumulation and narrows only at the final store.
+{% set castF32 = castF32 is defined and castF32 %}
+{% set isInt = scalar == "i32" or scalar == "u32" %}
+{% set accScalar = scalar if isInt else "f32" %}
+{% set xv = "vec4<f32>(" if castF32 else "" %}
+{% set vx = ")" if castF32 else "" %}
+{% set yv = "f16(" if castF32 else "" %}
+{% set vy = ")" if castF32 else "" %}
+enable subgroups;
+{% if usesF16Spec is defined and usesF16Spec %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
+const WG: u32 = {{ workgroupSize }}u;
+{%- if op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op, accScalar) }}
+{%- endif %}
+{%- if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{%- endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>,
+        @builtin(subgroup_invocation_id) sgLane: u32,
+        @builtin(subgroup_size) sgSize: u32) {
+  // Rows tile the folded workgroup grid; the lanes of one subgroup share a row.
+  let rowsPerWorkgroup = WG / sgSize;
+  let row = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * rowsPerWorkgroup + lid.x / sgSize;
+  // A subgroup past the last row folds an empty row and stores nothing, so
+  // every collective stays in uniform control flow.
+  let rowValid = row < params.rows;
+  let chunkLimit = select(0u, params.chunkCount, rowValid);
+  let base = row * params.chunkCount;
+{% if op == "logsumexp" %}
+  // The row stays in registers: the max census and the exp pass read the same
+  // words, so the row is fetched from memory once.
+  var localMax = F32_MIN;
+  var localNan = 0.0;
+  var localNanValue = 0.0;
+{% for i in range(vecsPerLane) %}
+  var v{{ i }} = vec4<f32>(0.0);
+  let c{{ i }} = sgLane + {{ i }}u * sgSize;
+  if (c{{ i }} < chunkLimit) {
+    v{{ i }} = {{ xv }}x[base + c{{ i }}]{{ vx }};
+{% for comp in ["x", "y", "z", "w"] %}
+    if (is_nan_f32(v{{ i }}.{{ comp }})) {
+      localNan = 1.0;
+      localNanValue = v{{ i }}.{{ comp }};
+    } else {
+      localMax = max(localMax, v{{ i }}.{{ comp }});
+    }
+{% endfor %}
+  }
+{% endfor %}
+  let rowMax = subgroupMax(localMax);
+  let nanCount = subgroupAdd(localNan);
+  let nanValue = subgroupAdd(localNanValue);
+  let hasPositiveInf = rowMax > F32_MAX;
+  let hasNan = nanCount > 0.0;
+  var acc = 0.0;
+{% for i in range(vecsPerLane) %}
+  if (c{{ i }} < chunkLimit) {
+    let e{{ i }} = select(exp(v{{ i }} - vec4<f32>(rowMax)), vec4<f32>(0.0), hasPositiveInf || hasNan);
+    acc = acc + ((e{{ i }}.x + e{{ i }}.y) + (e{{ i }}.z + e{{ i }}.w));
+  }
+{% endfor %}
+  let sum = subgroupAdd(acc);
+  if (rowValid && sgLane == 0u) {
+    let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
+    y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
+  }
+{%- else %}
+{% if op == "max" or op == "min" %}
+  let INIT: {{ accScalar }} = reduction_identity();
+{%- elif op == "prod" %}
+  let INIT: {{ accScalar }} = {{ "1.0" if not isInt else accScalar ~ "(1)" }};
+{%- else %}
+  let INIT: {{ accScalar }} = {{ "0.0" if not isInt else accScalar ~ "(0)" }};
+{%- endif %}
+  var acc4 = vec4<{{ accScalar }}>(INIT);
+  for (var c = sgLane; c < chunkLimit; c = c + sgSize) {
+    let v = {{ xv }}x[base + c]{{ vx }};
+{%- if op == "max" %}
+    acc4 = max(acc4, v);
+{%- elif op == "min" %}
+    acc4 = min(acc4, v);
+{%- elif op == "prod" %}
+    acc4 = acc4 * v;
+{%- elif op == "l1" %}
+    acc4 = acc4 + abs(v);
+{%- elif op == "l2" or op == "sumsquare" %}
+    acc4 = acc4 + v * v;
+{%- else %}
+    acc4 = acc4 + v;
+{%- endif %}
+  }
+{%- if op == "max" %}
+  let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
+  let total = subgroupMax(acc);
+{%- elif op == "min" %}
+  let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
+  let total = subgroupMin(acc);
+{%- elif op == "prod" %}
+  let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
+  let total = subgroupMul(acc);
+{%- else %}
+  let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
+  let total = subgroupAdd(acc);
+{%- endif %}
+  if (rowValid && sgLane == 0u) {
+{%- if op == "mean" and isInt %}
+    y[row] = total / {{ accScalar }}(params.chunkCount * 4u);
+{%- elif op == "mean" %}
+    y[row] = {{ yv }}total / f32(params.chunkCount * 4u){{ vy }};
+{%- elif op == "l2" and isInt %}
+    y[row] = {{ accScalar }}(sqrt(f32(total)));
+{%- elif op == "l2" %}
+    y[row] = {{ yv }}sqrt(total){{ vy }};
+{%- elif op == "logsum" %}
+    y[row] = {{ yv }}log(total){{ vy }};
+{%- else %}
+    y[row] = {{ yv }}total{{ vy }};
+{%- endif %}
+  }
+{%- endif %}
+}

build/webgpu/reduce-row-subgroup.wgsl.jinja CHANGED Viewed

@@ -1,30 +1,63 @@
-// Subgroup row reduction for a contiguous last axis. One workgroup owns each
 // output row. Threads reduce strided chunks, subgroup leaders deposit their
 // partials in workgroup memory, and thread 0 folds those slots and finalizes.
-//
-// Accumulators are f32 except for integer max/min. Mean, l2, and logsum apply
-// division, square root, or logarithm only at finalization. logsumexp uses
-// max-subtraction, bit-exact NaN detection, NaN propagation, and a +Inf
-// short-circuit. Runtime bitcasts provide max/min infinity identities because
-// WGSL constant evaluation rejects infinite constants. f16 storage is widened
-// before accumulation and narrowed only for the final store.
-{% set castF32 = source.castF32 is defined and source.castF32 %}
 {% set scalar = "f32" if castF32 else scalar %}
-{% set xv = ("vec4<f32>(" if source.vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
 enable subgroups;
-{% if source.usesF16 is defined and source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 var<workgroup> wgPartial: array<{{ scalar }}, WG>;
 {% macro emit_reduce(name, collective, combine) %}
@@ -44,26 +77,141 @@ fn {{ name }}(value: {{ scalar }}, sgLid: u32, sgId: u32, numSg: u32) -> {{ scal
   workgroupBarrier();
   return total;
 }
-{%- endmacro %}{{ emit_reduce("reduce_row", "subgroupAdd", "total = total + wgPartial[i];") }}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
         @builtin(subgroup_invocation_id) sgLid: u32,
         @builtin(subgroup_id) sgId: u32,
         @builtin(num_subgroups) numSg: u32) {
-  let row = wg.x + wg.y * nwg.x;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
-  let base = row * params.chunkCount;  let INIT: f32 = 0.0;{% if source.vec4 %}
   var acc4 = vec4<{{ scalar }}>(INIT);
   for (var c = tid; c < params.chunkCount; c = c + WG) {
-    let v = {{ xv }}x[base + c]{{ vx }};    acc4 = acc4 + v;  }  let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);{% else %}
   var acc = INIT;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
-    let v = {{ xv }}x[base + c]{{ vx }};    acc = acc + v;  }
 {%- endif %}
   let total = reduce_row(acc, sgLid, sgId, numSg);
-  if (tid == 0u) {    y[row] = {{ yv }}total{{ vy }};  }}

+// Subgroup reduction for a contiguous last axis. One workgroup owns each
 // output row. Threads reduce strided chunks, subgroup leaders deposit their
 // partials in workgroup memory, and thread 0 folds those slots and finalizes.
+{% if op == "mean" %}
+// The finalizer divides the f32 sum by the row width.
+{% elif op == "l2" %}
+// The finalizer takes the square root of the sum of squares.
+{% elif op == "logsum" %}
+// The finalizer takes the logarithm of the sum.
+{% elif op == "logsumexp" %}
+// Max subtraction, explicit NaN propagation, and a positive-infinity branch
+// define the stable log-sum-exp result.
+{% elif op == "max" or op == "min" %}
+// Integer values retain their native type; f32 identities are constructed from
+// IEEE-754 bit patterns because WGSL rejects infinite constants.
+{% endif %}
+// f16 storage widens before accumulation and narrows only at the final store.
+{% set castF32 = castF32 is defined and castF32 %}
 {% set scalar = "f32" if castF32 else scalar %}
+{% set xv = ("vec4<f32>(" if vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
 enable subgroups;
+{% if usesF16Spec is defined and usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
+{%- if op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op, scalar) }}
+{%- endif %}
+{%- if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{%- endif %}
 var<workgroup> wgPartial: array<{{ scalar }}, WG>;
 {% macro emit_reduce(name, collective, combine) %}
   workgroupBarrier();
   return total;
 }
+{%- endmacro %}
+{%- if op == "max" %}
+{{ emit_reduce("reduce_row", "subgroupMax", "total = max(total, wgPartial[i]);") }}
+{%- elif op == "min" %}
+{{ emit_reduce("reduce_row", "subgroupMin", "total = min(total, wgPartial[i]);") }}
+{%- elif op == "prod" %}
+{{ emit_reduce("reduce_row", "subgroupMul", "total = total * wgPartial[i];") }}
+{%- elif op == "logsumexp" %}
+{{ emit_reduce("reduce_row_add", "subgroupAdd", "total = total + wgPartial[i];") }}
+{{ emit_reduce("reduce_row_max", "subgroupMax", "total = max(total, wgPartial[i]);") }}
+{%- else %}
+{{ emit_reduce("reduce_row", "subgroupAdd", "total = total + wgPartial[i];") }}
+{%- endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
         @builtin(subgroup_invocation_id) sgLid: u32,
         @builtin(subgroup_id) sgId: u32,
         @builtin(num_subgroups) numSg: u32) {
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
+  let base = row * params.chunkCount;
+{%- if op == "logsumexp" %}
+  var localMax = F32_MIN;
+  var localNan = 0.0;
+  var localNanValue = 0.0;
+  for (var c = tid; c < params.chunkCount; c = c + WG) {
+    let v = {{ xv }}x[base + c]{{ vx }};
+{%- if vec4 %}
+{% for comp in ["x", "y", "z", "w"] %}
+    if (is_nan_f32(v.{{ comp }})) {
+      localNan = 1.0;
+      localNanValue = v.{{ comp }};
+    } else {
+      localMax = max(localMax, v.{{ comp }});
+    }
+{%- endfor %}
+{% else %}
+    if (is_nan_f32(v)) {
+      localNan = 1.0;
+      localNanValue = v;
+    } else {
+      localMax = max(localMax, v);
+    }
+{%- endif %}
+  }
+  let rowMax = reduce_row_max(localMax, sgLid, sgId, numSg);
+  let nanCount = reduce_row_add(localNan, sgLid, sgId, numSg);
+  let nanValue = reduce_row_add(localNanValue, sgLid, sgId, numSg);
+  let hasPositiveInf = rowMax > F32_MAX;
+  let hasNan = nanCount > 0.0;
+  var acc = 0.0;
+  for (var c = tid; c < params.chunkCount; c = c + WG) {
+    let v = {{ xv }}x[base + c]{{ vx }};
+{%- if vec4 %}
+    let e = select(exp(v - vec4<f32>(rowMax)), vec4<f32>(0.0), hasPositiveInf || hasNan);
+    acc = acc + ((e.x + e.y) + (e.z + e.w));
+{%- else %}
+    acc = acc + select(exp(v - rowMax), 0.0, hasPositiveInf || hasNan);
+{%- endif %}
+  }
+  let sum = reduce_row_add(acc, sgLid, sgId, numSg);
+  if (tid == 0u) {
+    let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
+    y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
+  }
+{%- else %}
+{% if op == "max" or op == "min" %}
+  let INIT: {{ scalar }} = reduction_identity();
+{%- elif op == "prod" %}
+  let INIT: f32 = 1.0;
+{%- else %}
+  let INIT: f32 = 0.0;
+{%- endif %}
+{% if vec4 %}
   var acc4 = vec4<{{ scalar }}>(INIT);
   for (var c = tid; c < params.chunkCount; c = c + WG) {
+    let v = {{ xv }}x[base + c]{{ vx }};
+{%- if op == "max" %}
+    acc4 = max(acc4, v);
+{%- elif op == "min" %}
+    acc4 = min(acc4, v);
+{%- elif op == "prod" %}
+    acc4 = acc4 * v;
+{%- elif op == "l1" %}
+    acc4 = acc4 + abs(v);
+{%- elif op == "l2" or op == "sumsquare" %}
+    acc4 = acc4 + v * v;
+{%- else %}
+    acc4 = acc4 + v;
+{%- endif %}
+  }
+{%- if op == "max" %}
+  let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
+{%- elif op == "min" %}
+  let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
+{%- elif op == "prod" %}
+  let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
+{%- else %}
+  let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
+{%- endif %}
+{% else %}
   var acc = INIT;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
+    let v = {{ xv }}x[base + c]{{ vx }};
+{%- if op == "max" %}
+    acc = max(acc, v);
+{%- elif op == "min" %}
+    acc = min(acc, v);
+{%- elif op == "prod" %}
+    acc = acc * v;
+{%- elif op == "l1" %}
+    acc = acc + abs(v);
+{%- elif op == "l2" or op == "sumsquare" %}
+    acc = acc + v * v;
+{%- else %}
+    acc = acc + v;
+{%- endif %}
+  }
 {%- endif %}
   let total = reduce_row(acc, sgLid, sgId, numSg);
+  if (tid == 0u) {
+{%- if op == "mean" %}
+    y[row] = {{ yv }}total / f32(params.cols){{ vy }};
+{%- elif op == "l2" %}
+    y[row] = {{ yv }}sqrt(total){{ vy }};
+{%- elif op == "logsum" %}
+    y[row] = {{ yv }}log(total){{ vy }};
+{%- else %}
+    y[row] = {{ yv }}total{{ vy }};
+{%- endif %}
+  }
+{%- endif %}
+}

build/webgpu/reduce-row-tree.wgsl.jinja CHANGED Viewed

@@ -1,70 +1,266 @@
-// Portable one-workgroup-per-row reduction for the Reduce value family.
 // Threads stride a contiguous row, accumulate locally, and fold their values
 // through a shared-memory tree without relying on subgroups.
-//
-// Max, min, and product use the native f32/i32/u32 value type. Additive
-// transforms accumulate float inputs in f32 but retain integer accumulation
-// for sum, L1, and sum-of-squares. Mean, L2, log-sum, and log-sum-exp are
-// f32-only because their finalizers divide, take a square root, or take a log.
 // f16 storage widens through f32 for both accumulation and the shared tree,
 // then narrows only at the final store.
-{% set castF32 = source.castF32 is defined and source.castF32 %}
 {% set scalar = "f32" if castF32 else scalar %}
-{% set xv = ("vec4<f32>(" if source.vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-{% if source.usesF16 is defined and source.usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
 const WG: u32 = {{ workgroupSize }}u;
 {% set is_int = scalar == "i32" or scalar == "u32" %}
-{% set intAdditive = is_int and (source.op == "sum" or source.op == "l1" or source.op == "sumsquare"
-                                 or source.op == "l2" or source.op == "mean") %}
-{% set accType = scalar if source.op == "max" or source.op == "min" or source.op == "prod" or intAdditive else "f32" %}
 fn identity() -> {{ accType }} { return {{ accType }}(0); }
 fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
   return a + b;
 }
 var<workgroup> partial: array<{{ accType }}, WG>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
-  let row = wg.x + wg.y * nwg.x;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
-{% if source.vec4 %}
   let base = row * params.chunkCount;
 {% else %}
   let base = row * params.cols;
 {% endif %}
-{% if source.vec4 %}
   var acc4 = vec4<{{ accType }}>(identity());
   for (var col = tid; col < params.chunkCount; col = col + WG) {
     let value = {{ xv }}x[base + col]{{ vx }};
     acc4 = acc4 + value;
   }
   let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
 {% else %}
   var acc = identity();
   for (var col = tid; col < params.cols; col = col + WG) {
     let value = {{ xv }}x[base + col]{{ vx }};
     acc = combine(acc, value);
   }
 {% endif %}
   partial[tid] = acc;
@@ -79,6 +275,19 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   if (tid == 0u) {
     let total = partial[0];
     y[row] = {{ yv }}total{{ vy }};
   }
 }

+// Portable one-workgroup-per-row reduction for the selected operation.
 // Threads stride a contiguous row, accumulate locally, and fold their values
 // through a shared-memory tree without relying on subgroups.
+{% if op == "max" or op == "min" or op == "prod" %}
+// Values accumulate in their native f32/i32/u32 type.
+{% elif op == "mean" %}
+// The tree sums values before the final division by the row width.
+{% elif op == "l2" %}
+// The tree sums squared values before the final square root.
+{% elif op == "logsum" %}
+// The tree sums values before the final logarithm.
+{% elif op == "logsumexp" %}
+// A max-subtracted exponential sum handles NaN and positive-infinity rows.
+{% endif %}
 // f16 storage widens through f32 for both accumulation and the shared tree,
 // then narrows only at the final store.
+{% set isVec4 = vec4 is defined and vec4 %}
+{% set rowIsEmpty = "params.chunkCount == 0u" if isVec4 else "params.cols == 0u" %}
+{% set castF32 = castF32 is defined and castF32 %}
 {% set scalar = "f32" if castF32 else scalar %}
+{% set xv = ("vec4<f32>(" if vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
+{% if scalar != "i32" and scalar != "u32" and (op == "logsum" or op == "logsumexp") %}
+fn negative_infinity() -> f32 {
+  var bits = 0xff800000u;
+  return bitcast<f32>(bits);
+}
+{% endif %}
 const WG: u32 = {{ workgroupSize }}u;
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+var<workgroup> partial: array<f32, WG>;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
+fn {{ name }}(value: f32, tid: u32) -> f32 {
+  {{ buffer }}[tid] = value;
+  workgroupBarrier();
+  // Ceil-halving keeps every lane when the workgroup size is not a power of
+  // two. For even n this matches the power-of-two tree order; for odd n, lanes
+  // [0, n-half) fold the upper tail while the middle lane carries forward.
+  var n: u32 = {{ wg }};
+  loop {
+    let half = (n + 1u) / 2u;
+    if (tid < n - half) {
+{% if mode == "max" %}
+      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
+{% else %}
+      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
+{% endif %}
+    }
+    workgroupBarrier();
+    n = half;
+    if (n == 1u) {
+      break;
+    }
+  }
+  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
+  // slot 0 here, so the next call's first store must not run until all lanes have read it.
+  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
+  let reduced = {{ buffer }}[0];
+{% if trailingBarrier %}
+  workgroupBarrier();
+{% endif %}
+  return reduced;
+}
+{% endmacro %}
+{{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
+{{ wgsl_tree_reduce_f32("reduce_max", "max", "partial", "WG") }}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u
+      && (bits & 0x007fffffu) != 0u;
+}
+{% else %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
+{% set intAdditive = is_int and (op == "sum" or op == "l1" or op == "sumsquare"
+                                 or op == "l2" or op == "mean") %}
+{% set accType = scalar if op == "max" or op == "min" or op == "prod" or intAdditive else "f32" %}
+{% if op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("identity", op, accType) }}
+{% elif op == "prod" and accType == "i32" %}
+fn identity() -> i32 { return 1i; }
+{% elif op == "prod" %}
+fn identity() -> {{ accType }} { return {{ accType }}(1); }
+{% else %}
 fn identity() -> {{ accType }} { return {{ accType }}(0); }
+{% endif %}
 fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
+{% if op == "max" %}
+  return max(a, b);
+{% elif op == "min" %}
+  return min(a, b);
+{% elif op == "prod" %}
+  return a * b;
+{% else %}
   return a + b;
+{% endif %}
 }
 var<workgroup> partial: array<{{ accType }}, WG>;
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
+{% if vec4 %}
   let base = row * params.chunkCount;
 {% else %}
   let base = row * params.cols;
 {% endif %}
+{% if op == "logsumexp" %}
+  if ({{ rowIsEmpty }}) {
+    if (tid == 0u) {
+      y[row] = {{ yv }}negative_infinity(){{ vy }};
+    }
+    return;
+  }
+  var localMax = F32_MIN;
+  var localNan = 0.0;
+  var localNanValue = 0.0;
+{% if vec4 %}
+  for (var col = tid; col < params.chunkCount; col = col + WG) {
+    let value = {{ xv }}x[base + col]{{ vx }};
+{% for component in ["x", "y", "z", "w"] %}
+    if (is_nan_f32(value.{{ component }})) {
+      localNan = 1.0;
+      localNanValue = value.{{ component }};
+    } else {
+      localMax = max(localMax, value.{{ component }});
+    }
+{% endfor %}
+  }
+{% else %}
+  for (var col = tid; col < params.cols; col = col + WG) {
+    let value = {{ xv }}x[base + col]{{ vx }};
+    if (is_nan_f32(value)) {
+      localNan = 1.0;
+      localNanValue = value;
+    } else {
+      localMax = max(localMax, value);
+    }
+  }
+{% endif %}
+  let rowMax = reduce_max(localMax, tid);
+  workgroupBarrier();
+  let nanCount = reduce_sum(localNan, tid);
+  workgroupBarrier();
+  let nanValue = reduce_sum(localNanValue, tid);
+  workgroupBarrier();
+  let hasPositiveInf = rowMax > F32_MAX;
+  let hasNan = nanCount > 0.0;
+  var acc = 0.0;
+{% if vec4 %}
+  for (var col = tid; col < params.chunkCount; col = col + WG) {
+    let value = {{ xv }}x[base + col]{{ vx }};
+    let exponentials = select(exp(value - vec4<f32>(rowMax)), vec4<f32>(0.0),
+                              hasPositiveInf || hasNan);
+    acc = acc + (exponentials.x + exponentials.y)
+              + (exponentials.z + exponentials.w);
+  }
+{% else %}
+  for (var col = tid; col < params.cols; col = col + WG) {
+    acc = acc + select(exp({{ xv }}x[base + col]{{ vx }} - rowMax), 0.0,
+                       hasPositiveInf || hasNan);
+  }
+{% endif %}
+  let sum = reduce_sum(acc, tid);
+  if (tid == 0u) {
+    let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
+    y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
+  }
+{% else %}
+{% if op == "mean" %}
+  if ({{ rowIsEmpty }}) {
+{% if is_int %}
+    if (tid == 0u) { y[row] = {{ scalar }}(0); }
+{% else %}
+    if (tid == 0u) { y[row] = {{ yv }}0.0{{ vy }}; }
+{% endif %}
+    return;
+  }
+{% elif op == "logsum" %}
+  if ({{ rowIsEmpty }}) {
+    if (tid == 0u) { y[row] = {{ yv }}negative_infinity(){{ vy }}; }
+    return;
+  }
+{% endif %}
+{% if vec4 %}
   var acc4 = vec4<{{ accType }}>(identity());
   for (var col = tid; col < params.chunkCount; col = col + WG) {
     let value = {{ xv }}x[base + col]{{ vx }};
+{% if op == "max" %}
+    acc4 = max(acc4, value);
+{% elif op == "min" %}
+    acc4 = min(acc4, value);
+{% elif op == "prod" %}
+    acc4 = acc4 * value;
+{% elif op == "l1" %}
+    acc4 = acc4 + abs(value);
+{% elif op == "l2" or op == "sumsquare" %}
+    acc4 = acc4 + value * value;
+{% else %}
     acc4 = acc4 + value;
+{% endif %}
   }
+{% if op == "max" %}
+  let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
+{% elif op == "min" %}
+  let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
+{% elif op == "prod" %}
+  let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
+{% else %}
   let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
+{% endif %}
 {% else %}
   var acc = identity();
   for (var col = tid; col < params.cols; col = col + WG) {
     let value = {{ xv }}x[base + col]{{ vx }};
+{% if op == "l1" %}
+    acc = combine(acc, abs(value));
+{% elif op == "l2" or op == "sumsquare" %}
+    acc = combine(acc, value * value);
+{% else %}
     acc = combine(acc, value);
+{% endif %}
   }
 {% endif %}
   partial[tid] = acc;
   if (tid == 0u) {
     let total = partial[0];
+{% if op == "l2" and is_int %}
+    y[row] = {{ scalar }}(sqrt(f32(total)));
+{% elif op == "l2" %}
+    y[row] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
+    y[row] = {{ yv }}log(total){{ vy }};
+{% elif op == "mean" and is_int %}
+    y[row] = total / {{ scalar }}(params.cols);
+{% elif op == "mean" %}
+    y[row] = {{ yv }}total / f32(params.cols){{ vy }};
+{% else %}
     y[row] = {{ yv }}total{{ vy }};
+{% endif %}
   }
+{% endif %}
 }

build/webgpu/reduce-serial-axis.wgsl.jinja CHANGED Viewed

@@ -1,87 +1,38 @@
-// Serial one-thread-per-output reduction for the no-feature tier. f16 storage
-// is widened before every accumulation and narrowed only for the final store.
-{% set castF32 = source.castF32 is defined and source.castF32 %}
-{% set intMode = source.intMode is defined and source.intMode %}
-{% set yv = "f16(" if castF32 else "" %}
-{% set vy = ")" if castF32 else "" %}
-{% if source.usesF16 is defined and source.usesF16 %}
-enable f16;
-{% endif %}
-{{ env.wgsl.resourceDeclarations }}
-/* Exact max/min reduction identities. WGSL rejects infinity during constant
- * evaluation, so f32 identities are constructed at runtime from their
- * IEEE-754 bit patterns. */
-{% if source.indexing == "rankn" %}
-fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
-  var rem = out_index;
-{% for axis in range(source.outputRank) %}
-{% set out_stride = namespace(value=1) %}
-{% for j in range(axis + 1, source.outputRank) %}
-{% set out_stride.value = out_stride.value * source.outputShape[j] %}
-{% endfor %}
-{% set safe_out_stride = 1 if out_stride.value == 0 else out_stride.value %}
-{% if not source.keepDims or axis != source.axis %}
-  let out_coord{{ axis }} = rem / {{ safe_out_stride }}u;
-{% endif %}
-  rem = rem % {{ safe_out_stride }}u;
-{% endfor %}
-{% for axis in range(source.rank) %}
-{% if axis == source.axis %}
-  let coord{{ axis }} = reduce_index;
-{% elif source.keepDims %}
-  let coord{{ axis }} = out_coord{{ axis }};
-{% elif axis < source.axis %}
-  let coord{{ axis }} = out_coord{{ axis }};
-{% else %}
-  let coord{{ axis }} = out_coord{{ axis - 1 }};
-{% endif %}
-{% endfor %}
-{% set src = namespace(value="coord0") %}
-{% for axis in range(1, source.rank) %}
-{% set src.value = "(" ~ src.value ~ " * " ~ source.dataShape[axis] ~ "u + coord" ~ axis ~ ")" %}
-{% endfor %}
-  return {{ src.value }};
-}
-{% endif %}
-{% if source.indexing == "multiaxis" %}
-{% set hasReducedAxis = namespace(value=false) %}
-{% for a in range(source.rank) %}{% if source.reduce[a] %}{% set hasReducedAxis.value = true %}{% endif %}{% endfor %}
 // One thread per output element walks the Cartesian product of the reduced axes,
 // linearized as reduce_linear. Specialized shapes make every input offset a sum
 // of coordinate-times-constant terms.
-fn input_offset(out_index: u32{% if hasReducedAxis.value %}, reduce_linear: u32{% endif %}) -> u32 {
   var rem = out_index;
-{% for oaxis in range(source.outputRank) %}
 {% set ostride = namespace(value=1) %}
-{% for j in range(oaxis + 1, source.outputRank) %}
-{% set ostride.value = ostride.value * source.outputShape[j] %}
 {% endfor %}
 {% set osafe = 1 if ostride.value == 0 else ostride.value %}
-{% if not source.keepDims or not source.reduce[oaxis] %}
   let out_coord{{ oaxis }} = rem / {{ osafe }}u;
 {% endif %}
   rem = rem % {{ osafe }}u;
 {% endfor %}
-{% if hasReducedAxis.value %}
   var rrem = reduce_linear;
 {% endif %}
-{% for a in range(source.rank) if source.reduce[a] %}
 {% set rstride = namespace(value=1) %}
-{% for b in range(a + 1, source.rank) if source.reduce[b] %}
-{% set rstride.value = rstride.value * source.dataShape[b] %}
 {% endfor %}
 {% set rsafe = 1 if rstride.value == 0 else rstride.value %}
   let red_coord{{ a }} = rrem / {{ rsafe }}u;
   rrem = rrem % {{ rsafe }}u;
 {% endfor %}
 {% set oc = namespace(i=0) %}
-{% for a in range(source.rank) %}
-{% if source.reduce[a] %}
   let coord{{ a }} = red_coord{{ a }};
-{% elif source.keepDims %}
   let coord{{ a }} = out_coord{{ a }};
 {% else %}
   let coord{{ a }} = out_coord{{ oc.i }};
@@ -89,54 +40,248 @@ fn input_offset(out_index: u32{% if hasReducedAxis.value %}, reduce_linear: u32{
 {% endif %}
 {% endfor %}
 {% set src = namespace(value="coord0") %}
-{% for a in range(1, source.rank) %}
-{% set src.value = "(" ~ src.value ~ " * " ~ source.dataShape[a] ~ "u + coord" ~ a ~ ")" %}
 {% endfor %}
   return {{ src.value }};
 }
 {% endif %}
-{% if source.indexing == "multiaxis" %}
 {% set mcount = namespace(value=1) %}
-{% for a in range(source.rank) if source.reduce[a] %}
-{% set mcount.value = mcount.value * source.dataShape[a] %}
 {% endfor %}
 {% set count = mcount.value ~ "u" %}
 {% if hasReducedAxis.value %}
 {% set at = "x[input_offset(i, r)]" %}
 {% else %}
 {% set at = "x[input_offset(i)]" %}
 {% endif %}
-{% elif source.indexing == "rankn" %}
 {% set count = "params.axisDim" %}
 {% set at = "x[input_offset(i, r)]" %}
 {% elif axis == 0 %}
 {% set count = "params.rows" %}
 {% set at = "x[r * params.cols + i]" %}
 {% else %}
 {% set count = "params.cols" %}
 {% set at = "x[i * params.cols + r]" %}
 {% endif %}
 {% if castF32 %}
 {% set at = "f32(" ~ at ~ ")" %}
 {% endif %}
 @compute @workgroup_size({{ reduceWorkgroupSize }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
-  let i = gid.x + gid.y * nwg.x * {{ reduceWorkgroupSize }}u;
   if (i >= params.outCount) {
     return;
   }
 {% if intMode %}
   // Integer reduction accumulates in the output type. sum/prod/l1/sumsquare stay
   // in integer arithmetic; l2/logsum widen for sqrt/log and truncate afterward.
-  var acc = {{ scalar }}({{ "1" if source.op == "prod" else "0" }});
 {% else %}
-  var acc = {{ "1.0" if source.op == "prod" else "0.0" }};
 {% endif %}
   for (var r = 0u; r < {{ count }}; r = r + 1u) {
     acc = acc + {{ at }};
   }
   y[i] = {{ yv }}acc{{ vy }};
 }

+{% macro reduce_multi_axis_offset(hasReduced) %}
 // One thread per output element walks the Cartesian product of the reduced axes,
 // linearized as reduce_linear. Specialized shapes make every input offset a sum
 // of coordinate-times-constant terms.
+fn input_offset(out_index: u32{% if hasReduced %}, reduce_linear: u32{% endif %}) -> u32 {
   var rem = out_index;
+{% for oaxis in range(outputRank) %}
 {% set ostride = namespace(value=1) %}
+{% for j in range(oaxis + 1, outputRank) %}
+{% set ostride.value = ostride.value * outputShape[j] %}
 {% endfor %}
 {% set osafe = 1 if ostride.value == 0 else ostride.value %}
+{% if not keepDims or not reduce[oaxis] %}
   let out_coord{{ oaxis }} = rem / {{ osafe }}u;
 {% endif %}
   rem = rem % {{ osafe }}u;
 {% endfor %}
+{% if hasReduced %}
   var rrem = reduce_linear;
 {% endif %}
+{% for a in range(rank) if reduce[a] %}
 {% set rstride = namespace(value=1) %}
+{% for b in range(a + 1, rank) if reduce[b] %}
+{% set rstride.value = rstride.value * dataShape[b] %}
 {% endfor %}
 {% set rsafe = 1 if rstride.value == 0 else rstride.value %}
   let red_coord{{ a }} = rrem / {{ rsafe }}u;
   rrem = rrem % {{ rsafe }}u;
 {% endfor %}
 {% set oc = namespace(i=0) %}
+{% for a in range(rank) %}
+{% if reduce[a] %}
   let coord{{ a }} = red_coord{{ a }};
+{% elif keepDims %}
   let coord{{ a }} = out_coord{{ a }};
 {% else %}
   let coord{{ a }} = out_coord{{ oc.i }};
 {% endif %}
 {% endfor %}
 {% set src = namespace(value="coord0") %}
+{% for a in range(1, rank) %}
+{% set src.value = "(" ~ src.value ~ " * " ~ dataShape[a] ~ "u + coord" ~ a ~ ")" %}
+{% endfor %}
+  return {{ src.value }};
+}
+{%- endmacro %}
+// Serial one-thread-per-output reduction for the no-feature tier. f16 storage
+// is widened before every accumulation and narrowed only for the final store.
+{% set castF32 = castF32 is defined and castF32 %}
+{% set logicalBool = logicalBool is defined and logicalBool %}
+{% set intMode = intMode is defined and intMode %}
+{% set yv = "f16(" if castF32 else "" %}
+{% set vy = ")" if castF32 else "" %}
+{% if usesF16Spec is defined and usesF16Spec %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
+{% if not intMode and (op == "logsum" or op == "logsumexp") %}
+fn negative_infinity() -> f32 {
+  var bits = 0xff800000u;
+  return bitcast<f32>(bits);
+}
+{% endif %}
+{% if op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("empty_value", op, "f32" if castF32 else scalar) }}
+{% elif op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+{% if not intMode %}
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% endif %}
+{% endif %}
+{% if indexing == "rankn" %}
+fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
+  var rem = out_index;
+{% for axis in range(outputRank) %}
+{% set out_stride = namespace(value=1) %}
+{% for j in range(axis + 1, outputRank) %}
+{% set out_stride.value = out_stride.value * outputShape[j] %}
+{% endfor %}
+{% set safe_out_stride = 1 if out_stride.value == 0 else out_stride.value %}
+{% if not keepDims or axis != axisSpec %}
+  let out_coord{{ axis }} = rem / {{ safe_out_stride }}u;
+{% endif %}
+  rem = rem % {{ safe_out_stride }}u;
+{% endfor %}
+{% for axis in range(rank) %}
+{% if axis == axisSpec %}
+  let coord{{ axis }} = reduce_index;
+{% elif keepDims %}
+  let coord{{ axis }} = out_coord{{ axis }};
+{% elif axis < axisSpec %}
+  let coord{{ axis }} = out_coord{{ axis }};
+{% else %}
+  let coord{{ axis }} = out_coord{{ axis - 1 }};
+{% endif %}
+{% endfor %}
+{% set src = namespace(value="coord0") %}
+{% for axis in range(1, rank) %}
+{% set src.value = "(" ~ src.value ~ " * " ~ dataShape[axis] ~ "u + coord" ~ axis ~ ")" %}
 {% endfor %}
   return {{ src.value }};
 }
 {% endif %}
+{% if indexing == "multiaxis" %}
+{% set hasReducedAxis = namespace(value=false) %}
+{% for a in range(rank) %}{% if reduce[a] %}{% set hasReducedAxis.value = true %}{% endif %}{% endfor %}
+{{- reduce_multi_axis_offset(hasReducedAxis.value) }}
+{% endif %}
+{% if indexing == "multiaxis" %}
 {% set mcount = namespace(value=1) %}
+{% for a in range(rank) if reduce[a] %}
+{% set mcount.value = mcount.value * dataShape[a] %}
 {% endfor %}
 {% set count = mcount.value ~ "u" %}
 {% if hasReducedAxis.value %}
 {% set at = "x[input_offset(i, r)]" %}
+{% set at_first = "x[input_offset(i, 0u)]" %}
 {% else %}
 {% set at = "x[input_offset(i)]" %}
+{% set at_first = "x[input_offset(i)]" %}
 {% endif %}
+{% elif indexing == "rankn" %}
 {% set count = "params.axisDim" %}
 {% set at = "x[input_offset(i, r)]" %}
+{% set at_first = "x[input_offset(i, 0u)]" %}
 {% elif axis == 0 %}
 {% set count = "params.rows" %}
 {% set at = "x[r * params.cols + i]" %}
+{% set at_first = "x[i]" %}
 {% else %}
 {% set count = "params.cols" %}
 {% set at = "x[i * params.cols + r]" %}
+{% set at_first = "x[i * params.cols]" %}
 {% endif %}
 {% if castF32 %}
 {% set at = "f32(" ~ at ~ ")" %}
+{% set at_first = "f32(" ~ at_first ~ ")" %}
 {% endif %}
 @compute @workgroup_size({{ reduceWorkgroupSize }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width (outputs > 16.7M elements).
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ reduceWorkgroupSize }}u;
   if (i >= params.outCount) {
     return;
   }
+{% if op == "logsumexp" %}
+{% if intMode %}
+  // Integer logsumexp widens each element for exp/log, then truncates the result
+  // back to the integer output type.
+  // Integers are never NaN, so the f32 NaN-propagation path is unnecessary here.
+  if ({{ count }} == 0u) {
+    y[i] = {{ scalar }}(F32_MIN);
+    return;
+  }
+  var m = F32_MIN;
+  for (var r = 0u; r < {{ count }}; r = r + 1u) {
+    m = max(m, f32({{ at }}));
+  }
+  var acc = 0.0;
+  for (var r = 0u; r < {{ count }}; r = r + 1u) {
+    acc = acc + exp(f32({{ at }}) - m);
+  }
+  y[i] = {{ scalar }}(m + log(acc));
+{% else %}
+  if ({{ count }} == 0u) {
+    y[i] = {{ yv }}negative_infinity(){{ vy }};
+    return;
+  }
+  var m = F32_MIN;
+  var has_nan = false;
+  var nan_value = 0.0;
+  for (var r = 0u; r < {{ count }}; r = r + 1u) {
+    let value = {{ at }};
+    if (is_nan_f32(value)) {
+      has_nan = true;
+      nan_value = value;
+    } else {
+      m = max(m, value);
+    }
+  }
+  if (has_nan) {
+    y[i] = {{ yv }}nan_value{{ vy }};
+    return;
+  }
+  if (m > F32_MAX) {
+    y[i] = {{ yv }}m{{ vy }};
+    return;
+  }
+  var acc = 0.0;
+  for (var r = 0u; r < {{ count }}; r = r + 1u) {
+    acc = acc + exp({{ at }} - m);
+  }
+  y[i] = {{ yv }}m + log(acc){{ vy }};
+{% endif %}
+{% elif op == "max" or op == "min" %}
+  if ({{ count }} == 0u) {
+    y[i] = {{ yv }}empty_value(){{ vy }};
+    return;
+  }
+  var acc = {{ at_first }};
+  for (var r = 1u; r < {{ count }}; r = r + 1u) {
+    acc = {{ op }}(acc, {{ at }});
+  }
+  y[i] = {{ yv }}acc{{ vy }};
+{% else %}
+{% if op == "mean" %}
+  if ({{ count }} == 0u) {
+{% if intMode %}
+    y[i] = {{ scalar }}(0);
+{% else %}
+    y[i] = {{ yv }}0.0{{ vy }};
+{% endif %}
+    return;
+  }
+{% elif op == "logsum" %}
+  if ({{ count }} == 0u) {
+    y[i] = {{ yv }}negative_infinity(){{ vy }};
+    return;
+  }
+{% endif %}
 {% if intMode %}
   // Integer reduction accumulates in the output type. sum/prod/l1/sumsquare stay
   // in integer arithmetic; l2/logsum widen for sqrt/log and truncate afterward.
+  var acc = {{ scalar }}({{ "1" if op == "prod" else "0" }});
 {% else %}
+  var acc = {{ "1.0" if op == "prod" else "0.0" }};
 {% endif %}
   for (var r = 0u; r < {{ count }}; r = r + 1u) {
+{% if op == "prod" %}
+    acc = acc * {{ at }};
+{% elif op == "l1" %}
+    acc = acc + abs({{ at }});
+{% elif op == "l2" or op == "sumsquare" %}
+    let value = {{ at }};
+    acc = acc + value * value;
+{% else %}
     acc = acc + {{ at }};
+{% endif %}
   }
+{% if op == "l2" %}
+{% if intMode %}
+  y[i] = {{ scalar }}(sqrt(f32(acc)));
+{% else %}
+  y[i] = {{ yv }}sqrt(acc){{ vy }};
+{% endif %}
+{% elif op == "logsum" %}
+  y[i] = {{ yv }}log(acc){{ vy }};
+{% elif op == "mean" %}
+{% if intMode %}
+  y[i] = acc / {{ scalar }}({{ count }});
+{% else %}
+  y[i] = {{ yv }}acc / f32({{ count }}){{ vy }};
+{% endif %}
+{% else %}
   y[i] = {{ yv }}acc{{ vy }};
+{% endif %}
+{% endif %}
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.ReduceSum",
   "fixtureArrays": {
     "rank3_axis1_middle_no_keepdims_input_x": [1, 2, 3, 4, -1, -2, -3, -4, 0.5, 1.5, 2.5, 3.5, 10, 20, 30, 40, -10, -20, -30, -40, 2, 4, 6, 8],
     "rank3_axis0_no_keepdims_input_x": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
@@ -8,7 +7,7 @@
     {
       "name": "all_axes_flat_rank1_boundary_8192",
       "provenance": {
-        "notes": "The parallel full-reduction threshold must supersede the rank1 serial and row-reduction fallbacks."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [8192], "data": { "kind": "constant", "value": 1.0 } } },
@@ -91,7 +90,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
-        "notes": "Tiled axis-0 companion: 64 finite subnormal values reduce to a larger finite subnormal in each column."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [64, 16], "data": { "kind": "constant", "value": 1e-40 } } },
@@ -169,7 +168,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
-        "notes": "Axis-0 companion: finite subnormal columns must reduce to finite subnormal sums rather than zero."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": {
@@ -212,7 +211,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
-        "notes": "Vec4 last-axis companion: finite subnormal rows should survive the subgroup/vectorized reducer."
       },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
@@ -240,7 +239,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
-        "notes": "Odd-width last-axis companion: finite subnormal rows should survive the non-vec4 subgroup reducer."
       },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
@@ -268,7 +267,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
-        "notes": "Rank-3 axis-1 companion: the middle-axis indexing path must preserve finite subnormal column sums."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
       "inputs": {
@@ -299,7 +298,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum_default_axes_do_not_keep_dims",
-        "notes": "Rank-3 default-axes companion: reducing all finite subnormal values to a scalar should preserve the finite subnormal total."
       },
       "attrs": { "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 3, 2], "data": { "kind": "constant", "value": 1e-40 } } },
@@ -316,7 +315,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum_default_axes_keepdims",
-        "notes": "Rank-3 default-axes keepdims companion: reducing all finite subnormals should preserve the finite total in shape [1,1,1]."
       },
       "attrs": { "keepdims": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 3, 2], "data": { "kind": "constant", "value": 1e-40 } } },
@@ -384,7 +383,7 @@
     {
       "name": "rank3_two_axes_trailing_empty_axis_identity_zero",
       "provenance": {
-        "notes": "Two reduced axes where the trailing one is empty. The serial walk divides the reduced linear index by the stride of the reduced axes after each one, and that stride is zero here; the divisor has a floor at one for exactly this shape. Every existing empty-reduce case reduces a single axis, so nothing had reached the floor."
       },
       "attrs": { "axes": [1, 2], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } } },
@@ -887,7 +886,7 @@
           "shape": [2, 2, 1024],
           "data": {
             "kind": "cycle",
-            "values": [1.0, -2.0, 0.5, 3.25, -1.5, 2.0, -0.75, 4.0, -3.5, 1.25, 0.0, -2.25, 5.0, -4.0, 2.75, -1.0]
           }
         }
       },
@@ -897,7 +896,7 @@
       "name": "dispatch_cliff_axis0_cols_1048577",
       "requires": { "limits": { "maxBufferSize": 268435712, "maxStorageBufferBindingSize": 268435712 } },
       "provenance": {
-        "notes": "rows>=64 selects axis0_tilecols (prio 20); cols=1048577 > 65535*TILE_COLS(16) forces the folded column-tile dispatch (x=65535, y=2). Validates the (wg.x + wg.y*nwg.x)*TILE_COLS reconstruction in the shared reduce-axis0-tilecols.wgsl.jinja (used by all Reduce*/ArgMax/ArgMin axis0 paths). Before the 2D fold the plan-time dispatch-limit guard throws. Its largest tensor is 268435712 bytes, so the case needs an adapter whose maxBufferSize and maxStorageBufferBindingSize both reach it — declared, because the WebGPU guaranteed minimums (256 MiB / 128 MiB) do not, and a device at them must report the case inapplicable rather than fail allocating it."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": {
@@ -1000,7 +999,7 @@
     {
       "name": "int32_lastaxis_tree_parallel_3x67",
       "provenance": {
-        "notes": "67 cols >= ROW_PARALLEL_MIN_COLS selects the axis1_parallel tree path for i32; values above 2^24 fail if the accumulator smuggles through f32 (regression: f32 combine vs i32 value type mismatch)."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
       "inputs": {
@@ -1037,7 +1036,7 @@
     {
       "name": "all_axes_flat_fullreduce_101x103_nonmul4_keepdims",
       "provenance": {
-        "notes": "numel=10403 (%4=3, >=8192) full-reduce: previously fell to the single-lane reduce-serial-axis kernel; now stays on all_axes_flat (f32 binding, manual vec4 groups + 3-element scalar tail). Cross-checks the flat path against the serial all_axes_keepdims reference — a dropped tail would change the result."
       },
       "attrs": { "keepdims": 1 },
       "inputs": {
@@ -1061,13 +1060,13 @@
       "inputs": { "x": { "dtype": "int32", "shape": [128, 64], "data": { "kind": "cycle", "values": [7000, -1, 2] } } },
       "outputs": { "y": { "dtype": "int32", "shape": [], "tolerance": 0 } },
       "provenance": {
-        "notes": "The cycle [1,-1,2,-2] summed to exactly zero across all 8192 elements, so this case could not distinguish an int32 result from an f32 one - on the very axis its name claims - and referenceReduceAllFloat was in fact returning a Float32Array here. A three-value cycle makes the total 19119729, which is ODD and above 2^24 where f32 has an ulp of 2, so the exact answer is unrepresentable in f32 and the dtype leak becomes visible. It also stays well inside int32, so nothing wraps."
       }
     },
     {
       "name": "axis0_narrow_f32_8192x3_splitk_guard_lock",
       "provenance": {
-        "notes": "Compact lock below the historical 16-column split-K guard. It selects the split-K axis0 path; constant ones make the partial reduction and reassociation exact."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [8192, 3], "data": { "kind": "constant", "value": 1.0 } } },
@@ -1230,10 +1229,13 @@
         "x": {
           "dtype": "float16",
           "shape": [8192, 8],
-          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2 }
         }
       },
-      "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.002 } }
     },
     {
       "name": "f16_last_axis_vec4_8x1024",
@@ -1302,10 +1304,13 @@
         "x": {
           "dtype": "float16",
           "shape": [4096, 64],
-          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2 }
         }
       },
-      "outputs": { "y": { "dtype": "float16", "shape": [64], "tolerance": 0.05, "relTolerance": 0.002 } }
     },
     {
       "name": "int32_axis0_tiled_64x32",
@@ -1318,6 +1323,173 @@
         }
       },
       "outputs": { "y": { "dtype": "int32", "shape": [32], "tolerance": 0 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "rank3_axis1_middle_no_keepdims_input_x": [1, 2, 3, 4, -1, -2, -3, -4, 0.5, 1.5, 2.5, 3.5, 10, 20, 30, 40, -10, -20, -30, -40, 2, 4, 6, 8],
     "rank3_axis0_no_keepdims_input_x": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
     {
       "name": "all_axes_flat_rank1_boundary_8192",
       "provenance": {
+        "notes": "Exactly 8,192 rank-1 elements exercise the inclusive lower boundary of the parallel full reduction."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [8192], "data": { "kind": "constant", "value": 1.0 } } },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
+        "notes": "An axis-0 reduction over 64 finite subnormal values produces a larger finite subnormal in each column."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [64, 16], "data": { "kind": "constant", "value": 1e-40 } } },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
+        "notes": "An axis-0 reduction of finite subnormal columns must produce finite subnormal sums rather than zero."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
+        "notes": "Finite subnormal rows should survive the vectorized last-axis reduction."
       },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
+        "notes": "Finite subnormal rows should survive an odd-width last-axis reduction."
       },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum",
+        "notes": "A rank-3 axis-1 reduction must preserve finite subnormal column sums through middle-axis indexing."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum_default_axes_do_not_keep_dims",
+        "notes": "A rank-3 default-axes reduction should preserve the finite subnormal scalar total."
       },
       "attrs": { "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 3, 2], "data": { "kind": "constant", "value": 1e-40 } } },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceSum_default_axes_keepdims",
+        "notes": "A rank-3 default-axes reduction with keepdims should preserve the finite subnormal total in shape [1,1,1]."
       },
       "attrs": { "keepdims": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 3, 2], "data": { "kind": "constant", "value": 1e-40 } } },
     {
       "name": "rank3_two_axes_trailing_empty_axis_identity_zero",
       "provenance": {
+        "notes": "Two axes are reduced and the trailing reduced axis is empty. The serial index walk must use its zero-stride safeguard and return the additive identity without dividing by zero."
       },
       "attrs": { "axes": [1, 2], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } } },
           "shape": [2, 2, 1024],
           "data": {
             "kind": "cycle",
+            "values": [1.0, -2.0, 0.5, 3.25, -1.5, 2.0, -0.75, 4.0, -3.5, 1.25, 0.0, -2.25, 5.0, -4.0, 2.75, -1.0, 6.5]
           }
         }
       },
       "name": "dispatch_cliff_axis0_cols_1048577",
       "requires": { "limits": { "maxBufferSize": 268435712, "maxStorageBufferBindingSize": 268435712 } },
       "provenance": {
+        "notes": "With 1,048,577 columns, the axis-0 reduction requires 65,536 column tiles and therefore a two-row workgroup dispatch. The final tile verifies folded workgroup indexing and the over-dispatch bound. The 268,435,712-byte tensor requires the declared storage limits because it exceeds WebGPU's guaranteed minimum storage-binding size."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": {
     {
       "name": "int32_lastaxis_tree_parallel_3x67",
       "provenance": {
+        "notes": "A 67-column int32 row exercises parallel last-axis reduction. Values above 2^24 make any unintended f32 accumulation observable."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
       "inputs": {
     {
       "name": "all_axes_flat_fullreduce_101x103_nonmul4_keepdims",
       "provenance": {
+        "notes": "A 10,403-element full reduction uses float32 vec4 groups plus a three-element scalar tail. The kept singleton dimensions and independently computed result make any dropped tail element observable."
       },
       "attrs": { "keepdims": 1 },
       "inputs": {
       "inputs": { "x": { "dtype": "int32", "shape": [128, 64], "data": { "kind": "cycle", "values": [7000, -1, 2] } } },
       "outputs": { "y": { "dtype": "int32", "shape": [], "tolerance": 0 } },
       "provenance": {
+        "notes": "A three-value int32 cycle over 8,192 elements sums to the odd value 19,119,729. The result is inside int32 range but not exactly representable as f32, exposing any unintended float accumulation."
       }
     },
     {
       "name": "axis0_narrow_f32_8192x3_splitk_guard_lock",
       "provenance": {
+        "notes": "An 8,192-by-3 axis-0 reduction exercises split-K with a narrow output. Constant ones make the partial reduction and reassociation exact."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [8192, 3], "data": { "kind": "constant", "value": 1.0 } } },
         "x": {
           "dtype": "float16",
           "shape": [8192, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 0.25 }
         }
       },
+      "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.002 } },
+      "provenance": {
+        "notes": "A 8192-row float16 axis-0 sum. Values oscillate about 0.25 so each column total is proportional to the reduction length, making a dropped, duplicated or mis-strided partial observable at float16 resolution."
+      }
     },
     {
       "name": "f16_last_axis_vec4_8x1024",
         "x": {
           "dtype": "float16",
           "shape": [4096, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2, "offset": 0.25 }
         }
       },
+      "outputs": { "y": { "dtype": "float16", "shape": [64], "tolerance": 0.05, "relTolerance": 0.002 } },
+      "provenance": {
+        "notes": "A 4096-row float16 axis-0 sum. Values oscillate about 0.25 so each column total is proportional to the reduction length, making a dropped, duplicated or mis-strided partial observable at float16 resolution."
+      }
     },
     {
       "name": "int32_axis0_tiled_64x32",
         }
       },
       "outputs": { "y": { "dtype": "int32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "subgroup_rows_last_axis_f32_96x256",
+      "attrs": { "axes": [-1], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [96, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [96], "tolerance": 0.0002, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "subgroup_rows_last_axis_f32_rank3_2x40x1024",
+      "attrs": { "axes": [-1], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 40, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.25 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 40], "tolerance": 0.001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "subgroup_rows_last_axis_f16_80x1024",
+      "attrs": { "axes": [1], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [80, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [80], "tolerance": 0.05, "relTolerance": 0.002 } }
+    },
+    {
+      "name": "int32_lastaxis_subgroup_rows_64x1024",
+      "provenance": {
+        "notes": "Sixty-four rows of 1,024 int32 values exercise subgroup-per-row reduction in the output type. A five-value cycle shifts phase every row so outputs differ while matching tree accumulation exactly."
+      },
+      "attrs": { "axes": [1], "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [64, 1024], "data": { "kind": "cycle", "values": [3, 4, 0, -2, 1] } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [64], "tolerance": 0 } }
+    },
+    {
+      "name": "f32_lastaxis_rowserial_band_subgroup_rows_8192x256",
+      "provenance": {
+        "notes": "Route lock for the subgroup-per-row last-axis reduce inside the row-serial band: 8192 rows of 256 f32 values sit at the ROW_SERIAL_MIN_ROWS floor with SUBGROUP_MIN_COLS columns, so the subgroup route outranks the serial one."
+      },
+      "attrs": { "axes": [1], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8192, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 1.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [8192], "tolerance": 0.001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "rank4_axes023_noncontiguous_multi_axis_2x4x3x3",
+      "provenance": {
+        "notes": "A rank-4 reduction over axes {0,2,3} keeps the middle channel axis. The reduced elements are not a contiguous suffix, so each output uses a serial fold."
+      },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4, 3, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.07, "scale": 0.2, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "multi_axis_rank4_coop_channel_reduce_axes023",
+      "provenance": {
+        "notes": "A rank-4 reduction over axes {0,2,3} leaves one output per channel and 512 reduced elements per output, exercising cooperative tree accumulation. The tolerance allows its f32 reassociation."
+      },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 1.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [8], "tolerance": 0.0005, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "multi_axis_rank4_coop_channel_reduce_axes023_keepdims",
+      "provenance": {
+        "notes": "A rank-4 reduction over axes {0,2,3} with keepdims leaves one output per channel and 512 reduced elements per output, exercising cooperative tree accumulation. The tolerance allows its f32 reassociation."
+      },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 8, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 1.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 8, 1, 1], "tolerance": 0.0005, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "multi_axis_rank3_coop_axes02",
+      "provenance": {
+        "notes": "A rank-3 reduction over axes {0,2} leaves one output per channel and 512 reduced elements per output, exercising cooperative tree accumulation. The tolerance allows its f32 reassociation."
+      },
+      "attrs": { "axes": [0, 2], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8, 6, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 1.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [6], "tolerance": 0.0005, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "multi_axis_rank4_coop_channel_reduce_axes023_f16",
+      "provenance": {
+        "notes": "A cooperative multi-axis channel reduction widens f16 storage to f32 for accumulation."
+      },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [2, 8, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 1.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [8], "tolerance": 0.05, "relTolerance": 0.002 } }
+    },
+    {
+      "name": "multi_axis_rank4_coop_channel_reduce_axes023_i32",
+      "provenance": {
+        "notes": "A cooperative multi-axis channel reduction accumulates int32 values in the output type."
+      },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 8, 16, 16],
+          "data": { "kind": "cycle", "values": [3, -1, 4, -1, 5, -9, 2] }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [8], "tolerance": 0 } }
+    },
+    {
+      "name": "coop_and_serial_one_large_many_small",
+      "provenance": {
+        "notes": "Each channel contains one 1e8 value and 2,047 ones. Cooperative reassociation at this reduction size must remain within the declared relative tolerance of the sequential sum."
+      },
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8, 64, 16, 16],
+          "data": { "kind": "cycle", "values": [100000000.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [64], "tolerance": 1, "relTolerance": 0.00001 } }
     }
   ]
 }