chore(array): use fearless_simd for AVX2 1- and 2-byte filter compress (#10324)

Signed-off-by: Claude <noreply@anthropic.com>

Author: joseph-isaacs

Cargo.lock

@@ -31,6 +31,8 @@ bytes = { workspace = true }
 cfg-if = { workspace = true }
 cudarc = { workspace = true, optional = true }
 enum-iterator = { workspace = true }
+fearless_simd = { workspace = true }
+fearless_simd_macros = { workspace = true }
 flatbuffers = { workspace = true }
 futures = { workspace = true, features = ["alloc", "async-await", "std"] }
 goldenfile = { workspace = true, optional = true }

vortex-array/Cargo.toml

@@ -0,0 +1,129 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Portable `fearless_simd` compress kernels for 1- and 2-byte elements.
+//!
+//! Each sub-word of 8 lanes is compacted with one 16-byte `swizzle_dyn` driven by a byte-index
+//! lookup table. Dispatch wraps the whole mask walk, so the per-word body inlines into one
+//! target-feature function rather than being called once per mask word.
+
+use std::ptr;
+
+use fearless_simd::Level;
+use fearless_simd::Simd;
+use fearless_simd::dispatch;
+use fearless_simd::prelude::*;
+use fearless_simd::u8x16;
+use fearless_simd_macros::simd;
+use vortex_mask::MaskValues;
+
+use super::super::slice::for_each_mask_word;
+use super::super::slice::low_bits_mask;
+use super::bulk_copy;
+use super::compress_lut;
+use super::compress_tail;
+
+static IDX_LUT_8: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 1);
+static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2);
+
+/// Compact one mask word of `ELEM`-byte elements, `LANES` elements per shuffle.
+///
+/// # 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"
+)]
+#[allow(clippy::too_many_arguments)]
+#[simd]
+unsafe fn compress_word<S: Simd, const IN_PLACE: bool, const ELEM: usize, const LANES: usize>(
+    simd: S,
+    idx_lut: &[[u8; 16]],
+    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::<IN_PLACE>(src, dst, word_start, word_len, write_pos, ELEM) };
+        return write_pos + word_len;
+    }
+
+    // Empty chunks still store garbage that the next chunk overwrites; branching here
+    // regresses masks near the density crossover.
+    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.
+        let chunk_ptr = unsafe { src.add((word_start + sub) * ELEM) };
+        // Materializing the bytes ends the source read before an overlapping in-place store.
+        let bytes: [u8; 16] = if ELEM * LANES == 16 {
+            // SAFETY: see above; the chunk is exactly 16 bytes.
+            unsafe { chunk_ptr.cast::<[u8; 16]>().read_unaligned() }
+        } else {
+            // SAFETY: see above; the chunk is exactly 8 bytes.
+            let half = unsafe { chunk_ptr.cast::<[u8; 8]>().read_unaligned() };
+            let mut bytes = [0u8; 16];
+            bytes[..8].copy_from_slice(&half);
+            bytes
+        };
+        let packed = u8x16::from_slice(simd, &bytes)
+            .swizzle_dyn(u8x16::from_slice(simd, &idx_lut[m]))
+            .to_array();
+        // SAFETY: out-of-place output has vector slack. In-place, the store ends within the
+        // source chunk already loaded, and later stores overwrite trailing garbage.
+        unsafe {
+            ptr::copy_nonoverlapping(packed.as_ptr(), dst.add(write_pos * ELEM), ELEM * LANES)
+        };
+        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::<IN_PLACE>(src, dst, bits, word_start + sub, write_pos, ELEM) };
+    }
+
+    write_pos
+}
+
+/// Generate a mask-walking entry point for one element width.
+macro_rules! generic_compress_kernel {
+    ($walk_fn:ident,elem_size: $elem_size:literal,lanes: $lanes:literal,idx_lut: $idx_lut:ident) => {
+        /// # Safety
+        ///
+        /// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) /
+        /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold.
+        pub(super) unsafe fn $walk_fn<const IN_PLACE: bool>(
+            src: *const u8,
+            dst: *mut u8,
+            mask: &MaskValues,
+        ) -> usize {
+            dispatch!(Level::new(), simd => {
+                let mut write_pos = 0;
+                for_each_mask_word(mask, |word, word_start, word_len| {
+                    // SAFETY: forwarded from the caller contract.
+                    write_pos = unsafe {
+                        compress_word::<_, IN_PLACE, $elem_size, $lanes>(
+                            simd, &$idx_lut, src, dst, word, word_start, word_len, write_pos,
+                        )
+                    };
+                });
+                write_pos
+            })
+        }
+    };
+}
+
+generic_compress_kernel!(compress_generic_8, elem_size: 1, lanes: 8, idx_lut: IDX_LUT_8);
+generic_compress_kernel!(compress_generic_16, elem_size: 2, lanes: 8, idx_lut: IDX_LUT_16);

vortex-array/src/arrays/filter/execute/simd_compress/generic.rs

@@ -31,6 +31,8 @@ use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_mask::MaskValues;
 
+#[cfg(all(target_arch = "x86_64", not(miri)))]
+mod generic;
 #[cfg(all(target_arch = "aarch64", not(miri)))]
 mod neon;
 #[cfg(test)]

vortex-array/src/arrays/filter/execute/simd_compress/mod.rs

@@ -105,7 +105,7 @@ fn declines_sparse_and_short_masks() {
     assert!(filter_slice_by_bitmap(&values[..32], short).is_none());
 }
 
-// AVX-512 machines need direct coverage of the otherwise-unselected AVX2 tier.
+// AVX-512 machines need direct coverage of the otherwise-unselected AVX2-tier kernels.
 #[cfg(all(target_arch = "x86_64", not(miri)))]
 #[test]
 fn avx2_kernels_match_scalar() {
@@ -152,14 +152,14 @@ fn avx2_kernels_match_scalar() {
                 let u32_values: Vec<u32> = (0..len as u32).collect();
                 let u64_values: Vec<u64> = (0..len as u64).collect();
                 check_kernel(
-                    x86::compress_pshufb_epi8::<false>,
-                    x86::compress_pshufb_epi8::<true>,
+                    generic::compress_generic_8::<false>,
+                    generic::compress_generic_8::<true>,
                     &u8_values,
                     mask,
                 );
                 check_kernel(
-                    x86::compress_pshufb_epi16::<false>,
-                    x86::compress_pshufb_epi16::<true>,
+                    generic::compress_generic_16::<false>,
+                    generic::compress_generic_16::<true>,
                     &u16_values,
                     mask,
                 );

vortex-array/src/arrays/filter/execute/simd_compress/tests.rs

@@ -1,16 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! AVX-512 `vpcompress`, AVX2 `vpermd`, and 128-bit `pshufb` compress kernels.
+//! AVX-512 `vpcompress` and AVX2 `vpermd` compress kernels. The 1- and 2-byte AVX2 paths use the
+//! portable kernels in [`generic`](super::generic).
 //!
 //! 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;
@@ -45,8 +40,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;
+use super::generic::compress_generic_8;
+use super::generic::compress_generic_16;
 
 /// Choose the widest available kernel above its benchmarked density crossover.
 ///
@@ -59,8 +54,8 @@ pub(super) fn select_kernel<T, const IN_PLACE: bool>(mask: &MaskValues) -> Optio
         4 if avx512f() => (compress_avx512_epi32::<IN_PLACE> as Kernel, 0.25),
         8 if avx512f() => (compress_avx512_epi64::<IN_PLACE> as Kernel, 0.30),
         // AVX-512F without VBMI2 (e.g. Skylake-X) falls through to these too.
-        1 if avx2() => (compress_pshufb_epi8::<IN_PLACE> as Kernel, 0.15),
-        2 if avx2() => (compress_pshufb_epi16::<IN_PLACE> as Kernel, 0.25),
+        1 if avx2() => (compress_generic_8::<IN_PLACE> as Kernel, 0.15),
+        2 if avx2() => (compress_generic_16::<IN_PLACE> as Kernel, 0.25),
         4 if avx2() => (compress_avx2_epi32::<IN_PLACE> as Kernel, 0.25),
         8 if avx2() => (compress_avx2_epi64::<IN_PLACE> as Kernel, 0.45),
         _ => return None,
@@ -451,126 +446,3 @@ 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 an AVX2 `pshufb` kernel for 1- or 2-byte elements.
-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
-    ) => {
-        /// # 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<const IN_PLACE: bool>(
-            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::<IN_PLACE>(src, dst, word_start, word_len, write_pos, $elem_size)
-                };
-                return write_pos + word_len;
-            }
-
-            // Empty chunks still store garbage that the next chunk overwrites; branching here
-            // regresses masks near the density crossover.
-            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: out-of-place output has vector slack. In-place, the store ends within
-                // the source chunk already loaded, and later stores overwrite trailing garbage.
-                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::<IN_PLACE>(
-                        src,
-                        dst,
-                        bits,
-                        word_start + sub,
-                        write_pos,
-                        $elem_size,
-                    )
-                };
-            }
-
-            write_pos
-        }
-
-        /// # 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<const IN_PLACE: bool>(
-            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::<IN_PLACE>(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
-);