perf: speed up RaBitQ 4-bit LUT distance on ARM by 16x (lance-format#6537)

Author: justinrmiller

rust/lance-index/src/vector/bq/storage.rs

@@ -348,17 +348,8 @@ impl DistCalculator for RabitDistCalculator<'_> {
         let id = id as usize;
         let code_len = self.dim * (self.num_bits as usize) / u8::BITS as usize;
         let num_vectors = self.codes.len() / code_len;
-        let code = get_rq_code(self.codes, id, num_vectors, code_len);
-        let dist = code
-            .zip(self.dist_table.chunks_exact(SEGMENT_NUM_CODES).tuples())
-            .map(|(code_byte, (dist_table, next_dist_table))| {
-                // code is a bit vector, we iterate over 8 bits at a time,
-                // every 4 bits is a sub-vector, we need to extract the bits
-                let current_code = (code_byte & 0x0F) as usize;
-                let next_code = (code_byte >> 4) as usize;
-                dist_table[current_code] + next_dist_table[next_code]
-            })
-            .sum::<f32>();
+        let dist =
+            compute_single_rq_distance(self.codes, id, num_vectors, code_len, &self.dist_table);
 
         // distance between quantized vector and query vector
         let dist_vq_qr = (2.0 * dist - self.sum_q) / self.sqrt_d;
@@ -799,6 +790,60 @@ impl QuantizerStorage for RabitQuantizationStorage {
     }
 }
 
+/// Compute the raw distance for a single vector without allocating.
+///
+/// Fuses code extraction from the packed layout with distance accumulation
+/// in a single pass, avoiding the intermediate `Vec` allocation that
+/// `get_rq_code` + iterator would require.
+#[inline]
+fn compute_single_rq_distance(
+    codes: &[u8],
+    id: usize,
+    num_vectors: usize,
+    num_code_bytes: usize,
+    dist_table: &[f32],
+) -> f32 {
+    let remainder = num_vectors % BATCH_SIZE;
+    let mut dist_table_iter = dist_table.chunks_exact(SEGMENT_NUM_CODES).tuples();
+
+    if id < num_vectors - remainder {
+        let batch_codes = &codes[id / BATCH_SIZE * BATCH_SIZE * num_code_bytes
+            ..(id / BATCH_SIZE + 1) * BATCH_SIZE * num_code_bytes];
+
+        let id_in_batch = id % BATCH_SIZE;
+        let idx = PERM0_INVERSE[id_in_batch % 16];
+        let is_lower = id_in_batch < 16;
+
+        let mut dist = 0.0f32;
+        for block in batch_codes.chunks_exact(BATCH_SIZE) {
+            let code_byte = if is_lower {
+                (block[idx] & 0xF) | (block[idx + 16] << 4)
+            } else {
+                (block[idx] >> 4) | (block[idx + 16] & 0xF0)
+            };
+            if let Some((current_dt, next_dt)) = dist_table_iter.next() {
+                let current_code = (code_byte & 0x0F) as usize;
+                let next_code = (code_byte >> 4) as usize;
+                dist += current_dt[current_code] + next_dt[next_code];
+            }
+        }
+        dist
+    } else {
+        let offset_id = id - (num_vectors - remainder);
+        let remainder_codes = &codes[(num_vectors - remainder) * num_code_bytes..];
+
+        let mut dist = 0.0f32;
+        for &code_byte in remainder_codes.iter().skip(offset_id).step_by(remainder) {
+            if let Some((current_dt, next_dt)) = dist_table_iter.next() {
+                let current_code = (code_byte & 0x0F) as usize;
+                let next_code = (code_byte >> 4) as usize;
+                dist += current_dt[current_code] + next_dt[next_code];
+            }
+        }
+        dist
+    }
+}
+
 #[inline]
 fn get_rq_code(
     codes: &[u8],

rust/lance-linalg/Cargo.toml

@@ -62,5 +62,9 @@ harness = false
 name = "norm_l2"
 harness = false
 
+[[bench]]
+name = "dist_table"
+harness = false
+
 [lints]
 workspace = true

rust/lance-linalg/benches/dist_table.rs

@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+//! Benchmark of 4-bit LUT distance table summation (RaBitQ inner loop).
+//!
+//! Measures both the dispatched path (NEON on ARM, AVX2 on x86) and the
+//! scalar fallback, so the speedup is visible in a single benchmark run.
+
+use std::iter::repeat_with;
+
+use criterion::{Criterion, black_box, criterion_group, criterion_main};
+use lance_linalg::simd::dist_table::{BATCH_SIZE, sum_4bit_dist_table, sum_4bit_dist_table_scalar};
+use rand::Rng;
+
+fn bench_sum_4bit_dist_table(c: &mut Criterion) {
+    let mut rng = rand::rng();
+
+    // code_len = dim / 8 for 1-bit quantization
+    for (label, n_vectors, code_len) in [
+        ("32vec_dim128", 32_usize, 16_usize),
+        ("32vec_dim1536", 32, 192),
+        ("32vec_dim4096", 32, 512),
+        ("32vec_dim65536", 32, 8192),
+        ("16Kvec_dim128", 16_000, 16),
+        ("16Kvec_dim1536", 16_000, 192),
+    ] {
+        let n = n_vectors.div_ceil(BATCH_SIZE) * BATCH_SIZE;
+
+        let codes: Vec<u8> = repeat_with(|| rng.random::<u8>())
+            .take(n * code_len)
+            .collect();
+
+        let dist_table: Vec<u8> = repeat_with(|| rng.random::<u8>())
+            .take(BATCH_SIZE * code_len)
+            .collect();
+
+        let mut dists = vec![0u16; n];
+
+        // Dispatched path (NEON on ARM, AVX2 on x86)
+        c.bench_function(&format!("sum_4bit_dist_table/simd/{}", label), |b| {
+            b.iter(|| {
+                dists.fill(0);
+                sum_4bit_dist_table(n, code_len, &codes, &dist_table, &mut dists);
+                black_box(&dists);
+            })
+        });
+
+        // Scalar reference path
+        c.bench_function(&format!("sum_4bit_dist_table/scalar/{}", label), |b| {
+            b.iter(|| {
+                dists.fill(0);
+                sum_4bit_dist_table_scalar(code_len, &codes, &dist_table, &mut dists);
+                black_box(&dists);
+            })
+        });
+    }
+}
+
+criterion_group!(
+    name = benches;
+    config = Criterion::default().significance_level(0.1).sample_size(10);
+    targets = bench_sum_4bit_dist_table
+);
+criterion_main!(benches);

rust/lance-linalg/src/simd/dist_table.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
 #[cfg(target_arch = "x86_64")]
 use std::arch::x86_64::*;
 
@@ -57,13 +59,28 @@ pub fn sum_4bit_dist_table(
                 )
             }
         },
+        #[cfg(target_arch = "aarch64")]
+        SimdSupport::Neon => unsafe {
+            for i in (0..n).step_by(BATCH_SIZE) {
+                sum_dist_table_32bytes_batch_neon(
+                    &codes[i * code_len..(i + BATCH_SIZE) * code_len],
+                    dist_table,
+                    &mut dists[i..i + BATCH_SIZE],
+                )
+            }
+        },
         _ => sum_4bit_dist_table_scalar(code_len, codes, dist_table, dists),
     }
 }
 
 #[inline]
 #[allow(unused)]
-fn sum_4bit_dist_table_scalar(code_len: usize, codes: &[u8], dist_table: &[u8], dists: &mut [u16]) {
+pub fn sum_4bit_dist_table_scalar(
+    code_len: usize,
+    codes: &[u8],
+    dist_table: &[u8],
+    dists: &mut [u16],
+) {
     for (vec_block_idx, blocks) in codes.chunks_exact(BATCH_SIZE * code_len).enumerate() {
         for (sub_vec_idx, block) in blocks.chunks_exact(BATCH_SIZE).enumerate() {
             let current_dist_table = &dist_table[sub_vec_idx * 2 * 16..(sub_vec_idx * 2 + 1) * 16];
@@ -159,6 +176,77 @@ unsafe fn sum_dist_table_32bytes_batch_avx2(codes: &[u8], dist_table: &[u8], dis
     _mm256_storeu_si256(dists.as_mut_ptr().add(16) as *mut __m256i, dis1);
 }
 
+#[cfg(target_arch = "aarch64")]
+#[inline]
+unsafe fn sum_dist_table_32bytes_batch_neon(codes: &[u8], dist_table: &[u8], dists: &mut [u16]) {
+    let low_mask = vdupq_n_u8(0x0f);
+
+    // 8 accumulators: 4 per 128-bit "lane" (lo = bytes 0..16, hi = bytes 16..32 of each block)
+    let mut accu0_lo = vdupq_n_u16(0);
+    let mut accu1_lo = vdupq_n_u16(0);
+    let mut accu2_lo = vdupq_n_u16(0);
+    let mut accu3_lo = vdupq_n_u16(0);
+    let mut accu0_hi = vdupq_n_u16(0);
+    let mut accu1_hi = vdupq_n_u16(0);
+    let mut accu2_hi = vdupq_n_u16(0);
+    let mut accu3_hi = vdupq_n_u16(0);
+
+    let codes_ptr = codes.as_ptr();
+    let dt_ptr = dist_table.as_ptr();
+
+    for i in (0..codes.len()).step_by(32) {
+        // Process lo lane: bytes [i..i+16]
+        let c_lo = vld1q_u8(codes_ptr.add(i));
+        let lut_lo = vld1q_u8(dt_ptr.add(i));
+
+        let lo_lo = vandq_u8(c_lo, low_mask);
+        let hi_lo = vshrq_n_u8::<4>(c_lo);
+
+        let res_lo_lo = vqtbl1q_u8(lut_lo, lo_lo);
+        let res_hi_lo = vqtbl1q_u8(lut_lo, hi_lo);
+
+        accu0_lo = vaddq_u16(accu0_lo, vreinterpretq_u16_u8(res_lo_lo));
+        accu1_lo = vaddq_u16(accu1_lo, vshrq_n_u16::<8>(vreinterpretq_u16_u8(res_lo_lo)));
+        accu2_lo = vaddq_u16(accu2_lo, vreinterpretq_u16_u8(res_hi_lo));
+        accu3_lo = vaddq_u16(accu3_lo, vshrq_n_u16::<8>(vreinterpretq_u16_u8(res_hi_lo)));
+
+        // Process hi lane: bytes [i+16..i+32]
+        let c_hi = vld1q_u8(codes_ptr.add(i + 16));
+        let lut_hi = vld1q_u8(dt_ptr.add(i + 16));
+
+        let lo_hi = vandq_u8(c_hi, low_mask);
+        let hi_hi = vshrq_n_u8::<4>(c_hi);
+
+        let res_lo_hi = vqtbl1q_u8(lut_hi, lo_hi);
+        let res_hi_hi = vqtbl1q_u8(lut_hi, hi_hi);
+
+        accu0_hi = vaddq_u16(accu0_hi, vreinterpretq_u16_u8(res_lo_hi));
+        accu1_hi = vaddq_u16(accu1_hi, vshrq_n_u16::<8>(vreinterpretq_u16_u8(res_lo_hi)));
+        accu2_hi = vaddq_u16(accu2_hi, vreinterpretq_u16_u8(res_hi_hi));
+        accu3_hi = vaddq_u16(accu3_hi, vshrq_n_u16::<8>(vreinterpretq_u16_u8(res_hi_hi)));
+    }
+
+    // Merge: clean even bytes by subtracting the odd-byte bleed
+    accu0_lo = vsubq_u16(accu0_lo, vshlq_n_u16::<8>(accu1_lo));
+    accu0_hi = vsubq_u16(accu0_hi, vshlq_n_u16::<8>(accu1_hi));
+
+    // Cross-lane merge: add lo and hi lane accumulators
+    // This is the NEON equivalent of AVX2's permute2f128 + blend + add
+    let dis0_even = vaddq_u16(accu0_lo, accu0_hi);
+    let dis0_odd = vaddq_u16(accu1_lo, accu1_hi);
+    vst1q_u16(dists.as_mut_ptr(), dis0_even);
+    vst1q_u16(dists.as_mut_ptr().add(8), dis0_odd);
+
+    // Same for hi-nibble accumulators (vectors 16..31)
+    accu2_lo = vsubq_u16(accu2_lo, vshlq_n_u16::<8>(accu3_lo));
+    accu2_hi = vsubq_u16(accu2_hi, vshlq_n_u16::<8>(accu3_hi));
+
+    let dis1_even = vaddq_u16(accu2_lo, accu2_hi);
+    let dis1_odd = vaddq_u16(accu3_lo, accu3_hi);
+    vst1q_u16(dists.as_mut_ptr().add(16), dis1_even);
+    vst1q_u16(dists.as_mut_ptr().add(24), dis1_odd);
+}
+
 // We implement the AVX512 version in C because AVX512 is not stable yet in Rust,
 // implement it in Rust once we upgrade rust to 1.89.0.
 unsafe extern "C" {
@@ -214,4 +302,83 @@ mod tests {
         // so the distance is 2 * (dist_table[0x6] + dist_table[0xb + 16]) = 2*(7 + 12) = 38
         assert_eq!(dists[1], 38);
     }
+
+    /// Test that the SIMD path (NEON on ARM, AVX2 on x86) produces identical
+    /// results to the scalar reference across a range of dimensions, including
+    /// very large ones (up to DIM=65536).
+    ///
+    /// Note: dist_table values are capped to avoid u16 overflow, matching
+    /// production behavior where values are quantized to a small range.
+    /// (The scalar path uses saturating_add while SIMD uses wrapping add,
+    /// so they diverge on overflow — but overflow never occurs with real
+    /// quantized data.)
+    #[test]
+    fn test_simd_matches_scalar_varied_dimensions() {
+        use rand::{Rng, SeedableRng};
+        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+
+        // code_len = dim / 8 for 1-bit quantization; we test various code_lens
+        // directly since that's what the function sees.
+        // code_len=16 → DIM=128, code_len=192 → DIM=1536,
+        // code_len=512 → DIM=4096, code_len=8192 → DIM=65536
+        for code_len in [2, 16, 96, 192, 512, 1024, 8192] {
+            let n = BATCH_SIZE; // 32 vectors per batch
+
+            // Each code byte produces 2 lookups; cap values so
+            // 2 * code_len * max_val < u16::MAX.
+            let max_val = (u16::MAX as usize / (2 * code_len)).min(255) as u8;
+
+            let codes: Vec<u8> = (0..n * code_len).map(|_| rng.random::<u8>()).collect();
+            let dist_table: Vec<u8> = (0..BATCH_SIZE * code_len)
+                .map(|_| rng.random_range(0..=max_val))
+                .collect();
+
+            let mut expected = vec![0u16; n];
+            sum_4bit_dist_table_scalar(code_len, &codes, &dist_table, &mut expected);
+
+            let mut actual = vec![0u16; n];
+            sum_4bit_dist_table(n, code_len, &codes, &dist_table, &mut actual);
+
+            assert_eq!(
+                actual,
+                expected,
+                "SIMD and scalar mismatch for code_len={} (DIM={})",
+                code_len,
+                code_len * 8,
+            );
+        }
+    }
+
+    /// Test with multiple batches to verify accumulation across batch boundaries.
+    #[test]
+    fn test_simd_matches_scalar_multi_batch() {
+        use rand::{Rng, SeedableRng};
+        let mut rng = rand::rngs::StdRng::seed_from_u64(123);
+
+        for code_len in [16, 192, 1024] {
+            let n = BATCH_SIZE * 10; // 320 vectors = 10 batches
+
+            let max_val = (u16::MAX as usize / (2 * code_len)).min(255) as u8;
+
+            let codes: Vec<u8> = (0..n * code_len).map(|_| rng.random::<u8>()).collect();
+            let dist_table: Vec<u8> = (0..BATCH_SIZE * code_len)
+                .map(|_| rng.random_range(0..=max_val))
+                .collect();
+
+            let mut expected = vec![0u16; n];
+            sum_4bit_dist_table_scalar(code_len, &codes, &dist_table, &mut expected);
+
+            let mut actual = vec![0u16; n];
+            sum_4bit_dist_table(n, code_len, &codes, &dist_table, &mut actual);
+
+            assert_eq!(
+                actual,
+                expected,
+                "SIMD and scalar mismatch for multi-batch code_len={} (DIM={}, n={})",
+                code_len,
+                code_len * 8,
+                n,
+            );
+        }
+    }
 }