From b98976fb88a2d8cb7de0990ff522176f854eedc5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 15 Oct 2026 16:06:56 +0100 Subject: [PATCH 1/4] Move generic simd filtering to vortex-buffer Signed-off-by: Robert Kruszewski --- Cargo.lock | 2 - .../src/bitpacking/compute/filter.rs | 13 +- vortex-array/Cargo.toml | 2 - .../src/arrays/filter/execute/buffer.rs | 126 +------ .../arrays/filter/execute/byte_compress.rs | 347 ------------------ vortex-array/src/arrays/filter/execute/mod.rs | 7 +- .../src/arrays/filter/execute/slice.rs | 272 -------------- vortex-array/src/arrays/filter/mod.rs | 3 - vortex-buffer/src/filter/bits.rs | 84 +++++ vortex-buffer/src/filter/byte_compress.rs | 108 ++++++ .../src/filter}/chunked.rs | 189 ++++------ vortex-buffer/src/filter/mod.rs | 329 +++++++++++++++++ vortex-buffer/src/filter/scalar.rs | 177 +++++++++ .../src/filter/simd}/generic.rs | 14 +- .../src/filter/simd}/mod.rs | 91 ++--- .../src/filter/simd}/neon.rs | 14 +- .../src/filter/simd}/tests.rs | 153 ++++---- .../src/filter/simd}/x86.rs | 26 +- vortex-buffer/src/filter/tests.rs | 184 ++++++++++ vortex-buffer/src/lib.rs | 1 + vortex-mask/src/lib.rs | 41 +++ vortex-mask/src/tests.rs | 47 +++ 22 files changed, 1191 insertions(+), 1039 deletions(-) delete mode 100644 vortex-array/src/arrays/filter/execute/byte_compress.rs delete mode 100644 vortex-array/src/arrays/filter/execute/slice.rs create mode 100644 vortex-buffer/src/filter/bits.rs create mode 100644 vortex-buffer/src/filter/byte_compress.rs rename {vortex-array/src/arrays/filter/execute => vortex-buffer/src/filter}/chunked.rs (75%) create mode 100644 vortex-buffer/src/filter/mod.rs create mode 100644 vortex-buffer/src/filter/scalar.rs rename {vortex-array/src/arrays/filter/execute/simd_compress => vortex-buffer/src/filter/simd}/generic.rs (90%) rename {vortex-array/src/arrays/filter/execute/simd_compress => vortex-buffer/src/filter/simd}/mod.rs (75%) rename {vortex-array/src/arrays/filter/execute/simd_compress => vortex-buffer/src/filter/simd}/neon.rs (92%) rename {vortex-array/src/arrays/filter/execute/simd_compress => vortex-buffer/src/filter/simd}/tests.rs (54%) rename {vortex-array/src/arrays/filter/execute/simd_compress => vortex-buffer/src/filter/simd}/x86.rs (94%) create mode 100644 vortex-buffer/src/filter/tests.rs diff --git a/Cargo.lock b/Cargo.lock index c16d8af9205..8072d44bcd5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11393,8 +11393,6 @@ dependencies = [ "cfg-if", "codspeed-divan-compat", "enum-iterator", - "fearless_simd", - "fearless_simd_macros", "flatbuffers", "futures", "goldenfile", diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index d2b7245adb6..4981dd8b2ac 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -9,14 +9,14 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::filter::ChunkDecoder; -use vortex_array::arrays::filter::FILTER_CHUNK_LEN; use vortex_array::arrays::filter::FilterKernel; -use vortex_array::arrays::filter::filter_chunked; use vortex_array::dtype::NativePType; use vortex_array::match_each_unsigned_integer_ptype; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::filter::ChunkDecoder; +use vortex_buffer::filter::FILTER_CHUNK_LEN; +use vortex_buffer::filter::filter_chunked; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_mask::MaskValues; @@ -87,7 +87,12 @@ fn filter_values( bit_width, packed_chunk_len: FL_CHUNK_SIZE * bit_width / (8 * size_of::()), }; - filter_chunked(&chunks, array.offset() as usize, mask, allocator) + filter_chunked( + &chunks, + array.offset() as usize, + mask.selection(), + allocator, + ) } /// The packed FastLanes chunks of a bit-packed array. diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 10733c48d8b..dbaf32a1ed9 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -29,8 +29,6 @@ arrow-buffer = { workspace = true } bytes = { workspace = true } cfg-if = { workspace = true } enum-iterator = { workspace = true } -fearless_simd = { workspace = true } -fearless_simd_macros = { workspace = true } flatbuffers = { workspace = true } futures = { workspace = true, features = ["alloc", "async-await", "std"] } goldenfile = { workspace = true, optional = true } diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index dfd9c76af9c..12856a3f010 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -1,36 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Buffer-level filter dispatch across cached, SIMD, and scalar strategies. +//! Buffer-level filtering by a [`MaskValues`]. //! -//! Selection is an ordered ladder, not a set of independent choices. The first matching row wins: -//! -//! | Priority | Condition | Out-of-place | In-place | -//! | -------- | --------------------------------------------------- | -------------- | ------------- | -//! | 1 | cached slices: one run, or average run length >= 8 | copy runs | move runs | -//! | 2 | cached indices and density <= 0.50 | gather indices | move indices | -//! | 3 | an eligible architecture kernel | SIMD compress | SIMD compress | -//! | 4 | density reaches [`byte_compress_density_threshold`] | byte compress | bitmap walk | -//! | 5 | otherwise | bitmap walk | bitmap walk | -//! -//! In-place dispatch is attempted only for a uniquely owned buffer with density >= 0.50. SIMD -//! eligibility is summarized in [`simd_compress`]. All-true, all-false, and contiguous masks are -//! handled before buffer dispatch. - -use std::mem::size_of; +//! The strategies live in [`vortex_buffer::filter`], and [`MaskValues::selection`] chooses the +//! mask representation that they use. All-true, all-false, and contiguous masks are handled +//! before buffer filtering. use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::filter; use vortex_mask::MaskValues; -use crate::arrays::filter::execute::byte_compress; -use crate::arrays::filter::execute::simd_compress; -use crate::arrays::filter::execute::slice; - -const CACHED_INDICES_MAX_DENSITY: f64 = 0.5; -const IN_PLACE_MIN_DENSITY: f64 = 0.5; -const MIN_SLICES_AVERAGE_RUN_LENGTH: usize = 8; - /// Filter a [`Buffer`] by [`MaskValues`], returning a new buffer. /// /// Dense uniquely owned buffers are compacted in place; other buffers allocate a new output. @@ -39,102 +20,7 @@ pub(crate) fn filter_buffer( mask: &MaskValues, allocator: &BufferAllocatorRef, ) -> Buffer { - assert_eq!(buffer.len(), mask.len()); - - let buffer = if mask.density() >= IN_PLACE_MIN_DENSITY { - match buffer.try_into_mut() { - Ok(mut buffer_mut) => { - let new_len = filter_slice_in_place(buffer_mut.as_mut_slice(), mask); - buffer_mut.truncate(new_len); - return buffer_mut.freeze(); - } - Err(buffer) => buffer, - } - } else { - buffer - }; - - filter_slice(buffer.as_slice(), mask, allocator) -} - -fn filter_slice( - values: &[T], - mask: &MaskValues, - allocator: &BufferAllocatorRef, -) -> Buffer { - if let Some(slices) = useful_cached_slices(mask) { - return slice::filter_slice_by_slices(values, slices, mask.true_count(), allocator); - } - - if let Some(indices) = useful_cached_indices(mask) { - return slice::filter_slice_by_indices(values, indices, allocator); - } - - if let Some(filtered) = simd_compress::filter_slice_by_bitmap(values, mask, allocator) { - return filtered; - } - - if mask.density() >= byte_compress_density_threshold::() { - byte_compress::filter_buffer(values, mask, allocator) - } else { - slice::filter_slice_by_bitmap(values, mask, allocator) - } -} - -fn filter_slice_in_place(values: &mut [T], mask: &MaskValues) -> usize { - if let Some(slices) = useful_cached_slices(mask) { - return slice::filter_slice_mut_by_slices(values, slices); - } - - if let Some(indices) = useful_cached_indices(mask) { - return slice::filter_slice_mut_by_indices(values, indices); - } - - if let Some(new_len) = simd_compress::filter_slice_mut_by_bitmap(values, mask) { - return new_len; - } - - slice::filter_slice_mut_by_bitmap(values, mask) -} - -pub(super) fn useful_cached_indices(mask: &MaskValues) -> Option<&[usize]> { - if mask.density() <= CACHED_INDICES_MAX_DENSITY { - mask.cached_indices() - } else { - None - } -} - -pub(super) fn useful_cached_slices(mask: &MaskValues) -> Option<&[(usize, usize)]> { - mask.cached_slices().filter(|slices| { - slices.len() == 1 || mask.true_count() / slices.len() >= MIN_SLICES_AVERAGE_RUN_LENGTH - }) -} - -pub(super) fn byte_compress_density_threshold() -> f64 { - let width = size_of::(); - - // A density at or above the table entry selects byte compress after the higher-priority - // strategies have declined the mask. These crossovers are benchmarked in - // `benches/filter_fixed_width.rs`. - // - // | Target | 1 byte | 2 bytes | 4 bytes | 8 bytes | other | - // | ------- | -----: | ------: | ------: | ------: | ----: | - // | aarch64 | 0.90 | 0.90 | 0.90 | 0.75 | 0.90 | - // | other | 0.00 | 0.50 | 0.50 | 0.75 | 0.875 | - if cfg!(target_arch = "aarch64") { - return match width { - 8 => 0.75, - _ => 0.9, - }; - } - - match width { - 1 => 0.0, - 2 | 4 => 0.5, - 8 => 0.75, - _ => 0.875, - } + filter::filter_buffer(buffer, mask.selection(), allocator) } /// Materialize sparse indices when enough sibling arrays will reuse the same mask. diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs deleted file mode 100644 index ae6345b9a83..00000000000 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ /dev/null @@ -1,347 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Byte-level compress for fixed-width filtering using a `1 << 8 = 256`-entry lookup table. -//! -//! For each byte of the mask (8 bits -> 8 source elements), a precomputed -//! permutation table compacts the selected bytes in a single indexed copy, -//! avoiding the overhead of materializing indices or slices. - -use std::mem::MaybeUninit; - -use vortex_buffer::Alignment; -use vortex_buffer::Buffer; -use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_mask::MaskValues; - -/// For each mask byte (0..256), stores the element indices to keep and the count. -/// -/// `BYTE_COMPRESS_LUT[mask_byte]` = `([i0, i1, ..., i7], popcount)` where -/// `i0..i_{popcount-1}` are the positions of set bits in `mask_byte`. -/// -/// Total size: 256 * 9 = 2304 bytes, which trivially fits in L1 cache. -static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{ - let mut lut = [([0u8; 8], 0u8); 256]; - let mut mask: usize = 0; - while mask < 256 { - let mut indices = [0u8; 8]; - let mut count: u8 = 0; - let mut bit: u8 = 0; - while bit < 8 { - if mask & (1 << bit) != 0 { - indices[count as usize] = bit; - count += 1; - } - bit += 1; - } - lut[mask] = (indices, count); - mask += 1; - } - lut -}; - -/// Filter a `Buffer` using the byte-compress LUT. -/// -/// Processes the mask one byte at a time (8 source elements per byte), -/// using a precomputed permutation to compact selected elements. -pub(crate) fn filter_buffer( - buffer: impl AsRef<[T]>, - mask: &MaskValues, - allocator: &BufferAllocatorRef, -) -> Buffer { - let src = buffer.as_ref(); - debug_assert_eq!(src.len(), mask.len()); - - let true_count = mask.true_count(); - - if true_count == 0 { - return BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()).freeze(); - } - - let mask_buffer = mask.bit_buffer(); - let mask_bytes = mask_buffer.inner().as_ref(); - let mask_offset = mask_buffer.offset(); - - filter_bitpacked(src, mask_bytes, mask_offset, true_count, allocator) -} - -fn filter_bitpacked( - src: &[T], - mask_bytes: &[u8], - mask_offset: usize, - true_count: usize, - allocator: &BufferAllocatorRef, -) -> Buffer { - let mut out = BufferMut::::with_capacity_in(true_count, allocator.clone()); - let mut write_pos: usize = 0; - - if mask_offset == 0 { - filter_aligned_into(src, mask_bytes, &mut out, &mut write_pos); - } else { - let head_len = (8 - mask_offset).min(src.len()); - let head_mask = (mask_bytes[0] >> mask_offset) & low_bits_mask(head_len); - filter_chunk_into(&src[..head_len], head_mask, &mut out, &mut write_pos); - filter_aligned_into(&src[head_len..], &mask_bytes[1..], &mut out, &mut write_pos); - } - - debug_assert_eq!(write_pos, true_count); - // SAFETY: we wrote exactly true_count elements. - unsafe { out.set_len(true_count) }; - out.freeze() -} - -/// Aligned fast path: mask bits start at byte boundary. -fn filter_aligned_into( - src: &[T], - mask_bytes: &[u8], - out: &mut BufferMut, - write_pos: &mut usize, -) { - let full_bytes = src.len() / 8; - let remainder = src.len() % 8; - - for i in 0..full_bytes { - let m = mask_bytes[i]; - if m == 0 { - continue; - } - let chunk = &src[i * 8..i * 8 + 8]; - filter_chunk_into(chunk, m, out, write_pos); - } - - // Handle the final partial chunk. - if remainder > 0 { - let m = mask_bytes[full_bytes] & low_bits_mask(remainder); - if m != 0 { - let chunk = &src[full_bytes * 8..]; - filter_chunk_into(chunk, m, out, write_pos); - } - } -} - -fn filter_chunk_into( - chunk: &[T], - mask_byte: u8, - out: &mut BufferMut, - write_pos: &mut usize, -) { - debug_assert_eq!(mask_byte & !low_bits_mask(chunk.len()), 0); - // SAFETY: the mask byte selects only elements of `chunk`, and the output has capacity for - // every selected element. - *write_pos += unsafe { - compress_byte( - chunk, - mask_byte, - out.spare_capacity_mut().get_unchecked_mut(*write_pos..), - ) - }; -} - -/// Copy the elements of `src` selected by the first `src.len()` bits of `words` to `dst` and -/// return the number copied. -/// -/// # Safety -/// -/// `words` must hold at least `src.len()` bits with any bits past `src.len()` cleared, and `dst` -/// must hold at least one value for each selected element. -pub(super) unsafe fn compress_words( - src: &[T], - words: &[u64], - dst: &mut [MaybeUninit], -) -> usize { - let mut written = 0; - for (word, src) in words.iter().zip(src.chunks(64)) { - let (chunks, tail) = src.as_chunks::<8>(); - let mask_bytes = word.to_le_bytes(); - - for (chunk, &mask_byte) in chunks.iter().zip(&mask_bytes) { - // SAFETY: the mask byte selects only elements of `chunk`, and `dst` has room for - // every selected element. - written += unsafe { compress_byte(chunk, mask_byte, dst.get_unchecked_mut(written..)) }; - } - - if !tail.is_empty() { - // SAFETY: the bits past `src.len()` are cleared, so the mask byte selects only - // elements of `tail`. - written += unsafe { - compress_byte( - tail, - mask_bytes[chunks.len()], - dst.get_unchecked_mut(written..), - ) - }; - } - } - written -} - -/// Copy the elements of `chunk` selected by `mask_byte` to `dst` and return the number copied. -/// -/// # Safety -/// -/// `mask_byte` must select only elements of `chunk`, and `dst` must hold at least one value for -/// each selected element. -#[inline] -unsafe fn compress_byte(chunk: &[T], mask_byte: u8, dst: &mut [MaybeUninit]) -> usize { - if mask_byte == 0 { - return 0; - } - - if chunk.len() == 8 && mask_byte == 0xFF { - // All 8 selected, so bulk copy. - dst[..8].write_copy_of_slice(chunk); - return 8; - } - - let (perm, count) = &BYTE_COMPRESS_LUT[mask_byte as usize]; - let count = *count as usize; - // SAFETY: perm indices are all below `chunk.len()`, and `dst` has room for `count` values. - unsafe { - for j in 0..count { - dst.get_unchecked_mut(j) - .write(*chunk.get_unchecked(*perm.get_unchecked(j) as usize)); - } - } - count -} - -fn low_bits_mask(bits: usize) -> u8 { - debug_assert!(bits <= 8); - if bits == 8 { - u8::MAX - } else { - (1u8 << bits) - 1 - } -} - -#[cfg(test)] -#[allow(clippy::unwrap_used, clippy::cast_possible_truncation)] -mod tests { - use vortex_buffer::BitBuffer; - use vortex_buffer::buffer; - use vortex_mask::Mask; - - use super::*; - - fn filter_buffer(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer { - super::filter_buffer(buffer, mask, &BufferAllocatorRef::statically_allocated()) - } - - fn mask_values(mask: &Mask) -> &MaskValues { - match mask { - Mask::Values(v) => v.as_ref(), - _ => panic!("expected Mask::Values"), - } - } - - fn offset_mask(values: [bool; N], offset: usize) -> Mask { - let bit_buffer = - BitBuffer::from_iter(std::iter::repeat_n(false, offset).chain(values.iter().copied())); - Mask::from_buffer(BitBuffer::new_with_offset( - bit_buffer.inner().clone(), - values.len(), - offset, - )) - } - - #[test] - fn test_filter_all_selected() { - let buf = buffer![1u8, 2, 3, 4, 5, 6, 7, 8, 9]; - let mask = Mask::from_iter([true, true, true, true, true, true, true, true, false]); - let result = filter_buffer(buf, mask_values(&mask)); - assert_eq!(result, buffer![1u8, 2, 3, 4, 5, 6, 7, 8]); - } - - #[test] - fn test_filter_mostly_false() { - let buf = buffer![1u8, 2, 3, 4, 5, 6, 7, 8, 9]; - let mask = Mask::from_iter([false, false, false, false, false, false, false, false, true]); - let result = filter_buffer(buf, mask_values(&mask)); - assert_eq!(result, buffer![9u8]); - } - - #[test] - fn test_filter_alternating() { - let buf = buffer![10u8, 20, 30, 40, 50, 60, 70, 80]; - let mask = Mask::from_iter([true, false, true, false, true, false, true, false]); - let result = filter_buffer(buf, mask_values(&mask)); - assert_eq!(result, buffer![10u8, 30, 50, 70]); - } - - #[test] - fn test_filter_with_remainder() { - let buf = buffer![1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10]; - let mask = Mask::from_iter([ - true, false, true, false, true, false, true, false, true, true, - ]); - let result = filter_buffer(buf, mask_values(&mask)); - assert_eq!(result, buffer![1u8, 3, 5, 7, 9, 10]); - } - - #[test] - fn test_filter_large() -> vortex_error::VortexResult<()> { - let data: Vec = (0..1000).map(|i| (i % 256) as u8).collect(); - let buf = Buffer::from(BufferMut::from_iter(data.iter().copied())); - let mask = Mask::from_iter((0..1000).map(|i| i % 3 == 0)); - let result = filter_buffer(buf, mask_values(&mask)); - let expected: Vec = data.iter().copied().step_by(3).collect(); - assert_eq!(result.as_slice(), &expected[..]); - Ok(()) - } - - #[test] - fn test_filter_signed_and_wider_integers() { - let mask = Mask::from_iter([true, false, true, true, false, true, false, true, true]); - - let i8_result = filter_buffer( - buffer![-5i8, -4, -3, -2, -1, 0, 1, 2, 3], - mask_values(&mask), - ); - assert_eq!(i8_result, buffer![-5i8, -3, -2, 0, 2, 3]); - - let u16_result = filter_buffer( - buffer![10u16, 20, 30, 40, 50, 60, 70, 80, 90], - mask_values(&mask), - ); - assert_eq!(u16_result, buffer![10u16, 30, 40, 60, 80, 90]); - - let i32_result = filter_buffer( - buffer![-100i32, -50, 0, 50, 100, 150, 200, 250, 300], - mask_values(&mask), - ); - assert_eq!(i32_result, buffer![-100i32, 0, 50, 150, 250, 300]); - } - - #[test] - fn test_filter_unaligned_byte_mask() { - let mask = offset_mask( - [ - false, false, true, false, false, false, false, false, true, false, false, - ], - 3, - ); - - let result = filter_buffer( - buffer![1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], - mask_values(&mask), - ); - assert_eq!(result, buffer![3u8, 9]); - } - - #[test] - fn test_filter_unaligned_wide_mask() { - let mask = offset_mask( - [ - true, false, true, true, false, true, false, true, true, false, false, true, - ], - 5, - ); - - let result = filter_buffer( - buffer![10u16, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120], - mask_values(&mask), - ); - assert_eq!(result, buffer![10u16, 30, 40, 60, 80, 90, 120]); - } -} diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 2ed36ab128e..8cc7653ec93 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -15,7 +15,8 @@ //! | mask selects one contiguous run | zero-copy slice | //! | child is entirely null | null constant array | //! -//! Fixed-width canonical arrays then use the strategy ladder in [`buffer`]. +//! Fixed-width canonical arrays then use the strategy ladder in [`vortex_buffer::filter`], with +//! the mask representation that [`buffer`] chooses. use std::ops::Range; @@ -48,12 +49,8 @@ use crate::validity::Validity; mod bitbuffer; mod bool; pub(crate) mod buffer; -pub(crate) mod byte_compress; -pub(crate) mod chunked; mod fixed_size_list; mod listview; -mod simd_compress; -mod slice; mod struct_; mod take; mod union; diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs deleted file mode 100644 index f603b512928..00000000000 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ /dev/null @@ -1,272 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Core slice-level filtering algorithms. -//! -//! Provides both immutable and mutable (in-place) filtering of typed slices by cached mask -//! representations or directly from the mask bitmap. - -use std::mem::MaybeUninit; - -use vortex_buffer::Buffer; -use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_mask::MaskValues; - -/// The mask bits that a word walker visits. -#[derive(Clone, Copy)] -pub(super) enum MaskBits<'a> { - /// The whole bitmap of a mask. - Mask(&'a MaskValues), - /// The first `len` bits of `words`, with any bits past `len` cleared. - Words { words: &'a [u64], len: usize }, -} - -/// Invoke `f` with each `(word, word_start, word_len)` of `bits`, where `word` holds the mask -/// bits for elements `word_start..word_start + word_len` in its low `word_len` bits. -#[inline] -pub(super) fn for_each_mask_word(bits: MaskBits<'_>, f: impl FnMut(u64, usize, usize)) { - match bits { - MaskBits::Mask(mask) => for_each_bitmap_word(mask, f), - MaskBits::Words { words, len } => for_each_slice_word(words, len, f), - } -} - -#[inline] -fn for_each_slice_word(words: &[u64], len: usize, mut f: impl FnMut(u64, usize, usize)) { - debug_assert_eq!(words.len(), len.div_ceil(64)); - - for (word_idx, &word) in words.iter().enumerate() { - let word_start = word_idx * 64; - f(word, word_start, (len - word_start).min(64)); - } -} - -#[inline] -fn for_each_bitmap_word(mask: &MaskValues, mut f: impl FnMut(u64, usize, usize)) { - let bits = mask.bit_buffer(); - let unaligned = bits.unaligned_chunks(); - let lead = unaligned.lead_padding(); - let mut base = 0; - - if let Some(prefix) = unaligned.prefix() { - let len = (64 - lead).min(mask.len()); - f(prefix >> lead, base, len); - base += len; - } - - for &word in unaligned.chunks() { - f(word, base, 64); - base += 64; - } - - if let Some(suffix) = unaligned.suffix() { - let len = mask.len() - base; - f(suffix, base, len); - base += len; - } - - debug_assert_eq!(base, mask.len()); -} - -/// A `u64` with the low `len` bits set. -#[inline] -pub(super) fn low_bits_mask(len: usize) -> u64 { - debug_assert!(len <= 64); - if len == 64 { - u64::MAX - } else { - (1u64 << len) - 1 - } -} - -/// Filter a slice from the mask bitmap without materializing indices or ranges. -pub(super) fn filter_slice_by_bitmap( - slice: &[T], - mask: &MaskValues, - allocator: &BufferAllocatorRef, -) -> Buffer { - assert_eq!( - mask.len(), - slice.len(), - "Selection mask length must equal the buffer length" - ); - - let output_len = mask.true_count(); - let mut out = BufferMut::::with_capacity_in(output_len, allocator.clone()); - // SAFETY: the mask selects only elements of `slice`, and the output was allocated for exactly - // `mask.true_count()` values. - let written = - unsafe { compact_by_bitmap(slice, MaskBits::Mask(mask), out.spare_capacity_mut()) }; - - debug_assert_eq!(written, output_len); - // SAFETY: every output slot was initialized exactly once above. - unsafe { out.set_len(output_len) }; - out.freeze() -} - -/// Copy the elements of `src` selected by `bits` to `dst` and return the number copied. -/// -/// # Safety -/// -/// `bits` must select only elements of `src`, and `dst` must hold at least one value for each -/// selected element. -#[inline] -pub(super) unsafe fn compact_by_bitmap( - src: &[T], - bits: MaskBits<'_>, - dst: &mut [MaybeUninit], -) -> usize { - let mut write_pos = 0; - - for_each_mask_word(bits, |word, word_start, word_len| { - let all_selected = low_bits_mask(word_len); - debug_assert_eq!(word & !all_selected, 0); - if word == all_selected { - dst[write_pos..][..word_len].write_copy_of_slice(&src[word_start..][..word_len]); - write_pos += word_len; - } else { - let mut selected = word; - while selected != 0 { - let index = word_start + selected.trailing_zeros() as usize; - // SAFETY: set bits index into `src`, and `dst` has room for every selected - // element. - unsafe { - dst.get_unchecked_mut(write_pos) - .write(*src.get_unchecked(index)) - }; - write_pos += 1; - selected &= selected - 1; - } - } - }); - - write_pos -} - -/// Copy the elements of `src` selected by `bits` to `dst` one run of set bits at a time, and -/// return the number copied. -/// -/// This is faster than [`compact_by_bitmap`] when the selected elements form long runs. -// Inlining this walk next to `compact_by_bitmap` makes the bit walk 1.5x slower on Apple M4. -#[inline(never)] -pub(super) fn compact_runs_by_bitmap( - src: &[T], - bits: MaskBits<'_>, - dst: &mut [MaybeUninit], -) -> usize { - let mut write_pos = 0; - - for_each_mask_word(bits, |mut word, word_start, _| { - while word != 0 { - let run_start = word.trailing_zeros() as usize; - let run_len = (!(word >> run_start)).trailing_zeros() as usize; - dst[write_pos..][..run_len] - .write_copy_of_slice(&src[word_start + run_start..][..run_len]); - write_pos += run_len; - - // Adding the lowest bit of the run carries through the run and clears it. - word &= word.wrapping_add(1 << run_start); - } - }); - - write_pos -} - -/// Filter a slice by a set of strictly increasing indices. -pub(super) fn filter_slice_by_indices( - slice: &[T], - indices: &[usize], - allocator: &BufferAllocatorRef, -) -> Buffer { - let mut out = BufferMut::::with_capacity_in(indices.len(), allocator.clone()); - - for (dst, &index) in out.spare_capacity_mut().iter_mut().zip(indices) { - // SAFETY: mask indices are validated when the mask is constructed. - dst.write(unsafe { *slice.get_unchecked(index) }); - } - - // SAFETY: the loop initialized every output slot. - unsafe { out.set_len(indices.len()) }; - out.freeze() -} - -/// Filter a slice by a set of strictly increasing `(start, end)` ranges. -pub(super) fn filter_slice_by_slices( - slice: &[T], - slices: &[(usize, usize)], - output_len: usize, - allocator: &BufferAllocatorRef, -) -> Buffer { - let mut out = BufferMut::::with_capacity_in(output_len, allocator.clone()); - for (start, end) in slices { - out.extend_from_slice(&slice[*start..*end]); - } - - out.freeze() -} - -/// Filter a mutable slice in-place from the mask bitmap, returning the new valid length. -pub(super) fn filter_slice_mut_by_bitmap(slice: &mut [T], mask: &MaskValues) -> usize { - assert_eq!( - slice.len(), - mask.len(), - "Mask length must equal the slice length" - ); - - let mut write_pos = 0; - - for_each_mask_word(MaskBits::Mask(mask), |word, word_start, word_len| { - let all_selected = low_bits_mask(word_len); - debug_assert_eq!(word & !all_selected, 0); - if word == all_selected { - if write_pos != word_start { - slice.copy_within(word_start..word_start + word_len, write_pos); - } - write_pos += word_len; - } else { - let mut selected = word; - while selected != 0 { - let index = word_start + selected.trailing_zeros() as usize; - // SAFETY: set bits are limited to `word_len` and stable compaction guarantees - // `write_pos <= index`. - unsafe { *slice.get_unchecked_mut(write_pos) = *slice.get_unchecked(index) }; - write_pos += 1; - selected &= selected - 1; - } - } - }); - - debug_assert_eq!(write_pos, mask.true_count()); - write_pos -} - -/// Filter a mutable slice in-place by strictly increasing indices. -pub(super) fn filter_slice_mut_by_indices(slice: &mut [T], indices: &[usize]) -> usize { - for (write_pos, &index) in indices.iter().enumerate() { - // SAFETY: mask indices are in bounds and stable compaction guarantees - // `write_pos <= index`. - unsafe { *slice.get_unchecked_mut(write_pos) = *slice.get_unchecked(index) }; - } - indices.len() -} - -/// Filter a mutable slice in-place by a set of `(start, end)` ranges, returning the new length. -pub(super) fn filter_slice_mut_by_slices( - slice: &mut [T], - slices: &[(usize, usize)], -) -> usize { - let mut write_pos = 0; - - for &(start, end) in slices { - let len = end - start; - - if write_pos != start { - slice.copy_within(start..end, write_pos); - } - - write_pos += len; - } - - write_pos -} diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index c79212c1d6e..fb7dc6a4141 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -12,9 +12,6 @@ pub use vtable::FilterArray; mod execute; pub(crate) use execute::buffer::filter_buffer; pub(crate) use execute::buffer::prepare_mask_for_reuse; -pub use execute::chunked::ChunkDecoder; -pub use execute::chunked::FILTER_CHUNK_LEN; -pub use execute::chunked::filter_chunked; pub(crate) use execute::filter_validity; mod kernel; diff --git a/vortex-buffer/src/filter/bits.rs b/vortex-buffer/src/filter/bits.rs new file mode 100644 index 00000000000..5b73b6a493c --- /dev/null +++ b/vortex-buffer/src/filter/bits.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Walks over the 64-bit words of a selection bitmap. + +use crate::BitBuffer; + +/// The selection bits that a word walker visits. +#[derive(Clone, Copy)] +pub(super) enum MaskBits<'a> { + /// All the bits of a bitmap. + Bitmap(&'a BitBuffer), + /// The first `len` bits of `words`, with any bits past `len` cleared. + Words { words: &'a [u64], len: usize }, +} + +/// Invoke `f` with each `(word, word_start, word_len)` of `bits`, where `word` holds the bits +/// for elements `word_start..word_start + word_len` in its low `word_len` bits. +#[inline] +pub(super) fn for_each_mask_word(bits: MaskBits<'_>, f: impl FnMut(u64, usize, usize)) { + match bits { + MaskBits::Bitmap(bits) => for_each_bitmap_word(bits, f), + MaskBits::Words { words, len } => for_each_slice_word(words, len, f), + } +} + +/// Returns whether `f` holds for every word of `bits`. Stops at the first word for which `f` +/// does not hold. +/// +/// Words do not align to element positions, so `f` must only look at properties of a word that +/// do not depend on its position, such as the number of runs of set bits. +pub(super) fn all_mask_words(bits: MaskBits<'_>, mut f: impl FnMut(u64) -> bool) -> bool { + match bits { + MaskBits::Bitmap(bits) => bits.unaligned_chunks().iter().all(f), + MaskBits::Words { words, .. } => words.iter().all(|&word| f(word)), + } +} + +#[inline] +fn for_each_slice_word(words: &[u64], len: usize, mut f: impl FnMut(u64, usize, usize)) { + debug_assert_eq!(words.len(), len.div_ceil(64)); + + for (word_idx, &word) in words.iter().enumerate() { + let word_start = word_idx * 64; + f(word, word_start, (len - word_start).min(64)); + } +} + +#[inline] +fn for_each_bitmap_word(bits: &BitBuffer, mut f: impl FnMut(u64, usize, usize)) { + let unaligned = bits.unaligned_chunks(); + let lead = unaligned.lead_padding(); + let mut base = 0; + + if let Some(prefix) = unaligned.prefix() { + let len = (64 - lead).min(bits.len()); + f(prefix >> lead, base, len); + base += len; + } + + for &word in unaligned.chunks() { + f(word, base, 64); + base += 64; + } + + if let Some(suffix) = unaligned.suffix() { + let len = bits.len() - base; + f(suffix, base, len); + base += len; + } + + debug_assert_eq!(base, bits.len()); +} + +/// A `u64` with the low `len` bits set. +#[inline] +pub(super) fn low_bits_mask(len: usize) -> u64 { + debug_assert!(len <= 64); + if len == 64 { + u64::MAX + } else { + (1u64 << len) - 1 + } +} diff --git a/vortex-buffer/src/filter/byte_compress.rs b/vortex-buffer/src/filter/byte_compress.rs new file mode 100644 index 00000000000..054c39351dd --- /dev/null +++ b/vortex-buffer/src/filter/byte_compress.rs @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Byte-level compress using a `1 << 8 = 256`-entry lookup table. +//! +//! For each byte of the mask (8 bits -> 8 source elements), a precomputed permutation table +//! compacts the selected elements in a single indexed copy, without materializing indices or +//! slices. + +use std::mem::MaybeUninit; + +use super::bits::MaskBits; +use super::bits::for_each_mask_word; + +/// For each mask byte (0..256), stores the element indices to keep and the count. +/// +/// `BYTE_COMPRESS_LUT[mask_byte]` = `([i0, i1, ..., i7], popcount)` where +/// `i0..i_{popcount-1}` are the positions of set bits in `mask_byte`. +/// +/// Total size: 256 * 9 = 2304 bytes, which trivially fits in L1 cache. +static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{ + let mut lut = [([0u8; 8], 0u8); 256]; + let mut mask: usize = 0; + while mask < 256 { + let mut indices = [0u8; 8]; + let mut count: u8 = 0; + let mut bit: u8 = 0; + while bit < 8 { + if mask & (1 << bit) != 0 { + indices[count as usize] = bit; + count += 1; + } + bit += 1; + } + lut[mask] = (indices, count); + mask += 1; + } + lut +}; + +/// Copy the elements of `src` selected by `bits` to `dst` and return the number copied. +/// +/// # Safety +/// +/// `bits` must select only elements of `src`, and `dst` must hold at least one value for each +/// selected element. +pub(super) unsafe fn compress( + src: &[T], + bits: MaskBits<'_>, + dst: &mut [MaybeUninit], +) -> usize { + let mut written = 0; + + for_each_mask_word(bits, |word, word_start, word_len| { + let (chunks, tail) = src[word_start..][..word_len].as_chunks::<8>(); + let mask_bytes = word.to_le_bytes(); + + for (chunk, &mask_byte) in chunks.iter().zip(&mask_bytes) { + // SAFETY: the mask byte selects only elements of `chunk`, and `dst` has room for + // every selected element. + written += unsafe { compress_byte(chunk, mask_byte, dst.get_unchecked_mut(written..)) }; + } + + if !tail.is_empty() { + // SAFETY: the bits past `word_len` are cleared, so the mask byte selects only + // elements of `tail`. + written += unsafe { + compress_byte( + tail, + mask_bytes[chunks.len()], + dst.get_unchecked_mut(written..), + ) + }; + } + }); + + written +} + +/// Copy the elements of `chunk` selected by `mask_byte` to `dst` and return the number copied. +/// +/// # Safety +/// +/// `mask_byte` must select only elements of `chunk`, and `dst` must hold at least one value for +/// each selected element. +#[inline] +unsafe fn compress_byte(chunk: &[T], mask_byte: u8, dst: &mut [MaybeUninit]) -> usize { + if mask_byte == 0 { + return 0; + } + + if chunk.len() == 8 && mask_byte == 0xFF { + // All 8 selected, so bulk copy. + dst[..8].write_copy_of_slice(chunk); + return 8; + } + + let (perm, count) = &BYTE_COMPRESS_LUT[mask_byte as usize]; + let count = *count as usize; + // SAFETY: perm indices are all below `chunk.len()`, and `dst` has room for `count` values. + unsafe { + for j in 0..count { + dst.get_unchecked_mut(j) + .write(*chunk.get_unchecked(*perm.get_unchecked(j) as usize)); + } + } + count +} diff --git a/vortex-array/src/arrays/filter/execute/chunked.rs b/vortex-buffer/src/filter/chunked.rs similarity index 75% rename from vortex-array/src/arrays/filter/execute/chunked.rs rename to vortex-buffer/src/filter/chunked.rs index d8f0edf02dd..b3a8ecd5593 100644 --- a/vortex-array/src/arrays/filter/execute/chunked.rs +++ b/vortex-buffer/src/filter/chunked.rs @@ -5,22 +5,16 @@ use std::mem::MaybeUninit; -use vortex_buffer::Buffer; -use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_mask::MaskValues; - -use super::buffer::byte_compress_density_threshold; -use super::buffer::useful_cached_indices; -use super::buffer::useful_cached_slices; -use super::byte_compress; -use super::simd_compress; -use super::simd_compress::SLACK_BYTES; -use super::slice::MaskBits; -use super::slice::compact_by_bitmap; -use super::slice::compact_runs_by_bitmap; -use super::slice::low_bits_mask; -use crate::dtype::NativePType; +use super::Repr; +use super::Selection; +use super::bits::MaskBits; +use super::bits::low_bits_mask; +use super::compact_bits; +use super::simd::SLACK_BYTES; +use crate::BitBuffer; +use crate::Buffer; +use crate::BufferAllocatorRef; +use crate::BufferMut; /// Number of values in each chunk of a [`ChunkDecoder`]. pub const FILTER_CHUNK_LEN: usize = 1024; @@ -43,10 +37,6 @@ const fn max_sparse_values() -> usize { } } -/// Chunks without a SIMD kernel whose runs of selected values average at least this many bytes -/// copy each run. Shorter runs are faster to compact with byte compress. -const MIN_RUN_BYTES: usize = 96; - /// Number of mask bytes read to assemble the words of a chunk at an arbitrary bit offset. const CHUNK_WORD_BYTES: usize = FILTER_CHUNK_LEN / 8 + 8; @@ -65,15 +55,15 @@ pub unsafe trait ChunkDecoder { fn decode_indices(&self, chunk_idx: usize, indices: &[usize], dst: &mut [MaybeUninit]); } -/// Filters the values of `decoder` by `mask`, and decodes only the chunks that hold selected -/// values. +/// Filters the values of `decoder` by `selection`, and decodes only the chunks that hold +/// selected values. /// /// Value `i` of the filtered array is value `offset + i` of the decoded chunks, so `offset` skips -/// the values at the start of the first chunk. +/// the values at the start of the first chunk. The decoded chunks must hold at least +/// `offset + selection.len()` values. /// -/// The mask is applied one chunk at a time. The filter uses the cached slices or the cached -/// indices of the mask when they are useful, in that order, and otherwise uses its bitmap. Each chunk takes the first strategy that -/// matches: +/// The selection is applied one chunk at a time, in the representation that `selection` holds. +/// Each chunk takes the first strategy that matches: /// /// | Selected values in the chunk | Strategy | /// | ---------------------------- | ------------------------------------------- | @@ -82,26 +72,25 @@ pub unsafe trait ChunkDecoder { /// | at most the sparse limit | decode only the selected values | /// | otherwise | decode to scratch, then compact the scratch | /// -/// The sparse limit is 16, 32, 64 or 160 for values of 1, 2, 4 or 8 bytes, because a wider -/// chunk costs more to decode. +/// The sparse limit is 16, 32, 64 or 160 for values of 1, 2, 4 or 8 or more bytes, because a +/// wider chunk costs more to decode. +/// +/// The scratch buffer stays in cache. The filter compacts it by copying the slices or gathering +/// the indices, or from a bitmap with the same strategy ladder as [`filter_buffer`]. /// -/// The scratch buffer stays in cache. The filter compacts it by copying the slices, or with the -/// SIMD kernel that a plain buffer filter chooses for the density of the chunk. When no SIMD -/// kernel applies, the filter copies each run of set bits if the runs are long, and otherwise -/// uses byte compress or the bit walk. -pub fn filter_chunked( +/// [`filter_buffer`]: super::filter_buffer +pub fn filter_chunked( decoder: &impl ChunkDecoder, offset: usize, - mask: &MaskValues, + selection: Selection<'_>, allocator: &BufferAllocatorRef, ) -> Buffer { - let mut filter = ChunkedFilter::new(decoder, mask.true_count(), allocator); - if let Some(slices) = useful_cached_slices(mask) { - filter.push_slices(offset, slices); - } else if let Some(indices) = useful_cached_indices(mask) { - filter.push_mask_indices(offset, indices); - } else { - filter.push_bitmap(offset, mask); + let mut filter = ChunkedFilter::new(decoder, selection.true_count, allocator); + + match selection.repr { + Repr::Bitmap(bits) => filter.push_bitmap(offset, bits), + Repr::Indices(indices) => filter.push_mask_indices(offset, indices), + Repr::Slices(slices) => filter.push_slices(offset, slices), } filter.finish() @@ -130,7 +119,7 @@ struct ChunkedFilter<'a, T, D> { scratch: Scratch, } -impl<'a, T: NativePType, D: ChunkDecoder> ChunkedFilter<'a, T, D> { +impl<'a, T: Copy, D: ChunkDecoder> ChunkedFilter<'a, T, D> { fn new(decoder: &'a D, true_count: usize, allocator: &BufferAllocatorRef) -> Self { Self { decoder, @@ -161,10 +150,9 @@ impl<'a, T: NativePType, D: ChunkDecoder> ChunkedFilter<'a, T, D> { self.written += indices.len(); } - /// Filters chunk by chunk from the bitmap of `mask`. - fn push_bitmap(&mut self, offset: usize, mask: &MaskValues) { - let len = mask.len(); - let bits = mask.bit_buffer(); + /// Filters chunk by chunk from the bitmap `bits`. + fn push_bitmap(&mut self, offset: usize, bits: &BitBuffer) { + let len = bits.len(); let bits_bytes = bits.inner().as_slice(); let mut words = [0u64; CHUNK_WORDS]; let mut indices = [0usize; MAX_SPARSE_VALUES]; @@ -204,9 +192,13 @@ impl<'a, T: NativePType, D: ChunkDecoder> ChunkedFilter<'a, T, D> { let src = &self.scratch.decode(self.decoder, chunk_idx)[chunk_offset..][..chunk_len]; let dst = &mut self.values.spare_capacity_mut()[self.written..]; + let bits = MaskBits::Words { + words, + len: chunk_len, + }; // SAFETY: the words select only values of `src` and clear any bits past // `chunk_len`, and the output has room for every selected value plus the slack. - let written = unsafe { compact_words(src, words, selected, dst) }; + let written = unsafe { compact_bits(src, bits, selected, dst) }; debug_assert_eq!(written, selected); self.written += selected; } @@ -313,64 +305,6 @@ impl<'a, T: NativePType, D: ChunkDecoder> ChunkedFilter<'a, T, D> { } } -/// Copies the `selected` values of `src` selected by `words` to `dst` and returns the number -/// copied. -/// -/// The kernel is the SIMD kernel that [`filter_buffer`](super::buffer::filter_buffer) chooses for -/// the density, if any. Otherwise, runs of set bits that average at least [`MIN_RUN_BYTES`] are -/// copied run by run, and shorter runs use byte compress or the bit walk. -/// -/// # Safety -/// -/// `words` must hold `src.len()` bits with any bits past `src.len()` cleared, and `dst` must hold -/// every selected value plus [`SLACK_BYTES`]. -unsafe fn compact_words( - src: &[T], - words: &[u64], - selected: usize, - dst: &mut [MaybeUninit], -) -> usize { - let bits = MaskBits::Words { - words, - len: src.len(), - }; - let density = selected as f64 / src.len() as f64; - - // SAFETY: forwarded from the caller contract. - unsafe { - if let Some(written) = simd_compress::compress_bits(src, bits, density, dst) { - written - } else if has_long_runs::(words, selected) { - compact_runs_by_bitmap(src, bits, dst) - } else if density >= byte_compress_density_threshold::() { - byte_compress::compress_words(src, words, dst) - } else { - compact_by_bitmap(src, bits, dst) - } - } -} - -/// Returns whether the runs of the `selected` set bits of `words` average at least -/// [`MIN_RUN_BYTES`] of `T`. -// Inlining this check into `compact_words` makes the bit walk up to 1.3x slower on x86-64. -#[inline(never)] -fn has_long_runs(words: &[u64], selected: usize) -> bool { - // The run walk copies at most one word of values at a time, so narrow values never qualify. - if 64 * size_of::() < MIN_RUN_BYTES { - return false; - } - - // A run that crosses a word boundary counts once in each word, as the run walk copies it. The - // count stops at the first word that takes it past `max_runs`, because a popcount is not a - // native instruction on baseline x86-64. - let max_runs = selected * size_of::() / MIN_RUN_BYTES; - let mut runs = 0; - words.iter().all(|&word| { - runs += (word & !(word << 1)).count_ones() as usize; - runs <= max_runs - }) -} - /// Loads the `len <= FILTER_CHUNK_LEN` mask bits starting at bit `start` of `bytes` into /// `words`, with any bits past `len` cleared. #[inline] @@ -427,16 +361,15 @@ mod tests { use rand::SeedableRng; use rand::rngs::StdRng; use rstest::rstest; - use vortex_buffer::BitBuffer; - use vortex_buffer::BufferAllocatorRef; - use vortex_mask::Mask; + use super::BufferAllocatorRef; use super::ChunkDecoder; use super::FILTER_CHUNK_LEN; use super::filter_chunked; - use crate::dtype::NativePType; + use crate::BitBuffer; + use crate::filter::Selection; - /// The representation that the test mask caches. + /// The representation of the test selection. #[derive(Clone, Copy, Debug)] enum Cached { Bitmap, @@ -461,7 +394,11 @@ mod tests { } } - fn check>(bits: &BitBuffer, offset: usize, cached: Cached) { + fn check>( + bits: &BitBuffer, + offset: usize, + cached: Cached, + ) { let padded_len = (offset + bits.len()).next_multiple_of(FILTER_CHUNK_LEN); let values = (0..padded_len) .map(|i| { @@ -474,18 +411,23 @@ mod tests { .filter_map(|(i, selected)| selected.then_some(values[offset + i])) .collect::>(); - let mask = match cached { - Cached::Bitmap => Mask::from_buffer(bits.clone()), - Cached::Slices => Mask::from_slices(bits.len(), bits.set_slices().collect()), - Cached::Indices => Mask::from_indices(bits.len(), bits.set_indices()), - }; - let Mask::Values(mask) = mask else { - unreachable!("the patterns select some but not all values") + let slices = bits.set_slices().collect::>(); + let indices = bits.set_indices().collect::>(); + // SAFETY: all three representations come from `bits`. + let selection = unsafe { + match cached { + Cached::Bitmap => Selection::from_bitmap_unchecked(bits, bits.true_count()), + Cached::Slices => { + Selection::from_slices_unchecked(bits.len(), &slices, bits.true_count()) + } + Cached::Indices => Selection::from_indices_unchecked(bits.len(), &indices), + } }; + let filtered = filter_chunked( &SliceDecoder(&values), offset, - &mask, + selection, BufferAllocatorRef::static_ref(), ); assert_eq!(filtered.as_slice(), expected); @@ -496,7 +438,7 @@ mod tests { #[values(0.001, 0.01, 0.05, 0.2, 0.5, 0.9, 0.999)] density: f64, #[values(0, 3, 1000)] offset: usize, #[values(0, 5)] bit_offset: usize, - #[values(Cached::Bitmap, Cached::Indices)] cached: Cached, + #[values(Cached::Bitmap, Cached::Slices, Cached::Indices)] cached: Cached, ) { let mut rng = StdRng::seed_from_u64(42); let bits = (0..5000 + bit_offset) @@ -508,9 +450,10 @@ mod tests { check::(&bits, offset, cached); check::(&bits, offset, cached); check::(&bits, offset, cached); + check::(&bits, offset, cached); } - /// Runs long enough that the cached slices are used rather than the bitmap. + /// Runs of selected values, some of which cross chunk boundaries. #[rstest] fn filter_chunked_runs_matches_expected( #[values(9, 12, 200, 3000)] run_len: usize, @@ -526,12 +469,12 @@ mod tests { check::(&bits, offset, cached); check::(&bits, offset, cached); check::(&bits, offset, cached); + check::(&bits, offset, cached); } #[rstest] fn filter_chunked_full_and_empty_chunks(#[values(0, 3)] offset: usize) { - // One whole chunk of selected values in every three chunks. The density stays low enough - // for the cached indices to be used. + // One whole chunk of selected values in every three chunks. let bits = (0..5000) .map(|i| ((i + offset) / FILTER_CHUNK_LEN).is_multiple_of(3)) .collect::(); diff --git a/vortex-buffer/src/filter/mod.rs b/vortex-buffer/src/filter/mod.rs new file mode 100644 index 00000000000..5f9c2d20315 --- /dev/null +++ b/vortex-buffer/src/filter/mod.rs @@ -0,0 +1,329 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Filtering of fixed-width values by a [`Selection`]. +//! +//! [`filter_buffer`] filters values that are already in memory, and [`filter_chunked`] filters +//! values that a [`ChunkDecoder`] decodes one chunk at a time. Both work for any `T: Copy` and +//! share one strategy ladder. The first matching row wins: +//! +//! | Priority | Selection or condition | Out-of-place | In-place | +//! | -------- | --------------------------------------------------- | -------------- | ------------- | +//! | 1 | slices | copy runs | move runs | +//! | 2 | indices | gather indices | move indices | +//! | 3 | bitmap, an eligible architecture kernel | SIMD compress | SIMD compress | +//! | 4 | bitmap, density reaches the byte compress threshold | byte compress | bitmap walk | +//! | 5 | bitmap, runs average at least 8 values and 96 bytes | copy runs | bitmap walk | +//! | 6 | bitmap, otherwise | bitmap walk | bitmap walk | +//! +//! [`filter_buffer`] compacts in place only a uniquely owned buffer with density >= 0.50. +//! [`filter_chunked`] applies the ladder to each chunk that it decodes, at the density of that +//! chunk. SIMD eligibility is summarized in the `simd` module. + +use std::mem::MaybeUninit; + +use crate::BitBuffer; +use crate::Buffer; +use crate::BufferAllocatorRef; +use crate::BufferMut; + +mod bits; +mod byte_compress; +mod chunked; +mod scalar; +mod simd; +#[cfg(test)] +mod tests; + +use bits::MaskBits; +pub use chunked::ChunkDecoder; +pub use chunked::FILTER_CHUNK_LEN; +pub use chunked::filter_chunked; + +const IN_PLACE_MIN_DENSITY: f64 = 0.5; + +/// Bitmaps without a SIMD kernel whose runs of selected values average at least this many bytes, +/// and at least [`MIN_RUN_LEN`] values, copy each run. +const MIN_RUN_BYTES: usize = 96; + +/// The smallest average run length, in values, for which a bitmap copies each run. Random masks +/// of wide values reach [`MIN_RUN_BYTES`] with runs of only a few values, and the bitmap walk is +/// faster for those. +const MIN_RUN_LEN: usize = 8; + +/// The positions that a filter keeps, out of `len` positions. +/// +/// A selection holds one representation of the positions. The caller chooses the +/// representation, because the cost to make one from another is higher than the cost to filter. +#[derive(Clone, Copy)] +pub struct Selection<'a> { + len: usize, + true_count: usize, + repr: Repr<'a>, +} + +#[derive(Clone, Copy)] +enum Repr<'a> { + Bitmap(&'a BitBuffer), + Indices(&'a [usize]), + Slices(&'a [(usize, usize)]), +} + +impl<'a> Selection<'a> { + /// Selects the set bits of `bits`. + /// + /// # Safety + /// + /// `true_count` must equal the number of set bits in `bits`. + pub unsafe fn from_bitmap_unchecked(bits: &'a BitBuffer, true_count: usize) -> Self { + debug_assert_eq!(bits.true_count(), true_count); + + Self { + len: bits.len(), + true_count, + repr: Repr::Bitmap(bits), + } + } + + /// Selects the positions in `indices`, out of `len` positions. + /// + /// # Safety + /// + /// `indices` must be strictly increasing, and each index must be less than `len`. + pub unsafe fn from_indices_unchecked(len: usize, indices: &'a [usize]) -> Self { + debug_assert!(indices.is_sorted_by(|a, b| a < b)); + debug_assert!(indices.last().is_none_or(|&last| last < len)); + + Self { + len, + true_count: indices.len(), + repr: Repr::Indices(indices), + } + } + + /// Selects the positions in the `(start, end)` ranges of `slices`, out of `len` positions. + /// + /// # Safety + /// + /// The ranges must be sorted and disjoint, each `start` must not be more than its `end`, the + /// last `end` must not be more than `len`, and `true_count` must equal the total length of + /// the ranges. + pub unsafe fn from_slices_unchecked( + len: usize, + slices: &'a [(usize, usize)], + true_count: usize, + ) -> Self { + debug_assert!(slices.iter().all(|&(start, end)| start <= end)); + debug_assert!(slices.is_sorted_by(|a, b| a.1 <= b.0)); + debug_assert!(slices.last().is_none_or(|&(_, end)| end <= len)); + debug_assert_eq!( + slices.iter().map(|(start, end)| end - start).sum::(), + true_count + ); + + Self { + len, + true_count, + repr: Repr::Slices(slices), + } + } + + /// Returns the number of positions that the selection keeps or drops. + pub fn len(&self) -> usize { + self.len + } + + /// Returns whether the selection has no positions. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Returns the number of positions that the selection keeps. + pub fn true_count(&self) -> usize { + self.true_count + } + + fn density(&self) -> f64 { + if self.len == 0 { + return 0.0; + } + + self.true_count as f64 / self.len as f64 + } +} + +/// Filters `buffer` by `selection` and returns the selected values. +/// +/// A uniquely owned buffer with density of at least 0.50 is compacted in place. Other buffers +/// are copied to a new allocation from `allocator`. +/// +/// # Panics +/// +/// Panics if `buffer` and `selection` have different lengths. +pub fn filter_buffer( + buffer: Buffer, + selection: Selection<'_>, + allocator: &BufferAllocatorRef, +) -> Buffer { + assert_eq!( + buffer.len(), + selection.len(), + "Selection length must equal the buffer length" + ); + + let buffer = if selection.density() >= IN_PLACE_MIN_DENSITY { + match buffer.try_into_mut() { + Ok(mut buffer_mut) => { + let new_len = filter_slice_in_place(buffer_mut.as_mut_slice(), selection); + buffer_mut.truncate(new_len); + return buffer_mut.freeze(); + } + Err(buffer) => buffer, + } + } else { + buffer + }; + + filter_slice(buffer.as_slice(), selection, allocator) +} + +/// Copies the values of `values` that `selection` keeps to a new buffer. +fn filter_slice( + values: &[T], + selection: Selection<'_>, + allocator: &BufferAllocatorRef, +) -> Buffer { + debug_assert_eq!(values.len(), selection.len()); + + let true_count = selection.true_count(); + // The SIMD kernels can store one vector past the last value that they write. + let mut out = BufferMut::::with_capacity_in( + true_count + simd::SLACK_BYTES / size_of::(), + allocator.clone(), + ); + let dst = out.spare_capacity_mut(); + + let written = match selection.repr { + Repr::Slices(slices) => scalar::copy_slices(values, slices, dst), + Repr::Indices(indices) => { + // SAFETY: the selection guarantees that each index is less than `values.len()`. + unsafe { scalar::gather_indices(values, indices, dst) }; + indices.len() + } + // SAFETY: the bitmap has one bit for each value, and the output has room for every + // selected value plus the slack. + Repr::Bitmap(bits) => unsafe { + compact_bits(values, MaskBits::Bitmap(bits), true_count, dst) + }, + }; + + debug_assert_eq!(written, true_count); + // SAFETY: the strategy above initialized the first `true_count` values. + unsafe { out.set_len(true_count) }; + out.freeze() +} + +/// Compacts the values of `values` that `selection` keeps to its front, and returns the new +/// length. +fn filter_slice_in_place(values: &mut [T], selection: Selection<'_>) -> usize { + debug_assert_eq!(values.len(), selection.len()); + + let written = match selection.repr { + Repr::Slices(slices) => scalar::compact_in_place_by_slices(values, slices), + // SAFETY: the selection guarantees that each index is less than `values.len()`. + Repr::Indices(indices) => unsafe { scalar::compact_in_place_by_indices(values, indices) }, + Repr::Bitmap(bits) => { + let bits = MaskBits::Bitmap(bits); + // SAFETY: the bitmap has one bit for each value. + unsafe { simd::compress_in_place(values, bits, selection.density()) } + .unwrap_or_else(|| scalar::compact_in_place_by_bitmap(values, bits)) + } + }; + + debug_assert_eq!(written, selection.true_count()); + written +} + +/// Copies the `selected` values of `src` that `bits` selects to `dst` and returns the number +/// copied. +/// +/// This is rows 3 to 6 of the [module](self) ladder: the SIMD kernel for the density, if any. +/// Otherwise, a density at or above the byte compress threshold uses byte compress, long runs +/// of set bits are copied run by run, and shorter runs use the bitmap walk. +/// +/// # Safety +/// +/// `bits` must select only values of `src`, and `dst` must hold every selected value plus +/// [`SLACK_BYTES`](simd::SLACK_BYTES). +unsafe fn compact_bits( + src: &[T], + bits: MaskBits<'_>, + selected: usize, + dst: &mut [MaybeUninit], +) -> usize { + let density = if src.is_empty() { + 0.0 + } else { + selected as f64 / src.len() as f64 + }; + + // SAFETY: forwarded from the caller contract. + unsafe { + if let Some(written) = simd::compress(src, bits, density, dst) { + written + } else if density >= byte_compress_density_threshold::() { + byte_compress::compress(src, bits, dst) + } else if has_long_runs::(bits, selected) { + scalar::compact_runs_by_bitmap(src, bits, dst) + } else { + scalar::compact_by_bitmap(src, bits, dst) + } + } +} + +/// Returns whether the runs of the `selected` set bits of `bits` average at least +/// [`MIN_RUN_LEN`] values and [`MIN_RUN_BYTES`] of `T`. +// Inlining this check into `compact_bits` makes the bit walk up to 1.3x slower on x86-64. +#[inline(never)] +fn has_long_runs(bits: MaskBits<'_>, selected: usize) -> bool { + // The run walk copies at most one word of values at a time, so narrow values never qualify. + if 64 * size_of::() < MIN_RUN_BYTES { + return false; + } + + // A run that crosses a word boundary counts once in each word, as the run walk copies it. The + // count stops at the first word that takes it past `max_runs`, because a popcount is not a + // native instruction on baseline x86-64. + let max_runs = (selected * size_of::() / MIN_RUN_BYTES).min(selected / MIN_RUN_LEN); + let mut runs = 0; + bits::all_mask_words(bits, |word| { + runs += (word & !(word << 1)).count_ones() as usize; + runs <= max_runs + }) +} + +/// Returns the density at or above which byte compress is used after the higher-priority +/// strategies have declined the bitmap. +fn byte_compress_density_threshold() -> f64 { + let width = size_of::(); + + // These crossovers are benchmarked in `vortex-array/benches/filter_fixed_width.rs`. + // + // | Target | 1 byte | 2 bytes | 4 bytes | 8 bytes | other | + // | ------- | -----: | ------: | ------: | ------: | ----: | + // | aarch64 | 0.90 | 0.90 | 0.90 | 0.75 | 0.90 | + // | other | 0.00 | 0.50 | 0.50 | 0.75 | 0.875 | + if cfg!(target_arch = "aarch64") { + return match width { + 8 => 0.75, + _ => 0.9, + }; + } + + match width { + 1 => 0.0, + 2 | 4 => 0.5, + 8 => 0.75, + _ => 0.875, + } +} diff --git a/vortex-buffer/src/filter/scalar.rs b/vortex-buffer/src/filter/scalar.rs new file mode 100644 index 00000000000..7dd33d8bd9f --- /dev/null +++ b/vortex-buffer/src/filter/scalar.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar compaction of slices by bits, indices, or ranges, both out of place and in place. + +use std::mem::MaybeUninit; + +use super::bits::MaskBits; +use super::bits::for_each_mask_word; +use super::bits::low_bits_mask; + +/// Copy the elements of `src` selected by `bits` to `dst` and return the number copied. +/// +/// # Safety +/// +/// `bits` must select only elements of `src`, and `dst` must hold at least one value for each +/// selected element. +#[inline] +pub(super) unsafe fn compact_by_bitmap( + src: &[T], + bits: MaskBits<'_>, + dst: &mut [MaybeUninit], +) -> usize { + let mut write_pos = 0; + + for_each_mask_word(bits, |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + + if word == all_selected { + dst[write_pos..][..word_len].write_copy_of_slice(&src[word_start..][..word_len]); + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + // SAFETY: set bits index into `src`, and `dst` has room for every selected + // element. + unsafe { + dst.get_unchecked_mut(write_pos) + .write(*src.get_unchecked(index)) + }; + write_pos += 1; + selected &= selected - 1; + } + } + }); + + write_pos +} + +/// Copy the elements of `src` selected by `bits` to `dst` one run of set bits at a time, and +/// return the number copied. +/// +/// This is faster than [`compact_by_bitmap`] when the selected elements form long runs. +// Inlining this walk next to `compact_by_bitmap` makes the bit walk 1.5x slower on Apple M4. +#[inline(never)] +pub(super) fn compact_runs_by_bitmap( + src: &[T], + bits: MaskBits<'_>, + dst: &mut [MaybeUninit], +) -> usize { + let mut write_pos = 0; + + for_each_mask_word(bits, |mut word, word_start, _| { + while word != 0 { + let run_start = word.trailing_zeros() as usize; + let run_len = (!(word >> run_start)).trailing_zeros() as usize; + dst[write_pos..][..run_len] + .write_copy_of_slice(&src[word_start + run_start..][..run_len]); + write_pos += run_len; + + // Adding the lowest bit of the run carries through the run and clears it. + word &= word.wrapping_add(1 << run_start); + } + }); + + write_pos +} + +/// Copy the elements of `src` at the strictly increasing `indices` to `dst`. +/// +/// # Safety +/// +/// Each index must be in bounds of `src`. +pub(super) unsafe fn gather_indices( + src: &[T], + indices: &[usize], + dst: &mut [MaybeUninit], +) { + for (dst, &index) in dst[..indices.len()].iter_mut().zip(indices) { + // SAFETY: forwarded from the caller contract. + dst.write(unsafe { *src.get_unchecked(index) }); + } +} + +/// Copy the `(start, end)` ranges of `src` to `dst` and return the number of elements copied. +pub(super) fn copy_slices( + src: &[T], + slices: &[(usize, usize)], + dst: &mut [MaybeUninit], +) -> usize { + let mut write_pos = 0; + + for &(start, end) in slices { + dst[write_pos..][..end - start].write_copy_of_slice(&src[start..end]); + write_pos += end - start; + } + + write_pos +} + +/// Compact the elements of `values` selected by `bits` to its front, and return the new length. +pub(super) fn compact_in_place_by_bitmap(values: &mut [T], bits: MaskBits<'_>) -> usize { + let mut write_pos = 0; + + for_each_mask_word(bits, |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + + if word == all_selected { + if write_pos != word_start { + values.copy_within(word_start..word_start + word_len, write_pos); + } + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + // SAFETY: set bits are limited to `word_len` and stable compaction guarantees + // `write_pos <= index`. + unsafe { *values.get_unchecked_mut(write_pos) = *values.get_unchecked(index) }; + write_pos += 1; + selected &= selected - 1; + } + } + }); + + write_pos +} + +/// Compact the elements of `values` at the strictly increasing `indices` to its front, and +/// return the new length. +/// +/// # Safety +/// +/// Each index must be in bounds of `values`. +pub(super) unsafe fn compact_in_place_by_indices( + values: &mut [T], + indices: &[usize], +) -> usize { + for (write_pos, &index) in indices.iter().enumerate() { + // SAFETY: the index is in bounds, and strictly increasing indices guarantee + // `write_pos <= index`. + unsafe { *values.get_unchecked_mut(write_pos) = *values.get_unchecked(index) }; + } + + indices.len() +} + +/// Compact the sorted, disjoint `(start, end)` ranges of `values` to its front, and return the +/// new length. +pub(super) fn compact_in_place_by_slices( + values: &mut [T], + slices: &[(usize, usize)], +) -> usize { + let mut write_pos = 0; + + for &(start, end) in slices { + if write_pos != start { + values.copy_within(start..end, write_pos); + } + write_pos += end - start; + } + + write_pos +} diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/generic.rs b/vortex-buffer/src/filter/simd/generic.rs similarity index 90% rename from vortex-array/src/arrays/filter/execute/simd_compress/generic.rs rename to vortex-buffer/src/filter/simd/generic.rs index 36a7172e010..ee8b2cb3880 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/generic.rs +++ b/vortex-buffer/src/filter/simd/generic.rs @@ -16,9 +16,9 @@ use fearless_simd::prelude::*; use fearless_simd::u8x16; use fearless_simd_macros::simd; -use super::super::slice::MaskBits; -use super::super::slice::for_each_mask_word; -use super::super::slice::low_bits_mask; +use super::super::bits::MaskBits; +use super::super::bits::for_each_mask_word; +use super::super::bits::low_bits_mask; use super::bulk_copy; use super::compress_lut; use super::compress_tail; @@ -30,8 +30,8 @@ static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2); /// /// # Safety /// -/// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / -/// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. +/// The pointer contract of [`compress`](super::compress) / +/// [`compress_in_place`](super::compress_in_place) must hold. #[expect( clippy::cast_possible_truncation, reason = "deliberate submask narrowing" @@ -102,8 +102,8 @@ macro_rules! generic_compress_kernel { ($walk_fn:ident,elem_size: $elem_size:literal,lanes: $lanes:literal,idx_lut: $idx_lut:ident) => { /// # Safety /// - /// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// The pointer contract of [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. pub(super) unsafe fn $walk_fn( src: *const u8, dst: *mut u8, diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs b/vortex-buffer/src/filter/simd/mod.rs similarity index 75% rename from vortex-array/src/arrays/filter/execute/simd_compress/mod.rs rename to vortex-buffer/src/filter/simd/mod.rs index c9c39bd999e..1ab44930975 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs +++ b/vortex-buffer/src/filter/simd/mod.rs @@ -22,17 +22,12 @@ //! \* The 1- and 2-byte fallbacks require AVX2. //! //! On x86, the widest available kernel is considered first for each width. When no entry matches, -//! control returns to [`buffer`](super::buffer) for scalar selection. +//! control returns to [`filter`](super) for scalar selection. use std::mem::MaybeUninit; use std::ptr; -use vortex_buffer::Buffer; -use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_mask::MaskValues; - -use super::slice::MaskBits; +use super::bits::MaskBits; #[cfg(all(target_arch = "x86_64", not(miri)))] mod generic; @@ -45,62 +40,12 @@ mod x86; const MIN_LEN: usize = 64; +/// The number of bytes that an out-of-place kernel can store past the last element it writes. pub(super) const SLACK_BYTES: usize = 64; /// `dst == src` for in-place kernels. type Kernel = unsafe fn(*const u8, *mut u8, MaskBits<'_>) -> usize; -/// Filter a slice with a SIMD compress kernel, if one applies. -/// -/// Returns `None` when the caller should use a scalar strategy. -#[inline] -pub(super) fn filter_slice_by_bitmap( - values: &[T], - mask: &MaskValues, - allocator: &BufferAllocatorRef, -) -> Option> { - debug_assert_eq!(values.len(), mask.len()); - let kernel = select_kernel::(mask.len(), mask.density())?; - - let true_count = mask.true_count(); - let mut out = BufferMut::::with_capacity_in( - true_count + SLACK_BYTES / size_of::(), - allocator.clone(), - ); - // SAFETY: `select_kernel` probed the kernel's target features; `values` holds `mask.len()` - // elements and the output has capacity for every selected element plus a full vector of - // slack, so each unmasked store stays in bounds. - let written = unsafe { - kernel( - values.as_ptr().cast(), - out.spare_capacity_mut().as_mut_ptr().cast(), - MaskBits::Mask(mask), - ) - }; - debug_assert_eq!(written, true_count); - // SAFETY: the kernel initialized the first `true_count` elements. - unsafe { out.set_len(true_count) }; - Some(out.freeze()) -} - -/// In-place variant of [`filter_slice_by_bitmap`]: compact the selected elements to the front -/// of `values` and return the new length, if a SIMD kernel applies. -pub(super) fn filter_slice_mut_by_bitmap( - values: &mut [T], - mask: &MaskValues, -) -> Option { - debug_assert_eq!(values.len(), mask.len()); - let kernel = select_kernel::(mask.len(), mask.density())?; - - let dst = values.as_mut_ptr().cast::(); - // SAFETY: `select_kernel` probed the kernel's target features; the in-place instantiation - // compacts forward (stores never pass the positions it has already read) and keeps partial - // tail chunks off the full-width store path, so all accesses stay inside `values`. - let written = unsafe { kernel(dst.cast_const(), dst, MaskBits::Mask(mask)) }; - debug_assert_eq!(written, mask.true_count()); - Some(written) -} - /// Compact the elements of `src` selected by `bits` to `dst` with a SIMD kernel, if one applies /// at `density`, and return the number of elements written. /// @@ -111,18 +56,40 @@ pub(super) fn filter_slice_mut_by_bitmap( /// `bits` must select only elements of `src`, and `dst` must hold every selected element plus /// [`SLACK_BYTES`]. #[inline] -pub(super) unsafe fn compress_bits( +pub(super) unsafe fn compress( src: &[T], bits: MaskBits<'_>, density: f64, dst: &mut [MaybeUninit], ) -> Option { let kernel = select_kernel::(src.len(), density)?; - // SAFETY: `select_kernel` probed the kernel's target features, and the caller upholds the - // out-of-place pointer contract of `filter_slice_by_bitmap`. + + // SAFETY: `select_kernel` probed the kernel's target features, `bits` selects only elements + // of `src`, and `dst` has room for every selected element plus a full vector of slack, so + // each unmasked store stays in bounds. Some(unsafe { kernel(src.as_ptr().cast(), dst.as_mut_ptr().cast(), bits) }) } +/// In-place variant of [`compress`]: compact the selected elements to the front of `values` and +/// return the new length, if a SIMD kernel applies at `density`. +/// +/// # Safety +/// +/// `bits` must select only elements of `values`. +pub(super) unsafe fn compress_in_place( + values: &mut [T], + bits: MaskBits<'_>, + density: f64, +) -> Option { + let kernel = select_kernel::(values.len(), density)?; + + let dst = values.as_mut_ptr().cast::(); + // SAFETY: `select_kernel` probed the kernel's target features; the in-place instantiation + // compacts forward (stores never pass the positions it has already read) and keeps partial + // tail chunks off the full-width store path, so all accesses stay inside `values`. + Some(unsafe { kernel(dst.cast_const(), dst, bits) }) +} + /// Choose the widest profitable kernel available for `T` to filter `len` elements at `density`. fn select_kernel(len: usize, density: f64) -> Option { if len < MIN_LEN { @@ -182,7 +149,7 @@ const fn compress_lut( /// /// # Safety /// -/// The pointer contract of [`filter_slice_by_bitmap`] / [`filter_slice_mut_by_bitmap`] must hold, +/// The pointer contract of [`compress`] / [`compress_in_place`] must hold, /// and `bits` must only select elements that are in bounds. #[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] #[allow(clippy::inline_always)] diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs b/vortex-buffer/src/filter/simd/neon.rs similarity index 92% rename from vortex-array/src/arrays/filter/execute/simd_compress/neon.rs rename to vortex-buffer/src/filter/simd/neon.rs index 171916be842..a1332adb834 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs +++ b/vortex-buffer/src/filter/simd/neon.rs @@ -12,9 +12,9 @@ use core::arch::aarch64::vst1_u8; use core::arch::aarch64::vst1q_u8; use core::arch::aarch64::vtbl1_u8; -use super::super::slice::MaskBits; -use super::super::slice::for_each_mask_word; -use super::super::slice::low_bits_mask; +use super::super::bits::MaskBits; +use super::super::bits::for_each_mask_word; +use super::super::bits::low_bits_mask; use super::Kernel; use super::bulk_copy; use super::compress_lut; @@ -64,8 +64,8 @@ macro_rules! neon_compress_kernel { ) => { /// # Safety /// - /// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// The pointer contract of [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. #[expect( clippy::cast_possible_truncation, reason = "deliberate submask narrowing" @@ -127,8 +127,8 @@ macro_rules! neon_compress_kernel { /// # Safety /// - /// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// The pointer contract of [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. pub(super) unsafe fn $walk_fn( src: *const u8, dst: *mut u8, diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs b/vortex-buffer/src/filter/simd/tests.rs similarity index 54% rename from vortex-array/src/arrays/filter/execute/simd_compress/tests.rs rename to vortex-buffer/src/filter/simd/tests.rs index 379f99459c6..dc5a5c6fbd9 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs +++ b/vortex-buffer/src/filter/simd/tests.rs @@ -4,31 +4,49 @@ #![allow(clippy::cast_possible_truncation)] use rstest::rstest; -use vortex_buffer::BitBuffer; -use vortex_mask::Mask; -use vortex_mask::MaskValues; use super::*; +use crate::BitBuffer; +use crate::filter::scalar::compact_by_bitmap; + +/// The scalar result of filtering `values` by `bits`. +fn expected(values: &[T], bits: &BitBuffer) -> Vec { + let mut out = Vec::with_capacity(values.len()); + // SAFETY: `bits` covers `values`, and the output has room for every element. + let written = + unsafe { compact_by_bitmap(values, MaskBits::Bitmap(bits), out.spare_capacity_mut()) }; + // SAFETY: the walk initialized the first `written` elements. + unsafe { out.set_len(written) }; + out +} -fn filter_slice_by_bitmap(values: &[T], mask: &MaskValues) -> Option> { - super::filter_slice_by_bitmap(values, mask, &BufferAllocatorRef::statically_allocated()) +/// Filter `values` out of place with the SIMD kernel, if one applies. +fn compress_out_of_place(values: &[T], bits: &BitBuffer) -> Option> { + let true_count = bits.true_count(); + let mut out = Vec::with_capacity(true_count + SLACK_BYTES / size_of::()); + // SAFETY: `bits` covers `values`, and the output has a vector of slack. + let written = unsafe { + compress( + values, + MaskBits::Bitmap(bits), + density(bits), + out.spare_capacity_mut(), + ) + }?; + assert_eq!(written, true_count); + // SAFETY: the kernel initialized the first `written` elements. + unsafe { out.set_len(written) }; + Some(out) } -fn mask_values(mask: &Mask) -> Option<&MaskValues> { - match mask { - Mask::Values(values) => Some(values.as_ref()), - _ => None, - } +fn density(bits: &BitBuffer) -> f64 { + bits.true_count() as f64 / bits.len() as f64 } -fn make_mask(len: usize, offset: usize, pattern: impl Fn(usize) -> bool) -> Mask { +fn make_bits(len: usize, offset: usize, pattern: impl Fn(usize) -> bool) -> BitBuffer { let backing = BitBuffer::from_iter(std::iter::repeat_n(false, offset).chain((0..len).map(pattern))); - Mask::from_buffer(BitBuffer::new_with_offset( - backing.inner().clone(), - len, - offset, - )) + BitBuffer::new_with_offset(backing.inner().clone(), len, offset) } type Pattern = fn(usize) -> bool; @@ -42,26 +60,18 @@ fn patterns() -> [Pattern; 4] { ] } -fn expected_values(values: &[T], mask: &MaskValues) -> Vec { - values - .iter() - .zip(mask.bit_buffer().iter()) - .filter_map(|(&value, selected)| selected.then_some(value)) - .collect() -} - -fn check(values: &[T], mask: &Mask) { - let Some(mask) = mask_values(mask) else { - return; - }; - let expected = expected_values(values, mask); +fn check(values: &[T], bits: &BitBuffer) { + let expected = expected(values, bits); - if let Some(actual) = filter_slice_by_bitmap(values, mask) { - assert_eq!(actual.as_slice(), expected.as_slice()); + if let Some(actual) = compress_out_of_place(values, bits) { + assert_eq!(actual, expected); } let mut compacted = values.to_vec(); - if let Some(new_len) = filter_slice_mut_by_bitmap(&mut compacted, mask) { + // SAFETY: `bits` covers `compacted`. + if let Some(new_len) = + unsafe { compress_in_place(&mut compacted, MaskBits::Bitmap(bits), density(bits)) } + { assert_eq!(&compacted[..new_len], expected.as_slice()); } } @@ -69,44 +79,48 @@ fn check(values: &[T], mask: &Mask) { #[rstest] fn simd_matches_scalar(#[values(0, 5)] offset: usize, #[values(64, 151)] len: usize) { for pattern in patterns() { - let mask = make_mask(len, offset, pattern); + let bits = make_bits(len, offset, pattern); let u8_values: Vec = (0..len).map(|i| i as u8).collect(); let u16_values: Vec = (0..len).map(|i| i as u16).collect(); let u32_values: Vec = (0..len).map(|i| i as u32).collect(); let u64_values: Vec = (0..len).map(|i| i as u64).collect(); - check(&u8_values, &mask); - check(&u16_values, &mask); - check(&u32_values, &mask); - check(&u64_values, &mask); + check(&u8_values, &bits); + check(&u16_values, &bits); + check(&u32_values, &bits); + check(&u64_values, &bits); } } #[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] fn check_engagement( values: &[T], - mask: &MaskValues, - expected: bool, + bits: &BitBuffer, + engaged: bool, ) { - let actual = filter_slice_by_bitmap(values, mask); - assert_eq!(actual.is_some(), expected); + let expected = expected(values, bits); + + let actual = compress_out_of_place(values, bits); + assert_eq!(actual.is_some(), engaged); if let Some(actual) = actual { - assert_eq!(actual.as_slice(), expected_values(values, mask)); + assert_eq!(actual, expected); } + let mut compacted = values.to_vec(); - let written = filter_slice_mut_by_bitmap(&mut compacted, mask); - assert_eq!(written.is_some(), expected); + // SAFETY: `bits` covers `compacted`. + let written = + unsafe { compress_in_place(&mut compacted, MaskBits::Bitmap(bits), density(bits)) }; + assert_eq!(written.is_some(), engaged); if let Some(written) = written { - assert_eq!(&compacted[..written], expected_values(values, mask)); + assert_eq!(&compacted[..written], expected.as_slice()); } } #[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] #[test] fn engages_on_supported_cpus() { - let mask = make_mask(256, 0, |i| i % 2 == 0); - let mask = mask_values(&mask).expect("alternating mask is mixed"); + let bits = make_bits(256, 0, |i| i % 2 == 0); #[cfg(target_arch = "x86_64")] let expected = is_x86_feature_detected!("avx2") && is_x86_feature_detected!("popcnt"); @@ -115,25 +129,23 @@ fn engages_on_supported_cpus() { check_engagement( &(0..256).map(|i| i as u8).collect::>(), - mask, + &bits, expected, ); - check_engagement(&(0u16..256).collect::>(), mask, expected); - check_engagement(&(0u32..256).collect::>(), mask, expected); - check_engagement(&(0u64..256).collect::>(), mask, expected); + check_engagement(&(0u16..256).collect::>(), &bits, expected); + check_engagement(&(0u32..256).collect::>(), &bits, expected); + check_engagement(&(0u64..256).collect::>(), &bits, expected); } #[test] fn declines_sparse_and_short_masks() { let values: Vec = (0..1024).collect(); - let sparse = make_mask(1024, 0, |i| i % 128 == 0); - let sparse = mask_values(&sparse).expect("sparse mask is mixed"); - assert!(filter_slice_by_bitmap(&values, sparse).is_none()); + let sparse = make_bits(1024, 0, |i| i % 128 == 0); + assert!(compress_out_of_place(&values, &sparse).is_none()); - let short = make_mask(32, 0, |i| i % 2 == 0); - let short = mask_values(&short).expect("alternating mask is mixed"); - assert!(filter_slice_by_bitmap(&values[..32], short).is_none()); + let short = make_bits(32, 0, |i| i % 2 == 0); + assert!(compress_out_of_place(&values[..32], &short).is_none()); } // AVX-512 machines need direct coverage of the otherwise-unselected AVX2-tier kernels. @@ -148,37 +160,34 @@ fn avx2_kernels_match_scalar() { kernel_out_of_place: Kernel, kernel_in_place: Kernel, values: &[T], - mask: &MaskValues, + bits: &BitBuffer, ) { - let expected = expected_values(values, mask); + let expected = expected(values, bits); - let mut out = vec![T::default(); mask.true_count() + SLACK_BYTES / size_of::()]; + let mut out = vec![T::default(); bits.true_count() + SLACK_BYTES / size_of::()]; // SAFETY: AVX2 and POPCNT were detected above and the output has a vector of slack. let written = unsafe { kernel_out_of_place( values.as_ptr().cast(), out.as_mut_ptr().cast(), - MaskBits::Mask(mask), + MaskBits::Bitmap(bits), ) }; - assert_eq!(written, mask.true_count()); + assert_eq!(written, bits.true_count()); assert_eq!(&out[..written], expected.as_slice()); let mut compacted = values.to_vec(); let ptr = compacted.as_mut_ptr().cast::(); // SAFETY: AVX2 and POPCNT were detected above; in-place compaction stays within the slice. - let written = unsafe { kernel_in_place(ptr.cast_const(), ptr, MaskBits::Mask(mask)) }; - assert_eq!(written, mask.true_count()); + let written = unsafe { kernel_in_place(ptr.cast_const(), ptr, MaskBits::Bitmap(bits)) }; + assert_eq!(written, bits.true_count()); assert_eq!(&compacted[..written], expected.as_slice()); } for pattern in patterns() { for len in [64, 151] { for offset in [0, 5] { - let mask = make_mask(len, offset, pattern); - let Some(mask) = mask_values(&mask) else { - continue; - }; + let bits = make_bits(len, offset, pattern); let u8_values: Vec = (0..len).map(|i| i as u8).collect(); let u16_values: Vec = (0..len).map(|i| i as u16).collect(); let u32_values: Vec = (0..len as u32).collect(); @@ -187,25 +196,25 @@ fn avx2_kernels_match_scalar() { generic::compress_generic_8::, generic::compress_generic_8::, &u8_values, - mask, + &bits, ); check_kernel( generic::compress_generic_16::, generic::compress_generic_16::, &u16_values, - mask, + &bits, ); check_kernel( x86::compress_avx2_epi32::, x86::compress_avx2_epi32::, &u32_values, - mask, + &bits, ); check_kernel( x86::compress_avx2_epi64::, x86::compress_avx2_epi64::, &u64_values, - mask, + &bits, ); } } diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs b/vortex-buffer/src/filter/simd/x86.rs similarity index 94% rename from vortex-array/src/arrays/filter/execute/simd_compress/x86.rs rename to vortex-buffer/src/filter/simd/x86.rs index cc6b79728ce..24d94ae6242 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs +++ b/vortex-buffer/src/filter/simd/x86.rs @@ -34,9 +34,9 @@ use std::arch::x86_64::_mm512_storeu_epi16; use std::arch::x86_64::_mm512_storeu_epi32; use std::arch::x86_64::_mm512_storeu_epi64; -use super::super::slice::MaskBits; -use super::super::slice::for_each_mask_word; -use super::super::slice::low_bits_mask; +use super::super::bits::MaskBits; +use super::super::bits::for_each_mask_word; +use super::super::bits::low_bits_mask; use super::Kernel; use super::bulk_copy; use super::generic::compress_generic_8; @@ -44,8 +44,8 @@ use super::generic::compress_generic_16; /// Choose the widest available kernel above its benchmarked density crossover. /// -/// Sparse masks stay on the scalar set-bit walk. See `benches/filter_fixed_width.rs` when -/// changing these thresholds. +/// Sparse masks stay on the scalar set-bit walk. See +/// `vortex-array/benches/filter_fixed_width.rs` when changing these thresholds. pub(super) fn select_kernel(density: f64) -> Option { let (kernel, min_density) = match size_of::() { 1 if avx512_vbmi2() => (compress_avx512_epi8:: as Kernel, 0.0), @@ -96,8 +96,8 @@ macro_rules! avx512_compress_kernel { /// # Safety /// /// The CPU must support the enabled target features and the pointer contract of - /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. // `allow` rather than `expect`: the cast is only lossy for sub-`u64` k-masks. #[allow(clippy::cast_possible_truncation)] #[target_feature(enable = $features)] @@ -165,8 +165,8 @@ macro_rules! avx512_compress_kernel { /// # Safety /// /// The CPU must support the enabled target features and the pointer contract of - /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. #[target_feature(enable = $features)] pub(super) unsafe fn $walk_fn( src: *const u8, @@ -323,8 +323,8 @@ macro_rules! avx2_compress_kernel { /// # Safety /// /// The CPU must support AVX2 and the pointer contract of - /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. #[expect( clippy::cast_possible_truncation, reason = "deliberate submask narrowing" @@ -407,8 +407,8 @@ macro_rules! avx2_compress_kernel { /// # Safety /// /// The CPU must support AVX2 and the pointer contract of - /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / - /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + /// [`compress`](super::compress) / + /// [`compress_in_place`](super::compress_in_place) must hold. #[target_feature(enable = "avx2,popcnt")] pub(super) unsafe fn $walk_fn( src: *const u8, diff --git a/vortex-buffer/src/filter/tests.rs b/vortex-buffer/src/filter/tests.rs new file mode 100644 index 00000000000..0ffe1b2971a --- /dev/null +++ b/vortex-buffer/src/filter/tests.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![allow(clippy::cast_possible_truncation)] + +use std::fmt::Debug; + +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rstest::rstest; + +use super::*; + +/// The representation of the test selection. +#[derive(Clone, Copy, Debug)] +enum Kind { + Bitmap, + Slices, + Indices, +} + +/// A selection of `len` positions, stored at bit `offset` of its bitmap. +fn make_bits(len: usize, offset: usize, pattern: impl Fn(usize) -> bool) -> BitBuffer { + let backing = + BitBuffer::from_iter(std::iter::repeat_n(false, offset).chain((0..len).map(pattern))); + BitBuffer::new_with_offset(backing.inner().clone(), len, offset) +} + +fn values>(len: usize) -> Vec { + (0..len).map(|i| T::from((i % 251) as u8)).collect() +} + +fn expected(values: &[T], bits: &BitBuffer) -> Vec { + values + .iter() + .zip(bits.iter()) + .filter_map(|(&value, selected)| selected.then_some(value)) + .collect() +} + +/// Filter `values` by `bits` in the representation `kind`, in place when `in_place` is set. +fn filter(values: &[T], bits: &BitBuffer, kind: Kind, in_place: bool) -> Buffer { + let slices = bits.set_slices().collect::>(); + let indices = bits.set_indices().collect::>(); + // SAFETY: all three representations come from `bits`. + let selection = unsafe { + match kind { + Kind::Bitmap => Selection::from_bitmap_unchecked(bits, bits.true_count()), + Kind::Slices => { + Selection::from_slices_unchecked(bits.len(), &slices, bits.true_count()) + } + Kind::Indices => Selection::from_indices_unchecked(bits.len(), &indices), + } + }; + + if in_place { + let mut values = values.to_vec(); + let new_len = filter_slice_in_place(&mut values, selection); + values.truncate(new_len); + Buffer::copy_from(values) + } else { + filter_slice(values, selection, BufferAllocatorRef::static_ref()) + } +} + +/// Compact `values` by `bits` with every scalar bitmap strategy, whatever the ladder chooses. +fn check_bitmap_strategies(values: &[T], bits: &BitBuffer) { + let expected = expected(values, bits); + let bits = MaskBits::Bitmap(bits); + + let mut out = Vec::with_capacity(values.len()); + // SAFETY: `bits` covers `values`, and the output has room for every value. + let written = unsafe { scalar::compact_by_bitmap(values, bits, out.spare_capacity_mut()) }; + // SAFETY: the walk initialized the first `written` values. + unsafe { out.set_len(written) }; + assert_eq!(out, expected); + + let mut out = Vec::with_capacity(values.len()); + let written = scalar::compact_runs_by_bitmap(values, bits, out.spare_capacity_mut()); + // SAFETY: the walk initialized the first `written` values. + unsafe { out.set_len(written) }; + assert_eq!(out, expected); + + let mut out = Vec::with_capacity(values.len()); + // SAFETY: `bits` covers `values`, and the output has room for every value. + let written = unsafe { byte_compress::compress(values, bits, out.spare_capacity_mut()) }; + // SAFETY: the walk initialized the first `written` values. + unsafe { out.set_len(written) }; + assert_eq!(out, expected); + + let mut compacted = values.to_vec(); + let new_len = scalar::compact_in_place_by_bitmap(&mut compacted, bits); + assert_eq!(&compacted[..new_len], expected.as_slice()); +} + +fn check>(bits: &BitBuffer) { + let values = values::(bits.len()); + let expected = expected(&values, bits); + + for kind in [Kind::Bitmap, Kind::Slices, Kind::Indices] { + for in_place in [false, true] { + assert_eq!( + filter(&values, bits, kind, in_place).as_slice(), + expected.as_slice(), + "{kind:?}, in place: {in_place}" + ); + } + } + + check_bitmap_strategies(&values, bits); +} + +#[rstest] +fn filter_matches_expected( + #[values(0.001, 0.05, 0.2, 0.5, 0.9, 0.999)] density: f64, + #[values(0, 3, 64)] offset: usize, + #[values(9, 151, 2000)] len: usize, +) { + let mut rng = StdRng::seed_from_u64(42); + let selected = (0..len) + .map(|_| rng.random_bool(density)) + .collect::>(); + let bits = make_bits(len, offset, |i| selected[i]); + + check::(&bits); + check::(&bits); + check::(&bits); + check::(&bits); + check::(&bits); +} + +#[rstest] +fn filter_runs_matches_expected( + #[values(1, 12, 200)] run_len: usize, + #[values(1, 16, 300)] gap: usize, + #[values(0, 5)] offset: usize, +) { + let bits = make_bits(2000, offset, |i| i % (run_len + gap) < run_len); + + check::(&bits); + check::(&bits); + check::(&bits); + check::(&bits); + check::(&bits); +} + +#[test] +fn filter_buffer_in_place_and_shared() { + let bits = make_bits(200, 3, |i| i % 7 != 0); + let expected = expected(&values::(200), &bits); + // SAFETY: the true count comes from `bits`. + let selection = unsafe { Selection::from_bitmap_unchecked(&bits, bits.true_count()) }; + + let unique = Buffer::copy_from(values::(200)); + let unique_ptr = unique.as_ptr(); + let filtered = filter_buffer(unique, selection, BufferAllocatorRef::static_ref()); + assert_eq!(filtered.as_slice(), expected.as_slice()); + assert_eq!( + filtered.as_ptr(), + unique_ptr, + "a dense unique buffer filters in place" + ); + + let shared = Buffer::copy_from(values::(200)); + let other = shared.clone(); + let filtered = filter_buffer(shared, selection, BufferAllocatorRef::static_ref()); + assert_eq!(filtered.as_slice(), expected.as_slice()); + assert_eq!(other.as_slice(), values::(200).as_slice()); +} + +#[test] +fn filter_buffer_empty() { + let bits = BitBuffer::from_iter(std::iter::empty::()); + // SAFETY: an empty bitmap has no set bits. + let selection = unsafe { Selection::from_bitmap_unchecked(&bits, 0) }; + + let filtered = filter_buffer( + Buffer::::empty(), + selection, + BufferAllocatorRef::static_ref(), + ); + assert!(filtered.is_empty()); +} diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index 33e68c1a327..05247463919 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -65,6 +65,7 @@ mod bytes; mod const_; mod debug; mod dispatch; +pub mod filter; mod macros; #[cfg(feature = "memmap2")] mod memmap2; diff --git a/vortex-mask/src/lib.rs b/vortex-mask/src/lib.rs index 14dc417e5bc..e8804057557 100644 --- a/vortex-mask/src/lib.rs +++ b/vortex-mask/src/lib.rs @@ -23,6 +23,7 @@ use itertools::Itertools; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::BitIterator; +use vortex_buffer::filter::Selection; use vortex_error::VortexResult; use vortex_error::vortex_panic; @@ -866,8 +867,48 @@ impl MaskValues { MaskIter::Indices(self.indices()) } } + + /// Returns the representation of the mask that is cheapest to filter with, for + /// [`vortex_buffer::filter`]. + /// + /// This uses the cached slices if there is one slice or their runs average at least 8 values, + /// then the cached indices if the density is at most 0.5, and otherwise the bitmap. It does + /// not build a representation that is not cached. + pub fn selection(&self) -> Selection<'_> { + // SAFETY: the mask validated its cached slices and indices when it created them, and it + // holds the true count of its bitmap. + unsafe { + if let Some(slices) = self.useful_cached_slices() { + Selection::from_slices_unchecked(self.len(), slices, self.true_count) + } else if let Some(indices) = self.useful_cached_indices() { + Selection::from_indices_unchecked(self.len(), indices) + } else { + Selection::from_bitmap_unchecked(&self.buffer, self.true_count) + } + } + } + + fn useful_cached_slices(&self) -> Option<&[(usize, usize)]> { + self.cached_slices().filter(|slices| { + slices.len() == 1 || self.true_count / slices.len() >= SELECTION_MIN_AVERAGE_RUN_LEN + }) + } + + fn useful_cached_indices(&self) -> Option<&[usize]> { + if self.density <= SELECTION_INDICES_MAX_DENSITY { + self.cached_indices() + } else { + None + } + } } +/// The smallest average run length for which [`MaskValues::selection`] uses the cached slices. +const SELECTION_MIN_AVERAGE_RUN_LEN: usize = 8; + +/// The largest density for which [`MaskValues::selection`] uses the cached indices. +const SELECTION_INDICES_MAX_DENSITY: f64 = 0.5; + /// Iterator over the indices or slices of a mask. pub enum MaskIter<'a> { /// Slice of pre-cached indices of a mask. diff --git a/vortex-mask/src/tests.rs b/vortex-mask/src/tests.rs index e8cf1fc1ee5..0920891cb91 100644 --- a/vortex-mask/src/tests.rs +++ b/vortex-mask/src/tests.rs @@ -3,6 +3,9 @@ use rstest::rstest; use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::filter::filter_buffer; use crate::AllOr; use crate::Mask; @@ -788,3 +791,47 @@ fn mask_iter_matches_value(#[case] mask: Mask, #[case] expected: Vec) { assert_eq!(it.len(), mask.len() - 1); } } + +/// The representation that a test mask caches before [`MaskValues::selection`] reads it. +#[derive(Clone, Copy, Debug)] +enum Cached { + Nothing, + Indices, + Slices, +} + +/// Every cached representation, at densities and run lengths on both sides of the selection +/// thresholds, filters to the same values as the bitmap. +#[rstest] +fn selection_filters_like_bitmap( + #[values(Cached::Nothing, Cached::Indices, Cached::Slices)] cached: Cached, + #[values(|i| i % 10 == 0, |i| i % 10 != 0, |i| i % 32 < 4, |i| i % 32 < 16)] + pattern: fn(usize) -> bool, +) { + let mask = Mask::from_buffer((0..1000).map(pattern).collect::()); + let Mask::Values(values) = &mask else { + unreachable!("the patterns select some but not all values") + }; + match cached { + Cached::Nothing => {} + Cached::Indices => { + values.indices(); + } + Cached::Slices => { + values.slices(); + } + } + + let selection = values.selection(); + assert_eq!(selection.len(), 1000); + assert_eq!(selection.true_count(), values.true_count()); + + let data = (0..1000u32).collect::>(); + let expected = data + .iter() + .zip(values.bit_buffer().iter()) + .filter_map(|(&value, selected)| selected.then_some(value)) + .collect::>(); + let filtered = filter_buffer(data, selection, BufferAllocatorRef::static_ref()); + assert_eq!(filtered.as_slice(), expected.as_slice()); +} From 86be634dbcd55346bf786bcbc4e5b7187ea74c08 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 28 Oct 2026 15:08:17 +0000 Subject: [PATCH 2/4] Filter runtime-width fixed-width records with vortex-buffer filter Add vortex_buffer::filter::filter_records, which filters records whose byte width is known only at run time. Power-of-two widths use filter_buffer, and other widths copy each run of the selection. The fixed-width array filter now uses it instead of its own bitmap loop. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/filter/mod.rs | 1 - vortex-array/src/arrays/fixed_width/filter.rs | 43 +--- vortex-buffer/src/filter/mod.rs | 5 +- vortex-buffer/src/filter/records.rs | 228 ++++++++++++++++++ 4 files changed, 234 insertions(+), 43 deletions(-) create mode 100644 vortex-buffer/src/filter/records.rs diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index fb7dc6a4141..3dd5968bf68 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -10,7 +10,6 @@ pub use array::FilterSlotsView; pub use vtable::FilterArray; mod execute; -pub(crate) use execute::buffer::filter_buffer; pub(crate) use execute::buffer::prepare_mask_for_reuse; pub(crate) use execute::filter_validity; diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs index 1c2b16b77d1..486afbcf410 100644 --- a/vortex-array/src/arrays/fixed_width/filter.rs +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -1,10 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; +use vortex_buffer::filter; use vortex_error::VortexExpect; use vortex_mask::MaskValues; use vortex_mask::MaskValuesRef; @@ -12,9 +11,7 @@ use vortex_mask::MaskValuesRef; use super::FixedWidthArray; use super::with_values; use crate::array::Array; -use crate::arrays::filter::filter_buffer; use crate::arrays::filter::filter_validity; -use crate::match_each_power_of_two_width; #[cfg(test)] #[expect(clippy::cast_possible_truncation)] @@ -48,41 +45,5 @@ fn filter_records( mask: &MaskValues, allocator: &BufferAllocatorRef, ) -> ByteBuffer { - let alignment = values.alignment(); - - match_each_power_of_two_width!( - byte_width, - |W| { - let records = Buffer::<[u8; W]>::from_byte_buffer(values); - // `filter_buffer` picks between in-place compaction, cached indices/slices, - // byte-compress, and bitmap iteration based on record width and mask density. - let filtered = filter_buffer(records, mask, allocator); - filtered.into_byte_buffer().aligned(alignment) - }, - _ => { - match values.try_into_mut() { - Ok(mut values) => { - let mut destination = 0; - mask.bit_buffer().for_each_set_index(|index| { - let source = index * byte_width; - values.copy_within(source..source + byte_width, destination); - destination += byte_width; - }); - values.truncate(destination); - values.freeze().into_byte_buffer().aligned(alignment) - } - Err(values) => { - let mut filtered = BufferMut::with_capacity_in( - mask.true_count() * byte_width, - allocator.clone(), - ); - mask.bit_buffer().for_each_set_index(|index| { - let start = index * byte_width; - filtered.extend_from_slice(&values[start..start + byte_width]); - }); - filtered.freeze().into_byte_buffer().aligned(alignment) - } - } - } - ) + filter::filter_records(values, byte_width, mask.selection(), allocator) } diff --git a/vortex-buffer/src/filter/mod.rs b/vortex-buffer/src/filter/mod.rs index 5f9c2d20315..98a66668f0a 100644 --- a/vortex-buffer/src/filter/mod.rs +++ b/vortex-buffer/src/filter/mod.rs @@ -5,7 +5,8 @@ //! //! [`filter_buffer`] filters values that are already in memory, and [`filter_chunked`] filters //! values that a [`ChunkDecoder`] decodes one chunk at a time. Both work for any `T: Copy` and -//! share one strategy ladder. The first matching row wins: +//! share one strategy ladder. [`filter_records`] filters records whose byte width is known only +//! at run time, and uses [`filter_buffer`] for power-of-two widths. The first matching row wins: //! //! | Priority | Selection or condition | Out-of-place | In-place | //! | -------- | --------------------------------------------------- | -------------- | ------------- | @@ -30,6 +31,7 @@ use crate::BufferMut; mod bits; mod byte_compress; mod chunked; +mod records; mod scalar; mod simd; #[cfg(test)] @@ -39,6 +41,7 @@ use bits::MaskBits; pub use chunked::ChunkDecoder; pub use chunked::FILTER_CHUNK_LEN; pub use chunked::filter_chunked; +pub use records::filter_records; const IN_PLACE_MIN_DENSITY: f64 = 0.5; diff --git a/vortex-buffer/src/filter/records.rs b/vortex-buffer/src/filter/records.rs new file mode 100644 index 00000000000..a25280a1f20 --- /dev/null +++ b/vortex-buffer/src/filter/records.rs @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Filtering of records whose byte width is known only at run time. See [`filter_records`]. + +use super::IN_PLACE_MIN_DENSITY; +use super::Repr; +use super::Selection; +use super::bits::MaskBits; +use super::bits::for_each_mask_word; +use super::filter_buffer; +use crate::Buffer; +use crate::BufferAllocatorRef; +use crate::ByteBuffer; +use crate::ByteBufferMut; +use crate::match_each_power_of_two_width; + +/// Filters the records of `width` bytes in `values` by `selection` and returns the selected +/// records. +/// +/// Power-of-two widths up to 32 bytes use [`filter_buffer`]. Other widths copy each run of +/// selected records with one copy. A uniquely owned buffer with density of at least 0.50 is +/// compacted in place, and other buffers are copied to a new allocation from `allocator`. The +/// result keeps the alignment of `values`. +/// +/// # Panics +/// +/// Panics if `values` does not hold one record of `width` bytes for each position of +/// `selection`. +pub fn filter_records( + values: ByteBuffer, + width: usize, + selection: Selection<'_>, + allocator: &BufferAllocatorRef, +) -> ByteBuffer { + assert_eq!( + values.len(), + selection.len() * width, + "Selection length must equal the number of records" + ); + + let alignment = values.alignment(); + + match_each_power_of_two_width!( + width, + |W| { + let records = Buffer::<[u8; W]>::from_byte_buffer(values); + filter_buffer(records, selection, allocator) + .into_byte_buffer() + .aligned(alignment) + }, + _ => { filter_runtime_width(values, width, selection, allocator).aligned(alignment) } + ) +} + +fn filter_runtime_width( + values: ByteBuffer, + width: usize, + selection: Selection<'_>, + allocator: &BufferAllocatorRef, +) -> ByteBuffer { + let values = if selection.density() >= IN_PLACE_MIN_DENSITY { + match values.try_into_mut() { + Ok(mut values_mut) => { + let new_len = compact_in_place(values_mut.as_mut_slice(), width, selection); + values_mut.truncate(new_len); + return values_mut.freeze(); + } + Err(values) => values, + } + } else { + values + }; + + let mut filtered = + ByteBufferMut::with_capacity_in(selection.true_count() * width, allocator.clone()); + for_each_run(selection, |start, len| { + filtered.extend_from_slice(&values[start * width..(start + len) * width]); + }); + + debug_assert_eq!(filtered.len(), selection.true_count() * width); + filtered.freeze() +} + +/// Compacts the records of `width` bytes that `selection` keeps to the front of `values`, and +/// returns the new length in bytes. +fn compact_in_place(values: &mut [u8], width: usize, selection: Selection<'_>) -> usize { + let mut write_pos = 0; + + for_each_run(selection, |start, len| { + let source = start * width; + let run_bytes = len * width; + if write_pos != source { + values.copy_within(source..source + run_bytes, write_pos); + } + write_pos += run_bytes; + }); + + debug_assert_eq!(write_pos, selection.true_count() * width); + write_pos +} + +/// Invokes `f` with the `(start, len)` of each run of positions that `selection` keeps, in +/// order. +/// +/// A run of set bits that crosses a word boundary of a bitmap is split at the boundary. +fn for_each_run(selection: Selection<'_>, mut f: impl FnMut(usize, usize)) { + match selection.repr { + Repr::Slices(slices) => { + for &(start, end) in slices { + f(start, end - start); + } + } + Repr::Indices(indices) => { + for &index in indices { + f(index, 1); + } + } + Repr::Bitmap(bits) => { + for_each_mask_word(MaskBits::Bitmap(bits), |mut word, word_start, _| { + while word != 0 { + let run_start = word.trailing_zeros() as usize; + let run_len = (!(word >> run_start)).trailing_zeros() as usize; + f(word_start + run_start, run_len); + + // Adding the lowest bit of the run carries through the run and clears it. + word &= word.wrapping_add(1 << run_start); + } + }); + } + } +} + +#[cfg(test)] +#[expect(clippy::cast_possible_truncation)] +mod tests { + use rstest::rstest; + + use super::*; + use crate::BitBuffer; + + /// A selection of `len` positions, stored at bit `offset` of its bitmap. + fn make_bits(len: usize, offset: usize, pattern: impl Fn(usize) -> bool) -> BitBuffer { + let backing = + BitBuffer::from_iter(std::iter::repeat_n(false, offset).chain((0..len).map(pattern))); + BitBuffer::new_with_offset(backing.inner().clone(), len, offset) + } + + fn records(len: usize, width: usize) -> Vec { + (0..len * width).map(|i| (i % 251) as u8).collect() + } + + fn expected(values: &[u8], width: usize, bits: &BitBuffer) -> Vec { + values + .chunks(width) + .zip(bits.iter()) + .filter_map(|(record, selected)| selected.then_some(record)) + .flatten() + .copied() + .collect() + } + + #[rstest] + fn filter_records_matches_expected( + #[values(3, 4, 12, 24, 48)] width: usize, + #[values(0, 5)] offset: usize, + #[values( + |i: usize| i.is_multiple_of(3), + |i: usize| !i.is_multiple_of(16), + |i: usize| (i / 70).is_multiple_of(2), + )] + pattern: fn(usize) -> bool, + #[values(false, true)] shared: bool, + ) { + const LEN: usize = 151; + + let bits = make_bits(LEN, offset, pattern); + let slices = bits.set_slices().collect::>(); + let indices = bits.set_indices().collect::>(); + let values = records(LEN, width); + let expected = expected(&values, width, &bits); + + // SAFETY: all three representations come from `bits`. + let selections = unsafe { + [ + Selection::from_bitmap_unchecked(&bits, bits.true_count()), + Selection::from_slices_unchecked(LEN, &slices, bits.true_count()), + Selection::from_indices_unchecked(LEN, &indices), + ] + }; + + for selection in selections { + let buffer = ByteBuffer::copy_from(values.as_slice()); + let _retained = shared.then(|| buffer.clone()); + let filtered = + filter_records(buffer, width, selection, BufferAllocatorRef::static_ref()); + assert_eq!(filtered.as_slice(), expected.as_slice()); + } + } + + #[test] + fn filter_records_in_place() { + let bits = make_bits(8, 0, |i| i % 4 != 0); + // SAFETY: the true count comes from `bits`. + let selection = unsafe { Selection::from_bitmap_unchecked(&bits, bits.true_count()) }; + + let values = ByteBuffer::copy_from(records(8, 3)); + let original_ptr = values.as_ptr(); + let filtered = filter_records(values, 3, selection, BufferAllocatorRef::static_ref()); + assert_eq!(filtered.as_ptr(), original_ptr); + assert_eq!(filtered.as_slice(), expected(&records(8, 3), 3, &bits)); + } + + #[test] + fn filter_zero_width_records() { + let bits = make_bits(10, 0, |i| i % 2 == 0); + // SAFETY: the true count comes from `bits`. + let selection = unsafe { Selection::from_bitmap_unchecked(&bits, bits.true_count()) }; + + let filtered = filter_records( + ByteBuffer::empty(), + 0, + selection, + BufferAllocatorRef::static_ref(), + ); + assert!(filtered.is_empty()); + } +} From d6bd74a7097240ed8d65ac520eba39972af62a97 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 28 Oct 2026 18:23:00 +0000 Subject: [PATCH 3/4] Inline per-word filter closures into large callers In `filter_records`, LLVM kept the per-word closures of byte compress and the bitmap walk out of line, because the word walk calls them from three places. Sparse masks then made one call for each mask word. Mark the closures `inline(always)`. Signed-off-by: Robert Kruszewski --- vortex-buffer/src/filter/byte_compress.rs | 5 ++++- vortex-buffer/src/filter/scalar.rs | 3 ++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/vortex-buffer/src/filter/byte_compress.rs b/vortex-buffer/src/filter/byte_compress.rs index 054c39351dd..c0b732f3827 100644 --- a/vortex-buffer/src/filter/byte_compress.rs +++ b/vortex-buffer/src/filter/byte_compress.rs @@ -51,7 +51,10 @@ pub(super) unsafe fn compress( ) -> usize { let mut written = 0; - for_each_mask_word(bits, |word, word_start, word_len| { + // The word walk calls this closure from three places. Without `inline(always)`, LLVM keeps + // it out of line in large callers such as `filter_records`, and sparse masks then pay for + // one call per word. + for_each_mask_word(bits, #[inline(always)] |word, word_start, word_len| { let (chunks, tail) = src[word_start..][..word_len].as_chunks::<8>(); let mask_bytes = word.to_le_bytes(); diff --git a/vortex-buffer/src/filter/scalar.rs b/vortex-buffer/src/filter/scalar.rs index 7dd33d8bd9f..7f3bfc8cdac 100644 --- a/vortex-buffer/src/filter/scalar.rs +++ b/vortex-buffer/src/filter/scalar.rs @@ -23,7 +23,8 @@ pub(super) unsafe fn compact_by_bitmap( ) -> usize { let mut write_pos = 0; - for_each_mask_word(bits, |word, word_start, word_len| { + // See `byte_compress::compress` for why this closure is always inlined. + for_each_mask_word(bits, #[inline(always)] |word, word_start, word_len| { let all_selected = low_bits_mask(word_len); debug_assert_eq!(word & !all_selected, 0); From 2664b5eb36b35021cab0c55a50c23a5bf34afedc Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 29 Oct 2026 11:10:22 +0000 Subject: [PATCH 4/4] format Signed-off-by: Robert Kruszewski --- vortex-buffer/src/filter/byte_compress.rs | 45 ++++++++++++---------- vortex-buffer/src/filter/scalar.rs | 46 ++++++++++++----------- 2 files changed, 50 insertions(+), 41 deletions(-) diff --git a/vortex-buffer/src/filter/byte_compress.rs b/vortex-buffer/src/filter/byte_compress.rs index c0b732f3827..99af2963f8e 100644 --- a/vortex-buffer/src/filter/byte_compress.rs +++ b/vortex-buffer/src/filter/byte_compress.rs @@ -54,28 +54,33 @@ pub(super) unsafe fn compress( // The word walk calls this closure from three places. Without `inline(always)`, LLVM keeps // it out of line in large callers such as `filter_records`, and sparse masks then pay for // one call per word. - for_each_mask_word(bits, #[inline(always)] |word, word_start, word_len| { - let (chunks, tail) = src[word_start..][..word_len].as_chunks::<8>(); - let mask_bytes = word.to_le_bytes(); + for_each_mask_word( + bits, + #[inline(always)] + |word, word_start, word_len| { + let (chunks, tail) = src[word_start..][..word_len].as_chunks::<8>(); + let mask_bytes = word.to_le_bytes(); - for (chunk, &mask_byte) in chunks.iter().zip(&mask_bytes) { - // SAFETY: the mask byte selects only elements of `chunk`, and `dst` has room for - // every selected element. - written += unsafe { compress_byte(chunk, mask_byte, dst.get_unchecked_mut(written..)) }; - } + for (chunk, &mask_byte) in chunks.iter().zip(&mask_bytes) { + // SAFETY: the mask byte selects only elements of `chunk`, and `dst` has room for + // every selected element. + written += + unsafe { compress_byte(chunk, mask_byte, dst.get_unchecked_mut(written..)) }; + } - if !tail.is_empty() { - // SAFETY: the bits past `word_len` are cleared, so the mask byte selects only - // elements of `tail`. - written += unsafe { - compress_byte( - tail, - mask_bytes[chunks.len()], - dst.get_unchecked_mut(written..), - ) - }; - } - }); + if !tail.is_empty() { + // SAFETY: the bits past `word_len` are cleared, so the mask byte selects only + // elements of `tail`. + written += unsafe { + compress_byte( + tail, + mask_bytes[chunks.len()], + dst.get_unchecked_mut(written..), + ) + }; + } + }, + ); written } diff --git a/vortex-buffer/src/filter/scalar.rs b/vortex-buffer/src/filter/scalar.rs index 7f3bfc8cdac..66a09022b66 100644 --- a/vortex-buffer/src/filter/scalar.rs +++ b/vortex-buffer/src/filter/scalar.rs @@ -24,28 +24,32 @@ pub(super) unsafe fn compact_by_bitmap( let mut write_pos = 0; // See `byte_compress::compress` for why this closure is always inlined. - for_each_mask_word(bits, #[inline(always)] |word, word_start, word_len| { - let all_selected = low_bits_mask(word_len); - debug_assert_eq!(word & !all_selected, 0); - - if word == all_selected { - dst[write_pos..][..word_len].write_copy_of_slice(&src[word_start..][..word_len]); - write_pos += word_len; - } else { - let mut selected = word; - while selected != 0 { - let index = word_start + selected.trailing_zeros() as usize; - // SAFETY: set bits index into `src`, and `dst` has room for every selected - // element. - unsafe { - dst.get_unchecked_mut(write_pos) - .write(*src.get_unchecked(index)) - }; - write_pos += 1; - selected &= selected - 1; + for_each_mask_word( + bits, + #[inline(always)] + |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + + if word == all_selected { + dst[write_pos..][..word_len].write_copy_of_slice(&src[word_start..][..word_len]); + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + // SAFETY: set bits index into `src`, and `dst` has room for every selected + // element. + unsafe { + dst.get_unchecked_mut(write_pos) + .write(*src.get_unchecked(index)) + }; + write_pos += 1; + selected &= selected - 1; + } } - } - }); + }, + ); write_pos }