diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index 553bb4c8dc3..2cf9f1f950b 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -6,6 +6,7 @@ mod avx2; mod records; mod scalar; mod slices; +mod small_table; #[cfg(test)] mod tests; @@ -61,7 +62,16 @@ impl TakeExecute for V { /// /// Implementors must have no uninitialized bytes. The shared AVX2 gather reads the complete /// representation through a same-width integer lane before writing those bytes back unchanged. -pub(crate) unsafe trait FixedWidthTakeValue: Copy {} +pub(crate) unsafe trait FixedWidthTakeValue: Copy { + /// Takes values using the kernel appropriate for this value type. + fn take( + values: &[Self], + indices: &[I], + allocator: &BufferAllocatorRef, + ) -> Buffer { + take_values_fallback(values, indices, allocator) + } +} macro_rules! impl_fixed_width_take_value { ($($ty:ty),+ $(,)?) => { @@ -72,15 +82,20 @@ macro_rules! impl_fixed_width_take_value { }; } -impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,); - -// SAFETY: Byte arrays have no padding and every byte is initialized. -unsafe impl FixedWidthTakeValue for [u8; N] {} +impl_fixed_width_take_value!(u16, u32, u64, i16, i32, i64, f16, f32, f64,); pub(crate) fn take_values( values: &[T], indices: &[I], allocator: &BufferAllocatorRef, +) -> Buffer { + T::take(values, indices, allocator) +} + +fn take_values_fallback( + values: &[T], + indices: &[I], + allocator: &BufferAllocatorRef, ) -> Buffer { #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] if *HAS_AVX2 { diff --git a/vortex-array/src/arrays/fixed_width/take/small_table.rs b/vortex-array/src/arrays/fixed_width/take/small_table.rs new file mode 100644 index 00000000000..cea9c429ab6 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Small byte-table take with a fallback for unsupported targets and inputs. + +cfg_if::cfg_if! { + if #[cfg(all(target_arch = "aarch64", target_endian = "little"))] { + mod neon; + } else { + mod fallback { + use super::super::FixedWidthTakeValue; + + // SAFETY: u8 has no padding or uninitialized bytes. + unsafe impl FixedWidthTakeValue for u8 {} + + // SAFETY: i8 has no padding or uninitialized bytes. + unsafe impl FixedWidthTakeValue for i8 {} + + // SAFETY: Byte arrays have no padding and every byte is initialized. + unsafe impl FixedWidthTakeValue for [u8; N] {} + } + } +} diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs new file mode 100644 index 00000000000..46b8450ce85 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! NEON byte-table take for `u8` codes and at most 16 one-byte values. + +use std::arch::aarch64::uint8x16_t; +use std::arch::aarch64::vdupq_n_u8; +use std::arch::aarch64::vld1q_u8; +use std::arch::aarch64::vmaxq_u8; +use std::arch::aarch64::vmaxvq_u8; +use std::arch::aarch64::vqtbl1q_u8; +use std::arch::aarch64::vst1q_u8; + +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; + +use super::super::FixedWidthTakeValue; +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 + { + 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()) }; + + let mut table = [values[0]; 16]; + table[..values.len()].copy_from_slice(values); + + 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::(); + // SAFETY: AArch64 always provides NEON. T is one byte with no uninitialized bytes, the table + // contains 16 values, and the output has capacity for every index. + let (offset, max_code) = unsafe { take_vectors(&table, indices, output_ptr) }; + + 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]); + } + assert!( + usize::from(max_code) < values.len(), + "take index {max_code} out of bounds for length {}", + values.len() + ); + // SAFETY: the vector loop and scalar remainder initialized every output value. + unsafe { output.set_len(indices.len()) }; + output.freeze() +} + +unsafe fn take_vectors( + table: &[T; 16], + indices: &[u8], + output: *mut u8, +) -> (usize, u8) { + let table = unsafe { vld1q_u8(table.as_ptr().cast::()) }; + let mut max_codes: uint8x16_t = unsafe { vdupq_n_u8(0) }; + let mut offset = 0; + while offset + 16 <= indices.len() { + let codes = unsafe { vld1q_u8(indices.as_ptr().add(offset)) }; + max_codes = unsafe { vmaxq_u8(max_codes, codes) }; + unsafe { vst1q_u8(output.add(offset), vqtbl1q_u8(table, codes)) }; + offset += 16; + } + (offset, unsafe { vmaxvq_u8(max_codes) }) +} diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index 1d97bde8d3a..98b46b7cef7 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -43,6 +43,37 @@ fn take_eight_byte_values() { assert_eq!(taken.as_slice(), &[20, 30, 10]); } +#[rstest] +#[case(64)] +#[case(65)] +#[case(79)] +#[case(128)] +fn take_small_byte_table(#[case] len: usize) { + let values = [10u8, 20, 30, 40]; + let indices = (0u8..4).cycle().take(len).collect::>(); + let expected = indices + .iter() + .map(|index| values[usize::from(*index)]) + .collect::>(); + assert_eq!( + take_values(&values, &indices, &allocator()).as_slice(), + expected.as_slice() + ); +} + +// 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")] +fn take_small_byte_table_rejects_out_of_bounds_index(#[case] offset: usize) { + let mut indices = vec![0u8; 129]; + indices[offset] = 4; + drop(take_values(&[10u8, 20, 30, 40], &indices, &allocator())); +} + #[rstest] #[case(1)] #[case(2)]