Shared subgroup digit match (GPU Multisplit Alg. 3); OneSweep's WLMS built on it

src/onesweep.mjs

@@ -728,11 +728,7 @@ ${this.wlmsWGSL()}
        */
       var offsets = array<u32, KEYS_PER_THREAD>();
       {
-        ${
-          this.wideSubgroups
-            ? "let lane_mask_lt = laneMaskBelow(sgid);"
-            : "let lane_mask_lt = (1u << sgid) - 1u;"
-        }
+        let lane_mask_lt = ${this.fnDeclarations.subgroupLaneMaskLt()};
         let s_offset = sid * RADIX;
         for (var k = 0u; k < KEYS_PER_THREAD; k += 1u) {
           /* offsets[k] reflects the number of digits seen by this subgroup
@@ -1183,61 +1179,11 @@ ${this.wlmsWGSL()}
    * WGSL is generated, so devices whose subgroups are at most 32 wide (Apple,
    * NVIDIA) get exactly the 32-bit WLMS and its register use. Emulated
    * subgroups are 32 wide. */
-  get wideSubgroups() {
-    return this.useSubgroups && (this.SUBGROUP_MAX_SIZE ?? 0) > 32;
-  }
-
-  /** WGSL for WLMS (warp-level multisplit): 32-bit lane masks, or vec4<u32>
-   * masks when wideSubgroups. */
+  /** WGSL for WLMS (warp-level multisplit). The ballot-based digit match is
+   * the shared subgroupDigitMatch (32-bit or, for subgroups that can be
+   * wider than 32, vec4<u32> lane masks). */
   wlmsWGSL() {
-    return this.wideSubgroups
-      ? /* wgsl */ `    /* warp-level multisplit function, for subgroups of up to 128 lanes */
-    /** Lane masks are vec4<u32> (lane i is bit i % 32 of word i / 32), as
-     * subgroupBallot returns them. Generated only when the adapter's subgroups
-     * can be wider than 32 (see wideSubgroups); otherwise WLMS is the 32-bit
-     * version, which needs fewer registers. */
-    fn laneMaskBelow(n: u32) -> vec4<u32> {
-      /* bits set for lanes 0..n-1 */
-      var m = vec4<u32>(0u);
-      for (var i = 0u; i < 4u; i += 1u) {
-        let lo = i * 32u;
-        m[i] = select(select((1u << (n - lo)) - 1u, 0xffffffffu, n >= lo + 32u), 0u, n <= lo);
-      }
-      return m;
-    }
-    fn WLMS(key: u32, shift: u32, sgid: u32, sgsz: u32, lane_mask_lt: vec4<u32>, s_offset: u32) -> u32 {
-      /* start from the lanes that exist, so ~ballot cannot mark lanes past sgsz */
-      var eq_mask = laneMaskBelow(sgsz);
-      /* loop through the bits in this digit place */
-      for (var k = 0u; k < RADIX_LOG; k += 1u) {
-        let curr_bit = 1u << (k + shift);
-        let pred = (key & curr_bit) != 0u;
-        let ballot = subgroupBallot(pred);
-        eq_mask &= select(~ballot, ballot, pred);
-      }
-      /* eq_mask marks the threads that have the same digit as me ... */
-      let below = eq_mask & lane_mask_lt;
-      var out = countOneBits(below.x) + countOneBits(below.y) + countOneBits(below.z) + countOneBits(below.w);
-      /* ... and out counts those threads in eq_mask whose ids are less than me */
-      /* what is the largest-sgid thread in my subgroup that has the same digit as me?
-       * (eq_mask includes me, so it is never zero) */
-      let top_word = select(select(select(0u, 1u, eq_mask.y != 0u), 2u, eq_mask.z != 0u), 3u, eq_mask.w != 0u);
-      let highest_rank_peer = top_word * 32u + firstLeadingBit(eq_mask[top_word]);
-      var pre_inc = 0u;
-      if (sgid == highest_rank_peer) {
-        /* add to wg_warpHist for my digit value: total number of threads with my digit value
-         * (the number ranked below me with my digit value, plus one for me) */
-        pre_inc = atomicAdd(&wg_warpHist[((key >> shift) & RADIX_MASK) + s_offset], out + 1u);
-      }
-
-      workgroupBarrier();
-      /* see the 32-bit version for what out means */
-      out += subgroupShuffle(pre_inc, highest_rank_peer);
-      return out;
-    }
-`
-      : /* wgsl */ `    /* warp-level multisplit function */
-    /* this appears to be limited to subgroup sizes <= 32 */
+    return /* wgsl */ `    /* warp-level multisplit function */
     /** more portable version, HLSL:
      * https://github.com/b0nes164/GPUSorting/blob/09a6081d964b682bbb58838b83ab75aebe7f4e05/GPUSortingD3D12/Shaders/SortCommon.hlsl#L400-L424
      * https://github.com/b0nes164/GPUSorting/blob/09a6081d964b682bbb58838b83ab75aebe7f4e05/GPUSortingD3D12/Shaders/SortCommon.hlsl#L359-L375
@@ -1248,20 +1194,9 @@ ${this.wlmsWGSL()}
      * doing more exotic techniques to prevent losing occupancy. Because basically you want a shared
      * mem histogram per warp."
      */
-    fn WLMS(key: u32, shift: u32, sgid: u32, sgsz: u32, lane_mask_lt: u32, s_offset: u32) -> u32 {
-      var eq_mask = 0xffffffffu;
-      /* loop through the bits in this digit place */
-      for (var k = 0u; k < RADIX_LOG; k += 1u) {
-        let curr_bit = 1u << (k + shift);
-        let pred = (key & curr_bit) != 0u;
-        let ballot = subgroupBallot(pred);
-        eq_mask &= select(~ballot.x, ballot.x, pred);
-      }
-      /* eq_mask marks the threads that have the same digit as me ... */
-      var out = countOneBits(eq_mask & lane_mask_lt);
-      /* ... and out marks those threads in eq_mask whose ids are less than me */
-      /* what is the largest-sgid thread in my subgroup that has the same digit as me? */
-      let highest_rank_peer = sgsz - countLeadingZeros(eq_mask) - 1u;
+${this.fnDeclarations.subgroupLaneMasks()}
+    fn WLMS(key: u32, shift: u32, sgid: u32, sgsz: u32, lane_mask_lt: ${this.wideSubgroups ? "vec4<u32>" : "u32"}, s_offset: u32) -> u32 {
+${this.fnDeclarations.subgroupDigitMatch({ key: "key", shift: "shift", bits: "RADIX_LOG", rank: "out", peer: "highest_rank_peer" })}
       var pre_inc = 0u;
       if (sgid == highest_rank_peer) {
         /* add to wg_warpHist for my digit value: total number of threads with my digit value */

src/primitive.mjs

@@ -223,6 +223,15 @@ export class BasePrimitive {
     return this.adapter?.info ?? this.device.adapterInfo;
   }
 
+  /** Can this device's subgroups be wider than 32 lanes (e.g., AMD RDNA,
+   * 32-64)? Then lane masks need more than one u32 (subgroupBallot returns
+   * up to 128 bits). Emulated subgroups are 32 wide. Known before the WGSL
+   * is generated, so devices whose subgroups are at most 32 wide get 32-bit
+   * code. */
+  get wideSubgroups() {
+    return this.useSubgroups && (this.SUBGROUP_MAX_SIZE ?? 0) > 32;
+  }
+
   makeParamString(params) {
     let str = " [";
     for (const param of params) {

src/wgslFunctions.mjs

@@ -463,6 +463,78 @@ fn ${fnName}(// in: ptr<storage, array<${env.datatype}>, read>,
 }`;
     return fn;
   }
+  /** Ballot-based digit matching across a subgroup: Algorithm 3 of Ashkiani
+   * et al., "GPU Multisplit" (TOPC 2017), and the core of OneSweep's WLMS.
+   * With one ballot per bit of the digit, each invocation finds the lanes of
+   * its subgroup whose digit equals its own, how many of them are below it,
+   * and the highest of them.
+   *
+   * Lane masks are u32, or vec4<u32> (lane i is bit i % 32 of word i / 32)
+   * when the device's subgroups can be wider than 32 (env.wideSubgroups);
+   * that choice is made when the WGSL is generated, so 32-wide devices get
+   * the cheaper 32-bit code.
+   *
+   *   ${this.fnDeclarations.subgroupLaneMasks()}     at module scope: fn laneMaskBelow (wide only)
+   *   let lane_mask_lt = ${this.fnDeclarations.subgroupLaneMaskLt()};   lanes below me
+   *   ${this.fnDeclarations.subgroupDigitMatch({ key: "key", shift: "shift", bits: "RADIX_LOG" })}
+   *
+   * subgroupDigitMatch emits statements (not a function), so a caller's
+   * generated code is the same as if written inline. They declare, in the
+   * caller's scope: eq_mask (the matching lanes, me included), var <rank>
+   * (default "rank": matching lanes below me) and let <peer> (default
+   * "peer": the highest matching lane). They read sgsz, and lane_mask_lt
+   * (or the name given as laneMaskLt). The digit is bits [shift, shift +
+   * bits) of key; every invocation of the subgroup must take part. */
+  subgroupLaneMasks(args = {}) {
+    const env = Object.assign(Object.create(this.env), args); // properties in args overwrite this.env
+    if (!env.wideSubgroups) return "";
+    return /* wgsl */ `
+    fn laneMaskBelow(n: u32) -> vec4<u32> {
+      /* bits set for lanes 0..n-1 */
+      var m = vec4<u32>(0u);
+      for (var i = 0u; i < 4u; i += 1u) {
+        let lo = i * 32u;
+        m[i] = select(select((1u << (n - lo)) - 1u, 0xffffffffu, n >= lo + 32u), 0u, n <= lo);
+      }
+      return m;
+    }`;
+  }
+  subgroupLaneMaskLt(args = {}) {
+    const env = Object.assign(Object.create(this.env), args); // properties in args overwrite this.env
+    return env.wideSubgroups ? "laneMaskBelow(sgid)" : "(1u << sgid) - 1u";
+  }
+  subgroupDigitMatch(args = {}) {
+    const env = Object.assign(Object.create(this.env), args); // properties in args overwrite this.env
+    const { key, shift, bits } = args;
+    const rank = args.rank ?? "rank";
+    const peer = args.peer ?? "peer";
+    const lt = args.laneMaskLt ?? "lane_mask_lt";
+    const loop = (word) => /* wgsl */ `
+      /* loop through the bits in this digit place */
+      for (var k = 0u; k < ${bits}; k += 1u) {
+        let curr_bit = 1u << (k + ${shift});
+        let pred = (${key} & curr_bit) != 0u;
+        let ballot = subgroupBallot(pred);
+        eq_mask &= select(~${word}, ${word}, pred);
+      }
+      /* eq_mask marks the lanes that have the same digit as me (me included) */`;
+    if (!env.wideSubgroups) {
+      return /* wgsl */ `
+      var eq_mask = 0xffffffffu;${loop("ballot.x")}
+      var ${rank} = countOneBits(eq_mask & ${lt});
+      /* the highest lane with my digit */
+      let ${peer} = sgsz - countLeadingZeros(eq_mask) - 1u;`;
+    }
+    return /* wgsl */ `
+      /* start from the lanes that exist, so ~ballot cannot mark lanes past sgsz */
+      var eq_mask = laneMaskBelow(sgsz);${loop("ballot")}
+      let below = eq_mask & ${lt};
+      var ${rank} = countOneBits(below.x) + countOneBits(below.y) + countOneBits(below.z) + countOneBits(below.w);
+      /* the highest lane with my digit (eq_mask includes me, so it is never zero) */
+      let top_word = select(select(select(0u, 1u, eq_mask.y != 0u), 2u, eq_mask.z != 0u), 3u, eq_mask.w != 0u);
+      let ${peer} = top_word * 32u + firstLeadingBit(eq_mask[top_word]);`;
+  }
+
   /** A histogram in workgroup memory: the shared core of Histogram,
    * MultisplitPerm's histogram pass and similar kernels. It needs no
    * subgroups, so both function classes use this version.