From 1a315882f9a68ed75f9e3ced8ab1a473ee0c91be Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:13:01 +0800 Subject: [PATCH 1/3] perf(index): add fused layered ex-dot kernels --- rust/lance-index/src/vector/bq.rs | 1 + rust/lance-index/src/vector/bq/ex_dot.rs | 749 ++++++++++++++++++++++ rust/lance-index/src/vector/bq/layered.rs | 33 + 3 files changed, 783 insertions(+) create mode 100644 rust/lance-index/src/vector/bq/layered.rs diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index 4e3ef461379..cbe1f43110f 100644 --- a/rust/lance-index/src/vector/bq.rs +++ b/rust/lance-index/src/vector/bq.rs @@ -20,6 +20,7 @@ use crate::vector::quantizer::QuantizerBuildParams; pub mod builder; pub(crate) mod dist_table_quant; pub mod ex_dot; +pub mod layered; pub mod prune; pub mod rotation; pub mod storage; diff --git a/rust/lance-index/src/vector/bq/ex_dot.rs b/rust/lance-index/src/vector/bq/ex_dot.rs index 1aeb83ba40c..0e44b10fc32 100644 --- a/rust/lance-index/src/vector/bq/ex_dot.rs +++ b/rust/lance-index/src/vector/bq/ex_dot.rs @@ -360,6 +360,214 @@ fn ex_dot_scalar(ex_query: &[f32], codes: &[u8]) -> f32 { sum } +/// `sum_d query[d] * (16 * hi[d] + lo[d])` for one row stored as two 4-bit +/// blocked planes (`hi` = high nibble, `lo` = low nibble of an 8-bit ex code), +/// i.e. the 9-bit RaBitQ ex dot product over a `1 + 4 + 4` layered layout. +/// +/// Both planes use the 4-bit blocked layout of [`ex_dot_kernel`] and must have +/// the same length; the query must cover a whole number of 64-dim blocks. +pub type ExDotLayeredFn = fn(&[f32], &[u8], &[u8]) -> f32; + +/// Resolve the layered `4 + 4` dot kernel once. +pub fn ex_dot_layered_kernel() -> ExDotLayeredFn { + static KERNEL: LazyLock = LazyLock::new(select_layered_kernel); + *KERNEL +} + +/// Fused kernel for a fixed layered index width. +pub fn ex_dot_prefix_kernel(num_bits: u8) -> ExDotLayeredFn { + match num_bits { + 5 => select_prefix_kernel::<2>(), + 7 => select_prefix_kernel::<4>(), + 9 => ex_dot_layered_kernel(), + _ => unreachable!("validated layered width"), + } +} + +/// Available implementations for validating and measuring a fixed prefix layout. +#[doc(hidden)] +pub fn ex_dot_prefix_kernel_variants(num_bits: u8) -> Vec<(&'static str, ExDotLayeredFn)> { + if num_bits == 9 { + return ex_dot_layered_kernel_variants(); + } + let mut variants: Vec<(&'static str, ExDotLayeredFn)> = + vec![("dispatched", ex_dot_prefix_kernel(num_bits))]; + match num_bits { + 5 => variants.push(("scalar", prefix_scalar::<2>)), + 7 => variants.push(("scalar", prefix_scalar::<4>)), + _ => unreachable!("validated layered width"), + } + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("avx512vbmi") + && std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx2") + { + variants.push(( + "avx512-vbmi", + if num_bits == 5 { + x86::prefix_vbmi_dispatch::<2> + } else { + x86::prefix_vbmi_dispatch::<4> + }, + )); + } + if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") + { + variants.push(( + "avx2", + if num_bits == 5 { + x86::prefix_avx2_dispatch::<2> + } else { + x86::prefix_avx2_dispatch::<4> + }, + )); + } + if std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx2") + { + variants.push(( + "avx512_two_chain", + if num_bits == 5 { + x86::prefix_avx512_two_chain::<2> + } else { + x86::prefix_avx512_two_chain::<4> + }, + )); + variants.push(( + "avx512", + if num_bits == 5 { + x86::prefix_avx512_dispatch::<2> + } else { + x86::prefix_avx512_dispatch::<4> + }, + )); + } + } + variants +} + +fn select_prefix_kernel() -> ExDotLayeredFn { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("avx512vbmi") + && std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx2") + { + return x86::prefix_vbmi_dispatch::; + } + if std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx2") + { + return x86::prefix_avx512_dispatch::; + } + if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") + { + return x86::prefix_avx2_dispatch::; + } + } + #[cfg(target_arch = "aarch64")] + { + return neon::prefix_neon_dispatch::; + } + #[allow(unreachable_code)] + prefix_scalar:: +} + +fn prefix_scalar(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + let mut sum = 0.0; + for (block, query) in q.chunks_exact(64).enumerate() { + let mut high = [0; 64]; + let mut low = [0; 64]; + unpack_group(2, &lo[block * 16..(block + 1) * 16], &mut low); + if H == 2 { + unpack_group(2, &hi[block * 16..(block + 1) * 16], &mut high); + } else { + for run in 0..4 { + let mut group = [0; 64]; + let start = block * 32 + run * 8; + unpack_group(4, &hi[start..start + 8], &mut group); + high[run * 16..(run + 1) * 16].copy_from_slice(&group[..16]); + } + } + for i in 0..64 { + sum += query[i] * ((high[i] << 2) | low[i]) as f32; + } + } + sum +} + +/// Every layered kernel implementation available on this machine, for +/// benchmarking the dispatch choice; the first entry is the dispatched one. +#[doc(hidden)] +pub fn ex_dot_layered_kernel_variants() -> Vec<(&'static str, ExDotLayeredFn)> { + #[allow(unused_mut)] + let mut variants: Vec<(&'static str, ExDotLayeredFn)> = + vec![("dispatched", ex_dot_layered_kernel())]; + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("avx512f") { + variants.push(("avx512-block", x86::dot_layered_avx512_dispatch)); + variants.push(( + "avx512-scalar-unpack", + x86::dot_layered_avx512_scalar_dispatch, + )); + } + if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") + { + variants.push(("avx2-block", x86::dot_layered_avx2_dispatch)); + variants.push(("avx2-scalar-unpack", x86::dot_layered_avx2_scalar_dispatch)); + } + } + variants.push(("scalar", ex_dot_layered_scalar)); + variants +} + +fn select_layered_kernel() -> ExDotLayeredFn { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("avx512f") { + return x86::dot_layered_avx512_dispatch; + } + if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") + { + return x86::dot_layered_avx2_dispatch; + } + } + #[cfg(target_arch = "aarch64")] + { + return neon::dot_layered_neon_dispatch; + } + #[allow(unreachable_code)] + ex_dot_layered_scalar +} + +/// Bytes per 16-dim group of a 4-bit blocked plane. +const LAYERED_GROUP_BYTES: usize = 8; + +fn ex_dot_layered_scalar(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(ex_query.len() % EX_DOT_BLOCK_DIMS, 0); + debug_assert_eq!(hi.len(), lo.len()); + debug_assert!(hi.len() * 2 <= ex_query.len()); + let mut sum = 0.0f32; + let mut hi_codes = [0u8; 64]; + let mut lo_codes = [0u8; 64]; + for (group, query) in ex_query.chunks_exact(16).enumerate() { + let start = group * LAYERED_GROUP_BYTES; + if start + LAYERED_GROUP_BYTES > hi.len() { + break; + } + unpack_group(4, &hi[start..start + LAYERED_GROUP_BYTES], &mut hi_codes); + unpack_group(4, &lo[start..start + LAYERED_GROUP_BYTES], &mut lo_codes); + for d in 0..16 { + sum += query[d] * (16 * hi_codes[d] as u32 + lo_codes[d] as u32) as f32; + } + } + sum +} + #[cfg(target_arch = "x86_64")] mod x86 { use super::ExDotFn; @@ -679,6 +887,390 @@ mod x86 { x86_dot_kernel_avx512!(dot_u6_avx512, dot_u6_avx512_dispatch, unpack_u6, 6, 4); x86_dot_kernel_avx512!(dot_u7_avx512, dot_u7_avx512_dispatch, unpack_u7, 7, 4); x86_dot_kernel_avx512!(dot_u8_avx512, dot_u8_avx512_dispatch, unpack_u8x16, 8, 1); + + /// 16 8-bit codes from one 16-dim group of two 4-bit planes, combined + /// with scalar integer ops (runs on the scalar ports, overlapping the + /// vector FMA pipeline) and moved into a vector register once. + #[inline] + #[target_feature(enable = "sse2")] + unsafe fn unpack_u4x2_scalar(hi: *const u8, lo: *const u8) -> __m128i { + let hi_word = unsafe { (hi as *const u64).read_unaligned() }; + let lo_word = unsafe { (lo as *const u64).read_unaligned() }; + let mask = 0x0f0f_0f0f_0f0f_0f0fu64; + let dims_0_7 = ((hi_word & mask) << 4) | (lo_word & mask); + let dims_8_15 = (((hi_word >> 4) & mask) << 4) | ((lo_word >> 4) & mask); + _mm_set_epi64x(dims_8_15 as i64, dims_0_7 as i64) + } + + /// Four 16-dim groups (one 64-dim block, 32 bytes per plane) combined + /// with 256-bit ops, returned as four 16-byte runs in natural dim order. + /// Plane byte `8g + b` holds dim `16g + b` (low nibble) and `16g + 8 + b` + /// (high nibble); nibble shifts inside 16-bit lanes cannot cross bytes + /// after masking because every masked nibble is < 16. + #[inline] + #[target_feature(enable = "avx2")] + unsafe fn unpack_u4x2_block(hi: *const u8, lo: *const u8) -> [__m128i; 4] { + let hi_plane = unsafe { _mm256_loadu_si256(hi as *const __m256i) }; + let lo_plane = unsafe { _mm256_loadu_si256(lo as *const __m256i) }; + let low_mask = _mm256_set1_epi8(0x0f); + let high_mask = _mm256_set1_epi8(0xf0u8 as i8); + // 64-bit lane `g` of each = dims 16g..16g+8 / 16g+8..16g+16. + let dims_lo = _mm256_or_si256( + _mm256_slli_epi16::<4>(_mm256_and_si256(hi_plane, low_mask)), + _mm256_and_si256(lo_plane, low_mask), + ); + let dims_hi = _mm256_or_si256( + _mm256_and_si256(hi_plane, high_mask), + _mm256_and_si256(_mm256_srli_epi16::<4>(lo_plane), low_mask), + ); + // Interleave 64-bit lanes: even groups in `even`, odd groups in `odd`. + let even = _mm256_unpacklo_epi64(dims_lo, dims_hi); // [g0 | g2] + let odd = _mm256_unpackhi_epi64(dims_lo, dims_hi); // [g1 | g3] + [ + _mm256_castsi256_si128(even), + _mm256_castsi256_si128(odd), + _mm256_extracti128_si256::<1>(even), + _mm256_extracti128_si256::<1>(odd), + ] + } + + #[inline] + #[target_feature(enable = "avx2")] + unsafe fn unpack_prefix(hi: *const u8, lo: *const u8) -> [__m128i; 4] { + let low = unsafe { _mm_loadu_si128(lo.cast()) }; + let mask = _mm_set1_epi8(3); + let lows = [ + _mm_and_si128(low, mask), + _mm_and_si128(_mm_srli_epi16::<2>(low), mask), + _mm_and_si128(_mm_srli_epi16::<4>(low), mask), + _mm_and_si128(_mm_srli_epi16::<6>(low), mask), + ]; + let highs = if H == 2 { + let high = unsafe { _mm_loadu_si128(hi.cast()) }; + let high_mask = _mm_set1_epi8(12); + [ + _mm_and_si128(_mm_slli_epi16::<2>(high), high_mask), + _mm_and_si128(high, high_mask), + _mm_and_si128(_mm_srli_epi16::<2>(high), high_mask), + _mm_and_si128(_mm_srli_epi16::<4>(high), high_mask), + ] + } else { + let high = unsafe { _mm256_loadu_si256(hi.cast()) }; + let nibble = _mm256_set1_epi8(15); + // Position the high codes before splitting into 16-dimensional runs. + // Masks remove any cross-byte carries from the 16-bit shifts. + let a = _mm256_slli_epi16::<2>(_mm256_and_si256(high, nibble)); + let b = _mm256_and_si256(_mm256_srli_epi16::<2>(high), _mm256_set1_epi8(60)); + let even = _mm256_unpacklo_epi64(a, b); + let odd = _mm256_unpackhi_epi64(a, b); + [ + _mm256_castsi256_si128(even), + _mm256_castsi256_si128(odd), + _mm256_extracti128_si256::<1>(even), + _mm256_extracti128_si256::<1>(odd), + ] + }; + std::array::from_fn(|i| _mm_or_si128(highs[i], lows[i])) + } + #[target_feature(enable = "avx2,fma")] + unsafe fn prefix_avx2(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(hi.len(), q.len() * H as usize / 8); + debug_assert_eq!(lo.len(), q.len() / 4); + let mut acc = [_mm256_setzero_ps(); 2]; + for block in 0..q.len() / 64 { + // SAFETY: complete 64-dimensional padded blocks cover both planes. + let runs = unsafe { + unpack_prefix::( + hi.as_ptr().add(block * H as usize * 8), + lo.as_ptr().add(block * 16), + ) + }; + for (run, codes) in runs.into_iter().enumerate() { + unsafe { fma16_avx2(codes, q.as_ptr().add(block * 64 + run * 16), &mut acc) }; + } + } + unsafe { reduce_add_avx2(acc) } + } + pub(super) fn prefix_avx2_dispatch(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: dispatcher checks CPU features. + unsafe { prefix_avx2::(q, hi, lo) } + } + #[target_feature(enable = "avx512f,avx2")] + unsafe fn prefix_avx512( + q: &[f32], + hi: &[u8], + lo: &[u8], + ) -> f32 { + debug_assert_eq!(hi.len(), q.len() * H as usize / 8); + debug_assert_eq!(lo.len(), q.len() / 4); + let mut acc = [_mm512_setzero_ps(); CHAINS]; + for block in 0..q.len() / 64 { + // SAFETY: complete 64-dimensional padded blocks cover both planes. + let runs = unsafe { + unpack_prefix::( + hi.as_ptr().add(block * H as usize * 8), + lo.as_ptr().add(block * 16), + ) + }; + for (run, codes) in runs.into_iter().enumerate() { + unsafe { + fma16_avx512( + codes, + q.as_ptr().add(block * 64 + run * 16), + &mut acc[run % CHAINS], + ) + }; + } + } + let sum = if CHAINS == 2 { + _mm512_add_ps(acc[0], acc[1]) + } else { + _mm512_add_ps(_mm512_add_ps(acc[0], acc[1]), _mm512_add_ps(acc[2], acc[3])) + }; + _mm512_reduce_add_ps(sum) + } + pub(super) fn prefix_avx512_dispatch(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: dispatcher checks CPU features. + // Match the accumulation order of the equivalent native width. + if H == 2 { + unsafe { prefix_avx512::(q, hi, lo) } + } else { + unsafe { prefix_avx512::(q, hi, lo) } + } + } + pub(super) fn prefix_avx512_two_chain(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: the profiling selector checks the required CPU features. + unsafe { prefix_avx512::(q, hi, lo) } + } + #[target_feature(enable = "avx512f,avx512bw,avx512vbmi,avx2")] + unsafe fn prefix_vbmi(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(hi.len(), q.len() * H as usize / 8); + debug_assert_eq!(lo.len(), q.len() / 4); + // Each word selects one bit field from each of its eight source bytes. + const BYTE_SHIFTS: i64 = 0x3830_2820_1810_0800; + const SHIFT_TWO: i64 = 0x0202_0202_0202_0202; + let low_control = _mm512_set_epi64( + BYTE_SHIFTS + 3 * SHIFT_TWO, + BYTE_SHIFTS + 3 * SHIFT_TWO, + BYTE_SHIFTS + 2 * SHIFT_TWO, + BYTE_SHIFTS + 2 * SHIFT_TWO, + BYTE_SHIFTS + SHIFT_TWO, + BYTE_SHIFTS + SHIFT_TWO, + BYTE_SHIFTS, + BYTE_SHIFTS, + ); + let high_control = if H == 2 { + low_control + } else { + _mm512_set_epi64( + BYTE_SHIFTS + 2 * SHIFT_TWO, + BYTE_SHIFTS, + BYTE_SHIFTS + 2 * SHIFT_TWO, + BYTE_SHIFTS, + BYTE_SHIFTS + 2 * SHIFT_TWO, + BYTE_SHIFTS, + BYTE_SHIFTS + 2 * SHIFT_TWO, + BYTE_SHIFTS, + ) + }; + let duplicate_words = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); + let low_mask = _mm512_set1_epi8(3); + let high_mask = _mm512_set1_epi8(if H == 2 { 12 } else { 60 }); + let mut acc = [_mm512_setzero_ps(); 4]; + for block in 0..q.len() / 64 { + // SAFETY: complete padded blocks contain 16 low bytes and H*8 high bytes. + let low_source = _mm512_broadcast_i32x4(unsafe { + _mm_loadu_si128(lo.as_ptr().add(block * 16).cast()) + }); + let high_source = if H == 2 { + _mm512_broadcast_i32x4(unsafe { + _mm_loadu_si128(hi.as_ptr().add(block * 16).cast()) + }) + } else { + let source = _mm512_broadcast_i64x4(unsafe { + _mm256_loadu_si256(hi.as_ptr().add(block * 32).cast()) + }); + _mm512_permutexvar_epi64(duplicate_words, source) + }; + let low = _mm512_multishift_epi64_epi8(low_control, low_source); + let high = _mm512_multishift_epi64_epi8(high_control, high_source); + // (high << 2) & high_mask | (low & 3), with one ternary merge. + let packed = _mm512_ternarylogic_epi32::<0xea>( + _mm512_slli_epi16::<2>(high), + high_mask, + _mm512_and_si512(low, low_mask), + ); + let runs = [ + _mm512_castsi512_si128(packed), + _mm512_extracti32x4_epi32::<1>(packed), + _mm512_extracti32x4_epi32::<2>(packed), + _mm512_extracti32x4_epi32::<3>(packed), + ]; + for (run, codes) in runs.into_iter().enumerate() { + unsafe { + fma16_avx512( + codes, + q.as_ptr().add(block * 64 + run * 16), + &mut acc[if H == 2 { run % 2 } else { run }], + ); + } + } + } + if H == 2 { + _mm512_reduce_add_ps(_mm512_add_ps(acc[0], acc[1])) + } else { + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(acc[0], acc[1]), + _mm512_add_ps(acc[2], acc[3]), + )) + } + } + pub(super) fn prefix_vbmi_dispatch(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: the dispatch and profiling selectors check all required features. + unsafe { prefix_vbmi::(q, hi, lo) } + } + + macro_rules! x86_layered_kernel_scalar_unpack { + ($name:ident, $dispatch:ident, avx2) => { + #[target_feature(enable = "avx2", enable = "fma")] + unsafe fn $name(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(ex_query.len() % super::EX_DOT_BLOCK_DIMS, 0); + debug_assert_eq!(hi.len(), lo.len()); + let groups = (ex_query.len() / 16).min(hi.len() / super::LAYERED_GROUP_BYTES); + let mut acc = [_mm256_setzero_ps(); 2]; + for group in 0..groups { + // SAFETY: `group < groups` keeps both planes and the query in bounds. + let codes16 = unsafe { + unpack_u4x2_scalar( + hi.as_ptr().add(group * super::LAYERED_GROUP_BYTES), + lo.as_ptr().add(group * super::LAYERED_GROUP_BYTES), + ) + }; + unsafe { fma16_avx2(codes16, ex_query.as_ptr().add(group * 16), &mut acc) }; + } + unsafe { reduce_add_avx2(acc) } + } + + pub(super) fn $dispatch(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: only selected when AVX2 and FMA were detected. + unsafe { $name(ex_query, hi, lo) } + } + }; + ($name:ident, $dispatch:ident, avx512) => { + #[target_feature(enable = "avx512f")] + unsafe fn $name(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(ex_query.len() % super::EX_DOT_BLOCK_DIMS, 0); + debug_assert_eq!(hi.len(), lo.len()); + let groups = (ex_query.len() / 16).min(hi.len() / super::LAYERED_GROUP_BYTES); + let mut acc = [_mm512_setzero_ps(); 2]; + for group in 0..groups { + // SAFETY: `group < groups` keeps both planes and the query in bounds. + let codes16 = unsafe { + unpack_u4x2_scalar( + hi.as_ptr().add(group * super::LAYERED_GROUP_BYTES), + lo.as_ptr().add(group * super::LAYERED_GROUP_BYTES), + ) + }; + unsafe { + fma16_avx512( + codes16, + ex_query.as_ptr().add(group * 16), + &mut acc[group % 2], + ) + }; + } + _mm512_reduce_add_ps(_mm512_add_ps(acc[0], acc[1])) + } + + pub(super) fn $dispatch(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: only selected when AVX-512F was detected. + unsafe { $name(ex_query, hi, lo) } + } + }; + } + + x86_layered_kernel_scalar_unpack!( + dot_layered_avx2_scalar, + dot_layered_avx2_scalar_dispatch, + avx2 + ); + x86_layered_kernel_scalar_unpack!( + dot_layered_avx512_scalar, + dot_layered_avx512_scalar_dispatch, + avx512 + ); + + /// Block-unpack variants: whole 64-dim blocks per iteration (both planes + /// are padded to whole blocks by `pack_blocked_row`). + #[target_feature(enable = "avx2", enable = "fma")] + unsafe fn dot_layered_avx2(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(ex_query.len() % super::EX_DOT_BLOCK_DIMS, 0); + debug_assert_eq!(hi.len(), lo.len()); + const BLOCK_BYTES: usize = super::EX_DOT_BLOCK_DIMS / 2; + let blocks = (ex_query.len() / super::EX_DOT_BLOCK_DIMS).min(hi.len() / BLOCK_BYTES); + let mut acc = [_mm256_setzero_ps(); 2]; + for block in 0..blocks { + // SAFETY: `block < blocks` keeps both planes and the query in bounds. + let runs = unsafe { + unpack_u4x2_block( + hi.as_ptr().add(block * BLOCK_BYTES), + lo.as_ptr().add(block * BLOCK_BYTES), + ) + }; + for (run, codes16) in runs.into_iter().enumerate() { + unsafe { + fma16_avx2( + codes16, + ex_query + .as_ptr() + .add(block * super::EX_DOT_BLOCK_DIMS + run * 16), + &mut acc, + ) + }; + } + } + unsafe { reduce_add_avx2(acc) } + } + + pub(super) fn dot_layered_avx2_dispatch(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: only selected when AVX2 and FMA were detected. + unsafe { dot_layered_avx2(ex_query, hi, lo) } + } + + #[target_feature(enable = "avx512f", enable = "avx2")] + unsafe fn dot_layered_avx512(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(ex_query.len() % super::EX_DOT_BLOCK_DIMS, 0); + debug_assert_eq!(hi.len(), lo.len()); + const BLOCK_BYTES: usize = super::EX_DOT_BLOCK_DIMS / 2; + let blocks = (ex_query.len() / super::EX_DOT_BLOCK_DIMS).min(hi.len() / BLOCK_BYTES); + let mut acc = [_mm512_setzero_ps(); 2]; + for block in 0..blocks { + // SAFETY: `block < blocks` keeps both planes and the query in bounds. + let runs = unsafe { + unpack_u4x2_block( + hi.as_ptr().add(block * BLOCK_BYTES), + lo.as_ptr().add(block * BLOCK_BYTES), + ) + }; + for (run, codes16) in runs.into_iter().enumerate() { + unsafe { + fma16_avx512( + codes16, + ex_query + .as_ptr() + .add(block * super::EX_DOT_BLOCK_DIMS + run * 16), + &mut acc[run % 2], + ) + }; + } + } + _mm512_reduce_add_ps(_mm512_add_ps(acc[0], acc[1])) + } + + pub(super) fn dot_layered_avx512_dispatch(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: only selected when AVX-512F (and therefore AVX2) was detected. + unsafe { dot_layered_avx512(ex_query, hi, lo) } + } } #[cfg(target_arch = "aarch64")] @@ -894,6 +1486,82 @@ mod neon { neon_dot_kernel!(dot_u6_neon, dot_u6_neon_dispatch, unpack_u6, 6, 4); neon_dot_kernel!(dot_u7_neon, dot_u7_neon_dispatch, unpack_u7, 7, 4); neon_dot_kernel!(dot_u8_neon, dot_u8_neon_dispatch, unpack_u8x16, 8, 1); + + #[inline] + #[target_feature(enable = "neon")] + unsafe fn unpack_u4x2(hi: *const u8, lo: *const u8) -> uint8x16_t { + // Plane byte `b` holds dim `b` (low nibble) and dim `8 + b` (high nibble). + let hi_plane = unsafe { vld1_u8(hi) }; + let lo_plane = unsafe { vld1_u8(lo) }; + let low_mask = vdup_n_u8(0x0f); + let dims_0_7 = vorr_u8( + vshl_n_u8::<4>(vand_u8(hi_plane, low_mask)), + vand_u8(lo_plane, low_mask), + ); + let dims_8_15 = vorr_u8(vand_u8(hi_plane, vdup_n_u8(0xf0)), vshr_n_u8::<4>(lo_plane)); + vcombine_u8(dims_0_7, dims_8_15) + } + + #[target_feature(enable = "neon")] + unsafe fn prefix_neon(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + let mut acc = [vdupq_n_f32(0.0); 4]; + for block in 0..q.len() / 64 { + // SAFETY: each plane stores whole 64-dimensional padded blocks. + let low = unsafe { vld1q_u8(lo.as_ptr().add(block * 16)) }; + let high = if H == 2 { + unsafe { vld1q_u8(hi.as_ptr().add(block * 16)) } + } else { + vdupq_n_u8(0) + }; + for run in 0..4 { + let shift = vdupq_n_s8(-(2 * run as i8)); + let l = vandq_u8(vshlq_u8(low, shift), vdupq_n_u8(3)); + let h = if H == 2 { + vandq_u8(vshlq_u8(high, shift), vdupq_n_u8(3)) + } else { + let packed = unsafe { vld1_u8(hi.as_ptr().add(block * 32 + run * 8)) }; + vcombine_u8(vand_u8(packed, vdup_n_u8(15)), vshr_n_u8::<4>(packed)) + }; + let codes = vorrq_u8(vshlq_n_u8::<2>(h), l); + unsafe { fma16_neon(codes, q.as_ptr().add(block * 64 + run * 16), &mut acc) }; + } + } + vaddvq_f32(vaddq_f32( + vaddq_f32(acc[0], acc[1]), + vaddq_f32(acc[2], acc[3]), + )) + } + pub(super) fn prefix_neon_dispatch(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: NEON is part of the aarch64 baseline. + unsafe { prefix_neon::(q, hi, lo) } + } + + #[target_feature(enable = "neon")] + unsafe fn dot_layered_neon(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + debug_assert_eq!(ex_query.len() % super::EX_DOT_BLOCK_DIMS, 0); + debug_assert_eq!(hi.len(), lo.len()); + let groups = (ex_query.len() / 16).min(hi.len() / super::LAYERED_GROUP_BYTES); + let mut acc = [vdupq_n_f32(0.0); 4]; + for group in 0..groups { + // SAFETY: `group < groups` keeps both planes and the query in bounds. + let codes16 = unsafe { + unpack_u4x2( + hi.as_ptr().add(group * super::LAYERED_GROUP_BYTES), + lo.as_ptr().add(group * super::LAYERED_GROUP_BYTES), + ) + }; + unsafe { fma16_neon(codes16, ex_query.as_ptr().add(group * 16), &mut acc) }; + } + vaddvq_f32(vaddq_f32( + vaddq_f32(acc[0], acc[1]), + vaddq_f32(acc[2], acc[3]), + )) + } + + pub(super) fn dot_layered_neon_dispatch(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { + // SAFETY: NEON is part of the aarch64 baseline. + unsafe { dot_layered_neon(ex_query, hi, lo) } + } } #[cfg(test)] @@ -979,6 +1647,87 @@ mod tests { } } + #[rstest] + fn test_prefix_kernels( + #[values(5, 7, 9)] bits: u8, + #[values(64, 128, 1024, 1536, 2048)] dim: usize, + ) { + let mut rng = SmallRng::seed_from_u64(1234); + let layout = crate::vector::bq::layered::RQLayout::try_new(bits).unwrap(); + let mask = ((1u16 << (bits - 1)) - 1) as u8; + let values: Vec = (0..dim).map(|_| rng.random::() & mask).collect(); + let query: Vec = (0..dim) + .map(|_| rng.random_range(-16..16) as f32 / 16.) + .collect(); + let high: Vec = values.iter().map(|v| v >> layout.low_bits).collect(); + let low: Vec = values + .iter() + .map(|v| v & ((1 << layout.low_bits) - 1)) + .collect(); + let mut packed = vec![0; blocked_ex_code_bytes(dim, bits - 1)]; + let mut hi = vec![0; blocked_ex_code_bytes(dim, layout.high_bits)]; + let mut lo = vec![0; blocked_ex_code_bytes(dim, layout.low_bits)]; + pack_blocked_row(&values, bits - 1, &mut packed); + pack_blocked_row(&high, layout.high_bits, &mut hi); + pack_blocked_row(&low, layout.low_bits, &mut lo); + let expected = ex_dot_kernel(bits - 1)(&query, &packed); + for (name, kernel) in ex_dot_prefix_kernel_variants(bits) { + assert_eq!( + kernel(&query, &hi, &lo), + expected, + "{name}, bits={bits}, dim={dim}" + ); + } + // Dyadic inputs above mask changes to floating-point reduction order. + // Dispatched kernels must also match their native-width accumulation + // for ordinary queries, particularly native u4 versus fused 2+2. + let query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); + assert_eq!( + ex_dot_prefix_kernel(bits)(&query, &hi, &lo).to_bits(), + ex_dot_kernel(bits - 1)(&query, &packed).to_bits(), + "dispatched bits={bits}, dim={dim}" + ); + } + + #[rstest] + fn test_layered_kernel_matches_8bit(#[values(64, 128, 1024, 1536, 2048)] dim: usize) { + let mut rng = SmallRng::seed_from_u64(dim as u64); + let values: Vec = (0..dim).map(|_| rng.random()).collect(); + let hi_values: Vec = values.iter().map(|v| v >> 4).collect(); + let lo_values: Vec = values.iter().map(|v| v & 0x0f).collect(); + let full = kernel_codes(&values, dim, 8); + let hi = kernel_codes(&hi_values, dim, 4); + let lo = kernel_codes(&lo_values, dim, 4); + let query: Vec = (0..padded_query_len(dim)) + .map(|d| { + if d < dim { + rng.random::() - 0.5 + } else { + 0.0 + } + }) + .collect(); + let expected = ex_dot_kernel(8)(&query, &full); + let two_pass = 16.0 * ex_dot_kernel(4)(&query, &hi) + ex_dot_kernel(4)(&query, &lo); + let tolerance = expected.abs().max(1.0) * 1e-4; + assert!( + (two_pass - expected).abs() <= tolerance, + "two-pass 4+4 dot {two_pass} != 8-bit dot {expected}" + ); + let mut kernels: Vec<(&str, ExDotLayeredFn)> = vec![ + ("scalar", ex_dot_layered_scalar), + ("dispatched", ex_dot_layered_kernel()), + ]; + kernels.extend(ex_dot_layered_kernel_variants()); + for (name, kernel) in kernels { + let actual = kernel(&query, &hi, &lo); + assert!( + (actual - expected).abs() <= tolerance, + "{name} layered dot {actual} != 8-bit dot {expected} (dim={dim})" + ); + } + } + #[rstest] fn test_unpack_group_roundtrip(#[values(1, 2, 3, 4, 5, 6, 7, 8)] ex_bits: u8) { let mut rng = SmallRng::seed_from_u64(7 + ex_bits as u64); diff --git a/rust/lance-index/src/vector/bq/layered.rs b/rust/lance-index/src/vector/bq/layered.rs new file mode 100644 index 00000000000..cb5e41afba2 --- /dev/null +++ b/rust/lance-index/src/vector/bq/layered.rs @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Fixed prefix layouts for opt-in layered RaBitQ indices. + +use lance_core::{Error, Result}; + +/// Plane widths for splitting a native full-precision code. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RQLayout { + pub high_bits: u8, + pub low_bits: u8, +} + +impl RQLayout { + /// Resolve the only supported layered layouts: 1+2+2, 1+4+2 and 1+4+4. + pub fn try_new(num_bits: u8) -> Result { + let (high_bits, low_bits) = match num_bits { + 5 => (2, 2), + 7 => (4, 2), + 9 => (4, 4), + _ => { + return Err(Error::invalid_input(format!( + "IVF_RQ layered requires num_bits=5, 7 or 9, got {num_bits}" + ))); + } + }; + Ok(Self { + high_bits, + low_bits, + }) + } +} From 696929909b9d977ede7db0d0aca3406af6ee83ca Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:10:44 +0800 Subject: [PATCH 2/3] chore: retrigger gatekeeper review From 639ded030e6cc461bd0bd531f3190fb140676c01 Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:07:50 +0800 Subject: [PATCH 3/3] chore: retrigger ci for fused RaBitQ kernels