diff --git a/src/simd.rs b/src/simd.rs index 01842d27..842b6b01 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -829,14 +829,19 @@ pub use crate::simd_masking_ops::{ mask_xor, mask_xor_assign, masked_group_count_u32, + masked_group_count_u32_pair, masked_group_count_u32_via, masked_group_max_i32, + masked_group_max_i32_pair, masked_group_max_i32_via, masked_group_min_i32, + masked_group_min_i32_pair, masked_group_min_i32_via, masked_group_sum_i32, + masked_group_sum_i32_pair, masked_group_sum_i32_via, masked_group_sum_sym_i32, + masked_group_sum_sym_i32_pair, masked_group_sum_sym_i32_via, masked_key_run_count_u32, masked_max_i32, diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index ff7b2591..bcf44b2c 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -1367,6 +1367,56 @@ pub fn masked_group_sum_i32_via(mask_words: &[u64], index: &[u32], table: &[u32] ); } +/// [`masked_group_sum_i32`] with the group of row `i` read as the composite +/// address `hi[i] * stride + lo[i]` — a two-column `GROUP BY (hi, lo)` +/// fused into one pass, no composite key lane ever materialised. +/// +/// **Zero-fallback**: `lo[i] >= stride` names no group at all (a minor key +/// outside its own declared cardinality) and drops row `i`; a resolved +/// composite key past `out.len()` drops it too — the same rule as +/// [`masked_group_sum_i32_via`]'s two hops. +/// +/// `out` is accumulated into (not zeroed by this function), same as +/// [`masked_group_sum_i32`]; overflow wraps the same way (widened to `i64`, +/// `wrapping_add`), and the mask tail is clamped identically. +/// +/// # Panics +/// +/// Panics if `hi.len() != lo.len()`, if `hi.len() != values.len()`, or if +/// `mask_words.len() < hi.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::masked_group_sum_i32_pair; +/// +/// let mask = [0b111u64]; +/// let hi = [0u32, 1, 0]; +/// let lo = [0u32, 0, 4]; // stride 4: row 2's lo (4) is out of range and dropped +/// let values = [10i32, 20, 999]; +/// let mut out = [0i64; 8]; // group (hi, lo) = hi * 4 + lo +/// masked_group_sum_i32_pair(&mask, &hi, &lo, 4, &values, &mut out); +/// assert_eq!(out[0], 10); // (0, 0) -> 0 +/// assert_eq!(out[4], 20); // (1, 0) -> 4 +/// ``` +#[inline] +pub fn masked_group_sum_i32_pair( + mask_words: &[u64], hi: &[u32], lo: &[u32], stride: u32, values: &[i32], out: &mut [i64], +) { + assert_eq!(hi.len(), lo.len(), "masked_group_sum_i32_pair: hi.len() != lo.len()"); + assert_eq!(hi.len(), values.len(), "masked_group_sum_i32_pair: hi/values length mismatch"); + // Accumulate: add into whatever `out` already holds. The caller zeroes + // once before the first call. + group_walk( + "masked_group_sum_i32_pair", + mask_words, + hi.len(), + GroupKeyAddr::Pair { hi, lo, stride }, + out, + |slot, i| *slot = slot.wrapping_add(values[i] as i64), + ); +} + /// The reserved code of the `_sym` family: the one `i64` a `_sym` fold /// never produces as a real value, marking "no row reached this group". /// @@ -1454,6 +1504,30 @@ pub fn masked_group_sum_sym_i32_via(mask_words: &[u64], index: &[u32], table: &[ ); } +/// [`masked_group_sum_sym_i32`] with the group of row `i` read as the +/// composite address `hi[i] * stride + lo[i]`, zero-fallback exactly as +/// [`masked_group_sum_i32_pair`]. Same reserved code, same row bound. +/// +/// # Panics +/// +/// Panics if `hi.len() != lo.len()`, if `hi.len() != values.len()`, or if +/// `mask_words.len() < hi.len().div_ceil(64)`. +#[inline] +pub fn masked_group_sum_sym_i32_pair( + mask_words: &[u64], hi: &[u32], lo: &[u32], stride: u32, values: &[i32], out: &mut [i64], +) { + assert_eq!(hi.len(), lo.len(), "masked_group_sum_sym_i32_pair: hi.len() != lo.len()"); + assert_eq!(hi.len(), values.len(), "masked_group_sum_sym_i32_pair: hi/values length mismatch"); + group_walk( + "masked_group_sum_sym_i32_pair", + mask_words, + hi.len(), + GroupKeyAddr::Pair { hi, lo, stride }, + out, + sym_sum_fold(values), + ); +} + /// The one `_sym` sum fold, shared by both key addresses. #[inline(always)] fn sym_sum_fold(values: &[i32]) -> impl FnMut(&mut i64, usize) + '_ { @@ -1487,18 +1561,33 @@ enum GroupKeyAddr<'a> { /// The key of row `i` is `table[index[i]]`. The indirection is fused: /// no remapped key lane is ever materialised between the two hops. Via { index: &'a [u32], table: &'a [u32] }, + /// The key of row `i` is `hi[i] * stride + lo[i]` — the composite + /// address of a two-column `GROUP BY`. Fused: no composite key lane is + /// ever materialised. + Pair { hi: &'a [u32], lo: &'a [u32], stride: u32 }, } impl GroupKeyAddr<'_> { /// The group of row `i`, or `None` when either hop names nothing: an - /// index past `table` (VIA only) or a key past the group universe. Both - /// drops are the zero-fallback contract of [`mask_gather_u32`] — an - /// unminted address is not a group — never an error. + /// index past `table` (VIA only), a minor key at or past `stride` + /// (Pair only — a minor key outside its own declared cardinality names + /// no group), or a key past the group universe. Every drop is the + /// zero-fallback contract of [`mask_gather_u32`] — an unminted address + /// is not a group — never an error. #[inline(always)] fn group_of(self, i: usize, groups: usize) -> Option { let k = match self { GroupKeyAddr::Resident(keys) => keys[i] as usize, GroupKeyAddr::Via { index, table } => *table.get(index[i] as usize)? as usize, + GroupKeyAddr::Pair { hi, lo, stride } => { + if lo[i] >= stride { + return None; + } + // hi and lo are both u32, stride is u32: widen to u64 first + // so the multiply-add cannot overflow. + let composite = hi[i] as u64 * stride as u64 + lo[i] as u64; + usize::try_from(composite).ok()? + } }; (k < groups).then_some(k) } @@ -1608,6 +1697,50 @@ pub fn masked_group_count_u32_via(mask_words: &[u64], index: &[u32], table: &[u3 ); } +/// [`masked_group_count_u32`] with the group of row `i` read as the +/// composite address `hi[i] * stride + lo[i]` — a two-column `GROUP BY +/// (hi, lo)` fused into one pass, no composite key lane ever materialised. +/// +/// **Zero-fallback**: `lo[i] >= stride` names no group at all (a minor key +/// outside its own declared cardinality) and drops row `i`; a resolved +/// composite key past `out.len()` drops it too — the same rule as +/// [`masked_group_sum_i32_via`]'s two hops. +/// +/// Same contract as [`masked_group_sum_i32_pair`] in every other respect: +/// `out` is accumulated into (the caller zeroes it once) with +/// `wrapping_add`, and the final mask word is clamped to `hi.len()`. +/// +/// # Panics +/// +/// Panics if `hi.len() != lo.len()`, or if `mask_words.len() < +/// hi.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::masked_group_count_u32_pair; +/// +/// let mask = [0b111u64]; +/// let hi = [0u32, 1, 0]; +/// let lo = [0u32, 0, 4]; // stride 4: row 2's lo (4) is out of range and dropped +/// let mut out = [0i64; 8]; // group (hi, lo) = hi * 4 + lo +/// masked_group_count_u32_pair(&mask, &hi, &lo, 4, &mut out); +/// assert_eq!(out[0], 1); // (0, 0) -> 0 +/// assert_eq!(out[4], 1); // (1, 0) -> 4 +/// ``` +#[inline] +pub fn masked_group_count_u32_pair(mask_words: &[u64], hi: &[u32], lo: &[u32], stride: u32, out: &mut [i64]) { + assert_eq!(hi.len(), lo.len(), "masked_group_count_u32_pair: hi.len() != lo.len()"); + group_walk( + "masked_group_count_u32_pair", + mask_words, + hi.len(), + GroupKeyAddr::Pair { hi, lo, stride }, + out, + |slot, _| *slot = slot.wrapping_add(1), + ); +} + /// Keyed group-minimum: for every row `i` selected by `mask_words`, lowers /// `out[keys[i]]` to `values[i]` (widened to `i64`) if that is smaller — /// `MIN(values) … GROUP BY key` in one pass. @@ -1679,6 +1812,44 @@ pub fn masked_group_min_i32_via(mask_words: &[u64], index: &[u32], table: &[u32] ); } +/// [`masked_group_min_i32`] with the group of row `i` read as the composite +/// address `hi[i] * stride + lo[i]`, zero-fallback exactly as +/// [`masked_group_count_u32_pair`]. Seed `out` with `i64::MAX`. +/// +/// # Panics +/// +/// Panics if `hi.len() != lo.len()`, if `hi.len() != values.len()`, or if +/// `mask_words.len() < hi.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::masked_group_min_i32_pair; +/// +/// let mask = [0b111u64]; +/// let hi = [0u32, 0, 0]; +/// let lo = [0u32, 0, 4]; // stride 4: row 2's lo (4) is out of range and dropped +/// let values = [4i32, -9, 999]; +/// let mut out = [i64::MAX; 1]; +/// masked_group_min_i32_pair(&mask, &hi, &lo, 4, &values, &mut out); +/// assert_eq!(out, [-9]); +/// ``` +#[inline] +pub fn masked_group_min_i32_pair( + mask_words: &[u64], hi: &[u32], lo: &[u32], stride: u32, values: &[i32], out: &mut [i64], +) { + assert_eq!(hi.len(), lo.len(), "masked_group_min_i32_pair: hi.len() != lo.len()"); + assert_eq!(hi.len(), values.len(), "masked_group_min_i32_pair: hi/values length mismatch"); + group_walk( + "masked_group_min_i32_pair", + mask_words, + hi.len(), + GroupKeyAddr::Pair { hi, lo, stride }, + out, + |slot, i| *slot = (*slot).min(values[i] as i64), + ); +} + /// Keyed group-maximum: the mirror of [`masked_group_min_i32`]. **The caller /// seeds `out` with `i64::MIN`**, which no `i32` can reach, so a slot still /// holding `i64::MIN` afterwards is a group no selected row named. @@ -1743,6 +1914,44 @@ pub fn masked_group_max_i32_via(mask_words: &[u64], index: &[u32], table: &[u32] ); } +/// [`masked_group_max_i32`] with the group of row `i` read as the composite +/// address `hi[i] * stride + lo[i]`, zero-fallback exactly as +/// [`masked_group_count_u32_pair`]. Seed `out` with `i64::MIN`. +/// +/// # Panics +/// +/// Panics if `hi.len() != lo.len()`, if `hi.len() != values.len()`, or if +/// `mask_words.len() < hi.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::masked_group_max_i32_pair; +/// +/// let mask = [0b111u64]; +/// let hi = [0u32, 0, 0]; +/// let lo = [0u32, 0, 4]; // stride 4: row 2's lo (4) is out of range and dropped +/// let values = [4i32, 99, 999]; +/// let mut out = [i64::MIN; 1]; +/// masked_group_max_i32_pair(&mask, &hi, &lo, 4, &values, &mut out); +/// assert_eq!(out, [99]); +/// ``` +#[inline] +pub fn masked_group_max_i32_pair( + mask_words: &[u64], hi: &[u32], lo: &[u32], stride: u32, values: &[i32], out: &mut [i64], +) { + assert_eq!(hi.len(), lo.len(), "masked_group_max_i32_pair: hi.len() != lo.len()"); + assert_eq!(hi.len(), values.len(), "masked_group_max_i32_pair: hi/values length mismatch"); + group_walk( + "masked_group_max_i32_pair", + mask_words, + hi.len(), + GroupKeyAddr::Pair { hi, lo, stride }, + out, + |slot, i| *slot = (*slot).max(values[i] as i64), + ); +} + /// Packs `index[i] < table.len() && table[index[i]] == v` into `out_words`, /// one bit per row `i < index.len()`, LSB-first — an equality predicate /// evaluated **through an index lane**, with no gathered mask and no @@ -7271,4 +7480,112 @@ mod group_family_tests { fn min_refuses_mismatched_lanes() { masked_group_min_i32(&[1], &[0, 0], &[1], &mut [0]); } + + // ── the `_pair` (composite two-column) address ────────────────────── + + /// Every `_pair` fold computes exactly the group the `Resident` form + /// computes when fed a precomputed `hi * stride + lo` key lane — proving + /// the fused composite address is not a new algorithm, just a different + /// route to the same key the caller could have materialised by hand. + #[test] + fn pair_matches_resident_on_a_precomputed_composite_key() { + let n = 1000usize; + let mut s = 0xC0FF_EE00_1234_5678u64 ^ n as u64; + let stride: u32 = 7; + let words = n.div_ceil(64); + // Mask words are left with their natural high bits past `n` — + // the same dirty-tail construction `fixture` uses above. + let mask: Vec = (0..words) + .map(|_| lcg(&mut s) ^ (lcg(&mut s) << 32)) + .collect(); + let hi: Vec = (0..n).map(|_| (lcg(&mut s) % 9) as u32).collect(); // 0..=8 + let lo: Vec = (0..n).map(|_| (lcg(&mut s) % 7) as u32).collect(); // 0..=6, always < stride + let values: Vec = (0..n).map(|_| (lcg(&mut s) % 101) as i32 - 50).collect(); // -50..=50 + let groups = 63usize; // 8 * 7 + 6 = 62 is the largest reachable key + let keys: Vec = hi.iter().zip(&lo).map(|(&h, &l)| h * stride + l).collect(); + + let mut count_pair = vec![0i64; groups]; + masked_group_count_u32_pair(&mask, &hi, &lo, stride, &mut count_pair); + let mut count_res = vec![0i64; groups]; + masked_group_count_u32(&mask, &keys, &mut count_res); + assert_eq!(count_pair, count_res, "count"); + assert!( + count_pair.iter().filter(|&&v| v != 0).count() >= 20, + "anti-vacuity: fewer than 20 distinct non-identity slots, the fixture barely exercises the fold" + ); + + let mut sum_pair = vec![0i64; groups]; + masked_group_sum_i32_pair(&mask, &hi, &lo, stride, &values, &mut sum_pair); + let mut sum_res = vec![0i64; groups]; + masked_group_sum_i32(&mask, &keys, &values, &mut sum_res); + assert_eq!(sum_pair, sum_res, "sum"); + + let mut sym_pair = vec![SYM_EMPTY_I64; groups]; + masked_group_sum_sym_i32_pair(&mask, &hi, &lo, stride, &values, &mut sym_pair); + let mut sym_res = vec![SYM_EMPTY_I64; groups]; + masked_group_sum_sym_i32(&mask, &keys, &values, &mut sym_res); + assert_eq!(sym_pair, sym_res, "sym sum"); + + let mut min_pair = vec![i64::MAX; groups]; + masked_group_min_i32_pair(&mask, &hi, &lo, stride, &values, &mut min_pair); + let mut min_res = vec![i64::MAX; groups]; + masked_group_min_i32(&mask, &keys, &values, &mut min_res); + assert_eq!(min_pair, min_res, "min"); + + let mut max_pair = vec![i64::MIN; groups]; + masked_group_max_i32_pair(&mask, &hi, &lo, stride, &values, &mut max_pair); + let mut max_res = vec![i64::MIN; groups]; + masked_group_max_i32(&mask, &keys, &values, &mut max_res); + assert_eq!(max_pair, max_res, "max"); + } + + /// A minor key at or past `stride` names no group at all, regardless of + /// `hi`. The universe is sized so the dropped rows WOULD alias a real + /// slot without the guard: `(hi=0, lo=4)` composes to 4, which is + /// exactly `(hi=1, lo=0)`'s group, and `(hi=1, lo=5)` composes to 9, + /// `(hi=2, lo=1)`'s. Only the `lo >= stride` drop keeps them out; the + /// group-universe drop cannot, since both composites are `< groups`. + /// Checked both ways: the kept rows ARE counted where they belong. + #[test] + fn pair_drops_a_minor_key_at_or_past_stride() { + let mask = [0b11111u64]; // rows 0..5 selected + let hi = [0u32, 0, 1, 1, 0]; + let lo = [3u32, 4, 5, 0, u32::MAX]; + let stride = 4u32; + let mut out = [0i64; 12]; // 3 majors x stride 4 + masked_group_count_u32_pair(&mask, &hi, &lo, stride, &mut out); + // Kept: row 0 -> 0*4+3 = 3, row 3 -> 1*4+0 = 4. + assert_eq!(out[3], 1, "(0, 3) is counted"); + assert_eq!(out[4], 1, "(1, 0) is counted once, and only by its own row"); + // Row 1 (lo=4 -> 4) and row 2 (lo=5 -> 9) would alias real slots. + assert_eq!(out[9], 0, "(1, 5) must not alias (2, 1)"); + assert_eq!(out.iter().sum::(), 2, "the lo >= stride rows contribute nothing anywhere"); + } + + /// A composite key that would address far past the group universe is + /// dropped like any other out-of-range key, and the `hi * stride + lo` + /// widening never panics even when both `hi` and `stride` are + /// `u32::MAX` — the multiply-add is done in `u64`, which cannot overflow. + #[test] + fn pair_drops_a_composite_past_the_group_universe_without_overflow() { + let mask = [0b11u64]; // rows 0, 1 + let hi = [u32::MAX, 0u32]; + let lo = [0u32, 1u32]; + let stride = u32::MAX; + let mut out = [0i64; 2]; + masked_group_count_u32_pair(&mask, &hi, &lo, stride, &mut out); + // Row 0's composite key (u32::MAX * u32::MAX + 0, ~1.8e19) is far + // past `out.len()` and is dropped without panicking; row 1's + // composite key (0 * u32::MAX + 1 = 1) lands cleanly in out[1]. + assert_eq!(out, [0, 1]); + } + + /// `hi` and `lo` of different lengths are refused before any row is + /// read, exactly the `keys/values`-style length guard the other members + /// of this family already carry. + #[test] + #[should_panic(expected = "hi.len() != lo.len()")] + fn pair_panics_on_mismatched_lengths() { + masked_group_count_u32_pair(&[1], &[0, 0], &[0], 4, &mut [0]); + } }