From a0aaec6c35273d10cf1ef8cece28c1dd8100ca81 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 18:49:17 +0100 Subject: [PATCH 1/4] perf: extend small byte dictionary take to 32 values Signed-off-by: Joe Isaacs --- vortex-array/Cargo.toml | 4 + vortex-array/benches/take_small_dict.rs | 85 +++++++++++++++++++ .../fixed_width/take/small_table/avx2.rs | 67 +++++++++++---- .../fixed_width/take/small_table/avx512.rs | 65 ++++++++++++++ .../fixed_width/take/small_table/tests.rs | 76 +++++++++++++++++ 5 files changed, 282 insertions(+), 15 deletions(-) create mode 100644 vortex-array/benches/take_small_dict.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/small_table/tests.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index ce18a7cc8d0..1e0b0edfc5a 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -259,6 +259,10 @@ harness = false name = "take_primitive" harness = false +[[bench]] +name = "take_small_dict" +harness = false + [[bench]] name = "take_chunked" harness = false diff --git a/vortex-array/benches/take_small_dict.rs b/vortex-array/benches/take_small_dict.rs new file mode 100644 index 00000000000..a30f3f9fee0 --- /dev/null +++ b/vortex-array/benches/take_small_dict.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Small enum dictionaries with byte codes. Run the same benchmark on the baseline and candidate +//! commits to include allocation, dispatch, and dictionary execution in both measurements. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +const LENGTHS: &[usize] = &[64, 1_024, 65_536, 1_000_000]; +const CARDINALITIES: &[u8] = &[2, 4, 8, 16, 32]; + +fn run(bencher: Bencher, len: usize, skewed: bool) { + // Non-identity values prevent a copy of the codes from looking like a correct lookup. + let values = (0..CARDINALITY) + .map(|code| code.wrapping_mul(37).wrapping_add(129)) + .collect::>(); + let mut rng = StdRng::seed_from_u64(0); + let codes = (0..len) + .map(|_| { + if skewed && rng.random_ratio(9, 10) { + 0 + } else { + rng.random_range(0..CARDINALITY) + } + }) + .collect::>(); + let expected = codes + .iter() + .map(|&code| values[usize::from(code)]) + .collect::>(); + let dict = DictArray::try_new( + PrimitiveArray::from_iter(codes).into_array(), + PrimitiveArray::from_iter(values).into_array(), + ) + .unwrap() + .into_array(); + + let actual = dict + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + assert_eq!(actual.as_slice::(), expected); + + bencher + .counter(ItemsCount::new(len)) + .with_inputs(|| (dict.clone(), SESSION.create_execution_ctx())) + .bench_values(|(dict, mut ctx)| dict.execute::(&mut ctx).unwrap()); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = LENGTHS, consts = CARDINALITIES)] +fn uniform(bencher: Bencher, len: usize) { + run::(bencher, len, false); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = LENGTHS, consts = CARDINALITIES)] +fn skewed(bencher: Bencher, len: usize) { + run::(bencher, len, true); +} 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 index 3cc9c83afec..be5eae2ee83 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs @@ -1,7 +1,9 @@ // 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. +//! Register-table take for `u8` codes and at most 32 one-byte values. + +use std::sync::LazyLock; cfg_if::cfg_if! { if #[cfg(target_arch = "x86")] { @@ -13,7 +15,9 @@ cfg_if::cfg_if! { use arch::__m256i; use arch::_mm_loadu_si128; +use arch::_mm256_blendv_epi8; use arch::_mm256_broadcastsi128_si256; +use arch::_mm256_cmpgt_epi8; use arch::_mm256_loadu_si256; use arch::_mm256_or_si256; use arch::_mm256_set1_epi8; @@ -33,6 +37,19 @@ use super::super::take_values_fallback; use crate::dtype::PType; use crate::dtype::UnsignedPType; +#[path = "avx512.rs"] +mod avx512; + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; + +static HAS_AVX512_VBMI: LazyLock = LazyLock::new(|| { + is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512bw") + && is_x86_feature_detected!("avx512vbmi") +}); + // SAFETY: u8 has no padding or uninitialized bytes. unsafe impl FixedWidthTakeValue for u8 { fn take( @@ -73,7 +90,7 @@ fn take( ) -> Buffer { if I::PTYPE != PType::U8 || values.is_empty() - || values.len() > 16 + || values.len() > 32 || size_of::() != 1 || indices.len() < 64 || !*HAS_AVX2 @@ -84,24 +101,37 @@ fn take( // 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()) }; + // The existing single-table AVX2 path is already competitive for 1..=16 values. + // Use VBMI where it also avoids the second table shuffle and blend. + if values.len() > 16 && *HAS_AVX512_VBMI { + // SAFETY: All required features were detected. T is one byte with initialized bytes, + // and the table contains between 1 and 32 values. + return unsafe { avx512::take(values, indices, allocator) }; + } // 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) } + // and the table contains between 1 and 32 values. Specializing the table count keeps the + // existing single-table loop free of a per-vector branch or second shuffle. + if values.len() <= 16 { + unsafe { take_avx2::(values, indices, allocator) } + } else { + unsafe { take_avx2::(values, indices, allocator) } + } } #[target_feature(enable = "avx2")] -unsafe fn take_avx2( +unsafe fn take_avx2( values: &[T], indices: &[u8], allocator: &BufferAllocatorRef, ) -> Buffer { - let mut table = [values[0]; 16]; + let mut table = [values[0]; 32]; 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); + // SAFETY: Each half contains 16 initialized one-byte values. + let low = _mm256_broadcastsi128_si256(unsafe { _mm_loadu_si128(table.as_ptr().cast()) }); + let high = + _mm256_broadcastsi128_si256(unsafe { _mm_loadu_si128(table.as_ptr().add(16).cast()) }); let limit = _mm256_set1_epi8( - i8::try_from(values.len() - 1).vortex_expect("table contains at most 16 values"), + i8::try_from(values.len() - 1).vortex_expect("table contains at most 32 values"), ); let mut invalid = _mm256_setzero_si256(); let mut output = BufferMut::::with_capacity_in(indices.len(), allocator.clone()); @@ -113,13 +143,20 @@ unsafe fn take_avx2( // 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), + let taken = _mm256_shuffle_epi8(low, codes); + let taken = if TWO_TABLES { + // VPSHUFB uses the low four bits within each 128-bit lane. Bit four selects + // which broadcast table supplies the result; the bounds check rejects other codes. + _mm256_blendv_epi8( + taken, + _mm256_shuffle_epi8(high, codes), + _mm256_cmpgt_epi8(codes, _mm256_set1_epi8(15)), ) + } else { + taken }; + // SAFETY: the output has capacity for every index and T is one byte wide. + unsafe { _mm256_storeu_si256(output_ptr.add(offset).cast::<__m256i>(), taken) }; offset += 32; } assert_eq!( diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs b/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs new file mode 100644 index 00000000000..95a48c4aada --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! AVX-512 VBMI byte lookup across a full register, including 32-entry dictionaries. + +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; + +use super::FixedWidthTakeValue; +use super::arch::_mm512_cmpgt_epu8_mask; +use super::arch::_mm512_loadu_si512; +use super::arch::_mm512_max_epu8; +use super::arch::_mm512_permutexvar_epi8; +use super::arch::_mm512_set1_epi8; +use super::arch::_mm512_setzero_si512; +use super::arch::_mm512_storeu_si512; + +/// Look up byte codes in a register table, checking every code against the dictionary length. +/// +/// # Safety +/// +/// Requires AVX-512F, BW, and VBMI, one-byte T, and between 1 and 32 values. +#[target_feature(enable = "avx512f,avx512bw,avx512vbmi")] +pub(super) unsafe fn take( + values: &[T], + indices: &[u8], + allocator: &BufferAllocatorRef, +) -> Buffer { + let mut table = [values[0]; 64]; + table[..values.len()].copy_from_slice(values); + // SAFETY: The padded table contains 64 initialized one-byte values. + let table = unsafe { _mm512_loadu_si512(table.as_ptr().cast()) }; + let mut max_codes = _mm512_setzero_si512(); + 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 + 64 <= indices.len() { + // SAFETY: The loop condition guarantees a complete 64-byte input and output vector. + let codes = unsafe { _mm512_loadu_si512(indices.as_ptr().add(offset).cast()) }; + max_codes = _mm512_max_epu8(max_codes, codes); + let taken = _mm512_permutexvar_epi8(codes, table); + // SAFETY: Output has capacity for every index and T is one byte wide. + unsafe { _mm512_storeu_si512(output_ptr.add(offset).cast(), taken) }; + offset += 64; + } + let limit = _mm512_set1_epi8( + i8::try_from(values.len() - 1).vortex_expect("table contains at most 32 values"), + ); + // VPERMB masks indices, so explicitly reject invalid codes rather than silently wrapping. + assert_eq!( + _mm512_cmpgt_epu8_mask(max_codes, limit), + 0, + "take index out of bounds" + ); + for offset in offset..indices.len() { + spare[offset].write(values[usize::from(indices[offset])]); + } + // SAFETY: The vector loop and checked scalar tail initialized every output element. + unsafe { output.set_len(indices.len()) }; + output.freeze() +} diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs b/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs new file mode 100644 index 00000000000..50befcb32c9 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; + +use super::HAS_AVX2; +use super::HAS_AVX512_VBMI; +use super::avx512; +use super::take_avx2; + +fn lookup(values: &[u8], codes: &[u8], vbmi: bool) -> Option> { + let allocator = BufferAllocatorRef::statically_allocated(); + if vbmi { + if !*HAS_AVX512_VBMI { + return None; + } + // SAFETY: Features are detected above; tests provide 1..=32 one-byte values. + Some(unsafe { avx512::take(values, codes, &allocator) }) + } else { + if !*HAS_AVX2 { + return None; + } + // SAFETY: AVX2 is detected above; tests provide 1..=32 one-byte values. + Some(unsafe { + if values.len() <= 16 { + take_avx2::<_, false>(values, codes, &allocator) + } else { + take_avx2::<_, true>(values, codes, &allocator) + } + }) + } +} + +#[rstest] +fn register_lookup_matches_scalar( + #[values(false, true)] vbmi: bool, + #[values(1, 2, 3, 4, 8, 15, 16, 17, 31, 32)] cardinality: u8, + #[values(0, 1, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1025)] len: usize, +) { + let values = (0..cardinality) + .map(|i| i.wrapping_mul(37).wrapping_add(129)) + .collect::>(); + // Offset slices exercise unaligned loads, with all valid codes and non-identity values. + let codes = (0..cardinality) + .rev() + .cycle() + .take(len + 1) + .collect::>(); + let codes = &codes[1..]; + let expected = codes + .iter() + .map(|&i| values[usize::from(i)]) + .collect::>(); + if let Some(actual) = lookup(&values, codes, vbmi) { + assert_eq!(actual.as_slice(), expected); + } +} + +#[rstest] +fn register_lookup_rejects_invalid_codes( + #[values(false, true)] vbmi: bool, + #[values(2, 16, 17, 32)] cardinality: u8, + #[values(0, 31, 32, 63, 64, 127, 128)] offset: usize, + #[values(0, 32, 127, 128, 255)] invalid: u8, +) { + if (vbmi && !*HAS_AVX512_VBMI) || (!vbmi && !*HAS_AVX2) { + return; + } + let values = (0..cardinality).collect::>(); + let mut codes = vec![0; 129]; + codes[offset] = invalid.max(cardinality); + let result = std::panic::catch_unwind(|| lookup(&values, &codes, vbmi)); + assert!(result.is_err()); +} From 5c569ce62ac29613a0f5a1e0a327a0eda4ff6d7e Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 18:52:10 +0100 Subject: [PATCH 2/4] refactor: align byte table kernels with existing SIMD modules Signed-off-by: Joe Isaacs --- .../arrays/fixed_width/take/small_table.rs | 3 ++ .../fixed_width/take/small_table/avx2.rs | 29 +++++------- .../fixed_width/take/small_table/avx512.rs | 44 ++++++++++++++----- .../fixed_width/take/small_table/tests.rs | 10 ++--- 4 files changed, 52 insertions(+), 34 deletions(-) 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 3f60bee1a26..c563e08d1d0 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -8,6 +8,9 @@ cfg_if::cfg_if! { mod neon; } else if #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] { mod avx2; + mod avx512; + #[cfg(test)] + mod tests; } 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 index be5eae2ee83..52a981325bb 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs @@ -1,9 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Register-table take for `u8` codes and at most 32 one-byte values. - -use std::sync::LazyLock; +//! AVX2 byte-table take for `u8` codes and at most 32 one-byte values. cfg_if::cfg_if! { if #[cfg(target_arch = "x86")] { @@ -34,22 +32,11 @@ use vortex_error::VortexExpect; use super::super::FixedWidthTakeValue; use super::super::HAS_AVX2; use super::super::take_values_fallback; +use super::avx512::HAS_AVX512_VBMI; +use super::avx512::take_avx512; use crate::dtype::PType; use crate::dtype::UnsignedPType; -#[path = "avx512.rs"] -mod avx512; - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; - -static HAS_AVX512_VBMI: LazyLock = LazyLock::new(|| { - is_x86_feature_detected!("avx512f") - && is_x86_feature_detected!("avx512bw") - && is_x86_feature_detected!("avx512vbmi") -}); - // SAFETY: u8 has no padding or uninitialized bytes. unsafe impl FixedWidthTakeValue for u8 { fn take( @@ -106,7 +93,7 @@ fn take( if values.len() > 16 && *HAS_AVX512_VBMI { // SAFETY: All required features were detected. T is one byte with initialized bytes, // and the table contains between 1 and 32 values. - return unsafe { avx512::take(values, indices, allocator) }; + return unsafe { take_avx512(values, indices, allocator) }; } // SAFETY: AVX2 was detected above. Values are one byte with no uninitialized bytes, // and the table contains between 1 and 32 values. Specializing the table count keeps the @@ -118,8 +105,14 @@ fn take( } } +/// Takes one-byte values using one or two AVX2 lookup tables. +/// +/// # Safety +/// +/// Requires AVX2, one-byte T, and between 1 and 32 values. When `TWO_TABLES` is false, +/// the dictionary must contain at most 16 values. #[target_feature(enable = "avx2")] -unsafe fn take_avx2( +pub(super) unsafe fn take_avx2( values: &[T], indices: &[u8], allocator: &BufferAllocatorRef, diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs b/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs index 95a48c4aada..d48c2ebfaa5 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/avx512.rs @@ -1,21 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! AVX-512 VBMI byte lookup across a full register, including 32-entry dictionaries. +//! AVX-512 VBMI byte-table take for `u8` codes and at most 32 one-byte values. +use std::sync::LazyLock; + +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::_mm512_cmpgt_epu8_mask; +use arch::_mm512_loadu_si512; +use arch::_mm512_max_epu8; +use arch::_mm512_permutexvar_epi8; +use arch::_mm512_set1_epi8; +use arch::_mm512_setzero_si512; +use arch::_mm512_storeu_si512; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; -use super::FixedWidthTakeValue; -use super::arch::_mm512_cmpgt_epu8_mask; -use super::arch::_mm512_loadu_si512; -use super::arch::_mm512_max_epu8; -use super::arch::_mm512_permutexvar_epi8; -use super::arch::_mm512_set1_epi8; -use super::arch::_mm512_setzero_si512; -use super::arch::_mm512_storeu_si512; +use super::super::FixedWidthTakeValue; + +pub(super) static HAS_AVX512_VBMI: LazyLock = LazyLock::new(|| { + is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512bw") + && is_x86_feature_detected!("avx512vbmi") +}); /// Look up byte codes in a register table, checking every code against the dictionary length. /// @@ -23,7 +39,7 @@ use super::arch::_mm512_storeu_si512; /// /// Requires AVX-512F, BW, and VBMI, one-byte T, and between 1 and 32 values. #[target_feature(enable = "avx512f,avx512bw,avx512vbmi")] -pub(super) unsafe fn take( +pub(super) unsafe fn take_avx512( values: &[T], indices: &[u8], allocator: &BufferAllocatorRef, @@ -57,7 +73,13 @@ pub(super) unsafe fn take( "take index out of bounds" ); for offset in offset..indices.len() { - spare[offset].write(values[usize::from(indices[offset])]); + 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 checked scalar tail initialized every output element. unsafe { output.set_len(indices.len()) }; diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs b/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs index 50befcb32c9..5c4ef7a6f79 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs @@ -5,10 +5,10 @@ use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; -use super::HAS_AVX2; -use super::HAS_AVX512_VBMI; -use super::avx512; -use super::take_avx2; +use super::super::HAS_AVX2; +use super::avx2::take_avx2; +use super::avx512::HAS_AVX512_VBMI; +use super::avx512::take_avx512; fn lookup(values: &[u8], codes: &[u8], vbmi: bool) -> Option> { let allocator = BufferAllocatorRef::statically_allocated(); @@ -17,7 +17,7 @@ fn lookup(values: &[u8], codes: &[u8], vbmi: bool) -> Option> { return None; } // SAFETY: Features are detected above; tests provide 1..=32 one-byte values. - Some(unsafe { avx512::take(values, codes, &allocator) }) + Some(unsafe { take_avx512(values, codes, &allocator) }) } else { if !*HAS_AVX2 { return None; From 3c0a33e1fc18d0b9a43220eafba52d3975b3e21c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 19:09:50 +0100 Subject: [PATCH 3/4] perf: add NEON lookup for 32-value byte dictionaries Signed-off-by: Joe Isaacs --- .../arrays/fixed_width/take/small_table.rs | 2 - .../fixed_width/take/small_table/avx2.rs | 2 +- .../fixed_width/take/small_table/neon.rs | 36 ++++++--- .../fixed_width/take/small_table/tests.rs | 76 ------------------- 4 files changed, 28 insertions(+), 88 deletions(-) delete mode 100644 vortex-array/src/arrays/fixed_width/take/small_table/tests.rs 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 c563e08d1d0..d626b854fb5 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -9,8 +9,6 @@ cfg_if::cfg_if! { } else if #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] { mod avx2; mod avx512; - #[cfg(test)] - mod tests; } 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 index 52a981325bb..0b64d611c8b 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/avx2.rs @@ -112,7 +112,7 @@ fn take( /// Requires AVX2, one-byte T, and between 1 and 32 values. When `TWO_TABLES` is false, /// the dictionary must contain at most 16 values. #[target_feature(enable = "avx2")] -pub(super) unsafe fn take_avx2( +unsafe fn take_avx2( values: &[T], indices: &[u8], allocator: &BufferAllocatorRef, 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 index 46b8450ce85..a516c7d0bde 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs @@ -1,14 +1,16 @@ // 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. +//! NEON byte-table take for `u8` codes and at most 32 one-byte values. use std::arch::aarch64::uint8x16_t; +use std::arch::aarch64::uint8x16x2_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::vqtbl2q_u8; use std::arch::aarch64::vst1q_u8; use vortex_buffer::Buffer; @@ -60,7 +62,7 @@ fn take( ) -> Buffer { if I::PTYPE != PType::U8 || values.is_empty() - || values.len() > 16 + || values.len() > 32 || size_of::() != 1 || indices.len() < 64 { @@ -71,15 +73,21 @@ fn take( let indices: &[u8] = unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) }; - let mut table = [values[0]; 16]; + let mut table = [values[0]; 32]; 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) }; + // contains 32 initialized values, and the output has capacity for every index. + let (offset, max_code) = unsafe { + if values.len() <= 16 { + take_vectors::(&table, indices, output_ptr) + } else { + take_vectors::(&table, indices, output_ptr) + } + }; for offset in offset..indices.len() { let code = usize::from(indices[offset]); @@ -100,18 +108,28 @@ fn take( output.freeze() } -unsafe fn take_vectors( - table: &[T; 16], +unsafe fn take_vectors( + table: &[T; 32], indices: &[u8], output: *mut u8, ) -> (usize, u8) { - let table = unsafe { vld1q_u8(table.as_ptr().cast::()) }; + // SAFETY: The caller supplies 32 initialized one-byte values. + let low = unsafe { vld1q_u8(table.as_ptr().cast::()) }; + let high = unsafe { vld1q_u8(table.as_ptr().add(16).cast::()) }; + let table = uint8x16x2_t(low, high); let mut max_codes: uint8x16_t = unsafe { vdupq_n_u8(0) }; let mut offset = 0; while offset + 16 <= indices.len() { + // SAFETY: The loop condition guarantees 16 input codes and output slots are in bounds. 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)) }; + let taken = if TWO_TABLES { + unsafe { vqtbl2q_u8(table, codes) } + } else { + unsafe { vqtbl1q_u8(low, codes) } + }; + // SAFETY: The output has capacity for every index, as guaranteed by the caller. + unsafe { vst1q_u8(output.add(offset), taken) }; offset += 16; } (offset, unsafe { vmaxvq_u8(max_codes) }) diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs b/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs deleted file mode 100644 index 5c4ef7a6f79..00000000000 --- a/vortex-array/src/arrays/fixed_width/take/small_table/tests.rs +++ /dev/null @@ -1,76 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use rstest::rstest; -use vortex_buffer::Buffer; -use vortex_buffer::BufferAllocatorRef; - -use super::super::HAS_AVX2; -use super::avx2::take_avx2; -use super::avx512::HAS_AVX512_VBMI; -use super::avx512::take_avx512; - -fn lookup(values: &[u8], codes: &[u8], vbmi: bool) -> Option> { - let allocator = BufferAllocatorRef::statically_allocated(); - if vbmi { - if !*HAS_AVX512_VBMI { - return None; - } - // SAFETY: Features are detected above; tests provide 1..=32 one-byte values. - Some(unsafe { take_avx512(values, codes, &allocator) }) - } else { - if !*HAS_AVX2 { - return None; - } - // SAFETY: AVX2 is detected above; tests provide 1..=32 one-byte values. - Some(unsafe { - if values.len() <= 16 { - take_avx2::<_, false>(values, codes, &allocator) - } else { - take_avx2::<_, true>(values, codes, &allocator) - } - }) - } -} - -#[rstest] -fn register_lookup_matches_scalar( - #[values(false, true)] vbmi: bool, - #[values(1, 2, 3, 4, 8, 15, 16, 17, 31, 32)] cardinality: u8, - #[values(0, 1, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1025)] len: usize, -) { - let values = (0..cardinality) - .map(|i| i.wrapping_mul(37).wrapping_add(129)) - .collect::>(); - // Offset slices exercise unaligned loads, with all valid codes and non-identity values. - let codes = (0..cardinality) - .rev() - .cycle() - .take(len + 1) - .collect::>(); - let codes = &codes[1..]; - let expected = codes - .iter() - .map(|&i| values[usize::from(i)]) - .collect::>(); - if let Some(actual) = lookup(&values, codes, vbmi) { - assert_eq!(actual.as_slice(), expected); - } -} - -#[rstest] -fn register_lookup_rejects_invalid_codes( - #[values(false, true)] vbmi: bool, - #[values(2, 16, 17, 32)] cardinality: u8, - #[values(0, 31, 32, 63, 64, 127, 128)] offset: usize, - #[values(0, 32, 127, 128, 255)] invalid: u8, -) { - if (vbmi && !*HAS_AVX512_VBMI) || (!vbmi && !*HAS_AVX2) { - return; - } - let values = (0..cardinality).collect::>(); - let mut codes = vec![0; 129]; - codes[offset] = invalid.max(cardinality); - let result = std::panic::catch_unwind(|| lookup(&values, &codes, vbmi)); - assert!(result.is_err()); -} From cac5eca381a2d3b785e16ebf9efc0b8979e0ff77 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 22:32:57 +0100 Subject: [PATCH 4/4] perf: extend NEON dictionary take across fixed-width values Signed-off-by: Joe Isaacs --- benchmarks/take-small-dict/Cargo.lock | 16 + benchmarks/take-small-dict/Cargo.toml | 17 + benchmarks/take-small-dict/README.md | 55 +++ benchmarks/take-small-dict/src/main.rs | 425 ++++++++++++++++++ vortex-array/benches/take_small_dict.rs | 41 +- .../fixed_width/take/small_table/neon.rs | 83 ++-- .../take/small_table/neon/planes.rs | 104 +++++ .../take/small_table/neon/table.rs | 95 ++++ 8 files changed, 778 insertions(+), 58 deletions(-) create mode 100644 benchmarks/take-small-dict/Cargo.lock create mode 100644 benchmarks/take-small-dict/Cargo.toml create mode 100644 benchmarks/take-small-dict/README.md create mode 100644 benchmarks/take-small-dict/src/main.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/small_table/neon/planes.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/small_table/neon/table.rs diff --git a/benchmarks/take-small-dict/Cargo.lock b/benchmarks/take-small-dict/Cargo.lock new file mode 100644 index 00000000000..2dfba93a3d5 --- /dev/null +++ b/benchmarks/take-small-dict/Cargo.lock @@ -0,0 +1,16 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "fearless-dict-take" +version = "0.1.0" +dependencies = [ + "fearless_simd", +] + +[[package]] +name = "fearless_simd" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3772b63c40606beea8fe1f7b06a60de27ceb9bfc2d9dad3427368ca41dc7e62" diff --git a/benchmarks/take-small-dict/Cargo.toml b/benchmarks/take-small-dict/Cargo.toml new file mode 100644 index 00000000000..15694b1204c --- /dev/null +++ b/benchmarks/take-small-dict/Cargo.toml @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +[package] +name = "fearless-dict-take" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +fearless_simd = "=1.0.0" + +# Keep the experimental comparison dependency outside the production workspace. +[workspace] + +[profile.release] +debug = false diff --git a/benchmarks/take-small-dict/README.md b/benchmarks/take-small-dict/README.md new file mode 100644 index 00000000000..8ffc1a74cd0 --- /dev/null +++ b/benchmarks/take-small-dict/README.md @@ -0,0 +1,55 @@ + + + +# Small dictionary take on NEON + +Run on AArch64 with NEON: + +```sh +cargo run --release --manifest-path benchmarks/take-small-dict/Cargo.toml > results.jsonl +``` + +This standalone experiment compares checked scalar lookup, native NEON, and generic +`fearless_simd` 1.0.0 kernels. It does not add a dependency to the Vortex workspace. +For widths 1, 2, 4, and 8, it includes the production NEON vector loop directly. Widths +16 and 32 use an experimental native loop to measure whether extending dispatch would help. +It deliberately measures combinations outside production dispatch limits as well. + +All variants use identical preallocated outputs and deterministic non-identity data. Timed +calls include table preparation, dispatch, bounds checks, lookup, and the scalar tail; +they exclude allocation and Vortex array execution. Each case rotates variant order over +31 samples. Each sample performs `clamp(2_000_000 / output_bytes, 1, 256)` iterations. +The JSONL output records the median, minimum, and maximum nanoseconds per call. + +Cases cover 1-, 2-, 4-, 8-, 16-, and 32-byte records, dictionaries through 256 total bytes, +non-power-of-two cardinalities, uniform and 90%-zero-code distributions, and 64, 1,024, +65,536, and 1,000,000 rows. Startup checks cover unaligned inputs, vector/tail boundaries, and invalid codes; +each measured case also compares every variant with a scalar reference. + +The generic kernel uses `Simd` functions dispatched through `Level::new()`. Precise byte +swizzles return zero outside each table bank, allowing multiple banks to be combined with +OR. For 2-, 4-, and 8-byte values, both native and generic kernels transpose the dictionary +into one table per byte position and interleave lookup results back into records. The generic +version uses wide swizzles and splits its vectors for the library's 128-bit interleaved stores. +The native version uses NEON TBL and ST2/ST4. The scalar baseline corresponds to the fallback algorithm, not the existing NEON +fast path for small byte dictionaries. + +For end-to-end measurements, including allocation and dictionary execution, build and save +the Vortex benchmark executable separately for the baseline and candidate revisions: + +```sh +cargo bench -p vortex-array --bench take_small_dict --no-run + --bench --sample-count 31 --sample-size 4 +``` + +That benchmark checks decoded values before timing and covers all primitive integer and +floating-point types. Keep the expanded benchmark source identical between revisions. +Run saved binaries serially in alternating order. Report results separately from this +kernel-only experiment; their timings have different boundaries and data generators. + +Production NEON dispatch accepts u8 codes and arbitrary cardinalities up to 256, 128, 64, +and 32 for 1-, 2-, 4-, and 8-byte values respectively. The 8-byte path requires at least +1,024 output rows to amortize setup; the other widths require 64. These physical widths +cover primitive integers, floats, and narrow decimal storage. Wider records and wider +code types retain their existing fallback. The x86 extension remains limited to byte +values and 32 dictionary entries. diff --git a/benchmarks/take-small-dict/src/main.rs b/benchmarks/take-small-dict/src/main.rs new file mode 100644 index 00000000000..5a271be9b44 --- /dev/null +++ b/benchmarks/take-small-dict/src/main.rs @@ -0,0 +1,425 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compare checked, allocation-free dictionary lookup kernels on AArch64. + +use std::arch::aarch64::*; +use std::hint::black_box; +use std::time::Instant; + +use fearless_simd::Level; +use fearless_simd::prelude::*; +use fearless_simd::u8x16; +use fearless_simd::u8x32; +use fearless_simd::u8x64; +use fearless_simd::u16x8; +use fearless_simd::u16x16; +use fearless_simd::u16x32; + +#[path = "../../../vortex-array/src/arrays/fixed_width/take/small_table/neon/table.rs"] +mod neon_table; + +#[path = "../../../vortex-array/src/arrays/fixed_width/take/small_table/neon/planes.rs"] +mod neon_planes; + +type Kernel = fn(&[[u8; W]], &[u8], &mut [[u8; W]]); + +#[inline(never)] +fn scalar(values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]]) { + assert_eq!(codes.len(), output.len()); + for (out, &code) in output.iter_mut().zip(codes) { + *out = values[usize::from(code)]; + } +} + +#[inline(always)] +fn bytes(values: &[[u8; W]]) -> &[u8] { + // SAFETY: Byte arrays have no padding; this covers exactly their initialized bytes. + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len() * W) } +} + +#[inline(always)] +fn bytes_mut(values: &mut [[u8; W]]) -> &mut [u8] { + // SAFETY: Byte arrays have no padding; the exclusive slice covers their initialized bytes. + unsafe { std::slice::from_raw_parts_mut(values.as_mut_ptr().cast(), values.len() * W) } +} + +macro_rules! portable_kernel { + ($name:ident, $vector:ident, $lanes:expr) => { + #[inline(always)] + fn $name( + simd: S, + values: &[[u8; W]], + codes: &[u8], + output: &mut [[u8; W]], + ) { + assert_eq!(codes.len(), output.len()); + let mut padded = [0u8; 256]; + padded[..values.len() * W].copy_from_slice(bytes(values)); + let a = $vector::from_slice(simd, &padded[..$lanes]); + let b = $vector::from_slice(simd, &padded[$lanes..2 * $lanes]); + let c = $vector::from_slice(simd, &padded[2 * $lanes..3 * $lanes]); + let d = $vector::from_slice(simd, &padded[3 * $lanes..4 * $lanes]); + let repeat = $vector::from_fn(simd, |i| (i / W) as u8); + let lane = $vector::from_fn(simd, |i| (i % W) as u8); + let width = $vector::splat(simd, W as u8); + let mut maximum = $vector::splat(simd, 0); + let out_bytes = bytes_mut(output); + let mut offset = 0; + while offset + $lanes <= codes.len() { + let code = $vector::from_slice(simd, &codes[offset..offset + $lanes]); + maximum = maximum.max(code); + let index = if W == 1 { code } else { + code.swizzle_dyn_precise(repeat) * width + lane + }; + let mut taken = a.swizzle_dyn_precise(index); + if BANKS >= 2 { + taken |= b.swizzle_dyn_precise(index - $vector::splat(simd, $lanes as u8)); + } + if BANKS >= 3 { + taken |= c.swizzle_dyn_precise(index - $vector::splat(simd, (2 * $lanes) as u8)); + } + if BANKS >= 4 { + taken |= d.swizzle_dyn_precise(index - $vector::splat(simd, (3 * $lanes) as u8)); + } + taken.store_slice(&mut out_bytes[offset * W..offset * W + $lanes]); + offset += $lanes / W; + } + assert!(usize::from(maximum.reduce_max()) < values.len()); + for i in offset..codes.len() { + output[i] = values[usize::from(codes[i])]; + } + } + }; +} + +portable_kernel!(portable16, u8x16, 16); +portable_kernel!(portable32, u8x32, 32); +portable_kernel!(portable64, u8x64, 64); + +#[inline(never)] +fn portable(values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]]) { + if W == 2 || W == 4 || W == 8 { + return portable_planar(values, codes, output); + } + let level = Level::new(); + match values.len() * W { + 1..=16 => fearless_simd::dispatch!(level, simd => portable16::<_, W, 1>(simd, values, codes, output)), + 17..=32 => fearless_simd::dispatch!(level, simd => portable32::<_, W, 1>(simd, values, codes, output)), + 33..=64 => fearless_simd::dispatch!(level, simd => portable64::<_, W, 1>(simd, values, codes, output)), + 65..=128 => fearless_simd::dispatch!(level, simd => portable64::<_, W, 2>(simd, values, codes, output)), + 129..=192 => fearless_simd::dispatch!(level, simd => portable64::<_, W, 3>(simd, values, codes, output)), + 193..=256 => fearless_simd::dispatch!(level, simd => portable64::<_, W, 4>(simd, values, codes, output)), + _ => scalar(values, codes, output), + } +} + +#[inline(always)] +fn store_bytes16(vectors: [u8x16; 4], dest: &mut [u8]) { + u8x16::store_four_interleaved(vectors, dest); +} + +#[inline(always)] +fn store_bytes32(vectors: [u8x32; 4], dest: &mut [u8]) { + let parts = vectors.map(|v| v.split()); + store_bytes16(parts.map(|p| p.0), &mut dest[..64]); + store_bytes16(parts.map(|p| p.1), &mut dest[64..]); +} + +#[inline(always)] +fn store_bytes64(vectors: [u8x64; 4], dest: &mut [u8]) { + let parts = vectors.map(|v| v.split()); + store_bytes32(parts.map(|p| p.0), &mut dest[..128]); + store_bytes32(parts.map(|p| p.1), &mut dest[128..]); +} + +#[inline(always)] +fn store_words8(vectors: [u16x8; 4], dest: &mut [u16]) { + u16x8::store_four_interleaved(vectors, dest); +} + +#[inline(always)] +fn store_words16(vectors: [u16x16; 4], dest: &mut [u16]) { + let parts = vectors.map(|v| v.split()); + store_words8(parts.map(|p| p.0), &mut dest[..32]); + store_words8(parts.map(|p| p.1), &mut dest[32..]); +} + +#[inline(always)] +fn store_words32(vectors: [u16x32; 4], dest: &mut [u16]) { + let parts = vectors.map(|v| v.split()); + store_words16(parts.map(|p| p.0), &mut dest[..64]); + store_words16(parts.map(|p| p.1), &mut dest[64..]); +} + +macro_rules! portable_planes { + ($name:ident, $vector:ident, $words:ident, $store_bytes:ident, $store_words:ident, $lanes:expr) => { + #[inline(always)] + fn $name( + simd: S, values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]], + ) { + assert_eq!(codes.len(), output.len()); + let mut planes = [[0u8; 128]; W]; + for (i, value) in values.iter().enumerate() { + for j in 0..W { planes[j][i] = value[j]; } + } + let first = std::array::from_fn::<_, W, _>(|j| $vector::from_slice(simd, &planes[j][..$lanes])); + let second = std::array::from_fn::<_, W, _>(|j| $vector::from_slice(simd, &planes[j][$lanes..2*$lanes])); + let mut maximum = $vector::splat(simd, 0); + let mut offset = 0; + while offset + $lanes <= codes.len() { + let code = $vector::from_slice(simd, &codes[offset..offset+$lanes]); + maximum = maximum.max(code); + let taken = std::array::from_fn::<_, W, _>(|j| { + let mut v = first[j].swizzle_dyn_precise(code); + if BANKS == 2 { + v |= second[j].swizzle_dyn_precise(code - $vector::splat(simd, $lanes as u8)); + } + v + }); + let dest = &mut bytes_mut(output)[offset*W..(offset+$lanes)*W]; + if W == 2 { + taken[0].zip_low(taken[1]).store_slice(&mut dest[..$lanes]); + taken[0].zip_high(taken[1]).store_slice(&mut dest[$lanes..]); + } else if W == 4 { + $store_bytes([taken[0], taken[1], taken[2], taken[3]], dest); + } else if W == 8 { + let low: [$words; 4] = std::array::from_fn(|j| taken[j*2].zip_low(taken[j*2+1]).bitcast()); + let high: [$words; 4] = std::array::from_fn(|j| taken[j*2].zip_high(taken[j*2+1]).bitcast()); + assert_eq!(dest.as_ptr().align_offset(align_of::()), 0); + // SAFETY: Alignment is checked, every u16 bit pattern is valid, and the + // exclusive slice covers exactly the initialized output bytes. + let words = unsafe { std::slice::from_raw_parts_mut(dest.as_mut_ptr().cast::(), dest.len()/2) }; + $store_words(low, &mut words[..2*$lanes]); + $store_words(high, &mut words[2*$lanes..]); + } + offset += $lanes; + } + assert!(usize::from(maximum.reduce_max()) < values.len()); + for i in offset..codes.len() { output[i] = values[usize::from(codes[i])]; } + } + }; +} +portable_planes!(planes16, u8x16, u16x8, store_bytes16, store_words8, 16); +portable_planes!(planes32, u8x32, u16x16, store_bytes32, store_words16, 32); +portable_planes!(planes64, u8x64, u16x32, store_bytes64, store_words32, 64); + +#[inline(never)] +fn portable_planar(values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]]) { + let level = Level::new(); + match values.len() { + 1..=16 => fearless_simd::dispatch!(level, simd => planes16::<_, W, 1>(simd, values, codes, output)), + 17..=32 => fearless_simd::dispatch!(level, simd => planes32::<_, W, 1>(simd, values, codes, output)), + 33..=64 => fearless_simd::dispatch!(level, simd => planes64::<_, W, 1>(simd, values, codes, output)), + 65..=128 => fearless_simd::dispatch!(level, simd => planes64::<_, W, 2>(simd, values, codes, output)), + _ => scalar(values, codes, output), + } +} + +#[inline(always)] +unsafe fn raw_inner( + values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]], +) { + assert_eq!(codes.len(), output.len()); + let mut padded = [0u8; 256]; + padded[..values.len() * W].copy_from_slice(bytes(values)); + // SAFETY: AArch64 provides NEON; all padded loads are within initialized table storage. + unsafe { + let a = vld1q_u8_x4(padded.as_ptr()); + let b = vld1q_u8_x4(padded.as_ptr().add(64)); + let c = vld1q_u8_x4(padded.as_ptr().add(128)); + let d = vld1q_u8_x4(padded.as_ptr().add(192)); + let repeat = std::array::from_fn::<_, 16, _>(|i| (i / W) as u8); + let lane = std::array::from_fn::<_, 16, _>(|i| (i % W) as u8); + let repeat = vld1q_u8(repeat.as_ptr()); + let lane = vld1q_u8(lane.as_ptr()); + let width = vdupq_n_u8(W as u8); + let mut maximum = vdupq_n_u8(0); + let output_ptr = output.as_mut_ptr().cast::(); + let mut offset = 0; + while offset + 16 <= codes.len() { + let code = vld1q_u8(codes.as_ptr().add(offset)); + maximum = vmaxq_u8(maximum, code); + let base = if W == 1 { code } else { + vaddq_u8(vmulq_u8(vqtbl1q_u8(code, repeat), width), lane) + }; + for half in 0..W.div_ceil(16) { + let index = vaddq_u8(base, vdupq_n_u8((half * 16) as u8)); + let mut taken = match TABLE { + 16 => vqtbl1q_u8(a.0, index), + 32 => vqtbl2q_u8(uint8x16x2_t(a.0, a.1), index), + 48 => vqtbl3q_u8(uint8x16x3_t(a.0, a.1, a.2), index), + _ => vqtbl4q_u8(a, index), + }; + if TABLE > 64 { + taken = vorrq_u8(taken, vqtbl4q_u8(b, vsubq_u8(index, vdupq_n_u8(64)))); + } + if TABLE > 128 { + taken = vorrq_u8(taken, vqtbl4q_u8(c, vsubq_u8(index, vdupq_n_u8(128)))); + } + if TABLE > 192 { + taken = vorrq_u8(taken, vqtbl4q_u8(d, vsubq_u8(index, vdupq_n_u8(192)))); + } + vst1q_u8(output_ptr.add(offset * W + half * 16), taken); + } + offset += (16 / W).max(1); + } + assert!(usize::from(vmaxvq_u8(maximum)) < values.len()); + for i in offset..codes.len() { + output[i] = values[usize::from(codes[i])]; + } + } +} + +#[inline(never)] +fn raw_wide(values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]]) { + // SAFETY: The output has one initialized record per code. Each specialization bounds its + // vector loads/stores, and checks indices independently of the wrapping byte permutations. + unsafe { + match values.len() * W { + 1..=16 => raw_inner::(values, codes, output), + 17..=32 => raw_inner::(values, codes, output), + 33..=48 => raw_inner::(values, codes, output), + 49..=64 => raw_inner::(values, codes, output), + 65..=128 => raw_inner::(values, codes, output), + 129..=192 => raw_inner::(values, codes, output), + 193..=256 => raw_inner::(values, codes, output), + _ => scalar(values, codes, output), + } + } +} + +#[inline(never)] +fn raw(values: &[[u8; W]], codes: &[u8], output: &mut [[u8; W]]) { + if W > 8 { + return raw_wide(values, codes, output); + } + assert_eq!(output.len(), codes.len()); + if W == 2 || W == 4 || W == 8 { + // SAFETY: The benchmark limits the dictionary to 256 bytes and reserves output for + // every code; W is one of the supported plane counts. + let (offset, maximum) = unsafe { + neon_planes::take_vectors::(bytes(values), codes, output.as_mut_ptr().cast()) + }; + assert!(usize::from(maximum) < values.len()); + for i in offset..codes.len() { + output[i] = values[usize::from(codes[i])]; + } + return; + } + let mut table = [0u8; 256]; + let value_bytes = bytes(values); + table[..value_bytes.len()].copy_from_slice(value_bytes); + let output_ptr = output.as_mut_ptr().cast::(); + // SAFETY: W is 1, 2, 4, or 8, table is initialized, and output has one record per code. + let (offset, maximum) = unsafe { + match value_bytes.len() { + 1..=16 => neon_table::take_vectors::<[u8; W], 16>(&table, codes, output_ptr), + 17..=32 => neon_table::take_vectors::<[u8; W], 32>(&table, codes, output_ptr), + 33..=48 => neon_table::take_vectors::<[u8; W], 48>(&table, codes, output_ptr), + 49..=64 => neon_table::take_vectors::<[u8; W], 64>(&table, codes, output_ptr), + 65..=128 => neon_table::take_vectors::<[u8; W], 128>(&table, codes, output_ptr), + 129..=192 => neon_table::take_vectors::<[u8; W], 192>(&table, codes, output_ptr), + _ => neon_table::take_vectors::<[u8; W], 256>(&table, codes, output_ptr), + } + }; + assert!(usize::from(maximum) < values.len()); + for i in offset..codes.len() { + output[i] = values[usize::from(codes[i])]; + } +} + +fn dataset(count: usize, len: usize, skewed: bool) -> (Vec<[u8; W]>, Vec) { + let values = (0..count).map(|i| std::array::from_fn(|j| { + ((i * 37 + j * 53 + 129) ^ (i >> 1)) as u8 + })).collect(); + let mut seed = 0x9e3779b97f4a7c15u64; + let codes = (0..len).map(|_| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + if skewed && seed % 10 != 0 { 0 } else { (seed % count as u64) as u8 } + }).collect(); + (values, codes) +} + +fn validate() { + for count in [1, 2, 3, 7, 16, 17, 31, 32, 33, 48, 63, 64, 65, 96, 127, 128, 192, 255, 256] { + if count * W > 256 { continue; } + if count < 256 { + for position in [0, 64, 128] { + for invalid in [count as u8, 255] { + let (values, mut codes) = dataset::(count, 129, false); + codes[position] = invalid; + for kernel in [raw:: as Kernel, portable::] { + let result = std::panic::catch_unwind(|| { + let mut output = vec![[0; W]; codes.len()]; + kernel(&values, &codes, &mut output); + }); + assert!(result.is_err(), "accepted invalid code: width={W}, count={count}"); + } + } + } + } + for len in [0, 1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1025] { + let (values, mut codes) = dataset::(count, len + 1, false); + for (i, code) in codes.iter_mut().enumerate() { *code = (i % count) as u8; } + let codes = &codes[1..]; + let expected: Vec<_> = codes.iter().map(|&c| values[usize::from(c)]).collect(); + for kernel in [raw:: as Kernel, portable::] { + let mut actual = vec![[0; W]; len]; + kernel(&values, codes, &mut actual); + assert_eq!(actual, expected, "width={W} count={count} len={len}"); + } + } + } +} + +fn measure() { + let kernels: [(&str, Kernel); 3] = [("scalar", scalar::), ("neon", raw::), ("fearless", portable::)]; + for count in [2, 4, 8, 16, 17, 24, 31, 32, 33, 48, 63, 64, 65, 96, 127, 128, 129, 192, 255, 256] { + if count * W > 256 { continue; } + for len in [64, 1024, 65_536, 1_000_000] { + for skewed in [false, true] { + let (values, codes) = dataset::(count, len, skewed); + let expected: Vec<_> = codes.iter().map(|&c| values[usize::from(c)]).collect(); + let mut output = vec![[0; W]; len]; + for (_, kernel) in kernels { + kernel(&values, &codes, &mut output); + assert_eq!(output, expected); + } + let iterations = (2_000_000 / (len * W)).clamp(1, 256); + let mut samples = [Vec::new(), Vec::new(), Vec::new()]; + for sample in 0..31 { + for position in 0..3 { + let k = (sample + position) % 3; + let kernel = kernels[k].1; + let start = Instant::now(); + for _ in 0..iterations { + kernel(black_box(&values), black_box(&codes), black_box(&mut output)); + black_box(&output); + } + samples[k].push(start.elapsed().as_secs_f64() * 1e9 / iterations as f64); + } + } + for k in 0..3 { + samples[k].sort_by(f64::total_cmp); + println!("{{\"width\":{W},\"count\":{count},\"len\":{len},\"skewed\":{skewed},\"kernel\":\"{}\",\"median_ns\":{},\"min_ns\":{},\"max_ns\":{},\"iterations\":{iterations}}}", + kernels[k].0, samples[k][15], samples[k][0], samples[k][30]); + } + } + } + } +} + +fn main() { + assert!(Level::new().as_neon().is_some()); + let hook = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + validate::<1>(); validate::<2>(); validate::<4>(); + validate::<8>(); validate::<16>(); validate::<32>(); + std::panic::set_hook(hook); + eprintln!("Correctness checks passed for all eligible widths, table sizes, and vector tails."); + measure::<1>(); measure::<2>(); measure::<4>(); + measure::<8>(); measure::<16>(); measure::<32>(); +} diff --git a/vortex-array/benches/take_small_dict.rs b/vortex-array/benches/take_small_dict.rs index a30f3f9fee0..0a7af8f2759 100644 --- a/vortex-array/benches/take_small_dict.rs +++ b/vortex-array/benches/take_small_dict.rs @@ -19,6 +19,9 @@ use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::half::f16; use vortex_session::VortexSession; #[global_allocator] @@ -32,12 +35,20 @@ fn main() { } const LENGTHS: &[usize] = &[64, 1_024, 65_536, 1_000_000]; -const CARDINALITIES: &[u8] = &[2, 4, 8, 16, 32]; +const CARDINALITIES: &[usize] = &[2, 4, 8, 16, 17, 24, 32, 33, 48, 64, 65, 96, 128, 129, 192, 256]; -fn run(bencher: Bencher, len: usize, skewed: bool) { +fn run(bencher: Bencher, len: usize, skewed: bool) { // Non-identity values prevent a copy of the codes from looking like a correct lookup. let values = (0..CARDINALITY) - .map(|code| code.wrapping_mul(37).wrapping_add(129)) + .map(|code| { + let byte = (code as u8).wrapping_mul(37).wrapping_add(129); + let value = if T::PTYPE == PType::I8 { + i16::from(byte as i8) + } else { + i16::from(byte) + }; + T::from(value).unwrap() + }) .collect::>(); let mut rng = StdRng::seed_from_u64(0); let codes = (0..len) @@ -45,7 +56,7 @@ fn run(bencher: Bencher, len: usize, skewed: bool) { if skewed && rng.random_ratio(9, 10) { 0 } else { - rng.random_range(0..CARDINALITY) + rng.random_range(0..CARDINALITY) as u8 } }) .collect::>(); @@ -64,7 +75,7 @@ fn run(bencher: Bencher, len: usize, skewed: bool) { .clone() .execute::(&mut SESSION.create_execution_ctx()) .unwrap(); - assert_eq!(actual.as_slice::(), expected); + assert_eq!(actual.as_slice::(), expected); bencher .counter(ItemsCount::new(len)) @@ -73,13 +84,21 @@ fn run(bencher: Bencher, len: usize, skewed: bool) { } #[vortex_bench_support::cpu_features] -#[divan::bench(args = LENGTHS, consts = CARDINALITIES)] -fn uniform(bencher: Bencher, len: usize) { - run::(bencher, len, false); +#[divan::bench( + types = [u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64], + args = LENGTHS, + consts = CARDINALITIES +)] +fn uniform(bencher: Bencher, len: usize) { + run::(bencher, len, false); } #[vortex_bench_support::cpu_features] -#[divan::bench(args = LENGTHS, consts = CARDINALITIES)] -fn skewed(bencher: Bencher, len: usize) { - run::(bencher, len, true); +#[divan::bench( + types = [u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64], + args = LENGTHS, + consts = CARDINALITIES +)] +fn skewed(bencher: Bencher, len: usize) { + run::(bencher, len, true); } 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 index a516c7d0bde..37d42ef31ba 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs @@ -1,17 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! NEON byte-table take for `u8` codes and at most 32 one-byte values. +//! NEON register-table take for small dictionaries with byte codes. -use std::arch::aarch64::uint8x16_t; -use std::arch::aarch64::uint8x16x2_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::vqtbl2q_u8; -use std::arch::aarch64::vst1q_u8; +mod planes; +mod table; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; @@ -60,10 +53,17 @@ fn take( indices: &[I], allocator: &BufferAllocatorRef, ) -> Buffer { + // The byte-table footprint and interleaving cost determine which widths beat scalar take. + let max_values = match size_of::() { + 1 => 256, + 2 => 128, + 4 => 64, + 8 if indices.len() >= 1_024 => 32, + _ => 0, + }; if I::PTYPE != PType::U8 || values.is_empty() - || values.len() > 32 - || size_of::() != 1 + || values.len() > max_values || indices.len() < 64 { return take_values_fallback(values, indices, allocator); @@ -73,19 +73,35 @@ fn take( let indices: &[u8] = unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) }; - let mut table = [values[0]; 32]; - table[..values.len()].copy_from_slice(values); + // SAFETY: FixedWidthTakeValue guarantees initialized bytes. The dispatch above bounds the + // complete dictionary to 256 bytes and excludes zero-width values. + let value_bytes = unsafe { + std::slice::from_raw_parts(values.as_ptr().cast::(), size_of_val(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 32 initialized values, and the output has capacity for every index. + // SAFETY: AArch64 provides NEON, the table is fully initialized, and output has capacity for + // every input code. Each specialization covers the dictionary's entire byte representation. let (offset, max_code) = unsafe { - if values.len() <= 16 { - take_vectors::(&table, indices, output_ptr) - } else { - take_vectors::(&table, indices, output_ptr) + match size_of::() { + 2 => planes::take_vectors::<2>(value_bytes, indices, output_ptr), + 4 => planes::take_vectors::<4>(value_bytes, indices, output_ptr), + 8 => planes::take_vectors::<8>(value_bytes, indices, output_ptr), + _ => { + let mut table = [0u8; 256]; + table[..value_bytes.len()].copy_from_slice(value_bytes); + match value_bytes.len() { + 1..=16 => table::take_vectors::(&table, indices, output_ptr), + 17..=32 => table::take_vectors::(&table, indices, output_ptr), + 33..=48 => table::take_vectors::(&table, indices, output_ptr), + 49..=64 => table::take_vectors::(&table, indices, output_ptr), + 65..=128 => table::take_vectors::(&table, indices, output_ptr), + 129..=192 => table::take_vectors::(&table, indices, output_ptr), + _ => table::take_vectors::(&table, indices, output_ptr), + } + } } }; @@ -107,30 +123,3 @@ fn take( unsafe { output.set_len(indices.len()) }; output.freeze() } - -unsafe fn take_vectors( - table: &[T; 32], - indices: &[u8], - output: *mut u8, -) -> (usize, u8) { - // SAFETY: The caller supplies 32 initialized one-byte values. - let low = unsafe { vld1q_u8(table.as_ptr().cast::()) }; - let high = unsafe { vld1q_u8(table.as_ptr().add(16).cast::()) }; - let table = uint8x16x2_t(low, high); - let mut max_codes: uint8x16_t = unsafe { vdupq_n_u8(0) }; - let mut offset = 0; - while offset + 16 <= indices.len() { - // SAFETY: The loop condition guarantees 16 input codes and output slots are in bounds. - let codes = unsafe { vld1q_u8(indices.as_ptr().add(offset)) }; - max_codes = unsafe { vmaxq_u8(max_codes, codes) }; - let taken = if TWO_TABLES { - unsafe { vqtbl2q_u8(table, codes) } - } else { - unsafe { vqtbl1q_u8(low, codes) } - }; - // SAFETY: The output has capacity for every index, as guaranteed by the caller. - unsafe { vst1q_u8(output.add(offset), taken) }; - offset += 16; - } - (offset, unsafe { vmaxvq_u8(max_codes) }) -} diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/neon/planes.rs b/vortex-array/src/arrays/fixed_width/take/small_table/neon/planes.rs new file mode 100644 index 00000000000..a8e376ecbfc --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon/planes.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Separate byte tables avoid repeating every multi-byte lookup for each output byte. + +use std::arch::aarch64::uint16x8x4_t; +use std::arch::aarch64::uint8x16x2_t; +use std::arch::aarch64::uint8x16x4_t; +use std::arch::aarch64::vdupq_n_u8; +use std::arch::aarch64::vld1q_u8; +use std::arch::aarch64::vld1q_u8_x4; +use std::arch::aarch64::vmaxq_u8; +use std::arch::aarch64::vmaxvq_u8; +use std::arch::aarch64::vorrq_u8; +use std::arch::aarch64::vqtbl1q_u8; +use std::arch::aarch64::vqtbl2q_u8; +use std::arch::aarch64::vqtbl4q_u8; +use std::arch::aarch64::vreinterpretq_u16_u8; +use std::arch::aarch64::vst2q_u8; +use std::arch::aarch64::vst4q_u16; +use std::arch::aarch64::vst4q_u8; +use std::arch::aarch64::vsubq_u8; +use std::arch::aarch64::vzip1q_u8; +use std::arch::aarch64::vzip2q_u8; + +/// Returns the initialized record count and maximum observed code. +/// +/// # Safety +/// +/// Requires NEON, W of 2, 4, or 8, and 1..=128 complete initialized records in `values`. +/// Output must have capacity for `indices.len()` records. The caller must check the returned +/// maximum and initialize the scalar tail before exposing output. +pub(super) unsafe fn take_vectors( + values: &[u8], + indices: &[u8], + output: *mut u8, +) -> (usize, u8) { + // SAFETY: The caller bounds the dictionary to 128 records and supplies enough output space. + unsafe { + match values.len() / W { + 1..=16 => take_inner::(values, indices, output), + 17..=32 => take_inner::(values, indices, output), + 33..=64 => take_inner::(values, indices, output), + _ => take_inner::(values, indices, output), + } + } +} + +unsafe fn take_inner( + values: &[u8], + indices: &[u8], + output: *mut u8, +) -> (usize, u8) { + let mut planes = [[0u8; 128]; W]; + for (i, value) in values.chunks_exact(W).enumerate() { + for j in 0..W { + planes[j][i] = value[j]; + } + } + // SAFETY: Plane loads are within initialized arrays. Each iteration reads 16 codes and + // writes exactly 16 complete records within the caller-provided output capacity. + unsafe { + let first = std::array::from_fn::<_, W, _>(|j| vld1q_u8_x4(planes[j].as_ptr())); + let second = std::array::from_fn::<_, W, _>(|j| vld1q_u8_x4(planes[j].as_ptr().add(64))); + let mut maximum = vdupq_n_u8(0); + let mut offset = 0; + while offset + 16 <= indices.len() { + let code = vld1q_u8(indices.as_ptr().add(offset)); + maximum = vmaxq_u8(maximum, code); + let taken = std::array::from_fn::<_, W, _>(|j| { + let a = first[j]; + let mut v = match N { + 16 => vqtbl1q_u8(a.0, code), + 32 => vqtbl2q_u8(uint8x16x2_t(a.0, a.1), code), + _ => vqtbl4q_u8(a, code), + }; + if N > 64 { + v = vorrq_u8(v, vqtbl4q_u8(second[j], vsubq_u8(code, vdupq_n_u8(64)))); + } + v + }); + let ptr = output.add(offset * W); + if W == 2 { + vst2q_u8(ptr, uint8x16x2_t(taken[0], taken[1])); + } else if W == 4 { + vst4q_u8(ptr, uint8x16x4_t(taken[0], taken[1], taken[2], taken[3])); + } else if W == 8 { + let low = std::array::from_fn::<_, 4, _>(|j| { + vreinterpretq_u16_u8(vzip1q_u8(taken[j * 2], taken[j * 2 + 1])) + }); + let high = std::array::from_fn::<_, 4, _>(|j| { + vreinterpretq_u16_u8(vzip2q_u8(taken[j * 2], taken[j * 2 + 1])) + }); + vst4q_u16(ptr.cast(), uint16x8x4_t(low[0], low[1], low[2], low[3])); + vst4q_u16( + ptr.add(64).cast(), + uint16x8x4_t(high[0], high[1], high[2], high[3]), + ); + } + offset += 16; + } + (offset, vmaxvq_u8(maximum)) + } +} diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/neon/table.rs b/vortex-array/src/arrays/fixed_width/take/small_table/neon/table.rs new file mode 100644 index 00000000000..77cb9b13fb1 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon/table.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::aarch64::uint8x16x2_t; +use std::arch::aarch64::uint8x16x3_t; +use std::arch::aarch64::vaddq_u8; +use std::arch::aarch64::vdupq_n_u8; +use std::arch::aarch64::vld1q_u8; +use std::arch::aarch64::vld1q_u8_x4; +use std::arch::aarch64::vmaxq_u8; +use std::arch::aarch64::vmaxvq_u8; +use std::arch::aarch64::vmulq_u8; +use std::arch::aarch64::vorrq_u8; +use std::arch::aarch64::vqtbl1q_u8; +use std::arch::aarch64::vqtbl2q_u8; +use std::arch::aarch64::vqtbl3q_u8; +use std::arch::aarch64::vqtbl4q_u8; +use std::arch::aarch64::vst1q_u8; +use std::arch::aarch64::vsubq_u8; + +/// Returns the number of initialized records and maximum observed code. +/// +/// # Safety +/// +/// Requires NEON, a value width of 1, 2, 4, or 8 bytes, and output capacity for `indices.len()` +/// records. `TABLE_BYTES` must be 16, 32, 48, 64, 128, 192, or 256 and cover every dictionary byte. +/// The caller must check the returned maximum and handle the scalar tail before exposing output. +pub(super) unsafe fn take_vectors( + table: &[u8; 256], + indices: &[u8], + output: *mut u8, +) -> (usize, u8) { + let width = size_of::(); + let records_per_vector = 16 / width; + let repeat = std::array::from_fn::<_, 16, _>(|i| (i / width) as u8); + let byte_offsets = std::array::from_fn::<_, 16, _>(|i| (i % width) as u8); + // SAFETY: All table and mask loads are within fully initialized arrays. The loop bound + // covers four input loads and output stores; each store writes 16 / width complete records. + unsafe { + let first = vld1q_u8_x4(table.as_ptr()); + let second = vld1q_u8_x4(table.as_ptr().add(64)); + let third = vld1q_u8_x4(table.as_ptr().add(128)); + let fourth = vld1q_u8_x4(table.as_ptr().add(192)); + let repeat = vld1q_u8(repeat.as_ptr()); + let byte_offsets = vld1q_u8(byte_offsets.as_ptr()); + let widths = vdupq_n_u8(width as u8); + // Independent accumulators avoid a dependency chain through every lookup. + let mut max_codes = [vdupq_n_u8(0); 4]; + let mut offset = 0; + while offset + 3 * records_per_vector + 16 <= indices.len() { + for maximum in &mut max_codes { + let codes = vld1q_u8(indices.as_ptr().add(offset)); + *maximum = vmaxq_u8(*maximum, codes); + let byte_indices = if width == 1 { + codes + } else { + vaddq_u8(vmulq_u8(vqtbl1q_u8(codes, repeat), widths), byte_offsets) + }; + let mut taken = match TABLE_BYTES { + 16 => vqtbl1q_u8(first.0, byte_indices), + 32 => vqtbl2q_u8(uint8x16x2_t(first.0, first.1), byte_indices), + 48 => vqtbl3q_u8(uint8x16x3_t(first.0, first.1, first.2), byte_indices), + _ => vqtbl4q_u8(first, byte_indices), + }; + // TBL returns zero outside its bank. Wrapping subtraction selects exactly one + // 64-byte bank for each byte index, so the partial lookups can be ORed together. + if TABLE_BYTES > 64 { + taken = vorrq_u8( + taken, + vqtbl4q_u8(second, vsubq_u8(byte_indices, vdupq_n_u8(64))), + ); + } + if TABLE_BYTES > 128 { + taken = vorrq_u8( + taken, + vqtbl4q_u8(third, vsubq_u8(byte_indices, vdupq_n_u8(128))), + ); + } + if TABLE_BYTES > 192 { + taken = vorrq_u8( + taken, + vqtbl4q_u8(fourth, vsubq_u8(byte_indices, vdupq_n_u8(192))), + ); + } + vst1q_u8(output.add(offset * width), taken); + offset += records_per_vector; + } + } + let maximum = vmaxq_u8( + vmaxq_u8(max_codes[0], max_codes[1]), + vmaxq_u8(max_codes[2], max_codes[3]), + ); + (offset, vmaxvq_u8(maximum)) + } +}