[jowens]

Step 2 of the GPU Multisplit plan: factor out the ballot-based digit match, Algorithm 3 of Ashkiani et al. and the core of OneSweep's WLMS, so that a stable multisplit can reuse it.

• subgroupDigitMatch({ key, shift, bits, rank, peer }) in wgslFunctions: one ballot per digit bit. It yields eq_mask (the lanes with my digit), <rank> (how many of them are below me) and <peer> (the highest of them). It emits statements rather than a function, so a caller's code is exactly what it would be if written inline.
• subgroupLaneMasks() / subgroupLaneMaskLt(): the lane masks are u32, or vec4<u32> when the device's subgroups can be wider than 32. The choice is made when the WGSL is generated, as in OneSweep: support subgroups wider than 32 (AMD RDNA wave64) #75.
• wideSubgroups moves from OneSweep to BasePrimitive.
• OneSweep's WLMS is now one template built on these helpers, instead of separate 32-bit and wide copies.

Verification:

• Generated code: I generated every OneSweep kernel before and after for subgroup ranges 32–32, 8–32, 32–64 and 4–128, with u32, i32, f32 and u64 keys, keys-only and key-value. With comments and whitespace removed, the WGSL is identical, so sort performance can't change.
• Tests: on the simulated RDNA4 (gfx1201), all 82 small tests pass.

---

Generated by Claude Code