From 396596f5fde85d967d973b51aaf648a8bbd3aeb4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 23:21:22 +0000 Subject: [PATCH] simd: masked_sum_i32 sums dense words branch-free, walks sparse ones Per 64-bit mask word: empty words are skipped, words with at most 16 set bits keep the set-bit walk, and full words with more than 16 set bits sum all 64 lanes branch-free, each widened to i64 and masked, so the i64 contract holds and the loop vectorizes. The threshold and the speed-up (2-5x, never slower than the walk) come from the D-MPF-2b measurement in lance-graph #1752. A partial final word always takes the walk. New test masked_sum_i32_both_word_routes_match_reference covers words at, below and above the threshold plus random densities, and asserts the fixture takes both routes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01G22yT6htkcdyXsihxxXdrg --- src/simd_masking_ops.rs | 60 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 55 insertions(+), 5 deletions(-) diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index fc3dff9c..0d8dda36 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -791,11 +791,12 @@ pub fn mask_ternlog_any(a: &[u64], b: &[u64], c: &[u64]) -> bool /// `reduce_sum()` — is **wrong**, and quietly so: `reduce_sum` on `I32x16` /// accumulates in `i32`, and 16 lanes near `i32::MAX` overflow it while the /// widened contract promises they cannot. Preserving the `i64` guarantee is -/// worth more than the lanes here, so the body walks set bits with -/// `u64::trailing_zeros` (one `TZCNT`/`RBIT+CLZ` per selected element, and -/// entire zero words skipped in one test). Cost is proportional to the -/// popcount, not the row count, which is the right shape for a selective -/// filter anyway. +/// worth more than the lanes here. Per 64-bit word: an empty word is +/// skipped, a sparse word (at most 16 set bits) is walked with +/// `u64::trailing_zeros`, and a dense full word is summed over all 64 lanes, +/// each widened to `i64` and masked branch-free. Cost is proportional to the +/// popcount for selective filters and to the row count, without a branch per +/// row, for dense ones. /// /// # Panics /// @@ -831,6 +832,19 @@ pub fn masked_sum_i32(values: &[i32], mask_words: &[u64]) -> i64 { if valid < 64 { bits &= (1u64 << valid) - 1; } + // Dense full word: sum all 64 lanes branch-free, each widened to i64 + // before the add, so the i64 contract holds and the loop vectorizes. + // A sparse word keeps the set-bit walk. The threshold is the D-MPF-2b + // measurement (lance-graph #1752): never slower than the walk. + if valid >= 64 && bits.count_ones() > MASKED_SUM_DENSE_THRESHOLD { + let chunk = &values[base..base + 64]; + let mut part: i64 = 0; + for (j, &v) in chunk.iter().enumerate() { + part = part.wrapping_add(i64::from(v) & -(((bits >> j) & 1) as i64)); + } + acc = acc.wrapping_add(part); + continue; + } while bits != 0 { let lane = bits.trailing_zeros() as usize; acc = acc.wrapping_add(values[base + lane] as i64); @@ -840,6 +854,10 @@ pub fn masked_sum_i32(values: &[i32], mask_words: &[u64]) -> i64 { acc } +/// Set bits per 64-bit word above which [`masked_sum_i32`] sums the whole word +/// branch-free instead of walking its set bits. +const MASKED_SUM_DENSE_THRESHOLD: u32 = 16; + /// Sum `a[i].wrapping_add(b[i])` where mask bit `i` is set, with the /// wrapped row value widened to `i64` before accumulation. /// @@ -6226,6 +6244,38 @@ mod tests { } } + #[test] + fn masked_sum_i32_both_word_routes_match_reference() { + // Words at, just below and just above the dense threshold, plus random + // words of every density, so the walk and the branch-free sum are both + // taken and both checked against the same independent reference. + let th = MASKED_SUM_DENSE_THRESHOLD; + let mut seed = 0x0DDB_A11C_AFE5_EED5; + let values: Vec = (0..64 * 40).map(|_| splitmix64(&mut seed) as i32).collect(); + let mut mask: Vec = Vec::new(); + for c in [th - 1, th, th + 1, 63, 64] { + mask.push(if c == 64 { u64::MAX } else { (1u64 << c) - 1 }); + } + while mask.len() < 40 { + let k = (splitmix64(&mut seed) % 4) as u32; + let mut w = splitmix64(&mut seed); + for _ in 0..k { + w &= splitmix64(&mut seed); + } + mask.push(w); + } + let dense = mask.iter().filter(|w| w.count_ones() > th).count(); + let sparse = mask.iter().filter(|w| w.count_ones() <= th).count(); + assert!(dense >= 5 && sparse >= 5, "fixture must take both routes: {dense} dense, {sparse} sparse"); + let expected: i64 = values + .iter() + .enumerate() + .filter(|(i, _)| mask[i / 64] >> (i % 64) & 1 == 1) + .map(|(_, &v)| v as i64) + .sum(); + assert_eq!(masked_sum_i32(&values, &mask), expected); + } + #[test] fn masked_sum_i32_widens_beyond_i32_range() { // 64 × i32::MAX = 137_438_953_408, which overflows i32 by ~64×. An