diff --git a/src/kernels/bitmask/mod.rs b/src/kernels/bitmask/mod.rs index f7c1383..ac3a6ad 100644 --- a/src/kernels/bitmask/mod.rs +++ b/src/kernels/bitmask/mod.rs @@ -224,6 +224,7 @@ pub fn merge_bitmasks_to_new( mod tests { use super::*; use crate::Bitmask; + use crate::enums::operators::{LogicalOperator, UnaryOperator}; #[test] fn test_words_for() { @@ -406,4 +407,164 @@ mod tests { assert!(out.get(i), "bit {i} past the window was disturbed"); } } + + /// Window offsets for the offset sweeps. They cover bit positions within a + /// byte and within a word, whole-word offsets, and offsets that start + /// several words into the mask. + pub(super) const SWEEP_OFFSETS: [usize; 13] = [0, 1, 2, 3, 7, 8, 13, 63, 64, 65, 67, 512, 515]; + + /// Window lengths for the offset sweeps at a SIMD width of `lanes` words. + /// The two longest lengths span one and three whole vectors plus a partial + /// word. + pub(super) fn sweep_lengths(lanes: usize) -> [usize; 10] { + // Bits in one vector of `lanes` words. + let v = 64 * lanes; + [1, 57, 61, 63, 64, 65, 130, 200, v + 37, 3 * v + 37] + } + + /// Sweep offsets for windows of `len` bits, each paired with the length of + /// the mask that contains the window. Every offset appears twice, once in a + /// mask that ends with the window and once in a mask that extends one whole + /// vector and a partial word past it. + pub(super) fn sweep_windows(len: usize, lanes: usize) -> impl Iterator { + SWEEP_OFFSETS + .into_iter() + .flat_map(move |off| [(off, off + len), (off, off + len + 64 * lanes + 13)]) + } + + /// Mask of `n` bits with bit `i` set when `i % a == 0 || i % b == 1`. The + /// sweeps use values of `a` and `b` for which the pattern does not repeat + /// every 8 bits. + pub(super) fn sweep_mask(n: usize, a: usize, b: usize) -> Bitmask { + let mut m = Bitmask::new_set_all(n, false); + for i in 0..n { + if i % a == 0 || i % b == 1 { + m.set(i, true); + } + } + m + } + + /// Checks a binary `_into` kernel against a bit-by-bit `Bitmask::get` + /// reference, over every pair of sweep offsets, the sweep lengths and + /// output offsets 0, 8 and 64. Output bits outside the window keep their + /// prior values. + pub(super) fn sweep_binop_into( + lanes: usize, + kernel: impl Fn(&mut Bitmask, usize, BitmaskVT<'_>, BitmaskVT<'_>, LogicalOperator), + ) { + let ops = [ + LogicalOperator::And, + LogicalOperator::Or, + LogicalOperator::Xor, + ]; + for len in sweep_lengths(lanes) { + for out_off in [0usize, 8, 64] { + let out_n = out_off + len + 70; + let before = sweep_mask(out_n, 4, 9); + for (lhs_off, lhs_n) in sweep_windows(len, lanes) { + let lhs = sweep_mask(lhs_n, 3, 7); + for (rhs_off, rhs_n) in sweep_windows(len, lanes) { + let rhs = sweep_mask(rhs_n, 5, 11); + for op in ops { + let mut out = before.clone(); + let lhs_window = (&lhs, lhs_off, len); + let rhs_window = (&rhs, rhs_off, len); + kernel(&mut out, out_off, lhs_window, rhs_window, op); + for i in 0..out_n { + let expected = if i < out_off || i >= out_off + len { + before.get(i) + } else { + let a = lhs.get(lhs_off + i - out_off); + let b = rhs.get(rhs_off + i - out_off); + match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + } + }; + assert_eq!( + out.get(i), + expected, + "{op:?}: lhs ({lhs_off}, {len}) of {lhs_n} bits, \ + rhs ({rhs_off}, {len}) of {rhs_n} bits, \ + out_off {out_off}, out bit {i}" + ); + } + } + } + } + } + } + } + + /// Checks a unary `_into` kernel against a bit-by-bit `Bitmask::get` + /// reference, over the sweep offsets, the sweep lengths and output offsets + /// 0, 8 and 64. Output bits outside the window keep their prior values. + pub(super) fn sweep_unop_into( + lanes: usize, + kernel: impl Fn(&mut Bitmask, usize, BitmaskVT<'_>, UnaryOperator), + ) { + for len in sweep_lengths(lanes) { + for out_off in [0usize, 8, 64] { + let out_n = out_off + len + 70; + let before = sweep_mask(out_n, 4, 9); + for (off, n) in sweep_windows(len, lanes) { + let src = sweep_mask(n, 3, 7); + let mut out = before.clone(); + kernel(&mut out, out_off, (&src, off, len), UnaryOperator::Not); + for i in 0..out_n { + let expected = if i < out_off || i >= out_off + len { + before.get(i) + } else { + !src.get(off + i - out_off) + }; + assert_eq!( + out.get(i), + expected, + "Not: src ({off}, {len}) of {n} bits, out_off {out_off}, out bit {i}" + ); + } + } + } + } + } + + /// Checks a popcount kernel against a bit-by-bit `Bitmask::get` count over + /// the sweep offsets and lengths. Rows past the end of the window are + /// excluded from the count. + pub(super) fn sweep_popcount(lanes: usize, kernel: impl Fn(BitmaskVT<'_>) -> usize) { + for len in sweep_lengths(lanes) { + for (off, n) in sweep_windows(len, lanes) { + let m = sweep_mask(n, 3, 7); + let expected = (off..off + len).filter(|&i| m.get(i)).count(); + assert_eq!( + kernel((&m, off, len)), + expected, + "window ({off}, {len}) of {n} bits" + ); + } + } + } + + /// The dispatched binary `_into` kernel reads windows at any input offset. + #[test] + fn binop_into_window_offsets() { + use crate::kernels::bitmask::dispatch::{W8, bitmask_binop_into}; + sweep_binop_into(W8, bitmask_binop_into); + } + + /// The dispatched unary `_into` kernel reads windows at any input offset. + #[test] + fn unop_into_window_offsets() { + use crate::kernels::bitmask::dispatch::{W8, bitmask_unop_into}; + sweep_unop_into(W8, bitmask_unop_into); + } + + /// The dispatched popcount counts the rows of windows at any offset. + #[test] + fn popcount_window_offsets() { + use crate::kernels::bitmask::dispatch::{W8, popcount_mask}; + sweep_popcount(W8, popcount_mask); + } } diff --git a/src/kernels/bitmask/simd.rs b/src/kernels/bitmask/simd.rs index bbfc6dc..e2ff7c4 100644 --- a/src/kernels/bitmask/simd.rs +++ b/src/kernels/bitmask/simd.rs @@ -56,13 +56,34 @@ //! if !eq_mask.all() { return false; } //! ``` -use core::simd::Simd; +use core::simd::{Simd, Swizzle}; use crate::kernels::arithmetic::simd::{W8, W16, W32, W64}; use crate::{Bitmask, BitmaskVT}; use crate::enums::operators::{LogicalOperator, UnaryOperator}; -use crate::kernels::bitmask::{bitmask_window_bytes, bitmask_window_bytes_mut, load_word}; +use crate::kernels::bitmask::{ + bitmask_window_bytes, bitmask_window_bytes_mut, load_word, mask_bits_as_words, +}; + +/// Swizzle index `[1, 2, ..., N]` over two concatenated vectors. +/// +/// Applied to consecutive vectors of mask words `cur` and `nxt`, the result +/// contains the word that follows each lane of `cur`. The swizzle is a register +/// permute, for example `valignq` on AVX-512, rather than a memory load. +struct NextWords; + +impl Swizzle for NextWords { + const INDEX: [usize; N] = { + let mut index = [0; N]; + let mut i = 0; + while i < N { + index[i] = i + 1; + i += 1; + } + index + }; +} /// Primitive bit ops @@ -141,11 +162,19 @@ pub fn bitmask_unop_simd(src: BitmaskVT<'_>, op: UnaryOperat /// Full 64-bit words lying entirely within `[0, len)` run word-wise, `LANES` /// words per SIMD step and then any remaining whole words scalar. The final /// `len % 64` bits run one bit at a time, so the write touches only valid -/// positions and never the trailing bits past `len`. With byte-aligned offsets -/// the words never straddle a window edge, so adjacent windows of a shared -/// output buffer stay independent. +/// positions and never the trailing bits past `len`. With a byte-aligned +/// `out_off` the words never straddle a window edge. Adjacent windows of a +/// shared output buffer therefore stay independent. +/// +/// When every input offset is a multiple of 64, whole words are loaded from +/// the window. Otherwise, for an input window that starts at bit `s` of mask +/// word `w0`, window word `k` is formed from mask words `w0 + k` and +/// `w0 + k + 1` as `(cur >> s) | (next << (64 - s))`. Both paths load whole +/// words only. Whole result words are written with aligned stores when +/// `out_off` is a multiple of 64, and as little-endian bytes otherwise. /// -/// 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,22 +193,17 @@ 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; - { + if lhs_off % 64 == 0 && rhs_off % 64 == 0 { let lhs_bytes = bitmask_window_bytes(lhs_mask, lhs_off, len); let rhs_bytes = bitmask_window_bytes(rhs_mask, rhs_off, len); let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + // SAFETY: Both input offsets are multiples of 64 and mask buffers start on a + // 64-byte boundary, which makes `lp` and `rp` aligned `u64` pointers. Words + // `0..full_words` lie inside each input window and the output window. Whole + // words are written through `dp` as `u64` only when `out_off` is a multiple + // of 64, and as 8 bytes otherwise, which needs no alignment. unsafe { let lp = lhs_bytes.as_ptr().cast::(); let rp = rhs_bytes.as_ptr().cast::(); @@ -195,20 +219,143 @@ pub fn bitmask_binop_simd_into( LogicalOperator::Or => a | b, LogicalOperator::Xor => a ^ b, }; - std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); + if out_off % 64 == 0 { + std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); + } else { + for (j, w) in r.as_array().iter().enumerate() { + dp.add(i + j).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } i += LANES; } // 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 { + let w = match op { LogicalOperator::And => a & b, LogicalOperator::Or => a | b, LogicalOperator::Xor => a ^ b, }; + if out_off % 64 == 0 { + *dp.add(k) = w; + } else { + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } } } + } else { + let (lw0, ls) = (lhs_off / 64, (lhs_off % 64) as u64); + let (rw0, rs) = (rhs_off / 64, (rhs_off % 64) as u64); + // SAFETY: Mask buffers start on a 64-byte boundary, which aligns the buffer + // start for `u64`, and `bits.len() / 8` whole words lie inside the buffer. + let lwords = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&lhs_mask.bits), lhs_mask.bits.len() / 8) + }; + // SAFETY: The same holds for the rhs mask buffer. + let rwords = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&rhs_mask.bits), rhs_mask.bits.len() / 8) + }; + let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + let dp = out_bytes.as_mut_ptr().cast::(); + // SIMD steps run while the following vector of each input lies inside + // its buffer. + let steps = (full_words / LANES) + .min(lwords.len().saturating_sub(lw0 + LANES) / LANES) + .min(rwords.len().saturating_sub(rw0 + LANES) / LANES); + let mut i = 0; + if steps > 0 { + let (lshift, lback) = (Simd::splat(ls), Simd::splat(64 - ls)); + let (rshift, rback) = (Simd::splat(rs), Simd::splat(64 - rs)); + let mut lcur = Simd::::from_slice(&lwords[lw0..lw0 + LANES]); + let mut rcur = Simd::::from_slice(&rwords[rw0..rw0 + LANES]); + for _ in 0..steps { + let lj = lw0 + i + LANES; + let rj = rw0 + i + LANES; + let lnxt = Simd::::from_slice(&lwords[lj..lj + LANES]); + let rnxt = Simd::::from_slice(&rwords[rj..rj + LANES]); + // Inputs at offsets that are multiples of 64 are used without + // a shift. + let a = if ls == 0 { + lcur + } else { + (lcur >> lshift) | (NextWords::concat_swizzle(lcur, lnxt) << lback) + }; + let b = if rs == 0 { + rcur + } else { + (rcur >> rshift) | (NextWords::concat_swizzle(rcur, rnxt) << rback) + }; + let r = match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }; + // SAFETY: `i + LANES <= full_words`, which keeps words `i..i + LANES` inside + // the output window. They are written as aligned `u64` words only when + // `out_off` is a multiple of 64, and as 8 bytes each otherwise. + unsafe { + if out_off % 64 == 0 { + std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); + } else { + for (j, w) in r.as_array().iter().enumerate() { + dp.add(i + j).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + } + lcur = lnxt; + rcur = rnxt; + i += LANES; + } + } + // Remaining whole words one at a time. Mask words that are not whole + // inside the buffer are read with `load_word`, which reads bytes past + // the end of the buffer as zero. + let lword = |j: usize| { + lwords + .get(j) + .copied() + .unwrap_or_else(|| load_word(lhs_mask, j * 64)) + }; + let rword = |j: usize| { + rwords + .get(j) + .copied() + .unwrap_or_else(|| load_word(rhs_mask, j * 64)) + }; + let mut lcur = lword(lw0 + i); + let mut rcur = rword(rw0 + i); + for k in i..full_words { + let lnext = lword(lw0 + k + 1); + let rnext = rword(rw0 + k + 1); + let a = if ls == 0 { + lcur + } else { + (lcur >> ls) | (lnext << (64 - ls)) + }; + let b = if rs == 0 { + rcur + } else { + (rcur >> rs) | (rnext << (64 - rs)) + }; + let w = match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }; + // SAFETY: `k < full_words`, which keeps word `k` inside the output window. It + // is written as an aligned `u64` only when `out_off` is a multiple of 64, + // and as 8 bytes otherwise. + unsafe { + if out_off % 64 == 0 { + *dp.add(k) = w; + } else { + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + lcur = lnext; + rcur = rnext; + } } // Final partial word written one bit at a time. let base = full_words * 64; @@ -233,7 +380,15 @@ 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. +/// When `offset` is a multiple of 64, whole words are loaded from the window. +/// Otherwise, for a window that starts at bit `s` of mask word `w0`, window +/// word `k` is formed from mask words `w0 + k` and `w0 + k + 1` as +/// `(cur >> s) | (next << (64 - s))`. Both paths load whole words only. Whole +/// result words are written with aligned stores when `out_off` is a multiple +/// of 64, and as little-endian bytes otherwise. +/// +/// 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,16 +405,16 @@ 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; - { + if offset % 64 == 0 { let src_bytes = bitmask_window_bytes(mask, offset, len); let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + // SAFETY: The input offset is a multiple of 64 and mask buffers start on a + // 64-byte boundary, which makes `sp` an aligned `u64` pointer. Words + // `0..full_words` lie inside the input window and the output window. Whole + // words are written through `dp` as `u64` only when `out_off` is a multiple + // of 64, and as 8 bytes otherwise, which needs no alignment. unsafe { let sp = src_bytes.as_ptr().cast::(); let dp = out_bytes.as_mut_ptr().cast::(); @@ -271,16 +426,94 @@ pub fn bitmask_unop_simd_into( UnaryOperator::Not => !a, _ => unreachable!(), }; - std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); + if out_off % 64 == 0 { + std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); + } else { + for (j, w) in r.as_array().iter().enumerate() { + dp.add(i + j).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } i += LANES; } // Scalar word tail when full_words is not a multiple of LANES. for k in i..full_words { - *dp.add(k) = match op { + let w = match op { UnaryOperator::Not => !*sp.add(k), _ => unreachable!(), }; + if out_off % 64 == 0 { + *dp.add(k) = w; + } else { + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + } + } else { + let (w0, s) = (offset / 64, (offset % 64) as u64); + // SAFETY: Mask buffers start on a 64-byte boundary, which aligns the buffer + // start for `u64`, and `bits.len() / 8` whole words lie inside the buffer. + let words = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&mask.bits), mask.bits.len() / 8) + }; + let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + let dp = out_bytes.as_mut_ptr().cast::(); + // SIMD steps run while the following vector lies inside the buffer. + let steps = (full_words / LANES).min(words.len().saturating_sub(w0 + LANES) / LANES); + let mut i = 0; + if steps > 0 { + let (shift, back) = (Simd::splat(s), Simd::splat(64 - s)); + let mut cur = Simd::::from_slice(&words[w0..w0 + LANES]); + for _ in 0..steps { + let at = w0 + i + LANES; + let nxt = Simd::::from_slice(&words[at..at + LANES]); + let a = (cur >> shift) | (NextWords::concat_swizzle(cur, nxt) << back); + let r = match op { + UnaryOperator::Not => !a, + _ => unreachable!(), + }; + // SAFETY: `i + LANES <= full_words`, which keeps words `i..i + LANES` inside + // the output window. They are written as aligned `u64` words only when + // `out_off` is a multiple of 64, and as 8 bytes each otherwise. + unsafe { + if out_off % 64 == 0 { + std::ptr::copy_nonoverlapping(r.as_array().as_ptr(), dp.add(i), LANES); + } else { + for (j, w) in r.as_array().iter().enumerate() { + dp.add(i + j).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + } + cur = nxt; + i += LANES; + } + } + // Remaining whole words one at a time. Mask words that are not whole + // inside the buffer are read with `load_word`, which reads bytes past + // the end of the buffer as zero. + let word = |j: usize| { + words + .get(j) + .copied() + .unwrap_or_else(|| load_word(mask, j * 64)) + }; + let mut cur = word(w0 + i); + for k in i..full_words { + let next = word(w0 + k + 1); + let w = match op { + UnaryOperator::Not => !((cur >> s) | (next << (64 - s))), + _ => unreachable!(), + }; + // SAFETY: `k < full_words`, which keeps word `k` inside the output window. It + // is written as an aligned `u64` only when `out_off` is a multiple of 64, + // and as 8 bytes otherwise. + unsafe { + if out_off % 64 == 0 { + *dp.add(k) = w; + } else { + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } } + cur = next; } } // Final partial word written one bit at a time. @@ -722,6 +955,13 @@ where /// /// Counts set bits in a bitmask window using SIMD popcount with horizontal reduction. /// +/// When `offset` is a multiple of 64, whole words are loaded from the window. +/// Otherwise, for a window that starts at bit `s` of mask word `w0`, window +/// word `k` is formed from mask words `w0 + k` and `w0 + k + 1` as +/// `(cur >> s) | (next << (64 - s))`, using whole-word loads only. On both +/// paths the final window word is masked to the window length, which excludes +/// rows past the end of the window. +/// /// # Type Parameters /// - `LANES`: Number of u64 lanes to process simultaneously /// @@ -739,14 +979,64 @@ where return 0; } + if offset % 64 != 0 { + use std::simd::prelude::SimdUint; + let (w0, s) = (offset / 64, (offset % 64) as u64); + // SAFETY: Mask buffers start on a 64-byte boundary, which aligns the buffer + // start for `u64`, and `bits.len() / 8` whole words lie inside the buffer. + let words = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&mask.bits), mask.bits.len() / 8) + }; + let full_words = len / 64; + // SIMD steps run while the following vector lies inside the buffer. + let steps = (full_words / LANES).min(words.len().saturating_sub(w0 + LANES) / LANES); + let mut acc = 0usize; + let mut i = 0; + if steps > 0 { + let (shift, back) = (Simd::splat(s), Simd::splat(64 - s)); + let mut cur = Simd::::from_slice(&words[w0..w0 + LANES]); + for _ in 0..steps { + let at = w0 + i + LANES; + let nxt = Simd::::from_slice(&words[at..at + LANES]); + let v = (cur >> shift) | (NextWords::concat_swizzle(cur, nxt) << back); + acc += v.count_ones().reduce_sum() as usize; + cur = nxt; + i += LANES; + } + } + // Remaining words one at a time. Mask words that are not whole inside + // the buffer are read with `load_word`, which reads bytes past the end + // of the buffer as zero. + let word = |j: usize| { + words + .get(j) + .copied() + .unwrap_or_else(|| load_word(mask, j * 64)) + }; + let mut cur = word(w0 + i); + for k in i..full_words { + let next = word(w0 + k + 1); + acc += ((cur >> s) | (next << (64 - s))).count_ones() as usize; + cur = next; + } + // Final window word, masked to the rows inside the window. + let tail_bits = len % 64; + if tail_bits != 0 { + let next = word(w0 + full_words + 1); + let w = (cur >> s) | (next << (64 - s)); + acc += (w & (u64::MAX >> (64 - tail_bits))).count_ones() as usize; + } + return acc; + } + 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 + // Whole u64 words inside the window via SIMD, scalar tail for + // partial u64 words, then byte-by-byte for the 0..8 bytes of the + // final partial word. 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_words = len / 64; + let tail_bytes = total_bytes - full_words * 8; let full_logical_bytes = len / 8; let last_bits = len & 7; let mut acc = 0usize; @@ -1059,6 +1349,7 @@ impl_simd_eq_mask!(simd_eq_mask_u64, u64, W64); #[cfg(test)] mod tests { + use crate::kernels::bitmask::tests::{sweep_binop_into, sweep_popcount, sweep_unop_into}; use crate::{Bitmask, BitmaskVT}; use super::*; @@ -1196,6 +1487,21 @@ mod tests { assert_eq!(pop, 4); } + #[test] + fn test_binop_simd_into_window_offsets() { + sweep_binop_into(LANES, bitmask_binop_simd_into::); + } + + #[test] + fn test_unop_simd_into_window_offsets() { + sweep_unop_into(LANES, bitmask_unop_simd_into::); + } + + #[test] + fn test_popcount_mask_simd_window_offsets() { + sweep_popcount(LANES, popcount_mask_simd::); + } + #[test] fn test_all_true_mask_simd_and_false() { let all_true = Bitmask::new_set_all(64 * LANES, true); diff --git a/src/kernels/bitmask/std.rs b/src/kernels/bitmask/std.rs index 7907450..ea89922 100644 --- a/src/kernels/bitmask/std.rs +++ b/src/kernels/bitmask/std.rs @@ -51,7 +51,9 @@ use crate::{Bitmask, BitmaskVT}; use crate::{ enums::operators::{LogicalOperator, UnaryOperator}, - kernels::bitmask::{bitmask_window_bytes, bitmask_window_bytes_mut, load_word}, + kernels::bitmask::{ + bitmask_window_bytes, bitmask_window_bytes_mut, load_word, mask_bits_as_words, + }, }; /// Performs bitwise binary operations (AND/OR/XOR) over two bitmask slices using word-level processing. @@ -79,11 +81,19 @@ pub fn bitmask_binop_std(lhs: BitmaskVT<'_>, rhs: BitmaskVT<'_>, op: LogicalOper /// /// Full 64-bit words lying entirely within `[0, len)` run word-wise. The final /// `len % 64` bits run one bit at a time, so the write touches only valid -/// positions and never the trailing bits past `len`. With byte-aligned offsets -/// the words never straddle a window edge, so adjacent windows of a shared -/// output buffer stay independent. +/// positions and never the trailing bits past `len`. With a byte-aligned +/// `out_off` the words never straddle a window edge. Adjacent windows of a +/// shared output buffer therefore stay independent. /// -/// All offsets are byte-aligned i.e. a multiple of 8 bits. +/// When every input offset is a multiple of 64, whole words are loaded from +/// the window. Otherwise, for an input window that starts at bit `s` of mask +/// word `w0`, window word `k` is formed from mask words `w0 + k` and +/// `w0 + k + 1` as `(cur >> s) | (next << (64 - s))`. Both paths load whole +/// words only. Whole result words are written with aligned stores when +/// `out_off` is a multiple of 64, and as little-endian bytes otherwise. +/// +/// 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,34 +112,113 @@ 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; - { + if lhs_off % 64 == 0 && rhs_off % 64 == 0 { let lhs_bytes = bitmask_window_bytes(lhs_mask, lhs_off, len); let rhs_bytes = bitmask_window_bytes(rhs_mask, rhs_off, len); let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + // SAFETY: Both input offsets are multiples of 64 and mask buffers start on a + // 64-byte boundary, which makes `lp` and `rp` aligned `u64` pointers. Words + // `0..full_words` lie inside each input window and the output window. Whole + // words are written through `dp` as `u64` only when `out_off` is a multiple + // of 64, and as 8 bytes otherwise, which needs no alignment. 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), - }; + if out_off % 64 == 0 { + 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), + }; + } + } else { + for k in 0..full_words { + let w = 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), + }; + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } } } + } else { + let (lw0, ls) = (lhs_off / 64, (lhs_off % 64) as u32); + let (rw0, rs) = (rhs_off / 64, (rhs_off % 64) as u32); + // SAFETY: Mask buffers start on a 64-byte boundary, which aligns the buffer + // start for `u64`, and `bits.len() / 8` whole words lie inside the buffer. + let lwords = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&lhs_mask.bits), lhs_mask.bits.len() / 8) + }; + // SAFETY: The same holds for the rhs mask buffer. + let rwords = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&rhs_mask.bits), rhs_mask.bits.len() / 8) + }; + let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + let dp = out_bytes.as_mut_ptr().cast::(); + // Window word formed from mask words `cur` and `next`. An input at an + // offset that is a multiple of 64 is used without a shift. + let join = |cur: u64, next: u64, s: u32| { + if s == 0 { + cur + } else { + (cur >> s) | (next << (64 - s)) + } + }; + let write = |k: usize, a: u64, b: u64| { + let w = match op { + LogicalOperator::And => a & b, + LogicalOperator::Or => a | b, + LogicalOperator::Xor => a ^ b, + }; + // SAFETY: `k < full_words`, which keeps word `k` inside the output window. It + // is written as an aligned `u64` only when `out_off` is a multiple of 64, + // and as 8 bytes otherwise. + unsafe { + if out_off % 64 == 0 { + *dp.add(k) = w; + } else { + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + }; + // Window words whose two mask words lie whole inside both buffers. + let whole = full_words + .min(lwords.len().saturating_sub(lw0 + 1)) + .min(rwords.len().saturating_sub(rw0 + 1)); + if whole > 0 { + let lspan = &lwords[lw0..=lw0 + whole]; + let rspan = &rwords[rw0..=rw0 + whole]; + for k in 0..whole { + let a = join(lspan[k], lspan[k + 1], ls); + let b = join(rspan[k], rspan[k + 1], rs); + write(k, a, b); + } + } + // Remaining words. Mask words that are not whole inside the buffer are + // read with `load_word`, which reads bytes past the end of the buffer + // as zero. + let lword = |j: usize| { + lwords + .get(j) + .copied() + .unwrap_or_else(|| load_word(lhs_mask, j * 64)) + }; + let rword = |j: usize| { + rwords + .get(j) + .copied() + .unwrap_or_else(|| load_word(rhs_mask, j * 64)) + }; + for k in whole..full_words { + let a = join(lword(lw0 + k), lword(lw0 + k + 1), ls); + let b = join(rword(rw0 + k), rword(rw0 + k + 1), rs); + write(k, a, b); + } } // Final partial word written one bit at a time. let base = full_words * 64; @@ -162,7 +251,15 @@ 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. +/// When `offset` is a multiple of 64, whole words are loaded from the window. +/// Otherwise, for a window that starts at bit `s` of mask word `w0`, window +/// word `k` is formed from mask words `w0 + k` and `w0 + k + 1` as +/// `(cur >> s) | (next << (64 - s))`. Both paths load whole words only. Whole +/// result words are written with aligned stores when `out_off` is a multiple +/// of 64, and as little-endian bytes otherwise. +/// +/// 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,26 +276,81 @@ 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; - { + if offset % 64 == 0 { let src_bytes = bitmask_window_bytes(mask, offset, len); let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + // SAFETY: The input offset is a multiple of 64 and mask buffers start on a + // 64-byte boundary, which makes `sp` an aligned `u64` pointer. Words + // `0..full_words` lie inside the input window and the output window. Whole + // words are written through `dp` as `u64` only when `out_off` is a multiple + // of 64, and as 8 bytes otherwise, which needs no alignment. 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 - }; + if out_off % 64 == 0 { + for k in 0..full_words { + *dp.add(k) = match op { + UnaryOperator::Not => !*sp.add(k), + _ => unreachable!(), // Positive Negative invalid for bools + }; + } + } else { + for k in 0..full_words { + let w = match op { + UnaryOperator::Not => !*sp.add(k), + _ => unreachable!(), + }; + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + } + } else { + let (w0, s) = (offset / 64, (offset % 64) as u32); + // SAFETY: Mask buffers start on a 64-byte boundary, which aligns the buffer + // start for `u64`, and `bits.len() / 8` whole words lie inside the buffer. + let words = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&mask.bits), mask.bits.len() / 8) + }; + let out_bytes = bitmask_window_bytes_mut(out, out_off, len); + let dp = out_bytes.as_mut_ptr().cast::(); + let write = |k: usize, cur: u64, next: u64| { + let w = match op { + UnaryOperator::Not => !((cur >> s) | (next << (64 - s))), + _ => unreachable!(), + }; + // SAFETY: `k < full_words`, which keeps word `k` inside the output window. It + // is written as an aligned `u64` only when `out_off` is a multiple of 64, + // and as 8 bytes otherwise. + unsafe { + if out_off % 64 == 0 { + *dp.add(k) = w; + } else { + dp.add(k).cast::<[u8; 8]>().write(w.to_le_bytes()); + } + } + }; + // Window words whose two mask words lie whole inside the buffer. + let whole = full_words.min(words.len().saturating_sub(w0 + 1)); + if whole > 0 { + let span = &words[w0..=w0 + whole]; + for k in 0..whole { + write(k, span[k], span[k + 1]); } } + // Remaining words. Mask words that are not whole inside the buffer are + // read with `load_word`, which reads bytes past the end of the buffer + // as zero. + let word = |j: usize| { + words + .get(j) + .copied() + .unwrap_or_else(|| load_word(mask, j * 64)) + }; + for k in whole..full_words { + write(k, word(w0 + k), word(w0 + k + 1)); + } } // Final partial word written one bit at a time. let base = full_words * 64; @@ -399,6 +551,13 @@ pub fn all_ne_mask(a: BitmaskVT<'_>, b: BitmaskVT<'_>) -> bool { /// The implementation processes data in 64-bit words and uses native CPU popcount instructions for /// optimal performance. /// +/// When `offset` is a multiple of 64, whole words are loaded from the window. +/// Otherwise, for a window that starts at bit `s` of mask word `w0`, window +/// word `k` is formed from mask words `w0 + k` and `w0 + k + 1` as +/// `(cur >> s) | (next << (64 - s))`, using whole-word loads only. On both +/// paths the final window word is masked to the window length, which excludes +/// rows past the end of the window. +/// /// # Parameters /// - `m`: Bitmask window as `(mask, offset, length)` tuple /// @@ -410,10 +569,49 @@ pub fn popcount_mask(m: BitmaskVT<'_>) -> usize { if len == 0 { return 0; } + if offset % 64 != 0 { + let (w0, s) = (offset / 64, (offset % 64) as u32); + // SAFETY: Mask buffers start on a 64-byte boundary, which aligns the buffer + // start for `u64`, and `bits.len() / 8` whole words lie inside the buffer. + let words = unsafe { + std::slice::from_raw_parts(mask_bits_as_words(&mask.bits), mask.bits.len() / 8) + }; + let full_words = len / 64; + let mut acc = 0usize; + // Window words whose two mask words lie whole inside the buffer. + let whole = full_words.min(words.len().saturating_sub(w0 + 1)); + if whole > 0 { + let span = &words[w0..=w0 + whole]; + for k in 0..whole { + acc += ((span[k] >> s) | (span[k + 1] << (64 - s))).count_ones() as usize; + } + } + // Remaining words. Mask words that are not whole inside the buffer are + // read with `load_word`, which reads bytes past the end of the buffer + // as zero. + let word = |j: usize| { + words + .get(j) + .copied() + .unwrap_or_else(|| load_word(mask, j * 64)) + }; + for k in whole..full_words { + acc += ((word(w0 + k) >> s) | (word(w0 + k + 1) << (64 - s))).count_ones() as usize; + } + // Final window word, masked to the rows inside the window. + let tail_bits = len % 64; + if tail_bits != 0 { + let k = full_words; + let w = (word(w0 + k) >> s) | (word(w0 + k + 1) << (64 - s)); + acc += (w & (u64::MAX >> (64 - tail_bits))).count_ones() as usize; + } + return acc; + } 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; + // Whole words inside the window, then the bytes of the final partial word. + let full_words = len / 64; + let tail_bytes = total_bytes - full_words * 8; let full_logical_bytes = len / 8; let last_bits = len & 7; let mut acc = 0usize; @@ -554,6 +752,7 @@ mod tests { use super::*; use crate::Bitmask; use crate::kernels::bitmask::clear_trailing_bits; + use crate::kernels::bitmask::tests::{sweep_binop_into, sweep_popcount, sweep_unop_into}; // Helper: Create a Bitmask from a bool slice. fn bm(bits: &[bool]) -> Bitmask { @@ -707,6 +906,23 @@ mod tests { assert_eq!(popcount_mask((&a, 0, a.len())), 4); } + // The scalar kernels have no SIMD width. The sweeps take their lengths + // from a width of 8 words. + #[test] + fn test_binop_std_into_window_offsets() { + sweep_binop_into(8, bitmask_binop_std_into); + } + + #[test] + fn test_unop_std_into_window_offsets() { + sweep_unop_into(8, bitmask_unop_std_into); + } + + #[test] + fn test_popcount_mask_window_offsets() { + sweep_popcount(8, popcount_mask); + } + #[test] fn test_all_true_mask() { let a = bm(&[true, true, true, true]);