diff --git a/.gitignore b/.gitignore index bcc8ef746ee..f9613807332 100644 --- a/.gitignore +++ b/.gitignore @@ -224,6 +224,7 @@ compile_commands.json # AI Agents .agents/settings.local.json +.agents/scheduled_tasks.lock .claude/settings.local.json .opencode diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index d68154b6f2b..7f9a4e803de 100644 --- a/encodings/bytebool/src/compute.rs +++ b/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)] diff --git a/vortex-array/src/arrays/decimal/compute/between.rs b/vortex-array/src/arrays/decimal/compute/between.rs index 75c2da923ea..50623b077ee 100644 --- a/vortex-array/src/arrays/decimal/compute/between.rs +++ b/vortex-array/src/arrays/decimal/compute/between.rs @@ -132,8 +132,10 @@ fn between_impl( ) -> ArrayRef { let buffer = arr.buffer::(); 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() diff --git a/vortex-array/src/arrays/primitive/compute/between.rs b/vortex-array/src/arrays/primitive/compute/between.rs index 9a6c7f60c1b..85f2b3d42e3 100644 --- a/vortex-array/src/arrays/primitive/compute/between.rs +++ b/vortex-array/src/arrays/primitive/compute/between.rs @@ -105,7 +105,7 @@ where { let slice = arr.as_slice::(); 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) diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index ae9d7e6cc05..3d072145c8a 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -48,3 +48,7 @@ harness = false [[bench]] name = "vortex_bitbuffer" harness = false + +[[bench]] +name = "collect_bool" +harness = false diff --git a/vortex-buffer/benches/collect_bool.rs b/vortex-buffer/benches/collect_bool.rs new file mode 100644 index 00000000000..3bfe9a55b69 --- /dev/null +++ b/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 { + 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 { + make_words(len).map(|w| (w >> 33) & 1 == 1).collect() +} + +fn make_u32s(len: usize) -> Vec { + 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)); +} diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index d229329fcd2..8dc5e38a42c 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/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 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 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, diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index 1964aa664d4..b3e7207a640 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -11,6 +11,7 @@ use crate::ByteBufferMut; use crate::bit::collect_bool_words; use crate::bit::get_bit_unchecked; use crate::bit::ops; +use crate::bit::pack::collect_bool_words_multiversioned; use crate::bit::set_bit_unchecked; use crate::bit::unset_bit_unchecked; use crate::buffer_mut; @@ -184,15 +185,56 @@ impl BitBufferMut { } /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBufferMut` + /// + /// `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 bool>(len: usize, f: F) -> Self { + Self::collect_words(len, |words| collect_bool_words(words, len, f)) + } + + /// 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`]. + #[inline] + pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { + Self::collect_words(len, |words| { + collect_bool_words_multiversioned(words, len, f) + }) + } + + /// Allocate a zero-copy word buffer for `len` bits, let `fill` populate it, and wrap it as a + /// `BitBufferMut`. + #[inline] + fn collect_words(len: usize, fill: impl FnOnce(&mut [u64])) -> Self { let num_words = len.div_ceil(64); let mut buffer: BufferMut = BufferMut::with_capacity(num_words); - // SAFETY: `collect_bool_words` writes every word in `0..num_words` below - // before any read; `u64` has no invalid bit patterns and the assignments - // inside `collect_bool_words` are pure writes. + // SAFETY: `fill` (a `collect_bool_words` variant) writes every word in `0..num_words` + // below before any read; `u64` has no invalid bit patterns and the assignments inside + // `collect_bool_words` are pure writes. unsafe { buffer.set_len(num_words) }; - collect_bool_words(buffer.as_mut_slice(), len, f); + fill(buffer.as_mut_slice()); let mut bytes = buffer.into_byte_buffer(); bytes.truncate(len.div_ceil(8)); @@ -601,14 +643,23 @@ impl Not for BitBufferMut { impl From<&[bool]> for BitBufferMut { fn from(value: &[bool]) -> Self { - BitBufferMut::collect_bool(value.len(), |i| value[i]) + // SAFETY: the predicate is invoked with indices `0..value.len()` only. + // Skipping the bounds check lets the gather loop vectorize. + BitBufferMut::collect_bool_multiversioned(value.len(), |i| unsafe { + *value.get_unchecked(i) + }) } } // allow building a buffer from a set of truthy byte values. impl From<&[u8]> for BitBufferMut { fn from(value: &[u8]) -> Self { - BitBufferMut::collect_bool(value.len(), |i| value[i] > 0) + // SAFETY: the predicate is invoked with indices `0..value.len()` only. + // Skipping the bounds check lets the gather loop vectorize. + BitBufferMut::collect_bool_multiversioned( + value.len(), + |i| unsafe { *value.get_unchecked(i) } > 0, + ) } } diff --git a/vortex-buffer/src/bit/mod.rs b/vortex-buffer/src/bit/mod.rs index a39c93e3f25..46d1bee89be 100644 --- a/vortex-buffer/src/bit/mod.rs +++ b/vortex-buffer/src/bit/mod.rs @@ -14,6 +14,7 @@ mod count_ones; mod macros; mod meta; mod ops; +mod pack; mod select; mod view; @@ -27,29 +28,45 @@ pub use arrow_buffer::bit_iterator::BitSliceIterator; pub use buf::*; pub use buf_mut::*; pub use meta::*; +pub use pack::*; pub use view::*; /// Packs up to 64 boolean values into a little-endian `u64` word. +/// +/// This is [`collect_bool_words`] for a single word: a full 64-bit word is materialized as a +/// `[bool; 64]` and packed with the baseline SIMD byte→bit instruction of the target; shorter +/// lengths fall back to the bit-at-a-time [`collect_bool_word_scalar`] loop. #[inline] -pub fn collect_bool_word(len: usize, mut f: F) -> u64 +pub fn collect_bool_word(len: usize, f: F) -> u64 where F: FnMut(usize) -> bool, { assert!(len <= 64, "cannot pack {len} bits into a u64 word"); - let mut packed = 0; - for bit_idx in 0..len { - packed |= (f(bit_idx) as u64) << bit_idx; - } - packed + let mut word = [0u64; 1]; + collect_bool_words_inline(&mut word, len, f); + word[0] } /// Pack `len` boolean values returned by `f` into the prefix of `words`, LSB-first, /// 64 bits per `u64`. `words` must have capacity for at least `len.div_ceil(64)` entries. /// +/// `f` is invoked exactly once per index, in ascending order `0..len`. +/// /// Writes via `=` (not `|=`), so the destination need not be zero-initialised. +/// +/// The word loop packs with the baseline SIMD kernel of the target (SSE2 on x86-64, NEON on +/// aarch64), which inlines fully into the caller together with the predicate and the +/// `[bool; 64]` materialization — wider kernels would sit behind a non-inlinable +/// `#[target_feature]` boundary that deoptimizes expensive predicates. See +/// [`BitBuffer::collect_bool`] for the performance note on +/// avoiding bounds checks in `f`. +/// +/// Prefer this entry point for every predicate; only switch to +/// [`collect_bool_words_multiversioned`] after carefully checking that your specific `f` +/// meets its contract. #[inline] -pub fn collect_bool_words(words: &mut [u64], len: usize, mut f: F) +pub fn collect_bool_words(words: &mut [u64], len: usize, f: F) where F: FnMut(usize) -> bool, { @@ -60,18 +77,7 @@ where words.len(), ); - 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(64, |bit_idx| f(offset + bit_idx)); - } - - if remainder != 0 { - let offset = full * 64; - words[full] = collect_bool_word(remainder, |bit_idx| f(offset + bit_idx)); - } + collect_bool_words_inline(words, len, f) } /// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 diff --git a/vortex-buffer/src/bit/pack.rs b/vortex-buffer/src/bit/pack.rs new file mode 100644 index 00000000000..565ce1e8345 --- /dev/null +++ b/vortex-buffer/src/bit/pack.rs @@ -0,0 +1,521 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Kernels for packing boolean values into bitmap words. +//! +//! `collect_bool` materializes each full 64-bit chunk as a `[bool; 64]` (a loop the +//! auto-vectorizer turns into vector stores for simple predicates) and then packs the 64 bytes +//! into a single `u64` with a byte→bit kernel: +//! +//! - x86-64 AVX-512BW: one `vptestmb` produces the full 64-bit mask. +//! - x86-64 AVX2: two `vpmovmskb` halves. +//! - x86-64 SSE2 (baseline): four `pmovmskb` quarters. +//! - aarch64 NEON (baseline): per-lane `ushl` by the bit position, then an `addp` reduction tree. +//! - elsewhere (and under Miri): a branch-free SWAR multiply. +//! +//! There are two tiers. The default ([`collect_bool_words_inline`]) compiles the loop once +//! with the widest *statically-enabled* kernel (SSE2 / NEON / SWAR on stock targets) and +//! inlines fully into the caller — safe for arbitrary predicates. The opt-in tier +//! ([`collect_bool_words_multiversioned`]) compiles the loop — with `f` inside — once per CPU +//! feature level and selects a clone by runtime detection; only for predicates small and +//! simple enough that the per-level duplication and its `#[target_feature]` call boundary pay +//! off. +//! +//! The bit-at-a-time loop lives on as [`collect_bool_word_scalar`], used for tail chunks and as +//! the reference implementation for tests and benchmarks. + +/// Packs up to 64 boolean values into a little-endian `u64` word one bit at a time. +/// +/// This is the scalar reference implementation behind +/// [`collect_bool_word`](crate::bit::collect_bool_word); prefer calling that entry point, which +/// takes the SIMD fast path for full 64-bit words. +#[inline] +pub fn collect_bool_word_scalar(len: usize, mut f: F) -> u64 +where + F: FnMut(usize) -> bool, +{ + assert!(len <= 64, "cannot pack {len} bits into a u64 word"); + + let mut packed = 0; + for bit_idx in 0..len { + packed |= (f(bit_idx) as u64) << bit_idx; + } + packed +} + +/// Body of [`collect_bool_words`](crate::bit::collect_bool_words) (and, via a one-word slice, +/// of [`collect_bool_word`](crate::bit::collect_bool_word)): the word loop with the widest +/// pack kernel enabled *at compile time* — SSE2 on stock x86-64 (AVX2/AVX-512BW when built +/// with e.g. `-C target-cpu=native`), NEON on aarch64, SWAR elsewhere and under Miri. +/// +/// Statically-enabled kernels are part of every function's feature set, so this loop +/// (predicate, the `[bool; 64]` materialization, and the pack) inlines fully into the caller +/// with no `#[target_feature]` boundary. That boundary is why *runtime*-detected wider kernels +/// are not used here: hiding an expensive, non-vectorizable predicate (e.g. FSST's per-string +/// DFA scan) behind a non-inlinable AVX-512 loop copy costs far more (~30% end to end) than +/// the wider pack saves — and an indirect call per word is worse still (~4x on cheap +/// predicates), since an opaque call target blocks fill/pack fusion regardless of how cheap +/// the kernel *selection* is. For provably cheap predicates, use [`collect_bool_words_multiversioned`]. +#[inline(always)] +pub(crate) fn collect_bool_words_inline(words: &mut [u64], len: usize, f: F) +where + F: FnMut(usize) -> bool, +{ + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx512f", + target_feature = "avx512bw", + not(miri) + ))] + { + // SAFETY: AVX-512F/BW are statically enabled for this build (e.g. -C + // target-cpu=native), so they are in every function's feature set and the kernel + // inlines here like any other function. + collect_bool_words_with(words, len, f, |bools| unsafe { + pack_bool_word_avx512(bools) + }) + } + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + not(all(target_feature = "avx512f", target_feature = "avx512bw")), + not(miri) + ))] + { + // SAFETY: AVX2 is statically enabled for this build. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_avx2(bools) }) + } + #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2"), not(miri)))] + { + // SAFETY: SSE2 is part of the x86-64 baseline instruction set. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_sse2(bools) }) + } + #[cfg(all(target_arch = "aarch64", not(miri)))] + { + // SAFETY: NEON is part of the aarch64 baseline instruction set. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_neon(bools) }) + } + #[cfg(any(not(any(target_arch = "x86_64", target_arch = "aarch64")), miri))] + collect_bool_words_with(words, len, f, pack_bool_word_swar) +} + +/// Word loop with the *widest* pack kernel the CPU offers (AVX-512BW, then AVX2, then the +/// baseline), for predicates known to be cheap. +/// +/// The wide loop copies live behind a `#[target_feature]` function boundary that cannot inline +/// into the caller, which deoptimizes expensive predicates (see the module docs and +/// `collect_bool_words_inline`). Only route a predicate here when its evaluation is trivial +/// — e.g. the bounds-check-free slice gathers in the `From<&[bool]>` / `From<&[u8]>` +/// conversions, or unchecked slice comparisons like the `between` kernels — where the fused +/// AVX-512 loop is worth another ~2x over the baseline kernel. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`, +/// exactly once per index in ascending order. +/// +/// Panics if `words` is too short. +#[inline] +pub fn collect_bool_words_multiversioned(words: &mut [u64], len: usize, f: F) +where + F: FnMut(usize) -> bool, +{ + let num_words = len.div_ceil(64); + assert!( + words.len() >= num_words, + "words slice has {} entries, need at least {num_words}", + words.len(), + ); + + // Without a full 64-bit word only the scalar tail would run; skip feature detection and + // the `#[target_feature]` call boundary entirely. + if len < 64 { + return collect_bool_words_inline(words, len, f); + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + // SAFETY: runtime detection guarantees the required target features. + return unsafe { collect_bool_words_avx512(words, len, f) }; + } + if is_x86_feature_detected!("avx2") { + // SAFETY: runtime detection guarantees the required target features. + return unsafe { collect_bool_words_avx2(words, len, f) }; + } + } + collect_bool_words_inline(words, len, f) +} + +/// Shared word loop: materialize each full 64-bit chunk as a `[bool; 64]` and pack it with +/// `pack`; the tail chunk goes through the scalar loop. +/// +/// Marked `#[inline(always)]` so each `#[target_feature]` wrapper gets its own fully-inlined +/// copy compiled with that feature set. +#[inline(always)] +fn collect_bool_words_with(words: &mut [u64], len: usize, mut f: F, pack: P) +where + F: FnMut(usize) -> bool, + P: Fn(&[bool; 64]) -> u64, +{ + let full = len / 64; + let remainder = len % 64; + + for (word_idx, word) in words[..full].iter_mut().enumerate() { + let offset = word_idx * 64; + let mut bools = [false; 64]; + for (bit_idx, b) in bools.iter_mut().enumerate() { + *b = f(offset + bit_idx); + } + *word = pack(&bools); + } + + if remainder != 0 { + let offset = full * 64; + words[full] = collect_bool_word_scalar(remainder, |bit_idx| f(offset + bit_idx)); + } +} + +/// SSE2 copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports SSE2 (always true on x86-64). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "sse2")] +pub unsafe fn collect_bool_words_sse2 bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees SSE2 support. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_sse2(bools) }) +} + +/// AVX2 copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX2. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +pub unsafe fn collect_bool_words_avx2 bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees AVX2 support. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_avx2(bools) }) +} + +/// AVX-512BW copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX-512F and AVX-512BW. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512bw")] +pub unsafe fn collect_bool_words_avx512 bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees AVX-512F and AVX-512BW support. + collect_bool_words_with(words, len, f, |bools| unsafe { + pack_bool_word_avx512(bools) + }) +} + +/// NEON copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports NEON (always true on aarch64). +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +pub unsafe fn collect_bool_words_neon bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees NEON support. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_neon(bools) }) +} + +/// Portable branch-free byte→bit pack kernel, used when no SIMD kernel is available. +/// +/// Reads the bools eight at a time as a `u64` and gathers the eight `0x00`/`0x01` bytes into +/// eight contiguous bits with a single multiply: byte `i` contributes `2^(8i)`, and the magic +/// constant `0x0102_0408_1020_4080 = Σ 2^(56-7i)` shifts each contribution to bit `56 + i` +/// without any cross-term collisions, so the mask falls out of the top byte of the product. +#[inline] +pub fn pack_bool_word_swar(bools: &[bool; 64]) -> u64 { + const MAGIC: u64 = 0x0102_0408_1020_4080; + + let (chunks, rest) = bools.as_chunks::<8>(); + debug_assert!(rest.is_empty()); + + let mut packed = 0u64; + for (chunk_idx, chunk) in chunks.iter().enumerate() { + let word = u64::from_le_bytes(chunk.map(|b| b as u8)); + packed |= (word.wrapping_mul(MAGIC) >> 56) << (8 * chunk_idx); + } + packed +} + +/// SSE2 byte→bit pack kernel: four 16-byte `pcmpeqb`-against-zero + `pmovmskb` rounds. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports SSE2 (always true on x86-64). +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "sse2")] +pub unsafe fn pack_bool_word_sse2(bools: &[bool; 64]) -> u64 { + use std::arch::x86_64::__m128i; + use std::arch::x86_64::_mm_cmpeq_epi8; + use std::arch::x86_64::_mm_loadu_si128; + use std::arch::x86_64::_mm_movemask_epi8; + use std::arch::x86_64::_mm_setzero_si128; + + let ptr = bools.as_ptr().cast::(); + let zero = _mm_setzero_si128(); + + let mut packed = 0u64; + for lane in 0..4 { + // SAFETY: `lane * 16 + 16 <= 64`, so the 16-byte load is in bounds. + let chunk = unsafe { _mm_loadu_si128(ptr.add(lane * 16).cast::<__m128i>()) }; + // `cmpeq` against zero sets 0xFF for *false* bytes; invert to get the truthy mask. + let zero_mask = _mm_movemask_epi8(_mm_cmpeq_epi8(chunk, zero)) as u32 as u64; + packed |= (!zero_mask & 0xFFFF) << (16 * lane); + } + packed +} + +/// AVX2 byte→bit pack kernel: two 32-byte `vpcmpeqb`-against-zero + `vpmovmskb` rounds. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX2. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx2")] +pub unsafe fn pack_bool_word_avx2(bools: &[bool; 64]) -> u64 { + use std::arch::x86_64::__m256i; + use std::arch::x86_64::_mm256_cmpeq_epi8; + use std::arch::x86_64::_mm256_loadu_si256; + use std::arch::x86_64::_mm256_movemask_epi8; + use std::arch::x86_64::_mm256_setzero_si256; + + let ptr = bools.as_ptr().cast::(); + let zero = _mm256_setzero_si256(); + + // SAFETY: both 32-byte loads are within the 64-byte array. + let lo = unsafe { _mm256_loadu_si256(ptr.cast::<__m256i>()) }; + // SAFETY: see above. + let hi = unsafe { _mm256_loadu_si256(ptr.add(32).cast::<__m256i>()) }; + + // `cmpeq` against zero sets 0xFF for *false* bytes; invert to get the truthy mask. + let lo_mask = !(_mm256_movemask_epi8(_mm256_cmpeq_epi8(lo, zero)) as u32) as u64; + let hi_mask = !(_mm256_movemask_epi8(_mm256_cmpeq_epi8(hi, zero)) as u32) as u64; + lo_mask | (hi_mask << 32) +} + +/// AVX-512BW byte→bit pack kernel: a single 64-byte `vptestmb` produces the whole word. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX-512F and AVX-512BW. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx512f,avx512bw")] +pub unsafe fn pack_bool_word_avx512(bools: &[bool; 64]) -> u64 { + use std::arch::x86_64::__m512i; + use std::arch::x86_64::_mm512_loadu_si512; + use std::arch::x86_64::_mm512_test_epi8_mask; + + // SAFETY: the 64-byte load covers exactly the `[bool; 64]` array. + let chunk = unsafe { _mm512_loadu_si512(bools.as_ptr().cast::<__m512i>()) }; + // Mask bit `i` is set iff byte `i` AND byte `i` is nonzero, i.e. iff `bools[i]`. + _mm512_test_epi8_mask(chunk, chunk) +} + +/// NEON byte→bit pack kernel: shift each `0x00`/`0x01` byte left by its bit position +/// (`ushl`), then fold the four vectors into one `u64` with a pairwise-add (`addp`) tree. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports NEON (always true on aarch64). +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +pub unsafe fn pack_bool_word_neon(bools: &[bool; 64]) -> u64 { + use std::arch::aarch64::vgetq_lane_u64; + use std::arch::aarch64::vld1q_s8; + use std::arch::aarch64::vld1q_u8; + use std::arch::aarch64::vpaddq_u8; + use std::arch::aarch64::vreinterpretq_u64_u8; + use std::arch::aarch64::vshlq_u8; + + const BIT_SHIFTS: [i8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7]; + + let ptr = bools.as_ptr().cast::(); + // SAFETY: loading 16 constant bytes from `BIT_SHIFTS`; the four 16-byte data loads below are + // all within the 64-byte array. + unsafe { + let shifts = vld1q_s8(BIT_SHIFTS.as_ptr()); + + // Byte j of each vector becomes `bools[16v + j] << (j % 8)`. + let m0 = vshlq_u8(vld1q_u8(ptr), shifts); + let m1 = vshlq_u8(vld1q_u8(ptr.add(16)), shifts); + let m2 = vshlq_u8(vld1q_u8(ptr.add(32)), shifts); + let m3 = vshlq_u8(vld1q_u8(ptr.add(48)), shifts); + + // Three rounds of pairwise adds sum each group of 8 weighted bytes into one mask byte, + // yielding the 8 mask bytes in order in the low half of the final vector. + let sum01 = vpaddq_u8(m0, m1); + let sum23 = vpaddq_u8(m2, m3); + let sum = vpaddq_u8(sum01, sum23); + let sum = vpaddq_u8(sum, sum); + vgetq_lane_u64::<0>(vreinterpretq_u64_u8(sum)) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::collect_bool_word_scalar; + use super::pack_bool_word_swar; + + fn patterns() -> Vec<[bool; 64]> { + let mut patterns = vec![ + [false; 64], + [true; 64], + std::array::from_fn(|i| i % 2 == 0), + std::array::from_fn(|i| i % 3 == 0), + std::array::from_fn(|i| i < 32), + std::array::from_fn(|i| i == 0 || i == 63), + ]; + // A few deterministic pseudo-random patterns. + let mut state = 0x9E37_79B9_7F4A_7C15u64; + for _ in 0..8 { + patterns.push(std::array::from_fn(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) & 1 == 1 + })); + } + patterns + } + + fn reference(bools: &[bool; 64]) -> u64 { + collect_bool_word_scalar(64, |i| bools[i]) + } + + #[test] + fn swar_matches_scalar() { + for bools in patterns() { + assert_eq!(pack_bool_word_swar(&bools), reference(&bools)); + } + } + + #[test] + fn dispatch_matches_scalar() { + for bools in patterns() { + assert_eq!( + crate::bit::collect_bool_word(64, |i| bools[i]), + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + #[test] + fn sse2_matches_scalar() { + for bools in patterns() { + // SAFETY: SSE2 is part of the x86-64 baseline instruction set. + assert_eq!( + unsafe { super::pack_bool_word_sse2(&bools) }, + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + #[test] + fn avx2_matches_scalar() { + if !is_x86_feature_detected!("avx2") { + return; + } + for bools in patterns() { + // SAFETY: runtime detection guarantees AVX2. + assert_eq!( + unsafe { super::pack_bool_word_avx2(&bools) }, + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + #[test] + fn avx512_matches_scalar() { + if !(is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw")) { + return; + } + for bools in patterns() { + // SAFETY: runtime detection guarantees AVX-512F and AVX-512BW. + assert_eq!( + unsafe { super::pack_bool_word_avx512(&bools) }, + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "aarch64", not(miri)))] + #[test] + fn neon_matches_scalar() { + for bools in patterns() { + // SAFETY: NEON is part of the aarch64 baseline instruction set. + assert_eq!( + unsafe { super::pack_bool_word_neon(&bools) }, + reference(&bools) + ); + } + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(63)] + #[case(64)] + #[case(65)] + #[case(200)] + fn multiversioned_matches_inline(#[case] len: usize) { + let pattern = |i: usize| i.is_multiple_of(3) || i.is_multiple_of(7); + let num_words = len.div_ceil(64); + let mut multiversioned = vec![0u64; num_words]; + super::collect_bool_words_multiversioned(&mut multiversioned, len, pattern); + let mut inline = vec![0u64; num_words]; + super::collect_bool_words_inline(&mut inline, len, pattern); + assert_eq!(multiversioned, inline); + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(5)] + #[case(63)] + #[case(64)] + fn collect_bool_word_partial_lens_match(#[case] len: usize) { + let expected = collect_bool_word_scalar(len, |i| i % 3 == 0); + assert_eq!(crate::bit::collect_bool_word(len, |i| i % 3 == 0), expected); + } +}