sync 6fdf6301e2bb

Browse files

Files changed (16) hide show

README.md +11 -4
build/webgpu/bench.json +27 -1
build/webgpu/manifest.json +98 -220
build/webgpu/metadata.json +17 -17
build/webgpu/reduce-axis-split-reduce.wgsl.jinja +3 -15
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +4 -22
build/webgpu/reduce-axis0-splitk-reduce.wgsl.jinja +1 -14
build/webgpu/reduce-axis0-tilecols.wgsl.jinja +4 -11
build/webgpu/reduce-flat-combine-logsumexp.wgsl.jinja +14 -10
build/webgpu/reduce-flat-partial-logsumexp.wgsl.jinja +11 -11
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 +13 -19
build/webgpu/test.json +26 -0

README.md CHANGED Viewed

@@ -44,6 +44,13 @@ Default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16`, `int32` |
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
@@ -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
 - [`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)
 - [`reduce-axis0-splitk-reduce.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-reduce.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.ReduceLogSumExp", { 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` |
+## 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
 - [`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)
 - [`reduce-axis0-splitk-reduce.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-reduce.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.ReduceLogSumExp", { 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

@@ -75,6 +75,32 @@
         ]
       }
     },
     {
       "name": "fullreduce-single-lane-r3-101x101x101-numel-not-mult4",
       "preset": "stress",
@@ -104,7 +130,7 @@
       "name": "reducelogsumexp-spatial-axes12-f32-16x256x256-low-lane",
       "preset": "stress",
       "provenance": {
-        "source": "repository-authored",
         "notes": "Low-lane rank-3 contiguous-suffix reduction that verifies the vec4 subgroup reducer and no-subgroup workgroup-tree fallback."
       },
       "vars": { "batch": 16, "height": 256, "width": 256 },

         ]
       }
     },
+    {
+      "name": "axis0-crossover-32768x512",
+      "preset": "smoke",
+      "vars": { "rows": 32768, "cols": 512 },
+      "attrs": { "axes": [0], "keepdims": 0 },
+      "inputs": { "x": { "shape": [32768, 512], "dtype": "float32", "dist": "normal", "scale": 0.2, "seed": 313 } },
+      "outputs": { "y": { "shape": [512], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "name": "stable max-plus-exp traffic", "value": "2 * args.rows * args.cols * 4" }
+        ]
+      }
+    },
+    {
+      "name": "axis0-crossover-32768x1024",
+      "preset": "smoke",
+      "vars": { "rows": 32768, "cols": 1024 },
+      "attrs": { "axes": [0], "keepdims": 0 },
+      "inputs": { "x": { "shape": [32768, 1024], "dtype": "float32", "dist": "normal", "scale": 0.2, "seed": 314 } },
+      "outputs": { "y": { "shape": [1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "name": "stable max-plus-exp traffic", "value": "2 * args.rows * args.cols * 4" }
+        ]
+      }
+    },
     {
       "name": "fullreduce-single-lane-r3-101x101x101-numel-not-mult4",
       "preset": "stress",
       "name": "reducelogsumexp-spatial-axes12-f32-16x256x256-low-lane",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic",
         "notes": "Low-lane rank-3 contiguous-suffix reduction that verifies the vec4 subgroup reducer and no-subgroup workgroup-tree fallback."
       },
       "vars": { "batch": 16, "height": 256, "width": 256 },

build/webgpu/manifest.json CHANGED Viewed

@@ -34,7 +34,10 @@
     "ROW_SERIAL_MAX_COLS": { "default": 1024 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "treeWorkgroupOk": "reduceWorkgroupSize > 0 and pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize and reduceWorkgroupSize * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize",
     "subgroupWorkgroupFloor": "min(reduceWorkgroupSize, max(1, device.adapterInfo.subgroupMaxSize))",
@@ -62,127 +65,117 @@
     "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) > 0 and numel(shapes.x) % numel(shapes.y) == 0 and numel(shapes.x) / numel(shapes.y) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.x == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.y) == dim(shapes.x, 0)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) == false and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0)))"
   },
   "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "y": { "buffer": "storage", "elementType": "$T" },
-    "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
         { "name": "chunkCount", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y) / tunables.VECTOR_WIDTH" }
       ]
     },
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$T" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
-        { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" }
       ]
     },
-    "params_3": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
     },
-    "params_5": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
     },
-    "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" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_9": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
     "partials": { "buffer": "storage", "elementType": "$partialElement" },
-    "params_10": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "inner", "type": "u32", "value": "axisSplitInner" },
         { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" }
       ]
     },
-    "partials_2": { "name": "partials", "buffer": "read-only-storage", "elementType": "$partialElement" },
-    "params_11": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }]
-    },
-    "params_12": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
       ]
     },
-    "params_13": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }]
-    },
-    "params_16": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_17": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
-        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
-    "params_18": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_19": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     }
@@ -203,15 +196,9 @@
           "id": "main",
           "name": "ReduceLogSumExp.ContiguousSuffixSubgroupVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y", "params"],
-          "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -229,13 +216,8 @@
           "id": "main",
           "name": "ReduceLogSumExp.ContiguousSuffixTreeVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y", "params"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
@@ -253,8 +235,7 @@
           "id": "main",
           "name": "ReduceLogSumExp.ContiguousSuffixTree",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "y", "params_2"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
@@ -270,7 +251,6 @@
           "name": "ReduceLogSumExp.MultiAxisRank3",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"logsumexp\"",
             "indexing": "\"multiaxis\"",
             "rank": 3,
             "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
@@ -278,11 +258,9 @@
             "outputShape": "shapes.y",
             "outputRank": "ranks.y",
             "keepDims": "attrs.keepdims != 0",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
           },
-          "bindings": ["x_2", "y", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -302,7 +280,6 @@
           "name": "ReduceLogSumExp.MultiAxisRank4",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"logsumexp\"",
             "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)"],
@@ -310,11 +287,9 @@
             "outputShape": "shapes.y",
             "outputRank": "ranks.y",
             "keepDims": "attrs.keepdims != 0",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
           },
-          "bindings": ["x_2", "y", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -337,7 +312,7 @@
           "shader": "reduce-i32-axes02.wgsl.jinja",
           "derive": { "op": "\"logsumexp\"", "workgroupSizeSpec": "axes02WorkgroupSize" },
           "bindings": [
-            "x_2",
             "y",
             {
               "name": "params",
@@ -368,7 +343,7 @@
           "name": "ReduceLogSumExp.NoopEmptyAxes",
           "shader": "reduce-noop-empty-axes.wgsl.jinja",
           "derive": { "op": "\"identity\"" },
-          "bindings": ["x_2", "y", "params_5"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -393,14 +368,12 @@
           "id": "main",
           "name": "ReduceLogSumExp.SubgroupRowsVec4",
           "shader": "reduce-row-subgroup-rows.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x", "y", "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"
         }
       ]
     },
@@ -419,13 +392,8 @@
           "id": "main",
           "name": "ReduceLogSumExp.TreeRowVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y", "params_6"],
           "dispatch": { "x": "min(lastAxisRows, 65535)", "y": "ceilDiv(lastAxisRows, 65535)", "z": 1 }
         }
       ]
@@ -441,14 +409,11 @@
           "name": "ReduceLogSumExp.Rank0Scalar",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"logsumexp\"",
             "indexing": "\"axis2d\"",
             "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
-          "bindings": ["x_2", "y", "params_7"],
           "dispatch": { "x": 1 }
         }
       ]
@@ -463,14 +428,11 @@
           "name": "ReduceLogSumExp.Rank1Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"logsumexp\"",
             "indexing": "\"axis2d\"",
             "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
-          "bindings": ["x_2", "y", "params_8"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -490,8 +452,7 @@
           "id": "main",
           "name": "ReduceLogSumExp.Axis1Parallel",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "y", "params_9"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -516,13 +477,8 @@
           "id": "split_reduce",
           "name": "ReduceLogSumExp.AxisSplitReduce",
           "shader": "reduce-axis-split-reduce.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "splitSpec": "splitCount",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_10"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -533,8 +489,8 @@
           "id": "combine",
           "name": "ReduceLogSumExp.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
@@ -561,14 +517,8 @@
           "id": "split_reduce",
           "name": "ReduceLogSumExp.AxisSplitTiledReduce",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "splitSpec": "splitCount",
-            "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_10"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -579,8 +529,8 @@
           "id": "combine",
           "name": "ReduceLogSumExp.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
@@ -593,6 +543,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\"",
@@ -605,13 +556,8 @@
           "id": "split_reduce",
           "name": "ReduceLogSumExp.Axis0SplitKReduce",
           "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "splitSpec": "splitCount",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_12"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -622,8 +568,8 @@
           "id": "combine",
           "name": "ReduceLogSumExp.Axis0SplitKCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y", "params_13"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
@@ -642,8 +588,8 @@
           "id": "main",
           "name": "ReduceLogSumExp.Axis0TileCols",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
-          "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "y", "params_12"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
@@ -657,7 +603,6 @@
       "priority": 31,
       "when": ["flatParallelCovered"],
       "derive": {
-        "scalar": "dtypes.T",
         "workgroupSize": "reduceWorkgroupSize",
         "flatScalar": "\"vec4<\" ~ dtypes.T ~ \">\" if numel(shapes.x) % tunables.VECTOR_WIDTH == 0 else dtypes.T",
         "split": "flatSplitCount"
@@ -668,14 +613,10 @@
           "id": "flat_partial",
           "name": "ReduceLogSumExp.AllAxesFlatPartial",
           "shader": "reduce-flat-partial-logsumexp.wgsl.jinja",
-          "derive": {
-            "vec4": "numel(shapes.x) % tunables.VECTOR_WIDTH == 0",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
           "bindings": [
             { "arg": "x", "elementType": "$flatScalar" },
-            { "name": "partials", "buffer": "storage", "elementType": "f32" },
             { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "flatItems" }] }
           ],
           "dispatch": { "x": "flatSplitCount" }
@@ -706,7 +647,6 @@
           "name": "ReduceLogSumExp.RankNSingleAxisGeneric",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"logsumexp\"",
             "indexing": "\"rankn\"",
             "rank": "ranks.x",
             "axisSpec": "reduceAxis",
@@ -714,11 +654,9 @@
             "outputShape": "shapes.y",
             "outputRank": "ranks.y",
             "keepDims": "attrs.keepdims != 0",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
           },
-          "bindings": ["x_2", "y", "params_16"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -742,19 +680,13 @@
           "id": "main",
           "name": "ReduceLogSumExp.SubgroupRowVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y", "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"
         }
       ]
     },
@@ -772,19 +704,13 @@
           "id": "main",
           "name": "ReduceLogSumExp.SubgroupRow",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "vec4": false,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y", "params_17"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -797,17 +723,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "axis": 0,
-            "op": "\"logsumexp\"",
-            "indexing": "\"axis2d\"",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y", "params_18"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -824,17 +743,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "axis1",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "axis": 1,
-            "op": "\"logsumexp\"",
-            "indexing": "\"axis2d\"",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y", "params_19"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -853,25 +765,8 @@
           "id": "main",
           "name": "ReduceLogSumExp.Rank3AllAxesKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "indexing": "\"axis2d\"",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": [
-            "x_2",
-            "y",
-            {
-              "name": "params",
-              "struct": [
-                { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
-                { "name": "cols", "type": "u32", "value": "1" },
-                { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
-              ]
-            }
-          ],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -890,25 +785,8 @@
           "id": "main",
           "name": "ReduceLogSumExp.Rank3AllAxesNoKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"logsumexp\"",
-            "indexing": "\"axis2d\"",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": [
-            "x_2",
-            "y",
-            {
-              "name": "params",
-              "struct": [
-                { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
-                { "name": "cols", "type": "u32", "value": "1" },
-                { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
-              ]
-            }
-          ],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",

     "ROW_SERIAL_MAX_COLS": { "default": 1024 }
   },
   "derive": {
+    "reduceOp": "\"logsumexp\"",
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "op": "reduceOp",
+    "castF32": "dtypes.T == \"f16\"",
     "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "treeWorkgroupOk": "reduceWorkgroupSize > 0 and pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize and reduceWorkgroupSize * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize",
     "subgroupWorkgroupFloor": "min(reduceWorkgroupSize, max(1, device.adapterInfo.subgroupMaxSize))",
     "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_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)" }]
     },
+    "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_out_count": {
+      "name": "params",
+      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
+    },
+    "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] },
+    "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" },
         { "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)" }
       ]
     },
     "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_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)" }
       ]
     }
           "id": "main",
           "name": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.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 }
         }
       ]
           "name": "ReduceLogSumExp.MultiAxisRank3",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "indexing": "\"multiaxis\"",
             "rank": 3,
             "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
             "outputShape": "shapes.y",
             "outputRank": "ranks.y",
             "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\""
           },
+          "bindings": ["x_t", "y", "params_out_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "name": "ReduceLogSumExp.MultiAxisRank4",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "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)"],
             "outputShape": "shapes.y",
             "outputRank": "ranks.y",
             "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\""
           },
+          "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-i32-axes02.wgsl.jinja",
           "derive": { "op": "\"logsumexp\"", "workgroupSizeSpec": "axes02WorkgroupSize" },
           "bindings": [
+            "x_t",
             "y",
             {
               "name": "params",
           "name": "ReduceLogSumExp.NoopEmptyAxes",
           "shader": "reduce-noop-empty-axes.wgsl.jinja",
           "derive": { "op": "\"identity\"" },
+          "bindings": ["x_t", "y", "params_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "id": "main",
           "name": "ReduceLogSumExp.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": "ReduceLogSumExp.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 }
         }
       ]
           "name": "ReduceLogSumExp.Rank0Scalar",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "indexing": "\"axis2d\"",
             "intMode": "dtypes.T == \"i32\"",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
+          "bindings": ["x_t", "y", "params_rank0_scalar"],
           "dispatch": { "x": 1 }
         }
       ]
           "name": "ReduceLogSumExp.Rank1Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "indexing": "\"axis2d\"",
             "intMode": "dtypes.T == \"i32\"",
             "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": "main",
           "name": "ReduceLogSumExp.Axis1Parallel",
           "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": "split_reduce",
           "name": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.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)",
       "priority": 31,
       "when": ["flatParallelCovered"],
       "derive": {
         "workgroupSize": "reduceWorkgroupSize",
         "flatScalar": "\"vec4<\" ~ dtypes.T ~ \">\" if numel(shapes.x) % tunables.VECTOR_WIDTH == 0 else dtypes.T",
         "split": "flatSplitCount"
           "id": "flat_partial",
           "name": "ReduceLogSumExp.AllAxesFlatPartial",
           "shader": "reduce-flat-partial-logsumexp.wgsl.jinja",
+          "derive": { "vec4": "numel(shapes.x) % tunables.VECTOR_WIDTH == 0" },
           "bindings": [
             { "arg": "x", "elementType": "$flatScalar" },
+            { "name": "partials", "elementType": "f32" },
             { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "flatItems" }] }
           ],
           "dispatch": { "x": "flatSplitCount" }
           "name": "ReduceLogSumExp.RankNSingleAxisGeneric",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
             "indexing": "\"rankn\"",
             "rank": "ranks.x",
             "axisSpec": "reduceAxis",
             "outputShape": "shapes.y",
             "outputRank": "ranks.y",
             "keepDims": "attrs.keepdims != 0",
+            "intMode": "dtypes.T == \"i32\""
           },
+          "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": "ReduceLogSumExp.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": "ReduceLogSumExp.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": "ReduceLogSumExp.Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"" },
+          "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": "ReduceLogSumExp.Axis1",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"" },
+          "bindings": ["x_t", "y", "params_axis1"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "id": "main",
           "name": "ReduceLogSumExp.Rank3AllAxesKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"" },
+          "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": "main",
           "name": "ReduceLogSumExp.Rank3AllAxesNoKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
+          "derive": { "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"" },
+          "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.ReduceLogSumExp",
-  "id": "_ai_onnx_reducelogsumexp_webgpu_fd6e0c7",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "msizeqfB88GjW8v7QU9KoKyjW3qc8lJJXaCPf6OejeI=",
-      "manifest.json": "ofeVPiTS0KOqgOnbSIdNK8ULQiV3GSHu3Sxay4FY9U0=",
-      "reduce-axis-split-reduce.wgsl.jinja": "4+ep9xH4pHZOfaZ8abJhA8SW5M4mUhDC+y1CtDz6vjY=",
-      "reduce-axis0-splitk-combine.wgsl.jinja": "Yz1hjK55R/kndUrw3ugPqgPaOBwKmYupqZoVdJTHO5Q=",
-      "reduce-axis0-splitk-reduce.wgsl.jinja": "jB2h58emn6rfKhd8ALzqylSDsmBbrrkEOhtIoNhcQM0=",
-      "reduce-axis0-tilecols.wgsl.jinja": "PjYkEUQJBeG70td3W2xxmexH9x6XNIfeSzXBY47HbaU=",
-      "reduce-flat-combine-logsumexp.wgsl.jinja": "vwwisgGJGC6hBX/WYMf/OEXKa71pMJKX6v2ggWuXWpM=",
-      "reduce-flat-partial-logsumexp.wgsl.jinja": "5ZVJnWsR0xWm4Oet3oqP0V7ePeqLea5wBs/rBsy9L20=",
       "reduce-i32-axes02.wgsl.jinja": "FsDSFjExTqueGACUuk9bPO7gN5yp8sJHtEXapuxter0=",
-      "reduce-noop-empty-axes.wgsl.jinja": "NNvXRO0Tt3Mrvk2Wdssndbea61oQN4tj+O08NVA43Ns=",
-      "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": "QyxQrLhFNbxc+mf7VOuv8RipIY2AlNZEhJZJvTQjuNw="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "contiguous_suffix_subgroup_vec4": ["reduce-row-subgroup.wgsl.jinja"],
       "contiguous_suffix_tree_vec4": ["reduce-row-tree.wgsl.jinja"],

 {
   "name": "ai.onnx.ReduceLogSumExp",
+  "id": "_ai_onnx_reducelogsumexp_webgpu_4aeeb2f",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "wXvsZOuaZHscg7oIb5NujBZvnlo8B9U6YLfa69+npYg=",
+      "manifest.json": "QcxF0HuWIGT+TZaAjJgOcvYBfwAKV98lh+fie3k41dM=",
+      "reduce-axis-split-reduce.wgsl.jinja": "xo4jwtx/MeHIZkh2KU31ZCR+Ak58xOe4Cg3HWrcvdqo=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "bY5/z/a29FortWA0aT28cnkQX6dzLGr1O1fV/ncfxLs=",
+      "reduce-axis0-splitk-reduce.wgsl.jinja": "rW4Vltn1kPBqmbVTl7Za5CWycPaUe7IwaqXoWH+72Tw=",
+      "reduce-axis0-tilecols.wgsl.jinja": "dyX/c20cWbPRfgqDdOi5XqfhvFtxyQD8KDJJ2gJ3fqc=",
+      "reduce-flat-combine-logsumexp.wgsl.jinja": "YxvwTOwRzXW5NTZ4tEtrt6syg32z4ELS2qW2/KKIHp4=",
+      "reduce-flat-partial-logsumexp.wgsl.jinja": "6OTGSyRzItrZtgKYi4OC+Ay3j8c/AoCu3W0tU/mKI5A=",
       "reduce-i32-axes02.wgsl.jinja": "FsDSFjExTqueGACUuk9bPO7gN5yp8sJHtEXapuxter0=",
+      "reduce-noop-empty-axes.wgsl.jinja": "I3eV+jyJXA4ojj9U6RPZnivKP4CD6R4HFmcM0i7v6vA=",
+      "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": "XZthzOtNawNeG2SXq528l1yO5sVPqTTK3Y9GVhJ46lk=",
+      "test.json": "IfasxXu3otrLuI7M2dRKeSFj/wZJcDf5OvKM7a9luNo="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "contiguous_suffix_subgroup_vec4": ["reduce-row-subgroup.wgsl.jinja"],
       "contiguous_suffix_tree_vec4": ["reduce-row-tree.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,20 @@
 // 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 +31,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" %}
@@ -101,9 +86,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 intMode = intMode is defined and intMode %}
 {% 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" %}
       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,7 +13,6 @@
 {% 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 %}
@@ -21,9 +20,6 @@
 {% 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;

 {% set rowEnd = "params.rows" %}
 {% set elem = "x[inputBase + row * params.cols + col]" %}
 {% endif %}
 {% set intMode = intMode is defined and intMode %}
 {% set scalar = "f32" if castF32 else scalar %}
 {% if castF32 %}
 {% 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;

build/webgpu/reduce-flat-combine-logsumexp.wgsl.jinja CHANGED Viewed

@@ -6,9 +6,6 @@
 // transcendental merges when the output has exactly one element.
 {% set yv = "f16(" if outputF16 else "" %}
 {% set vy = ")" if outputF16 else "" %}
-{% if outputF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -25,7 +22,6 @@ fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
   let tid = lid.x;
@@ -69,15 +65,23 @@ fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
   }
   workgroupBarrier();
-  stride = WG / 2u;
   loop {
-    if (stride == 0u) { break; }
-    if (tid < stride) {
-      wsum[tid] = wsum[tid] + wsum[tid + stride];
     }
-    stride = stride / 2u;
     workgroupBarrier();
-  }
   if (tid == 0u) {
     let has_positive_inf = global_max > F32_MAX;

 // transcendental merges when the output has exactly one element.
 {% set yv = "f16(" if outputF16 else "" %}
 {% set vy = ")" if outputF16 else "" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
   let tid = lid.x;
   }
   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) %}
+  {{ svar }} = {{ 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(["wsum"], idx="tid", wg="WG", form="head", breakInline=true, reuse=true) }}
   if (tid == 0u) {
     let has_positive_inf = global_max > F32_MAX;

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

@@ -4,13 +4,8 @@
 // Workgroups write three partial planes: segment maximum, shifted exponential
 // sum, and a NaN marker. The combine pass merges them stably and emits
 // globalMaximum + log(sum), preserving NaN and positive infinity.
-{% set vec4 = vec4 | default(true) %}
-{% 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 }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -25,7 +20,6 @@ fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
@@ -85,13 +79,19 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   }
   wsum[tid] = local_sum;
   workgroupBarrier();
-  stride = WG / 2u;
   loop {
-    if (stride == 0u) { break; }
-    if (tid < stride) { wsum[tid] = wsum[tid] + wsum[tid + stride]; }
-    stride = stride / 2u;
     workgroupBarrier();
-  }
   if (tid == 0u) {
     partials[wg.x] = seg_max;

 // Workgroups write three partial planes: segment maximum, shifted exponential
 // sum, and a NaN marker. The combine pass merges them stably and emits
 // globalMaximum + log(sum), preserving NaN and positive infinity.
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
   }
   wsum[tid] = local_sum;
   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) %}
+  {{ svar }} = {{ wg }} / 2u;
   loop {
+    if ({{ svar }} == 0u) { break; }
+    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
+    {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
+{{ wgsl_tree_fold(["wsum"], idx="tid", wg="WG", form="head", breakInline=true, bodyInline=true, reuse=true) }}
   if (tid == 0u) {
     partials[wg.x] = seg_max;

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];
   y[i] = 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];
   y[i] = 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" %}
@@ -131,7 +129,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,12 +165,7 @@ 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

 {% 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 logicalBool = logicalBool is defined and logicalBool %}
 {% 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" %}
 {% 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 intMode %}
   // Integer logsumexp widens each element for exp/log, then truncates the result

build/webgpu/test.json CHANGED Viewed

@@ -665,6 +665,19 @@
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.0001 } }
     },
     {
       "name": "axis0_splitk_8192x32",
       "attrs": { "axes": [0], "keepdims": 0 },
@@ -989,6 +1002,19 @@
       "outputs": {
         "y": { "dtype": "float32", "shape": [64], "tolerance": 0.0001, "allowNaN": true, "relTolerance": 0.0001 }
       }
     }
   ]
 }

       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0.0001 } }
     },
+    {
+      "name": "axis0_variable_subgroup_tile_selection_1024x64",
+      "attrs": { "axes": [0], "keepdims": 0 },
+      "tunables": { "AXIS0_SPLIT_MIN_ROWS": 512, "AXIS0_SPLIT_TARGET_ROWS": 256 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1024, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.07, "scale": 0.2 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [64], "tolerance": 0.0001 } }
+    },
     {
       "name": "axis0_splitk_8192x32",
       "attrs": { "axes": [0], "keepdims": 0 },
       "outputs": {
         "y": { "dtype": "float32", "shape": [64], "tolerance": 0.0001, "allowNaN": true, "relTolerance": 0.0001 }
       }
+    },
+    {
+      "name": "ort_empty_set_noop_with_empty_axes_identity",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc",
+        "test": "ReductionOpTest.EmptySetNoopWithEmptyAxes (ReduceLogSumExp)",
+        "notes": "Empty-set input combined with noop_with_empty_axes=1: the output must keep the input shape, not a -inf-filled reduced shape."
+      },
+      "attrs": { "keepdims": 1, "noop_with_empty_axes": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [2, 0, 4], "data": { "kind": "values", "values": [] } } },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 0, 4], "tolerance": 0, "data": { "kind": "values", "values": [] } }
+      }
     }
   ]
 }