Xenova HF Staff commited on 8 days ago

Commit

fe742ec

verified ·

1 Parent(s): 1577a99

sync 6fdf6301e2bb

Browse files

Files changed (17) hide show

README.md +6 -4
build/webgpu/bench.json +1 -1
build/webgpu/manifest.json +127 -251
build/webgpu/metadata.json +19 -19
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-partial.wgsl.jinja +18 -31
build/webgpu/reduce-multi-axis-coop.wgsl.jinja +24 -9
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/reduce-strided-axis.wgsl.jinja +7 -4
build/webgpu/test.json +2 -2

README.md CHANGED Viewed

@@ -48,6 +48,8 @@ Default values (overridable per request):
 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
 - `strided_axis_serial` — Flatten a single non-last reduction axis into outer/axis/inner geometry. Compile its strides and loop bound, keep one output per lane and float32 accumulation, and cap the workgroup by the device limits.
 ## Device requirements
@@ -59,7 +61,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)
@@ -76,7 +78,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.
@@ -95,7 +97,7 @@ import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/ai.onnx.ReduceMean", { 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" } },
 });
 ```

 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.
 - `strided_axis_serial` — Flatten a single non-last reduction axis into outer/axis/inner geometry. Compile its strides and loop bound, keep one output per lane and float32 accumulation, and cap the workgroup by the device limits.
 ## Device requirements
 - [`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.ReduceMean", { 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

@@ -130,7 +130,7 @@
       "name": "reducemean-spatial-axes23-f32-2x64x256x256-globalpool-pattern",
       "preset": "stress",
       "provenance": {
-        "source": "repository-authored",
         "notes": "Rank-4 contiguous-suffix reduction that verifies the vec4 subgroup reducer and portable workgroup-tree fallback."
       },
       "vars": { "batch": 2, "channels": 64, "height": 256, "width": 256 },

       "name": "reducemean-spatial-axes23-f32-2x64x256x256-globalpool-pattern",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic",
         "notes": "Rank-4 contiguous-suffix reduction that verifies the vec4 subgroup reducer and portable workgroup-tree fallback."
       },
       "vars": { "batch": 2, "channels": 64, "height": 256, "width": 256 },

build/webgpu/manifest.json CHANGED Viewed

@@ -36,6 +36,7 @@
     "MULTI_AXIS_COOP_MIN_REDUCED": { "default": 256 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
@@ -63,149 +64,147 @@
     "flatScratchBytes": "flatSplitCount * dtypeBytes(\"float32\")",
     "flatPathFits": "treeWorkgroupOk and flatSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize",
     "flatParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) == 1 and numel(shapes.x) >= tunables.FULL_REDUCE_MIN_ELEMENTS and flatPathFits",
-    "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) > 0 and numel(shapes.x) % numel(shapes.y) == 0 and numel(shapes.x) / numel(shapes.y) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.x == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.y) == dim(shapes.x, 0)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) == false and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0)))"
   },
   "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "y": { "buffer": "storage", "elementType": "$T" },
-    "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
         { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" },
         { "name": "chunkCount", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y) / tunables.VECTOR_WIDTH" }
       ]
     },
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$T" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
-        { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" }
       ]
     },
-    "params_3": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
     },
-    "params_4": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
     },
-    "params_5": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" },
-        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1) / tunables.VECTOR_WIDTH" }
       ]
     },
-    "params_6": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "1" },
         { "name": "cols", "type": "u32", "value": "1" },
         { "name": "outCount", "type": "u32", "value": "1" }
       ]
     },
-    "params_7": {
       "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_8": {
       "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_9": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "inner", "type": "u32", "value": "axisSplitInner" },
         { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" }
       ]
     },
-    "partials_2": { "name": "partials", "buffer": "read-only-storage", "elementType": "$partialElement" },
-    "params_10": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [
-        { "name": "rows", "type": "u32", "value": "axisSplitDim" },
-        { "name": "cols", "type": "u32", "value": "axisSplitOutputs" }
-      ]
-    },
-    "params_11": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
       ]
     },
-    "partials_3": { "name": "partials", "buffer": "storage", "elementType": "f32" },
-    "params_12": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "count4", "type": "u32", "value": "floor(numel(shapes.x) / tunables.VECTOR_WIDTH)" },
         { "name": "numel", "type": "u32", "value": "numel(shapes.x)" }
       ]
     },
-    "partials_4": { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
-    "params_13": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [
-        { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
-        { "name": "cols", "type": "u32", "value": "1" }
-      ]
-    },
-    "params_15": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_16": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
-        { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" },
-        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
-    "params_17": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_18": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
-        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     }
@@ -226,15 +225,9 @@
           "id": "main",
           "name": "ReduceMean.ContiguousSuffixSubgroupVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "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"
         }
       ]
     },
@@ -252,13 +245,8 @@
           "id": "main",
           "name": "ReduceMean.ContiguousSuffixTreeVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "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 }
         }
       ]
@@ -276,8 +264,7 @@
           "id": "main",
           "name": "ReduceMean.ContiguousSuffixTree",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "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 }
         }
       ]
@@ -292,7 +279,7 @@
           "id": "main",
           "name": "ReduceMean.NoopEmptyAxes",
           "shader": "reduce-noop-empty-axes.wgsl.jinja",
-          "bindings": ["x_2", "y", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -312,7 +299,6 @@
           "name": "ReduceMean.MultiAxisRank3Coop",
           "shader": "reduce-multi-axis-coop.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "indexing": "\"multiaxis\"",
             "rank": 3,
             "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
@@ -320,11 +306,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"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
@@ -340,7 +324,6 @@
           "name": "ReduceMean.MultiAxisRank4Coop",
           "shader": "reduce-multi-axis-coop.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "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)"],
@@ -348,11 +331,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"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
@@ -368,7 +349,6 @@
           "name": "ReduceMean.MultiAxisRank3",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "indexing": "\"multiaxis\"",
             "rank": 3,
             "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"],
@@ -376,11 +356,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_4"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -400,7 +378,6 @@
           "name": "ReduceMean.MultiAxisRank4",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "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)"],
@@ -408,11 +385,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_4"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -437,14 +412,12 @@
           "id": "main",
           "name": "ReduceMean.SubgroupRowsVec4",
           "shader": "reduce-row-subgroup-rows.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x", "y", "params_5"],
           "dispatch": {
             "x": "min(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
             "y": "ceilDiv(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -463,13 +436,8 @@
           "id": "main",
           "name": "ReduceMean.TreeRowVec4",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y", "params_5"],
           "dispatch": { "x": "min(lastAxisRows, 65535)", "y": "ceilDiv(lastAxisRows, 65535)", "z": 1 }
         }
       ]
@@ -485,14 +453,11 @@
           "name": "ReduceMean.Rank0Scalar",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "indexing": "\"axis2d\"",
             "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "logicalBool": "tensorDtypes.x == \"bool\""
           },
-          "bindings": ["x_2", "y", "params_6"],
           "dispatch": { "x": 1 }
         }
       ]
@@ -507,14 +472,11 @@
           "name": "ReduceMean.Rank1Axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "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": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -537,8 +499,7 @@
           "id": "main",
           "name": "ReduceMean.Axis1Parallel",
           "shader": "reduce-row-tree.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "y", "params_8"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
@@ -563,13 +524,8 @@
           "id": "split_reduce",
           "name": "ReduceMean.AxisSplitReduce",
           "shader": "reduce-axis-split-reduce.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "splitSpec": "splitCount",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_9"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -580,8 +536,8 @@
           "id": "combine",
           "name": "ReduceMean.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y", "params_10"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
@@ -608,14 +564,8 @@
           "id": "split_reduce",
           "name": "ReduceMean.AxisSplitTiledReduce",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "splitSpec": "splitCount",
-            "tileCols": "tunables.AXIS_SPLIT_TILE_COLS",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_9"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -626,8 +576,8 @@
           "id": "combine",
           "name": "ReduceMean.AxisSplitCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y", "params_10"],
           "dispatch": {
             "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)",
@@ -640,6 +590,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\"",
@@ -652,13 +603,8 @@
           "id": "split_reduce",
           "name": "ReduceMean.Axis0SplitKReduce",
           "shader": "reduce-axis0-splitk-reduce.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "splitSpec": "splitCount",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "partials", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)",
             "y": "splitCount",
@@ -669,8 +615,8 @@
           "id": "combine",
           "name": "ReduceMean.Axis0SplitKCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_2", "y", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)",
@@ -689,13 +635,8 @@
           "id": "main",
           "name": "ReduceMean.Axis0TileCols",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y", "params_11"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
@@ -708,23 +649,22 @@
       "id": "all_axes_flat",
       "priority": 31,
       "when": ["flatParallelCovered"],
-      "derive": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceMean.AllAxesFlatPartial",
           "shader": "reduce-flat-partial.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" },
-          "bindings": ["x_2", "partials_3", "params_12"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceMean.AllAxesFlatCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": { "op": "\"mean\"", "outputF16": "dtypes.T == \"f16\"" },
-          "bindings": ["partials_4", "y", "params_13"],
           "dispatch": { "x": 1 }
         }
       ]
@@ -777,7 +717,6 @@
           "name": "ReduceMean.RankNSingleAxisGeneric",
           "shader": "reduce-serial-axis.wgsl.jinja",
           "derive": {
-            "op": "\"mean\"",
             "indexing": "\"rankn\"",
             "rank": "ranks.x",
             "axisSpec": "reduceAxis",
@@ -785,11 +724,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_15"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -813,19 +750,13 @@
           "id": "main",
           "name": "ReduceMean.SubgroupRowVec4",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "vec4": true,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x", "y", "params_5"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -843,19 +774,13 @@
           "id": "main",
           "name": "ReduceMean.SubgroupRow",
           "shader": "reduce-row-subgroup.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "vec4": false,
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y", "params_16"],
           "dispatch": {
             "x": "min(rows(shapes.x, ranks.x - 1), 65535)",
             "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -868,17 +793,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "axis0",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "axis": 0,
-            "op": "\"mean\"",
-            "indexing": "\"axis2d\"",
-            "intMode": "dtypes.T == \"i32\"",
-            "castF32": "dtypes.T == \"f16\"",
-            "usesF16Spec": "dtypes.T == \"f16\""
-          },
-          "bindings": ["x_2", "y", "params_17"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
@@ -895,17 +813,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "axis1",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "axis": 1,
-            "op": "\"mean\"",
-            "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)",
@@ -915,34 +826,17 @@
       ]
     },
     {
-      "id": "all_axes_no_keepdims",
       "priority": 30,
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 0 and ranks.y == 0"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
-          "name": "ReduceMean.Rank3AllAxesNoKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "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)",
@@ -952,34 +846,17 @@
       ]
     },
     {
-      "id": "all_axes_keepdims",
       "priority": 30,
-      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 1 and ranks.y == ranks.x and numel(shapes.y) == 1"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
-          "name": "ReduceMean.Rank3AllAxesKeepdims",
           "shader": "reduce-serial-axis.wgsl.jinja",
-          "derive": {
-            "op": "\"mean\"",
-            "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)",
@@ -1003,7 +880,6 @@
             "axisDim": "axisSplitDim",
             "innerSize": "axisSplitInner",
             "usesF16": "dtypes.T == \"f16\"",
-            "op": "\"mean\"",
             "tailPaddingElements": "numel(shapes.y) % 2 if dtypes.T == \"f16\" else 0"
           },
           "bindings": [{ "arg": "x", "elementType": "$scalar" }, { "arg": "y", "elementType": "$scalar" }],

     "MULTI_AXIS_COOP_MIN_REDUCED": { "default": 256 }
   },
   "derive": {
+    "reduceOp": "\"mean\"",
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
     "flatScratchBytes": "flatSplitCount * dtypeBytes(\"float32\")",
     "flatPathFits": "treeWorkgroupOk and flatSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize",
     "flatParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) == 1 and numel(shapes.x) >= tunables.FULL_REDUCE_MIN_ELEMENTS and flatPathFits",
+    "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) > 0 and numel(shapes.x) % numel(shapes.y) == 0 and numel(shapes.x) / numel(shapes.y) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.x == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.y) == dim(shapes.x, 0)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) == false and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0)))",
+    "op": "reduceOp",
+    "castF32": "dtypes.T == \"f16\""
   },
   "bindings": {
+    "params_suffix_vec4": {
+      "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
         { "name": "cols", "type": "u32", "value": "numel(shapes.x) / 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": "cols", "type": "u32", "value": "dim(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": "rows", "type": "u32", "value": "axisSplitDim" },
+        { "name": "cols", "type": "u32", "value": "axisSplitOutputs" }
+      ]
     },
+    "params_axis0_combine": {
       "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }
+      ]
     },
+    "partials_flat": { "name": "partials", "buffer": "storage", "elementType": "f32" },
+    "partials_flat_read": { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+    "params_all_axes_flat": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
+        { "name": "cols", "type": "u32", "value": "1" }
+      ]
+    },
+    "params_rows_chunk_count": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" },
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" },
+        { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" }
       ]
     },
+    "x": { "elementType": "$vectorScalar" },
+    "y": { "elementType": "$T" },
+    "x_t": { "name": "x", "elementType": "$T" },
+    "params_main": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "numel(shapes.y)" },
+        { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" }
+      ]
+    },
+    "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] },
+    "params_out_count": {
+      "name": "params",
+      "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }]
+    },
+    "params_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_count4_numel": {
       "name": "params",
       "struct": [
         { "name": "count4", "type": "u32", "value": "floor(numel(shapes.x) / tunables.VECTOR_WIDTH)" },
         { "name": "numel", "type": "u32", "value": "numel(shapes.x)" }
       ]
     },
+    "params_axis_dim_out_count": {
       "name": "params",
       "struct": [
         { "name": "axisDim", "type": "u32", "value": "axisSplitDim" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
+    "params_axis0": {
       "name": "params",
       "struct": [
+        { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" },
+        { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
+    "params_axis1": {
       "name": "params",
       "struct": [
         { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
+    "params_all_axes": {
       "name": "params",
       "struct": [
+        { "name": "rows", "type": "u32", "value": "numel(shapes.x)" },
+        { "name": "cols", "type": "u32", "value": "1" },
         { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }
       ]
     }
           "id": "main",
           "name": "ReduceMean.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": "ReduceMean.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": "ReduceMean.ContiguousSuffixTree",
           "shader": "reduce-row-tree.wgsl.jinja",
+          "bindings": ["x_t", "y", "params_main"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
           "id": "main",
           "name": "ReduceMean.NoopEmptyAxes",
           "shader": "reduce-noop-empty-axes.wgsl.jinja",
+          "bindings": ["x_t", "y", "params_count"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)",
           "name": "ReduceMean.MultiAxisRank3Coop",
           "shader": "reduce-multi-axis-coop.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"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
           "name": "ReduceMean.MultiAxisRank4Coop",
           "shader": "reduce-multi-axis-coop.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"],
           "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }
         }
       ]
           "name": "ReduceMean.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": "ReduceMean.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)",
           "id": "main",
           "name": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.Axis0TileCols",
           "shader": "reduce-axis0-tilecols.wgsl.jinja",
+          "derive": { "intMode": "dtypes.T == \"i32\"" },
+          "bindings": ["x_t", "y", "params_split_reduce"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)",
       "id": "all_axes_flat",
       "priority": 31,
       "when": ["flatParallelCovered"],
+      "derive": { "workgroupSize": "reduceWorkgroupSize", "split": "flatSplitCount" },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[flatSplitCount]" }],
       "passes": [
         {
           "id": "flat_partial",
           "name": "ReduceMean.AllAxesFlatPartial",
           "shader": "reduce-flat-partial.wgsl.jinja",
+          "bindings": ["x_t", "partials_flat", "params_count4_numel"],
           "dispatch": { "x": "flatSplitCount" }
         },
         {
           "id": "combine",
           "name": "ReduceMean.AllAxesFlatCombine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials_flat_read", "y", "params_all_axes_flat"],
           "dispatch": { "x": 1 }
         }
       ]
           "name": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "ReduceMean.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": "all_axes_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 1", "ranks.y == ranks.x", "numel(shapes.y) == 1"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
+          "name": "ReduceMean.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": "all_axes_no_keepdims",
       "priority": 30,
+      "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 0", "ranks.y == 0"],
       "derive": { "axis": 0, "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
+          "name": "ReduceMean.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)",
             "axisDim": "axisSplitDim",
             "innerSize": "axisSplitInner",
             "usesF16": "dtypes.T == \"f16\"",
             "tailPaddingElements": "numel(shapes.y) % 2 if dtypes.T == \"f16\" else 0"
           },
           "bindings": [{ "arg": "x", "elementType": "$scalar" }, { "arg": "y", "elementType": "$scalar" }],

build/webgpu/metadata.json CHANGED Viewed

@@ -1,32 +1,32 @@
 {
   "name": "ai.onnx.ReduceMean",
-  "id": "_ai_onnx_reducemean_webgpu_3e440cd",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "pPlab/HIaUOReOgmFnA3jXKU0Vvj6XzHEsJ19hU56E8=",
-      "manifest.json": "DyrQEvKuUXsnVWpCbkUMBAXl+cjdPCJsxOc1nqmHEfA=",
-      "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": "I2BXFXhK2SWSwjfXB5BwDYyohcke/IPrl5SJE95kb5U=",
-      "reduce-flat-partial.wgsl.jinja": "+qToL+wFi9QOxvY887aBAEwZK6Xu/eLQkucKvJYNwSk=",
-      "reduce-multi-axis-coop.wgsl.jinja": "8v5G5dvcOUqPfLilBMg2ku67y2yYBggmsSqlBJuPUSE=",
-      "reduce-noop-empty-axes.wgsl.jinja": "5bKufrcfYt3BagZ6WxbUvqwKKKVMU7d+SUshgUFWTDI=",
-      "reduce-row-subgroup-rows.wgsl.jinja": "76u7rAvFoZZKrFDs2A2jk0vkB0uNrPL6twdOBE9b+v8=",
-      "reduce-row-subgroup.wgsl.jinja": "2mu9LEsk8HfaLvucBCfcB1/ENpXkD6ELiCtRt+5UqiU=",
-      "reduce-row-tree.wgsl.jinja": "Bwa5xcI0bTmKXb4r9Cc1bfVbM5rNqqpQVrWWVqcb8xA=",
-      "reduce-serial-axis.wgsl.jinja": "fvUV9htqzKzt4Pg05pYtRmup/5QIHYaUGQHJZnthXKo=",
-      "reduce-strided-axis.wgsl.jinja": "vbM7nz/UoIWB0tUb/rLmGalraF2JfjV9CT1RxP0Lfu8=",
-      "test.json": "86DHyI08xIMOwY+dUT7Kjh13d9OFGJ1WEc2OofnRroQ="
     }
   },
-  "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"],
@@ -52,8 +52,8 @@
       "subgroup_last_axis": ["reduce-row-subgroup.wgsl.jinja"],
       "axis0": ["reduce-serial-axis.wgsl.jinja"],
       "axis1": ["reduce-serial-axis.wgsl.jinja"],
-      "all_axes_no_keepdims": ["reduce-serial-axis.wgsl.jinja"],
       "all_axes_keepdims": ["reduce-serial-axis.wgsl.jinja"],
       "strided_axis_serial": ["reduce-strided-axis.wgsl.jinja"]
     }
   }

 {
   "name": "ai.onnx.ReduceMean",
+  "id": "_ai_onnx_reducemean_webgpu_b918009",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "bhV+fHo4TS6gK0PMa36JEwynO8XXhEGB9DlTDn1mBWk=",
+      "manifest.json": "rXpND7uYqiqCxsOFMhBDatls0p7CQ944DFSmMhomX/g=",
+      "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": "thqt34BS8Z2qD/Nyd67lSFruVudITEmfo2sNJZUeKQU=",
+      "reduce-flat-partial.wgsl.jinja": "gMQg8GcibV+NQadQLAprKk6K3w13GCOdfQvIdCvYDXs=",
+      "reduce-multi-axis-coop.wgsl.jinja": "mQldeEk3lvyh67eiW/mfZFaM2gpRc0fUUgpq7PI3SHo=",
+      "reduce-noop-empty-axes.wgsl.jinja": "MMZEXHeACtGwvDaH0nAyvMal+tAIjYrh2yxDuLfh1Lw=",
+      "reduce-row-subgroup-rows.wgsl.jinja": "5cMGE1dmrcawwHdoxQD8kPSurwA1wNHyoAQfxAOE7vU=",
+      "reduce-row-subgroup.wgsl.jinja": "7bRBWIQ7oHOq/E54hGrmXIzLznzVH9oTfYaf1jmBzV8=",
+      "reduce-row-tree.wgsl.jinja": "a+iEV/cBkrlsVz6NS0GY/TAsUSX5rO4N14wY1dKfDi0=",
+      "reduce-serial-axis.wgsl.jinja": "XZthzOtNawNeG2SXq528l1yO5sVPqTTK3Y9GVhJ46lk=",
+      "reduce-strided-axis.wgsl.jinja": "K+yHpuxWwKU75+/4M8BuQ8UU2e+gGKoklPj5I/uH5rU=",
+      "test.json": "1X6wSEe7Gq+J/wtGb7LTVh+F0Ifvw1FoCsLwOdgPMFQ="
     }
   },
+  "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"],
       "subgroup_last_axis": ["reduce-row-subgroup.wgsl.jinja"],
       "axis0": ["reduce-serial-axis.wgsl.jinja"],
       "axis1": ["reduce-serial-axis.wgsl.jinja"],
       "all_axes_keepdims": ["reduce-serial-axis.wgsl.jinja"],
+      "all_axes_no_keepdims": ["reduce-serial-axis.wgsl.jinja"],
       "strided_axis_serial": ["reduce-strided-axis.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-partial.wgsl.jinja CHANGED Viewed

@@ -3,33 +3,20 @@
 // in registers, reduce within each workgroup, and write one partial each. A
 // following combine pass folds the partials and applies the operation finalizer.
 //
-// Scalar f32 bindings keep arbitrary element counts legal. The grid-stride loop
-// manually assembles full vec4 groups from contiguous scalars, and one global
-// thread folds the final zero-to-three scalar elements exactly once.
 {% set intMode = intMode is defined and intMode %}
-{% set castF32 = castF32 is defined and castF32 %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
-{% if usesF16Spec is defined and usesF16Spec %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
  * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
-{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
 fn {{ name }}() -> {{ scalar }} {
-{% if scalar == "i32" %}
-  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
-{% elif scalar == "u32" %}
-  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
-{% else %}
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
-{% endif %}
-}
-{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 var<workgroup> red: array<{{ "i32" if intMode else "f32" }}, WG>;
@@ -100,23 +87,23 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   }
   red[tid] = acc;
   workgroupBarrier();
-  var stride: u32 = WG / 2u;
   loop {
-    if (stride == 0u) { break; }
-    if (tid < stride) {
-{% if op == "max" %}
-      red[tid] = max(red[tid], red[tid + stride]);
-{% elif op == "min" %}
-      red[tid] = min(red[tid], red[tid + stride]);
-{% elif op == "prod" %}
-      red[tid] = red[tid] * red[tid + stride];
-{% else %}
-      red[tid] = red[tid] + red[tid + stride];
-{% endif %}
     }
-    stride = stride / 2u;
     workgroupBarrier();
-  }
   if (tid == 0u) {
     partials[wg.x] = red[0];
   }

 // in registers, reduce within each workgroup, and write one partial each. A
 // following combine pass folds the partials and applies the operation finalizer.
 //
+// Scalar bindings support arbitrary element counts and a zero-to-three element
+// tail. An aligned packed binding loads each group directly, preserving the
+// component accumulation order and widening half storage before arithmetic.
 {% set intMode = intMode is defined and intMode %}
 {% set xa = "f32(" if castF32 else "" %}
 {% set ax = ")" if castF32 else "" %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
  * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 var<workgroup> red: array<{{ "i32" if intMode else "f32" }}, WG>;
   }
   red[tid] = acc;
   workgroupBarrier();
+{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" or op == "min" %}
+{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
+  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
+    if ({{ svar }} == 0u) { break; }
+    if ({{ idx }} < {{ svar }}) {
+{% for a in arrays %}
+      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
+{% endfor %}
     }
+    {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
+{{ wgsl_tree_fold(["red"], op=op, idx="tid", wg="WG", typed=true, form="head", breakInline=true) }}
   if (tid == 0u) {
     partials[wg.x] = red[0];
   }

build/webgpu/reduce-multi-axis-coop.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,8 +43,7 @@ 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 %}
 // Cooperative multi-axis reduction with one workgroup per output element.
 // Threads take strided shares of the reduced index range, then combine their
 // partials through a workgroup tree. The shared offset helper maps each reduced
@@ -53,17 +51,13 @@ fn input_offset(out_index: u32{% if hasReduced %}, reduce_linear: u32{% endif %}
 //
 // The tree changes f32 association relative to a left-to-right fold, so the two
 // accumulation orders need not be bit-identical.
-{% set castF32 = castF32 is defined and castF32 %}
-{% set intMode = intMode is defined and intMode %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-{% if usesF16Spec is defined and usesF16Spec %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set hasReducedAxis = namespace(value=false) %}
 {% for a in range(rank) %}{% if reduce[a] %}{% set hasReducedAxis.value = true %}{% endif %}{% endfor %}
-{{- reduce_multi_axis_offset(hasReducedAxis.value) }}
 {% set ACC = scalar if intMode else "f32" %}
 const WG: u32 = {{ workgroupSize }}u;
@@ -98,6 +92,7 @@ fn main(
     return;
   }
   let lid = lid3.x;
   if (REDUCED == 0u) {
 {% if intMode %}
     if (lid == 0u) { y[i] = {{ scalar }}(0); }
@@ -107,13 +102,23 @@ fn main(
     return;
   }
   var acc = {{ ACC }}(0);
   for (var r = lid; r < REDUCED; r = r + WG) {
 {% set at = "x[input_offset(i, r)]" %}
 {% if castF32 %}
 {% set at = "f32(" ~ at ~ ")" %}
 {% endif %}
     acc = combine(acc, {{ at }});
   }
   // Fold the per-thread accumulators. WORKGROUP_SIZE is a power of two, and every
@@ -134,9 +139,19 @@ fn main(
     return;
   }
   let total = partials[0];
 {% if intMode %}
   y[i] = total / {{ scalar }}(REDUCED);
 {% else %}
   y[i] = {{ yv }}total / f32(REDUCED){{ vy }};
 {% endif %}
 }

 {% macro reduce_multi_axis_offset(hasReduced) %}
 // One thread per output element walks the Cartesian product of the reduced axes,
 // linearized as reduce_linear. Specialized shapes make every input offset a sum
 // of coordinate-times-constant terms.
 {% set src.value = "(" ~ src.value ~ " * " ~ dataShape[a] ~ "u + coord" ~ a ~ ")" %}
 {% endfor %}
   return {{ src.value }};
+}{% endmacro %}
 // Cooperative multi-axis reduction with one workgroup per output element.
 // Threads take strided shares of the reduced index range, then combine their
 // partials through a workgroup tree. The shared offset helper maps each reduced
 //
 // The tree changes f32 association relative to a left-to-right fold, so the two
 // accumulation orders need not be bit-identical.
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
 {{ env.wgsl.resourceDeclarations }}
 {% set hasReducedAxis = namespace(value=false) %}
 {% for a in range(rank) %}{% if reduce[a] %}{% set hasReducedAxis.value = true %}{% endif %}{% endfor %}
+{{ reduce_multi_axis_offset(hasReducedAxis.value) }}
 {% set ACC = scalar if intMode else "f32" %}
 const WG: u32 = {{ workgroupSize }}u;
     return;
   }
   let lid = lid3.x;
+{% if op == "mean" %}
   if (REDUCED == 0u) {
 {% if intMode %}
     if (lid == 0u) { y[i] = {{ scalar }}(0); }
     return;
   }
   var acc = {{ ACC }}(0);
+{% else %}
+  var acc = {{ ACC }}(0);
+{% endif %}
   for (var r = lid; r < REDUCED; r = r + WG) {
 {% set at = "x[input_offset(i, r)]" %}
 {% if castF32 %}
 {% set at = "f32(" ~ at ~ ")" %}
 {% endif %}
+{% if op == "l1" %}
+    acc = combine(acc, abs({{ at }}));
+{% elif op == "l2" or op == "sumsquare" %}
+    let value = {{ at }};
+    acc = combine(acc, value * value);
+{% else %}
     acc = combine(acc, {{ at }});
+{% endif %}
   }
   // Fold the per-thread accumulators. WORKGROUP_SIZE is a power of two, and every
     return;
   }
   let total = partials[0];
+{% if op == "l2" %}
+{% if intMode %}
+  y[i] = {{ scalar }}(sqrt(f32(total)));
+{% else %}
+  y[i] = {{ yv }}sqrt(total){{ vy }};
+{% endif %}
+{% elif op == "mean" %}
 {% if intMode %}
   y[i] = total / {{ scalar }}(REDUCED);
 {% else %}
   y[i] = {{ yv }}total / f32(REDUCED){{ vy }};
 {% endif %}
+{% else %}
+  y[i] = {{ yv }}total{{ vy }};
+{% endif %}
 }

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

@@ -1,13 +1,16 @@
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ reduceWorkgroupSize }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // per-axis dispatch fold width (outputs > 16.7M elements).
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ reduceWorkgroupSize }}u;
-  if (i >= params.count) {
-    return;
-  }
   let v = x[i];
 {% if op == "abs" %}
   y[i] = abs(v);

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

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

@@ -15,18 +15,12 @@
 // The finalizer takes the logarithm of the sum.
 {% endif %}
 // f16 storage widens before accumulation and narrows only at the final store.
-{% set castF32 = castF32 is defined and castF32 %}
 {% set isInt = scalar == "i32" or scalar == "u32" %}
 {% set accScalar = scalar if isInt else "f32" %}
 {% set xv = "vec4<f32>(" if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-enable subgroups;
-{% if usesF16Spec is defined and usesF16Spec %}
-enable f16;
-{% endif %}
-{{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
  * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
@@ -40,15 +34,15 @@ fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
-}
-{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
-{%- if op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op, accScalar) }}
-{%- endif %}
-{%- if op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
@@ -57,7 +51,7 @@ fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
-{%- endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
@@ -110,58 +104,58 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
     y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
   }
-{%- else %}
 {% if op == "max" or op == "min" %}
   let INIT: {{ accScalar }} = reduction_identity();
-{%- elif op == "prod" %}
   let INIT: {{ accScalar }} = {{ "1.0" if not isInt else accScalar ~ "(1)" }};
-{%- else %}
   let INIT: {{ accScalar }} = {{ "0.0" if not isInt else accScalar ~ "(0)" }};
-{%- endif %}
   var acc4 = vec4<{{ accScalar }}>(INIT);
   for (var c = sgLane; c < chunkLimit; c = c + sgSize) {
     let v = {{ xv }}x[base + c]{{ vx }};
-{%- if op == "max" %}
     acc4 = max(acc4, v);
-{%- elif op == "min" %}
     acc4 = min(acc4, v);
-{%- elif op == "prod" %}
     acc4 = acc4 * v;
-{%- elif op == "l1" %}
     acc4 = acc4 + abs(v);
-{%- elif op == "l2" or op == "sumsquare" %}
     acc4 = acc4 + v * v;
-{%- else %}
     acc4 = acc4 + v;
-{%- endif %}
   }
-{%- if op == "max" %}
   let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
   let total = subgroupMax(acc);
-{%- elif op == "min" %}
   let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
   let total = subgroupMin(acc);
-{%- elif op == "prod" %}
   let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
   let total = subgroupMul(acc);
-{%- else %}
   let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
   let total = subgroupAdd(acc);
-{%- endif %}
   if (rowValid && sgLane == 0u) {
-{%- if op == "mean" and isInt %}
     y[row] = total / {{ accScalar }}(params.chunkCount * 4u);
-{%- elif op == "mean" %}
     y[row] = {{ yv }}total / f32(params.chunkCount * 4u){{ vy }};
-{%- elif op == "l2" and isInt %}
     y[row] = {{ accScalar }}(sqrt(f32(total)));
-{%- elif op == "l2" %}
     y[row] = {{ yv }}sqrt(total){{ vy }};
-{%- elif op == "logsum" %}
     y[row] = {{ yv }}log(total){{ vy }};
-{%- else %}
     y[row] = {{ yv }}total{{ vy }};
-{%- endif %}
   }
-{%- endif %}
 }

 // The finalizer takes the logarithm of the sum.
 {% endif %}
 // f16 storage widens before accumulation and narrows only at the final store.
 {% set isInt = scalar == "i32" or scalar == "u32" %}
 {% set accScalar = scalar if isInt else "f32" %}
 {% set xv = "vec4<f32>(" if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
  * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
+}{% endmacro %}
+enable subgroups;
+{{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+{% if op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op, accScalar) }}
+{% elif op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
     let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
     y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
   }
+{% else %}
 {% if op == "max" or op == "min" %}
   let INIT: {{ accScalar }} = reduction_identity();
+{% elif op == "prod" %}
   let INIT: {{ accScalar }} = {{ "1.0" if not isInt else accScalar ~ "(1)" }};
+{% else %}
   let INIT: {{ accScalar }} = {{ "0.0" if not isInt else accScalar ~ "(0)" }};
+{% endif %}
   var acc4 = vec4<{{ accScalar }}>(INIT);
   for (var c = sgLane; c < chunkLimit; c = c + sgSize) {
     let v = {{ xv }}x[base + c]{{ vx }};
+{% if op == "max" %}
     acc4 = max(acc4, v);
+{% elif op == "min" %}
     acc4 = min(acc4, v);
+{% elif op == "prod" %}
     acc4 = acc4 * v;
+{% elif op == "l1" %}
     acc4 = acc4 + abs(v);
+{% elif op == "l2" or op == "sumsquare" %}
     acc4 = acc4 + v * v;
+{% else %}
     acc4 = acc4 + v;
+{% endif %}
   }
+{% if op == "max" %}
   let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
   let total = subgroupMax(acc);
+{% elif op == "min" %}
   let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
   let total = subgroupMin(acc);
+{% elif op == "prod" %}
   let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
   let total = subgroupMul(acc);
+{% else %}
   let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
   let total = subgroupAdd(acc);
+{% endif %}
   if (rowValid && sgLane == 0u) {
+{% if op == "mean" and isInt %}
     y[row] = total / {{ accScalar }}(params.chunkCount * 4u);
+{% elif op == "mean" %}
     y[row] = {{ yv }}total / f32(params.chunkCount * 4u){{ vy }};
+{% elif op == "l2" and isInt %}
     y[row] = {{ accScalar }}(sqrt(f32(total)));
+{% elif op == "l2" %}
     y[row] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
     y[row] = {{ yv }}log(total){{ vy }};
+{% else %}
     y[row] = {{ yv }}total{{ vy }};
+{% endif %}
   }
+{% endif %}
 }

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

@@ -15,17 +15,11 @@
 // IEEE-754 bit patterns because WGSL rejects infinite constants.
 {% endif %}
 // f16 storage widens before accumulation and narrows only at the final store.
-{% set castF32 = castF32 is defined and castF32 %}
 {% set scalar = "f32" if castF32 else scalar %}
 {% set xv = ("vec4<f32>(" if vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-enable subgroups;
-{% if usesF16Spec is defined and usesF16Spec %}
-enable f16;
-{% endif %}
-{{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
  * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
@@ -39,15 +33,15 @@ fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
-}
-{%- endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
-{%- if op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op, scalar) }}
-{%- endif %}
-{%- if op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
@@ -56,7 +50,7 @@ fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
-{%- endif %}
 var<workgroup> wgPartial: array<{{ scalar }}, WG>;
@@ -77,20 +71,19 @@ fn {{ name }}(value: {{ scalar }}, sgLid: u32, sgId: u32, numSg: u32) -> {{ scal
   workgroupBarrier();
   return total;
 }
-{%- endmacro %}
-{%- if op == "max" %}
 {{ emit_reduce("reduce_row", "subgroupMax", "total = max(total, wgPartial[i]);") }}
-{%- elif op == "min" %}
 {{ emit_reduce("reduce_row", "subgroupMin", "total = min(total, wgPartial[i]);") }}
-{%- elif op == "prod" %}
 {{ emit_reduce("reduce_row", "subgroupMul", "total = total * wgPartial[i];") }}
-{%- elif op == "logsumexp" %}
 {{ emit_reduce("reduce_row_add", "subgroupAdd", "total = total + wgPartial[i];") }}
 {{ emit_reduce("reduce_row_max", "subgroupMax", "total = max(total, wgPartial[i]);") }}
-{%- else %}
 {{ emit_reduce("reduce_row", "subgroupAdd", "total = total + wgPartial[i];") }}
-{%- endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
@@ -103,13 +96,13 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   }
   let tid = lid.x;
   let base = row * params.chunkCount;
-{%- if op == "logsumexp" %}
   var localMax = F32_MIN;
   var localNan = 0.0;
   var localNanValue = 0.0;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
-{%- if vec4 %}
 {% for comp in ["x", "y", "z", "w"] %}
     if (is_nan_f32(v.{{ comp }})) {
       localNan = 1.0;
@@ -117,7 +110,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     } else {
       localMax = max(localMax, v.{{ comp }});
     }
-{%- endfor %}
 {% else %}
     if (is_nan_f32(v)) {
       localNan = 1.0;
@@ -125,7 +118,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     } else {
       localMax = max(localMax, v);
     }
-{%- endif %}
   }
   let rowMax = reduce_row_max(localMax, sgLid, sgId, numSg);
   let nanCount = reduce_row_add(localNan, sgLid, sgId, numSg);
@@ -135,83 +128,83 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   var acc = 0.0;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
-{%- if vec4 %}
     let e = select(exp(v - vec4<f32>(rowMax)), vec4<f32>(0.0), hasPositiveInf || hasNan);
     acc = acc + ((e.x + e.y) + (e.z + e.w));
-{%- else %}
     acc = acc + select(exp(v - rowMax), 0.0, hasPositiveInf || hasNan);
-{%- endif %}
   }
   let sum = reduce_row_add(acc, sgLid, sgId, numSg);
   if (tid == 0u) {
     let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
     y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
   }
-{%- else %}
 {% if op == "max" or op == "min" %}
   let INIT: {{ scalar }} = reduction_identity();
-{%- elif op == "prod" %}
   let INIT: f32 = 1.0;
-{%- else %}
   let INIT: f32 = 0.0;
-{%- endif %}
 {% if vec4 %}
   var acc4 = vec4<{{ scalar }}>(INIT);
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
-{%- if op == "max" %}
     acc4 = max(acc4, v);
-{%- elif op == "min" %}
     acc4 = min(acc4, v);
-{%- elif op == "prod" %}
     acc4 = acc4 * v;
-{%- elif op == "l1" %}
     acc4 = acc4 + abs(v);
-{%- elif op == "l2" or op == "sumsquare" %}
     acc4 = acc4 + v * v;
-{%- else %}
     acc4 = acc4 + v;
-{%- endif %}
   }
-{%- if op == "max" %}
   let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
-{%- elif op == "min" %}
   let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
-{%- elif op == "prod" %}
   let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
-{%- else %}
   let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
-{%- endif %}
 {% else %}
   var acc = INIT;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
-{%- if op == "max" %}
     acc = max(acc, v);
-{%- elif op == "min" %}
     acc = min(acc, v);
-{%- elif op == "prod" %}
     acc = acc * v;
-{%- elif op == "l1" %}
     acc = acc + abs(v);
-{%- elif op == "l2" or op == "sumsquare" %}
     acc = acc + v * v;
-{%- else %}
     acc = acc + v;
-{%- endif %}
   }
-{%- endif %}
   let total = reduce_row(acc, sgLid, sgId, numSg);
   if (tid == 0u) {
-{%- if op == "mean" %}
     y[row] = {{ yv }}total / f32(params.cols){{ vy }};
-{%- elif op == "l2" %}
     y[row] = {{ yv }}sqrt(total){{ vy }};
-{%- elif op == "logsum" %}
     y[row] = {{ yv }}log(total){{ vy }};
-{%- else %}
     y[row] = {{ yv }}total{{ vy }};
-{%- endif %}
   }
-{%- endif %}
 }

 // IEEE-754 bit patterns because WGSL rejects infinite constants.
 {% endif %}
 // f16 storage widens before accumulation and narrows only at the final store.
 {% set scalar = "f32" if castF32 else scalar %}
 {% set xv = ("vec4<f32>(" if vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
  * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
+}{% endmacro %}
+enable subgroups;
+{{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+{% if op == "max" or op == "min" %}
 {{ wgsl_minmax_identity("reduction_identity", op, scalar) }}
+{% elif op == "logsumexp" %}
 const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
+{% endif %}
 var<workgroup> wgPartial: array<{{ scalar }}, WG>;
   workgroupBarrier();
   return total;
 }
+{% endmacro %}
+{% if op == "max" %}
 {{ emit_reduce("reduce_row", "subgroupMax", "total = max(total, wgPartial[i]);") }}
+{% elif op == "min" %}
 {{ emit_reduce("reduce_row", "subgroupMin", "total = min(total, wgPartial[i]);") }}
+{% elif op == "prod" %}
 {{ emit_reduce("reduce_row", "subgroupMul", "total = total * wgPartial[i];") }}
+{% elif op == "logsumexp" %}
 {{ emit_reduce("reduce_row_add", "subgroupAdd", "total = total + wgPartial[i];") }}
 {{ emit_reduce("reduce_row_max", "subgroupMax", "total = max(total, wgPartial[i]);") }}
+{% else %}
 {{ emit_reduce("reduce_row", "subgroupAdd", "total = total + wgPartial[i];") }}
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
   }
   let tid = lid.x;
   let base = row * params.chunkCount;
+{% if op == "logsumexp" %}
   var localMax = F32_MIN;
   var localNan = 0.0;
   var localNanValue = 0.0;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
+{% if vec4 %}
 {% for comp in ["x", "y", "z", "w"] %}
     if (is_nan_f32(v.{{ comp }})) {
       localNan = 1.0;
     } else {
       localMax = max(localMax, v.{{ comp }});
     }
+{% endfor %}
 {% else %}
     if (is_nan_f32(v)) {
       localNan = 1.0;
     } else {
       localMax = max(localMax, v);
     }
+{% endif %}
   }
   let rowMax = reduce_row_max(localMax, sgLid, sgId, numSg);
   let nanCount = reduce_row_add(localNan, sgLid, sgId, numSg);
   var acc = 0.0;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
+{% if vec4 %}
     let e = select(exp(v - vec4<f32>(rowMax)), vec4<f32>(0.0), hasPositiveInf || hasNan);
     acc = acc + ((e.x + e.y) + (e.z + e.w));
+{% else %}
     acc = acc + select(exp(v - rowMax), 0.0, hasPositiveInf || hasNan);
+{% endif %}
   }
   let sum = reduce_row_add(acc, sgLid, sgId, numSg);
   if (tid == 0u) {
     let finiteOrInf = select(rowMax + log(sum), rowMax, hasPositiveInf);
     y[row] = {{ yv }}select(finiteOrInf, nanValue, hasNan){{ vy }};
   }
+{% else %}
 {% if op == "max" or op == "min" %}
   let INIT: {{ scalar }} = reduction_identity();
+{% elif op == "prod" %}
   let INIT: f32 = 1.0;
+{% else %}
   let INIT: f32 = 0.0;
+{% endif %}
 {% if vec4 %}
   var acc4 = vec4<{{ scalar }}>(INIT);
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
+{% if op == "max" %}
     acc4 = max(acc4, v);
+{% elif op == "min" %}
     acc4 = min(acc4, v);
+{% elif op == "prod" %}
     acc4 = acc4 * v;
+{% elif op == "l1" %}
     acc4 = acc4 + abs(v);
+{% elif op == "l2" or op == "sumsquare" %}
     acc4 = acc4 + v * v;
+{% else %}
     acc4 = acc4 + v;
+{% endif %}
   }
+{% if op == "max" %}
   let acc = max(max(acc4.x, acc4.y), max(acc4.z, acc4.w));
+{% elif op == "min" %}
   let acc = min(min(acc4.x, acc4.y), min(acc4.z, acc4.w));
+{% elif op == "prod" %}
   let acc = (acc4.x * acc4.y) * (acc4.z * acc4.w);
+{% else %}
   let acc = (acc4.x + acc4.y) + (acc4.z + acc4.w);
+{% endif %}
 {% else %}
   var acc = INIT;
   for (var c = tid; c < params.chunkCount; c = c + WG) {
     let v = {{ xv }}x[base + c]{{ vx }};
+{% if op == "max" %}
     acc = max(acc, v);
+{% elif op == "min" %}
     acc = min(acc, v);
+{% elif op == "prod" %}
     acc = acc * v;
+{% elif op == "l1" %}
     acc = acc + abs(v);
+{% elif op == "l2" or op == "sumsquare" %}
     acc = acc + v * v;
+{% else %}
     acc = acc + v;
+{% endif %}
   }
+{% endif %}
   let total = reduce_row(acc, sgLid, sgId, numSg);
   if (tid == 0u) {
+{% if op == "mean" %}
     y[row] = {{ yv }}total / f32(params.cols){{ vy }};
+{% elif op == "l2" %}
     y[row] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
     y[row] = {{ yv }}log(total){{ vy }};
+{% else %}
     y[row] = {{ yv }}total{{ vy }};
+{% endif %}
   }
+{% endif %}
 }

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

@@ -16,15 +16,11 @@
 // then narrows only at the final store.
 {% set isVec4 = vec4 is defined and vec4 %}
 {% set rowIsEmpty = "params.chunkCount == 0u" if isVec4 else "params.cols == 0u" %}
-{% set castF32 = castF32 is defined and castF32 %}
 {% set scalar = "f32" if castF32 else scalar %}
 {% set xv = ("vec4<f32>(" if vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-{% if usesF16Spec is defined and usesF16Spec %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
@@ -39,15 +35,12 @@ fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
-}
-{%- endmacro %}
 {% if scalar != "i32" and scalar != "u32" and (op == "logsum" or op == "logsumexp") %}
 fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 {% endif %}
 const WG: u32 = {{ workgroupSize }}u;
@@ -57,8 +50,8 @@ const F32_MIN: f32 = -3.4028234663852886e38;
 const F32_MAX: f32 = 3.4028234663852886e38;
 var<workgroup> partial: array<f32, WG>;
-{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
-fn {{ name }}(value: f32, tid: u32) -> f32 {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
@@ -84,20 +77,15 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
-{% if trailingBarrier %}
   workgroupBarrier();
-{% endif %}
   return reduced;
 }
 {% endmacro %}
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
 {{ wgsl_tree_reduce_f32("reduce_max", "max", "partial", "WG") }}
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
-  return (bits & 0x7f800000u) == 0x7f800000u
-      && (bits & 0x007fffffu) != 0u;
 }
 {% else %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
@@ -216,13 +204,14 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
 {% endif %}
     return;
   }
 {% elif op == "logsum" %}
   if ({{ rowIsEmpty }}) {
     if (tid == 0u) { y[row] = {{ yv }}negative_infinity(){{ vy }}; }
     return;
   }
-{% endif %}
 {% if vec4 %}
   var acc4 = vec4<{{ accType }}>(identity());
   for (var col = tid; col < params.chunkCount; col = col + WG) {

 // then narrows only at the final store.
 {% set isVec4 = vec4 is defined and vec4 %}
 {% set rowIsEmpty = "params.chunkCount == 0u" if isVec4 else "params.cols == 0u" %}
 {% set scalar = "f32" if castF32 else scalar %}
 {% set xv = ("vec4<f32>(" if vec4 else "f32(") if castF32 else "" %}
 {% set vx = ")" if castF32 else "" %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
+}{% endmacro %}
 {% if scalar != "i32" and scalar != "u32" and (op == "logsum" or op == "logsumexp") %}
 fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 {% endif %}
 const WG: u32 = {{ workgroupSize }}u;
 const F32_MAX: f32 = 3.4028234663852886e38;
 var<workgroup> partial: array<f32, WG>;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
   workgroupBarrier();
   return reduced;
 }
 {% endmacro %}
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
 {{ wgsl_tree_reduce_f32("reduce_max", "max", "partial", "WG") }}
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 {% else %}
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% endif %}
     return;
   }
 {% elif op == "logsum" %}
   if ({{ rowIsEmpty }}) {
     if (tid == 0u) { y[row] = {{ yv }}negative_infinity(){{ vy }}; }
     return;
   }
+{% endif %}
 {% if vec4 %}
   var acc4 = vec4<{{ accType }}>(identity());
   for (var col = tid; col < params.chunkCount; col = col + WG) {

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

@@ -1,5 +1,4 @@
 {% macro reduce_multi_axis_offset(hasReduced) %}
 // One thread per output element walks the Cartesian product of the reduced axes,
 // linearized as reduce_linear. Specialized shapes make every input offset a sum
 // of coordinate-times-constant terms.
@@ -44,18 +43,20 @@ fn input_offset(out_index: u32{% if hasReduced %}, reduce_linear: u32{% endif %}
 {% set src.value = "(" ~ src.value ~ " * " ~ dataShape[a] ~ "u + coord" ~ a ~ ")" %}
 {% endfor %}
   return {{ src.value }};
-}
-{%- endmacro %}
 // Serial one-thread-per-output reduction for the no-feature tier. f16 storage
 // is widened before every accumulation and narrowed only for the final store.
-{% set castF32 = castF32 is defined and castF32 %}
 {% set logicalBool = logicalBool is defined and logicalBool %}
-{% set intMode = intMode is defined and intMode %}
 {% set yv = "f16(" if castF32 else "" %}
 {% set vy = ")" if castF32 else "" %}
-{% if usesF16Spec is defined and usesF16Spec %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% macro wgsl_minmax_identity(name, op, scalar="f32") %}
 /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
@@ -70,15 +71,12 @@ fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
 {% endif %}
-}
-{%- endmacro %}
 {% if not intMode and (op == "logsum" or op == "logsumexp") %}
 fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 {% endif %}
 {% if op == "max" or op == "min" %}
@@ -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/reduce-strided-axis.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ workgroupSize }})
@@ -15,7 +12,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   var acc = 0.0;
   for (var r = 0u; r < {{ axisDim }}u; r++) {
     let value = {% if usesF16 %}f32({% endif %}x[base + r * {{ innerSize }}u]{% if usesF16 %}){% endif %};
     acc = acc + value;
   }
-  y[i] = {% if usesF16 %}f16({% endif %}acc / {{ axisDim }}.0{% if usesF16 %}){% endif %};
 }

 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ workgroupSize }})
   var acc = 0.0;
   for (var r = 0u; r < {{ axisDim }}u; r++) {
     let value = {% if usesF16 %}f32({% endif %}x[base + r * {{ innerSize }}u]{% if usesF16 %}){% endif %};
+{% if op == "l2" or op == "sumsquare" %}
+    acc = acc + value * value;
+{% elif op == "l1" %}
+    acc = acc + abs(value);
+{% else %}
     acc = acc + value;
+{% endif %}
   }
+  y[i] = {% if usesF16 %}f16({% endif %}{% if op == "l2" %}sqrt(acc){% elif op == "mean" %}acc / {{ axisDim }}.0{% else %}acc{% endif %}{% if usesF16 %}){% endif %};
 }

build/webgpu/test.json CHANGED Viewed

@@ -901,7 +901,7 @@
       "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.0001, "relTolerance": 0.00001 } }
     },
     {
-      "name": "axis0_narrow_f32_8192x3_splitk_guard_lock",
       "provenance": {
         "notes": "An 8,192-by-3 axis-0 mean exercises split-K with a narrow output. Constant ones verify that finalization remains exactly one."
       },
@@ -1179,7 +1179,7 @@
     {
       "name": "int32_lastaxis_subgroup_rows_64x1024",
       "provenance": {
-        "notes": "Route lock for the integer subgroup-per-row last-axis mean: 64 rows of 1024 int32 values sum in the output type and divide as integers, matching the tree kernel bit for bit; a five-value cycle shifts phase every row so the truncated means differ (11 or 12)."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
       "inputs": {

       "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.0001, "relTolerance": 0.00001 } }
     },
     {
+      "name": "axis0_narrow_f32_8192x3_splitk",
       "provenance": {
         "notes": "An 8,192-by-3 axis-0 mean exercises split-K with a narrow output. Constant ones verify that finalization remains exactly one."
       },
     {
       "name": "int32_lastaxis_subgroup_rows_64x1024",
       "provenance": {
+        "notes": "Sixty-four rows of 1024 int32 values check integer mean and truncation; a five-value cycle shifts phase so neighboring row results differ (11 or 12)."
       },
       "attrs": { "axes": [1], "keepdims": 0 },
       "inputs": {