diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index ae9d7e6cc05..3d4994e8d21 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -48,3 +48,7 @@ harness = false [[bench]] name = "vortex_bitbuffer" harness = false + +[[bench]] +name = "bitbuffer_count" +harness = false diff --git a/vortex-buffer/benches/bitbuffer_count.rs b/vortex-buffer/benches/bitbuffer_count.rs new file mode 100644 index 00000000000..f35af10324b --- /dev/null +++ b/vortex-buffer/benches/bitbuffer_count.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks comparing strategies for computing a bitwise-op result *and* its `true_count`: +//! +//! * `*_two_pass`: current behaviour, materialise the result buffer then run the (vectorised) +//! `count_ones` kernel as a separate pass (`&a & &b` followed by `.true_count()`). +//! * `*_fused`: single pass that combines the words and accumulates `count_ones` together +//! (`BitBuffer::bitand_with_true_count`). + +#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)] + +use divan::Bencher; +use divan::black_box; +use vortex_buffer::BitBuffer; + +fn main() { + divan::main(); +} + +// Lengths in bits, from L1-resident up to L2/L3. +const SIZES: &[usize] = &[ + 1_024, // 128 B + 8_192, // 1 KB + 65_536, // 8 KB + 524_288, // 64 KB +]; + +fn inputs(len: usize) -> (BitBuffer, BitBuffer) { + let lhs = + BitBuffer::from_iter((0..len).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) & 1 == 0)); + let rhs = BitBuffer::from_iter((0..len).map(|i| (i.wrapping_mul(40_503) >> 7) & 1 == 0)); + (lhs, rhs) +} + +#[divan::bench(args = SIZES)] +fn and_two_pass(bencher: Bencher, len: usize) { + let (lhs, rhs) = inputs(len); + bencher.bench_local(|| { + let result = black_box(&lhs) & black_box(&rhs); + let count = result.true_count(); + (result, count) + }); +} + +#[divan::bench(args = SIZES)] +fn and_fused(bencher: Bencher, len: usize) { + let (lhs, rhs) = inputs(len); + bencher.bench_local(|| black_box(&lhs).bitand_with_true_count(black_box(&rhs))); +} + +// Owned-LHS scenario: the current code can combine in-place (no allocation) then count +// separately; the fused path always allocates but does a single pass. A fresh owned, uniquely +// owned LHS is cloned per iteration so the in-place fast path is actually taken. +#[divan::bench(args = SIZES)] +fn and_owned_two_pass(bencher: Bencher, len: usize) { + let (lhs, rhs) = inputs(len); + bencher + .with_inputs(|| lhs.clone()) + .bench_local_values(|lhs| { + let result = lhs & black_box(&rhs); + let count = result.true_count(); + (result, count) + }); +} + +#[divan::bench(args = SIZES)] +fn and_owned_fused(bencher: Bencher, len: usize) { + let (lhs, rhs) = inputs(len); + bencher + .with_inputs(|| lhs.clone()) + .bench_local_values(|lhs| lhs.bitand_with_true_count(black_box(&rhs))); +} diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 457827f4ab9..06120a0b4a5 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -25,6 +25,7 @@ use crate::bit::collect_bool_word; use crate::bit::count_ones::count_ones; use crate::bit::get_bit_unchecked; use crate::bit::ops::bitwise_binary_op; +use crate::bit::ops::bitwise_binary_op_counted; use crate::bit::ops::bitwise_binary_op_lhs_owned; use crate::bit::ops::bitwise_unary_op; use crate::bit::ops::bitwise_unary_op_copy; @@ -541,6 +542,21 @@ impl BitBuffer { bitwise_binary_op_lhs_owned(self, rhs, |a, b| a & !b) } + /// Bitwise AND of two buffers, returning the result and its true count from a single pass. + pub fn bitand_with_true_count(&self, rhs: &BitBuffer) -> (BitBuffer, usize) { + bitwise_binary_op_counted(self, rhs, |a, b| a & b) + } + + /// Bitwise OR of two buffers, returning the result and its true count from a single pass. + pub fn bitor_with_true_count(&self, rhs: &BitBuffer) -> (BitBuffer, usize) { + bitwise_binary_op_counted(self, rhs, |a, b| a | b) + } + + /// Bitwise AND NOT (`self & !rhs`), returning the result and its true count from a single pass. + pub fn bitand_not_with_true_count(&self, rhs: &BitBuffer) -> (BitBuffer, usize) { + bitwise_binary_op_counted(self, rhs, |a, b| a & !b) + } + /// Iterate through bits in a buffer. /// /// # Arguments diff --git a/vortex-buffer/src/bit/ops.rs b/vortex-buffer/src/bit/ops.rs index 5006fbcd4cc..b0129f5eb7a 100644 --- a/vortex-buffer/src/bit/ops.rs +++ b/vortex-buffer/src/bit/ops.rs @@ -235,6 +235,113 @@ pub(super) fn bitwise_binary_op u64>( BitBuffer::new(bytes.freeze(), len) } +/// Combine `left` and `right` word-by-word with `op`, calling `comb` on each result word so a +/// caller can fold a reduction (e.g. `count_ones`) into the same pass. The word given to `comb` +/// has bits past `len` cleared; the buffer keeps `op`'s raw output in those positions. +pub(super) fn bitwise_binary_op_with( + left: &BitBuffer, + right: &BitBuffer, + op: F, + comb: C, +) -> BitBuffer +where + F: FnMut(u64, u64) -> u64, + C: FnMut(u64), +{ + assert_eq!(left.len(), right.len()); + let len = left.len(); + if len == 0 { + return BitBuffer::empty(); + } + + // Byte-aligned operands: logical bits map onto physical `u64` words, so we can read the backing + // bytes straight as words instead of shifting through `iter_padded`. + if left.offset().is_multiple_of(8) && right.offset().is_multiple_of(8) { + let n_bytes = len.div_ceil(8); + let l_start = left.offset() / 8; + let r_start = right.offset() / 8; + let lhs = &left.inner().as_slice()[l_start..l_start + n_bytes]; + let rhs = &right.inner().as_slice()[r_start..r_start + n_bytes]; + + let (lhs_words, lhs_tail) = lhs.as_chunks::<8>(); + let (rhs_words, rhs_tail) = rhs.as_chunks::<8>(); + + let words = lhs_words + .iter() + .zip(rhs_words) + .map(|(l, r)| (u64::from_le_bytes(*l), u64::from_le_bytes(*r))) + .chain((!lhs_tail.is_empty()).then(|| (read_u64_le(lhs_tail), read_u64_le(rhs_tail)))); + return combine_words(words, len, op, comb); + } + + // Sub-byte offset: `iter_padded` realigns the bits and zero-pads the final word. + let words = left + .chunks() + .iter_padded() + .zip(right.chunks().iter_padded()); + combine_words(words, len, op, comb) +} + +/// Combine an iterator of `(left, right)` words via `op` into a `len`-bit [`BitBuffer`], folding +/// `comb` over each in-range result word. Consumes exactly `ceil(len / 64)` words (any extra the +/// producer yields, e.g. `iter_padded`'s trailing pad word, is dropped). Only the final word can +/// hold bits past `len`, so it alone is masked for `comb`; the buffer keeps `op`'s raw output. +#[inline] +fn combine_words(words: I, len: usize, mut op: F, mut comb: C) -> BitBuffer +where + I: Iterator, + F: FnMut(u64, u64) -> u64, + C: FnMut(u64), +{ + let n_words = len.div_ceil(64); + let mut out: BufferMut = BufferMut::with_capacity(n_words); + let mut words = words.take(n_words); + + // Words fully within `len` need no masking; handling the final word separately keeps this loop + // uniform, with no per-word branch. + for (l, r) in words.by_ref().take(n_words - 1) { + let word = op(l, r); + comb(word); + out.push(word); + } + // Final word: clear bits past `len` for `comb` only. + if let Some((l, r)) = words.next() { + let word = op(l, r); + comb(word & last_word_mask(len)); + out.push(word); + } + + let mut bytes = out.into_byte_buffer(); + bytes.truncate(len.div_ceil(8)); + BitBuffer::new(bytes.freeze(), len) +} + +/// Mask of the in-range bits in the final word: the low `len % 64` bits, or all 64 when `len` is +/// a multiple of 64. +#[inline] +fn last_word_mask(len: usize) -> u64 { + let rem = len % 64; + if rem == 0 { + u64::MAX + } else { + (1u64 << rem) - 1 + } +} + +/// Like [`bitwise_binary_op`], but also returns the result's set-bit count, folded into the same +/// pass instead of re-scanning the buffer. +pub(super) fn bitwise_binary_op_counted u64>( + left: &BitBuffer, + right: &BitBuffer, + op: F, +) -> (BitBuffer, usize) { + let mut count = 0usize; + let buffer = bitwise_binary_op_with(left, right, op, |word| { + count += word.count_ones() as usize; + }); + (buffer, count) +} + #[cfg(test)] mod tests { use std::ops::Not; @@ -367,6 +474,44 @@ mod tests { } } + /// The fused path must yield the same buffer and true count as the reference + /// (`bitwise_binary_op` + a separate `true_count`), across all ops, offsets, and lengths, + /// including lengths like 127 whose final word holds out-of-range bits. + #[rstest] + #[case::aligned(0, 0)] + #[case::equal_nonzero(3, 3)] + #[case::byte_aligned(8, 16)] + #[case::mismatch(0, 5)] + fn counted_matches_reference(#[case] left_offset: usize, #[case] right_offset: usize) { + #[allow(clippy::cast_possible_truncation)] + let make = |offset: usize, len: usize, salt: u8| -> BitBuffer { + let bytes: ByteBufferMut = (0..(offset + len).div_ceil(8).max(1)) + .map(|i| (i as u8).wrapping_mul(31).wrapping_add(salt)) + .collect(); + BitBufferMut::from_buffer(bytes, offset, len).freeze() + }; + let ops: [fn(u64, u64) -> u64; 4] = + [|a, b| a & b, |a, b| a | b, |a, b| a ^ b, |a, b| a & !b]; + + for len in [1usize, 5, 8, 63, 64, 65, 127, 128, 129, 200, 256] { + let left = make(left_offset, len, 0xC3); + let right = make(right_offset, len, 0x5A); + for op in ops { + let (got_buf, got_count) = bitwise_binary_op_counted(&left, &right, op); + let expected_buf = bitwise_binary_op(&left, &right, op); + assert_eq!( + got_buf, expected_buf, + "buffer mismatch loff={left_offset} roff={right_offset} len={len}" + ); + assert_eq!( + got_count, + expected_buf.true_count(), + "count mismatch loff={left_offset} roff={right_offset} len={len}" + ); + } + } + } + /// Regression test for a bug where [`bitwise_unary_op`] produced corrupt results when /// the [`BitBuffer`]'s underlying byte pointer was not u64-aligned. Slicing a buffer by /// a non-multiple-of-8 number of bytes can cause this misalignment. The bug only diff --git a/vortex-mask/src/bitops.rs b/vortex-mask/src/bitops.rs index 4bebcf28c40..c2e01ab6cdc 100644 --- a/vortex-mask/src/bitops.rs +++ b/vortex-mask/src/bitops.rs @@ -23,7 +23,10 @@ impl BitAnd for &Mask { (_, AllOr::All) => self.clone(), (AllOr::None, _) => Mask::new_false(self.len()), (_, AllOr::None) => Mask::new_false(self.len()), - (AllOr::Some(lhs), AllOr::Some(rhs)) => Mask::from_buffer(lhs & rhs), + (AllOr::Some(lhs), AllOr::Some(rhs)) => { + let (buffer, true_count) = lhs.bitand_with_true_count(rhs); + Mask::from_buffer_with_true_count(buffer, true_count) + } } } } @@ -31,7 +34,6 @@ impl BitAnd for &Mask { impl BitAnd<&Mask> for Mask { type Output = Mask; - /// Owned-left AND: can reuse the left buffer in-place when possible. fn bitand(self, rhs: &Mask) -> Self::Output { if self.len() != rhs.len() { vortex_panic!("Masks must have the same length"); @@ -41,8 +43,9 @@ impl BitAnd<&Mask> for Mask { (AllOr::All, _) => rhs.clone(), (AllOr::None, _) | (_, AllOr::None) => Mask::new_false(self.len()), (_, AllOr::All) => self, - (AllOr::Some(_), AllOr::Some(rhs_buf)) => { - Mask::from_buffer(self.into_bit_buffer() & rhs_buf) + (AllOr::Some(lhs_buf), AllOr::Some(rhs_buf)) => { + let (buffer, true_count) = lhs_buf.bitand_with_true_count(rhs_buf); + Mask::from_buffer_with_true_count(buffer, true_count) } } } @@ -61,7 +64,10 @@ impl BitOr for &Mask { (_, AllOr::All) => Mask::new_true(self.len()), (AllOr::None, _) => rhs.clone(), (_, AllOr::None) => self.clone(), - (AllOr::Some(lhs), AllOr::Some(rhs)) => Mask::from_buffer(lhs | rhs), + (AllOr::Some(lhs), AllOr::Some(rhs)) => { + let (buffer, true_count) = lhs.bitor_with_true_count(rhs); + Mask::from_buffer_with_true_count(buffer, true_count) + } } } } @@ -69,7 +75,6 @@ impl BitOr for &Mask { impl BitOr<&Mask> for Mask { type Output = Mask; - /// Owned-left OR: can reuse the left buffer in-place when possible. fn bitor(self, rhs: &Mask) -> Self::Output { if self.len() != rhs.len() { vortex_panic!("Masks must have the same length"); @@ -79,8 +84,9 @@ impl BitOr<&Mask> for Mask { (AllOr::All, _) | (_, AllOr::All) => Mask::new_true(self.len()), (AllOr::None, _) => rhs.clone(), (_, AllOr::None) => self, - (AllOr::Some(_), AllOr::Some(rhs_buf)) => { - Mask::from_buffer(self.into_bit_buffer() | rhs_buf) + (AllOr::Some(lhs_buf), AllOr::Some(rhs_buf)) => { + let (buffer, true_count) = lhs_buf.bitor_with_true_count(rhs_buf); + Mask::from_buffer_with_true_count(buffer, true_count) } } } @@ -96,8 +102,9 @@ impl Mask { (AllOr::None, _) | (_, AllOr::All) => Mask::new_false(self.len()), (_, AllOr::None) => self, (AllOr::All, _) => !rhs, - (AllOr::Some(_), AllOr::Some(rhs_buf)) => { - Mask::from_buffer(self.into_bit_buffer().into_bitand_not(rhs_buf)) + (AllOr::Some(lhs_buf), AllOr::Some(rhs_buf)) => { + let (buffer, true_count) = lhs_buf.bitand_not_with_true_count(rhs_buf); + Mask::from_buffer_with_true_count(buffer, true_count) } } } @@ -118,7 +125,10 @@ impl Not for &Mask { match self.bit_buffer() { AllOr::All => Mask::new_false(self.len()), AllOr::None => Mask::new_true(self.len()), - AllOr::Some(buffer) => Mask::from_buffer(!buffer), + // Negation's true count is the complement of self's cached count; no popcount. + AllOr::Some(buffer) => { + Mask::from_buffer_with_true_count(!buffer, self.len() - self.true_count()) + } } } } @@ -126,10 +136,45 @@ impl Not for &Mask { #[cfg(test)] #[expect(clippy::many_single_char_names)] mod tests { + use rstest::rstest; use vortex_buffer::BitBuffer; use super::*; + /// Regression test for : the boolean + /// operators now cache a fused true count, which must match a brute-force count, including + /// lengths that aren't a multiple of 64, where the final word holds out-of-range bits. + #[rstest] + #[case(1)] + #[case(63)] + #[case(64)] + #[case(127)] + #[case(200)] + #[case(1000)] + fn test_boolean_ops_true_count(#[case] len: usize) { + let a_bits: Vec = (0..len).map(|i| (i * 7 + 1) % 5 < 3).collect(); + let b_bits: Vec = (0..len).map(|i| (i * 3 + 2) % 7 < 4).collect(); + let a = Mask::from_buffer(BitBuffer::from_iter(a_bits.iter().copied())); + let b = Mask::from_buffer(BitBuffer::from_iter(b_bits.iter().copied())); + + let count = |f: fn(bool, bool) -> bool| { + a_bits + .iter() + .zip(&b_bits) + .filter(|(x, y)| f(**x, **y)) + .count() + }; + + assert_eq!((&a & &b).true_count(), count(|x, y| x & y), "AND len={len}"); + assert_eq!((&a | &b).true_count(), count(|x, y| x | y), "OR len={len}"); + assert_eq!( + a.clone().bitand_not(&b).true_count(), + count(|x, y| x & !y), + "ANDNOT len={len}" + ); + assert_eq!((!&a).true_count(), count(|x, _| !x), "NOT len={len}"); + } + #[test] fn test_bitand_all_combinations() { let len = 5; diff --git a/vortex-mask/src/lib.rs b/vortex-mask/src/lib.rs index 401077f0b34..9ae50eebb81 100644 --- a/vortex-mask/src/lib.rs +++ b/vortex-mask/src/lib.rs @@ -187,8 +187,20 @@ impl Mask { /// Create a new [`Mask`] from a [`BitBuffer`]. pub fn from_buffer(buffer: BitBuffer) -> Self { - let len = buffer.len(); let true_count = buffer.true_count(); + Self::from_buffer_with_true_count(buffer, true_count) + } + + /// Create a [`Mask`] from a [`BitBuffer`] whose true count is already known, skipping the + /// popcount pass [`Mask::from_buffer`] would run. `true_count` must equal + /// `buffer.true_count()`; a wrong value yields a `Mask` with corrupt cached statistics. + pub(crate) fn from_buffer_with_true_count(buffer: BitBuffer, true_count: usize) -> Self { + debug_assert_eq!( + true_count, + buffer.true_count(), + "from_buffer_with_true_count given a count that disagrees with the buffer" + ); + let len = buffer.len(); if true_count == 0 { return Self::AllFalse(len);