diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index b37a9f5d6c8..c9389800995 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -51,6 +51,10 @@ _test-harness = ["dep:rand"] name = "bitpacking_take" harness = false +[[bench]] +name = "bitpacking_filter" +harness = false + [[bench]] name = "canonicalize_bench" harness = false diff --git a/encodings/fastlanes/benches/bitpacking_filter.rs b/encodings/fastlanes/benches/bitpacking_filter.rs new file mode 100644 index 00000000000..865a3d72ba2 --- /dev/null +++ b/encodings/fastlanes/benches/bitpacking_filter.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks filtering bit-packed arrays, both directly and as the elements of a +//! `FixedSizeList`, where the list-level selection expands into runs of selected elements. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::RngExt; +use rand::SeedableRng; +use rand::prelude::StdRng; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPackedData; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +const NUM_ELEMENTS: usize = 1 << 17; +const BIT_WIDTH: u8 = 16; + +fn bitpacked_i32(len: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(0); + let values = (0..len) + .map(|_| rng.random_range(0..(1i32 << BIT_WIDTH))) + .collect::>(); + let array = PrimitiveArray::new(values, Validity::NonNullable).into_array(); + BitPackedData::encode(&array, BIT_WIDTH, &mut SESSION.create_execution_ctx()) + .unwrap() + .into_array() +} + +fn random_bits(len: usize, density: f64) -> BitBuffer { + let mut rng = StdRng::seed_from_u64(1); + BitBuffer::from_iter((0..len).map(|_| rng.random_bool(density))) +} + +fn bench_filter(bencher: Bencher, array: &ArrayRef, bits: &BitBuffer) { + // Build a fresh mask for every iteration so that representations cached on the mask, such as + // its indices, are not reused across iterations. + bencher + .with_inputs(|| { + ( + array, + Mask::from_buffer(bits.clone()), + SESSION.create_execution_ctx(), + ) + }) + .bench_values(|(array, mask, mut ctx)| { + array + .filter(mask) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench(args = [0.01, 0.1])] +fn primitive_i32(bencher: Bencher, density: f64) { + let array = bitpacked_i32(NUM_ELEMENTS); + bench_filter(bencher, &array, &random_bits(NUM_ELEMENTS, density)); +} + +#[divan::bench(consts = [4, 16], args = [0.1])] +fn fsl_i32(bencher: Bencher, density: f64) { + let len = NUM_ELEMENTS / LIST_SIZE as usize; + let elements = bitpacked_i32(NUM_ELEMENTS); + let array = + FixedSizeListArray::new(elements, LIST_SIZE, Validity::NonNullable, len).into_array(); + bench_filter(bencher, &array, &random_bits(len, density)); +} diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 952e7728e52..37a289c65be 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -9,41 +9,26 @@ 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::dtype::PType; -use vortex_array::dtype::UnsignedPType; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_array::validity::Validity; use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_mask::MaskValuesRef; +use vortex_mask::MaskValues; -use super::chunked_indices; -use super::take::UNPACK_CHUNK_THRESHOLD; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitPackedData; use crate::BitWidthsView; -/// The threshold over which it is faster to fully unpack the entire [`BitPackedArray`](crate::BitPackedArray) and then -/// filter the result than to unpack only specific bitpacked values into the output buffer. -pub const fn unpack_then_filter_threshold(ptype: PType) -> f64 { - // TODO(connor): Where did these numbers come from? Add a public link after validating them. - // These numbers probably don't work for in-place filtering either. - match ptype.byte_width() { - 1 => 0.03, - 2 => 0.03, - 4 => 0.075, - _ => 0.09, - // >8 bytes may have a higher threshold. These numbers are derived from a GCP c2-standard-4 - // with a "Cascade Lake" CPU. - } -} - /// Kernel to execute filtering directly on a bit-packed array. +/// +/// [`filter_chunked`] unpacks only the FastLanes chunks that hold selected values, and compacts +/// each chunk while it is in cache, rather than unpack the whole array and then filter it. impl FilterKernel for BitPacked { fn filter( array: ArrayView<'_, Self>, @@ -60,142 +45,116 @@ impl FilterKernel for BitPacked { Mask::Values(values) => values, }; - // If the density is high enough, then we would rather decompress the whole array and then apply - // a filter over decompressing values one by one. - if values.density() > unpack_then_filter_threshold(array.dtype().as_ptype()) { - return Ok(None); - } - - // Filter and patch using the correct unsigned type for FastLanes, then cast to signed if needed. - let primitive = - match_each_unsigned_integer_ptype!(array.dtype().as_ptype().to_unsigned(), |U| { - let (buffer, validity) = - filter_primitive_without_patches::(array, bit_width, values)?; - // reinterpret_cast for signed types. - let primitive = PrimitiveArray::new(buffer, validity); - if array.dtype().as_ptype().is_signed_int() { - PrimitiveArray::from_buffer_handle( - primitive.buffer_handle().clone(), - array.dtype().as_ptype(), - primitive.validity()?, - ) - } else { - primitive - } - }); + // FastLanes only unpacks unsigned types, so filter as unsigned and reinterpret the + // resulting buffer with the array's (possibly signed) ptype. + let ptype = array.dtype().as_ptype(); + let validity = array.validity()?.filter(mask)?; + let buffer = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |U| { + filter_values::(array.data(), bit_width, values).into_byte_buffer() + }); + let primitive = PrimitiveArray::from_byte_buffer(buffer, ptype, validity); let patches = array .patches() - .map(|patches| patches.filter(&Mask::Values(MaskValuesRef::clone(values)), ctx)) + .map(|patches| patches.filter(mask, ctx)) .transpose()? .flatten(); - if let Some(patches) = patches { - let mut prim_array = primitive; - prim_array = prim_array.patch(&patches, ctx)?; - return Ok(Some(prim_array.into_array())); - } - - Ok(Some(primitive.into_array())) + Ok(Some(match patches { + Some(patches) => primitive.patch(&patches, ctx)?.into_array(), + None => primitive.into_array(), + })) } } -/// Specialized filter kernel for primitive bit-packed arrays. -/// -/// Because the FastLanes bit-packing kernels are only implemented for unsigned types, the provided -/// `U` should be promoted to the unsigned variant for any target bit width. -/// For example, if the array is bit-packed `i16`, this function should be called with `U = u16`. -/// -/// This function fully decompresses the array for all but the most selective masks because the -/// FastLanes decompression is so fast and the bookkeepping necessary to decompress individual -/// elements is relatively slow. +/// Unpacks the values of `array` selected by `mask`, ignoring patches and validity. /// -/// Returns a tuple of (values buffer, validity mask). -fn filter_primitive_without_patches( - array: ArrayView<'_, BitPacked>, - bit_width: u8, - selection: &MaskValuesRef, -) -> VortexResult<(Buffer, Validity)> { - let values = filter_with_indices(array.data(), bit_width, selection.indices()); - let validity = array - .validity()? - .filter(&Mask::Values(MaskValuesRef::clone(selection)))?; - - Ok((values.freeze(), validity)) -} - -fn filter_with_indices( +/// Because the FastLanes bit-packing kernels are only implemented for unsigned types, `T` must +/// be the unsigned variant of the array's ptype. +fn filter_values( array: &BitPackedData, bit_width: u8, - indices: &[usize], -) -> BufferMut { - let offset = array.offset() as usize; + mask: &MaskValues, +) -> Buffer { let bit_width = bit_width as usize; - let mut values = BufferMut::with_capacity(indices.len()); - - // Some re-usable memory to store per-chunk indices. - let mut unpacked = [const { MaybeUninit::::uninit() }; 1024]; - let packed_bytes = array.packed_slice::(); - - // Group the indices by the FastLanes chunk they belong to. - let chunk_size = 128 * bit_width / size_of::(); - - chunked_indices( - indices.iter().copied(), - offset, - |chunk_idx, indices_within_chunk| { - let packed = &packed_bytes[chunk_idx * chunk_size..][..chunk_size]; - - if indices_within_chunk.len() == 1024 { - // Unpack the entire chunk. - unsafe { - let values_len = values.len(); - values.set_len(values_len + 1024); - BitPacking::unchecked_unpack( - bit_width, - packed, - &mut values.as_mut_slice()[values_len..], - ); - } - } else if indices_within_chunk.len() > UNPACK_CHUNK_THRESHOLD { - // Unpack into a temporary chunk and then copy the values. - unsafe { - let dst: &mut [MaybeUninit] = &mut unpacked; - let dst: &mut [T] = std::mem::transmute(dst); - BitPacking::unchecked_unpack(bit_width, packed, dst); - } - values.extend_trusted( - indices_within_chunk - .iter() - .map(|&idx| unsafe { unpacked.get_unchecked(idx).assume_init() }), - ); - } else { - // Otherwise, unpack each element individually. - values.extend_trusted(indices_within_chunk.iter().map(|&idx| unsafe { - BitPacking::unchecked_unpack_single(bit_width, packed, idx) - })); - } - }, - ); + let chunks = PackedChunks { + packed: array.packed_slice::(), + bit_width, + packed_chunk_len: 128 * bit_width / size_of::(), + }; + filter_chunked(&chunks, array.offset() as usize, mask) +} - values +/// The packed FastLanes chunks of a bit-packed array. +struct PackedChunks<'a, T> { + packed: &'a [T], + bit_width: usize, + /// Number of `T` words that hold each packed chunk. + packed_chunk_len: usize, +} + +impl PackedChunks<'_, T> { + fn chunk(&self, chunk_idx: usize) -> &[T] { + &self.packed[chunk_idx * self.packed_chunk_len..][..self.packed_chunk_len] + } +} + +// SAFETY: the FastLanes unpack kernels initialize every value of `dst`. +unsafe impl ChunkDecoder for PackedChunks<'_, T> { + fn decode_chunk(&self, chunk_idx: usize, dst: &mut [MaybeUninit; FILTER_CHUNK_LEN]) { + // SAFETY: the unpack only writes to `dst`, and the packed chunk holds `FILTER_CHUNK_LEN` + // values of `bit_width` bits. + unsafe { + BitPacking::unchecked_unpack( + self.bit_width, + self.chunk(chunk_idx), + dst.assume_init_mut(), + ); + } + } + + fn decode_indices(&self, chunk_idx: usize, indices: &[usize], dst: &mut [MaybeUninit]) { + debug_assert!(indices.iter().all(|&index| index < FILTER_CHUNK_LEN)); + // SAFETY: every index is within the chunk, and `dst` holds one value for each index. + unsafe { + BitPacking::unchecked_unpack_indices( + self.bit_width, + self.chunk(chunk_idx), + indices, + dst, + ); + } + } } #[cfg(test)] mod tests { + use std::ops::Range; use std::sync::LazyLock; + use fastlanes::BitPacking; + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::slice::SliceKernel; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; + use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; + use super::filter_values; + use crate::BitPacked; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; @@ -359,4 +318,130 @@ mod tests { let expected: Vec = values[0..20].to_vec(); assert_arrays_eq!(filtered, PrimitiveArray::from_iter(expected), &mut ctx); } + + /// Selection patterns that exercise every chunk strategy of the kernel. + #[derive(Clone, Copy, Debug)] + enum Pattern { + /// Each value is selected independently with the given probability. + Random(f64), + /// Runs of 12 values every 64 values. + Runs, + /// Half of the values in the first 400 of every 1024 positions, which are too dense + /// within a chunk to unpack individually. + Clustered, + /// Everything except a single value, so most chunks are fully selected. + AllButOne, + } + + fn selection(pattern: Pattern, len: usize) -> BitBuffer { + let mut rng = StdRng::seed_from_u64(42); + BitBuffer::from_iter((0..len).map(|i| match pattern { + Pattern::Random(density) => rng.random_bool(density), + Pattern::Runs => i % 64 < 12, + Pattern::Clustered => i % 1024 < 400 && rng.random_bool(0.5), + Pattern::AllButOne => i != len / 2, + })) + } + + /// A bit-packed array of `len + 2000` values with some patches, sliced to `range`. + fn sliced_bitpacked(len: usize, range: Range) -> VortexResult<(ArrayRef, ArrayRef)> { + let mut ctx = SESSION.create_execution_ctx(); + let values = + PrimitiveArray::from_option_iter((0..len as u32).map(|i| { + (i % 97 != 0).then_some(if i % 501 == 0 { 100_000 + i } else { i % 1000 }) + })); + let bitpacked = BitPackedData::encode(&values.clone().into_array(), 10, &mut ctx)?; + assert!(bitpacked.patches().is_some()); + // Slicing patches needs execution, so slice with the kernel rather than lazily. + let sliced = + ::slice(bitpacked.as_view(), range.clone(), &mut ctx)? + .unwrap(); + Ok((values.into_array().slice(range)?, sliced)) + } + + #[rstest] + fn filter_matches_canonical( + #[values( + Pattern::Random(0.0005), + Pattern::Random(0.01), + Pattern::Random(0.05), + Pattern::Random(0.15), + Pattern::Random(0.5), + Pattern::Runs, + Pattern::Clustered, + Pattern::AllButOne + )] + pattern: Pattern, + #[values(0..5000, 3..5000, 1000..4099, 1024..3072)] range: Range, + #[values(false, true)] cached_slices: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let (values, bitpacked) = sliced_bitpacked(7000, range)?; + assert!(bitpacked.is::()); + + let bits = selection(pattern, values.len()); + let mask = if cached_slices { + Mask::from_slices(bits.len(), bits.set_slices().collect()) + } else { + Mask::from_buffer(bits) + }; + + let expected = values + .filter(mask.clone())? + .execute::(&mut ctx)?; + let actual = bitpacked + .filter(mask)? + .execute::(&mut ctx)?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// Checks every value width, since the number of selected values for which a chunk unpacks + /// only those values depends on the width. + #[rstest] + fn filter_values_matches_expected_for_every_width( + #[values( + Pattern::Random(0.01), + Pattern::Random(0.03), + Pattern::Random(0.08), + Pattern::Random(0.15), + Pattern::Runs, + Pattern::Clustered + )] + pattern: Pattern, + ) -> VortexResult<()> { + check_filter_values::(pattern)?; + check_filter_values::(pattern)?; + check_filter_values::(pattern)?; + check_filter_values::(pattern) + } + + fn check_filter_values>( + pattern: Pattern, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let range = 3..5000; + + let unpacked = + PrimitiveArray::from_iter((0..5000u32).map(|i| >::from((i % 100) as u8))); + let bitpacked = BitPackedData::encode(&unpacked.into_array(), 7, &mut ctx)? + .into_array() + .slice(range.clone())?; + let bitpacked = bitpacked.as_::(); + + let bits = selection(pattern, range.len()); + let expected: Vec = range + .zip(bits.iter()) + .filter_map(|(i, selected)| selected.then_some(>::from((i % 100) as u8))) + .collect(); + + let Mask::Values(mask) = Mask::from_buffer(bits) else { + unreachable!("patterns select some but not all values") + }; + assert_eq!( + filter_values::(bitpacked.data(), 7, &mask).as_slice(), + expected + ); + Ok(()) + } } diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index 6a80ee73385..e1d6d85a376 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -101,13 +101,13 @@ fn filter_slice_in_place(values: &mut [T], mask: &MaskValues) -> usize slice::filter_slice_mut_by_bitmap(values, mask) } -fn useful_cached_slices(mask: &MaskValues) -> Option<&[(usize, usize)]> { +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 }) } -fn byte_compress_density_threshold() -> f64 { +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 diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index c3d87622654..ae6345b9a83 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/vortex-array/src/arrays/filter/execute/byte_compress.rs @@ -7,6 +7,8 @@ //! 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; @@ -124,31 +126,84 @@ fn filter_chunk_into( 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; + return 0; } if chunk.len() == 8 && mask_byte == 0xFF { // All 8 selected, so bulk copy. - out.spare_capacity_mut()[*write_pos..][..8].write_copy_of_slice(chunk); - *write_pos += 8; - return; + dst[..8].write_copy_of_slice(chunk); + return 8; } - let out_ptr = out.spare_capacity_mut().as_mut_ptr(); let (perm, count) = &BYTE_COMPRESS_LUT[mask_byte as usize]; let count = *count as usize; - debug_assert_eq!(mask_byte & !low_bits_mask(chunk.len()), 0); - // SAFETY: perm indices are all < chunk.len(), write_pos + count <= capacity. + // SAFETY: perm indices are all below `chunk.len()`, and `dst` has room for `count` values. unsafe { for j in 0..count { - out_ptr - .add(*write_pos + j) - .cast::() + dst.get_unchecked_mut(j) .write(*chunk.get_unchecked(*perm.get_unchecked(j) as usize)); } } - *write_pos += count; + count } fn low_bits_mask(bits: usize) -> u8 { diff --git a/vortex-array/src/arrays/filter/execute/chunked.rs b/vortex-array/src/arrays/filter/execute/chunked.rs new file mode 100644 index 00000000000..95bfac6da1b --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/chunked.rs @@ -0,0 +1,478 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Filtering of values that decode in fixed-size chunks. +//! +//! The mask is applied one chunk at a time, from the mask's cached slices when +//! [`filter_buffer`](super::buffer::filter_buffer) would use them and otherwise from its bitmap. +//! Each chunk takes the first matching strategy: +//! +//! | Selected values in the chunk | Strategy | +//! | ---------------------------------- | ------------------------------------------------- | +//! | none | skip the chunk | +//! | all | decode the chunk straight to the output | +//! | at most 16 per byte of value width | decode only the selected values | +//! | otherwise | decode to scratch, then compact the scratch | +//! +//! The scratch buffer stays in cache. It is compacted by copying the slices, or with the same +//! SIMD kernels that `filter_buffer` chooses for the chunk's density. When no SIMD kernel applies, +//! 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. + +use std::mem::MaybeUninit; + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_mask::MaskValues; + +use super::buffer::byte_compress_density_threshold; +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; + +/// Number of values in each chunk of a [`ChunkDecoder`]. +pub const FILTER_CHUNK_LEN: usize = 1024; + +/// Number of 64-bit mask words that cover a chunk. +const CHUNK_WORDS: usize = FILTER_CHUNK_LEN / 64; + +/// Chunks with at most this many selected values per byte of value width decode only those +/// values rather than the whole chunk. Wider values make a whole chunk more expensive to decode, +/// so the threshold grows with the width. +const SPARSE_VALUES_PER_BYTE: usize = 16; + +/// The largest number of selected values for which a chunk decodes only those values. +const MAX_SPARSE_VALUES: usize = SPARSE_VALUES_PER_BYTE * 8; + +/// 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; + +/// Decodes the values of an array in chunks of [`FILTER_CHUNK_LEN`] values, for +/// [`filter_chunked`]. +/// +/// # Safety +/// +/// Each method must initialize every value of its `dst`. +pub unsafe trait ChunkDecoder { + /// Decodes all values of chunk `chunk_idx` to `dst`. + fn decode_chunk(&self, chunk_idx: usize, dst: &mut [MaybeUninit; FILTER_CHUNK_LEN]); + + /// Decodes the values at the sorted in-chunk `indices` of chunk `chunk_idx` to `dst`, which + /// holds one value for each index. + 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. +/// +/// 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. See the [module docs](self) for the strategy that +/// each chunk takes. +pub fn filter_chunked( + decoder: &impl ChunkDecoder, + offset: usize, + mask: &MaskValues, +) -> Buffer { + let mut filter = ChunkedFilter::new(decoder, mask.true_count()); + match useful_cached_slices(mask) { + Some(slices) => filter.push_slices(offset, slices), + None => filter.push_bitmap(offset, mask), + } + filter.finish() +} + +/// A chunk of decoded values, aligned to a cache line so that the decoder's vector stores never +/// cross one. +#[repr(align(64))] +struct Scratch([MaybeUninit; FILTER_CHUNK_LEN]); + +impl Scratch { + /// Decodes chunk `chunk_idx` of `decoder` and returns its values. + fn decode(&mut self, decoder: &impl ChunkDecoder, chunk_idx: usize) -> &[T] { + decoder.decode_chunk(chunk_idx, &mut self.0); + // SAFETY: the decoder initialized every value. + unsafe { self.0.assume_init_ref() } + } +} + +/// The output and scratch buffers of [`filter_chunked`]. +struct ChunkedFilter<'a, T, D> { + decoder: &'a D, + values: BufferMut, + true_count: usize, + written: usize, + max_sparse: usize, + scratch: Scratch, + indices: [usize; MAX_SPARSE_VALUES], +} + +impl<'a, T: NativePType, D: ChunkDecoder> ChunkedFilter<'a, T, D> { + fn new(decoder: &'a D, true_count: usize) -> Self { + Self { + decoder, + // The SIMD kernels can store one vector past the last value that they write. + values: BufferMut::with_capacity(true_count + SLACK_BYTES / size_of::()), + true_count, + written: 0, + max_sparse: (SPARSE_VALUES_PER_BYTE * size_of::()).min(MAX_SPARSE_VALUES), + scratch: Scratch([const { MaybeUninit::uninit() }; FILTER_CHUNK_LEN]), + indices: [0; MAX_SPARSE_VALUES], + } + } + + /// Decodes chunk `chunk_idx` straight to the output. + fn push_full_chunk(&mut self, chunk_idx: usize) { + let dst = self.values.spare_capacity_mut()[self.written..] + .first_chunk_mut::() + .unwrap_or_else(|| unreachable!("a fully selected chunk has output capacity")); + self.decoder.decode_chunk(chunk_idx, dst); + self.written += FILTER_CHUNK_LEN; + } + + /// Decodes the values at the first `len` of `self.indices` in chunk `chunk_idx` to the output. + fn push_indices(&mut self, chunk_idx: usize, len: usize) { + let dst = &mut self.values.spare_capacity_mut()[self.written..][..len]; + self.decoder + .decode_indices(chunk_idx, &self.indices[..len], dst); + self.written += 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(); + let bits_bytes = bits.inner().as_slice(); + let mut words = [0u64; CHUNK_WORDS]; + + for chunk_idx in 0..(offset + len).div_ceil(FILTER_CHUNK_LEN) { + // The range `start..end` of the array that falls within this chunk, which begins at + // index `chunk_offset` within the chunk. + let chunk_start = chunk_idx * FILTER_CHUNK_LEN; + let start = chunk_start.saturating_sub(offset); + let end = (chunk_start + FILTER_CHUNK_LEN - offset).min(len); + let chunk_offset = start + offset - chunk_start; + let chunk_len = end - start; + + load_chunk_words(bits_bytes, bits.offset() + start, chunk_len, &mut words); + let words = &words[..chunk_len.div_ceil(64)]; + let selected = words.iter().map(|w| w.count_ones() as usize).sum::(); + + if selected == 0 { + continue; + } + + if selected == FILTER_CHUNK_LEN { + self.push_full_chunk(chunk_idx); + } else if selected <= self.max_sparse { + let mut n = 0; + for (word_idx, &word) in words.iter().enumerate() { + let mut word = word; + while word != 0 { + self.indices[n] = + chunk_offset + word_idx * 64 + word.trailing_zeros() as usize; + n += 1; + word &= word - 1; + } + } + self.push_indices(chunk_idx, selected); + } else { + let src = + &self.scratch.decode(self.decoder, chunk_idx)[chunk_offset..][..chunk_len]; + let dst = &mut self.values.spare_capacity_mut()[self.written..]; + // 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) }; + debug_assert_eq!(written, selected); + self.written += selected; + } + } + } + + /// Filters chunk by chunk from the sorted, disjoint `slices`. + fn push_slices(&mut self, offset: usize, slices: &[(usize, usize)]) { + // `slices[next..]` holds the slices not yet fully written, and `resume` is the start of + // the first chunk not yet processed, from which `slices[next]` resumes if it began in an + // earlier chunk. Positions include `offset`. + let mut next = 0; + let mut resume = 0; + while let Some(&(first_start, _)) = slices.get(next) { + let chunk_idx = (first_start + offset).max(resume) / FILTER_CHUNK_LEN; + let chunk_start = chunk_idx * FILTER_CHUNK_LEN; + let chunk_end = chunk_start + FILTER_CHUNK_LEN; + + // The in-chunk ranges selected within this chunk. + let runs = slices[next..] + .iter() + .map(|&(start, end)| (start + offset, end + offset)) + .take_while(|&(start, _)| start < chunk_end) + .map(|(start, end)| { + ( + start.max(chunk_start) - chunk_start, + end.min(chunk_end) - chunk_start, + ) + }); + let selected = runs.clone().map(|(start, end)| end - start).sum::(); + + if selected == FILTER_CHUNK_LEN { + self.push_full_chunk(chunk_idx); + } else if selected <= self.max_sparse { + for (index, chunk_index) in self + .indices + .iter_mut() + .zip(runs.flat_map(|(start, end)| start..end)) + { + *index = chunk_index; + } + self.push_indices(chunk_idx, selected); + } else { + let src = self.scratch.decode(self.decoder, chunk_idx); + let dst = &mut self.values.spare_capacity_mut()[self.written..]; + let mut written = 0; + for (start, end) in runs { + dst[written..][..end - start].write_copy_of_slice(&src[start..end]); + written += end - start; + } + self.written += written; + } + + // Skip the slices that end within this chunk; the last one may continue into the + // next. + next += slices[next..] + .iter() + .take_while(|&&(_, end)| end + offset <= chunk_end) + .count(); + resume = chunk_end; + } + } + + fn finish(mut self) -> Buffer { + debug_assert_eq!(self.written, self.true_count); + // SAFETY: the first `true_count` output values were initialized above. + unsafe { self.values.set_len(self.true_count) }; + self.values.freeze() + } +} + +/// 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] +fn load_chunk_words(bytes: &[u8], start: usize, len: usize, words: &mut [u64; CHUNK_WORDS]) { + debug_assert!(len <= FILTER_CHUNK_LEN); + let byte_start = start / 8; + let shift = start % 8; + + if len == FILTER_CHUNK_LEN { + if shift == 0 + && let Some(src) = bytes[byte_start..].first_chunk::<{ FILTER_CHUNK_LEN / 8 }>() + { + for (word, src) in words.iter_mut().zip(src.as_chunks::<8>().0) { + *word = u64::from_le_bytes(*src); + } + return; + } + if let Some(src) = bytes[byte_start..].first_chunk::() { + assemble_words(src, shift, words); + return; + } + } + + // Near the end of the bitmap, copy into a zero-padded buffer so the words can still be + // assembled from fixed-size reads. + let byte_len = (shift + len).div_ceil(8); + let mut buf = [0u8; CHUNK_WORD_BYTES]; + buf[..byte_len].copy_from_slice(&bytes[byte_start..][..byte_len]); + assemble_words(&buf, shift, words); + for (i, word) in words.iter_mut().enumerate() { + *word &= low_bits_mask(len.saturating_sub(i * 64).min(64)); + } +} + +/// Assembles 64-bit words from `src`, starting `shift < 8` bits into its first byte. +#[inline] +fn assemble_words(src: &[u8; CHUNK_WORD_BYTES], shift: usize, words: &mut [u64; CHUNK_WORDS]) { + for (i, word) in words.iter_mut().enumerate() { + let (lo, rest) = src[i * 8..] + .split_first_chunk::<8>() + .unwrap_or_else(|| unreachable!()); + let lo = u64::from_le_bytes(*lo); + let hi = u64::from(rest[0]); + // Shifting `hi` in two steps keeps the shift amount below 64 when `shift == 0`. + *word = (lo >> shift) | ((hi << 1) << (63 - shift)); + } +} + +#[cfg(test)] +mod tests { + use std::mem::MaybeUninit; + + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + use rstest::rstest; + use vortex_buffer::BitBuffer; + use vortex_mask::Mask; + + use super::ChunkDecoder; + use super::FILTER_CHUNK_LEN; + use super::filter_chunked; + use crate::dtype::NativePType; + + /// Decodes chunks of a plain slice, which holds whole chunks. + struct SliceDecoder<'a, T>(&'a [T]); + + // SAFETY: both methods initialize every value of `dst`. + unsafe impl ChunkDecoder for SliceDecoder<'_, T> { + fn decode_chunk(&self, chunk_idx: usize, dst: &mut [MaybeUninit; FILTER_CHUNK_LEN]) { + dst.write_copy_of_slice(&self.0[chunk_idx * FILTER_CHUNK_LEN..][..FILTER_CHUNK_LEN]); + } + + fn decode_indices(&self, chunk_idx: usize, indices: &[usize], dst: &mut [MaybeUninit]) { + let chunk = &self.0[chunk_idx * FILTER_CHUNK_LEN..][..FILTER_CHUNK_LEN]; + for (dst, &index) in dst.iter_mut().zip(indices) { + dst.write(chunk[index]); + } + } + } + + fn check>(bits: &BitBuffer, offset: usize, slices: bool) { + let padded_len = (offset + bits.len()).next_multiple_of(FILTER_CHUNK_LEN); + let values = (0..padded_len) + .map(|i| { + >::from(u8::try_from(i % 251).unwrap_or_else(|_| unreachable!())) + }) + .collect::>(); + let expected = bits + .iter() + .enumerate() + .filter_map(|(i, selected)| selected.then_some(values[offset + i])) + .collect::>(); + + let mask = if slices { + Mask::from_slices(bits.len(), bits.set_slices().collect()) + } else { + Mask::from_buffer(bits.clone()) + }; + let Mask::Values(mask) = mask else { + unreachable!("the patterns select some but not all values") + }; + let filtered = filter_chunked(&SliceDecoder(&values), offset, &mask); + assert_eq!(filtered.as_slice(), expected); + } + + #[rstest] + fn filter_chunked_matches_expected( + #[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, + ) { + let mut rng = StdRng::seed_from_u64(42); + let bits = (0..5000 + bit_offset) + .map(|_| rng.random_bool(density)) + .collect::() + .slice(bit_offset..5000 + bit_offset); + + check::(&bits, offset, false); + check::(&bits, offset, false); + check::(&bits, offset, false); + check::(&bits, offset, false); + } + + /// Runs long enough that the cached slices are used rather than the bitmap. + #[rstest] + fn filter_chunked_runs_matches_expected( + #[values(9, 12, 200, 3000)] run_len: usize, + #[values(1, 16, 64, 1500)] gap: usize, + #[values(0, 3, 1000)] offset: usize, + #[values(false, true)] slices: bool, + ) { + let bits = (0..5000) + .map(|i| i % (run_len + gap) < run_len) + .collect::(); + + check::(&bits, offset, slices); + check::(&bits, offset, slices); + check::(&bits, offset, slices); + check::(&bits, offset, slices); + } + + #[rstest] + fn filter_chunked_full_and_empty_chunks(#[values(0, 3)] offset: usize) { + // Whole chunks of selected values, alternating with whole chunks of unselected values. + let bits = (0..5000) + .map(|i| ((i + offset) / FILTER_CHUNK_LEN).is_multiple_of(2)) + .collect::(); + + for slices in [false, true] { + check::(&bits, offset, slices); + check::(&bits, offset, slices); + } + } +} diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 4b357de43b6..3bf17f32a1f 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -49,6 +49,7 @@ 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; diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/generic.rs b/vortex-array/src/arrays/filter/execute/simd_compress/generic.rs index 2e288cd79cb..36a7172e010 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/generic.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/generic.rs @@ -15,8 +15,8 @@ use fearless_simd::dispatch; use fearless_simd::prelude::*; use fearless_simd::u8x16; use fearless_simd_macros::simd; -use vortex_mask::MaskValues; +use super::super::slice::MaskBits; use super::super::slice::for_each_mask_word; use super::super::slice::low_bits_mask; use super::bulk_copy; @@ -107,11 +107,11 @@ macro_rules! generic_compress_kernel { pub(super) unsafe fn $walk_fn( src: *const u8, dst: *mut u8, - mask: &MaskValues, + bits: MaskBits<'_>, ) -> usize { dispatch!(Level::new(), simd => { let mut write_pos = 0; - for_each_mask_word(mask, |word, word_start, word_len| { + for_each_mask_word(bits, |word, word_start, word_len| { // SAFETY: forwarded from the caller contract. write_pos = unsafe { compress_word::<_, IN_PLACE, $elem_size, $lanes>( diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs index 007bc733808..c9c39bd999e 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs @@ -24,6 +24,7 @@ //! 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. +use std::mem::MaybeUninit; use std::ptr; use vortex_buffer::Buffer; @@ -31,6 +32,8 @@ use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_mask::MaskValues; +use super::slice::MaskBits; + #[cfg(all(target_arch = "x86_64", not(miri)))] mod generic; #[cfg(all(target_arch = "aarch64", not(miri)))] @@ -42,10 +45,10 @@ mod x86; const MIN_LEN: usize = 64; -const SLACK_BYTES: usize = 64; +pub(super) const SLACK_BYTES: usize = 64; /// `dst == src` for in-place kernels. -type Kernel = unsafe fn(*const u8, *mut u8, &MaskValues) -> usize; +type Kernel = unsafe fn(*const u8, *mut u8, MaskBits<'_>) -> usize; /// Filter a slice with a SIMD compress kernel, if one applies. /// @@ -57,7 +60,7 @@ pub(super) fn filter_slice_by_bitmap( allocator: &BufferAllocatorRef, ) -> Option> { debug_assert_eq!(values.len(), mask.len()); - let kernel = select_kernel::(mask)?; + let kernel = select_kernel::(mask.len(), mask.density())?; let true_count = mask.true_count(); let mut out = BufferMut::::with_capacity_in( @@ -71,7 +74,7 @@ pub(super) fn filter_slice_by_bitmap( kernel( values.as_ptr().cast(), out.spare_capacity_mut().as_mut_ptr().cast(), - mask, + MaskBits::Mask(mask), ) }; debug_assert_eq!(written, true_count); @@ -87,34 +90,56 @@ pub(super) fn filter_slice_mut_by_bitmap( mask: &MaskValues, ) -> Option { debug_assert_eq!(values.len(), mask.len()); - let kernel = select_kernel::(mask)?; + 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, mask) }; + let written = unsafe { kernel(dst.cast_const(), dst, MaskBits::Mask(mask)) }; debug_assert_eq!(written, mask.true_count()); Some(written) } -/// Choose the widest profitable kernel available for `T`. -fn select_kernel(mask: &MaskValues) -> Option { - if mask.len() < MIN_LEN { +/// 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. +/// +/// Returns `None` when the caller should use a scalar strategy. +/// +/// # Safety +/// +/// `bits` must select only elements of `src`, and `dst` must hold every selected element plus +/// [`SLACK_BYTES`]. +#[inline] +pub(super) unsafe fn compress_bits( + 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`. + Some(unsafe { kernel(src.as_ptr().cast(), dst.as_mut_ptr().cast(), 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 { return None; } #[cfg(all(target_arch = "x86_64", not(miri)))] { - x86::select_kernel::(mask) + x86::select_kernel::(density) } #[cfg(all(target_arch = "aarch64", not(miri)))] { - neon::select_kernel::(mask) + neon::select_kernel::(density) } #[cfg(any(not(any(target_arch = "x86_64", target_arch = "aarch64")), miri))] { - let _ = mask; + let _ = density; None } } diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs index 5eb3376adf9..171916be842 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs @@ -12,8 +12,7 @@ use core::arch::aarch64::vst1_u8; use core::arch::aarch64::vst1q_u8; use core::arch::aarch64::vtbl1_u8; -use vortex_mask::MaskValues; - +use super::super::slice::MaskBits; use super::super::slice::for_each_mask_word; use super::super::slice::low_bits_mask; use super::Kernel; @@ -32,8 +31,8 @@ fn density_band() -> Option> { } /// Choose the kernel for this element width, if one applies to this mask. -pub(super) fn select_kernel(mask: &MaskValues) -> Option { - if !density_band::()?.contains(&mask.density()) { +pub(super) fn select_kernel(density: f64) -> Option { + if !density_band::()?.contains(&density) { return None; } @@ -133,10 +132,10 @@ macro_rules! neon_compress_kernel { pub(super) unsafe fn $walk_fn( src: *const u8, dst: *mut u8, - mask: &MaskValues, + bits: MaskBits<'_>, ) -> usize { let mut write_pos = 0; - for_each_mask_word(mask, |word, word_start, word_len| { + for_each_mask_word(bits, |word, word_start, word_len| { // SAFETY: forwarded from the caller contract. write_pos = unsafe { $word_fn::(src, dst, word, word_start, word_len, write_pos) diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs index 5ffc6f3c0c7..b8fb978f4e5 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs @@ -127,15 +127,20 @@ fn avx2_kernels_match_scalar() { let mut out = vec![T::default(); mask.true_count() + SLACK_BYTES / size_of::()]; // SAFETY: AVX2 was 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(), mask) }; + let written = unsafe { + kernel_out_of_place( + values.as_ptr().cast(), + out.as_mut_ptr().cast(), + MaskBits::Mask(mask), + ) + }; assert_eq!(written, mask.true_count()); assert_eq!(&out[..written], expected.as_slice()); let mut compacted = values.to_vec(); let ptr = compacted.as_mut_ptr().cast::(); // SAFETY: AVX2 was detected above; in-place compaction stays within the slice. - let written = unsafe { kernel_in_place(ptr.cast_const(), ptr, mask) }; + let written = unsafe { kernel_in_place(ptr.cast_const(), ptr, MaskBits::Mask(mask)) }; assert_eq!(written, mask.true_count()); assert_eq!(&compacted[..written], expected.as_slice()); } diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs index 90518580e00..f18d3597980 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs @@ -34,8 +34,7 @@ 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 vortex_mask::MaskValues; - +use super::super::slice::MaskBits; use super::super::slice::for_each_mask_word; use super::super::slice::low_bits_mask; use super::Kernel; @@ -47,7 +46,7 @@ use super::generic::compress_generic_16; /// /// Sparse masks stay on the scalar set-bit walk. See `benches/filter_fixed_width.rs` when /// changing these thresholds. -pub(super) fn select_kernel(mask: &MaskValues) -> Option { +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), 2 if avx512_vbmi2() => (compress_avx512_epi16:: as Kernel, 0.15), @@ -61,7 +60,7 @@ pub(super) fn select_kernel(mask: &MaskValues) -> Optio _ => return None, }; - (mask.density() >= min_density).then_some(kernel) + (density >= min_density).then_some(kernel) } fn avx512f() -> bool { @@ -171,10 +170,10 @@ macro_rules! avx512_compress_kernel { pub(super) unsafe fn $walk_fn( src: *const u8, dst: *mut u8, - mask: &MaskValues, + bits: MaskBits<'_>, ) -> usize { let mut write_pos = 0; - for_each_mask_word(mask, |word, word_start, word_len| { + for_each_mask_word(bits, |word, word_start, word_len| { // SAFETY: forwarded from the caller contract. write_pos = unsafe { $word_fn::(src, dst, word, word_start, word_len, write_pos) @@ -413,10 +412,10 @@ macro_rules! avx2_compress_kernel { pub(super) unsafe fn $walk_fn( src: *const u8, dst: *mut u8, - mask: &MaskValues, + bits: MaskBits<'_>, ) -> usize { let mut write_pos = 0; - for_each_mask_word(mask, |word, word_start, word_len| { + for_each_mask_word(bits, |word, word_start, word_len| { // SAFETY: forwarded from the caller contract. write_pos = unsafe { $word_fn::(src, dst, word, word_start, word_len, write_pos) diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index 0d1efa5c6c5..f603b512928 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -6,18 +6,44 @@ //! Provides both immutable and mutable (in-place) filtering of typed slices by cached mask //! representations or directly from the mask bitmap. -use std::ptr; +use std::mem::MaybeUninit; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_mask::MaskValues; -/// Invoke `f` with each `(word, word_start, word_len)` of the mask bitmap, where `word` holds -/// the mask bits for elements `word_start..word_start + word_len` in its low `word_len` bits. -#[allow(clippy::inline_always)] -#[inline(always)] -pub(super) fn for_each_mask_word(mask: &MaskValues, mut f: impl FnMut(u64, usize, usize)) { +/// 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(); @@ -68,37 +94,83 @@ pub(super) fn filter_slice_by_bitmap( let output_len = mask.true_count(); let mut out = BufferMut::::with_capacity_in(output_len, allocator.clone()); - let src_ptr = slice.as_ptr(); - let spare = out.spare_capacity_mut(); + // 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(mask, |word, word_start, word_len| { + 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 { - spare[write_pos..][..word_len].write_copy_of_slice(&slice[word_start..][..word_len]); + 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 are limited to `word_len`, and the output was allocated for - // exactly `mask.true_count()` values. + // SAFETY: set bits index into `src`, and `dst` has room for every selected + // element. unsafe { - spare - .get_unchecked_mut(write_pos) - .write(*src_ptr.add(index)); - } + dst.get_unchecked_mut(write_pos) + .write(*src.get_unchecked(index)) + }; write_pos += 1; selected &= selected - 1; } } }); - debug_assert_eq!(write_pos, output_len); - // SAFETY: every output slot was initialized exactly once above. - unsafe { out.set_len(output_len) }; - out.freeze() + 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. @@ -108,13 +180,10 @@ pub(super) fn filter_slice_by_indices( allocator: &BufferAllocatorRef, ) -> Buffer { let mut out = BufferMut::::with_capacity_in(indices.len(), allocator.clone()); - let src_ptr = slice.as_ptr(); - let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); - for (write_pos, &index) in indices.iter().enumerate() { - // SAFETY: mask indices are validated when the mask is constructed and the output has one - // slot allocated for every index. - unsafe { out_ptr.add(write_pos).write(*src_ptr.add(index)) }; + 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. @@ -145,28 +214,23 @@ pub(super) fn filter_slice_mut_by_bitmap(slice: &mut [T], mask: &MaskVa "Mask length must equal the slice length" ); - let ptr = slice.as_mut_ptr(); let mut write_pos = 0; - for_each_mask_word(mask, |word, word_start, word_len| { + 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 { - // SAFETY: source and destination are in bounds and may overlap while compacting - // toward the start of the same allocation. - unsafe { ptr::copy(ptr.add(word_start), ptr.add(write_pos), word_len) }; + 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; - if write_pos != index { - // SAFETY: set bits are limited to `word_len` and stable compaction guarantees - // `write_pos <= index`. - unsafe { ptr::copy(ptr.add(index), ptr.add(write_pos), 1) }; - } + // 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; } @@ -179,13 +243,10 @@ pub(super) fn filter_slice_mut_by_bitmap(slice: &mut [T], mask: &MaskVa /// Filter a mutable slice in-place by strictly increasing indices. pub(super) fn filter_slice_mut_by_indices(slice: &mut [T], indices: &[usize]) -> usize { - let ptr = slice.as_mut_ptr(); for (write_pos, &index) in indices.iter().enumerate() { - if write_pos != index { - // SAFETY: mask indices are in bounds and stable compaction guarantees - // `write_pos <= index`. - unsafe { ptr::copy(ptr.add(index), ptr.add(write_pos), 1) }; - } + // 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() } @@ -201,15 +262,7 @@ pub(super) fn filter_slice_mut_by_slices( let len = end - start; if write_pos != start { - // SAFETY: mask slices are in bounds and source and destination may overlap while - // compacting toward the start of the same allocation. - unsafe { - ptr::copy( - slice.as_ptr().add(start), - slice.as_mut_ptr().add(write_pos), - len, - ) - }; + slice.copy_within(start..end, write_pos); } write_pos += len; diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index fb7dc6a4141..c79212c1d6e 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -12,6 +12,9 @@ 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;