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..3dd5968bf68 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -10,11 +10,7 @@ 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 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-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/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..99af2963f8e --- /dev/null +++ b/vortex-buffer/src/filter/byte_compress.rs @@ -0,0 +1,116 @@ +// 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; + + // 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 (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..98a66668f0a --- /dev/null +++ b/vortex-buffer/src/filter/mod.rs @@ -0,0 +1,332 @@ +// 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. [`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 | +//! | -------- | --------------------------------------------------- | -------------- | ------------- | +//! | 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 records; +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; +pub use records::filter_records; + +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/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()); + } +} diff --git a/vortex-buffer/src/filter/scalar.rs b/vortex-buffer/src/filter/scalar.rs new file mode 100644 index 00000000000..66a09022b66 --- /dev/null +++ b/vortex-buffer/src/filter/scalar.rs @@ -0,0 +1,182 @@ +// 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; + + // 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; + } + } + }, + ); + + 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()); +}