@graphty/webgpu-graph-algorithms 0.5.0 → 0.6.0

package/src/wgsl/scan-add.wgsl.ts

@@ -0,0 +1,13 @@
+/**
+ * The `scan-add` kernel body (spec 6 row 2 step (c); P4-T2): adds the exclusive prefix of its block's sum,
+ * `blockOffsets[group]`, to every element of the block. Body only; normative text.
+ */
+export const scanAddWgsl = /* wgsl */ `
+@compute @workgroup_size(WG)
+fn scan_add(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    let g = group_id(wid);
+    let i = g * WG + lid.x;
+    if (i >= P.count) { return; }                                  // no barrier follows
+    out[i] = out[i] + blockOffsets[g];
+}
+`;

package/src/wgsl/scan-block.wgsl.ts

@@ -0,0 +1,29 @@
+/**
+ * The `scan-block` kernel body (spec 6 row 2; P4-T2): an exclusive prefix sum of one WG-wide block of `src` in
+ * workgroup memory (Hillis-Steele, log2(WG) rounds of two barriers each, every lane in uniform control flow) into
+ * `out`, and the block's inclusive total into `blockSums[group]`. u32 addition is exact in any order, so the
+ * output is bitwise the same on every adapter (PD-3). Body only (spec 3.5, D9); normative text.
+ */
+export const scanBlockWgsl = /* wgsl */ `
+var<workgroup> sh: array<u32, WG>;
+
+@compute @workgroup_size(WG)
+fn scan_block(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    let g = group_id(wid);
+    let i = g * WG + lid.x;
+    var v = 0u;
+    if (i < P.count) { v = src[i]; }
+    sh[lid.x] = v;
+    workgroupBarrier();
+    for (var s = 1u; s < WG; s = s * 2u) {                       // Hillis-Steele inclusive scan; uniform: every lane runs every round
+        var t = 0u;
+        if (lid.x >= s) { t = sh[lid.x - s]; }
+        workgroupBarrier();
+        sh[lid.x] = sh[lid.x] + t;
+        workgroupBarrier();
+    }
+    let inclusive = sh[lid.x];
+    if (i < P.count) { out[i] = inclusive - v; }                  // exclusive = inclusive - own value
+    if (lid.x == WG - 1u) { blockSums[g] = inclusive; }           // the block total (the last lane's inclusive sum)
+}
+`;

package/src/wgsl/segmented-reduce.wgsl.ts

@@ -1,28 +1,39 @@
 /**
- * The `segmented-reduce` kernel body (spec 6 row 3; contract 4.5), TIER 0 = thread-per-row: one invocation per row
- * of [P.start, P.end) folds the VALUE snippet over the row's arcs inside the bound window [P.arcBase, P.arcEnd) and
- * writes out[i] as f32; a row with no arcs gets the identity element (0 for sum, F32_MAX for min, -F32_MAX for
- * max). The mid / high degree tiers (TIER 1 / 2 over degreeOrder()) land at P4. The body is normative: a sabotage
- * mutation (test/helpers/sabotage.ts) is a textual edit of it, so it is not restyled.
+ * The `segmented-reduce` kernel body (spec 6 row 3; contract 4.5) in its three degree tiers (P4-T5, PD-6). Every
+ * tier folds the VALUE snippet over a row's arcs inside the bound window [P.arcBase, P.arcEnd) and writes out[i]
+ * as f32 (a row with no arcs gets the identity element: 0 for sum, F32_MAX for min, -F32_MAX for max); the row is
+ * `perm[row]` under USE_PERM and `row` otherwise, over the rows [P.start, P.end) of the dispatch. TIER 0 is one row
+ * per thread (no barrier, so its early return is legal); TIER 1 is 32 lanes per row and WG / 32 rows per workgroup,
+ * reduced by a five-step tree in workgroup memory (bitwise the same on every subgroup size); TIER 2 is one row per
+ * workgroup through `wg_reduce_f32` of the prelude (the subgroup variant when the feature exists, the workgroup twin
+ * otherwise -- hence `needs: ["subgroups"]`). The tier bodies are functions called under `if (TIER == 0u)`, an
+ * override, so every barrier is reached in uniform control flow. The body is normative: a sabotage mutation
+ * (test/helpers/sabotage.ts) is a textual edit of it, so it is not restyled.
+ *
+ * TIER 0 folds its row through `row_fold_dense`, a stride-one copy of `row_fold`, because the shader compiler emits
+ * `row_fold(i, 0u, 1u)` as a call and leaves the stride in a parameter: the loop then walks the row with a runtime
+ * step, which costs the strength-reduced addressing into colIdx / weights and the unrolling that keeps several loads
+ * in flight per thread. PageRank's out-weight pass, the only caller that passes no tiers today, runs TIER 0 alone,
+ * and under the tiers TIER 0 still folds the low-degree rows, which are most of them. The two folds spell their
+ * locals apart (`lo` / `hi` / `k` against `a0` / `a1`) so that each sabotage row names exactly one of them; the
+ * VALUE snippet sees the same `row`, `arc`, `nbr` and `weight` in both.
  */
 
 /**
- * Entry point `segmented_reduce`; overrides OP (0 sum, 1 min, 2 max) and TIER (0); snippet slot VALUE: statements
- * assigning `v` from `row`, `arc`, `nbr`, `weight` (`target` is a WGSL reserved word, hence `nbr`).
+ * Entry point `segmented_reduce`; overrides OP (0 sum, 1 min, 2 max) and TIER (0 thread-per-row, 1 32-lanes-per-row,
+ * 2 workgroup-per-row); snippet slot VALUE: statements assigning `v` from `row`, `arc`, `nbr`, `weight` (`target`
+ * is a WGSL reserved word, hence `nbr`).
  */
 export const segmentedReduceWgsl = /* wgsl */ `
 fn identity() -> f32 { if (OP == 1u) { return F32_MAX; } if (OP == 2u) { return -F32_MAX; } return 0.0; }
 fn comb(a: f32, b: f32) -> f32 { if (OP == 1u) { return min(a, b); } if (OP == 2u) { return max(a, b); } return a + b; }
-
-@compute @workgroup_size(WG)
-fn segmented_reduce(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
-    let row = linear_id(wid, lid.x) + P.start;
-    if (row >= P.end) { return; }
-    let i = select(row, perm[row], USE_PERM);
+fn row_node(row: u32) -> u32 { return select(row, perm[row], USE_PERM); }
+fn row_fold(i: u32, lane: u32, step: u32) -> f32 {              // the arcs of row i this lane walks inside the bound window [P.arcBase, P.arcEnd)
+    let row = i;                                                 // the CSR row the VALUE snippet may name (the node index, under USE_PERM too)
     let a0 = max(rowPtr[i], P.arcBase);
     let a1 = min(rowPtr[i + 1u], P.arcEnd);
     var acc = identity();
-    for (var arc = a0; arc < a1; arc = arc + 1u) {
+    for (var arc = a0 + lane; arc < a1; arc = arc + step) {
         let nbr = colIdx[arc - P.arcBase];               // the neighbour index (\`target\` is a WGSL reserved word)
         var weight = 1.0;
         if (HAS_WEIGHTS) { weight = weights[arc - P.arcBase]; }
@@ -30,6 +41,64 @@ fn segmented_reduce(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invoca
         //@@VALUE@@
         acc = comb(acc, v);
     }
-    out[i] = select(acc, comb(out[i], acc), P.accumulate == 1u);
+    return acc;
+}
+fn row_fold_dense(i: u32) -> f32 {                              // TIER 0's stride-one twin of row_fold (see the header)
+    let row = i;                                                 // the CSR row the VALUE snippet may name (as in row_fold)
+    let lo = max(rowPtr[i], P.arcBase);
+    let hi = min(rowPtr[i + 1u], P.arcEnd);
+    var acc = identity();
+    for (var arc = lo; arc < hi; arc = arc + 1u) {
+        let k = arc - P.arcBase;                         // the window-local index; this walk is contiguous
+        let nbr = colIdx[k];
+        var weight = 1.0;
+        if (HAS_WEIGHTS) { weight = weights[k]; }
+        var v = 0.0;
+        //@@VALUE@@
+        acc = comb(acc, v);
+    }
+    return acc;
+}
+fn finish(i: u32, acc: f32) { out[i] = select(acc, comb(out[i], acc), P.accumulate == 1u); }
+fn tier0(wid: vec3<u32>, lane: u32) {                            // TIER 0: one row per thread over [P.start, P.end); no barrier, so the early return is legal (3.5 rule 1)
+    let row = linear_id(wid, lane) + P.start;
+    if (row >= P.end) { return; }
+    let i = row_node(row);
+    finish(i, row_fold_dense(i));
+}
+
+var<workgroup> sh: array<f32, WG>;
+
+fn tiered(wid: vec3<u32>, lid: u32) {                            // TIER 1: 32 lanes per row, WG / 32 rows per workgroup; TIER 2: WG lanes per row (PD-6)
+    let g = group_id(wid);
+    var row = P.start + g;
+    var lane = lid;
+    var step = WG;
+    if (TIER == 1u) { row = P.start + g * (WG / 32u) + lid / 32u; lane = lid % 32u; step = 32u; }
+    let valid = row < P.end;
+    var i = 0u;
+    var acc = identity();
+    if (valid) { i = row_node(row); acc = row_fold(i, lane, step); }
+    if (TIER == 1u) {
+        sh[lid] = acc;
+        workgroupBarrier();
+        for (var s = 16u; s >= 1u; s = s / 2u) {                 // the five-step tree over each 32-lane group; every lane runs every step
+            var t = identity();
+            if (lane < s) { t = sh[lid + s]; }
+            workgroupBarrier();
+            sh[lid] = comb(sh[lid], t);
+            workgroupBarrier();
+        }
+        if (valid && lane == 0u) { finish(i, sh[lid]); }
+    }
+    if (TIER == 2u) {
+        let t = wg_reduce_f32(acc, lid, OP);                     // the workgroup tree of the prelude (subgroup variant when available)
+        if (valid && lid == 0u) { finish(i, t); }
+    }
+}
+
+@compute @workgroup_size(WG)
+fn segmented_reduce(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    if (TIER == 0u) { tier0(wid, lid.x); } else { tiered(wid, lid.x); }
 }
 `;

package/src/wgsl/spmv-pull.wgsl.ts

@@ -1,47 +1,121 @@
 /**
- * The `spmv-pull` kernel body (spec 6 row 9, 8.2; PD-1 of the M8b plan): one invocation per row of the REVERSE
- * adjacency, grid-stride over `[0, P.n)`, folding `weight * xNorm[nbr]` over the row's in-arcs in chunks of 64
- * terms (a two-level f32 sum: the chunk absorbs the rounding of 64 terms, the row total the rounding of the chunk
- * count, so a 10,000-arc hub row loses about 200 rounding steps instead of 10,000; Kahan compensation is not used
- * because Metal's shader compiler folds `((acc + term) - acc) - term` to zero whatever hides it) and writing
+ * The `spmv-pull` kernel body (spec 6 row 9, 8.2; PD-1 of the M8b plan) in its three in-degree tiers (P4-T5, PD-6):
+ * every tier folds `weight * xNorm[nbr]` over a row's in-arcs of the REVERSE adjacency in chunks of 64 terms (a
+ * two-level f32 sum: the chunk absorbs the rounding of 64 terms, the row total the rounding of the chunk count, so a
+ * 10,000-arc hub row loses about 200 rounding steps instead of 10,000; Kahan compensation is not used because
+ * Metal's shader compiler folds `((acc + term) - acc) - term` to zero whatever hides it) and writes
  * `rankOut[v] = beta * pv + alpha * (sum + danglingMass * pv)`, where `pv` is `personalization[v]` when
- * HAS_PERSONALIZATION and the uniform `P.uniformP` otherwise. PageRank sets alpha to the
- * damping factor, beta to `1 - alpha` and USE_DANGLING; HITS and eigenvector set alpha 1, beta 0, uniformP 0; Katz
- * sets alpha to the attenuation, beta to its constant and uniformP 1. The body is normative: a sabotage mutation
- * (test/helpers/sabotage.ts) is a textual edit of it, so it is not restyled.
+ * HAS_PERSONALIZATION and the uniform `P.uniformP` otherwise. The rows are [P.start, P.n) of the dispatch, `perm[row]`
+ * under USE_PERM: TIER 0 is grid-stride over them (one row per thread, `P.stride` from planGridStride); TIER 1 is 32
+ * lanes per row and WG / 32 rows per workgroup with a five-step tree in workgroup memory (each lane keeps its own
+ * 64-term chunking); TIER 2 is one row per workgroup through `wg_reduce_f32` (hence `needs: ["subgroups"]`). PageRank
+ * sets alpha to the damping factor, beta to `1 - alpha` and USE_DANGLING; HITS and eigenvector set alpha 1, beta 0,
+ * uniformP 0; Katz sets alpha to the attenuation, beta to its constant and uniformP 1. The body is normative: a
+ * sabotage mutation (test/helpers/sabotage.ts) is a textual edit of it, so it is not restyled.
+ *
+ * TIER 0 folds its row through `row_sum_dense`, a stride-one copy of `row_sum`, because the shader compiler emits
+ * `row_sum(v, 0u, 1u)` as a call and leaves the stride in a parameter: the loop then walks the row with a runtime
+ * step, which costs the strength-reduced addressing into colIdx / weights and the unrolling that keeps several loads
+ * in flight per thread. PageRank and every other caller that passes no in-degree tiers runs TIER 0 alone, and under
+ * the tiers it still folds the low-degree rows, which are most of them; the shared fold cost 38 per cent more GPU
+ * time at 100k nodes and 1M arcs on an RTX 4070 SUPER and about twice the time on a Tesla T4. The two folds spell
+ * their locals apart (`lo` / `hi` / `k` against `a0` / `a1`) so that each sabotage row names exactly one of them.
  */
 
-/** Entry point `spmv_pull`; overrides HAS_PERSONALIZATION and USE_DANGLING plus the standard USE_PERM / HAS_WEIGHTS. */
+/** Entry point `spmv_pull`; overrides HAS_PERSONALIZATION, USE_DANGLING and TIER (0 / 1 / 2) plus the standard USE_PERM / HAS_WEIGHTS. */
 export const spmvPullWgsl = /* wgsl */ `
-@compute @workgroup_size(WG)
-fn spmv_pull(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+fn row_node(row: u32) -> u32 { return select(row, perm[row], USE_PERM); }
+fn row_sum(v: u32, lane: u32, step: u32) -> f32 {               // this lane's arcs of row v inside the bound window, the two-level 64-term fold of the header
+    let a0 = max(rowPtr[v], P.arcBase);
+    let a1 = min(rowPtr[v + 1u], P.arcEnd);
+    var acc = 0.0;
+    var chunk = 0.0;
+    var inChunk = 0u;
+    for (var arc = a0 + lane; arc < a1; arc = arc + step) {
+        let nbr = colIdx[arc - P.arcBase];               // \`target\` is a WGSL reserved word (spec 16.2)
+        var weight = 1.0;
+        if (HAS_WEIGHTS) { weight = weights[arc - P.arcBase]; }
+        // two-level sum: 64 terms into chunk, chunk into acc (see the header; no compensation, no select)
+        chunk = chunk + (weight * xNorm[nbr]);
+        inChunk = inChunk + 1u;
+        if (inChunk == 64u) {
+            acc = acc + chunk;
+            chunk = 0.0;
+            inChunk = 0u;
+        }
+    }
+    acc = acc + chunk;
+    return acc;
+}
+fn row_sum_dense(v: u32) -> f32 {                                // TIER 0's stride-one twin of row_sum (see the header)
+    let lo = max(rowPtr[v], P.arcBase);
+    let hi = min(rowPtr[v + 1u], P.arcEnd);
+    var acc = 0.0;
+    var chunk = 0.0;
+    var inChunk = 0u;
+    for (var arc = lo; arc < hi; arc = arc + 1u) {
+        let k = arc - P.arcBase;                         // the window-local index; this walk is contiguous
+        let nbr = colIdx[k];                             // \`target\` is a WGSL reserved word (spec 16.2)
+        var weight = 1.0;
+        if (HAS_WEIGHTS) { weight = weights[k]; }
+        // two-level sum: 64 terms into chunk, chunk into acc (see the header; no compensation, no select)
+        chunk = chunk + (weight * xNorm[nbr]);
+        inChunk = inChunk + 1u;
+        if (inChunk == 64u) {
+            acc = acc + chunk;
+            chunk = 0.0;
+            inChunk = 0u;
+        }
+    }
+    acc = acc + chunk;
+    return acc;
+}
+fn finish(v: u32, acc: f32) {
     var dangling = 0.0;
     if (USE_DANGLING) { dangling = partials[0].danglingMass; }
-    let first = linear_id(wid, lid.x);
-    for (var row = first; row < P.n; row = row + P.stride) {
-        let v = select(row, perm[row], USE_PERM);
-        let a0 = max(rowPtr[v], P.arcBase);
-        let a1 = min(rowPtr[v + 1u], P.arcEnd);
-        var acc = 0.0;
-        var chunk = 0.0;
-        var inChunk = 0u;
-        for (var arc = a0; arc < a1; arc = arc + 1u) {
-            let nbr = colIdx[arc - P.arcBase];               // \`target\` is a WGSL reserved word (spec 16.2)
-            var weight = 1.0;
-            if (HAS_WEIGHTS) { weight = weights[arc - P.arcBase]; }
-            // two-level sum: 64 terms into chunk, chunk into acc (see the header; no compensation, no select)
-            chunk = chunk + (weight * xNorm[nbr]);
-            inChunk = inChunk + 1u;
-            if (inChunk == 64u) {
-                acc = acc + chunk;
-                chunk = 0.0;
-                inChunk = 0u;
-            }
+    var pv = P.uniformP;
+    if (HAS_PERSONALIZATION) { pv = personalization[v]; }
+    rankOut[v] = (P.beta * pv) + (P.alpha * (acc + (dangling * pv)));
+}
+fn tier0(wid: vec3<u32>, lane: u32) {                            // TIER 0: grid-stride over the rows [P.start, P.n); no barrier
+    for (var row = linear_id(wid, lane) + P.start; row < P.n; row = row + P.stride) {
+        let v = row_node(row);
+        finish(v, row_sum_dense(v));
+    }
+}
+
+var<workgroup> sh: array<f32, WG>;
+
+fn tiered(wid: vec3<u32>, lid: u32) {                            // TIER 1: 32 lanes per row; TIER 2: WG lanes per row; rows [P.start, P.n) (PD-6)
+    let g = group_id(wid);
+    var row = P.start + g;
+    var lane = lid;
+    var step = WG;
+    if (TIER == 1u) { row = P.start + g * (WG / 32u) + lid / 32u; lane = lid % 32u; step = 32u; }
+    let valid = row < P.n;
+    var v = 0u;
+    var acc = 0.0;
+    if (valid) { v = row_node(row); acc = row_sum(v, lane, step); }
+    if (TIER == 1u) {
+        sh[lid] = acc;
+        workgroupBarrier();
+        for (var s = 16u; s >= 1u; s = s / 2u) {                 // the five-step tree over each 32-lane group; every lane runs every step
+            var t = 0.0;
+            if (lane < s) { t = sh[lid + s]; }
+            workgroupBarrier();
+            sh[lid] = sh[lid] + t;
+            workgroupBarrier();
         }
-        acc = acc + chunk;
-        var pv = P.uniformP;
-        if (HAS_PERSONALIZATION) { pv = personalization[v]; }
-        rankOut[v] = (P.beta * pv) + (P.alpha * (acc + (dangling * pv)));
+        if (valid && lane == 0u) { finish(v, sh[lid]); }
+    }
+    if (TIER == 2u) {
+        let t = wg_reduce_f32(acc, lid, 0u);
+        if (valid && lid == 0u) { finish(v, t); }
     }
 }
+
+@compute @workgroup_size(WG)
+fn spmv_pull(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    if (TIER == 0u) { tier0(wid, lid.x); } else { tiered(wid, lid.x); }
+}
 `;