diff --git a/vortex-array/benches/filter_fixed_width.rs b/vortex-array/benches/filter_fixed_width.rs index 01334108848..b082e325133 100644 --- a/vortex-array/benches/filter_fixed_width.rs +++ b/vortex-array/benches/filter_fixed_width.rs @@ -15,6 +15,7 @@ use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -28,6 +29,9 @@ use vortex_buffer::BitBuffer; use vortex_mask::Mask; use vortex_session::VortexSession; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { LazyLock::force(&SESSION); divan::main(); diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index 73ee3026654..06610f5bf82 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -3,32 +3,138 @@ //! Buffer-level filter dispatch. //! -//! Provides [`filter_buffer`] which filters a [`Buffer`] by [`MaskValues`], attempting an -//! in-place filter when the buffer has exclusive ownership. +//! Reuses suitable cached mask representations and otherwise filters directly from the bitmap, +//! avoiding temporary index or range allocations for one-shot filters. + +use std::mem::size_of; use vortex_buffer::Buffer; 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. /// -/// This will attempt to filter in-place (via [`Buffer::try_into_mut`]) when the buffer has -/// exclusive ownership, avoiding an extra allocation. +/// Dense uniquely owned buffers are compacted in place. Shared and sparse buffers allocate an +/// exact-sized output and choose between cached indices, cached long ranges, SIMD compress, +/// bitmap iteration, and byte-compress based on the mask, element width, and CPU features. pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { - match buffer.try_into_mut() { - Ok(mut buffer_mut) => { - let new_len = slice::filter_slice_mut_by_mask_values(buffer_mut.as_mut_slice(), mask); - buffer_mut.truncate(new_len); - buffer_mut.freeze() + 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, } - // Otherwise, allocate a new buffer and fill it in. - Err(buffer) => slice::filter_slice_by_mask_values(buffer.as_slice(), mask), + } else { + buffer + }; + + filter_slice(buffer.as_slice(), mask) +} + +fn filter_slice(values: &[T], mask: &MaskValues) -> Buffer { + if let Some(slices) = useful_cached_slices(mask) { + return slice::filter_slice_by_slices(values, slices, mask.true_count()); + } + + if mask.density() <= CACHED_INDICES_MAX_DENSITY + && let Some(indices) = mask.cached_indices() + { + return slice::filter_slice_by_indices(values, indices); + } + + if let Some(filtered) = simd_compress::filter_slice_by_bitmap(values, mask) { + return filtered; + } + + if mask.density() >= byte_compress_density_threshold::() { + byte_compress::filter_buffer(values, mask) + } else { + slice::filter_slice_by_bitmap(values, mask) + } +} + +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 mask.density() <= CACHED_INDICES_MAX_DENSITY + && let Some(indices) = mask.cached_indices() + { + 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) +} + +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 { + let width = size_of::(); + + // The scalar byte-LUT has a later crossover on AArch64, where the word-at-a-time set-bit walk + // is particularly efficient. Wider values also favor the word walk until all-set bytes become + // common. Keep these cases covered by `benches/filter_fixed_width.rs`. + // + // On AArch64 the NEON kernels above already cover most of the range these thresholds used to + // govern, so this is reached only where `simd_compress::neon::density_band` declines: 8-byte + // values at density >= 0.8, and the 16/32-byte widths NEON has no kernel for. + 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, + } +} + +/// Materialize sparse indices when enough sibling arrays will reuse the same mask. +pub(super) fn prepare_mask_for_reuse(mask: &MaskValues, consumers: usize) { + if consumers <= 1 || mask.cached_indices().is_some() || mask.cached_slices().is_some() { + return; } + + let density_threshold = if consumers >= 3 { 0.1 } else { 0.05 }; + if mask.density() > density_threshold { + return; + } + + if super::contiguous_values_range(mask).is_some() { + return; + } + + let _ = mask.indices(); } #[cfg(test)] mod tests { + use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_buffer::buffer; use vortex_mask::Mask; @@ -80,4 +186,83 @@ mod tests { let result = filter_buffer(buf, mask_values(&mask)); assert_eq!(result, buffer![1u32, 2, 3, 4, 6, 7, 8, 10]); } + + #[test] + fn test_filter_shared_buffer_by_cached_indices() { + let buf = Buffer::from(BufferMut::from_iter(0u64..16)); + let shared = buf.clone(); + let mask = Mask::from_indices(16, [1, 5, 9, 15]); + + let result = filter_buffer(buf, mask_values(&mask)); + assert_eq!(result, buffer![1u64, 5, 9, 15]); + assert_eq!(shared.len(), 16); + } + + #[test] + fn test_filter_shared_buffer_by_cached_slices() { + let buf = Buffer::from(BufferMut::from_iter(0u32..32)); + let shared = buf.clone(); + let mask = Mask::from_slices(32, vec![(3, 15), (20, 30)]); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (3u32..15).chain(20..30).collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + assert_eq!(shared.len(), 32); + } + + #[test] + fn test_filter_shared_dense_bitmap() { + let buf = Buffer::from(BufferMut::from_iter(0u16..128)); + let shared = buf.clone(); + let mask = Mask::from_iter((0..128).map(|index| index != 63)); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (0u16..128).filter(|&value| value != 63).collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + assert_eq!(shared.len(), 128); + } + + #[test] + fn test_filter_unaligned_bitmap_words() { + const LEN: usize = 151; + const OFFSET: usize = 5; + + let backing = BitBuffer::from_iter( + std::iter::repeat_n(false, OFFSET).chain((0..LEN).map(|index| index % 7 == 2)), + ); + let mask = Mask::from_buffer(BitBuffer::new_with_offset( + backing.inner().clone(), + LEN, + OFFSET, + )); + let buf = Buffer::from(BufferMut::from_iter(0u64..LEN as u64)); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (0u64..LEN as u64) + .filter(|value| value % 7 == 2) + .collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + } + + #[test] + fn test_prepare_sparse_mask_for_sibling_reuse() { + let two_consumers = Mask::from_iter((0..100).map(|index| index % 20 == 0)); + let values = mask_values(&two_consumers); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 2); + assert!(values.cached_indices().is_some()); + + let three_consumers = Mask::from_iter((0..100).map(|index| index % 10 == 0)); + let values = mask_values(&three_consumers); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 2); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 3); + assert!(values.cached_indices().is_some()); + + let contiguous = Mask::from_iter((0..100).map(|index| (20..25).contains(&index))); + let values = mask_values(&contiguous); + prepare_mask_for_reuse(values, 3); + assert!(values.cached_indices().is_none()); + } } diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index 6c58ce4f33b..4ef8c520ce0 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/vortex-array/src/arrays/filter/execute/byte_compress.rs @@ -1,20 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Byte-level compress for primitive filtering using a `1 << 8 = 256`-entry lookup table. +//! 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::size_of; - use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_mask::MaskValues; -const BYTE_COMPRESS_DENSITY_THRESHOLD: f64 = 0.5; - /// 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 @@ -45,10 +41,10 @@ static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{ /// /// Processes the mask one byte at a time (8 source elements per byte), /// using a precomputed permutation to compact selected elements. -pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { - debug_assert_eq!(buffer.len(), mask.len()); +pub(super) fn filter_buffer(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer { + let src = buffer.as_ref(); + debug_assert_eq!(src.len(), mask.len()); - let src = buffer.as_slice(); let true_count = mask.true_count(); if true_count == 0 { @@ -59,14 +55,7 @@ pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Bu let mask_bytes = mask_buffer.inner().as_ref(); let mask_offset = mask_buffer.offset(); - // Fast path: byte-wide values benefit from avoiding index materialization more often. Wider - // values need enough selected values to justify scanning every mask byte directly. - if size_of::() == 1 || mask.density() >= BYTE_COMPRESS_DENSITY_THRESHOLD { - return filter_bitpacked(src, mask_bytes, mask_offset, true_count); - } - - // Slow path: lower-density wide values are better handled by the generic path. - super::slice::filter_slice_by_mask_values(src, mask) + filter_bitpacked(src, mask_bytes, mask_offset, true_count) } fn filter_bitpacked( diff --git a/vortex-array/src/arrays/filter/execute/decimal.rs b/vortex-array/src/arrays/filter/execute/decimal.rs index 5ffb4a963b2..11ff451ea0f 100644 --- a/vortex-array/src/arrays/filter/execute/decimal.rs +++ b/vortex-array/src/arrays/filter/execute/decimal.rs @@ -32,13 +32,48 @@ pub fn filter_decimal(array: &DecimalArray, mask: &Arc) -> DecimalAr } #[cfg(test)] -mod test { +mod tests { + use rstest::rstest; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::filter::execute::decimal::DecimalArray; use crate::compute::conformance::filter::test_filter_conformance; use crate::dtype::DecimalDType; + use crate::dtype::i256; + + #[rstest] + #[case(DecimalArray::from_iter( + [1i8, 2, 3, 4, 5], + DecimalDType::new(2, 0), + ))] + #[case(DecimalArray::from_iter( + [10i16, 20, 30, 40, 50], + DecimalDType::new(3, 0), + ))] + #[case(DecimalArray::from_iter( + [100i32, 200, 300, 400, 500], + DecimalDType::new(5, 0), + ))] + #[case(DecimalArray::from_iter( + [1_000i64, 2_000, 3_000, 4_000, 5_000], + DecimalDType::new(10, 0), + ))] + #[case(DecimalArray::from_iter( + [10_000i128, 20_000, 30_000, 40_000, 50_000], + DecimalDType::new(19, 0), + ))] + #[case(DecimalArray::from_iter( + [1i128, 2, 3, 4, 5].map(i256::from_i128), + DecimalDType::new(39, 0), + ))] + fn test_filter_decimal_physical_type_conformance(#[case] array: DecimalArray) { + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } #[test] fn test_filter_decimal128_conformance() { diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 2c712009d1d..29325c57c82 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -5,6 +5,7 @@ //! //! The main entrypoint is [`execute_filter`] which filters any [`Canonical`] array. +use std::ops::Range; use std::sync::Arc; use vortex_error::VortexExpect; @@ -30,6 +31,7 @@ use crate::validity::Validity; pub(crate) mod byte_compress; +mod simd_compress; mod slice; mod take; @@ -52,6 +54,50 @@ fn filter_validity(validity: Validity, mask: &Arc) -> Validity { .vortex_expect("Somehow unable to wrap filter around a validity array") } +pub(super) fn contiguous_filter_range(mask: &Mask) -> Option> { + match mask { + Mask::AllTrue(len) => (*len > 0).then_some(0..*len), + Mask::AllFalse(_) => None, + Mask::Values(values) => contiguous_values_range(values), + } +} + +fn contiguous_values_range(mask: &MaskValues) -> Option> { + if let Some(slices) = mask.cached_slices() { + return match slices { + [(start, end)] => Some(*start..*end), + _ => None, + }; + } + + if let Some(indices) = mask.cached_indices() { + let start = *indices.first()?; + let end = indices.last()?.checked_add(1)?; + return (end - start == indices.len()).then_some(start..end); + } + + let true_count = mask.true_count(); + let start = mask.bit_buffer().set_indices().next()?; + let end = start.checked_add(true_count)?; + if end > mask.len() { + return None; + } + + // Probe from the cheaper side: count the candidate run for sparse masks, or find the final + // set bit from the end for dense masks. This bounds the uncached work by the smaller half of + // the bitmap while retaining the zero-copy path. + let contiguous = if true_count <= mask.len() / 2 { + mask.bit_buffer().count_range(start, end) == true_count + } else { + mask.bit_buffer().last_set_index() == Some(end - 1) + }; + contiguous.then_some(start..end) +} + +pub(super) fn prepare_mask_for_reuse(mask: &MaskValues, consumers: usize) { + buffer::prepare_mask_for_reuse(mask, consumers); +} + /// Check for some fast-path execution conditions before calling [`execute_filter`]. pub(super) fn execute_filter_fast_paths( array: ArrayView<'_, Filter>, @@ -69,6 +115,11 @@ pub(super) fn execute_filter_fast_paths( return Ok(Some(array.child().clone())); } + // Filtering by one contiguous range is exactly a slice and can remain zero-copy. + if let Some(range) = contiguous_filter_range(array.filter_mask()) { + return array.child().slice(range).map(Some); + } + // Also check if the array itself is completely null, in which case we only care about the total // number of nulls, not the values. let child_arr = array.array(); @@ -125,3 +176,94 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca } } } + +#[cfg(test)] +mod tests { + use vortex_buffer::BitBuffer; + use vortex_error::VortexResult; + + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + + #[test] + fn contiguous_filter_executes_as_zero_copy_slice() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0i32..8); + let original = array.to_buffer::(); + let filtered = array + .into_array() + .filter(Mask::from_slices(8, vec![(2, 6)]))? + .execute::(&mut array_session().create_execution_ctx())?; + let filtered_values = filtered.to_buffer::(); + + assert_eq!(filtered_values.as_slice(), &[2, 3, 4, 5]); + assert_eq!(filtered_values.as_ptr(), original.as_ptr().wrapping_add(2)); + Ok(()) + } + + #[test] + fn uncached_contiguous_filter_executes_as_zero_copy_slice() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0i32..128); + let original = array.to_buffer::(); + let mask = Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (37..91).contains(&index)), + )); + assert!(mask.values().is_some_and(|values| { + values.cached_indices().is_none() && values.cached_slices().is_none() + })); + + let filtered = array + .into_array() + .filter(mask)? + .execute::(&mut array_session().create_execution_ctx())?; + let filtered_values = filtered.to_buffer::(); + + assert_eq!(filtered_values.as_slice(), &(37..91).collect::>()); + assert_eq!(filtered_values.as_ptr(), original.as_ptr().wrapping_add(37)); + Ok(()) + } + + #[test] + fn fragmented_filter_is_not_a_contiguous_range() { + let mask = Mask::from_indices(8, [1, 2, 5, 6]); + assert_eq!(contiguous_filter_range(&mask), None); + } + + #[test] + fn uncached_contiguous_range_handles_sparse_and_dense_masks() { + let cases = [ + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (37..41).contains(&index)), + )), + Some(37..41), + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (3..125).contains(&index)), + )), + Some(3..125), + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| matches!(index, 1 | 40 | 90)), + )), + None, + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| !matches!(index, 17 | 63)), + )), + None, + ), + ]; + + for (mask, expected) in cases { + assert!(mask.values().is_some_and(|values| { + values.cached_indices().is_none() && values.cached_slices().is_none() + })); + assert_eq!(contiguous_filter_range(&mask), expected); + } + } +} diff --git a/vortex-array/src/arrays/filter/execute/primitive.rs b/vortex-array/src/arrays/filter/execute/primitive.rs index b7abd8efa4b..08b22d2fd22 100644 --- a/vortex-array/src/arrays/filter/execute/primitive.rs +++ b/vortex-array/src/arrays/filter/execute/primitive.rs @@ -8,12 +8,8 @@ use vortex_mask::MaskValues; use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::buffer; -use crate::arrays::filter::execute::byte_compress; use crate::arrays::filter::execute::filter_validity; -use crate::dtype::NativePType; -use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::validity::Validity; pub fn filter_primitive(array: &PrimitiveArray, mask: &Arc) -> PrimitiveArray { let validity = array @@ -21,34 +17,13 @@ pub fn filter_primitive(array: &PrimitiveArray, mask: &Arc) -> Primi .vortex_expect("primitive validity should be derivable"); let filtered_validity = filter_validity(validity, mask); - match array.ptype() { - // Byte-compress avoids materializing indices/slices and processes 8 elements per mask byte. - PType::U8 => filter_byte_compress::(array, filtered_validity, mask), - PType::I8 => filter_byte_compress::(array, filtered_validity, mask), - PType::U16 => filter_byte_compress::(array, filtered_validity, mask), - PType::I16 => filter_byte_compress::(array, filtered_validity, mask), - PType::U32 => filter_byte_compress::(array, filtered_validity, mask), - PType::I32 => filter_byte_compress::(array, filtered_validity, mask), - _ => match_each_native_ptype!(array.ptype(), |T| { - let filtered_buffer = buffer::filter_buffer(array.to_buffer::(), mask.as_ref()); + match_each_native_ptype!(array.ptype(), |T| { + let filtered_buffer = buffer::filter_buffer(array.to_buffer::(), mask.as_ref()); - // SAFETY: We filter both the validity and the buffer with the same mask, so they must - // have the same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } - }), - } -} - -fn filter_byte_compress( - array: &PrimitiveArray, - filtered_validity: Validity, - mask: &Arc, -) -> PrimitiveArray { - let filtered_buffer = byte_compress::filter_buffer(array.to_buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must have the - // same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } + // SAFETY: We filter both the validity and the buffer with the same mask, so they must have + // the same length. + unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } + }) } #[cfg(test)] diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs new file mode 100644 index 00000000000..89fe9a80125 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! SIMD compress kernels for fixed-width filtering. +//! +//! Every architecture reduces the same problem — write the elements selected by a mask word to +//! the front of the output — to a per-(sub)word gather, then walks the mask with +//! [`for_each_mask_word`](super::slice::for_each_mask_word). What differs is how one chunk of +//! lanes is compacted: +//! +//! - AVX-512: `vpcompress[b/w/d/q]` consumes each mask (sub)word directly as the `k`-register +//! (AVX-512F for 4/8-byte elements, additionally VBMI2 for 1/2-byte elements). +//! - AVX2, 4/8-byte elements: a per-mask-byte (or nibble) lane-index LUT feeding `vpermd`, a +//! vectorized version of the scalar `BYTE_COMPRESS_LUT`. +//! - AVX2, 1/2-byte elements: a byte-index LUT feeding 128-bit `pshufb`. `pshufb` only shuffles +//! within a 128-bit lane, but eight lanes of a 1- or 2-byte element need at most a 16-byte +//! table, so one lane is all the table has to span. +//! - NEON: the same byte-index LUT construction feeding `tbl`, which unlike `pshufb` spans a +//! whole register, so it also serves 4- and 8-byte elements. +//! +//! The byte-shuffle kernels (`pshufb`/`tbl`) share [`compress_lut`] and [`compress_tail`]; only +//! the load/shuffle/store intrinsics differ. +//! +//! x86 selects a kernel per buffer by runtime feature detection (`is_x86_feature_detected!` +//! caches per feature, so this stays cheap); NEON is part of the aarch64 baseline and needs no +//! probe. Unlike `vortex_buffer::bit::pack`'s `collect_bool_words_multiversioned`, there is no +//! caller-supplied closure for a `#[target_feature]` boundary to deoptimize — the per-word body +//! is a fixed gather — so a supported element width routes here whenever the mask is dense +//! enough to amortize the per-chunk work. +//! +//! Compressed vectors are written with full-width unmasked stores: the out-of-place output is +//! over-allocated by one vector of slack so trailing garbage lands in spare capacity, which +//! also sidesteps `vpcompressstoreu`'s microcoded slowness on Zen 4. The in-place variant +//! bounds every unmasked store by the source position it has already consumed, and handles +//! partial tail chunks without a full-width store at all. + +use std::ptr; + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_mask::MaskValues; + +#[cfg(all(target_arch = "aarch64", not(miri)))] +mod neon; +#[cfg(test)] +mod tests; +#[cfg(all(target_arch = "x86_64", not(miri)))] +mod x86; + +/// Below one full mask word the scalar paths win; skip kernel selection entirely. +const MIN_LEN: usize = 64; + +/// One full AVX-512 vector of output slack, so every out-of-place store can be unmasked. +const SLACK_BYTES: usize = 64; + +/// A per-buffer compress kernel: `(src, dst, mask) -> written`, with `src`/`dst` pointing at +/// the first element. `dst == src` for the in-place instantiations. +type Kernel = unsafe fn(*const u8, *mut u8, &MaskValues) -> usize; + +/// Filter a slice with a SIMD compress kernel, if one applies. +/// +/// Returns `None` when no kernel applies (unsupported architecture or element width, missing +/// CPU features, or an input too short or too sparse to amortize the per-chunk work); the +/// caller then falls back to the scalar strategies. +pub(super) fn filter_slice_by_bitmap( + values: &[T], + mask: &MaskValues, +) -> Option> { + debug_assert_eq!(values.len(), mask.len()); + let kernel = select_kernel::(mask)?; + + let true_count = mask.true_count(); + let mut out = BufferMut::::with_capacity(true_count + SLACK_BYTES / size_of::()); + // 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(), + 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)?; + + 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) }; + debug_assert_eq!(written, mask.true_count()); + Some(written) +} + +/// Choose the widest kernel this build and CPU offer for `T`, if one beats the scalar walk for +/// this mask. +fn select_kernel(mask: &MaskValues) -> Option { + if mask.len() < MIN_LEN { + return None; + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + { + x86::select_kernel::(mask) + } + #[cfg(all(target_arch = "aarch64", not(miri)))] + { + neon::select_kernel::(mask) + } + #[cfg(any(not(any(target_arch = "x86_64", target_arch = "aarch64")), miri))] + { + let _ = mask; + None + } +} + +/// Build the byte-shuffle index rows for one element width: `lut[m]` gathers the lanes selected +/// by mask `m` into the front of the vector, leaving the trailing bytes as don't-care zeros. +/// +/// `BYTES` is the table width, which may exceed the `lanes * elem_size` bytes the mask covers — +/// `pshufb` always indexes a full 16-byte register even when only its low half holds lanes. Both +/// `ROWS == 1 << lanes` and the table fit are checked at compile time by the `static` +/// initializers. +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[expect( + clippy::cast_possible_truncation, + reason = "byte indices are bounded by the 8- or 16-byte table width" +)] +const fn compress_lut( + lanes: usize, + elem_size: usize, +) -> [[u8; BYTES]; ROWS] { + assert!(ROWS == 1 << lanes); + assert!(lanes * elem_size <= BYTES); + + let mut lut = [[0u8; BYTES]; ROWS]; + let mut m = 0; + while m < ROWS { + let mut out_lane = 0; + let mut lane = 0; + while lane < lanes { + if m & (1 << lane) != 0 { + let mut byte = 0; + while byte < elem_size { + lut[m][out_lane * elem_size + byte] = (lane * elem_size + byte) as u8; + byte += 1; + } + out_lane += 1; + } + lane += 1; + } + m += 1; + } + lut +} + +/// Compact the elements selected by `bits` (bit `i` selects element `base + i`) one at a time. +/// +/// Used for the final partial chunk of a mask word, where a full-width vector load would read +/// past the end of the source. A mask has at most two such chunks — one for an unaligned leading +/// word and one for the trailing word — so this never runs in the hot loop. +/// +/// # Safety +/// +/// The pointer contract of [`filter_slice_by_bitmap`] / [`filter_slice_mut_by_bitmap`] must hold, +/// and `bits` must only select elements that are in bounds. +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[inline(always)] +unsafe fn compress_tail( + src: *const u8, + dst: *mut u8, + mut bits: u64, + base: usize, + mut write_pos: usize, + elem_size: usize, +) -> usize { + while bits != 0 { + let index = base + bits.trailing_zeros() as usize; + // SAFETY: `index` is in bounds per the contract above and stable compaction guarantees + // `write_pos <= index`. + unsafe { bulk_copy::(src, dst, index, 1, write_pos, elem_size) }; + write_pos += 1; + bits &= bits - 1; + } + write_pos +} + +/// Copy `word_len` elements of `elem_size` bytes for a fully-set mask word. +/// +/// # Safety +/// +/// `word_len` source elements starting at `word_start` must be in bounds. Out-of-place, `dst` +/// must not overlap `src` and must have room at `write_pos`; in-place, forward compaction must +/// guarantee `write_pos <= word_start`. +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[inline(always)] +unsafe fn bulk_copy( + src: *const u8, + dst: *mut u8, + word_start: usize, + word_len: usize, + write_pos: usize, + elem_size: usize, +) { + // SAFETY: offsets are in bounds per the contract above. + let src_bytes = unsafe { src.add(word_start * elem_size) }; + // SAFETY: see above. + let dst_bytes = unsafe { dst.add(write_pos * elem_size) }; + if IN_PLACE { + if write_pos != word_start { + // SAFETY: source and destination may overlap while compacting forward. + unsafe { ptr::copy(src_bytes, dst_bytes, word_len * elem_size) }; + } + } else { + // SAFETY: out-of-place buffers are disjoint allocations. + unsafe { ptr::copy_nonoverlapping(src_bytes, dst_bytes, word_len * elem_size) }; + } +} diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs new file mode 100644 index 00000000000..da9d250a2b8 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs @@ -0,0 +1,252 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! NEON `tbl` compress kernels. +//! +//! `tbl` is an arbitrary byte shuffle over a register-sized table, so one instruction compacts +//! all the lanes of a single vector. Unlike AVX2's `vpermd` it is byte-granular, which means the +//! same construction serves every element width — the LUT just holds byte indices instead of +//! lane indices. What it does not have is AVX-512's `vpcompress`, so the permutation still comes +//! from a table indexed by the mask bits for one vector of lanes: +//! +//! | element | lanes per `tbl` | mask bits | LUT rows | LUT size | +//! | --- | --- | --- | --- | --- | +//! | 1 byte | 8 (`vtbl1_u8`, 64-bit) | 8 | 256 | 2 KiB | +//! | 2 bytes | 8 (`vqtbl1q_u8`) | 8 | 256 | 4 KiB | +//! | 4 bytes | 4 (`vqtbl1q_u8`) | 4 | 16 | 256 B | +//! | 8 bytes | 2 (`vqtbl1q_u8`) | 2 | 4 | 64 B | +//! +//! Every LUT fits comfortably in L1. The lanes-per-shuffle column sets both the payoff and the +//! density band each kernel is worth using over (see [`density_band`]), and is why wide elements +//! gain far less here than under `vpcompressq`: a `tbl` shuffles one register whatever the +//! element width, so the 8-byte kernel retires two elements per shuffle where AVX-512 retires +//! eight, leaving it worth ~1.1-1.6x over a narrow band while the 1/2-byte kernels reach 9x. +//! +//! See the [module docs](super) for how these fit the shared dispatch. NEON is part of the +//! aarch64 baseline, so there is nothing to probe and no `#[target_feature]` boundary to block +//! inlining. + +use core::arch::aarch64::vld1_u8; +use core::arch::aarch64::vld1q_u8; +use core::arch::aarch64::vqtbl1q_u8; +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::for_each_mask_word; +use super::super::slice::low_bits_mask; +use super::Kernel; +use super::bulk_copy; +use super::compress_lut; +use super::compress_tail; + +/// The mask-density band in which the kernel for `T` beats both scalar strategies. +/// +/// These kernels do a fixed amount of work per chunk, so their cost per element is flat in the +/// mask density; the bounds are just where the two scalar strategies cross that flat line. +/// +/// Below the lower bound the set-bit walk wins, because it steps from one set bit to the next +/// instead of paying a shuffle and store for lanes it discards. The fewer lanes a chunk holds +/// the longer that takes to amortize, so the bound rises with element width. +/// +/// Only 8-byte elements need an upper bound. A `tbl` shuffles one register whatever the element +/// width, so an 8-byte chunk is a mere two lanes and the kernel never gets far ahead; past ~0.8 +/// the byte-LUT's all-set bulk copies overtake it. Returning `None` is what lets +/// [`filter_buffer`](super::super::buffer) fall through to them. +/// +/// Bounds measured with `benches/filter_fixed_width.rs` on Apple M-series, at `LEN` 16384, +/// interleaved against the same binary with the kernels disabled. Speedups inside each band: +/// +/// | element | band | at floor | at 0.5 | at 0.9 | at 0.99 | +/// | --- | --- | --- | --- | --- | --- | +/// | 1 byte | 0.15.. | 1.4x | 5.7x | 10.1x | 2.4x | +/// | 2 bytes | 0.15.. | 1.3x | 5.1x | 9.2x | 2.0x | +/// | 4 bytes | 0.30.. | 1.2x | 2.9x | 5.1x | 1.1x | +/// | 8 bytes | 0.50..0.80 | 1.3x | 1.3x | — | — | +/// +/// The bounds sit just above the measured crossovers so a core with a different +/// shuffle-to-scalar balance still cannot regress. +fn density_band() -> Option> { + match size_of::() { + 1 | 2 => Some(0.15..f64::INFINITY), + 4 => Some(0.30..f64::INFINITY), + 8 => Some(0.50..0.80), + _ => None, + } +} + +/// 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()) { + return None; + } + + match size_of::() { + 1 => Some(compress_neon_8:: as Kernel), + 2 => Some(compress_neon_16:: as Kernel), + 4 => Some(compress_neon_32:: as Kernel), + 8 => Some(compress_neon_64:: as Kernel), + _ => None, + } +} + +static IDX_LUT_8: [[u8; 8]; 256] = compress_lut::<256, 8>(8, 1); +static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2); +static IDX_LUT_32: [[u8; 16]; 16] = compress_lut::<16, 16>(4, 4); +static IDX_LUT_64: [[u8; 16]; 4] = compress_lut::<4, 16>(2, 8); + +/// Generate a NEON `tbl` kernel for one element width: a per-word compress (`$word_fn`) plus +/// the mask-word walk (`$walk_fn`) that drives it. Each chunk of `$lanes` mask bits indexes +/// `$idx_lut` for the byte permutation that compacts one vector of lanes. +macro_rules! neon_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,elem_size: + $elem_size:literal,lanes: + $lanes:literal,idx_lut: + $idx_lut:ident,load: + $load:ident,tbl: + $tbl:ident,store: + $store:ident + ) => { + /// Compress the elements selected by one mask word. + /// + /// # 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. + #[expect( + clippy::cast_possible_truncation, + reason = "deliberate submask narrowing" + )] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::(src, dst, word_start, word_len, write_pos, $elem_size) + }; + return write_pos + word_len; + } + + // Deliberately branchless: an empty chunk still shuffles and stores, it just + // leaves `write_pos` where it was so the next chunk overwrites the garbage. A + // `m != 0` guard here is a large *pessimization*, because the mask bits of one + // chunk are exactly as random as the data — `P(m == 0) = (1 - density)^lanes` + // passes 0.5 somewhere inside the useful density range for every width, so the + // branch mispredicts near-maximally right where it would pay off. Measured on + // Apple M-series with `benches/filter_fixed_width.rs`, the guarded 4-byte kernel + // costs 1.95x at density 0.5 and the guarded 1-byte kernel 1.45x at 0.1. + let mut sub = 0; + while sub + $lanes <= word_len { + let m = ((word >> sub) & low_bits_mask($lanes)) as usize; + // SAFETY: the chunk holds `$lanes` in-bounds source elements, which is exactly + // the table width. + let chunk = unsafe { $load(src.add((word_start + sub) * $elem_size)) }; + // SAFETY: every LUT row is one table wide. + let idx = unsafe { $load($idx_lut[m].as_ptr()) }; + // SAFETY: the store covers `$lanes` elements at `write_pos`, of which the + // leading `m.count_ones()` are selected values and the rest are garbage that a + // later store overwrites, since stores resume at the advanced `write_pos`. + // Out-of-place, the trailing garbage of the final store lands in the vector of + // slack past `true_count`. In-place, `write_pos <= word_start + sub`, so the + // store ends at or before `word_start + sub + $lanes <= word_start + word_len + // <= len`, and it can only overwrite source elements already loaded into + // `chunk` or output positions not yet final. + unsafe { $store(dst.add(write_pos * $elem_size), $tbl(chunk, idx)) }; + write_pos += m.count_ones() as usize; + sub += $lanes; + } + + if sub < word_len { + let bits = (word >> sub) & low_bits_mask(word_len - sub); + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + compress_tail::( + src, + dst, + bits, + word_start + sub, + write_pos, + $elem_size, + ) + }; + } + + write_pos + } + + /// Walk the mask words of `mask`, compressing the selected elements. + /// + /// # 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. + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |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) + }; + }); + write_pos + } + }; +} + +neon_compress_kernel!( + compress_word_neon_8, compress_neon_8, + elem_size: 1, + lanes: 8, + idx_lut: IDX_LUT_8, + load: vld1_u8, + tbl: vtbl1_u8, + store: vst1_u8 +); + +neon_compress_kernel!( + compress_word_neon_16, compress_neon_16, + elem_size: 2, + lanes: 8, + idx_lut: IDX_LUT_16, + load: vld1q_u8, + tbl: vqtbl1q_u8, + store: vst1q_u8 +); + +neon_compress_kernel!( + compress_word_neon_32, compress_neon_32, + elem_size: 4, + lanes: 4, + idx_lut: IDX_LUT_32, + load: vld1q_u8, + tbl: vqtbl1q_u8, + store: vst1q_u8 +); + +neon_compress_kernel!( + compress_word_neon_64, compress_neon_64, + elem_size: 8, + lanes: 2, + idx_lut: IDX_LUT_64, + load: vld1q_u8, + tbl: vqtbl1q_u8, + store: vst1q_u8 +); diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs new file mode 100644 index 00000000000..c76da1dd80d --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![allow(clippy::cast_possible_truncation)] + +use rstest::rstest; +use vortex_buffer::BitBuffer; +use vortex_mask::Mask; +use vortex_mask::MaskValues; + +use super::super::slice; +use super::*; + +/// `None` when the mask normalized to `AllTrue`/`AllFalse`, which `filter_buffer` never sees +/// (the filter fast paths intercept those before buffer-level dispatch). +fn mask_values(mask: &Mask) -> Option<&MaskValues> { + match mask { + Mask::Values(values) => Some(values.as_ref()), + _ => None, + } +} + +fn make_mask(len: usize, offset: usize, pattern: impl Fn(usize) -> bool) -> Mask { + 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, + )) +} + +type Pattern = fn(usize) -> bool; + +fn patterns() -> Vec<(&'static str, Pattern)> { + vec![ + ("all_but_first", |i| i != 0), + ("only_first", |i| i == 0), + ("sparse", |i| i % 97 == 0), + ("mid", |i| i % 3 == 0), + ("dense", |i| i % 16 != 0), + ("alternating", |i| i % 2 == 0), + ("word_blocks", |i| (i / 64) % 2 == 0), + ("edges", |i| i < 3 || i % 61 == 60), + ] +} + +fn check(values: &[T], mask: &Mask) { + let Some(mask) = mask_values(mask) else { + return; + }; + let expected = slice::filter_slice_by_bitmap(values, mask); + + if let Some(actual) = filter_slice_by_bitmap(values, mask) { + assert_eq!(actual.as_slice(), expected.as_slice()); + } + + let mut compacted = values.to_vec(); + if let Some(new_len) = filter_slice_mut_by_bitmap(&mut compacted, mask) { + assert_eq!(&compacted[..new_len], expected.as_slice()); + } +} + +#[rstest] +fn simd_matches_scalar( + #[values(0, 3, 5)] offset: usize, + #[values(64, 100, 151, 1000, 1024)] len: usize, +) { + for (name, pattern) in patterns() { + let mask = make_mask(len, offset, pattern); + let mask_debug = format!("pattern={name} len={len} offset={offset}"); + + 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(); + + println!("checking {mask_debug}"); + check(&u8_values, &mask); + check(&u16_values, &mask); + check(&u32_values, &mask); + check(&u64_values, &mask); + } +} + +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[test] +fn engages_on_supported_cpus() { + let values: Vec = (0..256).collect(); + let mask = make_mask(256, 0, |i| i % 2 == 0); + let mask = mask_values(&mask).expect("alternating mask is mixed"); + + #[cfg(target_arch = "x86_64")] + let expected = is_x86_feature_detected!("avx2"); + // NEON is part of the aarch64 baseline. + #[cfg(target_arch = "aarch64")] + let expected = true; + + assert_eq!(filter_slice_by_bitmap(&values, mask).is_some(), expected); +} + +/// Every architecture declines a mask too sparse or too short to amortize a per-chunk compress, +/// which is what routes those to the scalar strategies. +#[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 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()); +} + +/// The dispatcher never selects the AVX2 tier on AVX-512 machines, so exercise those kernels +/// directly. Covers both the `vpermd` kernels for 4/8-byte elements and the `pshufb` kernels +/// for 1/2-byte elements. +#[cfg(all(target_arch = "x86_64", not(miri)))] +#[test] +fn avx2_kernels_match_scalar() { + if !is_x86_feature_detected!("avx2") { + return; + } + + fn check_kernel( + kernel_out_of_place: Kernel, + kernel_in_place: Kernel, + values: &[T], + mask: &MaskValues, + ) { + let expected = slice::filter_slice_by_bitmap(values, mask); + + // One vector of slack, mirroring the allocation in `filter_slice_by_bitmap`. + 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) }; + 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) }; + assert_eq!(written, mask.true_count()); + assert_eq!(&compacted[..written], expected.as_slice()); + } + + for (_, pattern) in patterns() { + for len in [64, 100, 151, 1000] { + for offset in [0, 5] { + let mask = make_mask(len, offset, pattern); + let Some(mask) = mask_values(&mask) else { + continue; + }; + 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(); + let u64_values: Vec = (0..len as u64).collect(); + check_kernel( + x86::compress_pshufb_epi8::, + x86::compress_pshufb_epi8::, + &u8_values, + mask, + ); + check_kernel( + x86::compress_pshufb_epi16::, + x86::compress_pshufb_epi16::, + &u16_values, + mask, + ); + check_kernel( + x86::compress_avx2_epi32::, + x86::compress_avx2_epi32::, + &u32_values, + mask, + ); + check_kernel( + x86::compress_avx2_epi64::, + x86::compress_avx2_epi64::, + &u64_values, + mask, + ); + } + } + } +} diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs new file mode 100644 index 00000000000..8cd994608b1 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs @@ -0,0 +1,634 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! AVX-512 `vpcompress`, AVX2 `vpermd`, and 128-bit `pshufb` compress kernels. +//! +//! See the [module docs](super) for how these fit the shared dispatch. + +use std::arch::x86_64::__m128i; +use std::arch::x86_64::_mm_loadl_epi64; +use std::arch::x86_64::_mm_loadu_si128; +use std::arch::x86_64::_mm_shuffle_epi8; +use std::arch::x86_64::_mm_storel_epi64; +use std::arch::x86_64::_mm_storeu_si128; +use std::arch::x86_64::_mm256_loadu_si256; +use std::arch::x86_64::_mm256_maskload_epi32; +use std::arch::x86_64::_mm256_maskload_epi64; +use std::arch::x86_64::_mm256_maskstore_epi32; +use std::arch::x86_64::_mm256_maskstore_epi64; +use std::arch::x86_64::_mm256_permutevar8x32_epi32; +use std::arch::x86_64::_mm256_storeu_si256; +use std::arch::x86_64::_mm512_loadu_epi8; +use std::arch::x86_64::_mm512_loadu_epi16; +use std::arch::x86_64::_mm512_loadu_epi32; +use std::arch::x86_64::_mm512_loadu_epi64; +use std::arch::x86_64::_mm512_mask_storeu_epi8; +use std::arch::x86_64::_mm512_mask_storeu_epi16; +use std::arch::x86_64::_mm512_mask_storeu_epi32; +use std::arch::x86_64::_mm512_mask_storeu_epi64; +use std::arch::x86_64::_mm512_maskz_compress_epi8; +use std::arch::x86_64::_mm512_maskz_compress_epi16; +use std::arch::x86_64::_mm512_maskz_compress_epi32; +use std::arch::x86_64::_mm512_maskz_compress_epi64; +use std::arch::x86_64::_mm512_maskz_loadu_epi8; +use std::arch::x86_64::_mm512_maskz_loadu_epi16; +use std::arch::x86_64::_mm512_maskz_loadu_epi32; +use std::arch::x86_64::_mm512_maskz_loadu_epi64; +use std::arch::x86_64::_mm512_storeu_epi8; +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::for_each_mask_word; +use super::super::slice::low_bits_mask; +use super::Kernel; +use super::bulk_copy; +use super::compress_lut; +use super::compress_tail; + +/// Choose the widest kernel the CPU supports for this element width, and the minimum mask +/// density at which it beats the scalar strategies. +/// +/// Below that density the scalar set-bit walk wins, because it steps from one set bit to the +/// next instead of paying a compress and a full-width store for lanes it discards. How long that +/// takes to amortize is set by how many lanes one compress covers, so the floor falls as the +/// kernel gets wider: `vpcompressb` consumes a whole 64-bit mask word at once and never loses, +/// while AVX2's 4-lane `vpermd` over 8-byte elements needs a nearly half-set mask to pay off. +/// +/// Measured with `benches/filter_fixed_width.rs` on Zen 5 (EPYC 9R05, AVX-512 VBMI2), at `LEN` +/// 16384, interleaved against the same binary with the kernels disabled. The AVX2 rows come from +/// the same machine with AVX-512 detection forced off: +/// +/// | kernel | lanes per compress | element | floor | at floor | best in band | +/// | --- | --- | --- | --- | --- | --- | +/// | `vpcompressb` | 64 | 1 byte | none | 2.1x at 0.05 | 5.7x at 0.8 | +/// | `vpcompressw` | 32 | 2 bytes | 0.15 | 1.1x | 3.9x at 0.8 | +/// | `vpcompressd` | 16 | 4 bytes | 0.25 | 1.1x | 2.8x at 0.8 | +/// | `vpcompressq` | 8 | 8 bytes | 0.30 | 1.1x | 3.1x at 0.65 | +/// | `pshufb` | 8 | 1 byte | 0.15 | 1.4x | 2.6x at 0.8 | +/// | `pshufb` | 8 | 2 bytes | 0.25 | 1.0x | 2.5x at 0.8 | +/// | `vpermd` | 8 | 4 bytes | 0.25 | 1.0x | 2.6x at 0.8 | +/// | `vpermd` | 4 | 8 bytes | 0.45 | 1.0x | 1.4x at 0.65 | +/// +/// The two `pshufb` rows have the same lane count but different floors because they compete +/// against different scalar strategies. `byte_compress_density_threshold` routes 1-byte elements +/// to the byte-LUT at *every* density on x86, and that LUT is weak on a sparse mask, so the +/// 1-byte kernel already wins at 0.15; 2-byte elements fall to the set-bit walk instead, which +/// holds out until 0.25. Retuning that threshold would raise the 1-byte floor to match. +/// +/// `is_x86_feature_detected!` caches per feature, so per-buffer selection stays cheap. +pub(super) fn select_kernel(mask: &MaskValues) -> 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), + 4 if avx512f() => (compress_avx512_epi32:: as Kernel, 0.25), + 8 if avx512f() => (compress_avx512_epi64:: as Kernel, 0.30), + // AVX-512F without VBMI2 (e.g. Skylake-X) falls through to these too. + 1 if avx2() => (compress_pshufb_epi8:: as Kernel, 0.15), + 2 if avx2() => (compress_pshufb_epi16:: as Kernel, 0.25), + 4 if avx2() => (compress_avx2_epi32:: as Kernel, 0.25), + 8 if avx2() => (compress_avx2_epi64:: as Kernel, 0.45), + _ => return None, + }; + + (mask.density() >= min_density).then_some(kernel) +} + +fn avx512f() -> bool { + is_x86_feature_detected!("avx512f") +} + +fn avx512_vbmi2() -> bool { + is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512bw") + && is_x86_feature_detected!("avx512vbmi2") +} + +fn avx2() -> bool { + is_x86_feature_detected!("avx2") +} + +/// Generate an AVX-512 `vpcompress` kernel for one element width: a per-word compress +/// (`$word_fn`) plus the mask-word walk (`$walk_fn`) that drives it. Each mask subword of +/// `$lanes` bits is used directly as the `k`-register of one compress. +macro_rules! avx512_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,features: + $features:literal,elem: + $elem:ty,lanes: + $lanes:literal,kmask: + $kmask:ty,loadu: + $loadu:ident,maskz_loadu: + $maskz_loadu:ident,maskz_compress: + $maskz_compress:ident,storeu: + $storeu:ident,mask_storeu: + $mask_storeu:ident + ) => { + /// Compress the elements selected by one mask word. + /// + /// # 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. + // `allow` rather than `expect`: the cast is only lossy for sub-`u64` k-masks. + #[allow(clippy::cast_possible_truncation)] + #[target_feature(enable = $features)] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::( + src, + dst, + word_start, + word_len, + write_pos, + size_of::<$elem>(), + ) + }; + return write_pos + word_len; + } + + let mut sub = 0; + while sub < word_len { + let count = (word_len - sub).min($lanes); + let k = ((word >> sub) & low_bits_mask(count)) as $kmask; + if k != 0 { + let src_ptr = unsafe { src.cast::<$elem>().add(word_start + sub) }; + let chunk = if count == $lanes { + // SAFETY: `$lanes` source elements starting at + // `word_start + sub` are in bounds. + unsafe { $loadu(src_ptr) } + } else { + // SAFETY: the load mask enables only the `count` in-bounds lanes. + unsafe { $maskz_loadu(low_bits_mask(count) as $kmask, src_ptr) } + }; + let packed = $maskz_compress(k, chunk); + let selected = k.count_ones() as usize; + let dst_ptr = unsafe { dst.cast::<$elem>().add(write_pos) }; + if !IN_PLACE || count == $lanes { + // SAFETY: out-of-place output has a vector of slack past + // `true_count`; in-place, `write_pos <= word_start + sub`, so + // `write_pos + $lanes <= word_start + sub + $lanes <= len`. + unsafe { $storeu(dst_ptr, packed) }; + } else { + // SAFETY: stores exactly the `selected` selected lanes, all + // within the `true_count` output positions. + unsafe { $mask_storeu(dst_ptr, low_bits_mask(selected) as $kmask, packed) }; + } + write_pos += selected; + } + sub += $lanes; + } + write_pos + } + + /// Walk the mask words of `mask`, compressing the selected elements. + /// + /// # 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. + #[target_feature(enable = $features)] + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |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) + }; + }); + write_pos + } + }; +} + +avx512_compress_kernel!( + compress_word_avx512_epi8, compress_avx512_epi8, + features: "avx512f,avx512bw,avx512vbmi2", + elem: i8, + lanes: 64, + kmask: u64, + loadu: _mm512_loadu_epi8, + maskz_loadu: _mm512_maskz_loadu_epi8, + maskz_compress: _mm512_maskz_compress_epi8, + storeu: _mm512_storeu_epi8, + mask_storeu: _mm512_mask_storeu_epi8 +); + +avx512_compress_kernel!( + compress_word_avx512_epi16, compress_avx512_epi16, + features: "avx512f,avx512bw,avx512vbmi2", + elem: i16, + lanes: 32, + kmask: u32, + loadu: _mm512_loadu_epi16, + maskz_loadu: _mm512_maskz_loadu_epi16, + maskz_compress: _mm512_maskz_compress_epi16, + storeu: _mm512_storeu_epi16, + mask_storeu: _mm512_mask_storeu_epi16 +); + +avx512_compress_kernel!( + compress_word_avx512_epi32, compress_avx512_epi32, + features: "avx512f", + elem: i32, + lanes: 16, + kmask: u16, + loadu: _mm512_loadu_epi32, + maskz_loadu: _mm512_maskz_loadu_epi32, + maskz_compress: _mm512_maskz_compress_epi32, + storeu: _mm512_storeu_epi32, + mask_storeu: _mm512_mask_storeu_epi32 +); + +avx512_compress_kernel!( + compress_word_avx512_epi64, compress_avx512_epi64, + features: "avx512f", + elem: i64, + lanes: 8, + kmask: u8, + loadu: _mm512_loadu_epi64, + maskz_loadu: _mm512_maskz_loadu_epi64, + maskz_compress: _mm512_maskz_compress_epi64, + storeu: _mm512_storeu_epi64, + mask_storeu: _mm512_mask_storeu_epi64 +); + +/// For each mask byte, `vpermd` lane indices compacting the selected 4-byte lanes to the front +/// (trailing lanes are don't-care). +static PERM_LUT_32: [[u32; 8]; 256] = { + let mut lut = [[0u32; 8]; 256]; + let mut m = 0; + while m < 256 { + let mut out_lane = 0; + let mut bit = 0; + while bit < 8 { + if m & (1 << bit) != 0 { + lut[m][out_lane] = bit as u32; + out_lane += 1; + } + bit += 1; + } + m += 1; + } + lut +}; + +/// For each mask nibble, `vpermd` lane indices compacting the selected 8-byte lanes (as pairs +/// of 4-byte lanes) to the front. +static PERM_LUT_64: [[u32; 8]; 16] = { + let mut lut = [[0u32; 8]; 16]; + let mut m = 0; + while m < 16 { + let mut out_lane = 0; + let mut bit = 0; + while bit < 4 { + if m & (1 << bit) != 0 { + lut[m][out_lane * 2] = (bit * 2) as u32; + lut[m][out_lane * 2 + 1] = (bit * 2 + 1) as u32; + out_lane += 1; + } + bit += 1; + } + m += 1; + } + lut +}; + +/// Lane-enable vectors for `vpmaskmov` loads/stores of the first `count` 4-byte lanes. +static LANE_MASK_32: [[i32; 8]; 9] = { + let mut lut = [[0i32; 8]; 9]; + let mut count = 0; + while count <= 8 { + let mut lane = 0; + while lane < count { + lut[count][lane] = -1; + lane += 1; + } + count += 1; + } + lut +}; + +/// Lane-enable vectors for `vpmaskmov` loads/stores of the first `count` 8-byte lanes. +static LANE_MASK_64: [[i64; 4]; 5] = { + let mut lut = [[0i64; 4]; 5]; + let mut count = 0; + while count <= 4 { + let mut lane = 0; + while lane < count { + lut[count][lane] = -1; + lane += 1; + } + count += 1; + } + lut +}; + +/// Generate an AVX2 permutation-LUT kernel for one element width, mirroring +/// [`avx512_compress_kernel`] with `vpermd` compaction per mask byte (or nibble). +macro_rules! avx2_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,elem: + $elem:ty,lanes: + $lanes:literal,perm_lut: + $perm_lut:ident,lane_masks: + $lane_masks:ident,maskload: + $maskload:ident,maskstore: + $maskstore:ident + ) => { + /// Compress the elements selected by one mask word. + /// + /// # 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. + #[expect( + clippy::cast_possible_truncation, + reason = "deliberate submask narrowing" + )] + #[target_feature(enable = "avx2")] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::( + src, + dst, + word_start, + word_len, + write_pos, + size_of::<$elem>(), + ) + }; + return write_pos + word_len; + } + + let mut sub = 0; + while sub < word_len { + let count = (word_len - sub).min($lanes); + let m = ((word >> sub) & low_bits_mask(count)) as usize; + if m != 0 { + let src_ptr = unsafe { src.cast::<$elem>().add(word_start + sub) }; + let chunk = if count == $lanes { + // SAFETY: `$lanes` source elements starting at + // `word_start + sub` are in bounds. + unsafe { _mm256_loadu_si256(src_ptr.cast()) } + } else { + // SAFETY: the lane mask enables only the `count` in-bounds lanes. + unsafe { + $maskload( + src_ptr, + _mm256_loadu_si256($lane_masks[count].as_ptr().cast()), + ) + } + }; + // SAFETY: LUT rows are 32 bytes. + let perm = unsafe { _mm256_loadu_si256($perm_lut[m].as_ptr().cast()) }; + let packed = _mm256_permutevar8x32_epi32(chunk, perm); + let selected = m.count_ones() as usize; + let dst_ptr = unsafe { dst.cast::<$elem>().add(write_pos) }; + if !IN_PLACE || count == $lanes { + // SAFETY: out-of-place output has a vector of slack past + // `true_count`; in-place, `write_pos <= word_start + sub`, so + // `write_pos + $lanes <= word_start + sub + $lanes <= len`. + unsafe { _mm256_storeu_si256(dst_ptr.cast(), packed) }; + } else { + // SAFETY: stores exactly the `selected` selected lanes, all + // within the `true_count` output positions. + unsafe { + $maskstore( + dst_ptr, + _mm256_loadu_si256($lane_masks[selected].as_ptr().cast()), + packed, + ) + }; + } + write_pos += selected; + } + sub += $lanes; + } + write_pos + } + + /// Walk the mask words of `mask`, compressing the selected elements. + /// + /// # 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. + #[target_feature(enable = "avx2")] + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |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) + }; + }); + write_pos + } + }; +} + +avx2_compress_kernel!( + compress_word_avx2_epi32, compress_avx2_epi32, + elem: i32, + lanes: 8, + perm_lut: PERM_LUT_32, + lane_masks: LANE_MASK_32, + maskload: _mm256_maskload_epi32, + maskstore: _mm256_maskstore_epi32 +); + +avx2_compress_kernel!( + compress_word_avx2_epi64, compress_avx2_epi64, + elem: i64, + lanes: 4, + perm_lut: PERM_LUT_64, + lane_masks: LANE_MASK_64, + maskload: _mm256_maskload_epi64, + maskstore: _mm256_maskstore_epi64 +); + +/// Byte-index rows for `pshufb`, which always indexes a full 16-byte register even though only +/// the low 8 (1-byte elements) or all 16 (2-byte elements) bytes hold lanes. +static SHUF_LUT_8: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 1); +static SHUF_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2); + +/// Generate a 128-bit `pshufb` kernel for one narrow element width, giving 1- and 2-byte +/// elements the vector compress that `vpermd` cannot express. +/// +/// This is the same construction as the NEON kernels — see +/// [`neon`](super::super::simd_compress) for the LUT layout and the reason the chunk loop is +/// branchless. `pshufb` shuffles only within a 128-bit lane, but eight lanes of a 1- or 2-byte +/// element need at most a 16-byte table, so a single lane spans the whole table. +/// +/// Enabling AVX2 rather than just SSSE3 costs nothing in reach worth having and gets the +/// VEX-encoded `vpshufb`, which avoids AVX/SSE transition penalties in mixed code. +macro_rules! pshufb_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,elem_size: + $elem_size:literal,idx_lut: + $idx_lut:ident,load: + $load:ident,store: + $store:ident + ) => { + /// Compress the elements selected by one mask word. + /// + /// # 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. + #[expect( + clippy::cast_possible_truncation, + reason = "deliberate submask narrowing" + )] + #[target_feature(enable = "avx2")] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::(src, dst, word_start, word_len, write_pos, $elem_size) + }; + return write_pos + word_len; + } + + // Branchless for the same reason as the NEON kernels: with eight lanes per chunk, + // `P(m == 0) = (1 - density)^8` crosses 0.5 near density 0.08, so an `m != 0` guard + // mispredicts hardest exactly where skipping would pay. + let mut sub = 0; + while sub + 8 <= word_len { + let m = ((word >> sub) & low_bits_mask(8)) as usize; + // SAFETY: the chunk holds 8 in-bounds source elements. + let chunk = unsafe { $load(src.add((word_start + sub) * $elem_size).cast()) }; + // SAFETY: every LUT row is 16 bytes. + let idx = unsafe { _mm_loadu_si128($idx_lut[m].as_ptr().cast()) }; + // SAFETY: the store covers 8 elements at `write_pos`, of which the leading + // `m.count_ones()` are selected values and the rest are garbage that a later + // store overwrites, since stores resume at the advanced `write_pos`. + // Out-of-place, the trailing garbage of the final store lands in the vector of + // slack past `true_count`. In-place, `write_pos <= word_start + sub`, so the + // store ends at or before `word_start + sub + 8 <= word_start + word_len + // <= len`, and it can only overwrite source elements already loaded into + // `chunk` or output positions not yet final. + unsafe { + $store( + dst.add(write_pos * $elem_size).cast::<__m128i>(), + _mm_shuffle_epi8(chunk, idx), + ) + }; + write_pos += m.count_ones() as usize; + sub += 8; + } + + if sub < word_len { + let bits = (word >> sub) & low_bits_mask(word_len - sub); + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + compress_tail::( + src, + dst, + bits, + word_start + sub, + write_pos, + $elem_size, + ) + }; + } + + write_pos + } + + /// Walk the mask words of `mask`, compressing the selected elements. + /// + /// # 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. + #[target_feature(enable = "avx2")] + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |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) + }; + }); + write_pos + } + }; +} + +pshufb_compress_kernel!( + compress_word_pshufb_epi8, compress_pshufb_epi8, + elem_size: 1, + idx_lut: SHUF_LUT_8, + load: _mm_loadl_epi64, + store: _mm_storel_epi64 +); + +pshufb_compress_kernel!( + compress_word_pshufb_epi16, compress_pshufb_epi16, + elem_size: 2, + idx_lut: SHUF_LUT_16, + load: _mm_loadu_si128, + store: _mm_storeu_si128 +); diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index 1528d272b28..c778be13fd2 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -3,47 +3,127 @@ //! Core slice-level filtering algorithms. //! -//! Provides both immutable and mutable (in-place) filtering of typed slices by various mask -//! representations: indices and ranges (slices). +//! 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 vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_mask::MaskIter; use vortex_mask::MaskValues; -// This is modeled after the constant with the equivalent name in arrow-rs. -pub(super) const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8; +/// 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. +#[inline(always)] +pub(super) fn for_each_mask_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 + } +} // --------------------------------------------------------------------------- // Immutable slice filtering // --------------------------------------------------------------------------- -/// Filter a slice by [`MaskValues`], dispatching to the indices or slices path based on a -/// selectivity threshold. -pub(super) fn filter_slice_by_mask_values(slice: &[T], mask: &MaskValues) -> Buffer { +/// Filter a slice from the mask bitmap without materializing indices or ranges. +pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> Buffer { assert_eq!( mask.len(), slice.len(), "Selection mask length must equal the buffer length" ); - match mask.threshold_iter(FILTER_SLICES_SELECTIVITY_THRESHOLD) { - MaskIter::Indices(indices) => filter_slice_by_indices(slice, indices), - MaskIter::Slices(slices) => filter_slice_by_slices(slice, slices), - } + let output_len = mask.true_count(); + let mut out = BufferMut::::with_capacity(output_len); + let src_ptr = slice.as_ptr(); + let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let mut write_pos = 0; + + for_each_mask_word(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 { + // SAFETY: a full mask word selects `word_len` in-bounds source values and the output + // was allocated for every selected value. + unsafe { + ptr::copy_nonoverlapping(src_ptr.add(word_start), out_ptr.add(write_pos), 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. + unsafe { + out_ptr.add(write_pos).write(*src_ptr.add(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() } /// Filter a slice by a set of strictly increasing indices. -fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { - Buffer::::from_trusted_len_iter(indices.iter().map(|&idx| slice[idx])) +pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { + let mut out = BufferMut::::with_capacity(indices.len()); + 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)) }; + } + + // 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. -fn filter_slice_by_slices(slice: &[T], slices: &[(usize, usize)]) -> Buffer { - let output_len: usize = slices.iter().map(|(start, end)| end - start).sum(); - +pub(super) fn filter_slice_by_slices( + slice: &[T], + slices: &[(usize, usize)], + output_len: usize, +) -> Buffer { let mut out = BufferMut::::with_capacity(output_len); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); @@ -56,40 +136,81 @@ fn filter_slice_by_slices(slice: &[T], slices: &[(usize, usize)]) -> Bu // Mutable (in-place) slice filtering // --------------------------------------------------------------------------- -/// Filter a mutable slice in-place by [`MaskValues`], returning the new valid length. -/// -/// We always use the slices path here because iterating over indices will have strictly more -/// loop iterations than slices (more branches), and the overhead of batched `ptr::copy(len)` is -/// not that high. -pub(super) fn filter_slice_mut_by_mask_values( - slice: &mut [T], - mask: &MaskValues, -) -> usize { +/// 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" ); - filter_slice_mut_by_slices(slice, mask.slices()) + let ptr = slice.as_mut_ptr(); + let mut write_pos = 0; + + for_each_mask_word(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) }; + } + 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) }; + } + 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 { + 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) }; + } + } + indices.len() } /// Filter a mutable slice in-place by a set of `(start, end)` ranges, returning the new length. -fn filter_slice_mut_by_slices(slice: &mut [T], slices: &[(usize, usize)]) -> usize { +pub(super) fn filter_slice_mut_by_slices( + slice: &mut [T], + slices: &[(usize, usize)], +) -> usize { let mut write_pos = 0; // For each range in the selection, copy all of the elements to the current write position. for &(start, end) in slices { let len = end - start; - // SAFETY: Slices should be within bounds. - unsafe { - ptr::copy( - slice.as_ptr().add(start), - slice.as_mut_ptr().add(write_pos), - len, - ) - }; + 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, + ) + }; + } write_pos += len; } diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index e4a8157d2e2..85866fde196 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -10,9 +10,12 @@ use vortex_mask::MaskValues; use crate::ArrayRef; use crate::arrays::StructArray; use crate::arrays::filter::execute::filter_validity; +use crate::arrays::filter::execute::prepare_mask_for_reuse; use crate::arrays::struct_::StructArrayExt; pub fn filter_struct(array: &StructArray, mask: &Arc) -> StructArray { + prepare_mask_for_reuse(mask, array.struct_fields().nfields()); + let filtered_validity = filter_validity( array .validity() diff --git a/vortex-array/src/arrays/filter/execute/take/rank.rs b/vortex-array/src/arrays/filter/execute/take/rank.rs index efbfa768df3..d0f966609b8 100644 --- a/vortex-array/src/arrays/filter/execute/take/rank.rs +++ b/vortex-array/src/arrays/filter/execute/take/rank.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_bail; use vortex_mask::AllOr; use vortex_mask::Mask; +use super::super::contiguous_filter_range; use super::small_take_rank_lookup_len; use crate::arrays::PrimitiveArray; use crate::dtype::IntegerPType; @@ -153,7 +154,5 @@ pub(in crate::arrays::filter) fn contiguous_sequential_take_range Option { - let start = filter.first()?; - let end = filter.last()?.checked_add(1)?; - (end - start == filter.true_count()).then_some(start) + contiguous_filter_range(filter).map(|range| range.start) } diff --git a/vortex-array/src/arrays/filter/rules.rs b/vortex-array/src/arrays/filter/rules.rs index d952a0df460..9798d6dd693 100644 --- a/vortex-array/src/arrays/filter/rules.rs +++ b/vortex-array/src/arrays/filter/rules.rs @@ -14,6 +14,7 @@ use crate::arrays::StructArray; use crate::arrays::filter::FilterArraySlotsExt; use crate::arrays::filter::FilterReduce; use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::filter::execute::prepare_mask_for_reuse; use crate::arrays::struct_::StructDataParts; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; @@ -66,6 +67,10 @@ impl ArrayReduceRule for FilterStructRule { .. } = struct_array.into_owned().into_data_parts(); + if let Some(values) = mask.values() { + prepare_mask_for_reuse(values, fields.len()); + } + let filtered_validity = validity.filter(mask)?; let filtered_fields = fields diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index cf69ce35449..f4e145b79aa 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -417,6 +417,36 @@ impl BitBuffer { bit_select(self.buffer.as_slice(), self.offset, self.len, nth) } + /// Returns the index of the last set bit, or `None` if every bit is unset. + /// + /// This scans from the end a word at a time, avoiding the full forward scan required by + /// [`Self::select`] when selecting the final set bit. + #[inline] + pub fn last_set_index(&self) -> Option { + let chunks = self.unaligned_chunks(); + let lead = chunks.lead_padding(); + let prefix_words = usize::from(chunks.prefix().is_some()); + + if let Some(word) = chunks.suffix() + && word != 0 + { + let word_index = prefix_words + chunks.chunks().len(); + return Some(word_index * 64 + 63 - word.leading_zeros() as usize - lead); + } + + for (index, &word) in chunks.chunks().iter().enumerate().rev() { + if word != 0 { + let word_index = prefix_words + index; + return Some(word_index * 64 + 63 - word.leading_zeros() as usize - lead); + } + } + + chunks.prefix().filter(|word| *word != 0).map(|word| { + debug_assert!(word.trailing_zeros() as usize >= lead); + 63 - word.leading_zeros() as usize - lead + }) + } + /// Get the number of unset bits in the buffer. #[inline] pub fn false_count(&self) -> usize { @@ -795,6 +825,39 @@ mod tests { BitBuffer::from_indices(5, [0, 5]); } + #[test] + fn last_set_index_handles_offsets_and_padding() { + type Pattern = fn(usize) -> bool; + + let patterns: [Pattern; 5] = [ + |_| false, + |index| index == 0, + |index| index % 97 == 0, + |index| index % 3 == 1, + |index| (64..96).contains(&index), + ]; + + for offset in [0, 3, 8, 13, 67] { + for len in [0, 1, 7, 8, 63, 64, 65, 127, 128, 151] { + for pattern in patterns { + let backing = BitBuffer::from_iter( + std::iter::repeat_n(true, offset) + .chain((0..len).map(pattern)) + .chain(std::iter::repeat_n(true, 7)), + ); + let buffer = BitBuffer::new_with_offset(backing.inner().clone(), len, offset); + let expected = (0..len).rfind(|&index| pattern(index)); + + assert_eq!( + buffer.last_set_index(), + expected, + "offset={offset} len={len}" + ); + } + } + } + } + #[rstest] #[case(5)] #[case(8)]