BitView Filter

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

Author: gatesn

vortex-compute/benches/filter_bitview.rs

@@ -10,7 +10,7 @@ use std::iter::Iterator;
 use divan::Bencher;
 use rand::prelude::StdRng;
 use rand::{Rng, SeedableRng};
-use vortex_buffer::{buffer_mut, BitBuffer};
+use vortex_buffer::{BitBuffer, buffer_mut};
 use vortex_compute::bench;
 use vortex_compute::filter::Filter;
 

vortex-compute/src/filter/slice/neon/mod.rs

@@ -12,19 +12,19 @@ mod neon_u32;
 mod neon_u8;
 
 use std::arch::is_aarch64_feature_detected;
+
 use vortex_buffer::BitView;
 use vortex_error::vortex_panic;
 
 /// Benchmark wrapper for [`filter_neon`].
 #[doc(hidden)]
 #[cfg(feature = "bench")]
-#[cfg(target_arch = "aarch64")]
 #[inline(never)]
 pub fn bench_filter_neon<const NB: usize, T: Copy>(bit_view: &BitView<NB>, slice: &mut [T]) {
-    if !is_aarch64_feature_detected!("neon") {
-        vortex_panic!("NEON not detected on this CPU");
+    if is_aarch64_feature_detected!("neon") {
+        unsafe { filter_neon(slice, bit_view) }
     }
-    unsafe { filter_neon(slice, bit_view) }
+    vortex_panic!("NEON not detected on this CPU");
 }
 
 /// Filters the given slice of items in place according to the provided BitView using neon
@@ -52,6 +52,7 @@ pub(super) unsafe fn filter_neon<const NB: usize, T: Copy>(slice: &mut [T], mask
 }
 
 #[cfg(test)]
+#[allow(clippy::cast_possible_truncation)]
 mod tests {
     use super::*;
 

vortex-compute/src/filter/slice/neon/neon_u16.rs

@@ -3,11 +3,13 @@
 
 #![allow(unsafe_op_in_unsafe_fn)]
 
-use crate::filter::slice::neon::neon_u8::SHUFFLE_MASKS;
 use std::arch::aarch64::*;
 use std::ptr;
+
 use vortex_buffer::BitView;
 
+use crate::filter::slice::neon::neon_u8::SHUFFLE_MASKS;
+
 /// For u16 types, we perform a similar strategy to u8 with a few key differences.
 ///
 /// When it comes to shuffling u16 elements, we load u16x8 values into a uint8x8x2 vector. This
@@ -57,10 +59,10 @@ pub(super) unsafe fn filter_neon_u16<const NB: usize>(data: *mut u16, mask: &Bit
                             let count = byte.count_ones() as usize;
                             let shuffle_vec = vld1_u8(SHUFFLE_MASKS[byte as usize].as_ptr());
                             // Shuffle both lower and higher byte vectors separately.
-                            let compressed = uint8x8x2_t {
-                                0: vtbl1_u8(values.0, shuffle_vec),
-                                1: vtbl1_u8(values.1, shuffle_vec),
-                            };
+                            let compressed = uint8x8x2_t(
+                                vtbl1_u8(values.0, shuffle_vec),
+                                vtbl1_u8(values.1, shuffle_vec),
+                            );
 
                             // Store all compressed values, and only increment write_ptr by count.
                             vst2_u8(write_ptr.cast(), compressed);

vortex-compute/src/filter/slice/neon/neon_u32.rs

@@ -3,11 +3,13 @@
 
 #![allow(unsafe_op_in_unsafe_fn)]
 
-use crate::filter::slice::neon::neon_u8::SHUFFLE_MASKS;
 use std::arch::aarch64::*;
 use std::ptr;
+
 use vortex_buffer::BitView;
 
+use crate::filter::slice::neon::neon_u8::SHUFFLE_MASKS;
+
 /// For u32 values we can only look at 4 values at a time (128 bits).
 /// Therefore, we have a very manageable 16 possible bitmask combinations (0..15) and therefore
 /// avoid the need for large lookup tables.
@@ -43,12 +45,12 @@ pub(super) unsafe fn filter_neon_u32<const NB: usize>(data: *mut u32, mask: &Bit
                     let shuffle_vec = vld1_u8(SHUFFLE_MASKS[byte as usize].as_ptr());
 
                     // Shuffle all four byte vectors separately.
-                    let compressed = uint8x8x4_t {
-                        0: vtbl1_u8(values.0, shuffle_vec),
-                        1: vtbl1_u8(values.1, shuffle_vec),
-                        2: vtbl1_u8(values.2, shuffle_vec),
-                        3: vtbl1_u8(values.3, shuffle_vec),
-                    };
+                    let compressed = uint8x8x4_t(
+                        vtbl1_u8(values.0, shuffle_vec),
+                        vtbl1_u8(values.1, shuffle_vec),
+                        vtbl1_u8(values.2, shuffle_vec),
+                        vtbl1_u8(values.3, shuffle_vec),
+                    );
 
                     // Store all compressed values, and only increment write_ptr by count.
                     vst4_u8(write_ptr.cast(), compressed);

vortex-compute/src/filter/slice/neon/neon_u8.rs

@@ -5,6 +5,7 @@
 
 use std::arch::aarch64::*;
 use std::ptr;
+
 use vortex_buffer::BitView;
 
 /// For u8 types, we use NEON's tbl lookup instruction to perform a shuffle based on a pre-computed

vortex-compute/src/filter/slice/scalar.rs

@@ -45,7 +45,7 @@ pub(super) fn filter_scalar<const NB: usize, T: Copy>(slice: &mut [T], mask: &Bi
                     let bit_pos = word.trailing_zeros();
                     word &= word - 1; // Clear the bit at `bit_pos`
                     let span = word.trailing_ones();
-                    word = word >> span;
+                    word >>= span;
                     unsafe {
                         ptr::copy(read_ptr.add(bit_pos as usize), write_ptr, span as usize);
                         write_ptr = write_ptr.add(span as usize);