From 0129184aa7a651b8190f20b6bb71bad0d05c8fc0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 24 Aug 2026 13:53:11 +0100 Subject: [PATCH] perf: vectorize small u8 table take with AVX2 Use AVX2 VPSHUFB for u8-coded tables with at most 16 one-byte values. Signed-off-by: Joseph Isaacs --- .../src/arrays/fixed_width/take/mod.rs | 22 +- .../arrays/fixed_width/take/small_table.rs | 197 +++++++++++++----- 2 files changed, 161 insertions(+), 58 deletions(-) diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index 78684e5b385..56ca69731c2 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -6,7 +6,11 @@ mod avx2; mod records; mod scalar; mod slices; -#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +#[cfg(any( + all(target_arch = "aarch64", target_endian = "little"), + target_arch = "x86_64", + target_arch = "x86" +))] mod small_table; #[cfg(test)] mod tests; @@ -38,7 +42,11 @@ use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +#[cfg(any( + all(target_arch = "aarch64", target_endian = "little"), + target_arch = "x86_64", + target_arch = "x86" +))] use crate::dtype::PType; use crate::dtype::UnsignedPType; use crate::dtype::half::f16; @@ -94,6 +102,16 @@ pub(crate) fn take_values( } } + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] + if I::PTYPE == PType::U8 && *HAS_AVX2 { + // SAFETY: the ptype dispatcher guarantees that `I::PTYPE == U8` is the concrete `u8` + // implementation, so these slices have identical layouts. + let indices = unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) }; + if let Some(taken) = small_table::take(values, indices) { + return taken; + } + } + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] if *HAS_AVX2 { // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte 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 2bb858ab721..f5d2fe53c53 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -1,74 +1,159 @@ // 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. +//! 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; +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +mod arch { + 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::BufferMut; + use vortex_buffer::Buffer; + use vortex_buffer::BufferMut; -use super::FixedWidthTakeValue; + use super::super::FixedWidthTakeValue; -pub(crate) fn take(values: &[T], indices: &[u8]) -> Option> { - if values.is_empty() || values.len() > 16 || size_of::() != 1 || indices.len() < 64 { - return None; - } + pub(crate) fn take(values: &[T], indices: &[u8]) -> Option> { + if values.is_empty() || values.len() > 16 || size_of::() != 1 || indices.len() < 64 { + return None; + } - let mut table = [0u8; 16]; - // SAFETY: one-byte values have the same representation as bytes, and the length is <= 16. - unsafe { - std::ptr::copy_nonoverlapping( - values.as_ptr().cast::(), - table.as_mut_ptr(), - values.len(), - ); - } + let mut table = [0u8; 16]; + // SAFETY: one-byte values have the same representation as bytes, and the length is <= 16. + unsafe { + std::ptr::copy_nonoverlapping( + values.as_ptr().cast::(), + table.as_mut_ptr(), + values.len(), + ); + } - let mut output = BufferMut::::with_capacity(indices.len()); - let output_ptr = output.spare_capacity_mut().as_mut_ptr().cast::(); - // SAFETY: AArch64 always provides NEON. Both pointers advance only by complete vectors. - let (offset, max_code) = unsafe { take_vectors(&table, indices, output_ptr) }; + let mut output = BufferMut::::with_capacity(indices.len()); + let output_ptr = output.spare_capacity_mut().as_mut_ptr().cast::(); + // SAFETY: AArch64 always provides NEON. Both pointers advance only by complete vectors. + let (offset, max_code) = unsafe { take_vectors(&table, indices, output_ptr) }; - for offset in offset..indices.len() { - let code = usize::from(indices[offset]); + 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() + ); + // SAFETY: the code was checked and this reserved output position is uninitialized. + unsafe { + output_ptr + .add(offset) + .write(*values.as_ptr().add(code).cast::()) + }; + } assert!( - code < values.len(), - "take index {code} out of bounds for length {}", + usize::from(max_code) < values.len(), + "take index {max_code} out of bounds for length {}", values.len() ); - // SAFETY: the code was checked and this reserved output position is uninitialized. - unsafe { - output_ptr - .add(offset) - .write(*values.as_ptr().add(code).cast::()) - }; + // SAFETY: the vector loop and scalar remainder initialized every output value. + unsafe { output.set_len(indices.len()) }; + Some(output.freeze()) + } + + unsafe fn take_vectors(table: &[u8; 16], indices: &[u8], output: *mut u8) -> (usize, u8) { + let table = unsafe { vld1q_u8(table.as_ptr()) }; + 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) }) } - 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()) }; - Some(output.freeze()) } -unsafe fn take_vectors(table: &[u8; 16], indices: &[u8], output: *mut u8) -> (usize, u8) { - let table = unsafe { vld1q_u8(table.as_ptr()) }; - 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; +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +mod arch { + use std::arch::x86_64::__m256i; + use std::arch::x86_64::_mm_loadu_si128; + use std::arch::x86_64::_mm256_broadcastsi128_si256; + use std::arch::x86_64::_mm256_loadu_si256; + use std::arch::x86_64::_mm256_or_si256; + use std::arch::x86_64::_mm256_set1_epi8; + use std::arch::x86_64::_mm256_setzero_si256; + use std::arch::x86_64::_mm256_shuffle_epi8; + use std::arch::x86_64::_mm256_storeu_si256; + use std::arch::x86_64::_mm256_subs_epu8; + use std::arch::x86_64::_mm256_testz_si256; + + use vortex_buffer::Buffer; + use vortex_buffer::BufferMut; + + use super::super::FixedWidthTakeValue; + + pub(crate) fn take(values: &[T], indices: &[u8]) -> Option> { + if values.is_empty() || values.len() > 16 || size_of::() != 1 || indices.len() < 64 { + return None; + } + // SAFETY: the caller detects AVX2 before entering this architecture-specific module. + Some(unsafe { take_avx2(values, indices) }) + } + + #[target_feature(enable = "avx2")] + unsafe fn take_avx2(values: &[T], indices: &[u8]) -> Buffer { + let mut table = [0u8; 16]; + unsafe { + std::ptr::copy_nonoverlapping( + values.as_ptr().cast::(), + table.as_mut_ptr(), + values.len(), + ); + } + let table = unsafe { _mm_loadu_si128(table.as_ptr().cast()) }; + let table = _mm256_broadcastsi128_si256(table); + let limit = _mm256_set1_epi8((values.len() - 1) as i8); + let mut invalid = _mm256_setzero_si256(); + let mut output = BufferMut::::with_capacity(indices.len()); + let output_ptr = output.spare_capacity_mut().as_mut_ptr().cast::(); + + let mut offset = 0; + while offset + 32 <= indices.len() { + let codes = unsafe { _mm256_loadu_si256(indices.as_ptr().add(offset).cast()) }; + invalid = _mm256_or_si256(invalid, _mm256_subs_epu8(codes, limit)); + 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() + ); + unsafe { + output_ptr + .add(offset) + .write(*values.as_ptr().add(code).cast::()) + }; + } + unsafe { output.set_len(indices.len()) }; + output.freeze() } - (offset, unsafe { vmaxvq_u8(max_codes) }) } + +pub(super) use arch::take;