sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +12 -9
build/webgpu/bench.json +1 -2
build/webgpu/manifest.json +44 -80
build/webgpu/metadata.json +11 -8
build/webgpu/pool-global-reduction.wgsl.jinja +41 -37
build/webgpu/pool-global-serial.wgsl.jinja +3 -6
build/webgpu/test.json +1 -2

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `GlobalMaxPool` spec](https://onnx.ai/onnx/operators/onnx__GlobalM
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Input tensor of shape `(N x C x D1 x ... x Dn)`, where `N` is the batch size and `C` is the number of channels. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `y` | `T` | same as `X` | — | Output tensor of shape `(N x C x 1 x ... x 1)`; the maximum value over each spatial region per channel. | required |
 ## Type constraints
@@ -36,7 +36,7 @@ See the [ONNX `GlobalMaxPool` spec](https://onnx.ai/onnx/operators/onnx__GlobalM
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -45,15 +45,18 @@ See the [ONNX `GlobalMaxPool` spec](https://onnx.ai/onnx/operators/onnx__GlobalM
 ## Use with `@huggingface/kernels`
-The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
-The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
-- `y`
-Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Input tensor of shape `(N x C x D1 x ... x Dn)`, where `N` is the batch size and `C` is the number of channels. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `T` | same as `x` | — | Output tensor of shape `(N x C x 1 x ... x 1)`; the maximum value over each spatial region per channel. | required |
 ## Type constraints
 ## 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
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
+This example supplies explicit metadata for:
+- `y`
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.GlobalMaxPool",
   "cases": [
     {
       "name": "nchw_1x2048x7x7",
@@ -47,7 +46,7 @@
       "name": "nchw_1x8x512x512_lowoccupancy_8planes",
       "preset": "stress",
       "provenance": {
-        "source": "authored for variant coverage",
         "notes": "Low-occupancy eight-plane global reduction that verifies the cooperative plane route; legacy serial paths remain reserved for shapes outside its feature-tier domain."
       },
       "vars": { "batch": 1, "channels": 8, "spatial": 262144 },

 {
   "cases": [
     {
       "name": "nchw_1x2048x7x7",
       "name": "nchw_1x8x512x512_lowoccupancy_8planes",
       "preset": "stress",
       "provenance": {
+        "source": "synthetic benchmark",
         "notes": "Low-occupancy eight-plane global reduction that verifies the cooperative plane route; legacy serial paths remain reserved for shapes outside its feature-tier domain."
       },
       "vars": { "batch": 1, "channels": 8, "spatial": 262144 },

build/webgpu/manifest.json CHANGED Viewed

@@ -2,109 +2,73 @@
   "domain": "ai.onnx",
   "name": "GlobalMaxPool",
   "sinceVersion": 1,
-  "description": "Applies max pooling across all spatial dimensions of `X`, producing one value per channel. Equivalent to MaxPool with kernel size equal to the full spatial extent of the input; output shape is `(N x C x 1 x ... x 1)`.",
-  "inputs": [
-    {
-      "role": "X",
-      "dtype": "T",
-      "description": "Input tensor of shape `(N x C x D1 x ... x Dn)`, where `N` is the batch size and `C` is the number of channels."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "Y",
-      "dtype": "T",
-      "rank": "ranks.X",
-      "description": "Output tensor of shape `(N x C x 1 x ... x 1)`; the maximum value over each spatial region per channel."
-    }
-  ],
   "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "Y", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "PLANE_PARALLEL_MIN_SPATIAL": 64,
-    "PLANE_PARALLEL_AMORTIZED_SPATIAL": 768,
-    "PLANE_PARALLEL_STARVED_PLANES": 1024
   },
   "derive": {
-    "baseContract": "f16Ok(dtypes.T) and ranks.X >= 3 and ranks.Y == ranks.X and dim(shapes.Y, 0) == dim(shapes.X, 0) and dim(shapes.Y, 1) == dim(shapes.X, 1) and inner(shapes.Y, 1) == 1",
-    "planeParallelEligible": "inner(shapes.X, 1) >= tunables.PLANE_PARALLEL_MIN_SPATIAL and (inner(shapes.X, 1) >= tunables.PLANE_PARALLEL_AMORTIZED_SPATIAL or numel(shapes.Y) <= tunables.PLANE_PARALLEL_STARVED_PLANES) and numel(shapes.Y) <= device.limits.maxComputeWorkgroupsPerDimension"
-  },
-  "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
-  "bindingSets": {
-    "parallelPlane": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "X",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$planeInput"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.X, 1)" }] }
-      }
-    ],
-    "serialPlane": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "spatial", "type": "u32", "value": "inner(shapes.X, 1)" },
-            { "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "wg_plane_parallel",
       "priority": 30,
-      "when": ["baseContract", "planeParallelEligible"],
-      "constants": { "planeInput": "\"vec4<\" ~ dtypes.T ~ \">\" if inner(shapes.X, 1) % 4 == 0 else dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "GlobalMaxPool.WgPlaneParallel",
-          "source": {
-            "shader": "pool-global-reduction.wgsl.jinja",
-            "inputs": {
-              "op": "\"max\"",
-              "vec4": "inner(shapes.X, 1) % 4 == 0",
-              "workgroupSize": "tunables.WORKGROUP_SIZE",
-              "useSubgroups": "device.features.has(\"subgroups\")"
-            }
           },
-          "bindings": "parallelPlane",
-          "dispatch": { "x": "numel(shapes.Y)" }
         }
       ]
     },
     {
       "id": "nchw",
-      "when": ["baseContract", "not planeParallelEligible"],
       "passes": [
         {
           "id": "main",
           "name": "GlobalMaxPool.Nchw",
-          "source": {
-            "shader": "pool-global-serial.wgsl.jinja",
-            "inputs": { "op": "\"max\"", "outputsPerInvocation": 1 }
-          },
-          "bindings": "serialPlane",
-          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "GlobalMaxPool",
   "sinceVersion": 1,
+  "inputs": { "x": { "onnx": "X", "dtype": "T" } },
+  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x" } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "PLANE_PARALLEL_MIN_SPATIAL": { "default": 64 },
+    "PLANE_PARALLEL_AMORTIZED_SPATIAL": { "default": 768 },
+    "PLANE_PARALLEL_STARVED_PLANES": { "default": 1024 }
   },
   "derive": {
+    "baseContract": "f16Ok(dtypes.T) and ranks.x >= 3 and ranks.y == ranks.x and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1) and inner(shapes.y, 1) == 1",
+    "planeParallelEligible": "inner(shapes.x, 1) >= tunables.PLANE_PARALLEL_MIN_SPATIAL and (inner(shapes.x, 1) >= tunables.PLANE_PARALLEL_AMORTIZED_SPATIAL or numel(shapes.y) <= tunables.PLANE_PARALLEL_STARVED_PLANES) and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "scalar": "dtypes.T"
   },
+  "when": ["baseContract"],
   "variants": [
     {
       "id": "wg_plane_parallel",
       "priority": 30,
+      "when": ["planeParallelEligible"],
+      "derive": { "planeInput": "\"vec4<\" ~ dtypes.T ~ \">\" if inner(shapes.x, 1) % 4 == 0 else dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "GlobalMaxPool.WgPlaneParallel",
+          "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\")"
           },
+          "bindings": [
+            { "arg": "x", "elementType": "$planeInput" },
+            "y",
+            { "name": "params", "struct": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.x, 1)" }] }
+          ],
+          "dispatch": { "x": "numel(shapes.y)" }
         }
       ]
     },
     {
       "id": "nchw",
+      "when": ["not planeParallelEligible"],
       "passes": [
         {
           "id": "main",
           "name": "GlobalMaxPool.Nchw",
+          "shader": "pool-global-serial.wgsl.jinja",
+          "derive": { "op": "\"max\"", "outputsPerInvocation": 1 },
+          "bindings": [
+            "x",
+            "y",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "spatial", "type": "u32", "value": "inner(shapes.x, 1)" },
+                { "name": "count", "type": "u32", "value": "numel(shapes.y)" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,19 +1,22 @@
 {
   "name": "ai.onnx.GlobalMaxPool",
-  "id": "_ai_onnx_globalmaxpool_webgpu_e1e2a9b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "Eet2dCeLHwGCclxUWpvzSTP1eo23iWxOYR0Hv0IevB8=",
-      "manifest.json": "Er9sIUKGJm7kZSVuRlRzkWCQjRnFZTlvPrketaONZXk=",
-      "pool-global-reduction.wgsl.jinja": "pY5tVjS1U8d/Hbabf04MwlObnCibUKA2Bjr0b4W3prk=",
-      "pool-global-serial.wgsl.jinja": "t2QTHN0c2huEUxRxnzd9dQZOcXvm89WwJiSzz9pOtzo=",
-      "test.json": "g0uQA0LM6LP5qlY6wRsAtFjX/IbNwTJMtiocyw157cI="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.GlobalMaxPool" }
 }

 {
   "name": "ai.onnx.GlobalMaxPool",
+  "id": "_ai_onnx_globalmaxpool_webgpu_3b227dd",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "ylL2awmbpw5AHQZ8F59dhSjg/cNeDcbRmweZNMUcmJI=",
+      "manifest.json": "wErFlwEMiglUVXVU3mkqkU91gMQEcJV+HlmHnZ58vj8=",
+      "pool-global-reduction.wgsl.jinja": "d0Kc8mGbdDfI0PElsqLqjpnRqfyQeaS5LqSO/f8tfTM=",
+      "pool-global-serial.wgsl.jinja": "1tV3ld6YisYBbl2khXWO6Nyh3dtgh3sEFmRvB1ODdek=",
+      "test.json": "mtjdotV4Sw0xAfrVdrLQvtL5emPC7ir7K4pqJOJf3iM="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": { "wg_plane_parallel": ["pool-global-reduction.wgsl.jinja"], "nchw": ["pool-global-serial.wgsl.jinja"] }
+  }
 }

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

@@ -1,7 +1,4 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
-{% if source.useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
@@ -9,17 +6,17 @@ enable subgroups;
 // Workgroup-parallel global pooling: one workgroup reduces one contiguous
 // (n, c) spatial plane (NC[D]HW layout makes the plane contiguous).
 // Threads stride the plane (vec4 loads when the plane size is divisible by 4),
-// then combine the per-thread partials. With subgroups: a subgroup collective
-// + a shared-memory fold of the per-subgroup partials. Without subgroups, a
-// full workgroup tree keeps workgroup-per-plane parallelism
-// instead of falling to the one-thread-per-plane serial scalar kernel. The
-// epilogue finalizes average / max / Lp. Sum reassociates;
-// max is exact.
-const WG: u32 = {{ source.workgroupSize }}u;
-{% if source.useSubgroups %}
-// WebGPU subgroups are at least 4 wide, so a workgroup holds at most WG / 4 subgroups.
-const MAX_SUBGROUPS: u32 = WG / 4u;
-var<workgroup> wg_partials: array<f32, MAX_SUBGROUPS>;
 {% else %}
 var<workgroup> wg_tree: array<f32, WG>;
 {% endif %}
@@ -28,7 +25,7 @@ fn negative_infinity() -> f32 {
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
-{% if source.useSubgroups %}
 fn quiet_nan() -> f32 {
   var bits = 0x7fc00000u;
@@ -54,23 +51,13 @@ 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 source.useSubgroups %},
         @builtin(subgroup_size) sg_size: u32
 {%- endif %}) {
   let plane = wid.x;
   let tid = lid.x;
-{% if source.useSubgroups %}
-  // Seed cross-subgroup slots with the reduction identity so unwritten slots
-  // (when the workgroup has fewer than MAX_SUBGROUPS subgroups) are harmless.
-  if (tid < MAX_SUBGROUPS) {
-    wg_partials[tid] = negative_infinity();
-  }
-  workgroupBarrier();
-{% endif %}
   var partial = negative_infinity();
-{% if source.vec4 %}
   let spatialVec = params.spatial / 4u;
   let base = plane * spatialVec;
   for (var i = tid; i < spatialVec; i = i + WG) {
@@ -88,21 +75,38 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   }
 {% endif %}
-{% if source.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));
-  if (subgroupElect()) {
-    wg_partials[tid / sg_size] = select(sg_max, quiet_nan(), sg_nan);
   }
-  workgroupBarrier();
   if (tid == 0u) {
-    let nsg = min((WG + sg_size - 1u) / sg_size, MAX_SUBGROUPS);
-    var total = negative_infinity();
-    for (var i = 0u; i < nsg; i = i + 1u) {
-      total = max_propagate_nan(total, wg_partials[i]);
-    }
     y[plane] = {{ scalar }}(total);
   }
 {% else %}

+{% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Workgroup-parallel global pooling: one workgroup reduces one contiguous
 // (n, c) spatial plane (NC[D]HW layout makes the plane contiguous).
 // Threads stride the plane (vec4 loads when the plane size is divisible by 4),
+// 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.
+const WG: u32 = {{ workgroupSizeSpec }}u;
+{% if useSubgroups %}
+// 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>;
 {% else %}
 var<workgroup> wg_tree: array<f32, WG>;
 {% endif %}
   var bits = 0xff800000u;
   return bitcast<f32>(bits);
 }
+{% if useSubgroups %}
 fn quiet_nan() -> f32 {
   var bits = 0x7fc00000u;
 @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();
+{% if vec4 %}
   let spatialVec = params.spatial / 4u;
   let base = plane * spatialVec;
   for (var i = tid; i < spatialVec; i = i + WG) {
   }
 {% 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 %}

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

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 fn negative_infinity() -> f32 {
@@ -31,11 +28,11 @@ fn pool_plane(out_index: u32) -> f32 {
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // Each invocation produces one or four channel planes. gid.y carries folded workgroups beyond
   // the device's per-dimension dispatch limit.
-  let first = (gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u) * {{ source.outputsPerInvocation }}u;
-{% for lane in range(source.outputsPerInvocation) %}
   if (first + {{ lane }}u < params.count) {
     y[first + {{ lane }}u] = {{ scalar }}(pool_plane(first + {{ lane }}u));
   }

 {{ env.wgsl.resourceDeclarations }}
 fn negative_infinity() -> f32 {
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // Each invocation produces one or four channel planes. gid.y carries folded workgroups beyond
   // the device's per-dimension dispatch limit.
+  let first = (gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u) * {{ outputsPerInvocation }}u;
+{% for lane in range(outputsPerInvocation) %}
   if (first + {{ lane }}u < params.count) {
     y[first + {{ lane }}u] = {{ scalar }}(pool_plane(first + {{ lane }}u));
   }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.GlobalMaxPool",
   "cases": [
     {
       "name": "dispatch_cliff_ncl1d",
@@ -64,7 +63,7 @@
     {
       "name": "f32_nchw2d_smallspatial_1x16x7x7",
       "provenance": {
-        "notes": "Compact f32 sibling for the 7x7 small-spatial vec4-output benchmark; keeps channels small while exercising the same low-spatial output path."
       },
       "inputs": {
         "x": {

 {
   "cases": [
     {
       "name": "dispatch_cliff_ncl1d",
     {
       "name": "f32_nchw2d_smallspatial_1x16x7x7",
       "provenance": {
+        "notes": "A compact float32 7x7 input with few channels exercises the small-spatial vec4-output path."
       },
       "inputs": {
         "x": {