OneSweep: fix global_hist layout when KEY_BITS < key width

[jowens]

OneSweep: fix global_hist histogram layout when KEY_BITS is less than the key width

Problem

global_hist always builds one 256-bucket histogram per 8-bit digit of the full key: 4 for 32-bit keys, 8 for 64-bit keys. The digit offsets are hard-coded (+256u, +512u, +768u, …). But the per-segment workgroup-memory histograms were laid out with stride ALL_RADIX = RADIX * SORT_PASSES, and hist[] was sized as SORT_PASSES * RADIX.

With KEY_BITS smaller than the key width (so SORT_PASSES < digits per key), segment 1 (threads 64–127) writes into segment 0's upper-digit region. The final reduction then adds the wrong counts into the low-digit histograms, and the global atomics into hist[] for the upper digits go out of bounds. The sort output is wrong. For example, a 32-bit keys-only sort with KEY_BITS: 16 produces invalid output.

The default configurations are unaffected, because there KEY_BITS equals the key width and SORT_PASSES == digits per key.

Fix

• Add DIGITS_PER_KEY = ceil(bitsPerElement / RADIX_LOG) in JS and as a WGSL constant, and set ALL_RADIX = RADIX * DIGITS_PER_KEY.
• Set histSize = DIGITS_PER_KEY * RADIX * 4. passHistSize still uses SORT_PASSES segments, which is all onesweep_scan/onesweep_pass read.
• In validate(), the CPU reference histogram now covers every digit, so it matches what global_hist writes.

For default configurations, every buffer size and the generated WGSL are unchanged, apart from the new constant and comments.

Testing

• Compilation. Under Dawn on lavapipe, the patched kernels compile for u32 keys-only with KEY_BITS 32 and 16, u32 key-value with 8, i32 keys-only with 24, and u64 with 64 and 40.
• End to end on an Apple M1 Max (Chrome). A subclass applying the same two changes was used in a hashed-key random-permutation benchmark. It sorts u32 keys-only with KEY_BITS 16 and 24 (2 and 3 passes) at n = 2^22. Over 8 seeds, every output was a valid permutation with keys correctly ordered by their low KEY_BITS bits. Without the change, all 2- and 3-pass runs were invalid; 4-pass runs were fine in both cases.
• Not run: the regression suite on a real GPU. Unmodified OneSweep also fails on lavapipe (subgroup size 8–32), so the software adapter can't be used to check sorts.

Follow-up (not in this PR)

validate() compares against a full sort, so it can't yet check a sort with KEY_BITS less than the key width. That needs a stable-sort-by-low-bits reference, and then a regression case for KEY_BITS < key width.

---

Generated by Claude Code