BitView Filter

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -18,6 +18,7 @@ all-features = true
 
 [features]
 arrow = []
+bench = []
 memmap2 = ["dep:memmap2"]
 serde = ["dep:serde"]
 warn-copy = ["dep:log"]
@@ -43,7 +44,13 @@ workspace = true
 [dev-dependencies]
 arrow-buffer = { workspace = true }
 divan = { workspace = true }
+rand = { workspace = true }
 rstest = { workspace = true }
+vortex-buffer = { path = ".", features = ["bench"] }
+
+[[bench]]
+name = "bitview_filter"
+harness = false
 
 [[bench]]
 name = "vortex_buffer"

vortex-buffer/Cargo.toml

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![allow(clippy::unwrap_used)]
+
+use std::hint::black_box;
+use std::iter::Iterator;
+
+use divan::Bencher;
+use rand::prelude::StdRng;
+use rand::{Rng, SeedableRng};
+use vortex_buffer::bench::bench_filter_in_place_scalar;
+use vortex_buffer::{buffer_mut, BitBuffer};
+
+fn main() {
+    divan::main();
+}
+
+// Focus on benchmarking for our known vector length.
+const N: usize = 1024;
+type BitView<'a> = vortex_buffer::BitView<'a, 128>;
+const MASK_DENSITY: &[f64] = &[
+    0.0, 0.01, 0.05, 0.1, 0.25, // 0.3,
+    // 0.4,
+    0.5,  // 0.6,
+    0.75, // 0.85,
+    0.9,  // 0.95,
+    0.99, 1.00,
+];
+
+#[divan::bench(
+    types = [u8, u16, u32, u64, u128],
+    args = MASK_DENSITY,
+)]
+fn filter_scalar_in_place<T: Default + Copy>(bencher: Bencher, mask_density: f64) {
+    let mut buffer = buffer_mut![T::default(); N];
+
+    let mut rng = StdRng::seed_from_u64(0);
+    let mask = (0..N)
+        .map(|_| rng.random_bool(mask_density))
+        .collect::<BitBuffer>();
+
+    bencher.bench_local(|| {
+        let view = BitView::new(mask.inner().as_ref().try_into().unwrap());
+        bench_filter_in_place_scalar(&view, &mut buffer);
+        black_box(&mut buffer);
+    });
+}

vortex-buffer/benches/bitview_filter.rs

@@ -52,8 +52,10 @@ impl<const NB: usize> BitView<'static, NB> {
 }
 
 impl<'a, const NB: usize> BitView<'a, NB> {
-    const N: usize = NB * 8;
-    const N_WORDS: usize = NB * 8 / (usize::BITS as usize);
+    /// The number of bits in the view.
+    pub const N: usize = NB * 8;
+    /// The number of machine words in the view.
+    pub const N_WORDS: usize = NB * 8 / (usize::BITS as usize);
 
     const _ASSERT_MULTIPLE_OF_8: () = assert!(
         NB % 8 == 0,
@@ -148,7 +150,7 @@ impl<'a, const NB: usize> BitView<'a, NB> {
     /// The words are loaded using unaligned loads to ensure correct bit ordering.
     /// For example, bit 0 is located in `word & 1 << 0`, bit 63 is located in `word & 1 << 63`,
     /// assuming the word size is 64 bits.
-    fn iter_words(&self) -> impl Iterator<Item = usize> + '_ {
+    pub fn iter_words(&self) -> impl Iterator<Item = usize> + '_ {
         let ptr = self.bits.as_ptr().cast::<usize>();
         // We use constant N_WORDS to trigger loop unrolling.
         (0..Self::N_WORDS).map(move |idx| unsafe { ptr.add(idx).read_unaligned() })

vortex-buffer/src/bit/view.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Implementations for filtering data using a [`BitView`] selection mask.
+//!
+//! These implementations are highly optimized as filtering is such a core operation within Vortex.
+
+pub(crate) mod scalar_in_place;
+
+use crate::BitView;
+
+impl<'a, const NB: usize> BitView<'a, NB> {
+    /// Filters the given slice of items in place.
+    ///
+    /// After calling this method, the first `self.true_count()` elements of `items`
+    /// will contain the filtered items. The remaining elements beyond that point are undefined.
+    pub fn filter_in_place<T: Copy>(&self, items: &mut [T]) {
+        match self.true_count() {
+            0 => {
+                // No items to keep; do nothing.
+            }
+            n if n == items.len() => {
+                // All items to keep; do nothing.
+            }
+            _ => {
+                // Some items to keep; do the filtering.
+                scalar_in_place::filter_in_place_scalar(self, items);
+            }
+        }
+    }
+}

vortex-buffer/src/filter/mod.rs

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::BitView;
+use std::ptr;
+
+/// Benchmark wrapper for [`filter_in_place_scalar`].
+#[doc(hidden)]
+#[cfg(feature = "bench")]
+pub fn bench_filter_in_place_scalar<const NB: usize, T: Copy>(
+    bit_view: &BitView<NB>,
+    items: &mut [T],
+) {
+    filter_in_place_scalar(bit_view, items);
+}
+
+/// Filters the given slice of items in place according to the provided BitView using scalar
+/// (non-SIMD) code.
+///
+/// The caller *should* handle where the BitView has zero or full true counts to avoid unnecessary
+/// work.
+pub(crate) fn filter_in_place_scalar<const NB: usize, T: Copy>(
+    bit_view: &BitView<NB>,
+    items: &mut [T],
+) {
+    let mut read_ptr = items.as_ptr();
+    let mut write_ptr = items.as_mut_ptr();
+
+    for mut word in bit_view.iter_words() {
+        match word {
+            usize::MAX => {
+                // All items => copy usize::BITS items.
+                unsafe {
+                    // We cannot guarantee non-overlapping, so use ptr::copy rather than
+                    // ptr::copy_nonoverlapping.
+                    ptr::copy(read_ptr, write_ptr, usize::BITS as usize);
+                    read_ptr = read_ptr.add(usize::BITS as usize);
+                    write_ptr = write_ptr.add(usize::BITS as usize);
+                }
+            }
+            _ => {
+                // Note this also handles word == 0 case by skipping the loop entirely.
+                while word != 0 {
+                    let bit_pos = word.trailing_zeros();
+                    unsafe {
+                        ptr::copy(read_ptr.add(bit_pos as usize), write_ptr, 1);
+                        write_ptr = write_ptr.add(1);
+                    }
+                    word &= word - 1; // Clear the bit at `bit_pos`
+                }
+                unsafe { read_ptr = read_ptr.add(usize::BITS as usize) };
+            }
+        }
+    }
+}

vortex-buffer/src/filter/scalar_in_place.rs

@@ -67,6 +67,7 @@ mod r#const;
 #[cfg(gpu_unstable)]
 mod cuda;
 mod debug;
+mod filter;
 mod macros;
 #[cfg(feature = "memmap2")]
 mod memmap2;
@@ -83,3 +84,9 @@ pub type ByteBufferMut = BufferMut<u8>;
 
 /// A const-aligned buffer of u8.
 pub type ConstByteBuffer<const A: usize> = ConstBuffer<u8, A>;
+
+/// Functions exported for benchmarking purposes.
+#[cfg(feature = "bench")]
+pub mod bench {
+    pub use super::filter::scalar_in_place::bench_filter_in_place_scalar;
+}

vortex-buffer/src/lib.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+mod scalar;
+
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_mask::{Mask, MaskIter};
 
@@ -202,7 +204,7 @@ mod tests {
         assert_eq!(result, buffer![1u32, 2, 5]);
     }
 
-    use vortex_buffer::{BufferMut, buffer_mut};
+    use vortex_buffer::{buffer_mut, BufferMut};
 
     #[test]
     fn test_filter_all_true() {