From d9dc33d2a4ab8c1c337066c9d576ba35f47e3d0c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sat, 11 Jul 2026 00:15:06 +0100 Subject: [PATCH 1/7] Better fixed-width buffer filtering Signed-off-by: Robert Kruszewski --- .../src/arrays/filter/execute/buffer.rs | 199 +++++++++++++++++- .../arrays/filter/execute/byte_compress.rs | 21 +- .../src/arrays/filter/execute/decimal.rs | 37 +++- vortex-array/src/arrays/filter/execute/mod.rs | 47 +++++ .../src/arrays/filter/execute/primitive.rs | 37 +--- .../src/arrays/filter/execute/slice.rs | 198 +++++++++++++---- .../src/arrays/filter/execute/struct_.rs | 3 + .../src/arrays/filter/execute/take/rank.rs | 5 +- vortex-array/src/arrays/filter/rules.rs | 5 + 9 files changed, 454 insertions(+), 98 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index 73ee3026654..b8244da0983 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -3,32 +3,130 @@ //! 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::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, bitmap iteration, +/// and byte-compress based on the mask and element width. 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 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); + } + + 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`. + 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; } + + let first = mask.bit_buffer().select(0); + let last = mask.bit_buffer().select(mask.true_count() - 1); + if first + .zip(last) + .is_some_and(|(first, last)| last - first + 1 == mask.true_count()) + { + 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 +178,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..fcaefb914ee 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; @@ -52,6 +53,16 @@ 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> { + let start = mask.first()?; + let end = mask.last()?.checked_add(1)?; + (end - start == mask.true_count()).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 +80,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 +141,34 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca } } } + +#[cfg(test)] +mod tests { + 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 fragmented_filter_is_not_a_contiguous_range() { + let mask = Mask::from_indices(8, [1, 2, 5, 6]); + assert_eq!(contiguous_filter_range(&mask), None); + } +} 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/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index 1528d272b28..d51de41244e 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -3,47 +3,130 @@ //! 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; +#[inline] +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()); +} + +#[inline] +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 { + debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); + + 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 { + debug_assert_eq!( + output_len, + slices.iter().map(|(start, end)| end - start).sum::() + ); let mut out = BufferMut::::with_capacity(output_len); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); @@ -56,40 +139,83 @@ 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 { + debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); + + 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 From 283d21d3fe0c5b0284331af3abb4659ec4590f91 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 13 Jul 2026 10:01:29 +0100 Subject: [PATCH 2/7] lessassert Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/filter/execute/slice.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index d51de41244e..d6ec94db062 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -100,8 +100,6 @@ pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> /// Filter a slice by a set of strictly increasing indices. pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { - debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); - 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::(); @@ -123,10 +121,6 @@ pub(super) fn filter_slice_by_slices( slices: &[(usize, usize)], output_len: usize, ) -> Buffer { - debug_assert_eq!( - output_len, - slices.iter().map(|(start, end)| end - start).sum::() - ); let mut out = BufferMut::::with_capacity(output_len); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); @@ -181,8 +175,6 @@ 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 { - debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); - let ptr = slice.as_mut_ptr(); for (write_pos, &index) in indices.iter().enumerate() { if write_pos != index { From 68af66b7ae88f56bbd3967b166c8bd3c4889b6aa Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 20 Jul 2026 13:32:25 +0100 Subject: [PATCH 3/7] more Signed-off-by: Robert Kruszewski --- .../src/arrays/filter/execute/buffer.rs | 13 +- vortex-array/src/arrays/filter/execute/mod.rs | 1 + .../arrays/filter/execute/simd_compress.rs | 761 ++++++++++++++++++ .../src/arrays/filter/execute/slice.rs | 9 +- 4 files changed, 779 insertions(+), 5 deletions(-) create mode 100644 vortex-array/src/arrays/filter/execute/simd_compress.rs diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index b8244da0983..e21ee3ba973 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -12,6 +12,7 @@ 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; @@ -21,8 +22,8 @@ const MIN_SLICES_AVERAGE_RUN_LENGTH: usize = 8; /// Filter a [`Buffer`] by [`MaskValues`], returning a new buffer. /// /// Dense uniquely owned buffers are compacted in place. Shared and sparse buffers allocate an -/// exact-sized output and choose between cached indices, cached long ranges, bitmap iteration, -/// and byte-compress based on the mask and element width. +/// 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 { assert_eq!(buffer.len(), mask.len()); @@ -53,6 +54,10 @@ fn filter_slice(values: &[T], mask: &MaskValues) -> Buffer { 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 { @@ -71,6 +76,10 @@ fn filter_slice_in_place(values: &mut [T], mask: &MaskValues) -> usize 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) } diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index fcaefb914ee..4b1b83b4fcd 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -31,6 +31,7 @@ use crate::validity::Validity; pub(crate) mod byte_compress; +mod simd_compress; mod slice; mod take; diff --git a/vortex-array/src/arrays/filter/execute/simd_compress.rs b/vortex-array/src/arrays/filter/execute/simd_compress.rs new file mode 100644 index 00000000000..eaa55327d2e --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress.rs @@ -0,0 +1,761 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Runtime-dispatched SIMD compress kernels for fixed-width filtering. +//! +//! This mirrors the multiversioning pattern of `vortex_buffer::bit::pack` +//! (`collect_bool_words_multiversioned`): a shared per-mask-word walk instantiated behind +//! per-feature-level `#[target_feature]` wrappers, selected once per buffer by runtime feature +//! detection. Unlike `collect_bool`, there is no caller-supplied closure for the +//! `#[target_feature]` boundary to deoptimize — the per-word body is a fixed gather — so every +//! supported element width routes here unconditionally and density thresholds only matter for +//! the scalar fallbacks. +//! +//! Two kernel tiers, dispatched on element width: +//! +//! - 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 only): a per-mask-byte (or nibble) permutation LUT feeding +//! `vpermd`, a vectorized version of the scalar `BYTE_COMPRESS_LUT`. AVX2 has no cross-lane +//! byte shuffle, so 1/2-byte elements stay on the scalar byte-LUT below AVX-512. +//! +//! 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 masks the +//! stores of partial tail chunks instead. + +use vortex_buffer::Buffer; +use vortex_mask::MaskValues; + +/// Filter a slice with a runtime-detected SIMD compress kernel, if one applies. +/// +/// Returns `None` when no kernel applies (non-x86 target, unsupported element width, missing +/// CPU features, or an input too short to amortize dispatch); the caller then falls back to +/// the scalar strategies. +pub(super) fn filter_slice_by_bitmap( + values: &[T], + mask: &MaskValues, +) -> Option> { + #[cfg(all(target_arch = "x86_64", not(miri)))] + { + return x86::filter(values, mask); + } + #[allow(unreachable_code)] + { + let _ = (values, mask); + None + } +} + +/// 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 { + #[cfg(all(target_arch = "x86_64", not(miri)))] + { + return x86::filter_mut(values, mask); + } + #[allow(unreachable_code)] + { + let _ = (values, mask); + None + } +} + +#[cfg(all(target_arch = "x86_64", not(miri)))] +mod x86 { + 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 std::ptr; + + use vortex_buffer::Buffer; + use vortex_buffer::BufferMut; + use vortex_mask::MaskValues; + + use super::super::slice::for_each_mask_word; + use super::super::slice::low_bits_mask; + + /// Below one full mask word the scalar paths win; skip feature detection entirely. + const MIN_LEN: usize = 64; + + /// Below this mask density the scalar set-bit walk beats vector compress for multi-byte + /// elements: with under ~3 set bits per mask word, the per-subword loop and full-width + /// stores are pure overhead. Single-byte elements consume a whole mask word per compress + /// and win at every density. Covered by `benches/filter_fixed_width.rs` (0.01 vs 0.5). + const MIN_DENSITY_WIDE: f64 = 0.05; + + /// One full 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; + + pub(super) fn filter(values: &[T], mask: &MaskValues) -> Option> { + debug_assert_eq!(values.len(), mask.len()); + if values.len() < MIN_LEN || too_sparse::(mask) { + return None; + } + let kernel = select_kernel::()?; + + 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()) + } + + pub(super) fn filter_mut(values: &mut [T], mask: &MaskValues) -> Option { + debug_assert_eq!(values.len(), mask.len()); + if values.len() < MIN_LEN || too_sparse::(mask) { + return None; + } + let kernel = select_kernel::()?; + + 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 masks partial tail stores, so all accesses stay inside `values`. + let written = unsafe { kernel(dst.cast_const(), dst, mask) }; + debug_assert_eq!(written, mask.true_count()); + Some(written) + } + + /// True when the mask is sparse enough that the scalar walk wins for this element width. + fn too_sparse(mask: &MaskValues) -> bool { + size_of::() > 1 && mask.density() < MIN_DENSITY_WIDE + } + + /// Choose the widest kernel the CPU supports for this element width, if any. + /// + /// `is_x86_feature_detected!` caches per feature, so per-buffer selection stays cheap. + fn select_kernel() -> Option { + match size_of::() { + 1 if avx512_vbmi2() => Some(compress_avx512_epi8:: as Kernel), + 2 if avx512_vbmi2() => Some(compress_avx512_epi16:: as Kernel), + 4 if avx512f() => Some(compress_avx512_epi32:: as Kernel), + 8 if avx512f() => Some(compress_avx512_epi64:: as Kernel), + 4 if avx2() => Some(compress_avx2_epi32:: as Kernel), + 8 if avx2() => Some(compress_avx2_epi64:: as Kernel), + _ => None, + } + } + + 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") + } + + /// 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`. + #[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) }; + } + } + + /// 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`]/[`filter_mut`] 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`]/[`filter_mut`] 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`]/[`filter_mut`] + /// 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`]/[`filter_mut`] + /// 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 + ); +} + +#[cfg(test)] +#[allow(clippy::cast_possible_truncation)] +mod tests { + use rstest::rstest; + use vortex_buffer::BitBuffer; + use vortex_mask::Mask; + + use super::super::slice; + #[cfg(all(target_arch = "x86_64", not(miri)))] + use super::x86; + 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(target_arch = "x86_64", 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"); + if is_x86_feature_detected!("avx2") { + assert!(filter_slice_by_bitmap(&values, mask).is_some()); + } + } + + /// On AVX-512 machines the dispatcher never selects the AVX2 tier for 4/8-byte elements, + /// so exercise those kernels directly. + #[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: unsafe fn(*const u8, *mut u8, &MaskValues) -> usize, + kernel_in_place: unsafe fn(*const u8, *mut u8, &MaskValues) -> usize, + values: &[T], + mask: &MaskValues, + ) { + let expected = slice::filter_slice_by_bitmap(values, mask); + + // One vector of slack, mirroring the allocation in `x86::filter`. + let mut out = vec![T::default(); mask.true_count() + 64 / 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 u32_values: Vec = (0..len as u32).collect(); + let u64_values: Vec = (0..len as u64).collect(); + 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/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index d6ec94db062..c778be13fd2 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -12,8 +12,10 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_mask::MaskValues; -#[inline] -fn for_each_mask_word(mask: &MaskValues, mut f: impl FnMut(u64, usize, usize)) { +/// 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(); @@ -39,8 +41,9 @@ fn for_each_mask_word(mask: &MaskValues, mut f: impl FnMut(u64, usize, usize)) { debug_assert_eq!(base, mask.len()); } +/// A `u64` with the low `len` bits set. #[inline] -fn low_bits_mask(len: usize) -> u64 { +pub(super) fn low_bits_mask(len: usize) -> u64 { debug_assert!(len <= 64); if len == 64 { u64::MAX From 6ac57693f1fab440de3880b8c2ac6e6e632dee74 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 20 Jul 2026 14:59:28 +0100 Subject: [PATCH 4/7] flags Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/filter/execute/simd_compress.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/simd_compress.rs b/vortex-array/src/arrays/filter/execute/simd_compress.rs index eaa55327d2e..a65c701f5ab 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress.rs @@ -39,9 +39,9 @@ pub(super) fn filter_slice_by_bitmap( ) -> Option> { #[cfg(all(target_arch = "x86_64", not(miri)))] { - return x86::filter(values, mask); + x86::filter(values, mask) } - #[allow(unreachable_code)] + #[cfg(any(not(target_arch = "x86_64"), miri))] { let _ = (values, mask); None @@ -56,9 +56,9 @@ pub(super) fn filter_slice_mut_by_bitmap( ) -> Option { #[cfg(all(target_arch = "x86_64", not(miri)))] { - return x86::filter_mut(values, mask); + x86::filter_mut(values, mask) } - #[allow(unreachable_code)] + #[cfg(any(not(target_arch = "x86_64"), miri))] { let _ = (values, mask); None From 5a00ebab336f6a46081a8856b0fccb16a81be7a9 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 20 Jul 2026 17:55:34 +0100 Subject: [PATCH 5/7] something Signed-off-by: Robert Kruszewski --- .../src/arrays/filter/execute/buffer.rs | 7 +- vortex-array/src/arrays/filter/execute/mod.rs | 100 +++++++++++++++++- vortex-buffer/src/bit/buf.rs | 63 +++++++++++ 3 files changed, 161 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index e21ee3ba973..967327e8ce8 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -121,12 +121,7 @@ pub(super) fn prepare_mask_for_reuse(mask: &MaskValues, consumers: usize) { return; } - let first = mask.bit_buffer().select(0); - let last = mask.bit_buffer().select(mask.true_count() - 1); - if first - .zip(last) - .is_some_and(|(first, last)| last - first + 1 == mask.true_count()) - { + if super::contiguous_values_range(mask).is_some() { return; } diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 4b1b83b4fcd..29325c57c82 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -55,9 +55,43 @@ fn filter_validity(validity: Validity, mask: &Arc) -> Validity { } pub(super) fn contiguous_filter_range(mask: &Mask) -> Option> { - let start = mask.first()?; - let end = mask.last()?.checked_add(1)?; - (end - start == mask.true_count()).then_some(start..end) + 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) { @@ -145,6 +179,7 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca #[cfg(test)] mod tests { + use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use super::*; @@ -167,9 +202,68 @@ mod tests { 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-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)] From a6fa5dd19665fc5ee87c2134f5e360cdee0d4f77 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 28 Jul 2026 15:42:16 +0100 Subject: [PATCH 6/7] NEON tbl compress kernels, and measured density bands for x86 Adds a NEON tier to the fixed-width filter compress kernels and retunes the x86 tier against real AVX-512 hardware. NEON has no `vpcompress`, but `tbl` *is* a cross-lane byte shuffle, so a byte-index LUT covering one vector of lanes serves every element width - unlike AVX2, which has no byte shuffle and so leaves 1/2-byte elements on the scalar byte-LUT. Lanes per shuffle is 8/8/4/2 for 1/2/4/8-byte elements. The chunk loop is deliberately branchless. An `m != 0` guard looks like a cheap skip but is a large pessimization: the mask bits of one chunk are as random as the data, so `P(m == 0) = (1 - density)^lanes` crosses 0.5 inside the useful density range for every width and the branch mispredicts near-maximally right where it would pay off. Guarding cost 1.95x at density 0.5 for 4-byte elements and 1.45x at 0.1 for 1-byte. Storing a full chunk for an empty mask is harmless because `write_pos` does not advance, so the next store overwrites the garbage. x86's `MIN_DENSITY_WIDE = 0.05` was set without AVX-512 hardware and with only two density points measured, so it never probed its own crossover. On Zen 5 it regressed 4-byte and 8-byte filtering by up to 32%, and AVX2 8-byte by 2.3x. Both tiers now carry per-(kernel, width) floors, since what sets the floor is how many lanes one compress covers. Bands are measured per architecture, interleaved against the same binary with the kernels disabled, and sit just above the crossover so a core with a different shuffle-to-scalar balance cannot regress. Verified no regression at any density on both architectures. Within band: NEON (Apple M-series): 1 byte 1.4-10.1x, 2 bytes 1.3-9.2x, 4 bytes 1.2-5.1x, 8 bytes 1.3x over 0.50..0.80 x86 (Zen 5, AVX-512): 1 byte 1.9-5.7x, 2 bytes 1.1-3.9x, 4 bytes 1.1-2.8x, 8 bytes 1.1-3.1x x86 (Zen 5, AVX2 only): 4 bytes 1.1-2.5x, 8 bytes 1.4-2.0x The 8-byte NEON kernel is the only one needing an upper bound: two lanes per shuffle never amortizes far, and past ~0.8 the byte-LUT's all-set bulk copies overtake it. Also promotes simd_compress to a directory module, and installs mimalloc as the global allocator in the benchmark. Without it the system allocator's behaviour at these output sizes dominates and destabilizes divan's per-iteration timing, swinging identical code by up to 20x. Checks: cargo +nightly fmt --all --check; cargo clippy -p vortex-array --all-targets --all-features; cargo nextest run -p vortex-array -p vortex-buffer -p vortex-mask (4058 passed); cargo test --doc; rustdoc with -D warnings. Full vortex-array suite also run on the AVX-512 machine (3180 passed), which is the first time the AVX-512 and AVX2 kernels have executed. Signed-off-by: Robert Kruszewski --- vortex-array/benches/filter_fixed_width.rs | 4 + .../src/arrays/filter/execute/buffer.rs | 4 + .../arrays/filter/execute/simd_compress.rs | 761 ------------------ .../filter/execute/simd_compress/mod.rs | 156 ++++ .../filter/execute/simd_compress/neon.rs | 317 ++++++++ .../filter/execute/simd_compress/tests.rs | 174 ++++ .../filter/execute/simd_compress/x86.rs | 472 +++++++++++ 7 files changed, 1127 insertions(+), 761 deletions(-) delete mode 100644 vortex-array/src/arrays/filter/execute/simd_compress.rs create mode 100644 vortex-array/src/arrays/filter/execute/simd_compress/mod.rs create mode 100644 vortex-array/src/arrays/filter/execute/simd_compress/neon.rs create mode 100644 vortex-array/src/arrays/filter/execute/simd_compress/tests.rs create mode 100644 vortex-array/src/arrays/filter/execute/simd_compress/x86.rs 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 967327e8ce8..06610f5bf82 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -95,6 +95,10 @@ fn byte_compress_density_threshold() -> f64 { // 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, diff --git a/vortex-array/src/arrays/filter/execute/simd_compress.rs b/vortex-array/src/arrays/filter/execute/simd_compress.rs deleted file mode 100644 index a65c701f5ab..00000000000 --- a/vortex-array/src/arrays/filter/execute/simd_compress.rs +++ /dev/null @@ -1,761 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Runtime-dispatched SIMD compress kernels for fixed-width filtering. -//! -//! This mirrors the multiversioning pattern of `vortex_buffer::bit::pack` -//! (`collect_bool_words_multiversioned`): a shared per-mask-word walk instantiated behind -//! per-feature-level `#[target_feature]` wrappers, selected once per buffer by runtime feature -//! detection. Unlike `collect_bool`, there is no caller-supplied closure for the -//! `#[target_feature]` boundary to deoptimize — the per-word body is a fixed gather — so every -//! supported element width routes here unconditionally and density thresholds only matter for -//! the scalar fallbacks. -//! -//! Two kernel tiers, dispatched on element width: -//! -//! - 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 only): a per-mask-byte (or nibble) permutation LUT feeding -//! `vpermd`, a vectorized version of the scalar `BYTE_COMPRESS_LUT`. AVX2 has no cross-lane -//! byte shuffle, so 1/2-byte elements stay on the scalar byte-LUT below AVX-512. -//! -//! 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 masks the -//! stores of partial tail chunks instead. - -use vortex_buffer::Buffer; -use vortex_mask::MaskValues; - -/// Filter a slice with a runtime-detected SIMD compress kernel, if one applies. -/// -/// Returns `None` when no kernel applies (non-x86 target, unsupported element width, missing -/// CPU features, or an input too short to amortize dispatch); the caller then falls back to -/// the scalar strategies. -pub(super) fn filter_slice_by_bitmap( - values: &[T], - mask: &MaskValues, -) -> Option> { - #[cfg(all(target_arch = "x86_64", not(miri)))] - { - x86::filter(values, mask) - } - #[cfg(any(not(target_arch = "x86_64"), miri))] - { - let _ = (values, mask); - None - } -} - -/// 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 { - #[cfg(all(target_arch = "x86_64", not(miri)))] - { - x86::filter_mut(values, mask) - } - #[cfg(any(not(target_arch = "x86_64"), miri))] - { - let _ = (values, mask); - None - } -} - -#[cfg(all(target_arch = "x86_64", not(miri)))] -mod x86 { - 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 std::ptr; - - use vortex_buffer::Buffer; - use vortex_buffer::BufferMut; - use vortex_mask::MaskValues; - - use super::super::slice::for_each_mask_word; - use super::super::slice::low_bits_mask; - - /// Below one full mask word the scalar paths win; skip feature detection entirely. - const MIN_LEN: usize = 64; - - /// Below this mask density the scalar set-bit walk beats vector compress for multi-byte - /// elements: with under ~3 set bits per mask word, the per-subword loop and full-width - /// stores are pure overhead. Single-byte elements consume a whole mask word per compress - /// and win at every density. Covered by `benches/filter_fixed_width.rs` (0.01 vs 0.5). - const MIN_DENSITY_WIDE: f64 = 0.05; - - /// One full 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; - - pub(super) fn filter(values: &[T], mask: &MaskValues) -> Option> { - debug_assert_eq!(values.len(), mask.len()); - if values.len() < MIN_LEN || too_sparse::(mask) { - return None; - } - let kernel = select_kernel::()?; - - 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()) - } - - pub(super) fn filter_mut(values: &mut [T], mask: &MaskValues) -> Option { - debug_assert_eq!(values.len(), mask.len()); - if values.len() < MIN_LEN || too_sparse::(mask) { - return None; - } - let kernel = select_kernel::()?; - - 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 masks partial tail stores, so all accesses stay inside `values`. - let written = unsafe { kernel(dst.cast_const(), dst, mask) }; - debug_assert_eq!(written, mask.true_count()); - Some(written) - } - - /// True when the mask is sparse enough that the scalar walk wins for this element width. - fn too_sparse(mask: &MaskValues) -> bool { - size_of::() > 1 && mask.density() < MIN_DENSITY_WIDE - } - - /// Choose the widest kernel the CPU supports for this element width, if any. - /// - /// `is_x86_feature_detected!` caches per feature, so per-buffer selection stays cheap. - fn select_kernel() -> Option { - match size_of::() { - 1 if avx512_vbmi2() => Some(compress_avx512_epi8:: as Kernel), - 2 if avx512_vbmi2() => Some(compress_avx512_epi16:: as Kernel), - 4 if avx512f() => Some(compress_avx512_epi32:: as Kernel), - 8 if avx512f() => Some(compress_avx512_epi64:: as Kernel), - 4 if avx2() => Some(compress_avx2_epi32:: as Kernel), - 8 if avx2() => Some(compress_avx2_epi64:: as Kernel), - _ => None, - } - } - - 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") - } - - /// 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`. - #[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) }; - } - } - - /// 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`]/[`filter_mut`] 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`]/[`filter_mut`] 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`]/[`filter_mut`] - /// 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`]/[`filter_mut`] - /// 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 - ); -} - -#[cfg(test)] -#[allow(clippy::cast_possible_truncation)] -mod tests { - use rstest::rstest; - use vortex_buffer::BitBuffer; - use vortex_mask::Mask; - - use super::super::slice; - #[cfg(all(target_arch = "x86_64", not(miri)))] - use super::x86; - 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(target_arch = "x86_64", 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"); - if is_x86_feature_detected!("avx2") { - assert!(filter_slice_by_bitmap(&values, mask).is_some()); - } - } - - /// On AVX-512 machines the dispatcher never selects the AVX2 tier for 4/8-byte elements, - /// so exercise those kernels directly. - #[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: unsafe fn(*const u8, *mut u8, &MaskValues) -> usize, - kernel_in_place: unsafe fn(*const u8, *mut u8, &MaskValues) -> usize, - values: &[T], - mask: &MaskValues, - ) { - let expected = slice::filter_slice_by_bitmap(values, mask); - - // One vector of slack, mirroring the allocation in `x86::filter`. - let mut out = vec![T::default(); mask.true_count() + 64 / 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 u32_values: Vec = (0..len as u32).collect(); - let u64_values: Vec = (0..len as u64).collect(); - 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/mod.rs b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs new file mode 100644 index 00000000000..04b490b6810 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs @@ -0,0 +1,156 @@ +// 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 only): a per-mask-byte (or nibble) permutation LUT feeding +//! `vpermd`, a vectorized version of the scalar `BYTE_COMPRESS_LUT`. AVX2 has no cross-lane +//! byte shuffle, so 1/2-byte elements stay on the scalar byte-LUT below AVX-512. +//! - NEON: `tbl` with a byte-index LUT. `tbl` *is* a cross-lane byte shuffle, so unlike AVX2 it +//! serves every element width, but its table is one register, so wide elements get few lanes +//! per shuffle. +//! +//! 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 + } +} + +/// 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..d14eb8ed2d3 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs @@ -0,0 +1,317 @@ +// 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; + +/// 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, + } +} + +/// Build the `tbl` 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 and `ROWS` must cover every mask value over the `BYTES / +/// elem_size` lanes it holds; both are checked at compile time by the `static` initializers. +#[expect( + clippy::cast_possible_truncation, + reason = "byte indices are bounded by the 8- or 16-byte table width" +)] +const fn compress_lut( + elem_size: usize, +) -> [[u8; BYTES]; ROWS] { + let lanes = BYTES / elem_size; + assert!(ROWS == 1 << lanes); + + 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 +} + +static IDX_LUT_8: [[u8; 8]; 256] = compress_lut::<256, 8>(1); +static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(2); +static IDX_LUT_32: [[u8; 16]; 16] = compress_lut::<16, 16>(4); +static IDX_LUT_64: [[u8; 16]; 4] = compress_lut::<4, 16>(8); + +/// 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`](super::filter_slice_by_bitmap) / +/// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold, and `bits` must +/// only select elements that are in bounds. +#[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 +} + +/// 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..ada294fe6b5 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs @@ -0,0 +1,174 @@ +// 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. +#[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 u32_values: Vec = (0..len as u32).collect(); + let u64_values: Vec = (0..len as u64).collect(); + 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..63f1b3231af --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs @@ -0,0 +1,472 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! AVX-512 `vpcompress` and AVX2 `vpermd` compress kernels. +//! +//! See the [module docs](super) for how these fit the shared dispatch. + +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; + +/// 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 | 1.9x at 0.05 | 5.5x at 0.8 | +/// | `vpcompressw` | 32 | 2 bytes | 0.15 | 1.1x | 3.8x 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.2x at 0.65 | +/// | `vpermd` | 8 | 4 bytes | 0.25 | 1.1x | 2.5x at 0.8 | +/// | `vpermd` | 4 | 8 bytes | 0.45 | 1.4x | 2.0x at 0.65 | +/// +/// `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), + 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 +); From b4714420d70bda53515182158c66e674e77c2b7c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 28 Jul 2026 16:49:35 +0100 Subject: [PATCH 7/7] pshufb compress kernels for 1/2-byte elements on x86 Gives 1- and 2-byte elements a vector compress on the AVX2 tier, where they previously had none and fell through to the scalar byte-LUT. The justification I committed for that gap was wrong. It claimed AVX2 has "no cross-lane byte shuffle, so 1/2-byte elements stay on the scalar byte-LUT" - but `pshufb` only has to span as many bytes as the chunk holds, and eight lanes of a 1- or 2-byte element need at most a 16-byte table. One 128-bit lane covers that, so the NEON `tbl` construction ports over unchanged: same LUT layout, same branchless chunk loop, only the load/shuffle/store intrinsics differ. `compress_lut` and `compress_tail` move up to the parent module and are now shared by both. Enabling AVX2 rather than SSSE3 gets the VEX-encoded `vpshufb`, avoiding AVX/SSE transition penalties in mixed code. AVX-512F parts without VBMI2 (Skylake-X) fall through to these too. Measured on Zen 5 with AVX-512 detection forced off, interleaved against the same binary with the kernels disabled: 1 byte, floor 0.15: 1.5x at floor, 2.6x at 0.8 2 bytes, floor 0.25: 1.0x at floor, 2.6x at 0.65 The floors differ despite the identical lane count because the two widths compete against different scalar strategies: `byte_compress_density_threshold` sends 1-byte elements to the byte-LUT at every density on x86, and that LUT is weak on a sparse mask, while 2-byte elements fall to the set-bit walk, which holds out to 0.25. A floor of 0.15 for 2-byte regressed 27%. Retuning that threshold would raise the 1-byte floor to match, and is noted in the docs. Also corrects two numbers in the previous commit that came from a single unreplicated run: the AVX2 8-byte `vpermd` kernel is 1.0x at its 0.45 floor and 1.4x at 0.65, not 1.4x and 2.0x. Two further runs agree. Verified no regression at any density on either x86 tier, and that the shared helper move left the NEON kernels' timings unchanged. `avx2_kernels_match_scalar` now covers the pshufb kernels directly, since a VBMI2 machine's dispatcher never selects them. Checks: cargo +nightly fmt --all --check; cargo clippy -p vortex-array --all-targets --all-features; cargo clippy --target x86_64-unknown-linux-gnu -p vortex-array --lib; cargo nextest run -p vortex-array -p vortex-buffer -p vortex-mask (4058 passed); rustdoc with -D warnings. On the AVX-512 machine: full vortex-array suite (3180 passed) and clippy --all-targets. Signed-off-by: Robert Kruszewski --- .../filter/execute/simd_compress/mod.rs | 87 ++++++++- .../filter/execute/simd_compress/neon.rs | 77 +------- .../filter/execute/simd_compress/tests.rs | 17 +- .../filter/execute/simd_compress/x86.rs | 174 +++++++++++++++++- 4 files changed, 271 insertions(+), 84 deletions(-) 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 04b490b6810..89fe9a80125 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs @@ -10,12 +10,16 @@ //! //! - 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 only): a per-mask-byte (or nibble) permutation LUT feeding -//! `vpermd`, a vectorized version of the scalar `BYTE_COMPRESS_LUT`. AVX2 has no cross-lane -//! byte shuffle, so 1/2-byte elements stay on the scalar byte-LUT below AVX-512. -//! - NEON: `tbl` with a byte-index LUT. `tbl` *is* a cross-lane byte shuffle, so unlike AVX2 it -//! serves every element width, but its table is one register, so wide elements get few lanes -//! per shuffle. +//! - 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 @@ -123,6 +127,77 @@ fn select_kernel(mask: &MaskValues) -> Option { } } +/// 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 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 d14eb8ed2d3..da9d250a2b8 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs @@ -39,6 +39,8 @@ 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. /// @@ -90,77 +92,10 @@ pub(super) fn select_kernel(mask: &MaskValues) -> Optio } } -/// Build the `tbl` 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 and `ROWS` must cover every mask value over the `BYTES / -/// elem_size` lanes it holds; both are checked at compile time by the `static` initializers. -#[expect( - clippy::cast_possible_truncation, - reason = "byte indices are bounded by the 8- or 16-byte table width" -)] -const fn compress_lut( - elem_size: usize, -) -> [[u8; BYTES]; ROWS] { - let lanes = BYTES / elem_size; - assert!(ROWS == 1 << lanes); - - 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 -} - -static IDX_LUT_8: [[u8; 8]; 256] = compress_lut::<256, 8>(1); -static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(2); -static IDX_LUT_32: [[u8; 16]; 16] = compress_lut::<16, 16>(4); -static IDX_LUT_64: [[u8; 16]; 4] = compress_lut::<4, 16>(8); - -/// 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`](super::filter_slice_by_bitmap) / -/// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold, and `bits` must -/// only select elements that are in bounds. -#[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 -} +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 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 ada294fe6b5..c76da1dd80d 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs @@ -115,7 +115,8 @@ fn declines_sparse_and_short_masks() { } /// The dispatcher never selects the AVX2 tier on AVX-512 machines, so exercise those kernels -/// directly. +/// 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() { @@ -154,8 +155,22 @@ fn avx2_kernels_match_scalar() { 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::, 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 63f1b3231af..8cd994608b1 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs @@ -1,10 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! AVX-512 `vpcompress` and AVX2 `vpermd` compress kernels. +//! 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; @@ -39,6 +45,8 @@ 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. @@ -55,12 +63,20 @@ use super::bulk_copy; /// /// | kernel | lanes per compress | element | floor | at floor | best in band | /// | --- | --- | --- | --- | --- | --- | -/// | `vpcompressb` | 64 | 1 byte | none | 1.9x at 0.05 | 5.5x at 0.8 | -/// | `vpcompressw` | 32 | 2 bytes | 0.15 | 1.1x | 3.8x at 0.8 | +/// | `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.2x at 0.65 | -/// | `vpermd` | 8 | 4 bytes | 0.25 | 1.1x | 2.5x at 0.8 | -/// | `vpermd` | 4 | 8 bytes | 0.45 | 1.4x | 2.0x at 0.65 | +/// | `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 { @@ -69,6 +85,9 @@ pub(super) fn select_kernel(mask: &MaskValues) -> Optio 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, @@ -470,3 +489,146 @@ avx2_compress_kernel!( 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 +);