diff --git a/crates/simd-masking-parity/src/lib.rs b/crates/simd-masking-parity/src/lib.rs index 4e746b37..29e48e78 100644 --- a/crates/simd-masking-parity/src/lib.rs +++ b/crates/simd-masking-parity/src/lib.rs @@ -40,15 +40,16 @@ use ndarray::simd::{ eq_u32_via_to_mask, eq_u64_to_mask, eq_u8_to_mask, ge_i32_to_mask, ge_i32_to_mask_under, ge_u64_to_mask, ge_u8_to_mask, gt_i32_to_mask, gt_i32_to_mask_under, gt_u64_to_mask, gt_u8_to_mask, le_i32_to_mask, le_i32_to_mask_under, le_u64_to_mask, le_u8_to_mask, lt_i32_to_mask, lt_i32_to_mask_under, lt_u64_to_mask, - lt_u8_to_mask, mask_all, mask_and, mask_and_assign, mask_andnot, mask_andnot_assign, mask_any, mask_gather_u32, - mask_gather_u32_under, mask_not, mask_not_assign, mask_or, mask_or_assign, mask_scatter_or_u32, mask_set_range, - mask_shift_morton, mask_ternlog, mask_ternlog_any, mask_ternlog_assign, mask_ternlog_popcount, mask_xor, - mask_xor_assign, masked_group_sum_i32, masked_group_sum_i32_via, masked_key_run_count_u32, masked_max_i32, - masked_min_i32, masked_strided_group_sum, masked_sum_i32, masked_sum_wrapping_add_i32, ne_i32_to_mask, - ne_i32_to_mask_under, ne_u32_to_mask, ne_u32_to_mask_under, ne_u64_to_mask, ne_u8_to_mask, - ternary_match_strided16_to_mask, ternary_match_strided_to_mask, ternary_match_u32_to_mask, - ternary_match_u32_to_mask_under, ternary_match_u64_to_mask, ternary_match_u64_to_mask_under, ternlog, F64x8, - I32x16, I8x16, KeyRunCarry, MortonDir, U32x16, U64x8, U8x16, + lt_u8_to_mask, mask_all, mask_and, mask_and_assign, mask_andnot, mask_andnot_assign, mask_any, + mask_gather_chain_u32, mask_gather_chain_u32_under, mask_gather_u32, mask_gather_u32_under, mask_not, + mask_not_assign, mask_or, mask_or_assign, mask_scatter_or_u32, mask_set_range, mask_shift_morton, mask_ternlog, + mask_ternlog_any, mask_ternlog_assign, mask_ternlog_popcount, mask_xor, mask_xor_assign, masked_group_sum_i32, + masked_group_sum_i32_via, masked_key_run_count_u32, masked_max_i32, masked_min_i32, masked_strided_group_sum, + masked_sum_i32, masked_sum_wrapping_add_i32, ne_i32_to_mask, ne_i32_to_mask_under, ne_u32_to_mask, + ne_u32_to_mask_under, ne_u64_to_mask, ne_u8_to_mask, ternary_match_strided16_to_mask, + ternary_match_strided_to_mask, ternary_match_u32_to_mask, ternary_match_u32_to_mask_under, + ternary_match_u64_to_mask, ternary_match_u64_to_mask_under, ternlog, ChainHop, ChainTest, F64x8, I32x16, I8x16, + KeyRunCarry, MortonDir, U32x16, U64x8, U8x16, }; /// Number of check groups [`run`] executes (for the log line only). @@ -1590,6 +1591,79 @@ fn check_gather_scatter_group() -> Result<(), u32> { return Err(0xD01); } + // ── mask_gather_chain_u32[_under] ──────────────────────────────── + // Two hops: index -> mid (live + one test, keyed onward through + // `mid_next`) -> far (live + one negated test). The reference walks + // the chain row by row from the same lanes; out-of-range keys at + // both depths exercise the "false by contract" arm. + let mid_rows = 30usize; + let far_rows = 20usize; + let mid_live: Vec = (0..words_for(mid_rows)).map(|_| rng.next()).collect(); + let far_live: Vec = (0..words_for(far_rows)).map(|_| rng.next()).collect(); + let mid_lane: Vec = (0..mid_rows).map(|_| (rng.next() % 3) as u32).collect(); + let far_lane: Vec = (0..far_rows).map(|_| (rng.next() % 3) as u32).collect(); + let mid_next: Vec = (0..mid_rows) + .map(|j| { + if j % 7 == 0 { + (far_rows + 2) as u32 + } else { + (rng.next() % far_rows as u64) as u32 + } + }) + .collect(); + let cindex: Vec = (0..n) + .map(|i| { + if i % 5 == 0 { + (mid_rows + 1 + i) as u32 + } else { + (rng.next() % mid_rows as u64) as u32 + } + }) + .collect(); + let mid_tests = [ChainTest { + lane: &mid_lane, + value: 1, + negate: false, + }]; + let far_tests = [ChainTest { + lane: &far_lane, + value: 2, + negate: true, + }]; + let hops = [ + ChainHop { + live: &mid_live, + rows: mid_rows, + tests: &mid_tests, + next: Some(&mid_next), + }, + ChainHop { + live: &far_live, + rows: far_rows, + tests: &far_tests, + next: None, + }, + ]; + let bit = |w: &[u64], k: usize| (w[k / 64] >> (k % 64)) & 1 == 1; + let holds = |i: usize| { + let a = cindex[i] as usize; + if a >= mid_rows || !bit(&mid_live, a) || mid_lane[a] != 1 { + return false; + } + let b = mid_next[a] as usize; + b < far_rows && bit(&far_live, b) && far_lane[b] != 2 + }; + let mut out = vec![u64::MAX; out_len]; + mask_gather_chain_u32(&cindex, &hops, &mut out); + if out != reference_mask(n, out_len, holds) { + return Err(0xD02); + } + let mut out = vec![u64::MAX; out_len]; + mask_gather_chain_u32_under(&cindex, &gate, &hops, &mut out); + if out != reference_mask(n, out_len, |i| bit(&gate, i) && holds(i)) { + return Err(0xD03); + } + // ── mask_scatter_or_u32 ────────────────────────────────────────── // `out2` starts explicitly zeroed (the caller's job now, not the // primitive's): the accumulate contract means a dirty prefill would diff --git a/src/simd.rs b/src/simd.rs index 3c618191..c0327b00 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -868,6 +868,8 @@ pub use crate::simd_masking_ops::{ mask_andnot, mask_andnot_assign, mask_any, + mask_gather_chain_u32, + mask_gather_chain_u32_under, mask_gather_u32, mask_gather_u32_under, mask_not, @@ -931,6 +933,8 @@ pub use crate::simd_masking_ops::{ widen_bounded_cross_power_sums, widen_bounded_power_sums, BoundedFoldError, + ChainHop, + ChainTest, CrossPowerSums, KeyRunCarry, MortonDir, diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index fc3dff9c..edf0a245 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -1237,6 +1237,202 @@ pub fn mask_gather_u32_under(src: &[u64], src_rows: usize, index: &[u32], under: } } +/// One equality test a [`ChainHop`] applies to the row it reached: +/// `lane[row] == value`, or `!=` when `negate`. +#[derive(Debug, Clone, Copy)] +pub struct ChainTest<'a> { + /// A value lane of the hop's table, at least `rows` long. + pub lane: &'a [u32], + /// The value the lane is compared against. + pub value: u32, + /// Pass on `!=` instead of `==`. + pub negate: bool, +} + +/// One hop of a [`mask_gather_chain_u32`]: the table the previous key names. +/// +/// The row reached must be in range (`< rows`), live in `live`, and pass +/// every test; then `next` names the row of the following hop. The last hop +/// has no `next`, every other hop has one. +#[derive(Debug, Clone, Copy)] +pub struct ChainHop<'a> { + /// The table's resident live plane, `rows.div_ceil(64)` words at least. + pub live: &'a [u64], + /// The table's row count; a key at or past it fails the row. + pub rows: usize, + /// The tests the reached row must pass. + pub tests: &'a [ChainTest<'a>], + /// The key lane naming the next hop's row, `rows` long; `None` on the + /// last hop. + pub next: Option<&'a [u32]>, +} + +/// Validate a chain's shape against the panics both entry points document. +fn check_chain(name: &str, hops: &[ChainHop<'_>]) { + assert!(!hops.is_empty(), "{name}: hops is empty"); + let last = hops.len() - 1; + for (j, h) in hops.iter().enumerate() { + assert!( + h.live.len() >= mask_words_for(h.rows), + "{name}: hop {j} live.len()={} < required {}", + h.live.len(), + mask_words_for(h.rows) + ); + assert!( + h.next.is_some() == (j < last), + "{name}: hop {j} must {} a next key", + if j < last { "have" } else { "not have" } + ); + if let Some(next) = h.next { + assert!(next.len() >= h.rows, "{name}: hop {j} next.len()={} < rows {}", next.len(), h.rows); + } + for t in h.tests { + assert!(t.lane.len() >= h.rows, "{name}: hop {j} test lane.len()={} < rows {}", t.lane.len(), h.rows); + } + } +} + +/// Whether the chain starting at key `k` reaches a row of every hop that is +/// in range, live and passes its tests. +#[inline(always)] +fn chain_holds(hops: &[ChainHop<'_>], k: u32) -> bool { + let mut cur = k as usize; + for h in hops { + if cur >= h.rows || (h.live[cur / 64] >> (cur % 64)) & 1 == 0 { + return false; + } + for t in h.tests { + if (t.lane[cur] == t.value) == t.negate { + return false; + } + } + if let Some(next) = h.next { + cur = next[cur] as usize; + } + } + true +} + +/// Chained gather: bit `i` of the result is set when the key `index[i]` +/// leads through every hop: each row reached is in range, live, and passes +/// that hop's tests, and each hop's `next` key names the following hop's row. +/// +/// This is a chain of key reads folded into one pass over `index`: the +/// targets' predicates are evaluated at the row each key reaches, in +/// registers, so no mask over any hop's table is ever built (the +/// `start[fk^k(i)]` shape of lance-graph's chained-hop co-research, option +/// (a)). With one hop and no tests it is [`mask_gather_u32`]. +/// +/// An out-of-range key at any depth fails the row, never panics. +/// +/// # Vectorisation, honestly +/// +/// **Scalar, and by necessity**, for the same reason as [`mask_gather_u32`]: +/// each hop's address is a value read by the previous hop, so the reads are +/// dependent and there is no bit gather to vectorise them with. Cost is +/// `O(index.len() × hops)` at most; a row stops at its first failing hop. +/// When the caller holds the rows it wants, [`mask_gather_chain_u32_under`] +/// reads only those. +/// +/// # Panics +/// +/// Panics if `hops` is empty, if a hop other than the last has no `next` or +/// the last has one, if a hop's `live`, `next` or test lane is shorter than +/// its `rows` require, or if `out_words.len() < index.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::{mask_gather_chain_u32, ChainHop}; +/// +/// // line -> partner -> country: partners 0..3, all live; countries 0..2, +/// // country 1 dead. +/// let partner_live = [0b111u64]; +/// let partner_country = [0u32, 1, 0]; +/// let country_live = [0b01u64]; +/// let hops = [ +/// ChainHop { live: &partner_live, rows: 3, tests: &[], next: Some(&partner_country) }, +/// ChainHop { live: &country_live, rows: 2, tests: &[], next: None }, +/// ]; +/// let line_partner = [0u32, 1, 2, 7]; // line 3's key dangles +/// let mut out = [u64::MAX]; +/// mask_gather_chain_u32(&line_partner, &hops, &mut out); +/// // lines 0 and 2 reach the live country 0; line 1 reaches the dead one. +/// assert_eq!(out[0], 0b0101); +/// ``` +#[inline] +pub fn mask_gather_chain_u32(index: &[u32], hops: &[ChainHop<'_>], out_words: &mut [u64]) { + let n = index.len(); + let words = mask_words_for(n); + assert!( + out_words.len() >= words, + "mask_gather_chain_u32: out_words.len()={} < required {}", + out_words.len(), + words + ); + check_chain("mask_gather_chain_u32", hops); + for w in out_words.iter_mut() { + *w = 0; + } + for (w, out_word) in out_words.iter_mut().enumerate().take(words) { + let base = w * 64; + let live = (n - base).min(64); + let mut acc = 0u64; + for lane in 0..live { + acc |= (chain_holds(hops, index[base + lane]) as u64) << lane; + } + *out_word = acc; + } +} + +/// Gated [`mask_gather_chain_u32`]: the result equals the chained gather +/// ANDed with `under`, and `index` is read **only** at rows whose `under` bit +/// is set, so the work is proportional to `popcount(under)`. +/// +/// `under` bits at or past `index.len()` are ignored. `out_words` is fully +/// overwritten. +/// +/// # Panics +/// +/// As [`mask_gather_chain_u32`], and if `under.len()` is below +/// `index.len().div_ceil(64)`. +#[inline] +pub fn mask_gather_chain_u32_under(index: &[u32], under: &[u64], hops: &[ChainHop<'_>], out_words: &mut [u64]) { + let n = index.len(); + let words = mask_words_for(n); + assert!( + out_words.len() >= words, + "mask_gather_chain_u32_under: out_words.len()={} < required {}", + out_words.len(), + words + ); + assert!( + under.len() >= words, + "mask_gather_chain_u32_under: under.len()={} < required {}", + under.len(), + words + ); + check_chain("mask_gather_chain_u32_under", hops); + for w in out_words.iter_mut() { + *w = 0; + } + for w in 0..words { + let base = w * 64; + let mut bits = under[w]; + let valid = n - base; + if valid < 64 { + bits &= (1u64 << valid) - 1; + } + let mut acc = 0u64; + while bits != 0 { + let lane = bits.trailing_zeros() as usize; + bits &= bits - 1; + acc |= (chain_holds(hops, index[base + lane]) as u64) << lane; + } + out_words[w] = acc; + } +} + /// Scatter-OR by row index: for every element `i` selected by `src` (bit `i` /// set, `i < index.len()`), sets bit `index[i]` of the result — provided /// `index[i] < out_rows`. @@ -6557,6 +6753,198 @@ mod tests { mask_gather_u32_under(&src, 1, &index, &[u64::MAX], &mut out); } + /// The row loop the chained gather must equal: follow the keys hop by + /// hop, failing on an out-of-range key, a dead row, or a failed test. + fn naive_chain(index: &[u32], gate: &[u64], hops: &[ChainHop<'_>]) -> Vec { + (0..index.len()) + .map(|i| { + if (gate[i / 64] >> (i % 64)) & 1 == 0 { + return false; + } + let mut cur = index[i] as usize; + for h in hops { + if cur >= h.rows || (h.live[cur / 64] >> (cur % 64)) & 1 == 0 { + return false; + } + for t in h.tests { + if (t.lane[cur] == t.value) == t.negate { + return false; + } + } + if let Some(next) = h.next { + cur = next[cur] as usize; + } + } + true + }) + .collect() + } + + #[test] + fn mask_gather_chain_u32_under_matches_the_row_loop() { + // Three tables after the anchor, each with dead rows, keys that run + // past the next table, and a test lane; gates with dirty tails. + let rows = [37usize, 23, 11]; + let mut seed = 0x0DDB_A11C_4A1E_5EEDu64; + let lanes = |r: usize, dom: u64, seed: &mut u64| -> Vec { + (0..r).map(|_| (splitmix(seed) % dom) as u32).collect() + }; + let live: Vec> = rows + .iter() + .map(|&r| { + (0..r.div_ceil(64)) + .map(|_| splitmix(&mut seed) | 0x5555_5555_5555_5555) + .collect() + }) + .collect(); + // Keys may name a row past the next table (out of range). + let next0 = lanes(rows[0], rows[1] as u64 + 3, &mut seed); + let next1 = lanes(rows[1], rows[2] as u64 + 2, &mut seed); + let val0 = lanes(rows[0], 3, &mut seed); + let val2 = lanes(rows[2], 2, &mut seed); + let t0 = [ChainTest { + lane: &val0, + value: 1, + negate: true, + }]; + let t2 = [ChainTest { + lane: &val2, + value: 0, + negate: false, + }]; + let hops = [ + ChainHop { + live: &live[0], + rows: rows[0], + tests: &t0, + next: Some(&next0), + }, + ChainHop { + live: &live[1], + rows: rows[1], + tests: &[], + next: Some(&next1), + }, + ChainHop { + live: &live[2], + rows: rows[2], + tests: &t2, + next: None, + }, + ]; + let mut kept_any = 0usize; + let mut total = 0usize; + for &n in &[0usize, 1, 63, 64, 65, 130, 1000] { + let index: Vec = (0..n) + .map(|_| (splitmix(&mut seed) % (rows[0] as u64 + 4)) as u32) + .collect(); + let words = n.div_ceil(64).max(1); + let gates: [Vec; 3] = + [vec![0u64; words], vec![u64::MAX; words], (0..words).map(|_| splitmix(&mut seed)).collect()]; + for gate in &gates { + let want = naive_chain(&index, gate, &hops); + kept_any += want.iter().filter(|&&b| b).count(); + total += n; + let want = bits_to_words(&want); + let mut out = vec![u64::MAX; words + 1]; + mask_gather_chain_u32_under(&index, gate, &hops, &mut out); + assert_eq!(&out[..want.len()], &want[..], "chain mismatch at n={n}"); + assert_eq!(out[words], 0, "surplus word must be cleared at n={n}"); + } + let mut ungated = vec![u64::MAX; words]; + mask_gather_chain_u32(&index, &hops, &mut ungated); + let all = vec![u64::MAX; words]; + let want = bits_to_words(&naive_chain(&index, &all, &hops)); + assert_eq!(&ungated[..want.len()], &want[..], "ungated chain mismatch at n={n}"); + } + // Anti-vacuity: the chain keeps some rows and drops most. + assert!(kept_any > 0 && kept_any * 3 < total, "{kept_any} of {total}"); + } + + #[test] + fn mask_gather_chain_u32_one_hop_is_the_gated_gather() { + let mut seed = 0x1234_5678_9ABC_DEF0u64; + let src_rows = 50usize; + let src: Vec = (0..src_rows.div_ceil(64)) + .map(|_| splitmix(&mut seed)) + .collect(); + let index: Vec = (0..200) + .map(|_| (splitmix(&mut seed) % 60) as u32) + .collect(); + let gate: Vec = (0..4).map(|_| splitmix(&mut seed)).collect(); + let hop = [ChainHop { + live: &src, + rows: src_rows, + tests: &[], + next: None, + }]; + let mut chain = [0u64; 4]; + let mut single = [0u64; 4]; + mask_gather_chain_u32_under(&index, &gate, &hop, &mut chain); + mask_gather_u32_under(&src, src_rows, &index, &gate, &mut single); + assert_eq!(chain, single); + } + + #[test] + #[should_panic(expected = "next")] + fn mask_gather_chain_u32_rejects_a_hop_without_its_key() { + let live = [u64::MAX]; + let hops = [ + ChainHop { + live: &live, + rows: 4, + tests: &[], + next: None, + }, + ChainHop { + live: &live, + rows: 4, + tests: &[], + next: None, + }, + ]; + mask_gather_chain_u32(&[0, 1], &hops, &mut [0u64; 1]); + } + + #[test] + #[should_panic(expected = "next")] + fn mask_gather_chain_u32_rejects_a_key_after_the_last_hop() { + let live = [u64::MAX]; + let key = [0u32; 4]; + let hops = [ChainHop { + live: &live, + rows: 4, + tests: &[], + next: Some(&key), + }]; + mask_gather_chain_u32(&[0, 1], &hops, &mut [0u64; 1]); + } + + #[test] + #[should_panic(expected = "lane")] + fn mask_gather_chain_u32_rejects_a_short_test_lane() { + let live = [u64::MAX]; + let lane = [0u32; 2]; + let t = [ChainTest { + lane: &lane, + value: 0, + negate: false, + }]; + let hops = [ChainHop { + live: &live, + rows: 4, + tests: &t, + next: None, + }]; + mask_gather_chain_u32(&[0, 1], &hops, &mut [0u64; 1]); + } + + #[test] + #[should_panic(expected = "hops")] + fn mask_gather_chain_u32_rejects_an_empty_chain() { + mask_gather_chain_u32(&[0, 1], &[], &mut [0u64; 1]); + } + fn naive_scatter(src_bits: &[bool], index: &[u32], out_rows: usize) -> Vec { let mut out = vec![false; out_rows]; for (i, &selected) in src_bits.iter().enumerate() {