diff --git a/src/kernels/bitmask/mod.rs b/src/kernels/bitmask/mod.rs index f7c1383..e3b89bd 100644 --- a/src/kernels/bitmask/mod.rs +++ b/src/kernels/bitmask/mod.rs @@ -406,4 +406,144 @@ mod tests { assert!(out.get(i), "bit {i} past the window was disturbed"); } } + + /// Window offsets for the bitmask kernel sweeps, covering starts inside a + /// byte, byte-aligned starts off a word boundary and word-aligned starts. + pub(super) const WINDOW_OFFSETS: [usize; 13] = [0, 1, 3, 5, 7, 8, 9, 13, 56, 63, 64, 65, 72]; + + /// Window lengths for the bitmask kernel sweeps, covering partial bytes, + /// partial final words, whole words and several whole SIMD batches. + pub(super) const WINDOW_LENS: [usize; 15] = + [1, 7, 8, 9, 56, 57, 60, 63, 64, 65, 127, 128, 129, 200, 1000]; + + /// Builds an `n`-bit mask with bit `i` set where `i % a == 0 || i % b == 1`. + /// + /// Coprime `a` and `b` give a pattern that does not repeat every 8 bits, + /// which exposes any read shifted by part of a byte. + pub(super) fn pattern_mask(n: usize, a: usize, b: usize) -> Bitmask { + let mut mask = Bitmask::new_set_all(n, false); + for i in 0..n { + if i % a == 0 || i % b == 1 { + mask.set(i, true); + } + } + mask + } + + /// Reads `len` bits from `offset` one at a time through `Bitmask::get`, as + /// the reference for the word-level kernels. + pub(super) fn window_bits(mask: &Bitmask, offset: usize, len: usize) -> Vec { + (0..len).map(|i| mask.get(offset + i)).collect() + } + + /// `popcount_mask` counts the window bits when the window starts inside a byte. + #[test] + fn popcount_counts_windows_starting_inside_a_byte() { + use crate::kernels::bitmask::dispatch::popcount_mask; + + let trailing_null = Bitmask::from_bools(&[true, true, true, true, true, true, true, false]); + let leading_null = Bitmask::from_bools(&[false, true, true, true, true, true, true, true]); + let all_set = Bitmask::new_set_all(24, true); + + for (mask, offset, len) in [ + (&trailing_null, 1, 7), + (&leading_null, 1, 7), + (&all_set, 5, 12), + ] { + let expected = window_bits(mask, offset, len) + .into_iter() + .filter(|&bit| bit) + .count(); + assert_eq!( + popcount_mask((mask, offset, len)), + expected, + "window ({offset}, {len})" + ); + } + } + + /// `popcount_mask` counts a window that starts on a byte off a word boundary. + #[test] + fn popcount_counts_windows_starting_off_a_word_boundary() { + use crate::kernels::bitmask::dispatch::popcount_mask; + + let mask = Bitmask::new_set_all(128, true); + let expected = window_bits(&mask, 8, 64) + .into_iter() + .filter(|&bit| bit) + .count(); + assert_eq!(popcount_mask((&mask, 8, 64)), expected); + } + + /// `popcount_mask` excludes the bits past the window in its final partial word. + #[test] + fn popcount_excludes_bits_past_the_window() { + use crate::kernels::bitmask::dispatch::popcount_mask; + + let mask = Bitmask::new_set_all(128, true); + let expected = window_bits(&mask, 0, 60) + .into_iter() + .filter(|&bit| bit) + .count(); + assert_eq!(popcount_mask((&mask, 0, 60)), expected); + } + + /// `and_masks` combines windows that start inside a byte and end within one word. + #[test] + fn and_masks_combines_windows_starting_inside_a_byte() { + use crate::kernels::bitmask::dispatch::and_masks; + + let mut a = Bitmask::new_set_all(16, true); + a.set(5, false); + let mut b = Bitmask::new_set_all(16, true); + b.set(9, false); + + let out = and_masks((&a, 3, 10), (&b, 3, 10)); + let expected: Vec = window_bits(&a, 3, 10) + .into_iter() + .zip(window_bits(&b, 3, 10)) + .map(|(x, y)| x && y) + .collect(); + assert_eq!(out.len(), 10); + assert_eq!(window_bits(&out, 0, 10), expected); + } + + /// `and_masks` combines whole words of windows that start inside a byte. + #[test] + fn and_masks_combines_whole_words_starting_inside_a_byte() { + use crate::kernels::bitmask::dispatch::and_masks; + + // Bit 2 lies before the window, which starts at bit 3. + let mut a = Bitmask::new_set_all(136, true); + a.set(2, false); + let b = Bitmask::new_set_all(136, true); + + let out = and_masks((&a, 3, 64), (&b, 3, 64)); + let expected: Vec = window_bits(&a, 3, 64) + .into_iter() + .zip(window_bits(&b, 3, 64)) + .map(|(x, y)| x && y) + .collect(); + assert_eq!(out.len(), 64); + assert_eq!(window_bits(&out, 0, 64), expected); + } + + /// `and_masks` combines windows that start on a byte off a word boundary. + #[test] + fn and_masks_combines_windows_starting_off_a_word_boundary() { + use crate::kernels::bitmask::dispatch::and_masks; + + let mut a = Bitmask::new_set_all(128, true); + a.set(20, false); + let b = Bitmask::new_set_all(128, true); + + let out = and_masks((&a, 8, 64), (&b, 8, 64)); + let expected: Vec = window_bits(&a, 8, 64) + .into_iter() + .zip(window_bits(&b, 8, 64)) + .map(|(x, y)| x && y) + .collect(); + assert_eq!(out.len(), 64); + assert_eq!(window_bits(&out, 0, 64), expected); + } } diff --git a/src/kernels/bitmask/simd.rs b/src/kernels/bitmask/simd.rs index bbfc6dc..9c5bda0 100644 --- a/src/kernels/bitmask/simd.rs +++ b/src/kernels/bitmask/simd.rs @@ -145,7 +145,8 @@ pub fn bitmask_unop_simd(src: BitmaskVT<'_>, op: UnaryOperat /// the words never straddle a window edge, so adjacent windows of a shared /// output buffer stay independent. /// -/// All offsets are byte-aligned i.e. a multiple of 8 bits. +/// Input offsets may be any bit position. `out_off` is byte-aligned i.e. a +/// multiple of 8 bits. #[inline(always)] pub fn bitmask_binop_simd_into( out: &mut Bitmask, @@ -164,51 +165,77 @@ pub fn bitmask_binop_simd_into( 0, "bitmask_binop_simd_into: out_off must be byte-aligned" ); - debug_assert_eq!( - lhs_off % 8, - 0, - "bitmask_binop_simd_into: lhs offset must be byte-aligned" - ); - debug_assert_eq!( - rhs_off % 8, - 0, - "bitmask_binop_simd_into: rhs offset must be byte-aligned" - ); let full_words = len / 64; let tail_bits = len % 64; { - let lhs_bytes = bitmask_window_bytes(lhs_mask, lhs_off, len); - let rhs_bytes = bitmask_window_bytes(rhs_mask, rhs_off, len); + // Each input is read from the byte holding its first window bit. For + // windows starting inside a byte, each word is joined with the low bits + // of the following word through a shift and an OR. Output words are + // copied in as little-endian bytes, which needs no word alignment of + // the output window. + let lhs_bytes: &[u8] = &lhs_mask.bits[lhs_off / 8..]; + let rhs_bytes: &[u8] = &rhs_mask.bits[rhs_off / 8..]; + let lhs_shift = (lhs_off % 8) as u64; + let rhs_shift = (rhs_off % 8) as u64; let out_bytes = bitmask_window_bytes_mut(out, out_off, len); - unsafe { - let lp = lhs_bytes.as_ptr().cast::(); - let rp = rhs_bytes.as_ptr().cast::(); - let dp = out_bytes.as_mut_ptr().cast::(); - let mut i = 0; - while i + LANES <= full_words { - let a = - Simd::::from_slice(std::slice::from_raw_parts(lp.add(i), LANES)); - let b = - Simd::::from_slice(std::slice::from_raw_parts(rp.add(i), LANES)); - let r = match op { - LogicalOperator::And => a & b, - LogicalOperator::Or => a | b, - LogicalOperator::Xor => a ^ b, - }; - std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); - i += LANES; + + let words_at = |bytes: &[u8], k: usize| { + let mut words = [0u64; LANES]; + for (word, chunk) in words + .iter_mut() + .zip(bytes[k * 8..(k + LANES) * 8].chunks_exact(8)) + { + *word = u64::from_le_bytes(chunk.try_into().unwrap()); } - // Scalar word tail when full_words is not a multiple of LANES. - for k in i..full_words { - let a = *lp.add(k); - let b = *rp.add(k); - *dp.add(k) = match op { - LogicalOperator::And => a & b, - LogicalOperator::Or => a | b, - LogicalOperator::Xor => a ^ b, - }; + Simd::::from_array(words) + }; + let window_batch = |bytes: &[u8], shift: u64, k: usize| { + if shift == 0 { + words_at(bytes, k) + } else { + (words_at(bytes, k) >> Simd::splat(shift)) + | (words_at(bytes, k + 1) << Simd::splat(64 - shift)) + } + }; + // Shifted batches also read the word after the batch, which has to lie + // within the buffer. + let batchable = |bytes: &[u8], shift: u64| { + if shift == 0 { + full_words + } else { + full_words.min((bytes.len() / 8).saturating_sub(1)) + } + }; + let n_batched = + batchable(lhs_bytes, lhs_shift).min(batchable(rhs_bytes, rhs_shift)) / LANES * LANES; + + for k in (0..n_batched).step_by(LANES) { + let a = window_batch(lhs_bytes, lhs_shift, k); + let b = window_batch(rhs_bytes, rhs_shift, k); + let r = match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }; + for (chunk, word) in out_bytes[k * 8..(k + LANES) * 8] + .chunks_exact_mut(8) + .zip(r.to_array()) + { + chunk.copy_from_slice(&word.to_le_bytes()); } } + // Scalar word tail when full_words is not a multiple of LANES, or when + // a shifted batch would read past the end of the buffer. + for k in n_batched..full_words { + let a = load_word(lhs_mask, lhs_off + k * 64); + let b = load_word(rhs_mask, rhs_off + k * 64); + let word = match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }; + out_bytes[k * 8..k * 8 + 8].copy_from_slice(&word.to_le_bytes()); + } } // Final partial word written one bit at a time. let base = full_words * 64; @@ -233,7 +260,8 @@ pub fn bitmask_binop_simd_into( /// remaining whole words scalar. The final `len % 64` bits run one bit at a /// time, so the write never touches the trailing bits past `len`. /// -/// All offsets are byte-aligned i.e. a multiple of 8 bits. +/// The input offset may be any bit position. `out_off` is byte-aligned i.e. a +/// multiple of 8 bits. #[inline(always)] pub fn bitmask_unop_simd_into( out: &mut Bitmask, @@ -250,38 +278,63 @@ pub fn bitmask_unop_simd_into( 0, "bitmask_unop_simd_into: out_off must be byte-aligned" ); - debug_assert_eq!( - offset % 8, - 0, - "bitmask_unop_simd_into: src offset must be byte-aligned" - ); let full_words = len / 64; let tail_bits = len % 64; { - let src_bytes = bitmask_window_bytes(mask, offset, len); + // The input is read from the byte holding its first window bit. For + // windows starting inside a byte, each word is joined with the low bits + // of the following word through a shift and an OR. Output words are + // copied in as little-endian bytes, which needs no word alignment of + // the output window. + let src_bytes: &[u8] = &mask.bits[offset / 8..]; + let shift = (offset % 8) as u64; let out_bytes = bitmask_window_bytes_mut(out, out_off, len); - unsafe { - let sp = src_bytes.as_ptr().cast::(); - let dp = out_bytes.as_mut_ptr().cast::(); - let mut i = 0; - while i + LANES <= full_words { - let a = - Simd::::from_slice(std::slice::from_raw_parts(sp.add(i), LANES)); - let r = match op { - UnaryOperator::Not => !a, - _ => unreachable!(), - }; - std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); - i += LANES; + + let words_at = |k: usize| { + let mut words = [0u64; LANES]; + for (word, chunk) in words + .iter_mut() + .zip(src_bytes[k * 8..(k + LANES) * 8].chunks_exact(8)) + { + *word = u64::from_le_bytes(chunk.try_into().unwrap()); } - // Scalar word tail when full_words is not a multiple of LANES. - for k in i..full_words { - *dp.add(k) = match op { - UnaryOperator::Not => !*sp.add(k), - _ => unreachable!(), - }; + Simd::::from_array(words) + }; + // Shifted batches also read the word after the batch, which has to lie + // within the buffer. + let batchable = if shift == 0 { + full_words + } else { + full_words.min((src_bytes.len() / 8).saturating_sub(1)) + }; + let n_batched = batchable / LANES * LANES; + + for k in (0..n_batched).step_by(LANES) { + let a = if shift == 0 { + words_at(k) + } else { + (words_at(k) >> Simd::splat(shift)) | (words_at(k + 1) << Simd::splat(64 - shift)) + }; + let r = match op { + UnaryOperator::Not => !a, + _ => unreachable!(), + }; + for (chunk, word) in out_bytes[k * 8..(k + LANES) * 8] + .chunks_exact_mut(8) + .zip(r.to_array()) + { + chunk.copy_from_slice(&word.to_le_bytes()); } } + // Scalar word tail when full_words is not a multiple of LANES, or when + // a shifted batch would read past the end of the buffer. + for k in n_batched..full_words { + let word = match op { + UnaryOperator::Not => !load_word(mask, offset + k * 64), + _ => unreachable!(), + }; + out_bytes[k * 8..k * 8 + 8].copy_from_slice(&word.to_le_bytes()); + } } // Final partial word written one bit at a time. let base = full_words * 64; @@ -429,48 +482,45 @@ where return Bitmask::new_set_all(0, false); } - // Check which boolean values are present in rhs. Process whole u64 words - // that fit entirely in the tight byte slice, then handle the 0..7 - // leftover bytes byte-by-byte. The partial last logical byte (when - // len & 7 != 0) is masked to its valid bits to avoid false positives - // from slack bits. - let rhs_bytes = bitmask_window_bytes(rhs_mask, rhs_off, len); - let total_bytes = rhs_bytes.len(); - let full_words = total_bytes / 8; - let tail_bytes = total_bytes % 8; - let full_logical_bytes = len / 8; - let last_bits = len & 7; + // Check which boolean values are present in rhs. Whole words are read from + // the byte holding the first window bit. For windows starting inside a + // byte, each word is joined with the low bits of the following byte. The + // final partial word is read through `load_word` and masked to the window, + // which keeps bits outside the window out of both checks. + let rhs_bytes: &[u8] = &rhs_mask.bits[rhs_off / 8..]; + let shift = rhs_off % 8; + let full_words = len / 64; let mut any_set = 0u64; let mut any_unset = 0u64; - unsafe { - let rp = rhs_bytes.as_ptr().cast::(); - for k in 0..full_words { - let w = *rp.add(k); - any_set |= w; - any_unset |= !w; - if any_set != 0 && any_unset != 0 { + let mut scan = |w: u64| { + any_set |= w; + any_unset |= !w; + any_set != 0 && any_unset != 0 + }; + if shift == 0 { + for chunk in rhs_bytes[..full_words * 8].chunks_exact(8) { + if scan(u64::from_le_bytes(chunk.try_into().unwrap())) { break; } } - } - if any_set == 0 || any_unset == 0 { - let base = full_words * 8; - for k in 0..tail_bytes { - let byte_index = base + k; - let b = rhs_bytes[byte_index]; - if byte_index < full_logical_bytes { - any_set |= b as u64; - any_unset |= (!b) as u64; - } else if last_bits != 0 { - let m = (1u8 << last_bits) - 1; - any_set |= (b & m) as u64; - any_unset |= ((!b) & m) as u64; - } - if any_set != 0 && any_unset != 0 { + } else { + // The bytes of every whole window word, and the byte after each, lie + // within the window. + let body = &rhs_bytes[..full_words * 8 + 1]; + for k in 0..full_words { + let lo = u64::from_le_bytes(body[k * 8..k * 8 + 8].try_into().unwrap()); + if scan((lo >> shift) | ((body[k * 8 + 8] as u64) << (64 - shift))) { break; } } } + let tail_bits = len % 64; + if tail_bits != 0 && (any_set == 0 || any_unset == 0) { + let in_window = u64::MAX >> (64 - tail_bits); + let w = load_word(rhs_mask, rhs_off + full_words * 64); + any_set |= w & in_window; + any_unset |= !w & in_window; + } let has_true = any_set != 0; let has_false = any_unset != 0; @@ -638,6 +688,8 @@ where pub fn all_eq_mask_simd(a: BitmaskVT<'_>, b: BitmaskVT<'_>) -> bool where { + use std::simd::prelude::SimdPartialEq; + let (am, ao, len) = a; let (bm, bo, blen) = b; debug_assert_eq!(len, blen, "BitWindow length mismatch in all_eq_mask"); @@ -653,66 +705,43 @@ where ); } - let a_bytes = bitmask_window_bytes(am, ao, len); - let b_bytes = bitmask_window_bytes(bm, bo, len); - // Whole u64 words within the tight byte slice, then byte-by-byte - // for the 0..7 leftover. The partial last logical byte is masked - // to its valid bits so slack differences don't cause false negatives. - let total_bytes = a_bytes.len(); - let full_words = total_bytes / 8; - let tail_bytes = total_bytes % 8; - let full_logical_bytes = len / 8; - let last_bits = len & 7; - - unsafe { - let ap = a_bytes.as_ptr().cast::(); - let bp = b_bytes.as_ptr().cast::(); - - #[cfg(feature = "simd")] + // Both windows start on a word boundary, and whole words are read from the + // window bytes. The final partial word is read through `load_word` and + // masked to the window, which keeps differences past the window from + // causing false negatives. + let a_bytes: &[u8] = &am.bits[ao / 8..]; + let b_bytes: &[u8] = &bm.bits[bo / 8..]; + let full_words = len / 64; + let n_batched = full_words / LANES * LANES; + let words_at = |bytes: &[u8], k: usize| { + let mut words = [0u64; LANES]; + for (word, chunk) in words + .iter_mut() + .zip(bytes[k * 8..(k + LANES) * 8].chunks_exact(8)) { - use std::simd::prelude::SimdPartialEq; - let mut i = 0; - while i + LANES <= full_words { - let sa = - Simd::::from_slice(std::slice::from_raw_parts(ap.add(i), LANES)); - let sb = - Simd::::from_slice(std::slice::from_raw_parts(bp.add(i), LANES)); - if !sa.simd_eq(sb).all() { - return false; - } - i += LANES; - } - for k in i..full_words { - if *ap.add(k) != *bp.add(k) { - return false; - } - } + *word = u64::from_le_bytes(chunk.try_into().unwrap()); } - #[cfg(not(feature = "simd"))] - { - for k in 0..full_words { - if *ap.add(k) != *bp.add(k) { - return false; - } - } + Simd::::from_array(words) + }; + + for k in (0..n_batched).step_by(LANES) { + if !words_at(a_bytes, k).simd_eq(words_at(b_bytes, k)).all() { + return false; } } - // Partial-byte tail: compare byte-by-byte, masking the partial - // last logical byte to its valid bits. - let base = full_words * 8; - for k in 0..tail_bytes { - let byte_index = base + k; - let av = a_bytes[byte_index]; - let bv = b_bytes[byte_index]; - if byte_index < full_logical_bytes { - if av != bv { - return false; - } - } else if last_bits != 0 { - let m = (1u8 << last_bits) - 1; - if (av & m) != (bv & m) { - return false; - } + for k in n_batched..full_words { + let a = u64::from_le_bytes(a_bytes[k * 8..k * 8 + 8].try_into().unwrap()); + let b = u64::from_le_bytes(b_bytes[k * 8..k * 8 + 8].try_into().unwrap()); + if a != b { + return false; + } + } + let tail_bits = len % 64; + if tail_bits != 0 { + let base = full_words * 64; + let in_window = u64::MAX >> (64 - tail_bits); + if (load_word(am, ao + base) ^ load_word(bm, bo + base)) & in_window != 0 { + return false; } } true @@ -734,60 +763,55 @@ where pub fn popcount_mask_simd(m: BitmaskVT<'_>) -> usize where { + use std::simd::prelude::SimdUint; + let (mask, offset, len) = m; if len == 0 { return 0; } - let bytes = bitmask_window_bytes(mask, offset, len); - // Whole u64 words within the tight byte slice via SIMD, scalar - // tail for partial u64 words, then byte-by-byte for the leftover - // 0..7 bytes. The partial last logical byte is masked to its - // valid bits so slack ones don't inflate the count. - let total_bytes = bytes.len(); - let full_words = total_bytes / 8; - let tail_bytes = total_bytes % 8; - let full_logical_bytes = len / 8; - let last_bits = len & 7; - let mut acc = 0usize; - - unsafe { - let words_ptr = bytes.as_ptr().cast::(); - - #[cfg(feature = "simd")] - { - use std::simd::prelude::SimdUint; - let mut i = 0; - while i + LANES <= full_words { - let v = Simd::::from_slice(std::slice::from_raw_parts( - words_ptr.add(i), - LANES, - )); - acc += v.count_ones().reduce_sum() as usize; - i += LANES; - } - for k in i..full_words { - acc += (*words_ptr.add(k)).count_ones() as usize; - } - } - #[cfg(not(feature = "simd"))] + // Whole words are read from the byte holding the first window bit. For + // windows starting inside a byte, each word is joined with the low bits of + // the following word through a shift and an OR. Whole batches of `LANES` + // words are reduced through SIMD. The remaining words are read through + // `load_word`, with the final partial word masked to the window, which + // keeps bits outside the window out of the count. + let bytes: &[u8] = &mask.bits[offset / 8..]; + let shift = (offset % 8) as u64; + let full_words = len / 64; + let words_at = |k: usize| { + let mut words = [0u64; LANES]; + for (word, chunk) in words + .iter_mut() + .zip(bytes[k * 8..(k + LANES) * 8].chunks_exact(8)) { - for k in 0..full_words { - acc += (*words_ptr.add(k)).count_ones() as usize; - } + *word = u64::from_le_bytes(chunk.try_into().unwrap()); } + Simd::::from_array(words) + }; + // Shifted batches also read the word after the batch, which has to lie + // within the buffer. + let batchable = if shift == 0 { + full_words + } else { + full_words.min((bytes.len() / 8).saturating_sub(1)) + }; + let n_batched = batchable / LANES * LANES; + + let mut acc = 0usize; + for k in (0..n_batched).step_by(LANES) { + let words = if shift == 0 { + words_at(k) + } else { + (words_at(k) >> Simd::splat(shift)) | (words_at(k + 1) << Simd::splat(64 - shift)) + }; + acc += words.count_ones().reduce_sum() as usize; } - // Partial-byte tail. - let base = full_words * 8; - for k in 0..tail_bytes { - let byte_index = base + k; - let b = bytes[byte_index]; - if byte_index < full_logical_bytes { - acc += b.count_ones() as usize; - } else if last_bits != 0 { - let m = (1u8 << last_bits) - 1; - acc += (b & m).count_ones() as usize; - } + for k in n_batched..len.div_ceil(64) { + let bit_base = k * 64; + let take = (len - bit_base).min(64); + let word = load_word(mask, offset + bit_base) & (u64::MAX >> (64 - take)); + acc += word.count_ones() as usize; } acc } @@ -1059,6 +1083,7 @@ impl_simd_eq_mask!(simd_eq_mask_u64, u64, W64); #[cfg(test)] mod tests { + use crate::kernels::bitmask::tests::{WINDOW_LENS, WINDOW_OFFSETS, pattern_mask, window_bits}; use crate::{Bitmask, BitmaskVT}; use super::*; @@ -1240,6 +1265,212 @@ mod tests { (0..len).filter(|&i| bits[offset + i]).collect(); assert_eq!(windowed, expected_windowed); } + + // The extra length fills one whole batch of `LANES` words and + // leaves a partial word after it. + fn sweep_lens() -> impl Iterator { + WINDOW_LENS.into_iter().chain([64 * LANES + 37]) + } + + #[test] + fn test_popcount_mask_simd_windows() { + for offset in WINDOW_OFFSETS { + for len in sweep_lens() { + let mask = pattern_mask(offset + len + 67, 3, 7); + let expected = window_bits(&mask, offset, len) + .into_iter() + .filter(|&bit| bit) + .count(); + assert_eq!( + popcount_mask_simd::((&mask, offset, len)), + expected, + "offset {offset} len {len}" + ); + } + } + } + + #[test] + fn test_bitmask_binop_simd_into_windows() { + let ops = [ + LogicalOperator::And, + LogicalOperator::Or, + LogicalOperator::Xor, + ]; + for (k, lhs_off) in WINDOW_OFFSETS.into_iter().enumerate() { + let rhs_off = WINDOW_OFFSETS[(k + 5) % WINDOW_OFFSETS.len()]; + for len in sweep_lens() { + let lhs = pattern_mask(lhs_off + len + 67, 3, 7); + let rhs = pattern_mask(rhs_off + len + 67, 5, 11); + let lhs_bits = window_bits(&lhs, lhs_off, len); + let rhs_bits = window_bits(&rhs, rhs_off, len); + for op in ops { + let expected: Vec = lhs_bits + .iter() + .zip(&rhs_bits) + .map(|(&a, &b)| match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }) + .collect(); + for out_off in [0usize, 8, 64, 128] { + let mut out = pattern_mask(out_off + len + 64, 2, 9); + let before = out.clone(); + bitmask_binop_simd_into::( + &mut out, + out_off, + (&lhs, lhs_off, len), + (&rhs, rhs_off, len), + op, + ); + let case = format!( + "{op:?} lhs {lhs_off} rhs {rhs_off} len {len} out {out_off}" + ); + let end = out_off + len; + assert_eq!(window_bits(&out, out_off, len), expected, "{case}"); + assert_eq!( + window_bits(&out, 0, out_off), + window_bits(&before, 0, out_off), + "{case}" + ); + assert_eq!( + window_bits(&out, end, out.len() - end), + window_bits(&before, end, before.len() - end), + "{case}" + ); + } + } + } + } + } + + #[test] + fn test_bitmask_unop_simd_into_windows() { + for offset in WINDOW_OFFSETS { + for len in sweep_lens() { + let src = pattern_mask(offset + len + 67, 3, 7); + let expected: Vec = window_bits(&src, offset, len) + .into_iter() + .map(|bit| !bit) + .collect(); + for out_off in [0usize, 8, 64, 128] { + let mut out = pattern_mask(out_off + len + 64, 2, 9); + let before = out.clone(); + bitmask_unop_simd_into::( + &mut out, + out_off, + (&src, offset, len), + UnaryOperator::Not, + ); + let case = format!("offset {offset} len {len} out {out_off}"); + let end = out_off + len; + assert_eq!(window_bits(&out, out_off, len), expected, "{case}"); + assert_eq!( + window_bits(&out, 0, out_off), + window_bits(&before, 0, out_off), + "{case}" + ); + assert_eq!( + window_bits(&out, end, out.len() - end), + window_bits(&before, end, before.len() - end), + "{case}" + ); + } + } + } + } + + #[test] + fn test_in_mask_simd_windows() { + for (k, lhs_off) in WINDOW_OFFSETS.into_iter().enumerate() { + let rhs_off = WINDOW_OFFSETS[(k + 5) % WINDOW_OFFSETS.len()]; + for len in sweep_lens() { + let lhs = pattern_mask(lhs_off + len + 67, 3, 7); + let n = rhs_off + len + 67; + // Right-hand sides holding both values, only true and + // only false within the window, with the opposite value + // outside the window for the uniform ones. + let mixed = pattern_mask(n, 5, 11); + let mut only_true = Bitmask::new_set_all(n, false); + let mut only_false = Bitmask::new_set_all(n, true); + for i in rhs_off..rhs_off + len { + only_true.set(i, true); + only_false.set(i, false); + } + for rhs in [&mixed, &only_true, &only_false] { + let rhs_bits = window_bits(rhs, rhs_off, len); + let has_true = rhs_bits.contains(&true); + let has_false = rhs_bits.contains(&false); + let expected: Vec = window_bits(&lhs, lhs_off, len) + .into_iter() + .map(|bit| if bit { has_true } else { has_false }) + .collect(); + let case = format!("lhs {lhs_off} rhs {rhs_off} len {len}"); + + let out = in_mask_simd::( + (&lhs, lhs_off, len), + (rhs, rhs_off, len), + ); + assert_eq!(window_bits(&out, 0, len), expected, "{case}"); + + let out = not_in_mask_simd::( + (&lhs, lhs_off, len), + (rhs, rhs_off, len), + ); + let expected_not: Vec = + expected.iter().map(|&bit| !bit).collect(); + assert_eq!(window_bits(&out, 0, len), expected_not, "{case}"); + } + } + } + } + + #[test] + fn test_all_eq_mask_simd_windows() { + for (lhs_off, rhs_off) in [(0usize, 0usize), (0, 64), (64, 0), (64, 64)] { + for len in sweep_lens() { + let n = 64 + len + 67; + let lhs = pattern_mask(n, 3, 7); + // Equal to `lhs` within the window and inverted outside it. + let mut rhs = Bitmask::new_set_all(n, false); + for i in 0..n { + let inside = i >= rhs_off && i < rhs_off + len; + let bit = if inside { + lhs.get(lhs_off + i - rhs_off) + } else { + !lhs.get(i) + }; + rhs.set(i, bit); + } + let case = format!("lhs {lhs_off} rhs {rhs_off} len {len}"); + + let expected = + window_bits(&lhs, lhs_off, len) == window_bits(&rhs, rhs_off, len); + assert_eq!( + all_eq_mask_simd::( + (&lhs, lhs_off, len), + (&rhs, rhs_off, len) + ), + expected, + "{case}" + ); + + let last = rhs_off + len - 1; + rhs.set(last, !rhs.get(last)); + let expected = + window_bits(&lhs, lhs_off, len) == window_bits(&rhs, rhs_off, len); + assert_eq!( + all_eq_mask_simd::( + (&lhs, lhs_off, len), + (&rhs, rhs_off, len) + ), + expected, + "{case}, last window bit flipped" + ); + } + } + } } }; } diff --git a/src/kernels/bitmask/std.rs b/src/kernels/bitmask/std.rs index 7907450..9514867 100644 --- a/src/kernels/bitmask/std.rs +++ b/src/kernels/bitmask/std.rs @@ -83,7 +83,8 @@ pub fn bitmask_binop_std(lhs: BitmaskVT<'_>, rhs: BitmaskVT<'_>, op: LogicalOper /// the words never straddle a window edge, so adjacent windows of a shared /// output buffer stay independent. /// -/// All offsets are byte-aligned i.e. a multiple of 8 bits. +/// Input offsets may be any bit position. `out_off` is byte-aligned i.e. a +/// multiple of 8 bits. #[inline(always)] pub fn bitmask_binop_std_into( out: &mut Bitmask, @@ -102,32 +103,80 @@ pub fn bitmask_binop_std_into( 0, "bitmask_binop_std_into: out_off must be byte-aligned" ); - debug_assert_eq!( - lhs_off % 8, - 0, - "bitmask_binop_std_into: lhs offset must be byte-aligned" - ); - debug_assert_eq!( - rhs_off % 8, - 0, - "bitmask_binop_std_into: rhs offset must be byte-aligned" - ); let full_words = len / 64; let tail_bits = len % 64; { - let lhs_bytes = bitmask_window_bytes(lhs_mask, lhs_off, len); - let rhs_bytes = bitmask_window_bytes(rhs_mask, rhs_off, len); + // Each input is read from the byte holding its first window bit. For + // windows starting inside a byte, each word is joined with the low bits + // of the following word through a shift and an OR. Output words are + // copied in as little-endian bytes, which needs no word alignment of + // the output window. + let lhs_bytes: &[u8] = &lhs_mask.bits[lhs_off / 8..]; + let rhs_bytes: &[u8] = &rhs_mask.bits[rhs_off / 8..]; + let lhs_shift = lhs_off % 8; + let rhs_shift = rhs_off % 8; let out_bytes = bitmask_window_bytes_mut(out, out_off, len); - unsafe { - let lp = lhs_bytes.as_ptr().cast::(); - let rp = rhs_bytes.as_ptr().cast::(); - let dp = out_bytes.as_mut_ptr().cast::(); - for k in 0..full_words { - *dp.add(k) = match op { - LogicalOperator::And => *lp.add(k) & *rp.add(k), - LogicalOperator::Or => *lp.add(k) | *rp.add(k), - LogicalOperator::Xor => *lp.add(k) ^ *rp.add(k), + let apply = |a: u64, b: u64| match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }; + + if lhs_shift == 0 && rhs_shift == 0 { + for ((chunk, a), b) in out_bytes[..full_words * 8] + .chunks_exact_mut(8) + .zip(lhs_bytes[..full_words * 8].chunks_exact(8)) + .zip(rhs_bytes[..full_words * 8].chunks_exact(8)) + { + let a = u64::from_le_bytes(a.try_into().unwrap()); + let b = u64::from_le_bytes(b.try_into().unwrap()); + chunk.copy_from_slice(&apply(a, b).to_le_bytes()); + } + } else { + // Where the following word lies within both buffers, input words are + // built from two whole-word reads. `(hi << 1) << (63 - shift)` moves + // the low `shift` bits of the following word to the top, and is zero + // for an input starting on a byte boundary. + let paired = full_words + .min((lhs_bytes.len() / 8).saturating_sub(1)) + .min((rhs_bytes.len() / 8).saturating_sub(1)); + if paired > 0 { + let join = |(lo, hi): (&[u8], &[u8]), shift: usize| { + let lo = u64::from_le_bytes(lo.try_into().unwrap()); + let hi = u64::from_le_bytes(hi.try_into().unwrap()); + (lo >> shift) | ((hi << 1) << (63 - shift)) }; + let lhs_words = lhs_bytes[..paired * 8] + .chunks_exact(8) + .zip(lhs_bytes[8..(paired + 1) * 8].chunks_exact(8)); + let rhs_words = rhs_bytes[..paired * 8] + .chunks_exact(8) + .zip(rhs_bytes[8..(paired + 1) * 8].chunks_exact(8)); + for ((chunk, a), b) in out_bytes[..paired * 8] + .chunks_exact_mut(8) + .zip(lhs_words) + .zip(rhs_words) + { + let word = apply(join(a, lhs_shift), join(b, rhs_shift)); + chunk.copy_from_slice(&word.to_le_bytes()); + } + } + // The remaining words are joined with the low bits of the following + // byte, which lies within the window. + let word_at = |bytes: &[u8], shift: usize, k: usize| { + let lo = u64::from_le_bytes(bytes[k * 8..k * 8 + 8].try_into().unwrap()); + if shift == 0 { + lo + } else { + (lo >> shift) | ((bytes[k * 8 + 8] as u64) << (64 - shift)) + } + }; + for k in paired..full_words { + let word = apply( + word_at(lhs_bytes, lhs_shift, k), + word_at(rhs_bytes, rhs_shift, k), + ); + out_bytes[k * 8..k * 8 + 8].copy_from_slice(&word.to_le_bytes()); } } } @@ -162,7 +211,8 @@ pub fn bitmask_unop_std(src: BitmaskVT<'_>, op: UnaryOperator) -> Bitmask { /// Full 64-bit words run word-wise. The final `len % 64` bits run one bit at a /// time, so the write never touches the trailing bits past `len`. /// -/// All offsets are byte-aligned i.e. a multiple of 8 bits. +/// The input offset may be any bit position. `out_off` is byte-aligned i.e. a +/// multiple of 8 bits. #[inline(always)] pub fn bitmask_unop_std_into( out: &mut Bitmask, @@ -179,24 +229,51 @@ pub fn bitmask_unop_std_into( 0, "bitmask_unop_std_into: out_off must be byte-aligned" ); - debug_assert_eq!( - offset % 8, - 0, - "bitmask_unop_std_into: src offset must be byte-aligned" - ); let full_words = len / 64; let tail_bits = len % 64; { - let src_bytes = bitmask_window_bytes(mask, offset, len); + // The input is read from the byte holding its first window bit. For + // windows starting inside a byte, each word is joined with the low bits + // of the following word through a shift and an OR. Output words are + // copied in as little-endian bytes, which needs no word alignment of + // the output window. + let src_bytes: &[u8] = &mask.bits[offset / 8..]; + let shift = offset % 8; let out_bytes = bitmask_window_bytes_mut(out, out_off, len); - unsafe { - let sp = src_bytes.as_ptr().cast::(); - let dp = out_bytes.as_mut_ptr().cast::(); - for k in 0..full_words { - *dp.add(k) = match op { - UnaryOperator::Not => !*sp.add(k), - _ => unreachable!(), // Positive Negative invalid for bools - }; + let apply = |src_word: u64| match op { + UnaryOperator::Not => !src_word, + _ => unreachable!(), // Positive Negative invalid for bools + }; + + if shift == 0 { + for (chunk, src) in out_bytes[..full_words * 8] + .chunks_exact_mut(8) + .zip(src_bytes[..full_words * 8].chunks_exact(8)) + { + let src_word = u64::from_le_bytes(src.try_into().unwrap()); + chunk.copy_from_slice(&apply(src_word).to_le_bytes()); + } + } else { + // Where the following word lies within the buffer, words are built + // from two whole-word reads. The remaining words are joined with the + // low bits of the following byte, which lies within the window. + let paired = full_words.min((src_bytes.len() / 8).saturating_sub(1)); + if paired > 0 { + for ((chunk, lo), hi) in out_bytes[..paired * 8] + .chunks_exact_mut(8) + .zip(src_bytes[..paired * 8].chunks_exact(8)) + .zip(src_bytes[8..(paired + 1) * 8].chunks_exact(8)) + { + let lo = u64::from_le_bytes(lo.try_into().unwrap()); + let hi = u64::from_le_bytes(hi.try_into().unwrap()); + let src_word = (lo >> shift) | (hi << (64 - shift)); + chunk.copy_from_slice(&apply(src_word).to_le_bytes()); + } + } + for k in paired..full_words { + let lo = u64::from_le_bytes(src_bytes[k * 8..k * 8 + 8].try_into().unwrap()); + let src_word = (lo >> shift) | ((src_bytes[k * 8 + 8] as u64) << (64 - shift)); + out_bytes[k * 8..k * 8 + 8].copy_from_slice(&apply(src_word).to_le_bytes()); } } } @@ -351,37 +428,27 @@ pub fn all_eq_mask(a: BitmaskVT<'_>, b: BitmaskVT<'_>) -> bool { ); } - let a_bytes = bitmask_window_bytes(am, ao, len); - let b_bytes = bitmask_window_bytes(bm, bo, len); - let total_bytes = a_bytes.len(); - let full_words = total_bytes / 8; - let tail_bytes = total_bytes % 8; - let full_logical_bytes = len / 8; - let last_bits = len & 7; - - unsafe { - let ap = a_bytes.as_ptr().cast::(); - let bp = b_bytes.as_ptr().cast::(); - for k in 0..full_words { - if *ap.add(k) != *bp.add(k) { - return false; - } + // Both windows start on a word boundary, and whole words are read from the + // window bytes. The final partial word is read through `load_word` and + // masked to the window, which keeps differences past the window from + // causing false negatives. + let full_words = len / 64; + let a_bytes: &[u8] = &am.bits[ao / 8..]; + let b_bytes: &[u8] = &bm.bits[bo / 8..]; + for (a, b) in a_bytes[..full_words * 8] + .chunks_exact(8) + .zip(b_bytes[..full_words * 8].chunks_exact(8)) + { + if a != b { + return false; } } - let base = full_words * 8; - for k in 0..tail_bytes { - let byte_index = base + k; - let av = a_bytes[byte_index]; - let bv = b_bytes[byte_index]; - if byte_index < full_logical_bytes { - if av != bv { - return false; - } - } else if last_bits != 0 { - let m = (1u8 << last_bits) - 1; - if (av & m) != (bv & m) { - return false; - } + let tail_bits = len % 64; + if tail_bits != 0 { + let base = full_words * 64; + let in_window = u64::MAX >> (64 - tail_bits); + if (load_word(am, ao + base) ^ load_word(bm, bo + base)) & in_window != 0 { + return false; } } true @@ -410,31 +477,44 @@ pub fn popcount_mask(m: BitmaskVT<'_>) -> usize { if len == 0 { return 0; } - let bytes = bitmask_window_bytes(mask, offset, len); - let total_bytes = bytes.len(); - let full_words = total_bytes / 8; - let tail_bytes = total_bytes % 8; - let full_logical_bytes = len / 8; - let last_bits = len & 7; + // Whole words are read from the byte holding the first window bit. For + // windows starting inside a byte, each word is joined with the low bits of + // the following word through a shift and an OR. The final partial word is + // read through `load_word` and masked to the window, which keeps bits + // outside the window out of the count. + let bytes: &[u8] = &mask.bits[offset / 8..]; + let shift = offset % 8; + let full_words = len / 64; let mut acc = 0usize; - - unsafe { - let words_ptr = bytes.as_ptr().cast::(); - for k in 0..full_words { - acc += (*words_ptr.add(k)).count_ones() as usize; + if shift == 0 { + for chunk in bytes[..full_words * 8].chunks_exact(8) { + acc += u64::from_le_bytes(chunk.try_into().unwrap()).count_ones() as usize; } - } - let base = full_words * 8; - for k in 0..tail_bytes { - let byte_index = base + k; - let b = bytes[byte_index]; - if byte_index < full_logical_bytes { - acc += b.count_ones() as usize; - } else if last_bits != 0 { - let m = (1u8 << last_bits) - 1; - acc += (b & m).count_ones() as usize; + } else { + // Where the following word lies within the buffer, words are built from + // two whole-word reads. The remaining words are joined with the low bits + // of the following byte, which lies within the window. + let paired = full_words.min((bytes.len() / 8).saturating_sub(1)); + if paired > 0 { + let lo_words = bytes[..paired * 8].chunks_exact(8); + let hi_words = bytes[8..(paired + 1) * 8].chunks_exact(8); + for (lo, hi) in lo_words.zip(hi_words) { + let lo = u64::from_le_bytes(lo.try_into().unwrap()); + let hi = u64::from_le_bytes(hi.try_into().unwrap()); + acc += ((lo >> shift) | (hi << (64 - shift))).count_ones() as usize; + } + } + for k in paired..full_words { + let lo = u64::from_le_bytes(bytes[k * 8..k * 8 + 8].try_into().unwrap()); + let word = (lo >> shift) | ((bytes[k * 8 + 8] as u64) << (64 - shift)); + acc += word.count_ones() as usize; } } + let tail_bits = len % 64; + if tail_bits != 0 { + let word = load_word(mask, offset + full_words * 64) & (u64::MAX >> (64 - tail_bits)); + acc += word.count_ones() as usize; + } acc } @@ -554,6 +634,7 @@ mod tests { use super::*; use crate::Bitmask; use crate::kernels::bitmask::clear_trailing_bits; + use crate::kernels::bitmask::tests::{WINDOW_LENS, WINDOW_OFFSETS, pattern_mask, window_bits}; // Helper: Create a Bitmask from a bool slice. fn bm(bits: &[bool]) -> Bitmask { @@ -744,4 +825,150 @@ mod tests { let window_mut = bitmask_window_bytes_mut(&mut b, 2, 4); assert_eq!(window_mut.len(), 1); } + + #[test] + fn test_popcount_mask_windows() { + for offset in WINDOW_OFFSETS { + for len in WINDOW_LENS { + let mask = pattern_mask(offset + len + 67, 3, 7); + let expected = window_bits(&mask, offset, len) + .into_iter() + .filter(|&bit| bit) + .count(); + assert_eq!( + popcount_mask((&mask, offset, len)), + expected, + "offset {offset} len {len}" + ); + } + } + } + + #[test] + fn test_bitmask_binop_std_into_windows() { + let ops = [ + LogicalOperator::And, + LogicalOperator::Or, + LogicalOperator::Xor, + ]; + for (k, lhs_off) in WINDOW_OFFSETS.into_iter().enumerate() { + let rhs_off = WINDOW_OFFSETS[(k + 5) % WINDOW_OFFSETS.len()]; + for len in WINDOW_LENS { + let lhs = pattern_mask(lhs_off + len + 67, 3, 7); + let rhs = pattern_mask(rhs_off + len + 67, 5, 11); + let lhs_bits = window_bits(&lhs, lhs_off, len); + let rhs_bits = window_bits(&rhs, rhs_off, len); + for op in ops { + let expected: Vec = lhs_bits + .iter() + .zip(&rhs_bits) + .map(|(&a, &b)| match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }) + .collect(); + for out_off in [0usize, 8, 64, 128] { + let mut out = pattern_mask(out_off + len + 64, 2, 9); + let before = out.clone(); + bitmask_binop_std_into( + &mut out, + out_off, + (&lhs, lhs_off, len), + (&rhs, rhs_off, len), + op, + ); + let case = + format!("{op:?} lhs {lhs_off} rhs {rhs_off} len {len} out {out_off}"); + let end = out_off + len; + assert_eq!(window_bits(&out, out_off, len), expected, "{case}"); + assert_eq!( + window_bits(&out, 0, out_off), + window_bits(&before, 0, out_off), + "{case}" + ); + assert_eq!( + window_bits(&out, end, out.len() - end), + window_bits(&before, end, before.len() - end), + "{case}" + ); + } + } + } + } + } + + #[test] + fn test_bitmask_unop_std_into_windows() { + for offset in WINDOW_OFFSETS { + for len in WINDOW_LENS { + let src = pattern_mask(offset + len + 67, 3, 7); + let expected: Vec = window_bits(&src, offset, len) + .into_iter() + .map(|bit| !bit) + .collect(); + for out_off in [0usize, 8, 64, 128] { + let mut out = pattern_mask(out_off + len + 64, 2, 9); + let before = out.clone(); + bitmask_unop_std_into( + &mut out, + out_off, + (&src, offset, len), + UnaryOperator::Not, + ); + let case = format!("offset {offset} len {len} out {out_off}"); + let end = out_off + len; + assert_eq!(window_bits(&out, out_off, len), expected, "{case}"); + assert_eq!( + window_bits(&out, 0, out_off), + window_bits(&before, 0, out_off), + "{case}" + ); + assert_eq!( + window_bits(&out, end, out.len() - end), + window_bits(&before, end, before.len() - end), + "{case}" + ); + } + } + } + } + + #[test] + fn test_all_eq_mask_windows() { + for (lhs_off, rhs_off) in [(0usize, 0usize), (0, 64), (64, 0), (64, 64)] { + for len in WINDOW_LENS { + let n = 64 + len + 67; + let lhs = pattern_mask(n, 3, 7); + // Equal to `lhs` within the window and inverted outside it. + let mut rhs = Bitmask::new_set_all(n, false); + for i in 0..n { + let inside = i >= rhs_off && i < rhs_off + len; + let bit = if inside { + lhs.get(lhs_off + i - rhs_off) + } else { + !lhs.get(i) + }; + rhs.set(i, bit); + } + let case = format!("lhs {lhs_off} rhs {rhs_off} len {len}"); + + let expected = window_bits(&lhs, lhs_off, len) == window_bits(&rhs, rhs_off, len); + assert_eq!( + all_eq_mask((&lhs, lhs_off, len), (&rhs, rhs_off, len)), + expected, + "{case}" + ); + + let last = rhs_off + len - 1; + rhs.set(last, !rhs.get(last)); + let expected = window_bits(&lhs, lhs_off, len) == window_bits(&rhs, rhs_off, len); + assert_eq!( + all_eq_mask((&lhs, lhs_off, len), (&rhs, rhs_off, len)), + expected, + "{case}, last window bit flipped" + ); + } + } + } }