diff --git a/vortex-array/src/arrays/fixed_width/take/small_table.rs b/vortex-array/src/arrays/fixed_width/take/small_table.rs index cea9c429ab6..3f60bee1a26 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -6,6 +6,8 @@ cfg_if::cfg_if! { if #[cfg(all(target_arch = "aarch64", target_endian = "little"))] { mod neon; + } else if #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] { + mod avx2; } else { mod fallback { use super::super::FixedWidthTakeValue; diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs b/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs new file mode 100644 index 00000000000..3cc9c83afec --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! AVX2 byte-table take for `u8` codes and at most 16 one-byte values. + +cfg_if::cfg_if! { + if #[cfg(target_arch = "x86")] { + use std::arch::x86 as arch; + } else { + use std::arch::x86_64 as arch; + } +} + +use arch::__m256i; +use arch::_mm_loadu_si128; +use arch::_mm256_broadcastsi128_si256; +use arch::_mm256_loadu_si256; +use arch::_mm256_or_si256; +use arch::_mm256_set1_epi8; +use arch::_mm256_setzero_si256; +use arch::_mm256_shuffle_epi8; +use arch::_mm256_storeu_si256; +use arch::_mm256_subs_epu8; +use arch::_mm256_testz_si256; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; + +use super::super::FixedWidthTakeValue; +use super::super::HAS_AVX2; +use super::super::take_values_fallback; +use crate::dtype::PType; +use crate::dtype::UnsignedPType; + +// SAFETY: u8 has no padding or uninitialized bytes. +unsafe impl FixedWidthTakeValue for u8 { + fn take( + values: &[Self], + indices: &[I], + allocator: &BufferAllocatorRef, + ) -> Buffer { + take(values, indices, allocator) + } +} + +// SAFETY: i8 has no padding or uninitialized bytes. +unsafe impl FixedWidthTakeValue for i8 { + fn take( + values: &[Self], + indices: &[I], + allocator: &BufferAllocatorRef, + ) -> Buffer { + take(values, indices, allocator) + } +} + +// SAFETY: Byte arrays have no padding and every byte is initialized. +unsafe impl FixedWidthTakeValue for [u8; N] { + fn take( + values: &[Self], + indices: &[I], + allocator: &BufferAllocatorRef, + ) -> Buffer { + take(values, indices, allocator) + } +} + +fn take( + values: &[T], + indices: &[I], + allocator: &BufferAllocatorRef, +) -> Buffer { + if I::PTYPE != PType::U8 + || values.is_empty() + || values.len() > 16 + || size_of::() != 1 + || indices.len() < 64 + || !*HAS_AVX2 + { + return take_values_fallback(values, indices, allocator); + } + + // SAFETY: the sealed index type is u8, as checked above. + let indices: &[u8] = + unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) }; + // SAFETY: AVX2 was detected above. Values are one byte with no uninitialized bytes, + // and the table contains between 1 and 16 values. + unsafe { take_avx2(values, indices, allocator) } +} + +#[target_feature(enable = "avx2")] +unsafe fn take_avx2( + values: &[T], + indices: &[u8], + allocator: &BufferAllocatorRef, +) -> Buffer { + let mut table = [values[0]; 16]; + table[..values.len()].copy_from_slice(values); + // SAFETY: the table contains 16 initialized one-byte values. + let table = unsafe { _mm_loadu_si128(table.as_ptr().cast()) }; + let table = _mm256_broadcastsi128_si256(table); + let limit = _mm256_set1_epi8( + i8::try_from(values.len() - 1).vortex_expect("table contains at most 16 values"), + ); + let mut invalid = _mm256_setzero_si256(); + let mut output = BufferMut::::with_capacity_in(indices.len(), allocator.clone()); + let spare = output.spare_capacity_mut(); + let output_ptr = spare.as_mut_ptr().cast::(); + + let mut offset = 0; + while offset + 32 <= indices.len() { + // SAFETY: the loop condition guarantees a complete 32-byte vector is in bounds. + let codes = unsafe { _mm256_loadu_si256(indices.as_ptr().add(offset).cast()) }; + invalid = _mm256_or_si256(invalid, _mm256_subs_epu8(codes, limit)); + // SAFETY: the output has capacity for every index and T is one byte wide. + unsafe { + _mm256_storeu_si256( + output_ptr.add(offset).cast::<__m256i>(), + _mm256_shuffle_epi8(table, codes), + ) + }; + offset += 32; + } + assert_eq!( + _mm256_testz_si256(invalid, invalid), + 1, + "take index out of bounds" + ); + + for offset in offset..indices.len() { + let code = usize::from(indices[offset]); + assert!( + code < values.len(), + "take index {code} out of bounds for length {}", + values.len() + ); + spare[offset].write(values[code]); + } + // SAFETY: the vector loop and scalar remainder initialized every output value. + unsafe { output.set_len(indices.len()) }; + output.freeze() +} diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index 98b46b7cef7..1d2b312826d 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -61,13 +61,10 @@ fn take_small_byte_table(#[case] len: usize) { ); } -// The bounds-check message is specific to the NEON table path; other targets reach a different -// fallback with its own message. -#[cfg(all(target_arch = "aarch64", target_endian = "little"))] #[rstest] #[case::vector(64)] #[case::tail(128)] -#[should_panic(expected = "take index")] +#[should_panic] fn take_small_byte_table_rejects_out_of_bounds_index(#[case] offset: usize) { let mut indices = vec![0u8; 129]; indices[offset] = 4;