diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index 9d246e910e7..65e49bfef4b 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(crate) mod pairwise; pub mod prune; pub mod rotation; diff --git a/rust/lance-index/src/vector/bq/ex_dot.rs b/rust/lance-index/src/vector/bq/ex_dot.rs index 8b2d29aae4e..73938049aae 100644 --- a/rust/lance-index/src/vector/bq/ex_dot.rs +++ b/rust/lance-index/src/vector/bq/ex_dot.rs @@ -396,6 +396,199 @@ fn ex_dot_scalar(ex_query: &[f32], codes: &[u8]) -> f32 { sum } +/// Dot product of an `f32` query and full codes stored as high/low blocked +/// planes. The selected kernel fixes their widths and combines the codes +/// before accumulation. The query is padded to whole 64-dimensional 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 layered index with a total width from 3 through 9 bits. +pub fn ex_dot_prefix_kernel(num_bits: u8) -> ExDotLayeredFn { + match num_bits { + 3 => select_prefix_kernel::<1, 1>(), + 4 => select_prefix_kernel::<2, 1>(), + 5 => select_prefix_kernel::<2, 2>(), + 6 => select_prefix_kernel::<4, 1>(), + 7 => select_prefix_kernel::<4, 2>(), + 8 => select_prefix_kernel::<4, 3>(), + 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)> { + match num_bits { + 3 => prefix_kernel_variants::<1, 1>(), + 4 => prefix_kernel_variants::<2, 1>(), + 5 => prefix_kernel_variants::<2, 2>(), + 6 => prefix_kernel_variants::<4, 1>(), + 7 => prefix_kernel_variants::<4, 2>(), + 8 => prefix_kernel_variants::<4, 3>(), + 9 => ex_dot_layered_kernel_variants(), + _ => unreachable!("validated layered width"), + } +} + +fn prefix_kernel_variants() -> Vec<(&'static str, ExDotLayeredFn)> { + let mut variants: Vec<(&'static str, ExDotLayeredFn)> = vec![ + ("dispatched", select_prefix_kernel::()), + ("scalar", prefix_scalar::), + ]; + #[cfg(target_arch = "x86_64")] + { + if L == 2 + && 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", x86::prefix_vbmi_dispatch::)); + } + if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") + { + variants.push(("avx2", x86::prefix_avx2_dispatch::)); + } + if std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx2") + { + variants.push(("avx512_two_chain", x86::prefix_avx512_two_chain::)); + variants.push(("avx512", x86::prefix_avx512_dispatch::)); + } + } + #[cfg(target_arch = "aarch64")] + variants.push(("neon", neon::prefix_neon_dispatch::)); + variants +} + +fn select_prefix_kernel() -> ExDotLayeredFn { + #[cfg(target_arch = "x86_64")] + { + if L == 2 + && 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.as_chunks::<64>().0.iter().enumerate() { + let mut high = [0; 64]; + let mut low = [0; 64]; + for (bits, plane, values) in [(H as u8, hi, &mut high), (L as u8, lo, &mut low)] { + let block_bytes = 8 * bits as usize; + let group_dims = group_dims(bits); + let mut group = [0; 64]; + for (group_codes, values) in plane[block * block_bytes..(block + 1) * block_bytes] + .chunks_exact(group_bytes(bits)) + .zip(values.chunks_exact_mut(group_dims)) + { + unpack_group(bits, group_codes, &mut group); + values.copy_from_slice(&group[..group_dims]); + } + } + for i in 0..64 { + sum += query[i] * ((high[i] << L) | 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.as_chunks::<16>().0.iter().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; @@ -776,6 +969,431 @@ 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), + ] + } + + /// Decode one plane block into four runs in natural dimension order. + #[inline] + #[target_feature(enable = "sse2")] + unsafe fn unpack_prefix_plane(ptr: *const u8) -> [__m128i; 4] { + match BITS { + 1 => std::array::from_fn(|run| unsafe { unpack_u1(ptr.add(run * 2))[0] }), + 2 => unsafe { unpack_u2(ptr) }, + 3 => unsafe { unpack_u3(ptr) }, + 4 => std::array::from_fn(|run| unsafe { unpack_u4(ptr.add(run * 8))[0] }), + _ => unreachable!("validated prefix width"), + } + } + + #[inline] + #[target_feature(enable = "avx2")] + unsafe fn unpack_prefix( + hi: *const u8, + lo: *const u8, + ) -> [__m128i; 4] { + if L != 2 { + let high = unsafe { unpack_prefix_plane::(hi) }; + let low = unsafe { unpack_prefix_plane::(lo) }; + return std::array::from_fn(|run| { + _mm_or_si128(_mm_slli_epi16::(high[run]), low[run]) + }); + } + 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() * L as usize / 8); + 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 * L as usize * 8), + ) + }; + 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() * L as usize / 8); + 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 * L as usize * 8), + ) + }; + for (run, codes) in runs.into_iter().enumerate() { + unsafe { + fma16_avx512( + codes, + q.as_ptr().add(block * 64 + run * 16), + &mut acc[if CHAINS == 4 { + run + } else if H + L == 4 { + run % 2 + } else { + (block + run) % 2 + }], + ) + }; + } + } + 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 + L == 2 || H + L == 6 { + 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")] @@ -991,6 +1609,88 @@ 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) + } + + #[inline] + #[target_feature(enable = "neon")] + unsafe fn unpack_prefix_plane(ptr: *const u8) -> [uint8x16_t; 4] { + match BITS { + 1 => std::array::from_fn(|run| unsafe { unpack_u1(ptr.add(run * 2))[0] }), + 2 => unsafe { unpack_u2(ptr) }, + 3 => unsafe { unpack_u3(ptr) }, + 4 => std::array::from_fn(|run| unsafe { unpack_u4(ptr.add(run * 8))[0] }), + _ => unreachable!("validated prefix width"), + } + } + + #[target_feature(enable = "neon")] + unsafe fn prefix_neon(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() * L as usize / 8); + 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 high = unsafe { unpack_prefix_plane::(hi.as_ptr().add(block * H as usize * 8)) }; + let low = unsafe { unpack_prefix_plane::(lo.as_ptr().add(block * L as usize * 8)) }; + for run in 0..4 { + let codes = vorrq_u8(vshlq_n_u8::(high[run]), low[run]); + 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)] @@ -1055,6 +1755,89 @@ mod tests { } } + #[rstest] + fn test_prefix_kernels( + #[values(3, 4, 5, 6, 7, 8, 9)] bits: u8, + #[values(7, 64, 72, 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 mut 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); + query.resize(padded_query_len(dim), 0.0); + 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 mut query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); + query.resize(padded_query_len(dim), 0.0); + 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..4e4b4fb4d5f --- /dev/null +++ b/rust/lance-index/src/vector/bq/layered.rs @@ -0,0 +1,38 @@ +// 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 a sign/high/low layout for a total width from 3 through 9 bits. + /// High uses 1, 2 or 4 bits; low holds the remaining full-code bits. + pub fn try_new(num_bits: u8) -> Result { + let (high_bits, low_bits) = match num_bits { + 3 => (1, 1), + 4 => (2, 1), + 5 => (2, 2), + 6 => (4, 1), + 7 => (4, 2), + 8 => (4, 3), + 9 => (4, 4), + _ => { + return Err(Error::invalid_input(format!( + "IVF_RQ layered requires num_bits in 3..=9, got {num_bits}" + ))); + } + }; + Ok(Self { + high_bits, + low_bits, + }) + } +}