feat[buffer]: specialized collect_bool per arch (#8749)

Currently LLVM cannot create a `bitcast <i1 x W> -> iW` instruction
(https://llvm.org/docs/LangRef.html#bitcast-to-instruction), this may be
added going forwards to support byte->bit packing case `collect_bool`.

This means we must add aarch specific collect_bool impls as done here.

Gives between 10 and 25x improvements

Signed-off-by: Claude <noreply@anthropic.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: joseph-isaacs

.gitignore

@@ -224,6 +224,7 @@ compile_commands.json
 
 # AI Agents
 .agents/settings.local.json
+.agents/scheduled_tasks.lock
 .claude/settings.local.json
 .opencode
 

encodings/bytebool/src/compute.rs

@@ -158,7 +158,12 @@ impl BooleanKernel for ByteBool {
 
 fn truthy_bit_buffer(array: ArrayView<'_, ByteBool>) -> BitBuffer {
     let bytes = array.truthy_bytes();
-    BitBuffer::collect_bool(bytes.len(), |idx| bytes[idx] != 0)
+    // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices `0..bytes.len()` only;
+    // the trivially cheap unchecked gather vectorizes inside the wide pack loop.
+    BitBuffer::collect_bool_multiversioned(
+        bytes.len(),
+        |idx| unsafe { *bytes.get_unchecked(idx) } != 0,
+    )
 }
 
 #[cfg(test)]

vortex-array/src/arrays/decimal/compute/between.rs

@@ -132,8 +132,10 @@ fn between_impl<T: NativeDecimalType>(
 ) -> ArrayRef {
     let buffer = arr.buffer::<T>();
     BoolArray::new(
-        BitBuffer::collect_bool(buffer.len(), |idx| {
-            let value = buffer[idx];
+        BitBuffer::collect_bool_multiversioned(buffer.len(), |idx| {
+            // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices
+            // `0..buffer.len()` only.
+            let value = unsafe { *buffer.get_unchecked(idx) };
             lower.is_none_or(|l| lower_op(l, value)) & upper.is_none_or(|u| upper_op(value, u))
         }),
         arr.validity()

vortex-array/src/arrays/primitive/compute/between.rs

@@ -105,7 +105,7 @@ where
 {
     let slice = arr.as_slice::<T>();
     BoolArray::new(
-        BitBuffer::collect_bool(slice.len(), |idx| {
+        BitBuffer::collect_bool_multiversioned(slice.len(), |idx| {
             // We only iterate upto arr len and |arr| == |slice|.
             let i = unsafe { *slice.get_unchecked(idx) };
             lower_fn(lower, i) & upper_fn(i, upper)

vortex-buffer/Cargo.toml

@@ -48,3 +48,7 @@ harness = false
 [[bench]]
 name = "vortex_bitbuffer"
 harness = false
+
+[[bench]]
+name = "collect_bool"
+harness = false

vortex-buffer/benches/collect_bool.rs

@@ -0,0 +1,209 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Benchmarks for `collect_bool` bit packing.
+//!
+//! Three groups, each with a scalar baseline that replicates the previous bit-at-a-time
+//! implementation. Under CodSpeed only the shipped entry points are tracked
+//! (`from_bool_slice`, `collect_bool_u32_gt`); the baselines, the dispatch loop,
+//! per-SIMD-level loops, and the portable SWAR kernel compile out and are for local A/B runs:
+//!
+//! - `pack_words_*`: the portable SWAR kernel vs the scalar loop, word by word. The SIMD
+//!   kernels are covered by `words_gather_*` instead, where they can inline.
+//! - `words_gather_*`: each multiversioned word loop with a bounds-check-free bool gather,
+//!   measuring the fused fill + pack pipeline per SIMD level.
+//! - `collect_bool_*` / `from_bool_slice`: the public entry points end to end, with a
+//!   boolean-gather predicate and a `u32` comparison predicate.
+
+use divan::Bencher;
+use vortex_buffer::BitBuffer;
+#[cfg(not(codspeed))]
+use vortex_buffer::collect_bool_word_scalar;
+#[cfg(not(codspeed))]
+use vortex_buffer::pack_bool_word_swar;
+
+fn main() {
+    // Pre-warm CPUID feature detection so the one-time probe cost is never
+    // included in any benchmark iteration.
+    #[cfg(target_arch = "x86_64")]
+    {
+        let _ = is_x86_feature_detected!("avx2");
+        let _ = is_x86_feature_detected!("avx512f");
+        let _ = is_x86_feature_detected!("avx512bw");
+    }
+
+    divan::main();
+}
+
+const INPUT_SIZE: &[usize] = &[1024, 65_536];
+
+/// Deterministic pseudo-random words (LCG), the source for all benchmark inputs.
+fn make_words(len: usize) -> impl Iterator<Item = u64> {
+    let mut state = 0x9E37_79B9_7F4A_7C15u64;
+    (0..len).map(move |_| {
+        state = state
+            .wrapping_mul(6364136223846793005)
+            .wrapping_add(1442695040888963407);
+        state
+    })
+}
+
+fn make_bools(len: usize) -> Vec<bool> {
+    make_words(len).map(|w| (w >> 33) & 1 == 1).collect()
+}
+
+fn make_u32s(len: usize) -> Vec<u32> {
+    make_words(len).map(|w| (w >> 32) as u32).collect()
+}
+
+/// Benchmark a per-word packing kernel over `len` pre-materialized bools.
+#[cfg(not(codspeed))]
+fn bench_pack_words(bencher: Bencher, len: usize, pack: impl Fn(&[bool; 64]) -> u64 + Sync) {
+    let bools = make_bools(len);
+    bencher
+        .with_inputs(|| vec![0u64; len / 64])
+        .bench_refs(|out| {
+            let (chunks, _) = bools.as_chunks::<64>();
+            for (word, chunk) in out.iter_mut().zip(chunks) {
+                *word = pack(chunk);
+            }
+        });
+}
+
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn pack_words_scalar(bencher: Bencher, len: usize) {
+    bench_pack_words(bencher, len, |chunk| {
+        collect_bool_word_scalar(64, |i| chunk[i])
+    });
+}
+
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn pack_words_swar(bencher: Bencher, len: usize) {
+    bench_pack_words(bencher, len, pack_bool_word_swar);
+}
+
+/// Benchmark a full multiversioned word loop with a bounds-check-free bool gather, measuring
+/// the packing pipeline itself (fill + pack fully inlined) rather than predicate evaluation.
+#[cfg(not(codspeed))]
+fn bench_words_gather(
+    bencher: Bencher,
+    len: usize,
+    collect: impl Fn(&mut [u64], usize, &[bool]) + Sync,
+) {
+    let bools = make_bools(len);
+    bencher
+        .with_inputs(|| vec![0u64; len.div_ceil(64)])
+        .bench_refs(|words| collect(words, len, &bools));
+}
+
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn words_gather_dispatch(bencher: Bencher, len: usize) {
+    bench_words_gather(bencher, len, |words, len, bools| {
+        // SAFETY: `collect_bool_words` invokes the predicate with indices `0..len` only.
+        vortex_buffer::collect_bool_words(words, len, |i| unsafe { *bools.get_unchecked(i) })
+    });
+}
+
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn words_gather_scalar(bencher: Bencher, len: usize) {
+    bench_words_gather(bencher, len, |words, len, bools| {
+        // SAFETY: `collect_bool_words_old` invokes the predicate with indices `0..len` only.
+        collect_bool_words_old(words, len, |i| unsafe { *bools.get_unchecked(i) })
+    });
+}
+
+#[cfg(target_arch = "x86_64")]
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn words_gather_sse2(bencher: Bencher, len: usize) {
+    bench_words_gather(bencher, len, |words, len, bools| {
+        // SAFETY: SSE2 is part of the x86-64 baseline; indices passed are `0..len`.
+        unsafe { vortex_buffer::collect_bool_words_sse2(words, len, |i| *bools.get_unchecked(i)) }
+    });
+}
+
+#[cfg(target_arch = "x86_64")]
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn words_gather_avx2(bencher: Bencher, len: usize) {
+    if !is_x86_feature_detected!("avx2") {
+        return;
+    }
+    bench_words_gather(bencher, len, |words, len, bools| {
+        // SAFETY: runtime detection guarantees AVX2; indices passed are `0..len`.
+        unsafe { vortex_buffer::collect_bool_words_avx2(words, len, |i| *bools.get_unchecked(i)) }
+    });
+}
+
+#[cfg(target_arch = "x86_64")]
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn words_gather_avx512(bencher: Bencher, len: usize) {
+    if !(is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw")) {
+        return;
+    }
+    bench_words_gather(bencher, len, |words, len, bools| {
+        // SAFETY: runtime detection guarantees AVX-512F/BW; indices passed are `0..len`.
+        unsafe { vortex_buffer::collect_bool_words_avx512(words, len, |i| *bools.get_unchecked(i)) }
+    });
+}
+
+#[cfg(target_arch = "aarch64")]
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn words_gather_neon(bencher: Bencher, len: usize) {
+    bench_words_gather(bencher, len, |words, len, bools| {
+        // SAFETY: NEON is part of the aarch64 baseline; indices passed are `0..len`.
+        unsafe { vortex_buffer::collect_bool_words_neon(words, len, |i| *bools.get_unchecked(i)) }
+    });
+}
+
+/// Faithful copy of the previous scalar-only `collect_bool_words` word loop, used as the
+/// baseline for the end-to-end comparison.
+#[cfg(not(codspeed))]
+fn collect_bool_words_old(words: &mut [u64], len: usize, mut f: impl FnMut(usize) -> bool) {
+    let full = len / 64;
+    let remainder = len % 64;
+    for word_idx in 0..full {
+        let offset = word_idx * 64;
+        words[word_idx] = collect_bool_word_scalar(64, |bit_idx| f(offset + bit_idx));
+    }
+    if remainder != 0 {
+        let offset = full * 64;
+        words[full] = collect_bool_word_scalar(remainder, |bit_idx| f(offset + bit_idx));
+    }
+}
+
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn from_bool_slice_old_scalar(bencher: Bencher, len: usize) {
+    let bools = make_bools(len);
+    bencher
+        .with_inputs(|| vec![0u64; len.div_ceil(64)])
+        .bench_refs(|words| collect_bool_words_old(words, len, |i| bools[i]));
+}
+
+#[divan::bench(args = INPUT_SIZE)]
+fn from_bool_slice(bencher: Bencher, len: usize) {
+    let bools = make_bools(len);
+    bencher.bench(|| vortex_buffer::BitBufferMut::from(bools.as_slice()));
+}
+
+#[cfg(not(codspeed))]
+#[divan::bench(args = INPUT_SIZE)]
+fn collect_bool_u32_gt_old_scalar(bencher: Bencher, len: usize) {
+    let values = make_u32s(len);
+    bencher
+        .with_inputs(|| vec![0u64; len.div_ceil(64)])
+        .bench_refs(|words| collect_bool_words_old(words, len, |i| values[i] > u32::MAX / 2));
+}
+
+#[divan::bench(args = INPUT_SIZE)]
+fn collect_bool_u32_gt(bencher: Bencher, len: usize) {
+    let values = make_u32s(len);
+    bencher.bench(|| BitBuffer::collect_bool(len, |i| values[i] > u32::MAX / 2));
+}

vortex-buffer/src/bit/buf.rs

@@ -185,11 +185,43 @@ impl BitBuffer {
     }
 
     /// Invokes `f` with indexes `0..len` collecting the boolean results into a new [`BitBuffer`].
+    ///
+    /// `f` is invoked exactly once per index, in ascending order, and the results are packed
+    /// with the baseline SIMD byte→bit instruction of the target.
+    ///
+    /// # Performance
+    ///
+    /// The packing is a few instructions per 64 bits, so evaluating `f` is usually the
+    /// bottleneck. In particular, a bounds-checked slice access in `f` (`|i| values[i] > x`)
+    /// blocks vectorization of the gather and can cost ~10x the packing itself. Since `f` only
+    /// ever sees indices `0..len`, callers reading from a slice with `len <= values.len()` may
+    /// soundly use `|i| unsafe { *values.get_unchecked(i) }`.
+    ///
+    /// Prefer this entry point for every predicate. Only switch to
+    /// [`Self::collect_bool_multiversioned`] after carefully checking that your specific `f`
+    /// meets its contract (a trivially cheap, bounds-check-free gather or comparison) —
+    /// ideally with a benchmark.
     #[inline]
     pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
         BitBufferMut::collect_bool(len, f).freeze()
     }
 
+    /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once
+    /// per CPU feature level (AVX-512BW/AVX2/baseline) and selects a clone by runtime feature
+    /// detection.
+    ///
+    /// Calling this asserts that `f` is small and simple enough (e.g. a bounds-check-free slice
+    /// gather or comparison) that duplicating it per feature level and paying a
+    /// `#[target_feature]` call boundary beats inlining it once into your function. For any
+    /// non-trivial `f` that assertion is false — the boundary deoptimizes the predicate — so
+    /// unless you have carefully checked (ideally benchmarked) that your specific `f`
+    /// qualifies, use [`Self::collect_bool`]. See
+    /// [`collect_bool_words_multiversioned`](crate::bit::collect_bool_words_multiversioned).
+    #[inline]
+    pub fn collect_bool_multiversioned<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
+        BitBufferMut::collect_bool_multiversioned(len, f).freeze()
+    }
+
     /// Maps over each bit in this buffer, calling `f(index, bit_value)` and collecting results.
     ///
     /// This is more efficient than `collect_bool` when you need to read the current bit value,