[codex] Experiment: Fearless SIMD byte compression regression

Cargo.toml

@@ -159,6 +159,7 @@ divan = { package = "codspeed-divan-compat", version = "5.0.0" }
 enum-iterator = "2.0.0"
 env_logger = "0.11"
 fastlanes = { version = "0.7.0", features = ["runtime"] }
+fearless_simd = "1.0.0"
 flatbuffers = "25.2.10"
 fsst-rs = "0.6.0"
 futures = { version = "0.3.31", default-features = false }

vortex-array/Cargo.toml

@@ -31,6 +31,7 @@ bytes = { workspace = true }
 cfg-if = { workspace = true }
 cudarc = { workspace = true, optional = true }
 enum-iterator = { workspace = true }
+fearless_simd = { workspace = true }
 flatbuffers = { workspace = true }
 futures = { workspace = true, features = ["alloc", "async-await", "std"] }
 goldenfile = { workspace = true, optional = true }
@@ -318,3 +319,7 @@ harness = false
 [[bench]]
 name = "probe"
 harness = false
+
+[[bench]]
+name = "filter_u8_simd"
+harness = false

vortex-array/benches/filter_u8_simd.rs

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Direct comparison of the production compression kernels, including in-place compaction.
+
+#![allow(clippy::unwrap_used)]
+
+use std::fmt::Debug;
+
+use divan::Bencher;
+use vortex_buffer::BitBuffer;
+use vortex_buffer::BufferAllocatorRef;
+use vortex_mask::Mask;
+
+// Compile the production sources into this benchmark so no public benchmark-only API is needed.
+#[allow(dead_code)]
+#[path = "../src/arrays/filter/execute/simd_compress/mod.rs"]
+mod simd_compress;
+#[allow(dead_code)]
+#[path = "../src/arrays/filter/execute/slice.rs"]
+mod slice;
+
+fn main() {
+    validate::<u8>();
+    divan::main();
+}
+
+fn values<T: From<u8>>(len: usize) -> Vec<T> {
+    (0..len)
+        .map(|i| T::from(u8::try_from(i % 251).unwrap()))
+        .collect()
+}
+
+fn mask(len: usize, offset: usize, density: f64) -> Mask {
+    let mut state = 0x1234_5678_9abc_def0u64;
+    let bits = BitBuffer::from_iter((0..len + offset).map(|_| {
+        state ^= state << 13;
+        state ^= state >> 7;
+        state ^= state << 17;
+        (state as f64 / u64::MAX as f64) < density
+    }));
+    Mask::from_buffer(bits.slice(offset..len + offset))
+}
+
+fn validate<T: From<u8> + Copy + PartialEq + Debug>() {
+    let allocator = BufferAllocatorRef::statically_allocated();
+    for offset in 0..8 {
+        for len in [64, 65, 127, 128, 151, 513] {
+            let bits = BitBuffer::collect_bool(len + offset, |i| i % 5 < 3);
+            let mask = Mask::from_buffer(bits.slice(offset..len + offset));
+            let mask = mask.values().unwrap();
+            let input = values::<T>(len);
+            let expected = slice::filter_slice_by_bitmap(&input, mask, &allocator);
+            let actual = simd_compress::filter_slice_by_bitmap(&input, mask, &allocator).unwrap();
+            assert_eq!(actual.as_slice(), expected.as_slice());
+            let mut compacted = input;
+            let written = simd_compress::filter_slice_mut_by_bitmap(&mut compacted, mask).unwrap();
+            assert_eq!(&compacted[..written], expected.as_slice());
+        }
+    }
+}
+
+#[divan::bench(types = [u8], args = [0.6, 0.75])]
+fn allocated<T: From<u8> + Copy + Send + Sync>(bencher: Bencher, density: f64) {
+    let values = values::<T>(65_536 / size_of::<T>());
+    let mask = mask(values.len(), 0, density);
+    let mask = mask.values().unwrap();
+    let allocator = BufferAllocatorRef::statically_allocated();
+    bencher.bench(|| {
+        simd_compress::filter_slice_by_bitmap(divan::black_box(&values), mask, &allocator).unwrap()
+    });
+}
+
+#[divan::bench(types = [u8], args = [0.6, 0.75])]
+fn in_place<T: From<u8> + Copy + Send + Sync>(bencher: Bencher, density: f64) {
+    let values = values::<T>(65_536 / size_of::<T>());
+    let mask = mask(values.len(), 0, density);
+    let mask = mask.values().unwrap();
+    bencher.with_inputs(|| values.clone()).bench_refs(|values| {
+        simd_compress::filter_slice_mut_by_bitmap(divan::black_box(values), mask).unwrap()
+    });
+}

vortex-array/src/arrays/filter/execute/simd_compress/fearless.rs

@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Compact eight byte lanes using a portable 16-byte shuffle.
+
+use std::ptr;
+
+use fearless_simd::prelude::*;
+use fearless_simd::u8x16;
+use vortex_mask::MaskValues;
+
+use super::super::slice::for_each_mask_word;
+use super::super::slice::low_bits_mask;
+use super::bulk_copy;
+use super::compress_lut;
+use super::compress_tail;
+
+static IDX_LUT: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 1);
+
+/// # Safety
+///
+/// The pointer contract of the parent module's filter entry points must hold.
+#[allow(clippy::inline_always)]
+#[inline(always)]
+unsafe fn compress_word<S: Simd, const IN_PLACE: bool>(
+    simd: S,
+    src: *const u8,
+    dst: *mut u8,
+    word: u64,
+    word_start: usize,
+    word_len: usize,
+    mut write_pos: usize,
+) -> usize {
+    if word == 0 {
+        return write_pos;
+    }
+    if word == low_bits_mask(word_len) {
+        // SAFETY: forwarded from the caller contract.
+        unsafe { bulk_copy::<IN_PLACE>(src, dst, word_start, word_len, write_pos, 1) };
+        return write_pos + word_len;
+    }
+
+    let mut sub = 0;
+    while sub + 8 <= word_len {
+        let mask = ((word >> sub) & 0xff) as usize;
+        // SAFETY: the chunk holds eight in-bounds bytes. Materializing the value ends
+        // the source borrow before an overlapping in-place store.
+        let bytes = unsafe { src.add(word_start + sub).cast::<[u8; 8]>().read_unaligned() };
+        let chunk = u8x16::from_slice(
+            simd,
+            &[
+                bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], 0,
+                0, 0, 0, 0, 0, 0, 0,
+            ],
+        );
+        let indices = u8x16::from_slice(simd, &IDX_LUT[mask]);
+        let packed = chunk.swizzle_dyn(indices).to_array();
+        // SAFETY: output has vector slack; an in-place store ends no later than the
+        // source chunk just loaded. Later stores overwrite unselected trailing bytes.
+        unsafe { ptr::copy_nonoverlapping(packed.as_ptr(), dst.add(write_pos), 8) };
+        write_pos += mask.count_ones() as usize;
+        sub += 8;
+    }
+
+    if sub < word_len {
+        let bits = (word >> sub) & low_bits_mask(word_len - sub);
+        // SAFETY: forwarded from the caller contract, with only in-bounds tail bits.
+        write_pos =
+            unsafe { compress_tail::<IN_PLACE>(src, dst, bits, word_start + sub, write_pos, 1) };
+    }
+    write_pos
+}
+
+/// # Safety
+///
+/// The pointer contract of the parent module's filter entry points must hold.
+pub(super) unsafe fn compress_fearless_8<const IN_PLACE: bool>(
+    src: *const u8,
+    dst: *mut u8,
+    mask: &MaskValues,
+) -> usize {
+    fearless_simd::dispatch!(fearless_simd::Level::new(), simd => {
+        let mut write_pos = 0;
+        for_each_mask_word(mask, |word, word_start, word_len| {
+            // SAFETY: forwarded from the caller contract.
+            write_pos = unsafe {
+                compress_word::<_, IN_PLACE>(simd, src, dst, word, word_start, word_len, write_pos)
+            };
+        });
+        write_pos
+    })
+}

vortex-array/src/arrays/filter/execute/simd_compress/mod.rs

@@ -31,6 +31,8 @@ use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_mask::MaskValues;
 
+#[cfg(all(target_arch = "aarch64", not(miri)))]
+mod fearless;
 #[cfg(all(target_arch = "aarch64", not(miri)))]
 mod neon;
 #[cfg(test)]

vortex-array/src/arrays/filter/execute/simd_compress/neon.rs

@@ -5,12 +5,9 @@
 //!
 //! Byte-index lookup tables drive `tbl` shuffles for 1-, 2-, 4-, and 8-byte elements.
 
-use core::arch::aarch64::vld1_u8;
 use core::arch::aarch64::vld1q_u8;
 use core::arch::aarch64::vqtbl1q_u8;
-use core::arch::aarch64::vst1_u8;
 use core::arch::aarch64::vst1q_u8;
-use core::arch::aarch64::vtbl1_u8;
 
 use vortex_mask::MaskValues;
 
@@ -38,15 +35,14 @@ pub(super) fn select_kernel<T, const IN_PLACE: bool>(mask: &MaskValues) -> Optio
     }
 
     match size_of::<T>() {
-        1 => Some(compress_neon_8::<IN_PLACE> as Kernel),
+        1 => Some(super::fearless::compress_fearless_8::<IN_PLACE> as Kernel),
         2 => Some(compress_neon_16::<IN_PLACE> as Kernel),
         4 => Some(compress_neon_32::<IN_PLACE> as Kernel),
         8 => Some(compress_neon_64::<IN_PLACE> as Kernel),
         _ => None,
     }
 }
 
-static IDX_LUT_8: [[u8; 8]; 256] = compress_lut::<256, 8>(8, 1);
 static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2);
 static IDX_LUT_32: [[u8; 16]; 16] = compress_lut::<16, 16>(4, 4);
 static IDX_LUT_64: [[u8; 16]; 4] = compress_lut::<4, 16>(2, 8);
@@ -147,16 +143,6 @@ macro_rules! neon_compress_kernel {
     };
 }
 
-neon_compress_kernel!(
-    compress_word_neon_8, compress_neon_8,
-    elem_size: 1,
-    lanes: 8,
-    idx_lut: IDX_LUT_8,
-    load: vld1_u8,
-    tbl: vtbl1_u8,
-    store: vst1_u8
-);
-
 neon_compress_kernel!(
     compress_word_neon_16, compress_neon_16,
     elem_size: 2,