sync 6fdf6301e2bb

Browse files

Files changed (15) hide show

README.md +11 -4
build/webgpu/bench.json +2 -2
build/webgpu/manifest.json +108 -232
build/webgpu/metadata.json +16 -16
build/webgpu/reduce-axis-split-reduce.wgsl.jinja +3 -15
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +4 -31
build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja +1 -14
build/webgpu/reduce-axis0-tilecols.wgsl.jinja +8 -20
build/webgpu/reduce-flat-partial.wgsl.jinja +20 -39
build/webgpu/reduce-noop-empty-axes.wgsl.jinja +9 -6
build/webgpu/reduce-row-subgroup-rows.wgsl.jinja +31 -37
build/webgpu/reduce-row-subgroup.wgsl.jinja +51 -58
build/webgpu/reduce-row-tree.wgsl.jinja +6 -17
build/webgpu/reduce-serial-axis.wgsl.jinja +14 -59
build/webgpu/test.json +4 -4

README.md CHANGED Viewed

@@ -44,6 +44,13 @@ Default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16`, `int32`, `uint32`, `int8`, `uint8`, `bool` |
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
@@ -53,7 +60,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 - [`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
 - [`datamove-flat-copy-x4-tail.wgsl.jinja`](build/webgpu/datamove-flat-copy-x4-tail.wgsl.jinja)
 - [`reduce-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-axis-split-reduce.wgsl.jinja)
 - [`reduce-axis0-splitk-combine.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja)
@@ -70,7 +77,7 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## 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.
@@ -89,7 +96,7 @@ import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/ai.onnx.ReduceMin", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
-const { y } = await kernel({ x: { data: xData, shape: [] } }, {
-  outputs: { y: { shape: [], dtype: "float32" } },
 });
 ```

 | --- | --- |
 | `T` | `float32`, `float16`, `int32`, `uint32`, `int8`, `uint8`, `bool` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `subgroup_last_axis_vec4` — Reduces each contiguous last-axis row with subgroup collectives and vec4-packed reads.
+- `subgroup_last_axis` — Reduces each contiguous last-axis row with subgroup collectives and scalar reads for an unaligned row width.
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 - [`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 cases
 - [`datamove-flat-copy-x4-tail.wgsl.jinja`](build/webgpu/datamove-flat-copy-x4-tail.wgsl.jinja)
 - [`reduce-axis-split-reduce.wgsl.jinja`](build/webgpu/reduce-axis-split-reduce.wgsl.jinja)
 - [`reduce-axis0-splitk-combine.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 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.
 const kernel = await getKernel("webgpu-kernels/ai.onnx.ReduceMin", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
+const { y } = await kernel({ x: { data: xData, shape: [3, 2, 2] } }, {
+  outputs: { y: { shape: [1, 1, 1], dtype: "float32" } },
 });
 ```

build/webgpu/bench.json CHANGED Viewed

@@ -73,7 +73,7 @@
       "name": "reducemin-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."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
@@ -85,7 +85,7 @@
       "name": "reducemin-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."
       },
       "attrs": { "axes": [2, 3], "keepdims": 1 },

       "name": "reducemin-rank3-spatial-axes12-f32-128x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic",
         "notes": "Rank-3 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [1, 2], "keepdims": 1 },
       "name": "reducemin-rank4-spatial-axes23-f32-2x64x256x256-pathology",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic",
         "notes": "Rank-4 multi-axis reduction with 128 output lanes, each serially scanning a 256x256 plane."
       },
       "attrs": { "axes": [2, 3], "keepdims": 1 },

build/webgpu/manifest.json CHANGED Viewed

@@ -33,6 +33,7 @@
     "ROW_SERIAL_MAX_COLS": { "default": 1024 }
   },
   "derive": {
     "emptyNarrowInputOk": "numel(shapes.x) > 0 or (tensorDtypes.x != \"int8\" and tensorDtypes.x != \"uint8\")",
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
@@ -59,62 +60,75 @@
     "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": {
-    "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] },
-    "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "y_2": { "name": "y", "buffer": "storage", "elementType": "$T" },
-    "params_2": {
       "name": "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_3": {
       "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_4": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
     },
-    "params_5": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
-        { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
-    "params_6": {
       "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_7": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "1" },
         { "name": "cols", "type": "u32", "value": "1" },
         { "name": "outCount", "type": "u32", "value": "1" }
       ]
     },
-    "params_8": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "1" },
@@ -122,75 +136,56 @@
       ]
     },
     "partials": { "buffer": "storage", "elementType": "$partialElement" },
-    "params_9": {
       "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_10": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }]
-    },
-    "params_11": {
-      "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
       ]
     },
-    "params_12": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }]
-    },
-    "partials_3": { "name": "partials", "buffer": "storage", "elementType": "f32" },
-    "params_13": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "count4", "type": "u32", "value": "floor(numel(shapes.x) / tunables.VECTOR_WIDTH)" },
         { "name": "numel", "type": "u32", "value": "numel(shapes.x)" }
       ]
     },
-    "partials_4": { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
-    "params_14": { "name": "params", "buffer": "uniform", "struct": [{ "name": "cols", "type": "u32", "value": "1" }] },
-    "params_15": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_16": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
-        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
-    "params_17": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_18": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     }
@@ -207,7 +202,7 @@
           "name": "ReduceMin.NarrowEmptyIdentity",
           "shader": "reduce-narrow-empty-identity.wgsl.jinja",
           "derive": { "identity": "1 if tensorDtypes.x == \"bool\" else 127 if tensorDtypes.x == \"int8\" else 255" },
-          "bindings": [{ "arg": "y", "elementType": "$scalar" }, "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -231,15 +226,9 @@
           "id": "main",
           "name": "ReduceMin.ContiguousSuffixSubgroupVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y_2", "params_2"],
-          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -257,13 +246,8 @@
           "id": "main",
           "name": "ReduceMin.ContiguousSuffixTreeVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y_2", "params_2"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
@@ -281,8 +265,7 @@
           "id": "main",
           "name": "ReduceMin.ContiguousSuffixTree",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": { "op": "\"min\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "y_2", "params_3"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
@@ -321,7 +304,7 @@
           "id": "main",
           "name": "ReduceMin.NoopEmptyAxes",
           "shader": "reduce-noop-empty-axes.wgsl.jinja",
-          "bindings": ["x_2", "y_2", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -342,10 +325,7 @@
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "reduceWorkgroupSize": "reduceWorkgroupSize",
-            "op": "\"min\"",
             "indexing": "\"multiaxis\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "rank": 3,
             "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
             "dataShape": "shapes.x",
@@ -354,7 +334,7 @@
             "keepDims": "attrs.keepdims != 0",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
-          "bindings": ["x_2", "y_2", "params_4"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -375,10 +355,7 @@
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "reduceWorkgroupSize": "reduceWorkgroupSize",
-            "op": "\"min\"",
             "indexing": "\"multiaxis\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "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",
@@ -387,7 +364,7 @@
             "keepDims": "attrs.keepdims != 0",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
-          "bindings": ["x_2", "y_2", "params_4"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -410,8 +387,7 @@
           "id": "main",
           "name": "ReduceMin.TreeLastAxis",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": { "op": "\"min\"" },
-          "bindings": ["x_2", "y_2", "params_5"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -436,14 +412,12 @@
           "id": "main",
           "name": "ReduceMin.SubgroupRowsVec4",
           "shader": "reduce-row-subgroup-rows.wgsl.jinja",
-          "derive": { "op": "\"min\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x", "y_2", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
             "y": "ceilDiv(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -462,13 +436,8 @@
           "id": "main",
           "name": "ReduceMin.TreeRowVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y_2", "params_6"],
           "dispatch": { "x": "min(lastAxisRows, 65535)", "y": "ceilDiv(lastAxisRows, 65535)", "z": 1 }
         }
       ]
@@ -483,14 +452,8 @@
           "id": "main",
           "name": "ReduceMin.Rank0Scalar",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "indexing": "\"axis2d\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
-            "logicalBool": "tensorDtypes.x == \"bool\""
-          },
-          "bindings": ["x_2", "y_2", "params_7"],
           "dispatch": { "x": 1 }
         }
       ]
@@ -504,14 +467,8 @@
           "id": "main",
           "name": "ReduceMin.Rank1Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "indexing": "\"axis2d\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
-            "logicalBool": "tensorDtypes.x == \"bool\""
-          },
-          "bindings": ["x_2", "y_2", "params_8"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -536,13 +493,8 @@
           "id": "split_reduce",
           "name": "ReduceMin.AxisSplitReduce",
           "shader": "reduce-axis-split-reduce.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "splitSpec": "splitCount",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_9"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -553,8 +505,8 @@
           "id": "combine",
           "name": "ReduceMin.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"min\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y_2", "params_10"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
@@ -581,14 +533,8 @@
           "id": "split_reduce",
           "name": "ReduceMin.AxisSplitTiledReduce",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "splitSpec": "splitCount",
-            "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_9"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -599,8 +545,8 @@
           "id": "combine",
           "name": "ReduceMin.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"min\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y_2", "params_10"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
@@ -613,6 +559,7 @@
       "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\"",
@@ -625,13 +572,8 @@
           "id": "split_reduce",
           "name": "ReduceMin.Axis0SplitKReduce",
           "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "splitSpec": "splitCount",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -642,8 +584,8 @@
           "id": "combine",
           "name": "ReduceMin.Axis0SplitKCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"min\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y_2", "params_12"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
@@ -662,8 +604,8 @@
           "id": "main",
           "name": "ReduceMin.Axis0TileCols",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
-          "derive": { "op": "\"min\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "y_2", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
@@ -676,23 +618,22 @@
       "id": "all_axes_flat",
       "priority": 31,
       "when": ["flatParallelCovered"],
-      "derive": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceMin.AllAxesFlatPartial",
           "shader": "reduce-flat-partial.wgsl.jinja",
-          "derive": { "op": "\"min\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "partials_3", "params_13"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceMin.AllAxesFlatCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"min\"", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_4", "y_2", "params_14"],
           "dispatch": { "x": 1 }
         }
       ]
@@ -710,10 +651,7 @@
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "reduceWorkgroupSize": "reduceWorkgroupSize",
-            "op": "\"min\"",
             "indexing": "\"rankn\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "rank": "ranks.x",
             "axisSpec": "reduceAxis",
             "dataShape": "shapes.x",
@@ -722,7 +660,7 @@
             "keepDims": "attrs.keepdims != 0",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
-          "bindings": ["x_2", "y_2", "params_15"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -746,19 +684,13 @@
           "id": "main",
           "name": "ReduceMin.SubgroupRowVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y_2", "params_6"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -776,19 +708,13 @@
           "id": "main",
           "name": "ReduceMin.SubgroupRow",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "vec4": false,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y_2", "params_16"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -801,18 +727,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "axis": 0,
-            "scalar": "dtypes.T",
-            "op": "\"min\"",
-            "indexing": "\"axis2d\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
-            "logicalBool": "tensorDtypes.x == \"bool\""
-          },
-          "bindings": ["x_2", "y_2", "params_17"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -829,18 +747,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "axis1",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "axis": 1,
-            "scalar": "dtypes.T",
-            "op": "\"min\"",
-            "indexing": "\"axis2d\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
-            "logicalBool": "tensorDtypes.x == \"bool\""
-          },
-          "bindings": ["x_2", "y_2", "params_18"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -852,32 +762,15 @@
     {
       "id": "all_axes_keepdims",
       "priority": 30,
-      "when": ["not flatParallelCovered", "emptyNarrowInputOk", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 1 and ranks.y == ranks.x and numel(shapes.y) == 1"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceMin.Rank3AllAxesKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "indexing": "\"axis2d\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
-            "logicalBool": "tensorDtypes.x == \"bool\""
-          },
-          "bindings": [
-            "x_2",
-            "y_2",
-            {
-              "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)",
@@ -889,32 +782,15 @@
     {
       "id": "all_axes_no_keepdims",
       "priority": 30,
-      "when": ["not flatParallelCovered", "emptyNarrowInputOk", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 0 and ranks.y == 0"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceMin.Rank3AllAxesNoKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"min\"",
-            "indexing": "\"axis2d\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
-            "logicalBool": "tensorDtypes.x == \"bool\""
-          },
-          "bindings": [
-            "x_2",
-            "y_2",
-            {
-              "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)",

     "ROW_SERIAL_MAX_COLS": { "default": 1024 }
   },
   "derive": {
+    "reduceOp": "\"min\"",
     "emptyNarrowInputOk": "numel(shapes.x) > 0 or (tensorDtypes.x != \"int8\" and tensorDtypes.x != \"uint8\")",
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "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)))",
+    "op": "reduceOp",
+    "castF32": "dtypes.T == \"f16\""
   },
   "bindings": {
+    "params_suffix_vec4": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
         { "name": "chunkCount", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y) / tunables.VECTOR_WIDTH" }
       ]
     },
+    "params_reduce_row_tree": {
       "name": "params",
       "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_axis_split_combine": {
       "name": "params",
+      "struct": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }]
     },
+    "params_axis0_combine": {
+      "name": "params",
+      "struct": [{ "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }]
+    },
+    "partials_flat": { "name": "partials", "buffer": "storage", "elementType": "f32" },
+    "partials_flat_read": { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+    "params_all_axes_flat": { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "1" }] },
+    "params_rows_chunk_count": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
+        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
+    "x": { "elementType": "$vectorScalar" },
+    "y": { "elementType": "$T" },
+    "x_t": { "name": "x", "elementType": "$T" },
+    "params_main": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
+        { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" }
+      ]
+    },
+    "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] },
+    "params_out_count": {
+      "name": "params",
+      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
+    },
+    "params_rows_cols": {
       "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)" }
       ]
     },
+    "params_rank0_scalar": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "1" },
         { "name": "cols", "type": "u32", "value": "1" },
         { "name": "outCount", "type": "u32", "value": "1" }
       ]
     },
+    "params_rank1_axis0": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "1" },
       ]
     },
     "partials": { "buffer": "storage", "elementType": "$partialElement" },
+    "params_axis_split": {
       "name": "params",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "inner", "type": "u32", "value": "axisSplitInner" },
         { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" }
       ]
     },
+    "partials_combine": { "name": "partials", "buffer": "read-only-storage", "elementType": "$partialElement" },
+    "params_split_reduce": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
       ]
     },
+    "params_count4_numel": {
       "name": "params",
       "struct": [
         { "name": "count4", "type": "u32", "value": "floor(numel(shapes.x) / tunables.VECTOR_WIDTH)" },
         { "name": "numel", "type": "u32", "value": "numel(shapes.x)" }
       ]
     },
+    "params_axis_dim_out_count": {
       "name": "params",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
+    "params_axis0": {
       "name": "params",
       "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_axis1": {
       "name": "params",
       "struct": [
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
+    "params_all_axes": {
       "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)" }
       ]
     }
           "name": "ReduceMin.NarrowEmptyIdentity",
           "shader": "reduce-narrow-empty-identity.wgsl.jinja",
           "derive": { "identity": "1 if tensorDtypes.x == \"bool\" else 127 if tensorDtypes.x == \"int8\" else 255" },
+          "bindings": [{ "arg": "y", "elementType": "$scalar" }, "params_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "id": "main",
           "name": "ReduceMin.ContiguousSuffixSubgroupVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": { "vec4": true },
+          "bindings": ["x", "y", "params_suffix_vec4"],
+          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
     },
           "id": "main",
           "name": "ReduceMin.ContiguousSuffixTreeVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": { "vec4": true },
+          "bindings": ["x", "y", "params_suffix_vec4"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
           "id": "main",
           "name": "ReduceMin.ContiguousSuffixTree",
           "shader": "reduce-row-tree.wgsl.jinja",
+          "bindings": ["x_t", "y", "params_main"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
           "id": "main",
           "name": "ReduceMin.NoopEmptyAxes",
           "shader": "reduce-noop-empty-axes.wgsl.jinja",
+          "bindings": ["x_t", "y", "params_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "reduceWorkgroupSize": "reduceWorkgroupSize",
             "indexing": "\"multiaxis\"",
             "rank": 3,
             "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
             "dataShape": "shapes.x",
             "keepDims": "attrs.keepdims != 0",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_t", "y", "params_out_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "reduceWorkgroupSize": "reduceWorkgroupSize",
             "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",
             "keepDims": "attrs.keepdims != 0",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_t", "y", "params_out_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "id": "main",
           "name": "ReduceMin.TreeLastAxis",
           "shader": "reduce-row-tree.wgsl.jinja",
+          "bindings": ["x_t", "y", "params_rows_cols"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
           "id": "main",
           "name": "ReduceMin.SubgroupRowsVec4",
           "shader": "reduce-row-subgroup-rows.wgsl.jinja",
+          "bindings": ["x", "y", "params_reduce_row_tree"],
           "dispatch": {
             "x": "min(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
             "y": "ceilDiv(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
             "z": 1
+          }
         }
       ]
     },
           "id": "main",
           "name": "ReduceMin.TreeRowVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
+          "derive": { "vec4": true },
+          "bindings": ["x", "y", "params_reduce_row_tree"],
           "dispatch": { "x": "min(lastAxisRows, 65535)", "y": "ceilDiv(lastAxisRows, 65535)", "z": 1 }
         }
       ]
           "id": "main",
           "name": "ReduceMin.Rank0Scalar",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "logicalBool": "tensorDtypes.x == \"bool\"" },
+          "bindings": ["x_t", "y", "params_rank0_scalar"],
           "dispatch": { "x": 1 }
         }
       ]
           "id": "main",
           "name": "ReduceMin.Rank1Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "logicalBool": "tensorDtypes.x == \"bool\"" },
+          "bindings": ["x_t", "y", "params_rank1_axis0"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "id": "split_reduce",
           "name": "ReduceMin.AxisSplitReduce",
           "shader": "reduce-axis-split-reduce.wgsl.jinja",
+          "derive": { "splitSpec": "splitCount" },
+          "bindings": ["x_t", "partials", "params_axis_split"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
           "id": "combine",
           "name": "ReduceMin.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_combine", "y", "params_axis_split_combine"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
           "id": "split_reduce",
           "name": "ReduceMin.AxisSplitTiledReduce",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": { "splitSpec": "splitCount" },
+          "bindings": ["x_t", "partials", "params_axis_split"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
           "id": "combine",
           "name": "ReduceMin.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_combine", "y", "params_axis_split_combine"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
       "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"],
+      "demoteWhen": ["dtypes.T == \"f32\" and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32 and axis0Rows >= 2 * tunables.AXIS0_SPLIT_MIN_ROWS and axis0Cols >= tunables.AXIS0_TILE_MIN_COLS and axis0TilePathFits and ceilDiv(axis0Cols, tunables.AXIS0_TILE_COLS) >= axis0SplitCount"],
       "derive": {
         "splitCount": "axis0SplitCount",
         "partialElement": "\"f32\"",
           "id": "split_reduce",
           "name": "ReduceMin.Axis0SplitKReduce",
           "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
+          "derive": { "splitSpec": "splitCount" },
+          "bindings": ["x_t", "partials", "params_split_reduce"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
           "id": "combine",
           "name": "ReduceMin.Axis0SplitKCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_combine", "y", "params_axis0_combine"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
           "id": "main",
           "name": "ReduceMin.Axis0TileCols",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": {},
+          "bindings": ["x_t", "y", "params_split_reduce"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
       "id": "all_axes_flat",
       "priority": 31,
       "when": ["flatParallelCovered"],
+      "derive": { "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceMin.AllAxesFlatPartial",
           "shader": "reduce-flat-partial.wgsl.jinja",
+          "bindings": ["x_t", "partials_flat", "params_count4_numel"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceMin.AllAxesFlatCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_flat_read", "y", "params_all_axes_flat"],
           "dispatch": { "x": 1 }
         }
       ]
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "reduceWorkgroupSize": "reduceWorkgroupSize",
             "indexing": "\"rankn\"",
             "rank": "ranks.x",
             "axisSpec": "reduceAxis",
             "dataShape": "shapes.x",
             "keepDims": "attrs.keepdims != 0",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_t", "y", "params_axis_dim_out_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "id": "main",
           "name": "ReduceMin.SubgroupRowVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": { "vec4": true },
+          "bindings": ["x", "y", "params_reduce_row_tree"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
+          }
         }
       ]
     },
           "id": "main",
           "name": "ReduceMin.SubgroupRow",
           "shader": "reduce-row-subgroup.wgsl.jinja",
+          "derive": { "vec4": false },
+          "bindings": ["x_t", "y", "params_rows_chunk_count"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
+          }
         }
       ]
     },
       "passes": [
         {
           "id": "main",
+          "name": "ReduceMin.Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "logicalBool": "tensorDtypes.x == \"bool\"" },
+          "bindings": ["x_t", "y", "params_axis0"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
       "passes": [
         {
           "id": "main",
+          "name": "ReduceMin.Axis1",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "logicalBool": "tensorDtypes.x == \"bool\"" },
+          "bindings": ["x_t", "y", "params_axis1"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
     {
       "id": "all_axes_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "emptyNarrowInputOk", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 1", "ranks.y == ranks.x", "numel(shapes.y) == 1"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceMin.Rank3AllAxesKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "logicalBool": "tensorDtypes.x == \"bool\"" },
+          "bindings": ["x_t", "y", "params_all_axes"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
     {
       "id": "all_axes_no_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "emptyNarrowInputOk", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 0", "ranks.y == 0"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "ReduceMin.Rank3AllAxesNoKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "logicalBool": "tensorDtypes.x == \"bool\"" },
+          "bindings": ["x_t", "y", "params_all_axes"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,32 +1,32 @@
 {
   "name": "ai.onnx.ReduceMin",
-  "id": "_ai_onnx_reducemin_webgpu_56651f5",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "9HSUFWy2zOka0uaIlARjO0hKvsiIkYJlVTOuwy7LlrU=",
       "datamove-flat-copy-x4-tail.wgsl.jinja": "1u3sJCuYJ8lM8QU7/oOPB6Q6A7vkcOxCR03RGiakHWo=",
-      "manifest.json": "Cwk7CFWuFgbjlIYAuOaGJieBrpOhc28p44FXo8+v7F4=",
-      "reduce-axis-split-reduce.wgsl.jinja": "4+ep9xH4pHZOfaZ8abJhA8SW5M4mUhDC+y1CtDz6vjY=",
-      "reduce-axis0-splitk-combine.wgsl.jinja": "Yz1hjK55R/kndUrw3ugPqgPaOBwKmYupqZoVdJTHO5Q=",
-      "reduce-axis0-splitk-reduce.wgsl.jinja": "iLqQIC+Sa+EUzZG4Z3BqZOT2XA1TQEuRcFR2JbkEa30=",
-      "reduce-axis0-tilecols.wgsl.jinja": "PjYkEUQJBeG70td3W2xxmexH9x6XNIfeSzXBY47HbaU=",
-      "reduce-flat-partial.wgsl.jinja": "+qToL+wFi9QOxvY887aBAEwZK6Xu/eLQkucKvJYNwSk=",
       "reduce-narrow-empty-identity.wgsl.jinja": "XmLNK4zp0ru14DEm8TL5CAL74RVu7aVPrFo3WSJX3wo=",
-      "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": "rOIoQuLOtT5w5RcBzucPfj30IkoeNSd+2TuyiE4cbk8="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "narrow_empty_identity": ["reduce-narrow-empty-identity.wgsl.jinja"],
       "contiguous_suffix_subgroup_vec4": ["reduce-row-subgroup.wgsl.jinja"],

 {
   "name": "ai.onnx.ReduceMin",
+  "id": "_ai_onnx_reducemin_webgpu_6673c73",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "t2isnzkULzGR9Ki2/j7Gg6PyxV3TDPc+ovEKMuVvYWI=",
       "datamove-flat-copy-x4-tail.wgsl.jinja": "1u3sJCuYJ8lM8QU7/oOPB6Q6A7vkcOxCR03RGiakHWo=",
+      "manifest.json": "vZJEw1TU0yAeZ9SzFGoIujiecANsX96IVN1/QXbvqTQ=",
+      "reduce-axis-split-reduce.wgsl.jinja": "xo4jwtx/MeHIZkh2KU31ZCR+Ak58xOe4Cg3HWrcvdqo=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "Zf7tz8nrapi2KMZIbv8cT2f4pliJoJHj4iwS4hc/6ms=",
+      "reduce-axis0-splitk-reduce.wgsl.jinja": "RVy0BuBzgen99cwoM0KNiZ6OPnc8ColMSsME3IZKLI4=",
+      "reduce-axis0-tilecols.wgsl.jinja": "bxt1WvxNRg3bgU4R+s5JUlHZjqH0THZZic++EJ7X+yA=",
+      "reduce-flat-partial.wgsl.jinja": "S573jG+5ZoOcPEzKnD8h3nvPNYOuM8+ku15d+15ypBA=",
       "reduce-narrow-empty-identity.wgsl.jinja": "XmLNK4zp0ru14DEm8TL5CAL74RVu7aVPrFo3WSJX3wo=",
+      "reduce-noop-empty-axes.wgsl.jinja": "MMZEXHeACtGwvDaH0nAyvMal+tAIjYrh2yxDuLfh1Lw=",
+      "reduce-row-subgroup-rows.wgsl.jinja": "5cMGE1dmrcawwHdoxQD8kPSurwA1wNHyoAQfxAOE7vU=",
+      "reduce-row-subgroup.wgsl.jinja": "7bRBWIQ7oHOq/E54hGrmXIzLznzVH9oTfYaf1jmBzV8=",
+      "reduce-row-tree.wgsl.jinja": "a+iEV/cBkrlsVz6NS0GY/TAsUSX5rO4N14wY1dKfDi0=",
+      "reduce-serial-axis.wgsl.jinja": "ZZ1kKKhlLr3cMnRDRBhWm0ZDMluJwcoqmpRmL9ukuYM=",
+      "test.json": "9zz9YIFSmjUPFmYqqDm3yRVIjl4dJQpSTgGyxTIeN7I="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "narrow_empty_identity": ["reduce-narrow-empty-identity.wgsl.jinja"],
       "contiguous_suffix_subgroup_vec4": ["reduce-row-subgroup.wgsl.jinja"],

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

@@ -8,29 +8,16 @@
 // 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;
@@ -42,6 +29,7 @@ fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% elif op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op) }}
 {% endif %}

 // Each output segment writes three partial planes: its maximum, the sum of
 // exp(x - maximum), and a packed NaN marker.
 {% endif %}
+{% set xa = ("vec" ~ columnVectorWidth ~ "<f32>(" if columnVectorWidth is defined else "f32(") if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
 {{ 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. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% elif op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op) }}
 {% endif %}

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

@@ -2,37 +2,19 @@
 // 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;
@@ -48,7 +30,9 @@ fn is_nan_f32(value: f32) -> bool {
 @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" %}
@@ -74,13 +58,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
     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" %}
@@ -89,7 +66,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
     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];
@@ -101,9 +77,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
       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" %}

 // 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 yv = "f16(" if outputF16 else "" %}
 {% set vy = ")" if outputF16 else "" %}
 {{ 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. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 {% if op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
+  // The start already folds gid.y in, so the stride must span every y row too;
+  // an x-only stride would send y = 0 lanes over columns the y >= 1 rows own.
+  let stride = nwg.x * nwg.y * 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" %}
     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 op == "max" %}
     var total = reduction_identity();
 {% elif op == "min" %}
     var total = 1.0;
 {% else %}
     var total = 0.0;
 {% endif %}
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];
       total = total + p;
 {% endif %}
     }
 {% if op == "l2" %}
     y[col] = {{ yv }}sqrt(total){{ vy }};
 {% elif op == "logsum" %}

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

@@ -2,29 +2,16 @@
 // 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;

 // 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 xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
 {{ 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. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;

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

@@ -13,17 +13,13 @@
 {% 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
@@ -38,15 +34,12 @@ fn {{ name }}() -> {{ scalar }} {
   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;
@@ -69,7 +62,6 @@ 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>) {
@@ -77,8 +69,9 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   let col_lane = tid % TILE_COLS;
   let row_lane = tid / TILE_COLS;
 {% if splitMode %}
-  // Narrow outputs: wg.x covers every column tile, wg.y is the axis segment.
-  let col = wg.x * TILE_COLS + col_lane;
   let outputIndex = col;
   let in_bounds = col < params.outputs;
   let seg = wg.y;
@@ -203,9 +196,9 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% 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) {
@@ -240,15 +233,10 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% 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 %}

 {% set rowEnd = "params.rows" %}
 {% set elem = "x[inputBase + row * params.cols + col]" %}
 {% endif %}
+{% set intMode = false %}
 {% 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 "" %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
   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;
   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 col_lane = tid % TILE_COLS;
   let row_lane = tid / TILE_COLS;
 {% if splitMode %}
+  // Narrow outputs: wg.x (with wg.z carrying its fold overflow) covers every
+  // column tile, wg.y is the axis segment.
+  let col = (wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u) * TILE_COLS + col_lane;
   let outputIndex = col;
   let in_bounds = col < params.outputs;
   let seg = wg.y;
 {% if op == "max" or op == "min" %}
   var acc = lane_identity();
 {% elif op == "prod" %}
+  var acc = 1.0;
 {% else %}
+  var acc = 0.0;
 {% endif %}
   if (in_bounds) {
     for (var row = {{ rowBegin }}; row < {{ rowEnd }}; row = row + ROW_LANES) {
 {% if splitMode %}
     partials[seg * params.outputs + outputIndex] = total;
 {% else %}
+{% if op == "l2" %}
     y[outputIndex] = {{ yv }}sqrt(total){{ vy }};
 {% elif op == "logsum" %}
     y[outputIndex] = {{ yv }}log(total){{ vy }};
 {% elif op == "mean" %}
     y[outputIndex] = {{ yv }}total / f32(params.rows){{ vy }};
 {% else %}

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

@@ -3,33 +3,20 @@
 // in registers, reduce within each workgroup, and write one partial each. A
 // following combine pass folds the partials and applies the operation finalizer.
 //
-// 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>;
@@ -44,13 +31,7 @@ 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 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();
@@ -100,23 +81,23 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   }
   red[tid] = acc;
   workgroupBarrier();
-  var stride: u32 = WG / 2u;
   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();
-  }
   if (tid == 0u) {
     partials[wg.x] = red[0];
   }

 // in registers, reduce within each workgroup, and write one partial each. A
 // following combine pass folds the partials and applies the operation finalizer.
 //
+// Scalar bindings support arbitrary element counts and a zero-to-three element
+// tail. An aligned packed binding loads each group directly, preserving the
+// component accumulation order and widening half storage before arithmetic.
+{% set intMode = false %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
 {{ 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. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 var<workgroup> red: array<{{ "i32" if intMode else "f32" }}, WG>;
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
   let gstride = nwg.x * WG;
+{% if op == "max" %}
   var acc = reduction_identity();
 {% elif op == "min" %}
   var acc = reduction_identity();
   }
   red[tid] = acc;
   workgroupBarrier();
+{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" or op == "min" %}
+{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
+  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
+    if ({{ svar }} == 0u) { break; }
+    if ({{ idx }} < {{ svar }}) {
+{% for a in arrays %}
+      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
+{% endfor %}
     }
+    {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
+{{ wgsl_tree_fold(["red"], op=op, idx="tid", wg="WG", typed=true, form="head", breakInline=true) }}
   if (tid == 0u) {
     partials[wg.x] = red[0];
   }

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

@@ -1,13 +1,16 @@
 {{ 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);

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ reduceWorkgroupSize }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(reduceWorkgroupSize) }}
   let v = x[i];
 {% if op == "abs" %}
   y[i] = abs(v);

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

@@ -15,18 +15,12 @@
 // 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. */
@@ -40,15 +34,15 @@ fn {{ name }}() -> {{ scalar }} {
   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;
@@ -57,7 +51,7 @@ 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>,
@@ -110,58 +104,58 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     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 %}
 }

 // The finalizer takes the logarithm of the sum.
 {% endif %}
 // f16 storage widens before accumulation and narrows only at the final store.
 {% 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 "" %}
 {% 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. */
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
+}{% endmacro %}
+enable subgroups;
+{{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+{% if op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op, accScalar) }}
+{% elif op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
   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>,
     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

@@ -15,17 +15,11 @@
 // 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. */
@@ -39,15 +33,15 @@ fn {{ name }}() -> {{ scalar }} {
   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;
@@ -56,7 +50,7 @@ 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>;
@@ -77,20 +71,19 @@ fn {{ name }}(value: {{ scalar }}, sgLid: u32, sgId: u32, numSg: u32) -> {{ scal
   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>,
@@ -103,13 +96,13 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   }
   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;
@@ -117,7 +110,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     } else {
       localMax = max(localMax, v.{{ comp }});
     }
-{%- endfor %}
 {% else %}
     if (is_nan_f32(v)) {
       localNan = 1.0;
@@ -125,7 +118,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     } else {
       localMax = max(localMax, v);
     }
-{%- endif %}
   }
   let rowMax = reduce_row_max(localMax, sgLid, sgId, numSg);
   let nanCount = reduce_row_add(localNan, sgLid, sgId, numSg);
@@ -135,83 +128,83 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   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 %}
 }

 // IEEE-754 bit patterns because WGSL rejects infinite constants.
 {% endif %}
 // f16 storage widens before accumulation and narrows only at the final store.
 {% 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 "" %}
 {% 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. */
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
+}{% endmacro %}
+enable subgroups;
+{{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+{% if op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op, scalar) }}
+{% elif op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
+{% endif %}
 var<workgroup> wgPartial: array<{{ scalar }}, WG>;
   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>,
   }
   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;
     } else {
       localMax = max(localMax, v.{{ comp }});
     }
+{% endfor %}
 {% else %}
     if (is_nan_f32(v)) {
       localNan = 1.0;
     } else {
       localMax = max(localMax, v);
     }
+{% endif %}
   }
   let rowMax = reduce_row_max(localMax, sgLid, sgId, numSg);
   let nanCount = reduce_row_add(localNan, sgLid, sgId, numSg);
   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

@@ -16,15 +16,11 @@
 // 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
@@ -39,15 +35,12 @@ fn {{ name }}() -> {{ scalar }} {
   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;
@@ -57,8 +50,8 @@ 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
@@ -84,20 +77,15 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
   // 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" %}
@@ -216,13 +204,14 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
 {% 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) {

 // 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 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 "" %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
   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;
 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, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
   // 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];
   workgroupBarrier();
   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" %}
 {% 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) {

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

@@ -1,5 +1,4 @@
 {% 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.
@@ -44,18 +43,20 @@ fn input_offset(out_index: u32{% if hasReduced %}, reduce_linear: u32{% endif %}
 {% 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
@@ -70,15 +71,12 @@ fn {{ name }}() -> {{ scalar }} {
   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" %}
@@ -86,7 +84,6 @@ fn negative_infinity() -> f32 {
 {% 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 {
@@ -94,7 +91,6 @@ fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
-{% endif %}
 {% if indexing == "rankn" %}
 fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
@@ -131,7 +127,8 @@ fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
 {% 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) %}
@@ -166,31 +163,8 @@ fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
 @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;
@@ -220,7 +194,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     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 }};
@@ -234,11 +207,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {% 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" %}
@@ -247,13 +216,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     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 }};
@@ -267,19 +230,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {% 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 %}

 {% 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.
 {% 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 intMode = false %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
   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" %}
 {% elif op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
 fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
 {% if indexing == "rankn" %}
 fn input_offset(out_index: u32, reduce_index: u32) -> u32 {
 {% 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) %}
 @compute @workgroup_size({{ reduceWorkgroupSize }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(reduceWorkgroupSize, bound="params.outCount") }}
 {% if op == "logsumexp" %}
   if ({{ count }} == 0u) {
     y[i] = {{ yv }}negative_infinity(){{ vy }};
     return;
     acc = acc + exp({{ at }} - m);
   }
   y[i] = {{ yv }}m + log(acc){{ vy }};
 {% elif op == "max" or op == "min" %}
   if ({{ count }} == 0u) {
     y[i] = {{ yv }}empty_value(){{ vy }};
 {% else %}
 {% if op == "mean" %}
   if ({{ count }} == 0u) {
     y[i] = {{ yv }}0.0{{ vy }};
     return;
   }
 {% elif op == "logsum" %}
     return;
   }
 {% endif %}
   var acc = {{ "1.0" if op == "prod" else "0.0" }};
   for (var r = 0u; r < {{ count }}; r = r + 1u) {
 {% if op == "prod" %}
     acc = acc * {{ at }};
 {% endif %}
   }
 {% if op == "l2" %}
   y[i] = {{ yv }}sqrt(acc){{ vy }};
 {% elif op == "logsum" %}
   y[i] = {{ yv }}log(acc){{ vy }};
 {% elif op == "mean" %}
   y[i] = {{ yv }}acc / f32({{ count }}){{ vy }};
 {% else %}
   y[i] = {{ yv }}acc{{ vy }};
 {% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -1107,7 +1107,7 @@
     {
       "name": "ort_int8_axis1_extrema",
       "provenance": {
-        "source": "ONNX Runtime CPUExecutionProvider",
         "notes": "Covers the signed narrow-integer reduction route."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
@@ -1121,7 +1121,7 @@
     {
       "name": "ort_uint8_axis0_extrema",
       "provenance": {
-        "source": "ONNX Runtime CPUExecutionProvider",
         "notes": "Covers the unsigned narrow-integer reduction route."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
@@ -1135,7 +1135,7 @@
     {
       "name": "ort_int8_empty_axis1_highest_identity",
       "provenance": {
-        "source": "ONNX Runtime CPUExecutionProvider",
         "notes": "An empty INT8 minimum is the logical dtype maximum (127), not the physical i32 storage maximum."
       },
       "attrs": { "axes": [1], "keepdims": 1 },
@@ -1152,7 +1152,7 @@
     {
       "name": "ort_uint8_empty_axis1_highest_identity",
       "provenance": {
-        "source": "ONNX Runtime CPUExecutionProvider",
         "notes": "An empty UINT8 minimum is the logical dtype maximum (255), not the physical u32 storage maximum."
       },
       "attrs": { "axes": [1], "keepdims": 1 },

     {
       "name": "ort_int8_axis1_extrema",
       "provenance": {
+        "source": "ONNX Runtime's CPU provider",
         "notes": "Covers the signed narrow-integer reduction route."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
     {
       "name": "ort_uint8_axis0_extrema",
       "provenance": {
+        "source": "ONNX Runtime's CPU provider",
         "notes": "Covers the unsigned narrow-integer reduction route."
       },
       "attrs": { "axes": [0], "keepdims": 0 },
     {
       "name": "ort_int8_empty_axis1_highest_identity",
       "provenance": {
+        "source": "ONNX Runtime's CPU provider",
         "notes": "An empty INT8 minimum is the logical dtype maximum (127), not the physical i32 storage maximum."
       },
       "attrs": { "axes": [1], "keepdims": 1 },
     {
       "name": "ort_uint8_empty_axis1_highest_identity",
       "provenance": {
+        "source": "ONNX Runtime's CPU provider",
         "notes": "An empty UINT8 minimum is the logical dtype maximum (255), not the physical u32 storage maximum."
       },
       "attrs": { "axes": [1], "keepdims": 1 },