sync 91d990483a17

Browse files

Files changed (9) hide show

README.md +15 -12
build/webgpu/bench.json +22 -1
build/webgpu/compress-scatter.wgsl.jinja +12 -15
build/webgpu/compress.wgsl.jinja +0 -3
build/webgpu/manifest.json +113 -246
build/webgpu/metadata.json +18 -10
build/webgpu/scan-block-prefix-u32.wgsl.jinja +48 -12
build/webgpu/scan-flags-block-exclusive.wgsl.jinja +53 -20
build/webgpu/test.json +3 -4

README.md CHANGED Viewed

@@ -18,16 +18,16 @@ See the [ONNX `Compress` spec](https://onnx.ai/onnx/operators/onnx__Compress.htm
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `input` | `input` | `T` | — | — | Input tensor of rank `r >= 1` to select from. | required |
-| `condition` | `condition` | `C` | `1` | — | Rank-1 boolean mask indicating which slices or elements to select; may be shorter than the axis dimension, in which case trailing slices are discarded. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `output` | `output` | `T` | derived | — | Selected slices with rank `r` when `axis` is specified, or rank 1 when the input is flattened. The selected dimension equals the number of true values in the inspected condition prefix. | required |
 ## Attributes
@@ -46,7 +46,7 @@ Attributes and default values (overridable per request):
 ## 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
@@ -57,15 +57,18 @@ Attributes and default values (overridable per request):
 ## 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:
-- `output`
-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 | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `input` | `T` | — | — | Input tensor of rank `r >= 1` to select from. | required |
+| `condition` | `C` | `1` | — | Rank-1 boolean mask indicating which slices or elements to select; may be shorter than the axis dimension, in which case trailing slices are discarded. | required |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | derived | — | Selected slices with rank `r` when `axis` is specified, or rank 1 when the input is flattened. The selected dimension equals the number of true values in the inspected condition prefix. | required |
 ## Attributes
 ## 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:
+- `output`
 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.Compress",
   "cases": [
     {
       "name": "compress-axis0-f32-64k-by-4",
@@ -99,6 +98,28 @@
           }
         ]
       }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "compress-axis0-f32-64k-by-4",
           }
         ]
       }
+    },
+    {
+      "name": "compress-axis1-f32-short-condition-128-of-2048",
+      "preset": "edge",
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [4, 2048, 768], "dist": "normal", "seed": 915, "scale": 1 },
+        "condition": { "dtype": "bool", "shape": [128], "dist": "linearMod", "mod": 2 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [4, 64, 768] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.output) * 2) * 4" }] }
+    },
+    {
+      "name": "compress-axis1-f32-full-condition-2048-control",
+      "preset": "edge",
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [4, 2048, 768], "dist": "normal", "seed": 915, "scale": 1 },
+        "condition": { "dtype": "bool", "shape": [2048], "dist": "linearMod", "mod": 2 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [4, 1024, 768] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.output) * 2) * 4" }] }
     }
   ]
 }

build/webgpu/compress-scatter.wgsl.jinja CHANGED Viewed

@@ -4,31 +4,29 @@
 // lower bits gives an item's rank within its lane. The flattened route walks a
 // lane's ITEMS elements in order, so its output position advances without a
 // second lookup per element.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 const ITEMS: u32 = {{ scanItems }}u;
-{% if source.axisMode %}
 const BLOCK_ITEMS: u32 = WG * ITEMS;
 {% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
-{% if source.axisMode %}
-  let input_index = (wg.x + wg.y * nwg.x) * WG + lid.x;
-  if (input_index >= params.inputCount) {
-    return;
-  }
-  let inner_index = input_index % params.inner;
-  let axis_index = (input_index / params.inner) % params.axisDim;
-  if (axis_index >= params.n) {
     return;
   }
   // ITEMS is a compile-time power of two, so these resolve to a shift and a mask.
   let lane = axis_index / ITEMS;
   let item = axis_index % ITEMS;
@@ -38,11 +36,10 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
   }
   let output_axis_index = blockPrefix[axis_index / BLOCK_ITEMS] + offsets[lane]
     + countOneBits(packed & ((1u << item) - 1u));
-  let outer_index = input_index / (params.inner * params.axisDim);
   let output_index = (outer_index * params.outputAxisDim + output_axis_index) * params.inner + inner_index;
   output[output_index] = input[input_index];
 {% else %}
-  let block = wg.x + wg.y * nwg.x;
   let lane = block * WG + lid.x;
   // base == block * BLOCK_ITEMS + tid * ITEMS, which is exactly lane * ITEMS.
   let base = lane * ITEMS;

 // lower bits gives an item's rank within its lane. The flattened route walks a
 // lane's ITEMS elements in order, so its output position advances without a
 // second lookup per element.
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 const ITEMS: u32 = {{ scanItems }}u;
+{% if axisMode %}
 const BLOCK_ITEMS: u32 = WG * ITEMS;
 {% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
+{% if axisMode %}
+  // Only the first `n` positions along the axis can be selected, so the grid
+  // covers (outer, n, inner) and the full-input address is rebuilt from it —
+  // the uninspected tail of the axis never costs an invocation.
+  let scan_index = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG + lid.x;
+  if (scan_index >= params.scanCount) {
     return;
   }
+  let inner_index = scan_index % params.inner;
+  let axis_index = (scan_index / params.inner) % params.n;
+  let outer_index = scan_index / (params.inner * params.n);
+  let input_index = (outer_index * params.axisDim + axis_index) * params.inner + inner_index;
   // ITEMS is a compile-time power of two, so these resolve to a shift and a mask.
   let lane = axis_index / ITEMS;
   let item = axis_index % ITEMS;
   }
   let output_axis_index = blockPrefix[axis_index / BLOCK_ITEMS] + offsets[lane]
     + countOneBits(packed & ((1u << item) - 1u));
   let output_index = (outer_index * params.outputAxisDim + output_axis_index) * params.inner + inner_index;
   output[output_index] = input[input_index];
 {% else %}
+  let block = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let lane = block * WG + lid.x;
   // base == block * BLOCK_ITEMS + tid * ITEMS, which is exactly lane * ITEMS.
   let base = lane * ITEMS;

build/webgpu/compress.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 @compute @workgroup_size(1)





1	{{ env.wgsl.resourceDeclarations }}
2
3	@compute @workgroup_size(1)

build/webgpu/manifest.json CHANGED Viewed

@@ -2,43 +2,20 @@
   "domain": "ai.onnx",
   "name": "Compress",
   "sinceVersion": 11,
-  "description": "Selects slices from the input tensor along an axis where the corresponding `condition` element is true. If `axis` is omitted, the input is flattened and elements are selected by position. The condition may be shorter than the selected dimension; values beyond its length are discarded. The data-dependent output extent must equal the number of true entries in the inspected prefix.",
-  "inputs": [
-    { "role": "input", "dtype": "T", "description": "Input tensor of rank `r >= 1` to select from." },
-    {
-      "role": "condition",
-      "dtype": "C",
-      "description": "Rank-1 boolean mask indicating which slices or elements to select; may be shorter than the axis dimension, in which case trailing slices are discarded.",
-      "rank": 1
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "description": "Selected slices with rank `r` when `axis` is specified, or rank 1 when the input is flattened. The selected dimension equals the number of true values in the inspected condition prefix.",
-      "rank": "ranks.input if has(attrs, \"axis\") else 1"
-    }
-  ],
-  "attributeDescriptions": {
-    "axis": "Axis along which to select slices; if omitted the input is flattened before selection. Negative values index from the end; accepted range is `[-r, r-1]`."
-  },
   "typeConstraints": { "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"], "C": ["bool"] },
-  "args": {
-    "input": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "condition": { "kind": "tensor", "semantic": "condition", "role": "input" },
-    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "SCAN_ITEMS_PER_THREAD": 16,
-    "PARALLEL_MIN_ELEMENTS": 512,
-    "MAX_PARALLEL_INPUT": 1073741823
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
-    "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
     "axisParam": "attrs.axis if has(attrs, \"axis\") else 2147483647",
     "normalizedAxis": "axisParam if axisParam >= 0 else axisParam + ranks.input",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
@@ -52,186 +29,33 @@
     "axisScanN": "min(dim(shapes.input, normalizedAxis), dim(shapes.condition, 0))",
     "axisScanBlocks": "ceilDiv(axisScanN, scanBlockItems)",
     "axisScanThreads": "axisScanBlocks * workgroupSize",
-    "axisParallelFits": "workgroupStorageOk and axisScanThreads * 4 <= storageBufferLimit and axisScanBlocks * 4 <= storageBufferLimit and axisScanBlocks <= foldedDispatchCapacity and ceilDiv(numel(shapes.input), workgroupSize) <= foldedDispatchCapacity"
   },
-  "constants": { "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
-  "bindingSets": {
-    "scanBlocks": [
-      {
-        "name": "src",
-        "arg": "condition",
-        "semantic": "condition",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "offsets", "semantic": "offsets", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "flags", "semantic": "flags", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "blockSums", "semantic": "blockSums", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "n", "type": "u32", "value": "scanN" }] }
-      }
-    ],
-    "scanPrefix": [
-      { "name": "blockSums", "semantic": "blockSums", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      { "name": "blockPrefix", "semantic": "blockPrefix", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "numBlocks", "type": "u32", "value": "scanBlocks" }] }
-      }
-    ],
-    "flatScatter": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "offsets", "semantic": "offsets", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      { "name": "flags", "semantic": "flags", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      {
-        "name": "blockPrefix",
-        "semantic": "blockPrefix",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "output",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "n", "type": "u32", "value": "scanN" }] }
-      }
-    ],
-    "axisScatter": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "offsets", "semantic": "offsets", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      { "name": "flags", "semantic": "flags", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
-      {
-        "name": "blockPrefix",
-        "semantic": "blockPrefix",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "output",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "inputCount", "type": "u32", "value": "numel(shapes.input)" },
-            { "name": "inner", "type": "u32", "value": "inner(shapes.input, normalizedAxis)" },
-            { "name": "axisDim", "type": "u32", "value": "dim(shapes.input, normalizedAxis)" },
-            { "name": "n", "type": "u32", "value": "scanN" },
-            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.output, normalizedAxis)" }
-          ]
-        }
-      }
-    ],
-    "flatSerial": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "condition",
-        "arg": "condition",
-        "semantic": "condition",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "output",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "inputCount", "type": "u32", "value": "numel(shapes.input)" },
-            { "name": "conditionCount", "type": "u32", "value": "dim(shapes.condition, 0)" }
-          ]
-        }
-      }
-    ],
-    "axisSerial": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "condition",
-        "arg": "condition",
-        "semantic": "condition",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "output",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "conditionCount", "type": "u32", "value": "dim(shapes.condition, 0)" },
-            { "name": "outer", "type": "u32", "value": "outer(shapes.input, normalizedAxis)" },
-            { "name": "axisDim", "type": "u32", "value": "dim(shapes.input, normalizedAxis)" },
-            { "name": "inner", "type": "u32", "value": "inner(shapes.input, normalizedAxis)" },
-            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.output, normalizedAxis)" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "flatten_parallel_scan",
       "priority": 15,
-      "when": ["normalizedAxis == 2147483647", "ranks.input >= 1", "ranks.condition == 1", "ranks.output == 1", "numel(shapes.output) <= numel(shapes.input)", "numel(shapes.input) >= tunables.PARALLEL_MIN_ELEMENTS", "numel(shapes.input) <= tunables.MAX_PARALLEL_INPUT", "dim(shapes.condition, 0) > 0", "flatParallelFits", "f16Ok(dtypes.T)"],
       "derive": { "scanN": "flatScanN", "scanBlocks": "flatScanBlocks", "scanThreads": "flatScanThreads" },
       "intermediates": [
         { "id": "offsets", "dtype": "uint32", "shape": "[scanThreads]" },
@@ -243,40 +67,39 @@
         {
           "id": "flag_scan",
           "name": "Compress.FlagBlockScan",
-          "source": {
-            "shader": "scan-flags-block-exclusive.wgsl.jinja",
-            "inputs": {
-              "predicate": "\"src[i] != 0u\"",
-              "itemsPerThread": "scanItems",
-              "useSubgroups": "device.features.has(\"subgroups\")"
-            }
           },
-          "bindings": "scanBlocks",
-          "dispatch": { "workgroups": "scanBlocks" }
         },
         {
           "id": "block_prefix",
           "name": "Compress.BlockPrefixScan",
-          "source": {
-            "shader": "scan-block-prefix-u32.wgsl.jinja",
-            "inputs": { "useSubgroups": "device.features.has(\"subgroups\")" }
-          },
-          "bindings": "scanPrefix",
           "dispatch": { "x": 1 }
         },
         {
           "id": "scatter",
           "name": "Compress.ParallelScatter",
-          "source": { "shader": "compress-scatter.wgsl.jinja", "inputs": { "axisMode": false } },
-          "bindings": "flatScatter",
-          "dispatch": { "workgroups": "scanBlocks" }
         }
       ]
     },
     {
       "id": "axis_parallel_scan",
       "priority": 15,
-      "when": ["normalizedAxis != 2147483647", "ranks.input >= 1", "ranks.condition == 1", "ranks.output == ranks.input", "outer(shapes.output, normalizedAxis) == outer(shapes.input, normalizedAxis)", "inner(shapes.output, normalizedAxis) == inner(shapes.input, normalizedAxis)", "dim(shapes.output, normalizedAxis) <= dim(shapes.input, normalizedAxis)", "numel(shapes.input) >= tunables.PARALLEL_MIN_ELEMENTS", "numel(shapes.input) <= tunables.MAX_PARALLEL_INPUT", "dim(shapes.condition, 0) > 0", "axisParallelFits", "f16Ok(dtypes.T)"],
       "derive": { "scanN": "axisScanN", "scanBlocks": "axisScanBlocks", "scanThreads": "axisScanThreads" },
       "intermediates": [
         { "id": "offsets", "dtype": "uint32", "shape": "[scanThreads]" },
@@ -288,60 +111,104 @@
         {
           "id": "flag_scan",
           "name": "Compress.FlagBlockScan",
-          "source": {
-            "shader": "scan-flags-block-exclusive.wgsl.jinja",
-            "inputs": {
-              "predicate": "\"src[i] != 0u\"",
-              "itemsPerThread": "scanItems",
-              "useSubgroups": "device.features.has(\"subgroups\")"
-            }
           },
-          "bindings": "scanBlocks",
-          "dispatch": { "workgroups": "scanBlocks" }
         },
         {
           "id": "block_prefix",
           "name": "Compress.BlockPrefixScan",
-          "source": {
-            "shader": "scan-block-prefix-u32.wgsl.jinja",
-            "inputs": { "useSubgroups": "device.features.has(\"subgroups\")" }
-          },
-          "bindings": "scanPrefix",
           "dispatch": { "x": 1 }
         },
         {
           "id": "scatter",
           "name": "Compress.ParallelScatterAxis",
-          "source": { "shader": "compress-scatter.wgsl.jinja", "inputs": { "axisMode": true } },
-          "bindings": "axisScatter",
-          "dispatch": { "threads": "numel(shapes.input)", "workgroupSize": "workgroupSize" }
         }
       ]
     },
     {
       "id": "flatten_serial",
-      "when": ["normalizedAxis == 2147483647", "ranks.condition == 1", "ranks.output == 1", "numel(shapes.output) <= numel(shapes.input)", "f16Ok(dtypes.T)"],
-      "constants": { "axisMode": false },
       "passes": [
         {
           "id": "main",
           "name": "Compress",
           "shader": "compress.wgsl.jinja",
-          "bindings": "flatSerial",
           "dispatch": { "x": 1 }
         }
       ]
     },
     {
       "id": "axis_serial",
-      "when": ["normalizedAxis != 2147483647", "ranks.input >= 1", "ranks.condition == 1", "ranks.output == ranks.input", "outer(shapes.output, normalizedAxis) == outer(shapes.input, normalizedAxis)", "inner(shapes.output, normalizedAxis) == inner(shapes.input, normalizedAxis)", "dim(shapes.output, normalizedAxis) <= dim(shapes.input, normalizedAxis)", "f16Ok(dtypes.T)"],
-      "constants": { "axisMode": true },
       "passes": [
         {
           "id": "main",
           "name": "Compress",
           "shader": "compress.wgsl.jinja",
-          "bindings": "axisSerial",
           "dispatch": { "x": 1 }
         }
       ]

   "domain": "ai.onnx",
   "name": "Compress",
   "sinceVersion": 11,
+  "inputs": { "input": { "dtype": "T" }, "condition": { "dtype": "C", "rank": 1 } },
+  "outputs": { "output": { "dtype": "T", "rank": "ranks.input if has(attrs, \"axis\") else 1" } },
+  "attributes": { "axis": {} },
   "typeConstraints": { "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"], "C": ["bool"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "SCAN_ITEMS_PER_THREAD": { "default": 16 },
+    "PARALLEL_MIN_ELEMENTS": { "default": 512 },
+    "MAX_PARALLEL_INPUT": { "default": 1073741823 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "axisParam": "attrs.axis if has(attrs, \"axis\") else 2147483647",
     "normalizedAxis": "axisParam if axisParam >= 0 else axisParam + ranks.input",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "axisScanN": "min(dim(shapes.input, normalizedAxis), dim(shapes.condition, 0))",
     "axisScanBlocks": "ceilDiv(axisScanN, scanBlockItems)",
     "axisScanThreads": "axisScanBlocks * workgroupSize",
+    "axisScatterThreads": "outer(shapes.input, normalizedAxis) * axisScanN * inner(shapes.input, normalizedAxis)",
+    "axisParallelFits": "workgroupStorageOk and axisScanThreads * 4 <= storageBufferLimit and axisScanBlocks * 4 <= storageBufferLimit and axisScanBlocks <= foldedDispatchCapacity and ceilDiv(axisScatterThreads, workgroupSize) <= foldedDispatchCapacity"
   },
+  "when": ["ranks.condition == 1", "f16Ok(dtypes.T)"],
+  "bindings": {
+    "src": { "arg": "condition", "buffer": "read-only-storage", "elementType": "u32" },
+    "offsets": { "buffer": "storage", "elementType": "u32" },
+    "flags": { "buffer": "storage", "elementType": "u32" },
+    "blockSums": { "buffer": "storage", "elementType": "u32" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "n", "type": "u32", "value": "scanN" }] },
+    "blockSums_2": { "name": "blockSums", "buffer": "read-only-storage", "elementType": "u32" },
+    "blockPrefix": { "buffer": "storage", "elementType": "u32" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "numBlocks", "type": "u32", "value": "scanBlocks" }]
+    },
+    "offsets_2": { "name": "offsets", "buffer": "read-only-storage", "elementType": "u32" },
+    "flags_2": { "name": "flags", "buffer": "read-only-storage", "elementType": "u32" },
+    "blockPrefix_2": { "name": "blockPrefix", "buffer": "read-only-storage", "elementType": "u32" },
+    "condition": { "buffer": "read-only-storage", "elementType": "u32" }
   },
   "variants": [
     {
       "id": "flatten_parallel_scan",
       "priority": 15,
+      "when": ["normalizedAxis == 2147483647", "ranks.input >= 1", "ranks.output == 1", "numel(shapes.output) <= numel(shapes.input)", "numel(shapes.input) >= tunables.PARALLEL_MIN_ELEMENTS", "numel(shapes.input) <= tunables.MAX_PARALLEL_INPUT", "dim(shapes.condition, 0) > 0", "flatParallelFits"],
       "derive": { "scanN": "flatScanN", "scanBlocks": "flatScanBlocks", "scanThreads": "flatScanThreads" },
       "intermediates": [
         { "id": "offsets", "dtype": "uint32", "shape": "[scanThreads]" },
         {
           "id": "flag_scan",
           "name": "Compress.FlagBlockScan",
+          "shader": "scan-flags-block-exclusive.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": {
+            "predicate": "\"src[i] != 0u\"",
+            "itemsPerThread": "scanItems",
+            "useSubgroups": "device.features.has(\"subgroups\")"
           },
+          "bindings": ["src", "offsets", "flags", "blockSums", "params"],
+          "dispatch": { "x": "min(scanBlocks, 65535)", "y": "ceilDiv(scanBlocks, 65535)", "z": 1 }
         },
         {
           "id": "block_prefix",
           "name": "Compress.BlockPrefixScan",
+          "shader": "scan-block-prefix-u32.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["blockSums_2", "blockPrefix", "params_2"],
           "dispatch": { "x": 1 }
         },
         {
           "id": "scatter",
           "name": "Compress.ParallelScatter",
+          "shader": "compress-scatter.wgsl.jinja",
+          "derive": { "axisMode": false },
+          "bindings": ["input", "offsets_2", "flags_2", "blockPrefix_2", "output", "params"],
+          "dispatch": { "x": "min(scanBlocks, 65535)", "y": "ceilDiv(scanBlocks, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "axis_parallel_scan",
       "priority": 15,
+      "when": ["normalizedAxis != 2147483647", "ranks.input >= 1", "ranks.output == ranks.input", "outer(shapes.output, normalizedAxis) == outer(shapes.input, normalizedAxis)", "inner(shapes.output, normalizedAxis) == inner(shapes.input, normalizedAxis)", "dim(shapes.output, normalizedAxis) <= dim(shapes.input, normalizedAxis)", "numel(shapes.input) >= tunables.PARALLEL_MIN_ELEMENTS", "numel(shapes.input) <= tunables.MAX_PARALLEL_INPUT", "dim(shapes.condition, 0) > 0", "axisParallelFits"],
       "derive": { "scanN": "axisScanN", "scanBlocks": "axisScanBlocks", "scanThreads": "axisScanThreads" },
       "intermediates": [
         { "id": "offsets", "dtype": "uint32", "shape": "[scanThreads]" },
         {
           "id": "flag_scan",
           "name": "Compress.FlagBlockScan",
+          "shader": "scan-flags-block-exclusive.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": {
+            "predicate": "\"src[i] != 0u\"",
+            "itemsPerThread": "scanItems",
+            "useSubgroups": "device.features.has(\"subgroups\")"
           },
+          "bindings": ["src", "offsets", "flags", "blockSums", "params"],
+          "dispatch": { "x": "min(scanBlocks, 65535)", "y": "ceilDiv(scanBlocks, 65535)", "z": 1 }
         },
         {
           "id": "block_prefix",
           "name": "Compress.BlockPrefixScan",
+          "shader": "scan-block-prefix-u32.wgsl.jinja",
+          "subgroupCollectivesWidth": "portable",
+          "derive": { "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["blockSums_2", "blockPrefix", "params_2"],
           "dispatch": { "x": 1 }
         },
         {
           "id": "scatter",
           "name": "Compress.ParallelScatterAxis",
+          "shader": "compress-scatter.wgsl.jinja",
+          "derive": { "axisMode": true },
+          "bindings": [
+            "input",
+            "offsets_2",
+            "flags_2",
+            "blockPrefix_2",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "scanCount", "type": "u32", "value": "axisScatterThreads" },
+                { "name": "inner", "type": "u32", "value": "inner(shapes.input, normalizedAxis)" },
+                { "name": "axisDim", "type": "u32", "value": "dim(shapes.input, normalizedAxis)" },
+                { "name": "n", "type": "u32", "value": "scanN" },
+                { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.output, normalizedAxis)" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((axisScatterThreads), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((axisScatterThreads), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "flatten_serial",
+      "when": ["normalizedAxis == 2147483647", "ranks.output == 1", "numel(shapes.output) <= numel(shapes.input)"],
+      "derive": { "axisMode": false },
       "passes": [
         {
           "id": "main",
           "name": "Compress",
           "shader": "compress.wgsl.jinja",
+          "bindings": [
+            "input",
+            "condition",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "inputCount", "type": "u32", "value": "numel(shapes.input)" },
+                { "name": "conditionCount", "type": "u32", "value": "dim(shapes.condition, 0)" }
+              ]
+            }
+          ],
           "dispatch": { "x": 1 }
         }
       ]
     },
     {
       "id": "axis_serial",
+      "when": ["normalizedAxis != 2147483647", "ranks.input >= 1", "ranks.output == ranks.input", "outer(shapes.output, normalizedAxis) == outer(shapes.input, normalizedAxis)", "inner(shapes.output, normalizedAxis) == inner(shapes.input, normalizedAxis)", "dim(shapes.output, normalizedAxis) <= dim(shapes.input, normalizedAxis)"],
+      "derive": { "axisMode": true },
       "passes": [
         {
           "id": "main",
           "name": "Compress",
           "shader": "compress.wgsl.jinja",
+          "bindings": [
+            "input",
+            "condition",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "conditionCount", "type": "u32", "value": "dim(shapes.condition, 0)" },
+                { "name": "outer", "type": "u32", "value": "outer(shapes.input, normalizedAxis)" },
+                { "name": "axisDim", "type": "u32", "value": "dim(shapes.input, normalizedAxis)" },
+                { "name": "inner", "type": "u32", "value": "inner(shapes.input, normalizedAxis)" },
+                { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.output, normalizedAxis)" }
+              ]
+            }
+          ],
           "dispatch": { "x": 1 }
         }
       ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,29 @@
 {
   "name": "ai.onnx.Compress",
-  "id": "_ai_onnx_compress_webgpu_cfbc4b5",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "eyLAKUGnwFcu+exPotgJl/ZViuzCu2hWDco3P8Cs9ps=",
-      "compress-scatter.wgsl.jinja": "T6PIeYVTixXpDnRTaQB/iTAZ0LHYZecAJqUnWr3Ww0I=",
-      "compress.wgsl.jinja": "2vtjJ6EHEsVl08tW1O1kOZTR0VuRxsw77mSMlYyuyyk=",
-      "manifest.json": "DsMO6552aeUXGp8WmmmTQMetnVzS/JXRiiEk7Hh0mLM=",
-      "scan-block-prefix-u32.wgsl.jinja": "UJRN3G2HczCxzkEaCkPruYSaNTw+dJJURbUiqmhTncU=",
-      "scan-flags-block-exclusive.wgsl.jinja": "KiTJI6k5iFHfkj8xV5C/8g/BLcMqPewng9bBqRJXNT0=",
-      "test.json": "MOnY26q5C+8aOvNgOQ3gWbu0CX7H1Huwg18qu4bot9M="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Compress" }
 }

 {
   "name": "ai.onnx.Compress",
+  "id": "_ai_onnx_compress_webgpu_ab772ca",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "+YEsjnASlQKDd+UVLHTmWVfYB68qJ7UzCl0HyYPytuA=",
+      "compress-scatter.wgsl.jinja": "/CDh8XVC6X8AL4VRCKYwsP4weVOSIA0jEggep0Opngw=",
+      "compress.wgsl.jinja": "H8cjLOeT5Nwk0Yxd0B/tjSTLtIV5OXbREitywwvQa2M=",
+      "manifest.json": "lhknKZZblyy6aWk5dGgODBL0X0i+KsP7etkPQbT7avk=",
+      "scan-block-prefix-u32.wgsl.jinja": "XxgrEAHcIrSr96cR8fvJxzmBGwSWaY2LNXsJM3ypxAI=",
+      "scan-flags-block-exclusive.wgsl.jinja": "xCIsrwQGNSLkCMtk1hS27aVhDGhdQOddrg86QCegerY=",
+      "test.json": "FroEGRUKQpjx/c0R8AMl6YB/dX5s3diS3YOCOZ4Fi90="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "flatten_parallel_scan": ["compress-scatter.wgsl.jinja", "scan-block-prefix-u32.wgsl.jinja", "scan-flags-block-exclusive.wgsl.jinja"],
+      "axis_parallel_scan": ["compress-scatter.wgsl.jinja", "scan-block-prefix-u32.wgsl.jinja", "scan-flags-block-exclusive.wgsl.jinja"],
+      "flatten_serial": ["compress.wgsl.jinja"],
+      "axis_serial": ["compress.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/scan-block-prefix-u32.wgsl.jinja CHANGED Viewed

@@ -3,24 +3,63 @@
 // scans each chunk with Hillis-Steele or subgroup collectives, and links chunks
 // with a running carry. blockPrefix[b] is therefore the number of set flags in
 // all blocks before b. The total count is the last prefix plus the last sum.
-{% if source.useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
-{% if source.useSubgroups %}
 var<workgroup> sgTotals: array<u32, WG>;
 {% else %}
 var<workgroup> wgScan: array<u32, WG>;
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(local_invocation_id) lid: vec3<u32>{% if source.useSubgroups %},
-        @builtin(subgroup_invocation_id) sgLane: u32,
-        @builtin(subgroup_size) sgSize: u32{% endif %}) {
   let tid = lid.x;
   var carry = 0u;
   let chunks = (params.numBlocks + WG - 1u) / WG;
   for (var c = 0u; c < chunks; c = c + 1u) {
@@ -29,14 +68,11 @@ fn main(@builtin(local_invocation_id) lid: vec3<u32>{% if source.useSubgroups %}
     if (j < params.numBlocks) {
       v = blockSums[j];
     }
-{% if source.useSubgroups %}
     let subgroupInclusive = subgroupInclusiveAdd(v);
     let subgroupTotal = subgroupAdd(v);
-    let safeSg = max(sgSize, 1u);
-    let subgroupId = tid / safeSg;
-    let subgroupCount = (WG + safeSg - 1u) / safeSg;
-    if (sgLane == 0u) {
-      sgTotals[subgroupId] = subgroupTotal;
     }
     workgroupBarrier();
     var subgroupOffset = 0u;
@@ -44,7 +80,7 @@ fn main(@builtin(local_invocation_id) lid: vec3<u32>{% if source.useSubgroups %}
     for (var s = 0u; s < subgroupCount; s = s + 1u) {
       let total = sgTotals[s];
       chunkTotal = chunkTotal + total;
-      if (s < subgroupId) {
         subgroupOffset = subgroupOffset + total;
       }
     }

 // scans each chunk with Hillis-Steele or subgroup collectives, and links chunks
 // with a running carry. blockPrefix[b] is therefore the number of set flags in
 // all blocks before b. The total count is the last prefix plus the last sum.
+{% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+{% if useSubgroups %}
+{% set skipLogicalLaneCount = skipLogicalLaneCount is defined and skipLogicalLaneCount %}
+var<workgroup> sgLaneClaims: atomic<u32>;
+struct SubgroupLogicalLanes {
+  ltid: u32,
+  ord: u32,
+  rank: u32,
+  count: u32,
+}
+fn subgroup_logical_lanes() -> SubgroupLogicalLanes {
+  let rank = subgroupExclusiveAdd(1u);
+  let count = subgroupAdd(1u);
+  // Every lane performs the atomic so no collective follows a lane guard (a
+  // subgroup op after a closed `if (rank == 0u)` block is the reconvergence
+  // hazard the render gate flags): only the rank-0 lane adds its subgroup's
+  // claim, every other lane adds 0 and discards its snapshot. The rank-0 lane is
+  // the lowest active lane, which is the lane `subgroupBroadcastFirst` reads.
+  let ticket = atomicAdd(&sgLaneClaims, select(0u, count | (1u << 16u), rank == 0u));
+  let claim = subgroupBroadcastFirst(ticket);
+  return SubgroupLogicalLanes((claim & 0xffffu) + rank, claim >> 16u, rank, count);
+}
+{% if not skipLogicalLaneCount %}
+fn subgroup_logical_count() -> u32 {
+  return atomicLoad(&sgLaneClaims) >> 16u;
+}
+{% endif %}
+// One partial per subgroup, addressed by the subgroup's logical ordinal, which
+// is below the number of subgroups and therefore below WG at any width.
 var<workgroup> sgTotals: array<u32, WG>;
 {% else %}
 var<workgroup> wgScan: array<u32, WG>;
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main({% if not useSubgroups %}@builtin(local_invocation_id) lid: vec3<u32>{% endif %}) {
+{% if useSubgroups %}
+  // Logical lane coordinates replace local_invocation_index: `ltid` is the data
+  // position this invocation owns and `ord` its subgroup's carry ordinal, so the
+  // scan order is independent of how the device partitions the workgroup.
+  let L = subgroup_logical_lanes();
+  let tid = L.ltid;
+  workgroupBarrier();
+  let subgroupCount = subgroup_logical_count();
+{% else %}
   let tid = lid.x;
+{% endif %}
   var carry = 0u;
   let chunks = (params.numBlocks + WG - 1u) / WG;
   for (var c = 0u; c < chunks; c = c + 1u) {
     if (j < params.numBlocks) {
       v = blockSums[j];
     }
+{% if useSubgroups %}
     let subgroupInclusive = subgroupInclusiveAdd(v);
     let subgroupTotal = subgroupAdd(v);
+    if (L.rank == 0u) {
+      sgTotals[L.ord] = subgroupTotal;
     }
     workgroupBarrier();
     var subgroupOffset = 0u;
     for (var s = 0u; s < subgroupCount; s = s + 1u) {
       let total = sgTotals[s];
       chunkTotal = chunkTotal + total;
+      if (s < L.ord) {
         subgroupOffset = subgroupOffset + total;
       }
     }

build/webgpu/scan-flags-block-exclusive.wgsl.jinja CHANGED Viewed

@@ -8,33 +8,68 @@
 // predicate mask, and blockSums stores the block population. After blockSums is
 // scanned, item k in lane t of block b lands at:
 // blockPrefix[b] + offsets[t] + countOneBits(flags[t] & ((1 << k) - 1)).
-{% if source.useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
-const ITEMS: u32 = {{ source.itemsPerThread }}u;
 const BLOCK_ITEMS: u32 = WG * ITEMS;
-{% if source.useSubgroups %}
-// Only the first ceil(WG / subgroup_size) entries are used. Keeping the array
-// device-independent lets one shader cover every legal subgroup width.
 var<workgroup> sgTotals: array<u32, WG>;
 {% else %}
 var<workgroup> wgScan: array<u32, WG>;
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>,
-        @builtin(local_invocation_id) lid: vec3<u32>{% if source.useSubgroups %},
-        @builtin(subgroup_invocation_id) sgLane: u32,
-        @builtin(subgroup_size) sgSize: u32{% endif %}) {
   let tid = lid.x;
-  // Linearize the 2D dispatch so numBlocks can exceed the maxComputeWorkgroupsPerDimension grid cap: the
-  // Dispatch folds blocks past maxComputeWorkgroupsPerDimension into y: block = x + y * gridWidth.
-  let block = wg.x + wg.y * nwg.x;
   let base = block * BLOCK_ITEMS + tid * ITEMS;
   var packed = 0u;
@@ -43,7 +78,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     let i = base + k;
     var flag = 0u;
     if (i < params.n) {
-      if ({{ source.predicate }}) {
         flag = 1u;
       }
     }
@@ -51,19 +86,17 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>,
     laneCount = laneCount + flag;
   }
-{% if source.useSubgroups %}
   // Hardware scans each subgroup, then one shared-memory rendezvous links the
   // handful of subgroup totals. This replaces log2(WG) two-barrier sweeps.
   let subgroupInclusive = subgroupInclusiveAdd(laneCount);
   let subgroupTotal = subgroupAdd(laneCount);
-  let safeSg = max(sgSize, 1u);
-  let subgroupId = tid / safeSg;
-  if (sgLane == 0u) {
-    sgTotals[subgroupId] = subgroupTotal;
   }
   workgroupBarrier();
   var subgroupOffset = 0u;
-  for (var s = 0u; s < subgroupId; s = s + 1u) {
     subgroupOffset = subgroupOffset + sgTotals[s];
   }
   let inclusive = subgroupOffset + subgroupInclusive;

 // predicate mask, and blockSums stores the block population. After blockSums is
 // scanned, item k in lane t of block b lands at:
 // blockPrefix[b] + offsets[t] + countOneBits(flags[t] & ((1 << k) - 1)).
+{% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+const ITEMS: u32 = {{ itemsPerThread }}u;
 const BLOCK_ITEMS: u32 = WG * ITEMS;
+{% if useSubgroups %}
+{% set skipLogicalLaneCount = true %}
+{% set skipLogicalLaneCount = skipLogicalLaneCount is defined and skipLogicalLaneCount %}
+var<workgroup> sgLaneClaims: atomic<u32>;
+struct SubgroupLogicalLanes {
+  ltid: u32,
+  ord: u32,
+  rank: u32,
+  count: u32,
+}
+fn subgroup_logical_lanes() -> SubgroupLogicalLanes {
+  let rank = subgroupExclusiveAdd(1u);
+  let count = subgroupAdd(1u);
+  // Every lane performs the atomic so no collective follows a lane guard (a
+  // subgroup op after a closed `if (rank == 0u)` block is the reconvergence
+  // hazard the render gate flags): only the rank-0 lane adds its subgroup's
+  // claim, every other lane adds 0 and discards its snapshot. The rank-0 lane is
+  // the lowest active lane, which is the lane `subgroupBroadcastFirst` reads.
+  let ticket = atomicAdd(&sgLaneClaims, select(0u, count | (1u << 16u), rank == 0u));
+  let claim = subgroupBroadcastFirst(ticket);
+  return SubgroupLogicalLanes((claim & 0xffffu) + rank, claim >> 16u, rank, count);
+}
+{% if not skipLogicalLaneCount %}
+fn subgroup_logical_count() -> u32 {
+  return atomicLoad(&sgLaneClaims) >> 16u;
+}
+{% endif %}
+// One partial per subgroup, addressed by its logical ordinal. Only as many
+// entries as there are subgroups are used; WG bounds that at every width.
 var<workgroup> sgTotals: array<u32, WG>;
 {% else %}
 var<workgroup> wgScan: array<u32, WG>;
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>{% if not useSubgroups %},
+        @builtin(local_invocation_id) lid: vec3<u32>{% endif %}) {
+{% if useSubgroups %}
+  // Logical lane coordinates replace local_invocation_index: `ltid` is the data
+  // position this invocation owns and `ord` its subgroup's scan ordinal, so the
+  // block scan order is independent of how the device partitions the workgroup.
+  let L = subgroup_logical_lanes();
+  let tid = L.ltid;
+{% else %}
   let tid = lid.x;
+{% endif %}
+  // Linearize the 2D dispatch so numBlocks can exceed the per-axis dispatch fold width grid cap: the
+  // Dispatch folds blocks past the per-axis dispatch fold width into y: block = x + y * gridWidth.
+  let block = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let base = block * BLOCK_ITEMS + tid * ITEMS;
   var packed = 0u;
     let i = base + k;
     var flag = 0u;
     if (i < params.n) {
+      if ({{ predicate }}) {
         flag = 1u;
       }
     }
     laneCount = laneCount + flag;
   }
+{% if useSubgroups %}
   // Hardware scans each subgroup, then one shared-memory rendezvous links the
   // handful of subgroup totals. This replaces log2(WG) two-barrier sweeps.
   let subgroupInclusive = subgroupInclusiveAdd(laneCount);
   let subgroupTotal = subgroupAdd(laneCount);
+  if (L.rank == 0u) {
+    sgTotals[L.ord] = subgroupTotal;
   }
   workgroupBarrier();
   var subgroupOffset = 0u;
+  for (var s = 0u; s < L.ord; s = s + 1u) {
     subgroupOffset = subgroupOffset + sgTotals[s];
   }
   let inclusive = subgroupOffset + subgroupInclusive;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Compress",
   "cases": [
     {
       "name": "int16_axis_copy_boundaries",
@@ -173,7 +172,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/compress_op.test.cc",
         "test": "CompressTest.Compress_default_axis_issue_9247_cumulative_sum_overflow",
-        "notes": "Extends ORT's >127-element cumulative-sum regression above 1024 elements so the WebGPU parallel prefix-scan path is exercised."
       },
       "inputs": {
         "input": { "dtype": "float32", "shape": [23, 50], "data": { "kind": "linspace", "start": 0.0, "end": 1149.0 } },
@@ -394,7 +393,7 @@
     {
       "name": "int8_payload_axis1_parallel_scatter",
       "provenance": {
-        "notes": "int8 activation payload through the axis_parallel_scan scatter (numel 12288 > 1024). int8 stores as i32 per scalarType, so the $scalar copy path is exercised with the full s8 range incl -128/127. ORT CPU runs Compress for int8."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -425,7 +424,7 @@
     {
       "name": "int32_large_magnitude_above_2pow24_flatten",
       "provenance": {
-        "notes": "int32 payload with magnitudes above 2^24 (incl ±2e9) through flatten_parallel_scan. Verifies the i32 scatter copy preserves values that would be lossy if carried through f32; comparison uses the int32 dtype so it stays exact. ORT CPU runs Compress for int32."
       },
       "inputs": {
         "input": {

 {
   "cases": [
     {
       "name": "int16_axis_copy_boundaries",
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/compress_op.test.cc",
         "test": "CompressTest.Compress_default_axis_issue_9247_cumulative_sum_overflow",
+        "notes": "An all-true condition over more than 1,024 elements exercises the parallel prefix scan while preserving every input value."
       },
       "inputs": {
         "input": { "dtype": "float32", "shape": [23, 50], "data": { "kind": "linspace", "start": 0.0, "end": 1149.0 } },
     {
       "name": "int8_payload_axis1_parallel_scatter",
       "provenance": {
+        "notes": "A 12,288-element int8 activation tensor exercises axis-parallel scan and scatter. The package stores logical int8 values in int32 slots, and the cyclic payload includes the full signed-byte range from -128 to 127."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "int32_large_magnitude_above_2pow24_flatten",
       "provenance": {
+        "notes": "int32 payload with magnitudes above 2^24 (incl ±2e9) through flatten_parallel_scan. Verifies the i32 scatter copy preserves values that would be lossy if carried through f32; comparison uses the int32 dtype so it stays exact."
       },
       "inputs": {
         "input": {