BitView Filter

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

Author: gatesn

Cargo.lock

@@ -18,7 +18,6 @@ all-features = true
 
 [features]
 arrow = []
-bench = []
 memmap2 = ["dep:memmap2"]
 serde = ["dep:serde"]
 warn-copy = ["dep:log"]
@@ -44,13 +43,7 @@ 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

@@ -156,6 +156,17 @@ impl<'a, const NB: usize> BitView<'a, NB> {
         (0..Self::N_WORDS).map(move |idx| unsafe { ptr.add(idx).read_unaligned() })
     }
 
+    /// Iterate the [`BitView`] in fixed-size words.
+    ///
+    /// 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.
+    pub fn iter_sized<W: 'static>(&self) -> impl Iterator<Item = W> + '_ {
+        let ptr = self.bits.as_ptr().cast::<W>();
+        // We use constant N_WORDS to trigger loop unrolling.
+        (0..(NB / size_of::<W>())).map(move |idx| unsafe { ptr.add(idx).read_unaligned() })
+    }
+
     /// Runs the provided function `f` for each index of a `true` bit in the view.
     pub fn iter_ones<F>(&self, mut f: F)
     where

vortex-buffer/src/bit/view.rs

@@ -67,7 +67,6 @@ mod r#const;
 #[cfg(gpu_unstable)]
 mod cuda;
 mod debug;
-mod filter;
 mod macros;
 #[cfg(feature = "memmap2")]
 mod memmap2;
@@ -84,9 +83,3 @@ 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/filter/mod.rs

@@ -34,9 +34,16 @@ num-traits = { workspace = true }
 [features]
 default = ["arrow"]
 arrow = ["dep:arrow-array", "dep:arrow-buffer", "dep:arrow-schema"]
+bench = []
 
 [dev-dependencies]
 divan = { workspace = true }
+rand = { workspace = true }
+vortex-compute = { path = ".", features = ["bench"] }
+
+[[bench]]
+name = "filter_bitview"
+harness = false
 
 [[bench]]
 name = "filter_buffer_mut"

vortex-buffer/src/lib.rs

@@ -9,8 +9,8 @@ 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};
+use vortex_compute::bench;
 
 fn main() {
     divan::main();
@@ -19,6 +19,7 @@ fn 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,
@@ -28,11 +29,23 @@ const MASK_DENSITY: &[f64] = &[
     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) {
+#[divan::bench(types = [u8, u16, u32, u64, u128],args = MASK_DENSITY)]
+fn filter_scalar<T: Default + Copy>(bencher: Bencher, mask_density: f64) {
+    bench_filter_fn(bencher, mask_density, bench::bench_filter_scalar::<_, T>)
+}
+
+#[cfg(target_arch = "aarch64")]
+#[divan::bench(types = [u8, u16, u32, u64, u128],args = MASK_DENSITY)]
+fn filter_neon<T: Default + Copy>(bencher: Bencher, mask_density: f64) {
+    if std::arch::is_aarch64_feature_detected!("neon") {
+        bench_filter_fn(bencher, mask_density, bench::bench_filter_neon::<_, T>)
+    }
+}
+
+fn bench_filter_fn<T: Default + Copy, F>(bencher: Bencher, mask_density: f64, f: F)
+where
+    F: Fn(&BitView, &mut [T]),
+{
     let mut buffer = buffer_mut![T::default(); N];
 
     let mut rng = StdRng::seed_from_u64(0);
@@ -42,7 +55,7 @@ fn filter_scalar_in_place<T: Default + Copy>(bencher: Bencher, mask_density: f64
 
     bencher.bench_local(|| {
         let view = BitView::new(mask.inner().as_ref().try_into().unwrap());
-        bench_filter_in_place_scalar(&view, &mut buffer);
+        f(&view, &mut buffer);
         black_box(&mut buffer);
     });
 }

vortex-compute/Cargo.toml

@@ -1,8 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-mod scalar;
-
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_mask::{Mask, MaskIter};
 

vortex-buffer/benches/bitview_filter.rs vortex-compute/benches/filter_bitview.rsvortex-buffer/benches/bitview_filter.rs renamed to vortex-compute/benches/filter_bitview.rs

@@ -8,6 +8,7 @@ use std::ops::Deref;
 mod bitbuffer;
 mod buffer;
 mod mask;
+pub(crate) mod slice;
 mod vector;
 
 /// Function for filtering based on a selection mask.

vortex-compute/src/filter/buffer.rs

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::filter::Filter;
+use std::arch::is_aarch64_feature_detected;
+use vortex_buffer::BitView;
+
+#[cfg(target_arch = "aarch64")]
+pub(crate) mod neon;
+pub(crate) mod scalar;
+
+impl<'a, const NB: usize, T: Copy> Filter<BitView<'a, NB>> for &mut [T] {
+    type Output = ();
+
+    fn filter(self, mask: &BitView<'a, NB>) -> Self::Output {
+        #[cfg(target_arch = "aarch64")]
+        {
+            if is_aarch64_feature_detected!("neon") {
+                return unsafe { neon::filter_neon(self, mask) };
+            }
+        }
+
+        // Otherwise, fall back to scalar implementation
+        scalar::filter_scalar(self, mask);
+    }
+}