perf: vectorize small u8 table take with NEON

Use NEON TBL for u8-coded tables with at most 16 one-byte values.

Signed-off-by: Joseph Isaacs <joseph-isaacs@users.noreply.github.com>

Author: joseph-isaacs

vortex-array/src/arrays/fixed_width/take/mod.rs

@@ -6,6 +6,8 @@ mod avx2;
 mod records;
 mod scalar;
 mod slices;
+#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
+mod small_table;
 #[cfg(test)]
 mod tests;
 
@@ -36,6 +38,8 @@ 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"))]
+use crate::dtype::PType;
 use crate::dtype::UnsignedPType;
 use crate::dtype::half::f16;
 use crate::match_each_unsigned_integer_ptype;
@@ -80,6 +84,16 @@ pub(crate) fn take_values<T: FixedWidthTakeValue, I: UnsignedPType>(
     values: &[T],
     indices: &[I],
 ) -> Buffer<T> {
+    #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
+    if I::PTYPE == PType::U8 {
+        // 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

vortex-array/src/arrays/fixed_width/take/small_table.rs

@@ -0,0 +1,74 @@
+// 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::BufferMut;
+
+use super::FixedWidthTakeValue;
+
+pub(crate) fn take<T: FixedWidthTakeValue>(values: &[T], indices: &[u8]) -> Option<Buffer<T>> {
+    if values.is_empty() || values.len() > 16 || size_of::<T>() != 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::<u8>(),
+            table.as_mut_ptr(),
+            values.len(),
+        );
+    }
+
+    let mut output = BufferMut::<T>::with_capacity(indices.len());
+    let output_ptr = output.spare_capacity_mut().as_mut_ptr().cast::<u8>();
+    // 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]);
+        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::<u8>())
+        };
+    }
+    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;
+    }
+    (offset, unsafe { vmaxvq_u8(max_codes) })
+}

vortex-array/src/arrays/fixed_width/take/tests.rs

@@ -38,6 +38,28 @@ fn take_eight_byte_values() {
     assert_eq!(taken.as_slice(), &[20, 30, 10]);
 }
 
+#[test]
+fn take_small_byte_table() {
+    let values = [10u8, 20, 30, 40];
+    let indices = (0..128).map(|index| index % 4).collect::<Vec<u8>>();
+    let expected = indices
+        .iter()
+        .map(|index| values[usize::from(*index)])
+        .collect::<Vec<_>>();
+    assert_eq!(
+        take_values(&values, &indices).as_slice(),
+        expected.as_slice()
+    );
+}
+
+#[test]
+#[should_panic(expected = "take index")]
+fn take_small_byte_table_rejects_out_of_bounds_index() {
+    let mut indices = vec![0u8; 128];
+    indices[64] = 4;
+    drop(take_values(&[10u8, 20, 30, 40], &indices));
+}
+
 #[rstest]
 #[case(1)]
 #[case(2)]