sync 6fdf6301e2bb

Browse files

Files changed (7) hide show

README.md +2 -2
build/webgpu/bench.json +3 -7
build/webgpu/manifest.json +4 -4
build/webgpu/metadata.json +8 -8
build/webgpu/pool-global-reduction.wgsl.jinja +31 -30
build/webgpu/pool-global-serial.wgsl.jinja +1 -0
build/webgpu/test.json +5 -7

README.md CHANGED Viewed

@@ -39,14 +39,14 @@ See the [ONNX `GlobalAveragePool` spec](https://onnx.ai/onnx/operators/onnx__Glo
 - [`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
 - [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
 - [`pool-global-serial.wgsl.jinja`](build/webgpu/pool-global-serial.wgsl.jinja)
 ## 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.

 - [`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
 - [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
 - [`pool-global-serial.wgsl.jinja`](build/webgpu/pool-global-serial.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.

build/webgpu/bench.json CHANGED Viewed

@@ -85,9 +85,7 @@
     },
     {
       "name": "nchw_1x1024x14x14_portable_gate_window",
-      "provenance": {
-        "notes": "Spatial 196 sits in the window the portable plane-parallel gate excludes (65 with subgroups, 1024 without), so a subgroupless device falls back to the serial kernel here."
-      },
       "preset": "smoke",
       "vars": { "batch": 1, "channels": 1024, "spatial": 196 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 1024, 14, 14], "dist": "normal", "seed": 216 } },
@@ -107,7 +105,7 @@
       "name": "nchw_32x128x16x16_manyplanes_spatial256",
       "preset": "smoke",
       "provenance": {
-        "notes": "4096 planes of 256 elements: the serial kernel already has an invocation per plane, so this is the regime where the plane-parallel workgroup has the least to add. GlobalLpPool's same shape prefers the serial route."
       },
       "vars": { "batch": 32, "channels": 128, "spatial": 256 },
       "inputs": { "x": { "dtype": "float32", "shape": [32, 128, 16, 16], "dist": "normal", "seed": 218 } },
@@ -150,9 +148,7 @@
     {
       "name": "nchw_1x65536x9x9_planes65536_spatial81",
       "preset": "smoke",
-      "provenance": {
-        "notes": "65536 planes: four outputs per invocation still leaves 16384 invocations, which is the regime the vec4-output serial kernel is written for."
-      },
       "vars": { "batch": 1, "channels": 65536, "spatial": 81 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 65536, 9, 9], "dist": "normal", "seed": 222 } },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 65536, 1, 1], "dist": "empty" } },

     },
     {
       "name": "nchw_1x1024x14x14_portable_gate_window",
+      "provenance": { "notes": "Measures global average pooling over 1024 channels of a 14x14 (196-element) plane." },
       "preset": "smoke",
       "vars": { "batch": 1, "channels": 1024, "spatial": 196 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 1024, 14, 14], "dist": "normal", "seed": 216 } },
       "name": "nchw_32x128x16x16_manyplanes_spatial256",
       "preset": "smoke",
       "provenance": {
+        "notes": "Measures global average pooling over 4096 planes (batch 32 x 128 channels) of a 256-element 16x16 spatial size."
       },
       "vars": { "batch": 32, "channels": 128, "spatial": 256 },
       "inputs": { "x": { "dtype": "float32", "shape": [32, 128, 16, 16], "dist": "normal", "seed": 218 } },
     {
       "name": "nchw_1x65536x9x9_planes65536_spatial81",
       "preset": "smoke",
+      "provenance": { "notes": "Measures global average pooling over 65,536 channels of a 9x9 (81-element) plane." },
       "vars": { "batch": 1, "channels": 65536, "spatial": 81 },
       "inputs": { "x": { "dtype": "float32", "shape": [1, 65536, 9, 9], "dist": "normal", "seed": 222 } },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 65536, 1, 1], "dist": "empty" } },

build/webgpu/manifest.json CHANGED Viewed

@@ -22,15 +22,14 @@
       "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": "GlobalAveragePool.WgPlaneParallel",
           "shader": "pool-global-reduction.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": {
-            "op": "\"average\"",
             "vec4": "inner(shapes.x, 1) % 4 == 0",
             "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
             "useSubgroups": "device.features.has(\"subgroups\")"
@@ -51,7 +50,7 @@
           "id": "main",
           "name": "GlobalAveragePool.Nchw",
           "shader": "pool-global-serial.wgsl.jinja",
-          "derive": { "op": "\"average\"", "outputsPerInvocation": 1 },
           "bindings": [
             "x",
             "y",
@@ -71,5 +70,6 @@
         }
       ]
     }
-  ]
 }

       "id": "wg_plane_parallel",
       "priority": 30,
       "when": ["planeParallelEligible"],
+      "demoteWhen": ["not device.features.has(\"subgroups\") and dtypes.T == \"f32\" and numel(shapes.y) >= tunables.PLANE_PARALLEL_STARVED_PLANES and inner(shapes.x, 1) <= 2 * tunables.PLANE_PARALLEL_AMORTIZED_SPATIAL"],
       "derive": { "planeInput": "\"vec4<\" ~ dtypes.T ~ \">\" if inner(shapes.x, 1) % 4 == 0 else dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "GlobalAveragePool.WgPlaneParallel",
           "shader": "pool-global-reduction.wgsl.jinja",
           "derive": {
             "vec4": "inner(shapes.x, 1) % 4 == 0",
             "workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
             "useSubgroups": "device.features.has(\"subgroups\")"
           "id": "main",
           "name": "GlobalAveragePool.Nchw",
           "shader": "pool-global-serial.wgsl.jinja",
+          "derive": { "outputsPerInvocation": 1 },
           "bindings": [
             "x",
             "y",
         }
       ]
     }
+  ],
+  "bindings": {}
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,22 +1,22 @@
 {
   "name": "ai.onnx.GlobalAveragePool",
-  "id": "_ai_onnx_globalaveragepool_webgpu_56a66d8",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "3Q7Hhu7A0Ya1DWVVrIWVxyW2qxLEFF9YtJEwjfLEOQE=",
-      "manifest.json": "13pXbztMIjCthx+34Tv/epFss6CetZgcwudxKej2M44=",
-      "pool-global-reduction.wgsl.jinja": "SKnTSUq9nU838U0t7VyvC7hgz4h5L+Ndv1rQCM9Vks8=",
-      "pool-global-serial.wgsl.jinja": "idkz50sNZFg0fXNECnKQxbRZK1xFDWzyuJyxIMaepEw=",
-      "test.json": "egIwDffBXas1eU87qZq5zYwzi9nVWXppV5JdlZc9PDE="
     }
   },
-  "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"] }
   }
 }

 {
   "name": "ai.onnx.GlobalAveragePool",
+  "id": "_ai_onnx_globalaveragepool_webgpu_545a553",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "popdR9Qi/BImt4Ux/qIVOhoPeNwySqGaBoF4YB6spn0=",
+      "manifest.json": "FHENwZk+f3tgm5aUoHx0/8TsVYl3ec5Hp436+Pue8a0=",
+      "pool-global-reduction.wgsl.jinja": "uCrSccgj/Fibd866kxpiYkprR8g6q1vofgQZZ/oxGkE=",
+      "pool-global-serial.wgsl.jinja": "pYR1xvboodN6Lty4Y9ytA9B1gNh2yyUZaCM0k8tRB/0=",
+      "test.json": "adZiAK+JRTpxqoklQoJU7ZP86KCq9GdZNUsyLpN1ruY="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": { "wg_plane_parallel": ["pool-global-reduction.wgsl.jinja"], "nchw": ["pool-global-serial.wgsl.jinja"] }
   }
 }

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

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

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

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

@@ -1,4 +1,5 @@
 {{ env.wgsl.resourceDeclarations }}
 fn pool_plane(out_index: u32) -> f32 {
   let base = out_index * params.spatial;
   var result = 0.0;

 {{ env.wgsl.resourceDeclarations }}
 fn pool_plane(out_index: u32) -> f32 {
   let base = out_index * params.spatial;
   var result = 0.0;

build/webgpu/test.json CHANGED Viewed

@@ -36,7 +36,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/pool_op_test.cc",
         "test": "PoolTest.GlobalAveragePool_Large_128",
-        "notes": "Large compact-cycle version selects the workgroup-parallel spatial-plane reduction variant while keeping the exact average subnormal."
       },
       "inputs": {
         "x": {
@@ -346,9 +346,7 @@
     },
     {
       "name": "f32_nchw2d_scalar_smallspatial_1x16x7x7",
-      "provenance": {
-        "notes": "A compact float32 7x7 input keeps the spatial size below the plane-parallel threshold and exercises scalar GlobalAveragePool."
-      },
       "inputs": {
         "x": {
           "dtype": "float32",
@@ -359,9 +357,9 @@
       "outputs": { "y": { "dtype": "float32", "shape": [1, 16, 1, 1], "tolerance": 0.000001 } }
     },
     {
-      "name": "f32_nchw2d_spatial64_selector_boundary_1x16x8x8",
       "provenance": {
-        "notes": "Pins the 64-element crossover: fixed subgroup-32 adapters favor the serial scalar plane kernel here, while variable-width subgroup adapters retain the workgroup-parallel route."
       },
       "inputs": {
         "x": {
@@ -403,7 +401,7 @@
     {
       "name": "f32_nchw2d_lowoccupancy_8planes_1x8x128x128",
       "provenance": {
-        "notes": "Eight independent float32 planes with a large spatial extent exercise low-occupancy GlobalAveragePool."
       },
       "inputs": {
         "x": {

       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/pool_op_test.cc",
         "test": "PoolTest.GlobalAveragePool_Large_128",
+        "notes": "A rank-three input with 1,024 equal subnormal values on its single spatial axis must average exactly to that subnormal."
       },
       "inputs": {
         "x": {
     },
     {
       "name": "f32_nchw2d_scalar_smallspatial_1x16x7x7",
+      "provenance": { "notes": "A compact float32 7x7 input checks spatial averaging over a short plane." },
       "inputs": {
         "x": {
           "dtype": "float32",
       "outputs": { "y": { "dtype": "float32", "shape": [1, 16, 1, 1], "tolerance": 0.000001 } }
     },
     {
+      "name": "f32_nchw2d_spatial64_boundary_1x16x8x8",
       "provenance": {
+        "notes": "A float32 global average pool reduces 16 channels of an 8x8 (64-element) plane to a single value each."
       },
       "inputs": {
         "x": {
     {
       "name": "f32_nchw2d_lowoccupancy_8planes_1x8x128x128",
       "provenance": {
+        "notes": "A float32 global average pool reduces 8 channels of a 128x128 (16,384-element) plane to a single value each."
       },
       "inputs": {
         "x": {