sync 6fdf6301e2bb

Browse files

Files changed (11) hide show

README.md +11 -2
build/webgpu/average-pool2d-nchw-horizontal-reuse.wgsl.jinja +7 -2
build/webgpu/average-pool2d-nchw-w3s1-reuse.wgsl.jinja +9 -4
build/webgpu/bench.json +86 -0
build/webgpu/manifest.json +176 -64
build/webgpu/metadata.json +15 -11
build/webgpu/pool-global-reduction.wgsl.jinja +31 -30
build/webgpu/pool-window-nd.wgsl.jinja +7 -2
build/webgpu/pool-window-reduction.wgsl.jinja +198 -0
build/webgpu/pool-window-unroll.wgsl.jinja +11 -9
build/webgpu/test.json +286 -4

README.md CHANGED Viewed

@@ -48,23 +48,32 @@ Attributes and default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16` |
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`average-pool2d-nchw-horizontal-reuse.wgsl.jinja`](build/webgpu/average-pool2d-nchw-horizontal-reuse.wgsl.jinja)
 - [`average-pool2d-nchw-w3s1-reuse.wgsl.jinja`](build/webgpu/average-pool2d-nchw-w3s1-reuse.wgsl.jinja)
 - [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
 - [`pool-window-nd.wgsl.jinja`](build/webgpu/pool-window-nd.wgsl.jinja)
 - [`pool-window-unroll.wgsl.jinja`](build/webgpu/pool-window-unroll.wgsl.jinja)
 - [`pool2d-nchw-separable.wgsl.jinja`](build/webgpu/pool2d-nchw-separable.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 | --- | --- |
 | `T` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `window_parallel_ncl1d` — Workgroup-cooperative average pooling over a 1-D window: one workgroup reduces a single output element's kernel window, its invocations striding the flattened taps. Chosen when the serial one-invocation-per-output route cannot fill the device and the window is long. Padded taps leave the divisor unless `count_include_pad` is set.
+- `window_parallel_nchw2d` — Workgroup-cooperative average pooling over a 2-D window: one workgroup reduces a single output element's kernel window, its invocations striding the flattened taps. Chosen when the serial one-invocation-per-output route cannot fill the device and the window is long. Padded taps leave the divisor unless `count_include_pad` is set.
+- `window_parallel_ncdhw3d` — Workgroup-cooperative average pooling over a 3-D window: one workgroup reduces a single output element's kernel window, its invocations striding the flattened taps. Chosen when the serial one-invocation-per-output route cannot fill the device and the window is long. Padded taps leave the divisor unless `count_include_pad` is set.
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`average-pool2d-nchw-horizontal-reuse.wgsl.jinja`](build/webgpu/average-pool2d-nchw-horizontal-reuse.wgsl.jinja)
 - [`average-pool2d-nchw-w3s1-reuse.wgsl.jinja`](build/webgpu/average-pool2d-nchw-w3s1-reuse.wgsl.jinja)
 - [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
 - [`pool-window-nd.wgsl.jinja`](build/webgpu/pool-window-nd.wgsl.jinja)
+- [`pool-window-reduction.wgsl.jinja`](build/webgpu/pool-window-reduction.wgsl.jinja)
 - [`pool-window-unroll.wgsl.jinja`](build/webgpu/pool-window-unroll.wgsl.jinja)
 - [`pool2d-nchw-separable.wgsl.jinja`](build/webgpu/pool2d-nchw-separable.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/average-pool2d-nchw-horizontal-reuse.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,9 @@
 {{ env.wgsl.resourceDeclarations }}
 // Four adjacent windows share source columns. For a K-wide, dilation-D window
@@ -22,8 +28,7 @@ fn load_value(validRow: bool, base: u32, col: i32) -> f32 {
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let q = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (q >= COUNT4) { return; }
   let ow0 = (q % OUT_W4) * 4u;
   var t = q / OUT_W4;

+{% 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 }}
 // Four adjacent windows share source columns. For a K-wide, dilation-D window
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "q", "COUNT4", guardInline=true) }}
   let ow0 = (q % OUT_W4) * 4u;
   var t = q / OUT_W4;

build/webgpu/average-pool2d-nchw-w3s1-reuse.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,11 @@
 {{ env.wgsl.resourceDeclarations }}
 // 2D NCHW AveragePool specialised for a 3-wide / stride-1 / pad-1 / dilation-1
@@ -32,10 +40,7 @@ fn col_valid(rowValid: bool, col: i32) -> u32 {
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let tile = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (tile >= params.tileCount) {
-    return;
-  }
   let colTile = tile % OUT_W_Q;     // which group-of-four columns in this row
   var t = tile / OUT_W_Q;

+{% 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 }}
 // 2D NCHW AveragePool specialised for a 3-wide / stride-1 / pad-1 / dilation-1
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "tile", "params.tileCount") }}
   let colTile = tile % OUT_W_Q;     // which group-of-four columns in this row
   var t = tile / OUT_W_Q;

build/webgpu/bench.json CHANGED Viewed

@@ -36,6 +36,17 @@
         "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 2 * 2" }]
       }
     },
     {
       "name": "average-pool-f32-b1c64-224x224-k7s2p3-generic",
       "preset": "smoke",
@@ -198,6 +209,81 @@
       "inputs": { "x": { "dtype": "float32", "shape": [8, 64, 56, 56] } },
       "outputs": { "y": { "dtype": "float32", "shape": [8, 64, 8, 8] } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
     }
   ]
 }

         "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 2 * 2" }]
       }
     },
+    {
+      "name": "nchw_f16_1x128x56x56_k3s1p1",
+      "preset": "model",
+      "vars": { "batch": 1, "channels": 128, "inH": 56, "inW": 56 },
+      "attrs": { "kernel_shape": [3, 3], "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": { "x": { "dtype": "float16", "shape": [1, 128, 56, 56], "dist": "normal", "seed": 207 } },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 128, 56, 56], "dist": "empty" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 2 * 2" }]
+      }
+    },
     {
       "name": "average-pool-f32-b1c64-224x224-k7s2p3-generic",
       "preset": "smoke",
       "inputs": { "x": { "dtype": "float32", "shape": [8, 64, 56, 56] } },
       "outputs": { "y": { "dtype": "float32", "shape": [8, 64, 8, 8] } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "average-pool-f32-b1c32-24x24-k12s1p2-window-parallel-starved",
+      "preset": "stress",
+      "provenance": {
+        "source": "synthetic benchmark",
+        "notes": "A 144-tap padded window produces 9,248 outputs across 32 channels. This brackets the region where output parallelism exceeds the number of taps."
+      },
+      "attrs": { "count_include_pad": 0, "kernel_shape": [12, 12], "strides": [1, 1], "pads": [2, 2, 2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 32, 24, 24], "dist": "normal", "seed": 4301, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 32, 17, 17], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "average-pool-f32-b1c64-24x24-k12s1p2-serial-side-of-gate",
+      "preset": "stress",
+      "provenance": {
+        "source": "synthetic benchmark",
+        "notes": "Measures average pooling with a 12x12 window and padding 2 over 64 channels of a 24x24 input, producing a 17x17 (18,496-element) output."
+      },
+      "attrs": { "count_include_pad": 0, "kernel_shape": [12, 12], "strides": [1, 1], "pads": [2, 2, 2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 64, 24, 24], "dist": "normal", "seed": 4302, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 17, 17], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "average-pool-f32-b1c16-56x56-k28s28-window-parallel-extreme-starved",
+      "preset": "stress",
+      "provenance": {
+        "source": "synthetic benchmark",
+        "notes": "A 784-tap non-overlapping window leaving 64 outputs: the serial route dispatches a single workgroup and loops 784 taps on each of 64 invocations."
+      },
+      "attrs": { "count_include_pad": 0, "kernel_shape": [28, 28], "strides": [28, 28], "pads": [0, 0, 0, 0] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 16, 56, 56], "dist": "normal", "seed": 4303, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 16, 2, 2], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "average-pool-f32-b1c64-14x14-k14s1p6-padded-near-global",
+      "preset": "stress",
+      "provenance": {
+        "source": "synthetic benchmark",
+        "notes": "A padded near-global pool from a classifier head: a 196-tap window over a 14x14 map with pad 6 leaves 10816 outputs, 43 serial dispatch groups."
+      },
+      "attrs": { "count_include_pad": 0, "kernel_shape": [14, 14], "strides": [1, 1], "pads": [6, 6, 6, 6] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 64, 14, 14], "dist": "normal", "seed": 4304, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 13, 13], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "occupancy-window_parallel_nchw2d_count_include_pad_k12p2_small_output",
+      "preset": "stress",
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [2, 2],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 1,
+        "ceil_mode": 0
+      },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 1, 24, 24], "dist": "normal", "seed": 5482, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 9, 9] } }
+    },
+    {
+      "name": "occupancy-window_parallel_outputs_to_taps_boundary_128",
+      "preset": "stress",
+      "attrs": { "kernel_shape": [128], "strides": [1] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 1, 255], "dist": "normal", "seed": 5482, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 128] } }
+    },
+    {
+      "name": "occupancy-window_parallel_outputs_to_taps_boundary_129",
+      "preset": "stress",
+      "attrs": { "kernel_shape": [128], "strides": [1] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 1, 256], "dist": "normal", "seed": 5482, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 129] } }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -20,10 +20,16 @@
     "kernel_shape": { "required": true }
   },
   "typeConstraints": { "T": ["float32", "float16"] },
-  "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "SEPARABLE_MIN_WINDOW": { "default": 25 } },
   "derive": {
     "spatialRank": "ranks.x - 2",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "kernelShapeLengthOk": "not has(attrs, \"kernel_shape\") or (attrs.kernel_shape | length) == spatialRank",
     "stridesLengthOk": "not has(attrs, \"strides\") or (attrs.strides | length) == spatialRank",
     "dilationsLengthOk": "not has(attrs, \"dilations\") or (attrs.dilations | length) == spatialRank",
@@ -43,7 +49,6 @@
     "padBack": "attrs.pads[spatialRank] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0",
     "padBottom": "attrs.pads[2 * spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0",
     "padRight": "attrs.pads[2 * spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0",
-    "poolBaseOk": "f16Ok(dtypes.T) and ranks.x >= 3 and ranks.x <= 5 and ranks.y == ranks.x and has(attrs, \"kernel_shape\") and kernelShapeLengthOk and stridesLengthOk and dilationsLengthOk and padsLengthOk and kernelD > 0 and kernelH > 0 and kernelW > 0 and strideD > 0 and strideH > 0 and strideW > 0 and dilationD > 0 and dilationH > 0 and dilationW > 0 and padFront >= 0 and padTop >= 0 and padLeft >= 0 and padBack >= 0 and padBottom >= 0 and padRight >= 0 and (attrs.auto_pad == \"NOTSET\" or not has(attrs, \"pads\")) and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1)",
     "explicitOutW": "min(ceil((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1, ceil((dim(shapes.x, ranks.x - 1) + padLeft) / strideW)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1",
     "explicitOutH": "min(ceil((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1, ceil((dim(shapes.x, ranks.x - 2) + padTop) / strideH)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1",
     "explicitOutD": "(min(ceil((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1, ceil((dim(shapes.x, 2) + padFront) / strideD)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1) if ranks.x == 5 else 0",
@@ -54,20 +59,78 @@
     "autoPadTotalD": "max(0, (dim(shapes.y, 2) - 1) * strideD + ((kernelD - 1) * dilationD + 1) - dim(shapes.x, 2)) if ranks.x == 5 else 0",
     "autoPadTotalH": "max(0, (dim(shapes.y, ranks.x - 2) - 1) * strideH + ((kernelH - 1) * dilationH + 1) - dim(shapes.x, ranks.x - 2)) if ranks.x >= 4 else 0",
     "autoPadTotalW": "max(0, (dim(shapes.y, ranks.x - 1) - 1) * strideW + ((kernelW - 1) * dilationW + 1) - dim(shapes.x, ranks.x - 1))",
-    "effectivePadFront": "(floor(autoPadTotalD / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalD - floor(autoPadTotalD / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padFront",
-    "effectivePadBack": "autoPadTotalD - effectivePadFront if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBack",
-    "effectivePadTop": "(floor(autoPadTotalH / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalH - floor(autoPadTotalH / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padTop",
-    "effectivePadLeft": "(floor(autoPadTotalW / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalW - floor(autoPadTotalW / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padLeft",
-    "effectivePadBottom": "autoPadTotalH - effectivePadTop if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBottom",
-    "effectivePadRight": "autoPadTotalW - effectivePadLeft if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padRight",
-    "narrowF32W3UnrolledPreferred": "narrowSubgroupRange and dtypes.T == \"f32\"",
-    "globalWindowParallelCovered": "poolBaseOk and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1) and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.auto_pad == \"NOTSET\" and strideD == 1 and strideH == 1 and strideW == 1 and dilationD == 1 and dilationH == 1 and dilationW == 1 and padFront == 0 and padTop == 0 and padLeft == 0 and padBack == 0 and padBottom == 0 and padRight == 0 and ((ranks.x == 3 and kernelW == dim(shapes.x, 2)) or (ranks.x == 4 and kernelH == dim(shapes.x, 2) and kernelW == dim(shapes.x, 3)) or (ranks.x == 5 and kernelD == dim(shapes.x, 2) and kernelH == dim(shapes.x, 3) and kernelW == dim(shapes.x, 4))) and inner(shapes.x, 1) >= 256",
     "separableTmpElements": "dim(shapes.x, 0) * dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.y, 3)",
     "separableFits": "separableTmpElements * 4 <= min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) and ceilDiv(separableTmpElements, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "wholeWindowDenominator": "attrs.ceil_mode == 0 and (attrs.count_include_pad == 1 or (effectivePadTop == 0 and effectivePadLeft == 0 and effectivePadBottom == 0 and effectivePadRight == 0))",
-    "scalar": "dtypes.T"
   },
-  "bindings": { "y_2": { "name": "y", "buffer": "storage", "elementType": "$outputElement" } },
   "variants": [
     {
       "id": "global_window_parallel",
@@ -81,7 +144,6 @@
           "id": "main",
           "name": "AveragePool.GlobalWindowParallel",
           "shader": "pool-global-reduction.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": {
             "op": "\"average\"",
             "vec4": "inner(shapes.x, 1) % 4 == 0",
@@ -97,10 +159,85 @@
         }
       ]
     },
     {
       "id": "nchw2d_k7s1_reuse_tail",
       "priority": 34,
-      "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadTop == 3", "effectivePadBottom == 3", "effectivePadLeft == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
@@ -116,16 +253,11 @@
             "dilationWSpec": 1,
             "padTopSpec": 3,
             "padLeftSpec": 3,
-            "inH": "dim(shapes.x, 2)",
-            "inW": "dim(shapes.x, 3)",
-            "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
-            "outH": "dim(shapes.y, 2)",
-            "outW": "dim(shapes.y, 3)",
             "outW4": "ceilDiv(dim(shapes.y, 3), 4)",
             "count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)",
             "scalarOutput": true
           },
-          "bindings": ["x", "y_2"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -137,8 +269,8 @@
     {
       "id": "nchw2d_w3s1_reuse",
       "priority": 30,
-      "when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadLeft == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadTop + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
-      "demoteWhen": ["narrowF32W3UnrolledPreferred"],
       "passes": [
         {
           "id": "main",
@@ -147,12 +279,7 @@
           "derive": {
             "kernelHSpec": "kernelH",
             "dilationHSpec": "dilationH",
-            "padTopSpec": "effectivePadTop",
-            "inH": "dim(shapes.x, 2)",
-            "inW": "dim(shapes.x, 3)",
-            "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
-            "outH": "dim(shapes.y, 2)",
-            "outW": "dim(shapes.y, 3)",
             "outWQuarter": "floor(dim(shapes.y, 3) / 4)"
           },
           "bindings": [
@@ -171,7 +298,7 @@
     {
       "id": "nchw2d_w3s1_reuse_tail",
       "priority": 29,
-      "when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadLeft == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadTop + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
@@ -185,18 +312,13 @@
             "strideWSpec": 1,
             "dilationHSpec": "dilationH",
             "dilationWSpec": 1,
-            "padTopSpec": "effectivePadTop",
             "padLeftSpec": 1,
-            "inH": "dim(shapes.x, 2)",
-            "inW": "dim(shapes.x, 3)",
-            "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
-            "outH": "dim(shapes.y, 2)",
-            "outW": "dim(shapes.y, 3)",
             "outW4": "ceilDiv(dim(shapes.y, 3), 4)",
             "count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)",
             "scalarOutput": true
           },
-          "bindings": ["x", "y_2"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -209,7 +331,7 @@
       "id": "nchw2d_separable",
       "priority": 22,
       "when": ["poolBaseOk", "ranks.x == 4", "ranks.y == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "strideH < kernelH", "wholeWindowDenominator", "separableFits"],
-      "derive": { "scalar": "dtypes.T", "pIs1": false, "pIs2": false },
       "intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }],
       "passes": [
         {
@@ -219,7 +341,7 @@
           "derive": { "axisSpec": "\"w\"", "op": "\"avg\"", "windowArea": "kernelH * kernelW" },
           "bindings": [
             "x",
-            { "name": "tmp", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
@@ -229,7 +351,7 @@
                 { "name": "kernelW", "type": "u32", "value": "kernelW" },
                 { "name": "strideW", "type": "u32", "value": "strideW" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
-                { "name": "padW", "type": "i32", "value": "effectivePadLeft" },
                 { "name": "tmpCount", "type": "u32", "value": "separableTmpElements" }
               ]
             }
@@ -257,7 +379,7 @@
                 { "name": "kernelH", "type": "u32", "value": "kernelH" },
                 { "name": "strideH", "type": "u32", "value": "strideH" },
                 { "name": "dilationH", "type": "u32", "value": "dilationH" },
-                { "name": "padH", "type": "i32", "value": "effectivePadTop" },
                 { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
               ]
             }
@@ -287,8 +409,8 @@
             "strideWSpec": "strideW",
             "dilationHSpec": "dilationH",
             "dilationWSpec": "dilationW",
-            "padTopSpec": "effectivePadTop",
-            "padLeftSpec": "effectivePadLeft",
             "inH": "dim(shapes.x, 2)",
             "inW": "dim(shapes.x, 3)",
             "outH": "dim(shapes.y, 2)",
@@ -332,7 +454,7 @@
                 { "name": "kernelW", "type": "u32", "value": "kernelW" },
                 { "name": "strideW", "type": "u32", "value": "strideW" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
-                { "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
                 { "name": "padRight", "type": "i32", "value": "effectivePadRight" },
                 { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
                 { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
@@ -353,7 +475,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "nchw2d",
           "shader": "pool-window-nd.wgsl.jinja",
           "derive": { "op": "\"average\"", "spatialRankSpec": 2 },
           "bindings": [
@@ -373,8 +495,8 @@
                 { "name": "strideW", "type": "u32", "value": "strideW" },
                 { "name": "dilationH", "type": "u32", "value": "dilationH" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
-                { "name": "padTop", "type": "i32", "value": "effectivePadTop" },
-                { "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
                 { "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
                 { "name": "padRight", "type": "i32", "value": "effectivePadRight" },
                 { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
@@ -423,9 +545,9 @@
                 { "name": "dilationD", "type": "u32", "value": "dilationD" },
                 { "name": "dilationH", "type": "u32", "value": "dilationH" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
-                { "name": "padFront", "type": "i32", "value": "effectivePadFront" },
-                { "name": "padTop", "type": "i32", "value": "effectivePadTop" },
-                { "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
                 { "name": "padBack", "type": "i32", "value": "effectivePadBack" },
                 { "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
                 { "name": "padRight", "type": "i32", "value": "effectivePadRight" },
@@ -445,7 +567,7 @@
     {
       "id": "nchw2d_k7s2_reuse",
       "priority": 36,
-      "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "effectivePadTop == 3", "effectivePadBottom == 3", "effectivePadLeft == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 3) % 4 == 0", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
@@ -461,15 +583,10 @@
             "dilationWSpec": 1,
             "padTopSpec": 3,
             "padLeftSpec": 3,
-            "inH": "dim(shapes.x, 2)",
-            "inW": "dim(shapes.x, 3)",
-            "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
-            "outH": "dim(shapes.y, 2)",
-            "outW": "dim(shapes.y, 3)",
             "outW4": "dim(shapes.y, 3) / 4",
             "count4": "numel(shapes.y) / 4"
           },
-          "bindings": ["x", "y_2"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -481,7 +598,7 @@
     {
       "id": "nchw2d_k7s1_reuse",
       "priority": 35,
-      "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadTop == 3", "effectivePadBottom == 3", "effectivePadLeft == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
@@ -497,15 +614,10 @@
             "dilationWSpec": 1,
             "padTopSpec": 3,
             "padLeftSpec": 3,
-            "inH": "dim(shapes.x, 2)",
-            "inW": "dim(shapes.x, 3)",
-            "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
-            "outH": "dim(shapes.y, 2)",
-            "outW": "dim(shapes.y, 3)",
             "outW4": "dim(shapes.y, 3) / 4",
             "count4": "numel(shapes.y) / 4"
           },
-          "bindings": ["x", "y_2"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",

     "kernel_shape": { "required": true }
   },
   "typeConstraints": { "T": ["float32", "float16"] },
+  "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "SEPARABLE_MIN_WINDOW": { "default": 25 },
+    "WINDOW_PARALLEL_MIN_TAPS": { "default": 128 },
+    "WINDOW_PARALLEL_MAX_SERIAL_GROUPS": { "default": 64 }
+  },
   "derive": {
     "spatialRank": "ranks.x - 2",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
     "kernelShapeLengthOk": "not has(attrs, \"kernel_shape\") or (attrs.kernel_shape | length) == spatialRank",
     "stridesLengthOk": "not has(attrs, \"strides\") or (attrs.strides | length) == spatialRank",
     "dilationsLengthOk": "not has(attrs, \"dilations\") or (attrs.dilations | length) == spatialRank",
     "padBack": "attrs.pads[spatialRank] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0",
     "padBottom": "attrs.pads[2 * spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0",
     "padRight": "attrs.pads[2 * spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0",
     "explicitOutW": "min(ceil((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1, ceil((dim(shapes.x, ranks.x - 1) + padLeft) / strideW)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1",
     "explicitOutH": "min(ceil((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1, ceil((dim(shapes.x, ranks.x - 2) + padTop) / strideH)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1",
     "explicitOutD": "(min(ceil((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1, ceil((dim(shapes.x, 2) + padFront) / strideD)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1) if ranks.x == 5 else 0",
     "autoPadTotalD": "max(0, (dim(shapes.y, 2) - 1) * strideD + ((kernelD - 1) * dilationD + 1) - dim(shapes.x, 2)) if ranks.x == 5 else 0",
     "autoPadTotalH": "max(0, (dim(shapes.y, ranks.x - 2) - 1) * strideH + ((kernelH - 1) * dilationH + 1) - dim(shapes.x, ranks.x - 2)) if ranks.x >= 4 else 0",
     "autoPadTotalW": "max(0, (dim(shapes.y, ranks.x - 1) - 1) * strideW + ((kernelW - 1) * dilationW + 1) - dim(shapes.x, ranks.x - 1))",
+    "effectivePadD": "(floor(autoPadTotalD / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalD - floor(autoPadTotalD / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padFront",
+    "effectivePadH": "(floor(autoPadTotalH / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalH - floor(autoPadTotalH / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padTop",
+    "effectivePadW": "(floor(autoPadTotalW / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalW - floor(autoPadTotalW / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padLeft",
+    "windowTaps": "kernelD * kernelH * kernelW",
+    "windowParallelEligible": "numel(shapes.y) > 0 and windowTaps >= tunables.WINDOW_PARALLEL_MIN_TAPS and ceilDiv(numel(shapes.y), tunables.WORKGROUP_SIZE) < tunables.WINDOW_PARALLEL_MAX_SERIAL_GROUPS and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and numel(shapes.y) <= windowTaps",
     "separableTmpElements": "dim(shapes.x, 0) * dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.y, 3)",
     "separableFits": "separableTmpElements * 4 <= min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) and ceilDiv(separableTmpElements, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "poolBaseOk": "f16Ok(dtypes.T) and ranks.x >= 3 and ranks.x <= 5 and ranks.y == ranks.x and has(attrs, \"kernel_shape\") and kernelShapeLengthOk and stridesLengthOk and dilationsLengthOk and padsLengthOk and kernelD > 0 and kernelH > 0 and kernelW > 0 and strideD > 0 and strideH > 0 and strideW > 0 and dilationD > 0 and dilationH > 0 and dilationW > 0 and padFront >= 0 and padTop >= 0 and padLeft >= 0 and padBack >= 0 and padBottom >= 0 and padRight >= 0 and (attrs.auto_pad == \"NOTSET\" or not has(attrs, \"pads\")) and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1)",
+    "effectivePadBack": "autoPadTotalD - effectivePadD if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBack",
+    "effectivePadBottom": "autoPadTotalH - effectivePadH if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBottom",
+    "effectivePadRight": "autoPadTotalW - effectivePadW if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padRight",
+    "narrowF32W3UnrolledPreferred": "narrowSubgroupRange and dtypes.T == \"f32\"",
+    "globalWindowParallelCovered": "poolBaseOk and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1) and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.auto_pad == \"NOTSET\" and strideD == 1 and strideH == 1 and strideW == 1 and dilationD == 1 and dilationH == 1 and dilationW == 1 and padFront == 0 and padTop == 0 and padLeft == 0 and padBack == 0 and padBottom == 0 and padRight == 0 and ((ranks.x == 3 and kernelW == dim(shapes.x, 2)) or (ranks.x == 4 and kernelH == dim(shapes.x, 2) and kernelW == dim(shapes.x, 3)) or (ranks.x == 5 and kernelD == dim(shapes.x, 2) and kernelH == dim(shapes.x, 3) and kernelW == dim(shapes.x, 4))) and inner(shapes.x, 1) >= 256",
+    "wholeWindowDenominator": "attrs.ceil_mode == 0 and (attrs.count_include_pad == 1 or (effectivePadH == 0 and effectivePadW == 0 and effectivePadBottom == 0 and effectivePadRight == 0))",
+    "scalar": "dtypes.T",
+    "inH": "dim(shapes.x, 2)",
+    "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
+    "inW": "dim(shapes.x, 3)",
+    "outH": "dim(shapes.y, 2)",
+    "outW": "dim(shapes.y, 3)"
+  },
+  "bindings": {
+    "y_main": { "name": "y", "elementType": "$outputElement" },
+    "params_window_ncl1d": {
+      "name": "params",
+      "struct": [
+        { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "inW", "type": "u32", "value": "dim(shapes.x, 2)" },
+        { "name": "outW", "type": "u32", "value": "dim(shapes.y, 2)" },
+        { "name": "strideW", "type": "u32", "value": "strideW" },
+        { "name": "dilationW", "type": "u32", "value": "dilationW" },
+        { "name": "padW", "type": "i32", "value": "effectivePadW" }
+      ]
+    },
+    "params_window_nchw2d": {
+      "name": "params",
+      "struct": [
+        { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
+        { "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" },
+        { "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" },
+        { "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" },
+        { "name": "strideH", "type": "u32", "value": "strideH" },
+        { "name": "strideW", "type": "u32", "value": "strideW" },
+        { "name": "dilationH", "type": "u32", "value": "dilationH" },
+        { "name": "dilationW", "type": "u32", "value": "dilationW" },
+        { "name": "padH", "type": "i32", "value": "effectivePadH" },
+        { "name": "padW", "type": "i32", "value": "effectivePadW" }
+      ]
+    },
+    "params_window_ncdhw3d": {
+      "name": "params",
+      "struct": [
+        { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "inD", "type": "u32", "value": "dim(shapes.x, 2)" },
+        { "name": "inH", "type": "u32", "value": "dim(shapes.x, 3)" },
+        { "name": "inW", "type": "u32", "value": "dim(shapes.x, 4)" },
+        { "name": "outD", "type": "u32", "value": "dim(shapes.y, 2)" },
+        { "name": "outH", "type": "u32", "value": "dim(shapes.y, 3)" },
+        { "name": "outW", "type": "u32", "value": "dim(shapes.y, 4)" },
+        { "name": "strideD", "type": "u32", "value": "strideD" },
+        { "name": "strideH", "type": "u32", "value": "strideH" },
+        { "name": "strideW", "type": "u32", "value": "strideW" },
+        { "name": "dilationD", "type": "u32", "value": "dilationD" },
+        { "name": "dilationH", "type": "u32", "value": "dilationH" },
+        { "name": "dilationW", "type": "u32", "value": "dilationW" },
+        { "name": "padD", "type": "i32", "value": "effectivePadD" },
+        { "name": "padH", "type": "i32", "value": "effectivePadH" },
+        { "name": "padW", "type": "i32", "value": "effectivePadW" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "global_window_parallel",
           "id": "main",
           "name": "AveragePool.GlobalWindowParallel",
           "shader": "pool-global-reduction.wgsl.jinja",
           "derive": {
             "op": "\"average\"",
             "vec4": "inner(shapes.x, 1) % 4 == 0",
         }
       ]
     },
+    {
+      "id": "window_parallel_ncl1d",
+      "priority": 28,
+      "when": ["poolBaseOk", "ranks.x == 3", "dim(shapes.y, 2) == inferredOutW", "attrs.count_include_pad == 0 or attrs.ceil_mode == 0", "windowParallelEligible"],
+      "supersededBy": ["global_window_parallel"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "AveragePool.WindowParallelNcl1d",
+          "shader": "pool-window-reduction.wgsl.jinja",
+          "derive": {
+            "op": "\"average\"",
+            "spatialRankSpec": 1,
+            "kernelDSpec": "kernelD",
+            "kernelHSpec": "kernelH",
+            "kernelWSpec": "kernelW",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
+            "useSubgroups": "device.features.has(\"subgroups\")",
+            "countIncludePad": "attrs.count_include_pad != 0"
+          },
+          "bindings": ["x", "y", "params_window_ncl1d"],
+          "dispatch": { "x": "numel(shapes.y)" }
+        }
+      ]
+    },
+    {
+      "id": "window_parallel_nchw2d",
+      "priority": 28,
+      "when": ["poolBaseOk", "ranks.x == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "attrs.count_include_pad == 0 or attrs.ceil_mode == 0", "windowParallelEligible"],
+      "supersededBy": ["global_window_parallel"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "AveragePool.WindowParallelNchw2d",
+          "shader": "pool-window-reduction.wgsl.jinja",
+          "derive": {
+            "op": "\"average\"",
+            "spatialRankSpec": 2,
+            "kernelDSpec": "kernelD",
+            "kernelHSpec": "kernelH",
+            "kernelWSpec": "kernelW",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
+            "useSubgroups": "device.features.has(\"subgroups\")",
+            "countIncludePad": "attrs.count_include_pad != 0"
+          },
+          "bindings": ["x", "y", "params_window_nchw2d"],
+          "dispatch": { "x": "numel(shapes.y)" }
+        }
+      ]
+    },
+    {
+      "id": "window_parallel_ncdhw3d",
+      "priority": 28,
+      "when": ["poolBaseOk", "ranks.x == 5", "dim(shapes.y, 2) == inferredOutD", "dim(shapes.y, 3) == inferredOutH", "dim(shapes.y, 4) == inferredOutW", "attrs.count_include_pad == 0 or attrs.ceil_mode == 0", "windowParallelEligible"],
+      "supersededBy": ["global_window_parallel"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "AveragePool.WindowParallelNcdhw3d",
+          "shader": "pool-window-reduction.wgsl.jinja",
+          "derive": {
+            "op": "\"average\"",
+            "spatialRankSpec": 3,
+            "kernelDSpec": "kernelD",
+            "kernelHSpec": "kernelH",
+            "kernelWSpec": "kernelW",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
+            "useSubgroups": "device.features.has(\"subgroups\")",
+            "countIncludePad": "attrs.count_include_pad != 0"
+          },
+          "bindings": ["x", "y", "params_window_ncdhw3d"],
+          "dispatch": { "x": "numel(shapes.y)" }
+        }
+      ]
+    },
     {
       "id": "nchw2d_k7s1_reuse_tail",
       "priority": 34,
+      "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadH == 3", "effectivePadBottom == 3", "effectivePadW == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
             "dilationWSpec": 1,
             "padTopSpec": 3,
             "padLeftSpec": 3,
             "outW4": "ceilDiv(dim(shapes.y, 3), 4)",
             "count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)",
             "scalarOutput": true
           },
+          "bindings": ["x", "y_main"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
     {
       "id": "nchw2d_w3s1_reuse",
       "priority": 30,
+      "when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadW == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadH + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
+      "demoteWhen": ["narrowF32W3UnrolledPreferred", "variableSubgroup16To32 and dtypes.T == \"f16\""],
       "passes": [
         {
           "id": "main",
           "derive": {
             "kernelHSpec": "kernelH",
             "dilationHSpec": "dilationH",
+            "padTopSpec": "effectivePadH",
             "outWQuarter": "floor(dim(shapes.y, 3) / 4)"
           },
           "bindings": [
     {
       "id": "nchw2d_w3s1_reuse_tail",
       "priority": 29,
+      "when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadW == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadH + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
             "strideWSpec": 1,
             "dilationHSpec": "dilationH",
             "dilationWSpec": 1,
+            "padTopSpec": "effectivePadH",
             "padLeftSpec": 1,
             "outW4": "ceilDiv(dim(shapes.y, 3), 4)",
             "count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)",
             "scalarOutput": true
           },
+          "bindings": ["x", "y_main"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
       "id": "nchw2d_separable",
       "priority": 22,
       "when": ["poolBaseOk", "ranks.x == 4", "ranks.y == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "strideH < kernelH", "wholeWindowDenominator", "separableFits"],
+      "derive": { "scalar": "dtypes.T" },
       "intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }],
       "passes": [
         {
           "derive": { "axisSpec": "\"w\"", "op": "\"avg\"", "windowArea": "kernelH * kernelW" },
           "bindings": [
             "x",
+            { "name": "tmp", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
                 { "name": "kernelW", "type": "u32", "value": "kernelW" },
                 { "name": "strideW", "type": "u32", "value": "strideW" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
+                { "name": "padW", "type": "i32", "value": "effectivePadW" },
                 { "name": "tmpCount", "type": "u32", "value": "separableTmpElements" }
               ]
             }
                 { "name": "kernelH", "type": "u32", "value": "kernelH" },
                 { "name": "strideH", "type": "u32", "value": "strideH" },
                 { "name": "dilationH", "type": "u32", "value": "dilationH" },
+                { "name": "padH", "type": "i32", "value": "effectivePadH" },
                 { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
               ]
             }
             "strideWSpec": "strideW",
             "dilationHSpec": "dilationH",
             "dilationWSpec": "dilationW",
+            "padTopSpec": "effectivePadH",
+            "padLeftSpec": "effectivePadW",
             "inH": "dim(shapes.x, 2)",
             "inW": "dim(shapes.x, 3)",
             "outH": "dim(shapes.y, 2)",
                 { "name": "kernelW", "type": "u32", "value": "kernelW" },
                 { "name": "strideW", "type": "u32", "value": "strideW" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
+                { "name": "padLeft", "type": "i32", "value": "effectivePadW" },
                 { "name": "padRight", "type": "i32", "value": "effectivePadRight" },
                 { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
                 { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
       "passes": [
         {
           "id": "main",
+          "name": "AveragePool.Nchw2d",
           "shader": "pool-window-nd.wgsl.jinja",
           "derive": { "op": "\"average\"", "spatialRankSpec": 2 },
           "bindings": [
                 { "name": "strideW", "type": "u32", "value": "strideW" },
                 { "name": "dilationH", "type": "u32", "value": "dilationH" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
+                { "name": "padTop", "type": "i32", "value": "effectivePadH" },
+                { "name": "padLeft", "type": "i32", "value": "effectivePadW" },
                 { "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
                 { "name": "padRight", "type": "i32", "value": "effectivePadRight" },
                 { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
                 { "name": "dilationD", "type": "u32", "value": "dilationD" },
                 { "name": "dilationH", "type": "u32", "value": "dilationH" },
                 { "name": "dilationW", "type": "u32", "value": "dilationW" },
+                { "name": "padFront", "type": "i32", "value": "effectivePadD" },
+                { "name": "padTop", "type": "i32", "value": "effectivePadH" },
+                { "name": "padLeft", "type": "i32", "value": "effectivePadW" },
                 { "name": "padBack", "type": "i32", "value": "effectivePadBack" },
                 { "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
                 { "name": "padRight", "type": "i32", "value": "effectivePadRight" },
     {
       "id": "nchw2d_k7s2_reuse",
       "priority": 36,
+      "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "effectivePadH == 3", "effectivePadBottom == 3", "effectivePadW == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 3) % 4 == 0", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
             "dilationWSpec": 1,
             "padTopSpec": 3,
             "padLeftSpec": 3,
             "outW4": "dim(shapes.y, 3) / 4",
             "count4": "numel(shapes.y) / 4"
           },
+          "bindings": ["x", "y_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
     {
       "id": "nchw2d_k7s1_reuse",
       "priority": 35,
+      "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadH == 3", "effectivePadBottom == 3", "effectivePadW == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
       "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
             "dilationWSpec": 1,
             "padTopSpec": 3,
             "padLeftSpec": 3,
             "outW4": "dim(shapes.y, 3) / 4",
             "count4": "numel(shapes.y) / 4"
           },
+          "bindings": ["x", "y_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,28 +1,32 @@
 {
   "name": "ai.onnx.AveragePool",
-  "id": "_ai_onnx_averagepool_webgpu_1a52fa9",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "average-pool2d-nchw-horizontal-reuse.wgsl.jinja": "WhpEal2kC3XudqPetAZeL6v/IGK1olc+Fb4HzzDj890=",
-      "average-pool2d-nchw-w3s1-reuse.wgsl.jinja": "IsI2gMcZbY0eemYs5z0pTC86cgaDw605Xfn2iQDBJkw=",
-      "bench.json": "7vXmd7ZA5xBCZjqXJYtMzO7vq/0qgzf5gNZ1S6J5Zps=",
-      "manifest.json": "a0sXaOaq8613yIAtQBB1R+xrbvA48WTKgkPo14ruXhA=",
-      "pool-global-reduction.wgsl.jinja": "SKnTSUq9nU838U0t7VyvC7hgz4h5L+Ndv1rQCM9Vks8=",
-      "pool-window-nd.wgsl.jinja": "2shuAamdoFjQCOq2LK2XC6pKsa10zad6bDhqY1h2zBI=",
-      "pool-window-unroll.wgsl.jinja": "iSbZ3lQkW4r/CZzsBitQBpHgBLRJ83MekQxyVqbInwI=",
       "pool2d-nchw-separable.wgsl.jinja": "GgafQWV6Pn5SqSn5gMns0EJEHGqk2WkdCdVurt3MgA8=",
-      "test.json": "2Xnd9YqkiREFwBJdjyqFIbqnT2KGPNdXYpCezSLDXkA="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "global_window_parallel": ["pool-global-reduction.wgsl.jinja"],
       "nchw2d_k7s1_reuse_tail": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],
       "nchw2d_w3s1_reuse": ["average-pool2d-nchw-w3s1-reuse.wgsl.jinja"],
       "nchw2d_w3s1_reuse_tail": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],

 {
   "name": "ai.onnx.AveragePool",
+  "id": "_ai_onnx_averagepool_webgpu_6441176",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "average-pool2d-nchw-horizontal-reuse.wgsl.jinja": "heeYKFbnmrXdEohH7tUG5PzItrk5yPz3hEgnuNihBWg=",
+      "average-pool2d-nchw-w3s1-reuse.wgsl.jinja": "D4razLs8qLkk5zubCU+wUa9OVgFisNZ95cGwXipOvwI=",
+      "bench.json": "ZvF9gqBTship90W+TiMfWNyrdsJYyEAwZAdZ5r6/vc0=",
+      "manifest.json": "/fTFl6NMtKYnbxuSbihY039sqo+P8xhGXVjFtj2gA6c=",
+      "pool-global-reduction.wgsl.jinja": "uCrSccgj/Fibd866kxpiYkprR8g6q1vofgQZZ/oxGkE=",
+      "pool-window-nd.wgsl.jinja": "QCtdlEJy4LmSU9g7oUOR0YbZ9cZuxH6P1ERT3m7YACY=",
+      "pool-window-reduction.wgsl.jinja": "Jv5CY3dfoWUzgoxb4H35nlKspu9GKe10Z1bRjGQ4DPk=",
+      "pool-window-unroll.wgsl.jinja": "UWhk8iANx6euYlJV5W31BP0D4FaJEC+lV/3y95T2wOQ=",
       "pool2d-nchw-separable.wgsl.jinja": "GgafQWV6Pn5SqSn5gMns0EJEHGqk2WkdCdVurt3MgA8=",
+      "test.json": "zw4aG7ZsPYmGpRG2g8Yi++TwV5Fvas/zLR5Gesmxys8="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "global_window_parallel": ["pool-global-reduction.wgsl.jinja"],
+      "window_parallel_ncl1d": ["pool-window-reduction.wgsl.jinja"],
+      "window_parallel_nchw2d": ["pool-window-reduction.wgsl.jinja"],
+      "window_parallel_ncdhw3d": ["pool-window-reduction.wgsl.jinja"],
       "nchw2d_k7s1_reuse_tail": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],
       "nchw2d_w3s1_reuse": ["average-pool2d-nchw-w3s1-reuse.wgsl.jinja"],
       "nchw2d_w3s1_reuse_tail": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],

build/webgpu/pool-global-reduction.wgsl.jinja CHANGED Viewed

@@ -22,13 +22,11 @@ var<workgroup> wg_partials: array<f32, WG>;
 var<workgroup> wg_tree: array<f32, WG>;
 {% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
-        @builtin(local_invocation_id) lid: vec3<u32>
-{%- if useSubgroups %},
-        @builtin(subgroup_size) sg_size: u32
-{%- endif %}) {
   let plane = wid.x;
   let tid = lid.x;
   var partial = 0.0;
@@ -50,32 +48,35 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
 {% endif %}
 {% if useSubgroups %}
-  var total = subgroupAdd(partial);
-  // Cross-subgroup fold that assumes nothing about which invocations share a
-  // subgroup, how many subgroups there are, or which of a subgroup's lanes are
-  // active: every invocation owns the slot at its own index, the elected lane
-  // publishes its subgroup partial there and every other lane publishes the
-  // reduction identity. Each subgroup then folds all WG slots — lane `rank`,
-  // its dense position among the active lanes, walks slots rank, rank + count,
-  // ... — and one more collective merges the lane partials, so every slot is
-  // merged exactly once at any legal width and partition. A one-subgroup
-  // workgroup skips shared memory entirely; that test reads the `subgroup_size`
-  // builtin, which is uniform, because a collective's result is not uniform to
-  // WGSL's analysis and may not guard a barrier.
-  if (sg_size != WG) {
-    let rank = subgroupExclusiveAdd(1u);
-    let count = subgroupAdd(1u);
-    wg_partials[tid] = select(0.0, total, rank == 0u);
-    workgroupBarrier();
-    var acc = 0.0;
-    for (var i = rank; i < WG; i = i + count) {
-      acc = acc + wg_partials[i];
-    }
-    total = subgroupAdd(acc);
   }
   if (tid == 0u) {
-    y[plane] = {{ scalar }}(total / f32(params.spatial));
-  }
 {% else %}
   // No-subgroups fallback: full workgroup-tree reduction of the per-thread partials.
   wg_tree[tid] = partial;
@@ -88,7 +89,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   }
   if (tid == 0u) {
     let total = wg_tree[0];
     y[plane] = {{ scalar }}(total / f32(params.spatial));
   }
-{% endif %}
 }

 var<workgroup> wg_tree: array<f32, WG>;
 {% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+        @builtin(subgroup_invocation_id) sg_lane: u32,
+        @builtin(subgroup_size) sg_size: u32{% endif %}) {
   let plane = wid.x;
   let tid = lid.x;
   var partial = 0.0;
 {% endif %}
 {% if useSubgroups %}
+{% filter indent(2, true) %}
+var total = subgroupAdd(partial);
+// Cross-subgroup fold that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are (it does require the uniform
+// control flow this entry point already has): every invocation owns the slot
+// at its own index, the elected lane
+// publishes its subgroup partial there and every other lane publishes the
+// reduction identity. Each subgroup then folds all WG slots — lane `sg_lane`
+// walks slots sg_lane, sg_lane + sg_size, ... — and one more collective merges
+// the lane partials, so every slot is merged exactly once at any legal width.
+// The coordinates are the BUILTINS, never `subgroupExclusiveAdd(1u)` /
+// `subgroupAdd(1u)`: those agree with them under this contract, but a driver
+// in the wild answers a claim derived from them with zero. A one-subgroup
+// workgroup skips shared memory entirely; that test reads the `subgroup_size`
+// builtin, which is uniform, because a collective's result is not uniform to
+// WGSL's analysis and may not guard a barrier.
+if (sg_size != WG) {
+  let rank = sg_lane;
+  let lanes = sg_size;
+  wg_partials[tid] = select(0.0, total, rank == 0u);
+  workgroupBarrier();
+  var acc = 0.0;
+  for (var i = rank; i < WG; i = i + lanes) {
+    acc = acc + wg_partials[i];
   }
+  total = subgroupAdd(acc);
+}
+{% endfilter %}
   if (tid == 0u) {
 {% else %}
   // No-subgroups fallback: full workgroup-tree reduction of the per-thread partials.
   wg_tree[tid] = partial;
   }
   if (tid == 0u) {
     let total = wg_tree[0];
+{% endif %}
     y[plane] = {{ scalar }}(total / f32(params.spatial));
   }
 }

build/webgpu/pool-window-nd.wgsl.jinja CHANGED Viewed

@@ -1,10 +1,15 @@
 {{ env.wgsl.resourceDeclarations }}
 // Direct spatial-pooling traversal for a channels-first tensor.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (index >= params.count) { return; }
   let ow = index % params.outW;
   var t = index / params.outW;

+{% 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 }}
 // Direct spatial-pooling traversal for a channels-first tensor.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "index", guardInline=true) }}
   let ow = index % params.outW;
   var t = index / params.outW;

build/webgpu/pool-window-reduction.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,198 @@

+{% set needsCount = op == "max" or (op == "average" and not countIncludePad) %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Workgroup-cooperative windowed pooling: one workgroup reduces the kernel
+// window of ONE output element. The serial route gives each output element a
+// single invocation that loops the whole window, so its only parallelism is the
+// output element count; when that count launches too few workgroups and the window is
+// long, the window itself is the parallelism worth spending. Threads here stride
+// the flattened tap index [kD x] kH x kW, each tap resolving its own padded,
+// strided and dilated input coordinate, then combine the per-thread partials.
+// With subgroups: a subgroup collective, then one shared-memory slot per
+// invocation (the subgroup partial from its elected lane, the identity from
+// every other lane) that every subgroup folds. Without subgroups, a full
+// workgroup tree combines the per-thread partials.
+{% if countIncludePad %}
+// count_include_pad == 1 with ceil_mode == 0 makes every tap of every window
+// land inside the padded extent, so the divisor is the baked tap count and no
+// second reduction is needed. The tree changes the association of additions.
+{% else %}
+// count_include_pad == 0 divides by the in-bounds tap count, which is reduced
+// across the workgroup alongside the sum. An output whose window is entirely
+// out of bounds writes 0, matching the serial route. The tree changes the
+// association of additions.
+{% endif %}
+const WG: u32 = {{ workgroupSizeSpec }}u;
+const KW: u32 = {{ kernelWSpec }}u;
+{% if spatialRankSpec >= 2 %}
+const KH: u32 = {{ kernelHSpec }}u;
+{% endif %}
+// Flattened tap count of the whole window; the loop bound and, for
+// count_include_pad == 1, the divisor.
+const TAPS: u32 = {{ kernelDSpec * kernelHSpec * kernelWSpec }}u;
+// One slot per invocation: the cross-subgroup fold below indexes by
+// local_invocation_id, never by a subgroup ordinal, so no subgroup width or
+// partition can overflow or alias it.
+var<workgroup> wg_partials: array<f32, WG>;
+{% if needsCount %}
+var<workgroup> wg_counts: array<f32, WG>;
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+        @builtin(subgroup_invocation_id) sg_lane: u32,
+        @builtin(subgroup_size) sg_size: u32{% endif %}) {
+  // One workgroup per output element: the dispatch is exactly numel(y), so no
+  // bounds guard is needed here.
+  let index = wid.x;
+  let tid = lid.x;
+  let ow = index % params.outW;
+  var t = index / params.outW;
+{% if spatialRankSpec >= 2 %}
+  let oh = t % params.outH;
+  t /= params.outH;
+{% endif %}
+{% if spatialRankSpec == 3 %}
+  let od = t % params.outD;
+  t /= params.outD;
+{% endif %}
+  let channel = t % params.channels;
+  let batch = t / params.channels;
+  let startW = i32(ow * params.strideW) - params.padW;
+{% if spatialRankSpec >= 2 %}
+  let startH = i32(oh * params.strideH) - params.padH;
+{% endif %}
+{% if spatialRankSpec == 3 %}
+  let startD = i32(od * params.strideD) - params.padD;
+{% endif %}
+{% if spatialRankSpec == 1 %}
+  let planeBase = (batch * params.channels + channel) * params.inW;
+{% elif spatialRankSpec == 2 %}
+  let planeBase = (batch * params.channels + channel) * params.inH * params.inW;
+{% else %}
+  let planeBase = (batch * params.channels + channel) * params.inD * params.inH * params.inW;
+{% endif %}
+  var partial = 0.0;
+{% if needsCount %}
+  var counted = 0.0;
+{% endif %}
+  for (var tap = tid; tap < TAPS; tap = tap + WG) {
+    let kw = tap % KW;
+{% if spatialRankSpec >= 2 %}
+    let rest = tap / KW;
+    let kh = rest % KH;
+{% endif %}
+{% if spatialRankSpec == 3 %}
+    let kd = rest / KH;
+{% endif %}
+    let iw = startW + i32(kw * params.dilationW);
+{% if spatialRankSpec >= 2 %}
+    let ih = startH + i32(kh * params.dilationH);
+{% endif %}
+{% if spatialRankSpec == 3 %}
+    let id = startD + i32(kd * params.dilationD);
+{% endif %}
+{% if spatialRankSpec == 1 %}
+    let inBounds = iw >= 0 && iw < i32(params.inW);
+{% elif spatialRankSpec == 2 %}
+    let inBounds = iw >= 0 && iw < i32(params.inW) && ih >= 0 && ih < i32(params.inH);
+{% else %}
+    let inBounds = iw >= 0 && iw < i32(params.inW) && ih >= 0 && ih < i32(params.inH) && id >= 0 && id < i32(params.inD);
+{% endif %}
+    if (inBounds) {
+{% if spatialRankSpec == 1 %}
+      let value = f32(x[planeBase + u32(iw)]);
+{% elif spatialRankSpec == 2 %}
+      let value = f32(x[planeBase + u32(ih) * params.inW + u32(iw)]);
+{% else %}
+      let value = f32(x[planeBase + (u32(id) * params.inH + u32(ih)) * params.inW + u32(iw)]);
+{% endif %}
+      partial = partial + value;
+{% if needsCount %}
+      counted = counted + 1.0;
+{% endif %}
+    }
+  }
+{% if useSubgroups %}
+{% filter indent(2, true) %}
+var total = subgroupAdd(partial);
+{% if needsCount %}
+var countTotal = subgroupAdd(counted);
+{% endif %}
+// Cross-subgroup fold that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are (it does require the uniform
+// control flow this entry point already has): every invocation owns the slot
+// at its own index, the elected lane
+// publishes its subgroup partial there and every other lane publishes the
+// reduction identity. Each subgroup then folds all WG slots — lane `sg_lane`
+// walks slots sg_lane, sg_lane + sg_size, ... — and one more collective merges
+// the lane partials, so every slot is merged exactly once at any legal width.
+// The coordinates are the BUILTINS, never `subgroupExclusiveAdd(1u)` /
+// `subgroupAdd(1u)`: those agree with them under this contract, but a driver
+// in the wild answers a claim derived from them with zero. A one-subgroup
+// workgroup skips shared memory entirely; that test reads the `subgroup_size`
+// builtin, which is uniform, because a collective's result is not uniform to
+// WGSL's analysis and may not guard a barrier.
+if (sg_size != WG) {
+  let rank = sg_lane;
+  let lanes = sg_size;
+  wg_partials[tid] = select(0.0, total, rank == 0u);
+{% if needsCount %}
+  wg_counts[tid] = select(0.0, countTotal, rank == 0u);
+{% endif %}
+  workgroupBarrier();
+  var acc = 0.0;
+{% if needsCount %}
+  var countAcc = 0.0;
+{% endif %}
+  for (var i = rank; i < WG; i = i + lanes) {
+    acc = acc + wg_partials[i];
+{% if needsCount %}
+    countAcc = countAcc + wg_counts[i];
+{% endif %}
+  }
+  total = subgroupAdd(acc);
+{% if needsCount %}
+  countTotal = subgroupAdd(countAcc);
+{% endif %}
+}
+{% endfilter %}
+{% else %}
+  // No-subgroups fallback: full workgroup-tree reduction of the per-thread
+  // partials, carrying the unpaired entry forward for odd active widths.
+  wg_partials[tid] = partial;
+{% if needsCount %}
+  wg_counts[tid] = counted;
+{% endif %}
+  workgroupBarrier();
+  for (var liveCount = WG; liveCount > 1u; liveCount = (liveCount + 1u) / 2u) {
+    let step = (liveCount + 1u) / 2u;
+    if (tid + step < liveCount) {
+      wg_partials[tid] = wg_partials[tid] + wg_partials[tid + step];
+{% if needsCount %}
+      wg_counts[tid] = wg_counts[tid] + wg_counts[tid + step];
+{% endif %}
+    }
+    workgroupBarrier();
+  }
+  let total = wg_partials[0];
+{% if needsCount %}
+  let countTotal = wg_counts[0];
+{% endif %}
+{% endif %}
+  if (tid == 0u) {
+{% if countIncludePad %}
+    y[index] = {{ scalar }}(total / f32(TAPS));
+{% else %}
+    y[index] = {{ scalar }}(select(0.0, total / countTotal, countTotal > 0.0));
+{% endif %}
+  }
+}

build/webgpu/pool-window-unroll.wgsl.jinja CHANGED Viewed

@@ -1,8 +1,13 @@
 // Compile-time-unrolled 2D NCHW pooling for small windows. Accumulation remains
 // kh-major/kw-minor, matching the loop-based implementations.
-{% set countIncludePad = countIncludePad is defined and countIncludePad %}
-{% set padBottomSpec = padBottomSpec | default(0) %}
-{% set padRightSpec = padRightSpec | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 const OUT_W: u32 = {{ outW }}u;
@@ -24,11 +29,7 @@ const PAD_RIGHT: i32 = {{ padRightSpec }};
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
-  // Fold very large flat outputs into the dispatch's second dimension.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (index >= params.count) {
-    return;
-  }
   let ow = index % OUT_W;
   var plane = index / OUT_W;
@@ -58,8 +59,8 @@ fn main(
 {% endfor %}
   }
 {% endfor %}
 {% if op == "average" and countIncludePad %}
   var denominator = 0u;
 {% for kh in range(kernelHSpec) %}
   if (ih_{{ kh }} >= -PAD_TOP && ih_{{ kh }} < IN_H + PAD_BOTTOM) {
@@ -71,6 +72,7 @@ fn main(
   }
 {% endfor %}
 {% else %}
   let denominator = seen;
 {% endif %}

 // Compile-time-unrolled 2D NCHW pooling for small windows. Accumulation remains
 // kh-major/kw-minor, matching the loop-based implementations.
+{% 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 }}
 const OUT_W: u32 = {{ outW }}u;
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "index") }}
   let ow = index % OUT_W;
   var plane = index / OUT_W;
 {% endfor %}
   }
 {% endfor %}
 {% if op == "average" and countIncludePad %}
   var denominator = 0u;
 {% for kh in range(kernelHSpec) %}
   if (ih_{{ kh }} >= -PAD_TOP && ih_{{ kh }} < IN_H + PAD_BOTTOM) {
   }
 {% endfor %}
 {% else %}
   let denominator = seen;
 {% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -93,7 +93,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/pool_op_test.cc",
         "test": "PoolTest.AveragePool",
-        "notes": "A 4x1 kernel bypasses the small-window unrolled path and covers the generic 2D AveragePool kernel with the same valid subnormal-average edge."
       },
       "attrs": { "count_include_pad": 0, "kernel_shape": [4, 1], "strides": [1, 1], "pads": [0, 0, 0, 0] },
       "inputs": {
@@ -985,7 +985,7 @@
     {
       "name": "f16_generic_global7x7_b1c256_nchw",
       "provenance": {
-        "notes": "Exercises the generic f16 NCHW 2D route with a full 7x7 reduction from [1,256,7,7] to [1,256,1,1]; kernelW=7 excludes the small-window unrolled route."
       },
       "attrs": { "count_include_pad": 0, "kernel_shape": [7, 7], "strides": [1, 1], "pads": [0, 0, 0, 0] },
       "inputs": {
@@ -1000,7 +1000,7 @@
     {
       "name": "f16_generic_count_include_pad_k5s1p2_b1c32_16x16",
       "provenance": {
-        "notes": "A realistic [1,32,16,16] float16 k5/s1/p2 pool selects the generic NCHW path. With count_include_pad=1, edge windows accumulate in-bounds values but divide by the full padded window size."
       },
       "attrs": { "count_include_pad": 1, "kernel_shape": [5, 5], "strides": [1, 1], "pads": [2, 2, 2, 2] },
       "inputs": {
@@ -1129,7 +1129,7 @@
     {
       "name": "k7s2p3_horizontal_reuse_f32",
       "provenance": {
-        "notes": "Compact route fixture for the ImageNet-style 7x7/stride-2/pad-3 pool. It locks the four-output horizontal reuse path and exclude-pad edge divisors."
       },
       "attrs": { "count_include_pad": 0, "kernel_shape": [7, 7], "strides": [2, 2], "pads": [3, 3, 3, 3] },
       "inputs": {
@@ -1170,6 +1170,288 @@
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 8, 10], "tolerance": 0.00001 } }
     }
   ]
 }

       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/pool_op_test.cc",
         "test": "PoolTest.AveragePool",
+        "notes": "A 4x1 pooling window checks a valid subnormal average over a non-square reduction."
       },
       "attrs": { "count_include_pad": 0, "kernel_shape": [4, 1], "strides": [1, 1], "pads": [0, 0, 0, 0] },
       "inputs": {
     {
       "name": "f16_generic_global7x7_b1c256_nchw",
       "provenance": {
+        "notes": "A full 7x7 float16 reduction from [1,256,7,7] to [1,256,1,1] checks global pooling across every channel."
       },
       "attrs": { "count_include_pad": 0, "kernel_shape": [7, 7], "strides": [1, 1], "pads": [0, 0, 0, 0] },
       "inputs": {
     {
       "name": "f16_generic_count_include_pad_k5s1p2_b1c32_16x16",
       "provenance": {
+        "notes": "A [1,32,16,16] float16 pool with a 5x5 window, stride one, padding two, and count_include_pad=1; edge windows divide by the full padded window size."
       },
       "attrs": { "count_include_pad": 1, "kernel_shape": [5, 5], "strides": [1, 1], "pads": [2, 2, 2, 2] },
       "inputs": {
     {
       "name": "k7s2p3_horizontal_reuse_f32",
       "provenance": {
+        "notes": "A 7x7, stride-2, pad-3 pool checks horizontal output positions and exclude-pad edge divisors."
       },
       "attrs": { "count_include_pad": 0, "kernel_shape": [7, 7], "strides": [2, 2], "pads": [3, 3, 3, 3] },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 8, 10], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "ort_avgpool_large_kernel_small_output_pad2",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/nn/pool_op_test.cc",
+        "test": "PoolTest.AveragePool_LargeKernelMultiElementOutput_WebGpu",
+        "notes": "A 12x12 average-pool window with padding 2 over a 24x24 input produces a 17x17 output, a large window behind a small output; padding is excluded from the divisor. Transcribed from onnxruntime's PoolTest.AveragePool_LargeKernelMultiElementOutput_WebGpu."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "pads": [2, 2, 2, 2],
+        "strides": [1, 1],
+        "count_include_pad": 0,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 2, 24, 24], "data": { "kind": "linspace", "start": 0.0, "end": 11.51 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 17, 17], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_count_include_pad_k12p2",
+      "provenance": {
+        "notes": "A 12x12 window with padding two over a 24x24 input and count_include_pad=1 produces 867 outputs; edge windows divide by the full padded tap count."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [1, 1],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 1,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 24, 24],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.053, "scale": 0.7 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3, 17, 17], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_ceil_mode_dilation_stride_overhang",
+      "provenance": {
+        "notes": "ceil_mode=1 with dilation 2, stride 5 and asymmetric pads: the last window in each axis starts inside the input and runs past both the input and the padded extent, so the cooperative reduction must reduce the in-bounds tap count as well as the sum."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [5, 5],
+        "pads": [3, 2, 4, 3],
+        "dilations": [2, 2],
+        "count_include_pad": 0,
+        "ceil_mode": 1
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 4, 30, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.029, "cosStep": 0.047, "scale": 0.8 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 4, 4, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_ncl1d_k128_at_min_taps",
+      "provenance": {
+        "notes": "A 512-element row is reduced by a 128-element average-pool window to 4 outputs per channel; a 127-element sibling window in this file reduces the same row by one fewer tap."
+      },
+      "attrs": { "kernel_shape": [128], "strides": [128], "pads": [0, 0] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 8, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.023, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 8, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_ncl1d_k127_below_min_taps",
+      "provenance": {
+        "notes": "A 512-element row is reduced by a 127-element average-pool window to 4 outputs per channel; its 128-element sibling in this file uses one more tap over the same row."
+      },
+      "attrs": { "kernel_shape": [127], "strides": [127], "pads": [0, 0] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 8, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.023, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 8, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_ncdhw3d_k8_pads_512_taps",
+      "provenance": {
+        "notes": "A 3-D 8x8x8 window with one pad on every face: 512 taps behind 32 outputs, the shape the serial route parallelises worst."
+      },
+      "attrs": { "kernel_shape": [8, 8, 8], "strides": [8, 8, 8], "pads": [1, 1, 1, 1, 1, 1], "count_include_pad": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.019, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 4, 2, 2, 2], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_serial_groups_at_gate",
+      "provenance": { "notes": "A 16,128-element output checks pooling across many independent windows." },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [1, 1],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 0,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 63, 23, 23],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.037, "scale": 0.4 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 63, 16, 16], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_serial_groups_over_gate",
+      "provenance": {
+        "notes": "One channel more than its sibling puts the output at exactly 64 serial dispatch groups, so the serial route is assumed to fill the device on its own and the cooperative reduction is rejected."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [1, 1],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 0,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 64, 23, 23],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.037, "scale": 0.4 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 16, 16], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_gate_ceil_mode_with_count_include_pad",
+      "provenance": {
+        "notes": "A 12x12 average-pool window with dilation 2, stride 5 and padding [3,2,4,3] combines ceil_mode and count_include_pad, both enabled, over a 30x30 input to produce a 4x4 output."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [5, 5],
+        "pads": [3, 2, 4, 3],
+        "dilations": [2, 2],
+        "count_include_pad": 1,
+        "ceil_mode": 1
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 4, 30, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.029, "cosStep": 0.047, "scale": 0.8 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 4, 4, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_f16_k12p2",
+      "provenance": {
+        "notes": "float16 input and output through the cooperative window reduction, which loads and accumulates in f32 and narrows once at the store."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [1, 1],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 0,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 3, 24, 24],
+          "data": { "kind": "fillFloat32", "sinStep": 0.041, "cosStep": 0.057, "scale": 1.1 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 3, 17, 17], "tolerance": 0.003 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_count_include_pad_k12p2_small_output",
+      "provenance": {
+        "notes": "A 12x12 window with stride two and padding two over a 24x24 input and count_include_pad=1 produces 81 outputs; edge windows divide by the full padded tap count."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [2, 2],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 1,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 24, 24],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.053, "scale": 0.7 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 9, 9], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "window_parallel_nchw2d_f16_k12p2_small_output",
+      "provenance": {
+        "notes": "float16 input and output through the cooperative window reduction, which loads and accumulates in f32 and narrows once at the store."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [2, 2],
+        "pads": [2, 2, 2, 2],
+        "count_include_pad": 0,
+        "ceil_mode": 0
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 1, 24, 24],
+          "data": { "kind": "fillFloat32", "sinStep": 0.041, "cosStep": 0.057, "scale": 1.1 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 1, 9, 9], "tolerance": 0.003 } }
+    },
+    {
+      "name": "window_parallel_outputs_to_taps_boundary_128",
+      "attrs": { "kernel_shape": [128], "strides": [1] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 1, 255], "data": { "kind": "fillFloat32", "scale": 0.3 } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 128] } }
+    },
+    {
+      "name": "window_parallel_outputs_to_taps_boundary_129",
+      "attrs": { "kernel_shape": [128], "strides": [1] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 1, 256], "data": { "kind": "fillFloat32", "scale": 0.3 } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 129] } }
+    },
+    {
+      "name": "window_parallel_genuinely_empty_dilated_window",
+      "attrs": { "kernel_shape": [128], "dilations": [2], "pads": [127, 127] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "constant", "value": 7.0 } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 1] } }
+    },
+    {
+      "name": "window_parallel_arbitrary_workgroup_96",
+      "tunables": { "WORKGROUP_SIZE": 96 },
+      "attrs": { "kernel_shape": [128], "strides": [128] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 128],
+          "data": {
+            "kind": "values",
+            "values": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 17.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 1] } }
+    },
+    {
+      "name": "window_parallel_arbitrary_workgroup_127",
+      "tunables": { "WORKGROUP_SIZE": 127 },
+      "attrs": { "kernel_shape": [128], "strides": [128] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 128],
+          "data": {
+            "kind": "values",
+            "values": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 17.0, 1.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 1] } }
     }
   ]
 }