sync 6fdf6301e2bb

Browse files

Files changed (15) hide show

README.md +11 -2
build/webgpu/bench.json +58 -1
build/webgpu/manifest.json +157 -39
build/webgpu/max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja +9 -7
build/webgpu/max-pool2d-nchw-u32.wgsl.jinja +9 -6
build/webgpu/max-pool3d-ncdhw-k5-tiled.wgsl.jinja +0 -2
build/webgpu/metadata.json +19 -15
build/webgpu/pool-global-reduction.wgsl.jinja +39 -37
build/webgpu/pool-ncl1d-x4.wgsl.jinja +17 -13
build/webgpu/pool-window-nd.wgsl.jinja +13 -6
build/webgpu/pool-window-reduction.wgsl.jinja +198 -0
build/webgpu/pool-window-unroll.wgsl.jinja +9 -8
build/webgpu/pool2d-nchw-k2s2-vec4.wgsl.jinja +10 -5
build/webgpu/pool2d-nchw-separable.wgsl.jinja +6 -4
build/webgpu/test.json +290 -8

README.md CHANGED Viewed

@@ -50,18 +50,27 @@ Attributes and default values (overridable per request):
 | `T` | `float32`, `float16`, `int8`, `uint8` |
 | `I` | `int64` |
 ## 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
 - [`max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja`](build/webgpu/max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja)
 - [`max-pool2d-nchw-u32.wgsl.jinja`](build/webgpu/max-pool2d-nchw-u32.wgsl.jinja)
 - [`max-pool3d-ncdhw-k5-tiled.wgsl.jinja`](build/webgpu/max-pool3d-ncdhw-k5-tiled.wgsl.jinja)
 - [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
 - [`pool-ncl1d-x4.wgsl.jinja`](build/webgpu/pool-ncl1d-x4.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-k2s2-vec4.wgsl.jinja`](build/webgpu/pool2d-nchw-k2s2-vec4.wgsl.jinja)
 - [`pool2d-nchw-separable.wgsl.jinja`](build/webgpu/pool2d-nchw-separable.wgsl.jinja)
@@ -69,7 +78,7 @@ Attributes and default values (overridable per request):
 ## 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`, `int8`, `uint8` |
 | `I` | `int64` |
+## 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 max 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. NaNs propagate through the fold; an all-out-of-bounds window writes 0.
+- `window_parallel_nchw2d` — Workgroup-cooperative max 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. NaNs propagate through the fold; an all-out-of-bounds window writes 0.
+- `window_parallel_ncdhw3d` — Workgroup-cooperative max 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. NaNs propagate through the fold; an all-out-of-bounds window writes 0.
 ## 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
 - [`max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja`](build/webgpu/max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja)
 - [`max-pool2d-nchw-u32.wgsl.jinja`](build/webgpu/max-pool2d-nchw-u32.wgsl.jinja)
 - [`max-pool3d-ncdhw-k5-tiled.wgsl.jinja`](build/webgpu/max-pool3d-ncdhw-k5-tiled.wgsl.jinja)
 - [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
 - [`pool-ncl1d-x4.wgsl.jinja`](build/webgpu/pool-ncl1d-x4.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-k2s2-vec4.wgsl.jinja`](build/webgpu/pool2d-nchw-k2s2-vec4.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/bench.json CHANGED Viewed

@@ -166,7 +166,7 @@
     {
       "name": "maxpool-f32-b1c512-20x20-k5s1p2-separable-floor",
       "provenance": {
-        "notes": "A YOLO-class detector shape with a 5x5 window exercises the minimum window size admitted by the separable route."
       },
       "attrs": { "kernel_shape": [5, 5], "strides": [1, 1], "pads": [2, 2, 2, 2] },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 512, 20, 20], "dist": "uniform" } },
@@ -187,6 +187,63 @@
       "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" }] }
     }
   ],
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }

     {
       "name": "maxpool-f32-b1c512-20x20-k5s1p2-separable-floor",
       "provenance": {
+        "notes": "Measures max pooling with a 5x5 window and padding 2 over 512 channels of a 20x20 YOLO-class detector feature map, keeping the spatial size unchanged."
       },
       "attrs": { "kernel_shape": [5, 5], "strides": [1, 1], "pads": [2, 2, 2, 2] },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 512, 20, 20], "dist": "uniform" } },
       "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": "max-pool-f32-b1c16-24x24-k12s1p2-window-parallel-starved",
+      "preset": "stress",
+      "provenance": {
+        "source": "synthetic benchmark",
+        "notes": "Measures max pooling with a 12x12 window and padding 2 over 16 channels of a 24x24 input, producing a 17x17 (289-element) output per channel."
+      },
+      "attrs": { "kernel_shape": [12, 12], "strides": [1, 1], "pads": [2, 2, 2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 16, 24, 24], "dist": "normal", "seed": 4311, "scale": 0.2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 16, 17, 17], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "max-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 separable route needs strideH < kernelH, so the serial route is the only alternative here."
+      },
+      "attrs": { "kernel_shape": [28, 28], "strides": [28, 28], "pads": [0, 0, 0, 0] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 16, 56, 56], "dist": "normal", "seed": 4312, "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": "max-pool-f32-b1c32-24x24-k12s1p2-serial-side-of-gate",
+      "preset": "stress",
+      "provenance": {
+        "source": "synthetic benchmark",
+        "notes": "A 9,248-element output measures pooling over many independent windows."
+      },
+      "attrs": { "kernel_shape": [12, 12], "strides": [1, 1], "pads": [2, 2, 2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 32, 24, 24], "dist": "normal", "seed": 4313, "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": "occupancy-window_parallel_nchw2d_small_output",
+      "preset": "stress",
+      "attrs": { "kernel_shape": [12, 12], "strides": [2, 2], "pads": [2, 2, 2, 2], "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] } }
     }
   ],
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }

build/webgpu/manifest.json CHANGED Viewed

@@ -36,7 +36,9 @@
     "TILED_H": { "default": 8 },
     "TILED_W": { "default": 8 },
     "TILED_LANES_D": { "default": 4 },
-    "SEPARABLE_MIN_WINDOW": { "default": 25 }
   },
   "derive": {
     "spatialRank": "ranks.x - 2",
@@ -59,7 +61,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",
-    "poolShapeOk": "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",
@@ -67,25 +68,27 @@
     "inferredOutH": "ceil(dim(shapes.x, ranks.x - 2) / strideH) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutH",
     "inferredOutD": "ceil(dim(shapes.x, 2) / strideD) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutD",
     "poolOutputShape": "[dim(shapes.x, 0), dim(shapes.x, 1), inferredOutW] if ranks.x == 3 else ([dim(shapes.x, 0), dim(shapes.x, 1), inferredOutH, inferredOutW] if ranks.x == 4 else [dim(shapes.x, 0), dim(shapes.x, 1), inferredOutD, inferredOutH, inferredOutW])",
-    "pool2dOutputShapeOk": "(dim(shapes.y, 2) == explicitOutH and dim(shapes.y, 3) == explicitOutW) or ((attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") and dim(shapes.y, 2) == ceil(dim(shapes.x, 2) / strideH) and dim(shapes.y, 3) == ceil(dim(shapes.x, 3) / strideW))",
     "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))",
     "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",
     "floatDtypeOk": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)",
     "globalWindowParallelCovered": "not present.indices and floatDtypeOk and poolShapeOk 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",
     "tiledStorageBytes": "(tunables.TILED_D + kernelD - 1) * (tunables.TILED_H + kernelH - 1) * (tunables.TILED_W + kernelW - 1) * 4",
     "tiledDeviceOk": "tiledStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and tunables.TILED_W * tunables.TILED_H * tunables.TILED_LANES_D <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILED_W <= device.limits.maxComputeWorkgroupSizeX and tunables.TILED_H <= device.limits.maxComputeWorkgroupSizeY and tunables.TILED_LANES_D <= device.limits.maxComputeWorkgroupSizeZ",
-    "tiledDispatchOk": "ceilDiv(dim(shapes.y, 4), tunables.TILED_W) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.y, 3), tunables.TILED_H) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.x, 0) * dim(shapes.x, 1) * ceilDiv(dim(shapes.y, 2), tunables.TILED_D) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "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)",
-    "autoPadTotalD": "max(0, (dim(shapes.y, 2) - 1) * strideD + ((kernelD - 1) * dilationD + 1) - dim(shapes.x, 2)) if ranks.x == 5 else 0",
-    "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"
   },
   "bindings": {
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$xVec4" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "inW", "type": "u32", "value": "dim(shapes.x, 2)" },
@@ -97,9 +100,8 @@
         { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
-    "params_7": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
@@ -116,6 +118,54 @@
         { "name": "padW", "type": "i32", "value": "effectivePadW" },
         { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
       ]
     }
   },
   "variants": [
@@ -132,9 +182,7 @@
           "id": "main",
           "name": "MaxPool.GlobalWindowParallel",
           "shader": "pool-global-reduction.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": {
-            "op": "\"max\"",
             "vec4": "inner(shapes.x, 1) % 4 == 0",
             "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
             "useSubgroups": "device.features.has(\"subgroups\")"
@@ -148,11 +196,83 @@
         }
       ]
     },
     {
       "id": "nchw2d_k3s2p1_vec4",
       "priority": 35,
       "when": ["not present.indices", "f16Ok(dtypes.T)", "dtypes.T == \"f16\"", "poolShapeOk", "ranks.x == 4", "ranks.y == 4", "attrs.auto_pad == \"NOTSET\"", "kernelH == 3", "kernelW == 3", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "padTop == 1", "padLeft == 1", "padBottom == 1", "padRight == 1", "attrs.ceil_mode == 0", "dim(shapes.x, 2) % 2 == 0", "dim(shapes.x, 3) % 8 == 0", "dim(shapes.y, 0) == dim(shapes.x, 0)", "dim(shapes.y, 1) == dim(shapes.x, 1)", "dim(shapes.y, 2) == dim(shapes.x, 2) / 2", "dim(shapes.y, 3) == dim(shapes.x, 3) / 2", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) >= 4"],
-      "derive": { "scalar": "dtypes.T", "xVec4": "\"vec4<f16>\"", "yVec4": "\"vec4<f16>\"" },
       "passes": [
         {
           "id": "main",
@@ -166,7 +286,7 @@
             "outBlocks": "dim(shapes.y, 3) / 4"
           },
           "bindings": [
-            "x_2",
             { "arg": "y", "elementType": "$yVec4" },
             { "name": "params", "struct": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.y) / 4" }] }
           ],
@@ -189,7 +309,6 @@
           "name": "MaxPool.Nchw2dK2s2Vec4",
           "shader": "pool2d-nchw-k2s2-vec4.wgsl.jinja",
           "derive": {
-            "op": "\"max\"",
             "outH": "dim(shapes.y, 2)",
             "outW": "dim(shapes.y, 3)",
             "outWHalf": "floor(dim(shapes.y, 3) / 2)",
@@ -197,7 +316,7 @@
             "inPlaneVec4": "floor(dim(shapes.x, 2) * dim(shapes.x, 3) / 4)"
           },
           "bindings": [
-            "x_2",
             "y",
             { "name": "params", "struct": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.y) / 2" }] }
           ],
@@ -220,7 +339,6 @@
           "name": "MaxPool.Nchw2dUnrolled",
           "shader": "pool-window-unroll.wgsl.jinja",
           "derive": {
-            "op": "\"max\"",
             "kernelHSpec": "kernelH",
             "kernelWSpec": "kernelW",
             "strideHSpec": "strideH",
@@ -257,9 +375,9 @@
           "id": "main",
           "name": "MaxPool.Ncl1dK2s2X4",
           "shader": "pool-ncl1d-x4.wgsl.jinja",
-          "derive": { "op": "\"max\"", "inW": "dim(shapes.x, 2)", "outW": "dim(shapes.y, 2)" },
           "bindings": [
-            "x_2",
             "y",
             {
               "name": "params",
@@ -290,7 +408,7 @@
           "id": "main",
           "name": "MaxPool.Ncl1d",
           "shader": "pool-window-nd.wgsl.jinja",
-          "derive": { "op": "\"max\"", "spatialRankSpec": 1 },
           "bindings": ["x", "y", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -303,13 +421,13 @@
     {
       "id": "nchw2d_uint8",
       "priority": 20,
-      "when": ["not present.indices", "tensorDtypes.x == \"uint8\"", "tensorDtypes.y == \"uint8\"", "poolShapeOk", "ranks.x == 4", "dim(shapes.y, 2) == explicitOutH", "dim(shapes.y, 3) == explicitOutW"],
       "passes": [
         {
           "id": "main",
-          "name": "nchw2d_uint8",
           "shader": "max-pool2d-nchw-u32.wgsl.jinja",
-          "bindings": ["x", "y", "params_7"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -322,17 +440,17 @@
       "id": "nchw2d_separable",
       "priority": 22,
       "when": ["not present.indices", "f16Ok(dtypes.T) and dtypes.T != \"u32\"", "poolShapeOk", "ranks.x == 4", "ranks.y == 4", "pool2dOutputShapeOk", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "strideH < kernelH", "separableFits"],
-      "derive": { "scalar": "dtypes.T", "pIs1": false, "pIs2": false },
       "intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }],
       "passes": [
         {
           "id": "reduce_w",
           "name": "MaxPool.SeparableW",
           "shader": "pool2d-nchw-separable.wgsl.jinja",
-          "derive": { "axisSpec": "\"w\"", "op": "\"max\"", "windowArea": "kernelH * kernelW" },
           "bindings": [
             "x",
-            { "name": "tmp", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
@@ -361,7 +479,7 @@
           "id": "reduce_h",
           "name": "MaxPool.SeparableH",
           "shader": "pool2d-nchw-separable.wgsl.jinja",
-          "derive": { "axisSpec": "\"h\"", "op": "\"max\"", "windowArea": "kernelH * kernelW" },
           "bindings": [
             { "name": "tmp", "buffer": "read-only-storage", "elementType": "f32" },
             "y",
@@ -399,10 +517,10 @@
       "passes": [
         {
           "id": "main",
-          "name": "nchw2d",
           "shader": "pool-window-nd.wgsl.jinja",
-          "derive": { "op": "\"max\"", "spatialRankSpec": 2 },
-          "bindings": ["x", "y", "params_7"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -454,7 +572,7 @@
           "id": "main",
           "name": "MaxPool.Ncdhw3d",
           "shader": "pool-window-nd.wgsl.jinja",
-          "derive": { "op": "\"max\"", "spatialRankSpec": 3 },
           "bindings": [
             "x",
             "y",
@@ -502,7 +620,7 @@
           "id": "main",
           "name": "MaxPool.Ncl1dIndices",
           "shader": "pool-window-nd.wgsl.jinja",
-          "derive": { "op": "\"max\"", "spatialRankSpec": 1 },
           "bindings": ["x", "y", "indices", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -515,15 +633,15 @@
     {
       "id": "nchw2d_uint8_indices",
       "priority": 16,
-      "when": ["present.indices", "tensorDtypes.x == \"uint8\"", "tensorDtypes.y == \"uint8\"", "poolShapeOk", "ranks.x == 4", "dim(shapes.y, 2) == explicitOutH", "dim(shapes.y, 3) == explicitOutW"],
       "derive": { "scalar": "\"u32\"", "hasIndices": true, "storageOrder": "attrs.storage_order" },
       "passes": [
         {
           "id": "main",
           "name": "MaxPool.Nchw2dUint8Indices",
           "shader": "pool-window-nd.wgsl.jinja",
-          "derive": { "op": "\"max\"", "spatialRankSpec": 2 },
-          "bindings": ["x", "y", "indices", "params_7"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -542,8 +660,8 @@
           "id": "main",
           "name": "MaxPool.Nchw2dIndices",
           "shader": "pool-window-nd.wgsl.jinja",
-          "derive": { "op": "\"max\"", "spatialRankSpec": 2 },
-          "bindings": ["x", "y", "indices", "params_7"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",

     "TILED_H": { "default": 8 },
     "TILED_W": { "default": 8 },
     "TILED_LANES_D": { "default": 4 },
+    "SEPARABLE_MIN_WINDOW": { "default": 25 },
+    "WINDOW_PARALLEL_MIN_TAPS": { "default": 128 },
+    "WINDOW_PARALLEL_MAX_SERIAL_GROUPS": { "default": 64 }
   },
   "derive": {
     "spatialRank": "ranks.x - 2",
     "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",
     "inferredOutH": "ceil(dim(shapes.x, ranks.x - 2) / strideH) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutH",
     "inferredOutD": "ceil(dim(shapes.x, 2) / strideD) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutD",
     "poolOutputShape": "[dim(shapes.x, 0), dim(shapes.x, 1), inferredOutW] if ranks.x == 3 else ([dim(shapes.x, 0), dim(shapes.x, 1), inferredOutH, inferredOutW] if ranks.x == 4 else [dim(shapes.x, 0), dim(shapes.x, 1), inferredOutD, inferredOutH, inferredOutW])",
+    "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)",
+    "poolShapeOk": "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)",
+    "pool2dOutputShapeOk": "(dim(shapes.y, 2) == explicitOutH and dim(shapes.y, 3) == explicitOutW) or ((attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") and dim(shapes.y, 2) == ceil(dim(shapes.x, 2) / strideH) and dim(shapes.y, 3) == ceil(dim(shapes.x, 3) / strideW))",
     "floatDtypeOk": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)",
     "globalWindowParallelCovered": "not present.indices and floatDtypeOk and poolShapeOk 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",
     "tiledStorageBytes": "(tunables.TILED_D + kernelD - 1) * (tunables.TILED_H + kernelH - 1) * (tunables.TILED_W + kernelW - 1) * 4",
     "tiledDeviceOk": "tiledStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and tunables.TILED_W * tunables.TILED_H * tunables.TILED_LANES_D <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILED_W <= device.limits.maxComputeWorkgroupSizeX and tunables.TILED_H <= device.limits.maxComputeWorkgroupSizeY and tunables.TILED_LANES_D <= device.limits.maxComputeWorkgroupSizeZ",
+    "tiledDispatchOk": "ceilDiv(dim(shapes.y, 4), tunables.TILED_W) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.y, 3), tunables.TILED_H) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.x, 0) * dim(shapes.x, 1) * ceilDiv(dim(shapes.y, 2), tunables.TILED_D) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
   },
   "bindings": {
+    "x_main": { "name": "x", "elementType": "$xVec4" },
     "params": {
       "struct": [
         { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "inW", "type": "u32", "value": "dim(shapes.x, 2)" },
         { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
       ]
     },
+    "params_main": {
       "name": "params",
       "struct": [
         { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
         { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
         { "name": "padW", "type": "i32", "value": "effectivePadW" },
         { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
       ]
+    },
+    "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": "main",
           "name": "MaxPool.GlobalWindowParallel",
           "shader": "pool-global-reduction.wgsl.jinja",
           "derive": {
             "vec4": "inner(shapes.x, 1) % 4 == 0",
             "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
             "useSubgroups": "device.features.has(\"subgroups\")"
         }
       ]
     },
+    {
+      "id": "window_parallel_ncl1d",
+      "priority": 28,
+      "when": ["not present.indices", "floatDtypeOk", "poolShapeOk", "ranks.x == 3", "dim(shapes.y, 2) == inferredOutW", "windowParallelEligible"],
+      "supersededBy": ["global_window_parallel"],
+      "derive": { "scalar": "dtypes.T" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "MaxPool.WindowParallelNcl1d",
+          "shader": "pool-window-reduction.wgsl.jinja",
+          "derive": {
+            "spatialRankSpec": 1,
+            "kernelDSpec": "kernelD",
+            "kernelHSpec": "kernelH",
+            "kernelWSpec": "kernelW",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
+            "useSubgroups": "device.features.has(\"subgroups\")"
+          },
+          "bindings": ["x", "y", "params_window_ncl1d"],
+          "dispatch": { "x": "numel(shapes.y)" }
+        }
+      ]
+    },
+    {
+      "id": "window_parallel_nchw2d",
+      "priority": 28,
+      "when": ["not present.indices", "floatDtypeOk", "poolShapeOk", "ranks.x == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "windowParallelEligible"],
+      "supersededBy": ["global_window_parallel"],
+      "derive": { "scalar": "dtypes.T" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "MaxPool.WindowParallelNchw2d",
+          "shader": "pool-window-reduction.wgsl.jinja",
+          "derive": {
+            "spatialRankSpec": 2,
+            "kernelDSpec": "kernelD",
+            "kernelHSpec": "kernelH",
+            "kernelWSpec": "kernelW",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
+            "useSubgroups": "device.features.has(\"subgroups\")"
+          },
+          "bindings": ["x", "y", "params_window_nchw2d"],
+          "dispatch": { "x": "numel(shapes.y)" }
+        }
+      ]
+    },
+    {
+      "id": "window_parallel_ncdhw3d",
+      "priority": 28,
+      "when": ["not present.indices", "floatDtypeOk", "poolShapeOk", "ranks.x == 5", "dim(shapes.y, 2) == inferredOutD", "dim(shapes.y, 3) == inferredOutH", "dim(shapes.y, 4) == inferredOutW", "windowParallelEligible"],
+      "supersededBy": ["global_window_parallel"],
+      "derive": { "scalar": "dtypes.T" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "MaxPool.WindowParallelNcdhw3d",
+          "shader": "pool-window-reduction.wgsl.jinja",
+          "derive": {
+            "spatialRankSpec": 3,
+            "kernelDSpec": "kernelD",
+            "kernelHSpec": "kernelH",
+            "kernelWSpec": "kernelW",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
+            "useSubgroups": "device.features.has(\"subgroups\")"
+          },
+          "bindings": ["x", "y", "params_window_ncdhw3d"],
+          "dispatch": { "x": "numel(shapes.y)" }
+        }
+      ]
+    },
     {
       "id": "nchw2d_k3s2p1_vec4",
       "priority": 35,
       "when": ["not present.indices", "f16Ok(dtypes.T)", "dtypes.T == \"f16\"", "poolShapeOk", "ranks.x == 4", "ranks.y == 4", "attrs.auto_pad == \"NOTSET\"", "kernelH == 3", "kernelW == 3", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "padTop == 1", "padLeft == 1", "padBottom == 1", "padRight == 1", "attrs.ceil_mode == 0", "dim(shapes.x, 2) % 2 == 0", "dim(shapes.x, 3) % 8 == 0", "dim(shapes.y, 0) == dim(shapes.x, 0)", "dim(shapes.y, 1) == dim(shapes.x, 1)", "dim(shapes.y, 2) == dim(shapes.x, 2) / 2", "dim(shapes.y, 3) == dim(shapes.x, 3) / 2", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) >= 4"],
+      "derive": { "xVec4": "\"vec4<f16>\"", "yVec4": "\"vec4<f16>\"" },
       "passes": [
         {
           "id": "main",
             "outBlocks": "dim(shapes.y, 3) / 4"
           },
           "bindings": [
+            "x_main",
             { "arg": "y", "elementType": "$yVec4" },
             { "name": "params", "struct": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.y) / 4" }] }
           ],
           "name": "MaxPool.Nchw2dK2s2Vec4",
           "shader": "pool2d-nchw-k2s2-vec4.wgsl.jinja",
           "derive": {
             "outH": "dim(shapes.y, 2)",
             "outW": "dim(shapes.y, 3)",
             "outWHalf": "floor(dim(shapes.y, 3) / 2)",
             "inPlaneVec4": "floor(dim(shapes.x, 2) * dim(shapes.x, 3) / 4)"
           },
           "bindings": [
+            "x_main",
             "y",
             { "name": "params", "struct": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.y) / 2" }] }
           ],
           "name": "MaxPool.Nchw2dUnrolled",
           "shader": "pool-window-unroll.wgsl.jinja",
           "derive": {
             "kernelHSpec": "kernelH",
             "kernelWSpec": "kernelW",
             "strideHSpec": "strideH",
           "id": "main",
           "name": "MaxPool.Ncl1dK2s2X4",
           "shader": "pool-ncl1d-x4.wgsl.jinja",
+          "derive": { "inW": "dim(shapes.x, 2)", "outW": "dim(shapes.y, 2)" },
           "bindings": [
+            "x_main",
             "y",
             {
               "name": "params",
           "id": "main",
           "name": "MaxPool.Ncl1d",
           "shader": "pool-window-nd.wgsl.jinja",
+          "derive": { "spatialRankSpec": 1 },
           "bindings": ["x", "y", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
     {
       "id": "nchw2d_uint8",
       "priority": 20,
+      "when": ["not present.indices", "tensorDtypes.x == \"uint8\"", "tensorDtypes.y == \"uint8\"", "poolShapeOk", "ranks.x == 4", "pool2dOutputShapeOk"],
       "passes": [
         {
           "id": "main",
+          "name": "MaxPool.Nchw2dUint8",
           "shader": "max-pool2d-nchw-u32.wgsl.jinja",
+          "bindings": ["x", "y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
       "id": "nchw2d_separable",
       "priority": 22,
       "when": ["not present.indices", "f16Ok(dtypes.T) and dtypes.T != \"u32\"", "poolShapeOk", "ranks.x == 4", "ranks.y == 4", "pool2dOutputShapeOk", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "strideH < kernelH", "separableFits"],
+      "derive": { "scalar": "dtypes.T" },
       "intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }],
       "passes": [
         {
           "id": "reduce_w",
           "name": "MaxPool.SeparableW",
           "shader": "pool2d-nchw-separable.wgsl.jinja",
+          "derive": { "axisSpec": "\"w\"", "windowArea": "kernelH * kernelW" },
           "bindings": [
             "x",
+            { "name": "tmp", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
           "id": "reduce_h",
           "name": "MaxPool.SeparableH",
           "shader": "pool2d-nchw-separable.wgsl.jinja",
+          "derive": { "axisSpec": "\"h\"", "windowArea": "kernelH * kernelW" },
           "bindings": [
             { "name": "tmp", "buffer": "read-only-storage", "elementType": "f32" },
             "y",
       "passes": [
         {
           "id": "main",
+          "name": "MaxPool.Nchw2d",
           "shader": "pool-window-nd.wgsl.jinja",
+          "derive": { "spatialRankSpec": 2 },
+          "bindings": ["x", "y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
           "id": "main",
           "name": "MaxPool.Ncdhw3d",
           "shader": "pool-window-nd.wgsl.jinja",
+          "derive": { "spatialRankSpec": 3 },
           "bindings": [
             "x",
             "y",
           "id": "main",
           "name": "MaxPool.Ncl1dIndices",
           "shader": "pool-window-nd.wgsl.jinja",
+          "derive": { "spatialRankSpec": 1 },
           "bindings": ["x", "y", "indices", "params"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
     {
       "id": "nchw2d_uint8_indices",
       "priority": 16,
+      "when": ["present.indices", "tensorDtypes.x == \"uint8\"", "tensorDtypes.y == \"uint8\"", "poolShapeOk", "ranks.x == 4", "pool2dOutputShapeOk"],
       "derive": { "scalar": "\"u32\"", "hasIndices": true, "storageOrder": "attrs.storage_order" },
       "passes": [
         {
           "id": "main",
           "name": "MaxPool.Nchw2dUint8Indices",
           "shader": "pool-window-nd.wgsl.jinja",
+          "derive": { "spatialRankSpec": 2 },
+          "bindings": ["x", "y", "indices", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
           "id": "main",
           "name": "MaxPool.Nchw2dIndices",
           "shader": "pool-window-nd.wgsl.jinja",
+          "derive": { "spatialRankSpec": 2 },
+          "bindings": ["x", "y", "indices", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",

build/webgpu/max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja CHANGED Viewed

@@ -6,7 +6,14 @@
 // The first loaded vector normally contributes only its final lane. At the left
 // edge that lane is padding, so the load is skipped. A NaN in any selected input
 // propagates to the corresponding output lane.
-enable f16;
 {{ env.wgsl.resourceDeclarations }}
 const IN_H: i32 = {{ inH }};
@@ -19,13 +26,11 @@ fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
   if (is_nan_f32(a)) {
     return a;
@@ -74,10 +79,7 @@ fn combine_rows(a: vec4<f32>, b: vec4<f32>) -> vec4<f32> {
 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 block = tile % OUT_BLOCKS;
   var outer = tile / OUT_BLOCKS;

 // The first loaded vector normally contributes only its final lane. At the left
 // edge that lane is padding, so the load is skipped. A NaN in any selected input
 // propagates to the corresponding output lane.
+{% 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 IN_H: i32 = {{ inH }};
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
   if (is_nan_f32(a)) {
     return a;
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "tile", "params.tileCount") }}
   let block = tile % OUT_BLOCKS;
   var outer = tile / OUT_BLOCKS;

build/webgpu/max-pool2d-nchw-u32.wgsl.jinja CHANGED Viewed

@@ -1,13 +1,16 @@
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 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 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 }}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "index") }}
   let ow = index % params.outW;
   var t = index / params.outW;

build/webgpu/max-pool3d-ncdhw-k5-tiled.wgsl.jinja CHANGED Viewed

@@ -5,13 +5,11 @@ fn negative_infinity() -> f32 {
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
   if (is_nan_f32(a)) {
     return a;

   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
   if (is_nan_f32(a)) {
     return a;

build/webgpu/metadata.json CHANGED Viewed

@@ -1,31 +1,35 @@
 {
   "name": "ai.onnx.MaxPool",
-  "id": "_ai_onnx_maxpool_webgpu_0edff93",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "AOgKxSadb2wkwmpf8SkHhMdwM1IpGC/TbJxTvDQZdNQ=",
-      "manifest.json": "5TTbqWe4MvoPE69DBi3YswL01bPFczoJuMndbfJkd2Q=",
-      "max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja": "NhbuTmSlBus0wOCYKCGshztjvOa5Ql6fscn6l5+9plc=",
-      "max-pool2d-nchw-u32.wgsl.jinja": "iLeF9bnjdoEU7ZeANJl9SED9RLEjB3Uo1/0MR5ktRCI=",
-      "max-pool3d-ncdhw-k5-tiled.wgsl.jinja": "Zwv1VailOkU7xlslNrbnUz9zyTXUSPGefhG71ybRVIk=",
-      "pool-global-reduction.wgsl.jinja": "d0Kc8mGbdDfI0PElsqLqjpnRqfyQeaS5LqSO/f8tfTM=",
-      "pool-ncl1d-x4.wgsl.jinja": "wADuoaBQjOxDwie14Tmpl/I5T6TzYIKKjepOAPZNwEE=",
-      "pool-window-nd.wgsl.jinja": "FtY2LWLhpfgAj3ye2GP+8uez3BF7qpIncbH7yX+bauo=",
-      "pool-window-unroll.wgsl.jinja": "/7wF7g88WmvgGS45Lg8NikpwzGdUoIX7VkTGRWjQy9w=",
-      "pool2d-nchw-k2s2-vec4.wgsl.jinja": "9LgEqP1IYbTszmUOHFYrIPpobfhunkVDsVITvSPdn58=",
-      "pool2d-nchw-separable.wgsl.jinja": "1v8HRM4mORp/jBxY1NUObJ3vjeV5fgR6PD3mBSmyIUM=",
-      "test.json": "1TZGbirb+TNnfoZRenB4AaQ1b5pqHnIrKUf9SGrs/Ac="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "global_window_parallel": ["pool-global-reduction.wgsl.jinja"],
       "nchw2d_k3s2p1_vec4": ["max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja"],
       "nchw2d_k2s2_vec4": ["pool2d-nchw-k2s2-vec4.wgsl.jinja"],
       "nchw2d_unrolled": ["pool-window-unroll.wgsl.jinja"],

 {
   "name": "ai.onnx.MaxPool",
+  "id": "_ai_onnx_maxpool_webgpu_9061839",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "CzLX4BrHv1R1JjxoSPkcWW10fd41NBG/vCioxDvTPAQ=",
+      "manifest.json": "2OsDYv2oeu4uHXX/5CQaYRkUph0hiuGYIYO5knHpIy4=",
+      "max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja": "507C/27t9YHR/gBvMJ9utJoVSQd8VUbm5oFs2jmESbo=",
+      "max-pool2d-nchw-u32.wgsl.jinja": "gS06j2pL8ITtuIdsLfQ3dZ+xLtrT6PJ1uH8sZEOGLMk=",
+      "max-pool3d-ncdhw-k5-tiled.wgsl.jinja": "JK3OfWsbQgJwIPKzQ1pZF4apPOHxA5RjnqzvoeqStVk=",
+      "pool-global-reduction.wgsl.jinja": "zsm141V8sZPAMwW+7yukfWPSpZsTLpdPoJkJHMf6Lr8=",
+      "pool-ncl1d-x4.wgsl.jinja": "r5hgQdHvV4WNJl1fUIbXORTuyWuf2EMDwDbbjHVIh9M=",
+      "pool-window-nd.wgsl.jinja": "mEftaXkif8tNyclNzIvkTfDMKBnkhS0TWRDTqT/5fLQ=",
+      "pool-window-reduction.wgsl.jinja": "b5BNYyK2mJ3qtdps6TxFRdGsf750RUcLEWhsyGh0YcY=",
+      "pool-window-unroll.wgsl.jinja": "1xyl/O2Zf1q8AWYD7obTxfJ0bpJxidnsaL4FociOdpY=",
+      "pool2d-nchw-k2s2-vec4.wgsl.jinja": "1LKfPAT8o8rGEUjdyMiBTZWGcYpLtVJQXmZxuWOtkTY=",
+      "pool2d-nchw-separable.wgsl.jinja": "2ulFFdVpiOnpnkJiCHA2TRUwDzfORoVh9EwIFEM42FI=",
+      "test.json": "8jRLh4Ogf3wHcAFtHvcBX8EmDrEYUqUCT2YlGXRfOTM="
     }
   },
+  "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_k3s2p1_vec4": ["max-pool2d-nchw-k3s2p1-vec4.wgsl.jinja"],
       "nchw2d_k2s2_vec4": ["pool2d-nchw-k2s2-vec4.wgsl.jinja"],
       "nchw2d_unrolled": ["pool-window-unroll.wgsl.jinja"],

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

@@ -31,8 +31,8 @@ fn quiet_nan() -> f32 {
   var bits = 0x7fc00000u;
   return bitcast<f32>(bits);
 }
 {% endif %}
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
@@ -50,10 +50,9 @@ fn max_propagate_nan(a: f32, b: f32) -> f32 {
 @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 = negative_infinity();
@@ -76,39 +75,42 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
 {% endif %}
 {% if useSubgroups %}
-  // subgroupMax does not propagate NaN; carry an explicit NaN flag alongside.
-  let lane_nan = is_nan_f32(partial);
-  let sg_nan = subgroupAny(lane_nan);
-  let sg_max = subgroupMax(select(partial, negative_infinity(), lane_nan));
-  var total = select(sg_max, quiet_nan(), sg_nan);
-  // 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(negative_infinity(), total, rank == 0u);
-    workgroupBarrier();
-    var acc = negative_infinity();
-    for (var i = rank; i < WG; i = i + count) {
-      acc = max_propagate_nan(acc, wg_partials[i]);
-    }
-    let acc_nan = is_nan_f32(acc);
-    let fold_nan = subgroupAny(acc_nan);
-    let fold_max = subgroupMax(select(acc, negative_infinity(), acc_nan));
-    total = select(fold_max, quiet_nan(), fold_nan);
   }
   if (tid == 0u) {
-    y[plane] = {{ scalar }}(total);
-  }
 {% else %}
   // No-subgroups fallback: full workgroup-tree reduction of the per-thread partials.
   wg_tree[tid] = partial;
@@ -121,7 +123,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   }
   if (tid == 0u) {
     let total = wg_tree[0];
     y[plane] = {{ scalar }}(total);
   }
-{% endif %}
 }

   var bits = 0x7fc00000u;
   return bitcast<f32>(bits);
 }
 {% endif %}
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 @compute @workgroup_size(WG)
 fn main(@builtin(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 = negative_infinity();
 {% endif %}
 {% if useSubgroups %}
+{% filter indent(2, true) %}
+// subgroupMax does not propagate NaN; carry an explicit NaN flag alongside.
+let lane_nan = is_nan_f32(partial);
+let sg_nan = subgroupAny(lane_nan);
+let sg_max = subgroupMax(select(partial, negative_infinity(), lane_nan));
+var total = select(sg_max, quiet_nan(), sg_nan);
+// 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(negative_infinity(), total, rank == 0u);
+  workgroupBarrier();
+  var acc = negative_infinity();
+  for (var i = rank; i < WG; i = i + lanes) {
+    acc = max_propagate_nan(acc, wg_partials[i]);
   }
+  let acc_nan = is_nan_f32(acc);
+  let fold_nan = subgroupAny(acc_nan);
+  let fold_max = subgroupMax(select(acc, negative_infinity(), acc_nan));
+  total = select(fold_max, quiet_nan(), fold_nan);
+}
+{% 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);
   }
 }

build/webgpu/pool-ncl1d-x4.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,9 @@
 {{ env.wgsl.resourceDeclarations }}
 // Four adjacent NCL max-pool outputs per invocation for non-overlapping k2/s2
@@ -11,23 +17,21 @@ fn is_nan_f32(value: f32) -> bool {
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
-  if (is_nan_f32(a)) { return a; }
-  if (is_nan_f32(b)) { return b; }
   return max(a, b);
 }
-fn pair_max(a: f32, b: f32) -> f32 {
-  return max_propagate_nan(a, b);
-}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>,
 ) {
-  let block = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (block >= params.blockCount) { return; }
   let plane = block / BLOCKS_PER_PLANE;
   let outputBlock = block - plane * BLOCKS_PER_PLANE;
@@ -38,8 +42,8 @@ fn main(
   let b = vec4<f32>(x[inputVec + 1u]);
   // k2/s2: two aligned vec4 reads cover four independent windows exactly.
-  y[outputBase] = {{ scalar }}(pair_max(a.x, a.y));
-  y[outputBase + 1u] = {{ scalar }}(pair_max(a.z, a.w));
-  y[outputBase + 2u] = {{ scalar }}(pair_max(b.x, b.y));
-  y[outputBase + 3u] = {{ scalar }}(pair_max(b.z, b.w));
 }

+{% 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 NCL max-pool outputs per invocation for non-overlapping k2/s2
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
+  if (is_nan_f32(a)) {
+    return a;
+  }
+  if (is_nan_f32(b)) {
+    return b;
+  }
   return max(a, b);
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>,
 ) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "block", "params.blockCount", guardInline=true) }}
   let plane = block / BLOCKS_PER_PLANE;
   let outputBlock = block - plane * BLOCKS_PER_PLANE;
   let b = vec4<f32>(x[inputVec + 1u]);
   // k2/s2: two aligned vec4 reads cover four independent windows exactly.
+  y[outputBase] = {{ scalar }}(max_propagate_nan(a.x, a.y));
+  y[outputBase + 1u] = {{ scalar }}(max_propagate_nan(a.z, a.w));
+  y[outputBase + 2u] = {{ scalar }}(max_propagate_nan(b.x, b.y));
+  y[outputBase + 3u] = {{ scalar }}(max_propagate_nan(b.z, b.w));
 }

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

@@ -1,5 +1,11 @@
 {% set hasIndices = hasIndices is defined and hasIndices %}
 {% set storageOrder = storageOrder | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 // Direct spatial-pooling traversal for a channels-first tensor.
@@ -7,23 +13,24 @@ fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
-  if (is_nan_f32(a)) { return a; }
-  if (is_nan_f32(b)) { return b; }
   return max(a, b);
 }
 @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;

 {% set hasIndices = hasIndices is defined and hasIndices %}
 {% set storageOrder = storageOrder | default(0) %}
+{% 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.
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
+  if (is_nan_f32(a)) {
+    return a;
+  }
+  if (is_nan_f32(b)) {
+    return b;
+  }
   return max(a, b);
 }
 @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 @@

+{% 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.
+// The fold propagates NaNs while taking the maximum. The in-bounds tap count is
+// reduced alongside it so an output whose window is entirely out of bounds
+// writes 0 rather than -inf, matching the serial route.
+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>;
+var<workgroup> wg_counts: array<f32, WG>;
+fn negative_infinity() -> f32 {
+  var bits = 0xff800000u;
+  return bitcast<f32>(bits);
+}
+{% if useSubgroups %}
+fn quiet_nan() -> f32 {
+  var bits = 0x7fc00000u;
+  return bitcast<f32>(bits);
+}
+{% endif %}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+fn max_propagate_nan(a: f32, b: f32) -> f32 {
+  if (is_nan_f32(a)) {
+    return a;
+  }
+  if (is_nan_f32(b)) {
+    return b;
+  }
+  return max(a, b);
+}
+@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 = negative_infinity();
+  var counted = 0.0;
+  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 = max_propagate_nan(partial, value);
+      counted = counted + 1.0;
+    }
+  }
+{% if useSubgroups %}
+{% filter indent(2, true) %}
+// subgroupMax does not propagate NaN; carry an explicit NaN flag alongside.
+let lane_nan = is_nan_f32(partial);
+let sg_nan = subgroupAny(lane_nan);
+let sg_max = subgroupMax(select(partial, negative_infinity(), lane_nan));
+var total = select(sg_max, quiet_nan(), sg_nan);
+var countTotal = subgroupAdd(counted);
+// 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(negative_infinity(), total, rank == 0u);
+  wg_counts[tid] = select(0.0, countTotal, rank == 0u);
+  workgroupBarrier();
+  var acc = negative_infinity();
+  var countAcc = 0.0;
+  for (var i = rank; i < WG; i = i + lanes) {
+    acc = max_propagate_nan(acc, wg_partials[i]);
+    countAcc = countAcc + wg_counts[i];
+  }
+  let acc_nan = is_nan_f32(acc);
+  let fold_nan = subgroupAny(acc_nan);
+  let fold_max = subgroupMax(select(acc, negative_infinity(), acc_nan));
+  total = select(fold_max, quiet_nan(), fold_nan);
+  countTotal = subgroupAdd(countAcc);
+}
+{% 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;
+  wg_counts[tid] = counted;
+  workgroupBarrier();
+  for (var liveCount = WG; liveCount > 1u; liveCount = (liveCount + 1u) / 2u) {
+    let step = (liveCount + 1u) / 2u;
+    if (tid + step < liveCount) {
+      wg_partials[tid] = max_propagate_nan(wg_partials[tid], wg_partials[tid + step]);
+      wg_counts[tid] = wg_counts[tid] + wg_counts[tid + step];
+    }
+    workgroupBarrier();
+  }
+  let total = wg_partials[0];
+  let countTotal = wg_counts[0];
+{% endif %}
+  if (tid == 0u) {
+    y[index] = {{ scalar }}(select(0.0, total, countTotal > 0.0));
+  }
+}

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

@@ -1,5 +1,13 @@
 // Compile-time-unrolled 2D NCHW pooling for small windows. Accumulation remains
 // kh-major/kw-minor, matching the loop-based implementations.
 {{ env.wgsl.resourceDeclarations }}
 const OUT_W: u32 = {{ outW }}u;
@@ -17,13 +25,11 @@ fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
   if (is_nan_f32(a)) {
     return a;
@@ -38,11 +44,7 @@ fn max_propagate_nan(a: f32, b: f32) -> f32 {
 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;
@@ -73,7 +75,6 @@ fn main(
   }
 {% endfor %}
   let output = select(0.0, result, seen != 0u);
   y[index] = {{ scalar }}(output);
 }

 // 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;
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
   if (is_nan_f32(a)) {
     return a;
 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 %}
   let output = select(0.0, result, seen != 0u);
   y[index] = {{ scalar }}(output);
 }

build/webgpu/pool2d-nchw-k2s2-vec4.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,11 @@
 {{ env.wgsl.resourceDeclarations }}
 // 2D NCHW MaxPool specialised for the dense 2x2 / stride-2 / no-pad /
@@ -24,11 +32,8 @@ fn max_nan(a: {{ scalar }}, b: {{ scalar }}) -> {{ scalar }} {
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat tile index (one tile = two output columns).
-  let tile = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (tile >= params.tileCount) {
-    return;
-  }
   let tileCol = tile % OUT_W_HALF;        // pair-of-columns within the row
   var t = tile / OUT_W_HALF;

+{% 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 MaxPool specialised for the dense 2x2 / stride-2 / no-pad /
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // One flat index per tile (one tile = two output columns).
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "tile", "params.tileCount") }}
   let tileCol = tile % OUT_W_HALF;        // pair-of-columns within the row
   var t = tile / OUT_W_HALF;

build/webgpu/pool2d-nchw-separable.wgsl.jinja CHANGED Viewed

@@ -12,16 +12,18 @@ fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
-  if (is_nan_f32(a)) { return a; }
-  if (is_nan_f32(b)) { return b; }
   return max(a, b);
 }

   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
 fn is_nan_f32(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
 }
 fn max_propagate_nan(a: f32, b: f32) -> f32 {
+  if (is_nan_f32(a)) {
+    return a;
+  }
+  if (is_nan_f32(b)) {
+    return b;
+  }
   return max(a, b);
 }

build/webgpu/test.json CHANGED Viewed

@@ -1178,7 +1178,7 @@
     {
       "name": "f16_nchw_realistic_k3s2_pad1",
       "provenance": {
-        "notes": "A realistic float16 ResNet-stem shape [1,64,56,56] with k3/s2/pad1 bypasses the k2 vec4 specialization and exercises padded windows in the generic NCHW loop."
       },
       "attrs": { "kernel_shape": [3, 3], "strides": [2, 2], "pads": [1, 1, 1, 1] },
       "inputs": {
@@ -1237,9 +1237,7 @@
     },
     {
       "name": "ncdhw3d_f32_k5s1pad2_tiled",
-      "provenance": {
-        "notes": "Compact correctness lock for the separable workgroup-tiled 5x5x5 stride-1/pad-2 path used by the realistic 48-cubed stress case."
-      },
       "attrs": { "kernel_shape": [5, 5, 5], "pads": [2, 2, 2, 2, 2, 2] },
       "inputs": {
         "x": {
@@ -1253,7 +1251,7 @@
     {
       "name": "ncdhw3d_f32_k5s1pad2_tiled_partial_depth",
       "provenance": {
-        "notes": "Correctness lock for a final tiled depth brick containing one valid output plane instead of requiring outD to be divisible by 24."
       },
       "attrs": { "kernel_shape": [5, 5, 5], "pads": [2, 2, 2, 2, 2, 2] },
       "inputs": {
@@ -1429,7 +1427,7 @@
     {
       "name": "separable_k6_nopad_f16",
       "provenance": {
-        "notes": "A float16 6x6 window is the smallest square reaching SEPARABLE_MIN_WINDOW=36 and exercises the separable two-dimensional pooling path."
       },
       "attrs": { "kernel_shape": [6, 6] },
       "inputs": {
@@ -1444,7 +1442,7 @@
     {
       "name": "autopad_same_upper_rank3_route",
       "provenance": {
-        "notes": "A 1-D SAME_UPPER pool is a legal ONNX graph whose output extent the port itself infers, and the generic rank-3 selector compared that extent against the explicit-pad arithmetic, which SAME cannot satisfy. The pinned values are exact: padding contributes negative infinity, so the maxima over the three windows are 2, 4 and 5."
       },
       "attrs": { "auto_pad": "SAME_UPPER", "kernel_shape": [3], "strides": [2] },
       "inputs": {
@@ -1466,7 +1464,7 @@
     {
       "name": "autopad_same_upper_rank5_route",
       "provenance": {
-        "notes": "A rank-5 SAME_UPPER pool derives output extents from automatic padding on all three spatial axes; volumetric path eligibility must use those derived extents."
       },
       "attrs": { "auto_pad": "SAME_UPPER", "kernel_shape": [3, 3, 3], "strides": [2, 2, 2] },
       "inputs": {
@@ -1477,6 +1475,290 @@
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2, 2, 2], "tolerance": 0 } }
     }
   ]
 }

     {
       "name": "f16_nchw_realistic_k3s2_pad1",
       "provenance": {
+        "notes": "A float16 [1,64,56,56] pooling case with a 3x3 window, stride two, and padding one checks padded edge windows."
       },
       "attrs": { "kernel_shape": [3, 3], "strides": [2, 2], "pads": [1, 1, 1, 1] },
       "inputs": {
     },
     {
       "name": "ncdhw3d_f32_k5s1pad2_tiled",
+      "provenance": { "notes": "Compact 5x5x5 stride-one, pad-two pooling checks all spatial boundaries." },
       "attrs": { "kernel_shape": [5, 5, 5], "pads": [2, 2, 2, 2, 2, 2] },
       "inputs": {
         "x": {
     {
       "name": "ncdhw3d_f32_k5s1pad2_tiled_partial_depth",
       "provenance": {
+        "notes": "The final depth region contains one valid output plane, checking a non-multiple-of-24 output depth."
       },
       "attrs": { "kernel_shape": [5, 5, 5], "pads": [2, 2, 2, 2, 2, 2] },
       "inputs": {
     {
       "name": "separable_k6_nopad_f16",
       "provenance": {
+        "notes": "A float16 max pool with a 6x6 (36-element) window and no padding reduces a 16x16 input to an 11x11 output over 2 channels."
       },
       "attrs": { "kernel_shape": [6, 6] },
       "inputs": {
     {
       "name": "autopad_same_upper_rank3_route",
       "provenance": {
+        "notes": "A 1-D SAME_UPPER pool infers its output extent from implicit padding. Padding contributes negative infinity, so the three window maxima are exactly 2, 4 and 5."
       },
       "attrs": { "auto_pad": "SAME_UPPER", "kernel_shape": [3], "strides": [2] },
       "inputs": {
     {
       "name": "autopad_same_upper_rank5_route",
       "provenance": {
+        "notes": "A rank-5 (3-D) max pool with a 3x3x3 window, stride 2 and SAME_UPPER automatic padding reduces a 3x3x3 input to a 2x2x2 output."
       },
       "attrs": { "auto_pad": "SAME_UPPER", "kernel_shape": [3, 3, 3], "strides": [2, 2, 2] },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "uint8_same_upper",
+      "provenance": {
+        "notes": "Asymmetric automatic padding with uint8 storage; increasing inputs pin each maximum and its flattened index."
+      },
+      "attrs": { "auto_pad": "SAME_UPPER", "kernel_shape": [3, 3], "strides": [2, 2] },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [1, 1, 4, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ceil_mode_partial_windows_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "uint8",
+          "shape": [1, 1, 2, 2],
+          "data": { "kind": "values", "values": [11, 12, 15, 16] },
+          "tolerance": 0
+        }
+      }
+    },
+    {
+      "name": "uint8_same_upper_indices",
+      "provenance": {
+        "notes": "Asymmetric automatic padding with uint8 storage; increasing inputs pin each maximum and its flattened index."
+      },
+      "attrs": { "auto_pad": "SAME_UPPER", "kernel_shape": [3, 3], "strides": [2, 2] },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [1, 1, 4, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ceil_mode_partial_windows_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "uint8",
+          "shape": [1, 1, 2, 2],
+          "data": { "kind": "values", "values": [11, 12, 15, 16] },
+          "tolerance": 0
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [1, 1, 2, 2],
+          "data": { "kind": "values", "values": [10, 11, 14, 15] },
+          "tolerance": 0
+        }
+      }
+    },
+    {
+      "name": "uint8_same_lower",
+      "provenance": {
+        "notes": "Asymmetric automatic padding with uint8 storage; increasing inputs pin each maximum and its flattened index."
+      },
+      "attrs": { "auto_pad": "SAME_LOWER", "kernel_shape": [3, 3], "strides": [2, 2] },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [1, 1, 4, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ceil_mode_partial_windows_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "uint8",
+          "shape": [1, 1, 2, 2],
+          "data": { "kind": "values", "values": [6, 8, 14, 16] },
+          "tolerance": 0
+        }
+      }
+    },
+    {
+      "name": "uint8_same_lower_indices",
+      "provenance": {
+        "notes": "Asymmetric automatic padding with uint8 storage; increasing inputs pin each maximum and its flattened index."
+      },
+      "attrs": { "auto_pad": "SAME_LOWER", "kernel_shape": [3, 3], "strides": [2, 2] },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [1, 1, 4, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ceil_mode_partial_windows_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "uint8",
+          "shape": [1, 1, 2, 2],
+          "data": { "kind": "values", "values": [6, 8, 14, 16] },
+          "tolerance": 0
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [1, 1, 2, 2],
+          "data": { "kind": "values", "values": [5, 7, 13, 15] },
+          "tolerance": 0
+        }
+      }
+    },
+    {
+      "name": "ort_maxpool_large_kernel_small_output_pad2",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/nn/pool_op_test.cc",
+        "test": "PoolTest.MaxPool_LargeKernelMultiElementOutput_WebGpu",
+        "notes": "A 12x12 max-pool window with padding 2 over a 24x24 input produces a 17x17 output, a large window behind a small output. Transcribed from onnxruntime's PoolTest.MaxPool_LargeKernelMultiElementOutput_WebGpu."
+      },
+      "attrs": { "kernel_shape": [12, 12], "pads": [2, 2, 2, 2], "strides": [1, 1], "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_nan_window_f32",
+      "provenance": {
+        "notes": "NaNs inside a 144-tap window must survive both the per-invocation fold and the cross-invocation reduction, which carries an explicit NaN flag because subgroupMax does not propagate one."
+      },
+      "attrs": { "kernel_shape": [12, 12], "strides": [1, 1], "pads": [2, 2, 2, 2], "ceil_mode": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 24, 24],
+          "data": {
+            "kind": "fillFloat32",
+            "sinStep": 0.031,
+            "cosStep": 0.053,
+            "scale": 0.7,
+            "nanStart": 300.0,
+            "nanCount": 4.0
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 17, 17], "tolerance": 0, "allowNaN": true } }
+    },
+    {
+      "name": "window_parallel_ncl1d_k256",
+      "provenance": {
+        "notes": "A 1-D 256-tap window over 64 outputs: the serial route would give each output one invocation looping 256 taps."
+      },
+      "attrs": { "kernel_shape": [256], "strides": [256], "pads": [0, 0] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 16, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.023, "scale": 0.5 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 16, 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: out-of-bounds taps must leave the maximum untouched and still be excluded from the in-bounds tap count."
+      },
+      "attrs": { "kernel_shape": [8, 8, 8], "strides": [8, 8, 8], "pads": [1, 1, 1, 1, 1, 1] },
+      "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_f16_k12p2",
+      "provenance": { "notes": "float16 max pooling through the cooperative window reduction." },
+      "attrs": { "kernel_shape": [12, 12], "strides": [1, 1], "pads": [2, 2, 2, 2], "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_dilation_stride_ceil_overhang",
+      "provenance": {
+        "notes": "ceil_mode=1 with dilation 2 and stride 5: the last window overhangs the input, and an output whose window holds no in-bounds tap must still write 0 rather than negative infinity."
+      },
+      "attrs": {
+        "kernel_shape": [12, 12],
+        "strides": [5, 5],
+        "pads": [3, 2, 4, 3],
+        "dilations": [2, 2],
+        "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_nan_window_f32_small_output",
+      "provenance": {
+        "notes": "NaNs inside a 144-tap window must survive both the per-invocation fold and the cross-invocation reduction, which carries an explicit NaN flag because subgroupMax does not propagate one."
+      },
+      "attrs": { "kernel_shape": [12, 12], "strides": [2, 2], "pads": [2, 2, 2, 2], "ceil_mode": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 24, 24],
+          "data": {
+            "kind": "fillFloat32",
+            "sinStep": 0.031,
+            "cosStep": 0.053,
+            "scale": 0.7,
+            "nanStart": 300.0,
+            "nanCount": 4.0
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 9, 9], "tolerance": 0, "allowNaN": true } }
+    },
+    {
+      "name": "window_parallel_nchw2d_f16_k12p2_small_output",
+      "provenance": { "notes": "float16 max pooling through the cooperative window reduction." },
+      "attrs": { "kernel_shape": [12, 12], "strides": [2, 2], "pads": [2, 2, 2, 2], "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] } }
     }
   ]
 }