sync 91d990483a17

Browse files

Files changed (16) hide show

README.md +14 -10
build/webgpu/bench.json +4 -7
build/webgpu/manifest.json +660 -784
build/webgpu/metadata.json +48 -16
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-i32-axes02.wgsl.jinja +3 -4
build/webgpu/reduce-noop-empty-axes.wgsl.jinja +11 -3
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 +226 -21
build/webgpu/reduce-serial-axis.wgsl.jinja +218 -73
build/webgpu/test.json +74 -16

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `ReduceProd` spec](https://onnx.ai/onnx/operators/onnx__ReduceProd
 ## 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 product-reduced output tensor. | required |
 ## Attributes
@@ -34,9 +34,9 @@ Default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `keepdims` | `1` | If 1, retains reduced dimensions with size 1 in the output; if 0, the reduced dimensions are removed. |
 | `noop_with_empty_axes` | `0` | If 0 (default), empty or absent axes trigger reduction over all axes; if 1, empty axes make the op a no-op identity pass-through. |
-| `axes` | `[]` | Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`. |
 ## Type constraints
@@ -50,7 +50,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## 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
@@ -61,21 +61,25 @@ Some implementation variants require `subgroups`. These are route-specific capab
 - [`reduce-flat-partial.wgsl.jinja`](build/webgpu/reduce-flat-partial.wgsl.jinja)
 - [`reduce-i32-axes02.wgsl.jinja`](build/webgpu/reduce-i32-axes02.wgsl.jinja)
 - [`reduce-noop-empty-axes.wgsl.jinja`](build/webgpu/reduce-noop-empty-axes.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 product-reduced output tensor. | 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, retains reduced dimensions with size 1 in the output; if 0, the reduced dimensions are removed. |
 | `noop_with_empty_axes` | `0` | If 0 (default), empty or absent axes trigger reduction over all axes; if 1, empty axes make the op a no-op identity pass-through. |
 ## Type constraints
 ## 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-flat-partial.wgsl.jinja`](build/webgpu/reduce-flat-partial.wgsl.jinja)
 - [`reduce-i32-axes02.wgsl.jinja`](build/webgpu/reduce-i32-axes02.wgsl.jinja)
 - [`reduce-noop-empty-axes.wgsl.jinja`](build/webgpu/reduce-noop-empty-axes.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.ReduceProd",
   "cases": [
     {
       "name": "reduce-prod-axis0-1024x512",
@@ -96,9 +95,7 @@
       "preset": "smoke",
       "vars": { "dtype": "int32", "rows": 256, "cols": 4096 },
       "attrs": { "axes": [1], "keepdims": 0 },
-      "inputs": {
-        "x": { "shape": [256, 4096], "dtype": "int32", "dist": "uniform", "seed": 423, "scale": 1, "offset": 1 }
-      },
       "outputs": { "y": { "shape": [256], "dtype": "int32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
     },
@@ -156,7 +153,7 @@
       "outputs": { "y": { "shape": [4096], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] },
       "provenance": {
-        "notes": "One row-count notch outside the row-serial region, where the row-per-workgroup tree wins by 1.7x. Pairs with the 8192x1024 case to bracket ROW_SERIAL_MIN_ROWS."
       }
     },
     {
@@ -174,7 +171,7 @@
       "name": "reduceprod-rank3-spatial-axes12-f32-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
-        "source": "authored for branch coverage",
         "notes": "Rank-3 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane. Inputs stay close to one so the long product remains finite."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
@@ -195,7 +192,7 @@
       "name": "reduceprod-rank4-spatial-axes23-f32-2x64x256x256-pathology",
       "preset": "stress",
       "provenance": {
-        "source": "authored for branch coverage",
         "notes": "Rank-4 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane. Inputs stay close to one so the long product remains finite."
       },
       "attrs": { "axes": [2, 3], "keepdims": 1 },

 {
   "cases": [
     {
       "name": "reduce-prod-axis0-1024x512",
       "preset": "smoke",
       "vars": { "dtype": "int32", "rows": 256, "cols": 4096 },
       "attrs": { "axes": [1], "keepdims": 0 },
+      "inputs": { "x": { "shape": [256, 4096], "dtype": "int32", "dist": "uniform", "seed": 423, "min": 0, "max": 2 } },
       "outputs": { "y": { "shape": [256], "dtype": "int32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] }
     },
       "outputs": { "y": { "shape": [4096], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype)" }] },
       "provenance": {
+        "notes": "A [4096, 1024] tensor lies below ROW_SERIAL_MIN_ROWS, so each output row uses a workgroup-tree reduction. Together with `reduceprod-rowserial-boundary-8192x1024`, this brackets the row-count selector transition."
       }
     },
     {
       "name": "reduceprod-rank3-spatial-axes12-f32-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "repository-authored",
         "notes": "Rank-3 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane. Inputs stay close to one so the long product remains finite."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
       "name": "reduceprod-rank4-spatial-axes23-f32-2x64x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "repository-authored",
         "notes": "Rank-4 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane. Inputs stay close to one so the long product remains finite."
       },
       "attrs": { "axes": [2, 3], "keepdims": 1 },

build/webgpu/manifest.json CHANGED Viewed

@@ -2,478 +2,242 @@
   "domain": "ai.onnx",
   "name": "ReduceProd",
   "sinceVersion": 18,
-  "description": "Computes the product of elements along specified axes of the input tensor. The output rank matches the input when `keepdims` is 1; reduced dimensions are pruned when `keepdims` is 0. Reduction over an empty set of values yields 1.",
-  "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 product-reduced output tensor."
     }
-  ],
-  "attributes": { "keepdims": 1, "noop_with_empty_axes": 0, "axes": [] },
-  "attributeDescriptions": {
-    "keepdims": "If 1, retains reduced dimensions with size 1 in the output; if 0, the reduced dimensions are removed.",
-    "noop_with_empty_axes": "If 0 (default), empty or absent axes trigger reduction over all axes; if 1, empty axes make the op a no-op identity pass-through.",
-    "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,
-    "AXES02_WORKGROUP_SIZE": 64,
-    "CONTIGUOUS_SUFFIX_MIN_COLS": 256,
-    "ROW_SERIAL_MIN_ROWS": 8192,
-    "ROW_SERIAL_MAX_COLS": 1024
   },
-  "bindingSets": {
-    "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)" }] }
-      }
-    ],
-    "axes02": [
-      { "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": "dim0", "type": "u32", "value": "dim(shapes.data, 0)" },
-            { "name": "dim1", "type": "u32", "value": "dim(shapes.data, 1)" },
-            { "name": "dim2", "type": "u32", "value": "dim(shapes.data, 2)" }
-          ]
-        }
-      }
-    ],
-    "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)" }
-          ]
-        }
-      }
-    ],
-    "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)" }
-          ]
-        }
-      }
-    ],
-    "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)" }] }
-      }
-    ],
-    "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": "ReduceProd.ContiguousSuffixSubgroupVec4",
-          "source": {
-            "shader": "reduce-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.ContiguousSuffixTreeVec4",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "vec4": true,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "suffixVec4",
-          "dispatch": { "workgroups": "numel(shapes.reduced)" }
         }
       ]
     },
@@ -481,142 +245,204 @@
       "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": "ReduceProd.ContiguousSuffixTree",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": { "op": "\"prod\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "suffixScalar",
-          "dispatch": { "workgroups": "numel(shapes.reduced)" }
         }
       ]
     },
     {
       "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": "ReduceProd.NoopEmptyAxes",
-          "source": { "shader": "reduce-noop-empty-axes.wgsl.jinja" },
-          "bindings": "elementwise",
-          "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "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": "ReduceProd.MultiAxisRank3",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.MultiAxisRank4",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "int32_rank3_axes02_keepdims",
       "priority": 35,
-      "when": ["dtypes.T == \"i32\"", "ranks.data == 3", "attrs.keepdims == 1", "hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 2, 3)", "hasAxis(attrs.axes, 1, 3) == false", "dim(shapes.data, 0) > 0", "dim(shapes.data, 2) > 0", "ranks.reduced == 3", "dim(shapes.reduced, 0) == 1", "dim(shapes.reduced, 1) == dim(shapes.data, 1)", "dim(shapes.reduced, 2) == 1"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Int32Rank3Axes02Keepdims",
-          "bindings": "axes02",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "axes02WorkgroupSize" },
-          "source": {
-            "shader": "reduce-i32-axes02.wgsl.jinja",
-            "inputs": { "op": "\"prod\"", "workgroupSize": "axes02WorkgroupSize" }
           }
         }
-      ],
-      "derive": {
-        "axes02WorkgroupSize": "min(tunables.AXES02_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)"
-      }
     },
     {
       "id": "tree_last_axis",
       "priority": 20,
       "demoteWhen": ["rowSerialPreferred"],
-      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"i32\")", "ranks.data >= 2", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "treeWorkgroupOk"],
-      "constants": {
         "scalar": "dtypes.T",
-        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.data, ranks.data - 1)))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.TreeLastAxis",
-          "source": { "shader": "reduce-row-tree.wgsl.jinja", "inputs": { "op": "\"prod\"" } },
-          "bindings": "lastAxisScalar",
-          "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
         }
       ]
     },
     {
       "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)))"
@@ -625,116 +451,116 @@
         {
           "id": "main",
           "name": "ReduceProd.TreeRowVec4",
-          "source": {
-            "shader": "reduce-row-tree.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.Rank0Scalar",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.Rank1Axis0",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "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": "ReduceProd.AxisSplitReduce",
-          "source": {
-            "shader": "reduce-axis-split-reduce.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "axisSplitReduce",
-          "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AxisSplitCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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",
@@ -746,175 +572,177 @@
         {
           "id": "split_reduce",
           "name": "ReduceProd.AxisSplitTiledReduce",
-          "source": {
-            "shader": "reduce-axis0-tilecols.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.AxisSplitCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.Axis0SplitKReduce",
-          "source": {
-            "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.Axis0SplitKCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount",
-              "outputF16": "dtypes.T == \"f16\""
-            }
-          },
-          "bindings": "axis0SplitCombine",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize" }
         }
       ]
     },
     {
       "id": "axis0_splitk_i32",
       "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": "ReduceProd.Axis0SplitKReduceI32",
-          "source": {
-            "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "intMode": true,
-              "workgroupSize": "reduceWorkgroupSize",
-              "split": "splitCount"
-            }
           },
-          "bindings": "axis0SplitReduce",
-          "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceProd.Axis0SplitKCombineI32",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd.Axis0TileCols",
-          "source": {
-            "shader": "reduce-axis0-tilecols.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "intMode": "dtypes.T == \"i32\"",
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "bindings": "axis0Parallel",
-          "dispatch": { "workgroups": "ceilDiv((dim(shapes.data, 1)), (constants.tileCols))" }
         }
       ]
     },
     {
       "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": "ReduceProd.AllAxesFlatPartial",
-          "source": {
-            "shader": "reduce-flat-partial.wgsl.jinja",
-            "inputs": { "op": "\"prod\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "flatPartial",
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AllAxesFlatCombine",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": { "op": "\"prod\"", "outputF16": "dtypes.T == \"f16\"" }
-          },
-          "bindings": "flatCombine",
           "dispatch": { "x": 1 }
         }
       ]
@@ -922,25 +750,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": "ReduceProd.AllAxesFlatPartialI32",
-          "source": { "shader": "reduce-flat-partial.wgsl.jinja", "inputs": { "op": "\"prod\"", "intMode": true } },
-          "bindings": "flatPartial",
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AllAxesFlatCombineI32",
-          "source": {
-            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-            "inputs": { "op": "\"prod\"", "intMode": true, "workgroupSize": 1, "split": "flatSplitCount" }
-          },
-          "bindings": "flatCombine",
           "dispatch": { "x": 1 }
         }
       ]
@@ -948,41 +775,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", "tree_last_axis"],
-      "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": "ReduceProd.RankNSingleAxisGeneric",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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",
       "priority": 25,
-      "requires": { "features": ["subgroups"] },
       "demoteWhen": ["rowSerialPreferred"],
-      "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)", "dim(shapes.data, ranks.data - 1) >= 4", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH == 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)"],
-      "constants": {
         "scalar": "dtypes.T",
         "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH))))"
@@ -991,28 +820,30 @@
         {
           "id": "main",
           "name": "ReduceProd.SubgroupRowVec4",
-          "source": {
-            "shader": "reduce-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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",
       "priority": 24,
-      "requires": { "features": ["subgroups"] },
       "demoteWhen": ["rowSerialPreferred"],
-      "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)", "dim(shapes.data, ranks.data - 1) > 0", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH != 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)"],
-      "constants": {
         "scalar": "dtypes.T",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))"
       },
@@ -1020,141 +851,186 @@
         {
           "id": "main",
           "name": "ReduceProd.SubgroupRow",
-          "source": {
-            "shader": "reduce-row-subgroup.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "vec4": false,
-              "castF32": "dtypes.T == \"f16\"",
-              "usesF16": "dtypes.T == \"f16\""
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "lastAxisScalarSubgroup",
-          "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" }
         }
       ]
     },
     {
       "id": "axis0",
       "priority": 0,
       "supersededBy": ["axis_split_tiled_narrow", "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": "\"prod\"",
-              "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": "\"prod\"",
-              "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": "rank2_all_axes_keepdims",
       "priority": 30,
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2 and attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == 1"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Rank2AllAxesKeepdims",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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 and attrs.keepdims == 1 and ranks.reduced == ranks.data and numel(shapes.reduced) == 1"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Rank3AllAxesKeepdims",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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_no_keepdims",
       "priority": 30,
-      "constants": { "axis": 0, "scalar": "dtypes.T" },
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3 and attrs.keepdims == 0 and ranks.reduced == 0"],
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Rank3AllAxesNoKeepdims",
-          "source": {
-            "shader": "reduce-serial-axis.wgsl.jinja",
-            "inputs": {
-              "op": "\"prod\"",
-              "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": "ReduceProd",
   "sinceVersion": 18,
+  "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 },
+    "AXES02_WORKGROUP_SIZE": { "default": 64 },
+    "CONTIGUOUS_SUFFIX_MIN_COLS": { "default": 256 },
+    "ROW_SERIAL_MIN_ROWS": { "default": 8192 },
+    "ROW_SERIAL_MAX_COLS": { "default": 1024 }
   },
   "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": "count", "type": "u32", "value": "numel(shapes.y)" }]
+    },
+    "params_4": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "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": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1) / tunables.VECTOR_WIDTH" }
+      ]
+    },
+    "params_8": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "1" },
+        { "name": "cols", "type": "u32", "value": "1" },
+        { "name": "outCount", "type": "u32", "value": "1" }
+      ]
+    },
+    "params_9": {
+      "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)" }
+      ]
+    },
+    "partials": { "buffer": "storage", "elementType": "$partialElement" },
+    "params_10": {
+      "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_11": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }]
+    },
+    "params_12": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
+      ]
+    },
+    "params_13": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }]
+    },
+    "params_14": {
+      "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_15": { "name": "params", "buffer": "uniform", "struct": [{ "name": "cols", "type": "u32", "value": "1" }] },
+    "params_16": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
+        { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
+      ]
+    },
+    "params_17": {
+      "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_18": {
+      "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_19": {
+      "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": "ReduceProd.ContiguousSuffixSubgroupVec4",
+          "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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": "ReduceProd.ContiguousSuffixTreeVec4",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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": "ReduceProd.ContiguousSuffixTree",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "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": "noop_empty_axes",
       "priority": 40,
+      "when": ["dtypes.T == \"f32\"", "attrs.noop_with_empty_axes == 1", "(attrs.axes | length) == 0", "sameShape(shapes.x, shapes.y)"],
+      "derive": { "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.NoopEmptyAxes",
+          "shader": "reduce-noop-empty-axes.wgsl.jinja",
+          "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_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": "ReduceProd.MultiAxisRank3",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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_4"],
+          "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": "ReduceProd.MultiAxisRank4",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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_4"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
+      ]
     },
     {
       "id": "int32_rank3_axes02_keepdims",
       "priority": 35,
+      "when": ["dtypes.T == \"i32\"", "ranks.x == 3", "attrs.keepdims == 1", "hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 2, 3)", "hasAxis(attrs.axes, 1, 3) == false", "dim(shapes.x, 0) > 0", "dim(shapes.x, 2) > 0", "ranks.y == 3", "dim(shapes.y, 0) == 1", "dim(shapes.y, 1) == dim(shapes.x, 1)", "dim(shapes.y, 2) == 1"],
+      "derive": {
+        "axes02WorkgroupSize": "min(tunables.AXES02_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)"
+      },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Int32Rank3Axes02Keepdims",
+          "shader": "reduce-i32-axes02.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "workgroupSizeSpec": "axes02WorkgroupSize" },
+          "bindings": [
+            "x_2",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "dim0", "type": "u32", "value": "dim(shapes.x, 0)" },
+                { "name": "dim1", "type": "u32", "value": "dim(shapes.x, 1)" },
+                { "name": "dim2", "type": "u32", "value": "dim(shapes.x, 2)" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (axes02WorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (axes02WorkgroupSize)), 65535)",
+            "z": 1
           }
         }
+      ]
     },
     {
       "id": "tree_last_axis",
       "priority": 20,
+      "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"i32\")", "ranks.x >= 2", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "treeWorkgroupOk"],
       "demoteWhen": ["rowSerialPreferred"],
+      "derive": {
         "scalar": "dtypes.T",
+        "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.x, ranks.x - 1)))"
       },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.TreeLastAxis",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": { "op": "\"prod\"" },
+          "bindings": [
+            "x_2",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
+                { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
+            "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 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": "ReduceProd.SubgroupRowsVec4",
+          "shader": "reduce-row-subgroup-rows.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": ["x", "y", "params_7"],
+          "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": "ReduceProd.TreeRowVec4",
+          "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "vec4": true,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x", "y", "params_7"],
+          "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": "ReduceProd.Rank0Scalar",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_2", "y", "params_8"],
           "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": "ReduceProd.Rank1Axis0",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_2", "y", "params_9"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 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": "ReduceProd.AxisSplitReduce",
+          "shader": "reduce-axis-split-reduce.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "splitSpec": "splitCount",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "partials", "params_10"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AxisSplitCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_11"],
+          "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": "ReduceProd.AxisSplitTiledReduce",
+          "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "splitSpec": "splitCount",
+            "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "partials", "params_10"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)",
+            "y": "splitCount",
+            "z": "ceilDiv(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)"
+          }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AxisSplitCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_11"],
+          "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": "ReduceProd.Axis0SplitKReduce",
+          "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "splitSpec": "splitCount",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "partials", "params_12"],
+          "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": "ReduceProd.Axis0SplitKCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_13"],
+          "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": "ReduceProd.Axis0SplitKReduceI32",
+          "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "intMode": true,
+            "workgroupSizeSpec": "reduceWorkgroupSize",
+            "splitSpec": "splitCount"
           },
+          "bindings": ["x_2", "partials", "params_12"],
+          "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": "ReduceProd.Axis0SplitKCombineI32",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "intMode": true,
+            "workgroupSizeSpec": "reduceWorkgroupSize",
+            "splitSpec": "splitCount"
           },
+          "bindings": ["partials_2", "y", "params_13"],
+          "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": "ReduceProd.Axis0TileCols",
+          "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_12"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
+            "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": "ReduceProd.AllAxesFlatPartial",
+          "shader": "reduce-flat-partial.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
+          "bindings": ["x_2", "partials", "params_14"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AllAxesFlatCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_2", "y", "params_15"],
           "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": "ReduceProd.AllAxesFlatPartialI32",
+          "shader": "reduce-flat-partial.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "intMode": true },
+          "bindings": ["x_2", "partials", "params_14"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceProd.AllAxesFlatCombineI32",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"prod\"", "intMode": true, "splitSpec": "flatSplitCount" },
+          "bindings": ["partials_2", "y", "params_15"],
           "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", "tree_last_axis"],
+      "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.RankNSingleAxisGeneric",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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_16"],
+          "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)", "dim(shapes.x, ranks.x - 1) >= 4", "dim(shapes.x, ranks.x - 1) % tunables.VECTOR_WIDTH == 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)"],
       "demoteWhen": ["rowSerialPreferred"],
+      "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": "ReduceProd.SubgroupRowVec4",
+          "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "vec4": true,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
+          },
+          "bindings": ["x", "y", "params_7"],
+          "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)", "dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) % tunables.VECTOR_WIDTH != 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)"],
       "demoteWhen": ["rowSerialPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "derive": {
         "scalar": "dtypes.T",
         "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))"
       },
         {
           "id": "main",
           "name": "ReduceProd.SubgroupRow",
+          "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "vec4": false,
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
+          },
+          "bindings": ["x_2", "y", "params_17"],
+          "dispatch": {
+            "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
+            "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
+            "z": 1
           },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "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", "axis0_tilecols"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "axis0",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "axis": 0,
+            "op": "\"prod\"",
+            "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": "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": "\"prod\"",
+            "indexing": "\"axis2d\"",
+            "intMode": "dtypes.T == \"i32\"",
+            "castF32": "dtypes.T == \"f16\"",
+            "usesF16Spec": "dtypes.T == \"f16\""
           },
+          "bindings": ["x_2", "y", "params_19"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "rank2_all_axes_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 2 and attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == 1"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Rank2AllAxesKeepdims",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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 and attrs.keepdims == 1 and ranks.y == ranks.x and numel(shapes.y) == 1"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Rank3AllAxesKeepdims",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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_no_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3 and attrs.keepdims == 0 and ranks.y == 0"],
+      "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceProd.Rank3AllAxesNoKeepdims",
+          "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": {
+            "op": "\"prod\"",
+            "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,27 +1,59 @@
 {
   "name": "ai.onnx.ReduceProd",
-  "id": "_ai_onnx_reduceprod_webgpu_af6578f",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "VznixYylucv/AFi4jPmrJwhC0TMoDKr+7rMdIVHsW+s=",
-      "manifest.json": "RVeo8qczKuuOGfvWE9Qf77fp9IgC4/Vx/tCYmwo86+E=",
-      "reduce-axis-split-reduce.wgsl.jinja": "sKtJglLU0XOIB5OZ3ayiQ3hyBozHELXbS2bRxPdDBtQ=",
-      "reduce-axis0-splitk-combine.wgsl.jinja": "CP4QmClXgzXKKmy/ynbQdaWsMEkI8ETCiaHU0P1mW1s=",
-      "reduce-axis0-splitk-reduce.wgsl.jinja": "qArZcLmfHk+cK8BAtnBAhfUlN85EljIXgqfGotutV0g=",
-      "reduce-axis0-tilecols.wgsl.jinja": "If5gLW0KSm7P9fRhWiwiQUbPt9jAj+lDnVLQu98vTe8=",
-      "reduce-flat-partial.wgsl.jinja": "iFF+WqB7whD91yu/pT/ESQegWBUlHq5r5b6DSHGEVHY=",
-      "reduce-i32-axes02.wgsl.jinja": "lyV7HWjK6vjRO2ekdha9khkcBkeSUa7xPPOA4YyByHI=",
-      "reduce-noop-empty-axes.wgsl.jinja": "IO2MEzyBGJr3Z8AVP7ELJXqfhwcG/vtBfa0wCdkbuwo=",
-      "reduce-row-subgroup.wgsl.jinja": "tCCSXf8DaUkcXj2vUD4bwfhvaYxXMN3bEE8q8L90ZrE=",
-      "reduce-row-tree.wgsl.jinja": "cLIGeHfVYpn8YH3HHQg84DuMNe/3brWgwLUi2zCFmOw=",
-      "reduce-serial-axis.wgsl.jinja": "1GJeXytXvJYFn3JhcEk8sbpze8q8+nFVHQrf0lbSbXI=",
-      "test.json": "w6OEmrn5Bdf1M23l36SbchwXIui/PnEzFjHVeHJUH1I="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.ReduceProd" }
 }

 {
   "name": "ai.onnx.ReduceProd",
+  "id": "_ai_onnx_reduceprod_webgpu_eba8bbf",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "agHVWGoEFusniUoQri0U27LrIk898f6FnHV0DXI6Ld8=",
+      "manifest.json": "MDsk7/aIS4exv9qL/BRnNp40y3k4RQB0UQzWnzSpwQw=",
+      "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-i32-axes02.wgsl.jinja": "t6kVFEoZLe2G9Cj58xMyYVyRe267yfpbVA0vA78id7I=",
+      "reduce-noop-empty-axes.wgsl.jinja": "5bKufrcfYt3BagZ6WxbUvqwKKKVMU7d+SUshgUFWTDI=",
+      "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": "AAcj4UFyWhoSVnZFFqVG7DSWAadBNDMcvZQc3t3XZDk="
     }
   },
+  "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"],
+      "noop_empty_axes": ["reduce-noop-empty-axes.wgsl.jinja"],
+      "multi_axis_rank3": ["reduce-serial-axis.wgsl.jinja"],
+      "multi_axis_rank4": ["reduce-serial-axis.wgsl.jinja"],
+      "int32_rank3_axes02_keepdims": ["reduce-i32-axes02.wgsl.jinja"],
+      "tree_last_axis": ["reduce-row-tree.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"],
+      "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"],
+      "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"],
+      "axis0": ["reduce-serial-axis.wgsl.jinja"],
+      "axis1": ["reduce-serial-axis.wgsl.jinja"],
+      "rank2_all_axes_keepdims": ["reduce-serial-axis.wgsl.jinja"],
+      "all_axes_keepdims": ["reduce-serial-axis.wgsl.jinja"],
+      "all_axes_no_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 = 1.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 = 1i;
 {% else %}
     var total = 1.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 = 1i;
 {% else %}
   var acc = 1.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 }}(1){% else %}1.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 = 1i;
 {% else %}
   var acc = 1.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-i32-axes02.wgsl.jinja CHANGED Viewed

@@ -1,12 +1,11 @@
 // Rank-3 i32 reduction over axes {0, 2}, leaving the middle dimension.
 {{ env.wgsl.resourceDeclarations }}
-const WG: u32 = {{ source.workgroupSize }}u;
 @compute @workgroup_size(WG)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  let outputIndex = gid.x + gid.y * nwg.x * WG;
   let d0 = params.dim0;
   let d1 = params.dim1;
   let d2 = params.dim2;

 // Rank-3 i32 reduction over axes {0, 2}, leaving the middle dimension.
 {{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSizeSpec }}u;
 @compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let outputIndex = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let d0 = params.dim0;
   let d1 = params.dim1;
   let d2 = params.dim2;

build/webgpu/reduce-noop-empty-axes.wgsl.jinja CHANGED Viewed

@@ -1,13 +1,21 @@
 {{ env.wgsl.resourceDeclarations }}
 @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.count) {
     return;
   }
   let v = x[i];
   y[i] = v;
 }

 {{ env.wgsl.resourceDeclarations }}
 @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.count) {
     return;
   }
   let v = x[i];
+{% if op == "abs" %}
+  y[i] = abs(v);
+{% elif op == "square" %}
+  y[i] = v * v;
+{% elif op == "log" %}
+  y[i] = log(v);
+{% else %}
   y[i] = v;
+{% endif %}
 }

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", "subgroupMul", "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 = 1.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,74 +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" %}
-{% if source.op == "prod" and accType == "i32" %}
 fn identity() -> i32 { return 1i; }
-{% else %}
 fn identity() -> {{ accType }} { return {{ accType }}(1); }
 {% endif %}
 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;
@@ -83,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.ReduceProd",
   "fixtureArrays": {
     "rank3_axis1_middle_keepdims_input_x": [1, -2, 0.5, 4, -3, 2, -1, 0.25, 2, -0.5, 3, -2, -1, 1.5, 2, -4, 0.5, -2, -3, 1, -2, 4, 0.25, -0.5]
   },
@@ -7,7 +6,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 } } },
@@ -104,7 +103,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceProd",
-        "notes": "Axis-0 companion: normal inputs produce finite subnormal products, which must not be flushed to zero."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": {
@@ -596,7 +595,7 @@
         "x": {
           "dtype": "float32",
           "shape": [2, 2, 1024],
-          "data": { "kind": "cycle", "values": [1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, 1.0] }
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.00001 } }
@@ -608,7 +607,7 @@
         "x": {
           "dtype": "int32",
           "shape": [2, 2, 1024],
-          "data": { "kind": "cycle", "values": [1, -1, 1, -1, 1, -1, 1, 1] }
         }
       },
       "outputs": { "y": { "dtype": "int32", "shape": [2, 2], "tolerance": 0.00001 } }
@@ -702,7 +701,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceProd",
-        "notes": "Float edge-value coverage on the working vec4 last-axis path (cols=4, %4==0 -> subgroup_last_axis_vec4). Row0 contains 0*inf -> NaN; row1 overflows the product to +inf. Metal handles inf/NaN in arithmetic (only subnormals are FTZ-flushed), so this should PASS on GPU. Fills the no-inf/NaN-input coverage gap noted in coverage."
       },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
@@ -736,7 +735,11 @@
       "name": "tree_lastaxis_dispatch_fold_f32_70000x64",
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
-        "x": { "dtype": "float32", "shape": [70000, 64], "data": { "kind": "cycle", "values": [1.0, -1.0, 0.5, 2.0] } }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [70000], "tolerance": 0 } }
     },
@@ -744,11 +747,7 @@
       "name": "tree_lastaxis_i32_dispatch_fold_70000x64",
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
-        "x": {
-          "dtype": "int32",
-          "shape": [70000, 64],
-          "data": { "kind": "cycle", "values": [1, -1, 1, -1, 1, -1, 1, 1] }
-        }
       },
       "outputs": { "y": { "dtype": "int32", "shape": [70000], "tolerance": 0 } }
     },
@@ -765,7 +764,11 @@
       "attrs": { "axes": [-1], "keepdims": 0 },
       "requires": { "features": ["subgroups"] },
       "inputs": {
-        "x": { "dtype": "float32", "shape": [16384, 64], "data": { "kind": "cycle", "values": [1.0, -1.0, 0.5, 2.0] } }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [16384], "tolerance": 0 } }
     },
@@ -778,7 +781,7 @@
     {
       "name": "axis0_narrow_f32_8192x3_splitk_guard_lock",
       "provenance": {
-        "notes": "Compact lock below the historical 16-column split-K guard. Multiplicative identity input validates partial and combine identities exactly."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [8192, 3], "data": { "kind": "constant", "value": 1.0 } } },
@@ -923,7 +926,10 @@
           "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "offset": 1.0, "scale": 0.001 }
         }
       },
-      "outputs": { "y": { "dtype": "float16", "shape": [], "tolerance": 0.05, "relTolerance": 0.01 } }
     },
     {
       "name": "f16_suffix_vec4_4x8x128",
@@ -972,11 +978,63 @@
     {
       "name": "axis0_splitk_i32_8192x5",
       "provenance": {
-        "notes": "Integer ReduceProd on the axis-0 split-K route, which ReduceProd did not have until now. A three-value cycle over five columns is coprime with the row stride, so the count of negative factors differs per column and so do the signs of the results; keeping every factor at +/-1 makes the product exact."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "int32", "shape": [8192, 5], "data": { "kind": "cycle", "values": [1, 1, -1] } } },
       "outputs": { "y": { "dtype": "int32", "shape": [5], "tolerance": 0 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "rank3_axis1_middle_keepdims_input_x": [1, -2, 0.5, 4, -3, 2, -1, 0.25, 2, -0.5, 3, -2, -1, 1.5, 2, -4, 0.5, -2, -3, 1, -2, 4, 0.25, -0.5]
   },
     {
       "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.ReduceProd",
+        "notes": "An axis-0 reduction of normal inputs produces finite subnormal products that must not be flushed to zero."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": {
         "x": {
           "dtype": "float32",
           "shape": [2, 2, 1024],
+          "data": { "kind": "cycle", "values": [-1.0, -1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] }
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.00001 } }
         "x": {
           "dtype": "int32",
           "shape": [2, 2, 1024],
+          "data": { "kind": "cycle", "values": [-1, -1, 1, 1, 1, 1, 1, 1, 1] }
         }
       },
       "outputs": { "y": { "dtype": "int32", "shape": [2, 2], "tolerance": 0.00001 } }
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
         "test": "ReductionOpTest.ReduceProd",
+        "notes": "Four columns exercise the vectorized last-axis reduction. Row 0 contains 0 multiplied by infinity and must produce NaN; row 1 overflows to positive infinity."
       },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
       "name": "tree_lastaxis_dispatch_fold_f32_70000x64",
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [70000, 64],
+          "data": { "kind": "cycle", "values": [1.0, -1.0, 0.5, 2.0, -1.0] }
+        }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [70000], "tolerance": 0 } }
     },
       "name": "tree_lastaxis_i32_dispatch_fold_70000x64",
       "attrs": { "axes": [-1], "keepdims": 0 },
       "inputs": {
+        "x": { "dtype": "int32", "shape": [70000, 64], "data": { "kind": "cycle", "values": [-1, -1, 1, 1, 1] } }
       },
       "outputs": { "y": { "dtype": "int32", "shape": [70000], "tolerance": 0 } }
     },
       "attrs": { "axes": [-1], "keepdims": 0 },
       "requires": { "features": ["subgroups"] },
       "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [16384, 64],
+          "data": { "kind": "cycle", "values": [1.0, -1.0, 0.5, 2.0, -1.0] }
+        }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [16384], "tolerance": 0 } }
     },
     {
       "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. Multiplicative identity input validates partial and combine identities exactly."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "float32", "shape": [8192, 3], "data": { "kind": "constant", "value": 1.0 } } },
           "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "offset": 1.0, "scale": 0.001 }
         }
       },
+      "outputs": { "y": { "dtype": "float16", "shape": [], "tolerance": 0.01, "relTolerance": 0.01 } },
+      "provenance": {
+        "notes": "A 65,543-element float16 product over factors near 1. The tolerance is scaled to the 0.25 result so a dropped or duplicated factor is observable."
+      }
     },
     {
       "name": "f16_suffix_vec4_4x8x128",
     {
       "name": "axis0_splitk_i32_8192x5",
       "provenance": {
+        "notes": "An integer axis-0 split-K reduction uses a three-value cycle over five columns. The coprime cycle gives each column a different count of negative factors, while values of plus or minus one keep every product exact."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
       "inputs": { "x": { "dtype": "int32", "shape": [8192, 5], "data": { "kind": "cycle", "values": [1, 1, -1] } } },
       "outputs": { "y": { "dtype": "int32", "shape": [5], "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, "offset": 1.0, "scale": 0.05 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [96], "tolerance": 0.00001, "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, "offset": 1.0, "scale": 0.001 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [80], "tolerance": 0.05, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "int32_lastaxis_subgroup_rows_64x1024",
+      "provenance": {
+        "notes": "Route lock for the integer subgroup-per-row last-axis product: 64 rows of 1024 int32 values fold in the output type; a seven-value cycle of +-1 shifts phase every row, so rows alternate sign."
+      },
+      "attrs": { "axes": [1], "keepdims": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [64, 1024], "data": { "kind": "cycle", "values": [1, -1, 1, 1, -1, 1, -1] } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [64], "tolerance": 0 } }
+    },
+    {
+      "name": "multi_axis_alternating_extremes_serial_fold_only",
+      "attrs": { "axes": [0, 2, 3], "keepdims": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "cycle", "values": [1e+30, 1e-30, 1e+30, 1e-30, 1.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4], "tolerance": 0.001, "relTolerance": 0.001, "allowNaN": true }
+      },
+      "provenance": {
+        "notes": "Alternating 1e30 and 1e-30 values remain finite under the required left-to-right fold. A pairwise tree could instead form positive infinity and zero in separate partials, whose product is NaN, so this shape requires serial fold order."
+      }
     }
   ]
 }