sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +2 -2
build/webgpu/cumprod-tiled-short-rows.wgsl.jinja +2 -1
build/webgpu/manifest.json +11 -24
build/webgpu/metadata.json +9 -9
build/webgpu/scan-row-chunked-vec4.wgsl.jinja +13 -17
build/webgpu/scan-row-chunked.wgsl.jinja +14 -17
build/webgpu/scan-thread-row.wgsl.jinja +9 -6
build/webgpu/test.json +5 -7

README.md CHANGED Viewed

@@ -54,7 +54,7 @@ Default values (overridable per request):
 - [`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
 - [`cumprod-tiled-short-rows.wgsl.jinja`](build/webgpu/cumprod-tiled-short-rows.wgsl.jinja)
 - [`scan-row-chunked-vec4.wgsl.jinja`](build/webgpu/scan-row-chunked-vec4.wgsl.jinja)
 - [`scan-row-chunked.wgsl.jinja`](build/webgpu/scan-row-chunked.wgsl.jinja)
@@ -63,7 +63,7 @@ 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.

 - [`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
 - [`cumprod-tiled-short-rows.wgsl.jinja`](build/webgpu/cumprod-tiled-short-rows.wgsl.jinja)
 - [`scan-row-chunked-vec4.wgsl.jinja`](build/webgpu/scan-row-chunked-vec4.wgsl.jinja)
 - [`scan-row-chunked.wgsl.jinja`](build/webgpu/scan-row-chunked.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/cumprod-tiled-short-rows.wgsl.jinja CHANGED Viewed

@@ -15,7 +15,8 @@ var<workgroup> tile: array<{{ accType }}, {{ tileElements }}>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
-  let rowBase = wg.x * ROWS_PER_WG;
   if (rowBase >= params.rows) {
     return;
   }

 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-folded tile index: wg.y carries the high bits past the dispatch limit.
+  let rowBase = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * ROWS_PER_WG;
   if (rowBase >= params.rows) {
     return;
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -15,34 +15,32 @@
     "SHORT_ROW_TILE_STORAGE_TARGET": { "default": 16384 }
   },
   "derive": {
     "shortRowsPerWorkgroup": "max(1, min(tunables.WORKGROUP_SIZE, floor(min(tunables.SHORT_ROW_TILE_STORAGE_TARGET, device.limits.maxComputeWorkgroupStorageSize) / ((dim(shapes.x, args.axis) + 1) * 4))))",
     "shortRowTileScan": "inner(shapes.x, args.axis) == 1 and dim(shapes.x, args.axis) >= 2 and dim(shapes.x, args.axis) <= tunables.SERIAL_SHORT_DIM_MAX and rows(shapes.x, args.axis) >= tunables.SERIAL_MIN_ROWS and ceilDiv(rows(shapes.x, args.axis), shortRowsPerWorkgroup) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "cooperativeRowScan": "inner(shapes.x, args.axis) == 1 and dim(shapes.x, args.axis) > 32 and rows(shapes.x, args.axis) >= 1 and rows(shapes.x, args.axis) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and not shortRowTileScan"
   },
   "when": ["f16Ok(dtypes.T)", "numel(shapes.x) == numel(shapes.y)", "ranks.x >= 1", "ranks.y == ranks.x"],
   "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$scalar" },
-    "y": { "buffer": "storage", "elementType": "$scalar" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, args.axis)" },
         { "name": "dim", "type": "u32", "value": "dim(shapes.x, args.axis)" }
       ]
     },
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$vec4Scalar" },
-    "y_2": { "name": "y", "buffer": "storage", "elementType": "$vec4Scalar" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, args.axis)" },
         { "name": "dim4", "type": "u32", "value": "dim(shapes.x, args.axis) / 4" }
       ]
     },
-    "params_3": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, args.axis)" },
         { "name": "dim", "type": "u32", "value": "dim(shapes.x, args.axis)" },
@@ -58,7 +56,6 @@
       "derive": {
         "scalar": "dtypes.T",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
         "workgroupSize": "tunables.WORKGROUP_SIZE",
         "rowsPerWorkgroup": "shortRowsPerWorkgroup",
         "rowStride": "dim(shapes.x, args.axis) + 1",
@@ -87,7 +84,6 @@
         "scalar": "dtypes.T",
         "vec4Scalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
         "workgroupSize": "tunables.WORKGROUP_SIZE"
       },
       "passes": [
@@ -95,13 +91,8 @@
           "id": "main",
           "name": "CumProd.WorkgroupRowScanVec4",
           "shader": "scan-row-chunked-vec4.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "derive": {
-            "op": "\"mul\"",
-            "exclusive": "attrs.exclusive != 0",
-            "useSubgroups": "device.features.has(\"subgroups\")"
-          },
-          "bindings": ["x_2", "y_2", "params_2"],
           "dispatch": { "x": "rows(shapes.x, args.axis)" }
         }
       ]
@@ -113,7 +104,6 @@
       "derive": {
         "scalar": "dtypes.T",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
         "workgroupSize": "tunables.WORKGROUP_SIZE"
       },
       "passes": [
@@ -121,9 +111,7 @@
           "id": "main",
           "name": "CumProd.WorkgroupRowScan",
           "shader": "scan-row-chunked.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
           "derive": {
-            "op": "\"mul\"",
             "exclusive": "attrs.exclusive != 0",
             "reverse": "attrs.reverse != 0",
             "useSubgroups": "device.features.has(\"subgroups\")"
@@ -140,7 +128,6 @@
       "derive": {
         "scalar": "dtypes.T",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
         "workgroupSize": "tunables.WORKGROUP_SIZE"
       },
       "passes": [
@@ -148,8 +135,8 @@
           "id": "main",
           "name": "CumProd.ThreadRowScan",
           "shader": "scan-thread-row.wgsl.jinja",
-          "derive": { "op": "\"mul\"", "exclusive": "attrs.exclusive != 0", "reverse": "attrs.reverse != 0" },
-          "bindings": ["x", "y", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((rows(shapes.x, args.axis)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((rows(shapes.x, args.axis)), (tunables.WORKGROUP_SIZE)), 65535)",

     "SHORT_ROW_TILE_STORAGE_TARGET": { "default": 16384 }
   },
   "derive": {
+    "op": "\"mul\"",
     "shortRowsPerWorkgroup": "max(1, min(tunables.WORKGROUP_SIZE, floor(min(tunables.SHORT_ROW_TILE_STORAGE_TARGET, device.limits.maxComputeWorkgroupStorageSize) / ((dim(shapes.x, args.axis) + 1) * 4))))",
     "shortRowTileScan": "inner(shapes.x, args.axis) == 1 and dim(shapes.x, args.axis) >= 2 and dim(shapes.x, args.axis) <= tunables.SERIAL_SHORT_DIM_MAX and rows(shapes.x, args.axis) >= tunables.SERIAL_MIN_ROWS and ceilDiv(rows(shapes.x, args.axis), shortRowsPerWorkgroup) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "cooperativeRowScan": "inner(shapes.x, args.axis) == 1 and dim(shapes.x, args.axis) > 32 and rows(shapes.x, args.axis) >= 1 and rows(shapes.x, args.axis) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and not shortRowTileScan"
   },
   "when": ["f16Ok(dtypes.T)", "numel(shapes.x) == numel(shapes.y)", "ranks.x >= 1", "ranks.y == ranks.x"],
   "bindings": {
+    "x": { "elementType": "$scalar" },
+    "y": { "elementType": "$scalar" },
     "params": {
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, args.axis)" },
         { "name": "dim", "type": "u32", "value": "dim(shapes.x, args.axis)" }
       ]
     },
+    "x_main": { "name": "x", "elementType": "$vec4Scalar" },
+    "y_main": { "name": "y", "elementType": "$vec4Scalar" },
+    "params_main": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, args.axis)" },
         { "name": "dim4", "type": "u32", "value": "dim(shapes.x, args.axis) / 4" }
       ]
     },
+    "params__uniform": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "rows(shapes.x, args.axis)" },
         { "name": "dim", "type": "u32", "value": "dim(shapes.x, args.axis)" },
       "derive": {
         "scalar": "dtypes.T",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
         "workgroupSize": "tunables.WORKGROUP_SIZE",
         "rowsPerWorkgroup": "shortRowsPerWorkgroup",
         "rowStride": "dim(shapes.x, args.axis) + 1",
         "scalar": "dtypes.T",
         "vec4Scalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
         "workgroupSize": "tunables.WORKGROUP_SIZE"
       },
       "passes": [
           "id": "main",
           "name": "CumProd.WorkgroupRowScanVec4",
           "shader": "scan-row-chunked-vec4.wgsl.jinja",
+          "derive": { "exclusive": "attrs.exclusive != 0", "useSubgroups": "device.features.has(\"subgroups\")" },
+          "bindings": ["x_main", "y_main", "params_main"],
           "dispatch": { "x": "rows(shapes.x, args.axis)" }
         }
       ]
       "derive": {
         "scalar": "dtypes.T",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
         "workgroupSize": "tunables.WORKGROUP_SIZE"
       },
       "passes": [
           "id": "main",
           "name": "CumProd.WorkgroupRowScan",
           "shader": "scan-row-chunked.wgsl.jinja",
           "derive": {
             "exclusive": "attrs.exclusive != 0",
             "reverse": "attrs.reverse != 0",
             "useSubgroups": "device.features.has(\"subgroups\")"
       "derive": {
         "scalar": "dtypes.T",
         "accType": "\"f32\" if (dtypes.T == \"f32\" or dtypes.T == \"f16\") else dtypes.T",
         "workgroupSize": "tunables.WORKGROUP_SIZE"
       },
       "passes": [
           "id": "main",
           "name": "CumProd.ThreadRowScan",
           "shader": "scan-thread-row.wgsl.jinja",
+          "derive": { "exclusive": "attrs.exclusive != 0", "reverse": "attrs.reverse != 0" },
+          "bindings": ["x", "y", "params__uniform"],
           "dispatch": {
             "x": "min(ceilDiv((rows(shapes.x, args.axis)), (tunables.WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((rows(shapes.x, args.axis)), (tunables.WORKGROUP_SIZE)), 65535)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "ai.onnx.CumProd",
-  "id": "_ai_onnx_cumprod_webgpu_feb0f63",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,17 +8,17 @@
     "algorithm": "sha256",
     "files": {
       "bench.json": "TsJ08po3IF5cQJUCRPrMoIhyLOTSsyuddPCr6fwUmMY=",
-      "cumprod-tiled-short-rows.wgsl.jinja": "zXgZLN3glP7YaO+4AVVr/b/xOX5jvkeUb+2bxES00DU=",
-      "manifest.json": "goNH2mdRcylLp6WBDSFf+2rT7cryeWckBuH/GMfZeiQ=",
-      "scan-row-chunked-vec4.wgsl.jinja": "9KkZD/utgA1D7Ns/f9SlJXlJYAkIYeGg0XJhUQsfzeM=",
-      "scan-row-chunked.wgsl.jinja": "oPnVIjf8h3JeHcmFRCGFrN8oP6xlEixAES9S7vKzVlk=",
-      "scan-thread-row.wgsl.jinja": "YDNpvK2v9J6jLxurSlT+vfgfUyrEn1m1kQ/bWL95pY8=",
-      "test.json": "L08e2ck+A4AmlJOEsJ90/G/QX4H9HHctuYqvVVzS2nI="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "tiled_short_rows": ["cumprod-tiled-short-rows.wgsl.jinja"],
       "wg_row_scan_vec4": ["scan-row-chunked-vec4.wgsl.jinja"],

 {
   "name": "ai.onnx.CumProd",
+  "id": "_ai_onnx_cumprod_webgpu_7b33632",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "TsJ08po3IF5cQJUCRPrMoIhyLOTSsyuddPCr6fwUmMY=",
+      "cumprod-tiled-short-rows.wgsl.jinja": "UerkBDU7fq9KungxlAQv+GKFy2l+U3AfB73nyavLe6g=",
+      "manifest.json": "D1sJBLkaP1YYTvDvGZer83zCF4TrdGApLVGGm4S93jo=",
+      "scan-row-chunked-vec4.wgsl.jinja": "UP1vcu2Akr7EoVd7rLEMKuX4TwqvkDCJ8jCDTK/BkBQ=",
+      "scan-row-chunked.wgsl.jinja": "YvnQVPH5vCCEQjxZrUDEMDmg9NOuKivLA7gWJOS21NE=",
+      "scan-thread-row.wgsl.jinja": "y4xCv3ymxbBws1At0MSD2CSd5L/NWrBsmXsIiPrGH2c=",
+      "test.json": "xW22ETyy/XBxMFQhxEmZNC7De9w9caIDTsysYfEIvAY="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "tiled_short_rows": ["cumprod-tiled-short-rows.wgsl.jinja"],
       "wg_row_scan_vec4": ["scan-row-chunked-vec4.wgsl.jinja"],

build/webgpu/scan-row-chunked-vec4.wgsl.jinja CHANGED Viewed

@@ -13,11 +13,10 @@
 // so it is the reduction of every element strictly before this thread's block.
 // Inclusive output k is combine(tbase, p_k); exclusive output k uses p_{k-1},
 // with p_{-1} equal to the identity.
-{% set useSubgroups = useSubgroups %}
 {% if useSubgroups %}
 enable subgroups;
-{%- set sg_exc = "subgroupExclusiveMul" if op == "mul" else "subgroupExclusiveAdd" %}
-{% set sg_tot = "subgroupMul" if op == "mul" else "subgroupAdd" %}
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
@@ -29,7 +28,6 @@ fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
 }
 {% if useSubgroups %}
-{% set skipLogicalLaneCount = skipLogicalLaneCount is defined and skipLogicalLaneCount %}
 var<workgroup> sgLaneClaims: atomic<u32>;
 struct SubgroupLogicalLanes {
@@ -39,30 +37,28 @@ struct SubgroupLogicalLanes {
   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 %}
 {% endif %}
 var<workgroup> sgTotals: array<{{ accType }}, WG>;
 @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 %}) {
   let row = wg.x;
 {% if not useSubgroups %}
@@ -75,7 +71,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>{% if not useSubgroups %},
 {% if useSubgroups %}
   // `ltid` is the vec4 block this invocation owns and `ord` its subgroup's carry
   // ordinal; both are dense and independent of the device's partition.
-  let L = subgroup_logical_lanes();
   let tid = L.ltid;
   workgroupBarrier();
   let numSub = subgroup_logical_count();

 // so it is the reduction of every element strictly before this thread's block.
 // Inclusive output k is combine(tbase, p_k); exclusive output k uses p_{k-1},
 // with p_{-1} equal to the identity.
+{% set sg_exc = "subgroupExclusiveMul" %}
+{% set sg_tot = "subgroupMul" %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 }
 {% if useSubgroups %}
 var<workgroup> sgLaneClaims: atomic<u32>;
 struct SubgroupLogicalLanes {
   count: u32,
 }
+fn subgroup_logical_lanes(rank: u32, count: u32) -> SubgroupLogicalLanes {
+  // Every lane performs the atomic so no collective follows a lane guard: a
+  // subgroup op after a closed `if (rank == 0u)` block is a reconvergence
+  // hazard, since every subgroup operation must be reached in uniform control
+  // flow. Only the lane-0 lane adds its subgroup's claim, every other lane
+  // adds 0 and discards its snapshot. Lane 0 is the lowest 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);
 }
 fn subgroup_logical_count() -> u32 {
   return atomicLoad(&sgLaneClaims) >> 16u;
 }
 {% endif %}
 var<workgroup> sgTotals: array<{{ accType }}, WG>;
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>{% if useSubgroups %},
+        @builtin(subgroup_invocation_id) sgLane: u32,
+        @builtin(subgroup_size) sgWidth: u32{% else %},
         @builtin(local_invocation_id) lid: vec3<u32>{% endif %}) {
   let row = wg.x;
 {% if not useSubgroups %}
 {% if useSubgroups %}
   // `ltid` is the vec4 block this invocation owns and `ord` its subgroup's carry
   // ordinal; both are dense and independent of the device's partition.
+  let L = subgroup_logical_lanes(sgLane, sgWidth);
   let tid = L.ltid;
   workgroupBarrier();
   let numSub = subgroup_logical_count();

build/webgpu/scan-row-chunked.wgsl.jinja CHANGED Viewed

@@ -12,9 +12,9 @@
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
-{% set sg_inc = "subgroupInclusiveMul" if op == "mul" else "subgroupInclusiveAdd" %}
-{% set sg_exc = "subgroupExclusiveMul" if op == "mul" else "subgroupExclusiveAdd" %}
-{% set sg_tot = "subgroupMul" if op == "mul" else "subgroupAdd" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -25,7 +25,6 @@ fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
 }
 {% if useSubgroups %}
-{% set skipLogicalLaneCount = skipLogicalLaneCount is defined and skipLogicalLaneCount %}
 var<workgroup> sgLaneClaims: atomic<u32>;
 struct SubgroupLogicalLanes {
@@ -35,39 +34,37 @@ struct SubgroupLogicalLanes {
   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. A workgroup of WG
 // invocations holds at most WG subgroups at any adapter width, so WG is a sound
 // bound with no width guard.
 var<workgroup> sgTotals: array<{{ accType }}, WG>;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(workgroup_id) wg: vec3<u32>) {
   let row = wg.x;
   if (row >= params.rows) {
     return;
   }
   // `ltid` is the row position this invocation owns and `ord` its subgroup's
   // carry ordinal; both are dense and independent of the device's partition.
-  let L = subgroup_logical_lanes();
   let tid = L.ltid;
   workgroupBarrier();
   let numSub = subgroup_logical_count();

 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
+{% set sg_inc = "subgroupInclusiveMul" %}
+{% set sg_exc = "subgroupExclusiveMul" %}
+{% set sg_tot = "subgroupMul" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 }
 {% if useSubgroups %}
 var<workgroup> sgLaneClaims: atomic<u32>;
 struct SubgroupLogicalLanes {
   count: u32,
 }
+fn subgroup_logical_lanes(rank: u32, count: u32) -> SubgroupLogicalLanes {
+  // Every lane performs the atomic so no collective follows a lane guard: a
+  // subgroup op after a closed `if (rank == 0u)` block is a reconvergence
+  // hazard, since every subgroup operation must be reached in uniform control
+  // flow. Only the lane-0 lane adds its subgroup's claim, every other lane
+  // adds 0 and discards its snapshot. Lane 0 is the lowest 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);
 }
 fn subgroup_logical_count() -> u32 {
   return atomicLoad(&sgLaneClaims) >> 16u;
 }
 // One partial per subgroup, addressed by its logical ordinal. A workgroup of WG
 // invocations holds at most WG subgroups at any adapter width, so WG is a sound
 // bound with no width guard.
 var<workgroup> sgTotals: array<{{ accType }}, WG>;
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(subgroup_invocation_id) sgLane: u32,
+        @builtin(subgroup_size) sgWidth: u32) {
   let row = wg.x;
   if (row >= params.rows) {
     return;
   }
   // `ltid` is the row position this invocation owns and `ord` its subgroup's
   // carry ordinal; both are dense and independent of the device's partition.
+  let L = subgroup_logical_lanes(sgLane, sgWidth);
   let tid = L.ltid;
   workgroupBarrier();
   let numSub = subgroup_logical_count();

build/webgpu/scan-thread-row.wgsl.jinja CHANGED Viewed

@@ -3,6 +3,14 @@
 // is O(numel), versus O(numel * dim) for a naive per-output-element kernel.
 // At each scan step, consecutive threads access consecutive inner indices. The
 // accumulation order is element-by-element in scan order.
 {{ env.wgsl.resourceDeclarations }}
 const IDENTITY: {{ accType }} = {% if accType == "f32" %}1.0{% elif accType == "i32" %}1i{% else %}1u{% endif %};
@@ -13,12 +21,7 @@ fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
 @compute @workgroup_size({{ workgroupSize }}, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat row index: gid.y carries the high bits past the
-  // per-axis dispatch fold width (rows > 16.7M).
-  let r = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ workgroupSize }}u;
-  if (r >= params.rows) {
-    return;
-  }
   let outer_i = r / params.inner;
   let inner_i = r % params.inner;
   let base = outer_i * params.dim * params.inner + inner_i;

 // is O(numel), versus O(numel * dim) for a naive per-output-element kernel.
 // At each scan step, consecutive threads access consecutive inner indices. The
 // accumulation order is element-by-element in scan order.
+{% 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 IDENTITY: {{ accType }} = {% if accType == "f32" %}1.0{% elif accType == "i32" %}1i{% else %}1u{% endif %};
 @compute @workgroup_size({{ workgroupSize }}, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(workgroupSize, "r", "params.rows") }}
   let outer_i = r / params.inner;
   let inner_i = r % params.inner;
   let base = outer_i * params.dim * params.inner + inner_i;

build/webgpu/test.json CHANGED Viewed

@@ -12,7 +12,7 @@
     {
       "name": "many_short_rows_axis1_portable_tiled_f32",
       "provenance": {
-        "notes": "Many independent 64-element rows with nontrivial signed products exercise the capability-selected short-row route. Variable-width and minimum-limit tiers use the coalesced tiled path; compatible fixed-width subgroups use their native row scan."
       },
       "args": { "axis": 1 },
       "inputs": {
@@ -718,7 +718,7 @@
     {
       "name": "wg_scan_multichunk_pow2_f32",
       "provenance": {
-        "notes": "Axis length 600 forces the chunk-carry path (3 chunks of 256) in the workgroup row-scan."
       },
       "args": { "axis": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 600], "data": { "kind": "cycle", "values": [2.0, 0.5] } } },
@@ -848,7 +848,7 @@
     {
       "name": "clean_f16_last_axis_dim48_wg_scan",
       "provenance": {
-        "notes": "f16 innermost-axis cumprod, dim=48 (>32, %4==0) selects wg_row_scan (no f16 vec4 path; accType=f32). Realistic f16 model doing a cumulative product over the last dim; near-1 pow2-friendly values keep the f32-accumulate-then-f16-store result within f16 tolerance of the f64 reference."
       },
       "args": { "axis": 1 },
       "inputs": {
@@ -862,9 +862,7 @@
     },
     {
       "name": "clean_f32_last_axis_nonmult4_dim50_scalar",
-      "provenance": {
-        "notes": "A 50-element innermost axis exceeds the short-row threshold but is not divisible by four, selecting the scalar workgroup row scan."
-      },
       "args": { "axis": 1 },
       "inputs": {
         "x": {
@@ -878,7 +876,7 @@
     {
       "name": "clean_rank4_middle_axis_reverse_f32",
       "provenance": {
-        "notes": "A rank-4 reverse CumProd over a middle axis exercises outer/dimension/inner index decomposition. With inner=32, workgroup row paths are ineligible and serial accumulation has an exact expected result."
       },
       "attrs": { "reverse": 1 },
       "args": { "axis": 1 },

     {
       "name": "many_short_rows_axis1_portable_tiled_f32",
       "provenance": {
+        "notes": "Many independent 64-element rows with nontrivial signed products check inclusive CumProd across device tiers."
       },
       "args": { "axis": 1 },
       "inputs": {
     {
       "name": "wg_scan_multichunk_pow2_f32",
       "provenance": {
+        "notes": "Axis length 600 (last axis) spans three 256-element chunks, two full and one of 88, requiring the running product to carry across chunk boundaries. Exact powers of two (2.0, 0.5) keep the cumulative product exact in float32."
       },
       "args": { "axis": 1 },
       "inputs": { "x": { "dtype": "float32", "shape": [2, 600], "data": { "kind": "cycle", "values": [2.0, 0.5] } } },
     {
       "name": "clean_f16_last_axis_dim48_wg_scan",
       "provenance": {
+        "notes": "Float16 innermost-axis cumulative product over 48 values cycling through 1, 0.5, and 2; the running product stays at 1 or 0.5, so every output is exact in float16."
       },
       "args": { "axis": 1 },
       "inputs": {
     },
     {
       "name": "clean_f32_last_axis_nonmult4_dim50_scalar",
+      "provenance": { "notes": "A 50-element innermost axis checks a partial four-lane tail in inclusive CumProd." },
       "args": { "axis": 1 },
       "inputs": {
         "x": {
     {
       "name": "clean_rank4_middle_axis_reverse_f32",
       "provenance": {
+        "notes": "A rank-4 reverse cumulative product over axis 1, with inner size 32 (the product of the two trailing dimensions), decomposes indices into outer, axis, and inner components. Powers-of-two values (1, 2, 0.5) give an exact expected result in float32."
       },
       "attrs": { "reverse": 1 },
       "args": { "axis": 1 },