1use crate::lloyd;
5use crate::rhdh::{padded_dim, Rhdh};
6use std::collections::HashMap;
7
8pub struct VecqIndex {
23 pub(crate) dim: usize,
24 padded: usize,
25 pub(crate) seed: u64,
26 transform: Rhdh,
27 pub(crate) codes: Vec<u8>, pub(crate) scales: Vec<f32>, pub(crate) n: usize, keys: Vec<Option<u64>>, key_to_slot: HashMap<u64, usize>,
32 alive: Vec<bool>, live: usize, }
35
36impl VecqIndex {
37 pub fn new(dim: usize, seed: u64) -> Self {
40 let padded = padded_dim(dim);
41 Self {
42 dim,
43 padded,
44 seed,
45 transform: Rhdh::new(dim, seed),
46 codes: Vec::new(),
47 scales: Vec::new(),
48 n: 0,
49 keys: Vec::new(),
50 key_to_slot: HashMap::new(),
51 alive: Vec::new(),
52 live: 0,
53 }
54 }
55
56 pub fn len(&self) -> usize {
58 self.live
59 }
60
61 pub fn is_empty(&self) -> bool {
62 self.live == 0
63 }
64
65 pub fn slots(&self) -> usize {
68 self.n
69 }
70
71 pub fn tombstones(&self) -> usize {
73 self.n - self.live
74 }
75
76 pub fn dim(&self) -> usize {
77 self.dim
78 }
79
80 pub fn seed(&self) -> u64 {
81 self.seed
82 }
83
84 #[cfg(test)]
85 pub(crate) fn padded(&self) -> usize {
86 self.padded
87 }
88
89 pub(crate) fn padded_dim(&self) -> usize {
92 self.padded
93 }
94
95 pub(crate) fn live_slots(&self) -> usize {
96 self.live
97 }
98
99 pub(crate) fn slot_alive(&self, slot: usize) -> bool {
100 self.alive[slot]
101 }
102
103 pub(crate) fn slot_scale(&self, slot: usize) -> f32 {
104 self.scales[slot]
105 }
106
107 pub(crate) fn slot_codes(&self, slot: usize, bpv: usize) -> &[u8] {
108 &self.codes[slot * bpv..(slot + 1) * bpv]
109 }
110
111 pub(crate) fn init_dense(&mut self, count: usize) {
114 self.keys = vec![None; count];
115 self.alive = vec![true; count];
116 self.key_to_slot.clear();
117 self.live = count;
118 }
119
120 pub fn add(&mut self, v: &[f32]) -> usize {
124 self.append_slot(v, None)
125 }
126
127 pub fn add_keyed(&mut self, key: u64, v: &[f32]) -> usize {
133 if let Some(&slot) = self.key_to_slot.get(&key) {
134 let scale = self.encode_into(slot * (self.padded / 2), v);
135 self.scales[slot] = scale;
136 return slot;
137 }
138 self.append_slot(v, Some(key))
139 }
140
141 pub fn remove_keyed(&mut self, key: u64) -> bool {
145 match self.key_to_slot.remove(&key) {
146 Some(slot) => {
147 self.alive[slot] = false;
148 self.keys[slot] = None;
149 self.live -= 1;
150 true
151 }
152 None => false,
153 }
154 }
155
156 pub fn key_of(&self, slot: usize) -> Option<u64> {
159 self.keys.get(slot).copied().flatten()
160 }
161
162 pub fn contains_key(&self, key: u64) -> bool {
164 self.key_to_slot.contains_key(&key)
165 }
166
167 pub fn compact(&mut self) {
172 if self.live == self.n {
173 return;
174 }
175 let bpv = self.padded / 2;
176 let mut codes = Vec::with_capacity(self.live * bpv);
177 let mut scales = Vec::with_capacity(self.live);
178 let mut keys = Vec::with_capacity(self.live);
179 let mut alive = Vec::with_capacity(self.live);
180 self.key_to_slot.clear();
181 for slot in 0..self.n {
182 if self.alive[slot] {
183 codes.extend_from_slice(&self.codes[slot * bpv..(slot + 1) * bpv]);
184 scales.push(self.scales[slot]);
185 if let Some(key) = self.keys[slot] {
186 self.key_to_slot.insert(key, keys.len());
187 }
188 keys.push(self.keys[slot]);
189 alive.push(true);
190 }
191 }
192 self.codes = codes;
193 self.scales = scales;
194 self.keys = keys;
195 self.alive = alive;
196 self.n = self.live;
197 }
198
199 fn append_slot(&mut self, v: &[f32], key: Option<u64>) -> usize {
201 let slot = self.n;
202 let scale = self.encode_into(slot * (self.padded / 2), v);
203 self.scales.push(scale);
204 self.keys.push(key);
205 if let Some(key) = key {
206 self.key_to_slot.insert(key, slot);
207 }
208 self.alive.push(true);
209 self.n += 1;
210 self.live += 1;
211 slot
212 }
213
214 fn encode_into(&mut self, base: usize, v: &[f32]) -> f32 {
217 assert_eq!(v.len(), self.dim, "vector dim mismatch");
218 let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
219 assert!(norm > 0.0, "zero vector");
220 let unit: Vec<f32> = v.iter().map(|x| x / norm).collect();
221
222 let mut rotated = Vec::with_capacity(self.padded);
223 self.transform.apply(&unit, &mut rotated);
224
225 let bytes_per_vec = self.padded / 2;
227 if self.codes.len() < base + bytes_per_vec {
228 self.codes.resize(base + bytes_per_vec, 0);
229 }
230 let mut sum_sq = 0f32;
231 for (i, &x) in rotated.iter().enumerate() {
232 let code = lloyd::quantize_4bit(x);
233 let b = base + i / 2;
234 let byte = if i % 2 == 0 {
235 (self.codes[b] & 0xF0) | code
236 } else {
237 (self.codes[b] & 0x0F) | (code << 4)
238 };
239 self.codes[b] = byte;
240 sum_sq += lloyd::dequantize_4bit(code).powi(2);
241 }
242
243 1.0 / sum_sq.sqrt()
246 }
247
248 pub fn prepare_query(&self, q: &[f32]) -> PreparedQuery {
250 assert_eq!(q.len(), self.dim);
251 let norm: f32 = q.iter().map(|x| x * x).sum::<f32>().sqrt();
252 let unit: Vec<f32> = q.iter().map(|x| x / norm).collect();
253 let mut rotated = Vec::with_capacity(self.padded);
254 self.transform.apply(&unit, &mut rotated);
255 let rnorm: f32 = rotated.iter().map(|x| x * x).sum::<f32>().sqrt();
258 for x in rotated.iter_mut() {
259 *x /= rnorm;
260 }
261 let mut lut = [0f32; 16];
264 for (c, slot) in lut.iter_mut().enumerate() {
265 *slot = lloyd::dequantize_4bit(c as u8);
266 }
267 PreparedQuery { rotated, lut, norm }
268 }
269
270 #[inline]
281 pub fn score(&self, pq: &PreparedQuery, idx: usize) -> f32 {
282 let base = idx * (self.padded / 2);
283 let codes = &self.codes[base..base + self.padded / 2];
284 let q = &pq.rotated[..self.padded];
285 #[cfg(target_arch = "aarch64")]
286 {
287 let raw = unsafe { neon::score_neon(codes, q, &pq.lut) };
289 raw * self.scales[idx]
290 }
291 #[cfg(target_arch = "x86_64")]
292 {
293 if avx2::available() {
294 let raw = unsafe { avx2::score_avx2(codes, q, &pq.lut) };
296 raw * self.scales[idx]
297 } else {
298 score_scalar(codes, q, &pq.lut) * self.scales[idx]
299 }
300 }
301 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
302 {
303 score_scalar(codes, q, &pq.lut) * self.scales[idx]
304 }
305 }
306
307 pub fn search(&self, q: &[f32], k: usize) -> Vec<(usize, f32)> {
314 self.search_slots(q, k)
315 }
316
317 pub fn search_keyed(&self, q: &[f32], k: usize) -> Vec<(u64, f32)> {
320 self.search_slots(q, k)
321 .into_iter()
322 .filter_map(|(slot, s)| self.key_of(slot).map(|key| (key, s)))
323 .collect()
324 }
325
326 fn search_slots(&self, q: &[f32], k: usize) -> Vec<(usize, f32)> {
327 use std::cmp::Reverse;
328 use std::collections::BinaryHeap;
329
330 let pq = self.prepare_query(q);
331 let k = k.min(self.live).max(1);
332 let bpv = self.padded / 2;
333 let key = |s: f32| -> u32 {
336 let b = s.to_bits();
337 if b & 0x8000_0000 != 0 {
338 !b
339 } else {
340 b ^ 0x8000_0000
341 }
342 };
343 let mut heap: BinaryHeap<Reverse<(u32, usize)>> = BinaryHeap::with_capacity(k + 1);
344 let consider = |s: f32, idx: usize, heap: &mut BinaryHeap<Reverse<(u32, usize)>>| {
345 let ks = key(s);
346 if heap.len() < k {
347 heap.push(Reverse((ks, idx)));
348 } else if ks > heap.peek().map(|r| r.0 .0).unwrap_or(0) {
349 heap.push(Reverse((ks, idx)));
350 heap.pop();
351 }
352 };
353 #[cfg(target_arch = "aarch64")]
354 let q_rot = &pq.rotated[..self.padded];
355 #[cfg(target_arch = "x86_64")]
356 let use_avx2 = avx2::available();
357 let mut idx = 0;
358 #[cfg(target_arch = "aarch64")]
359 {
360 while idx + 4 <= self.n {
364 let codes4 = &self.codes[idx * bpv..(idx + 4) * bpv];
365 let raw = unsafe { neon::score_neon4(codes4, q_rot, &pq.lut) };
366 for (v, &r) in raw.iter().enumerate() {
367 let si = idx + v;
368 if self.alive[si] {
369 consider(r * self.scales[si], si, &mut heap);
370 }
371 }
372 idx += 4;
373 }
374 }
375 #[cfg(target_arch = "x86_64")]
376 {
377 if use_avx2 {
378 while idx + 4 <= self.n {
382 let codes4 = &self.codes[idx * bpv..(idx + 4) * bpv];
383 let raw =
385 unsafe { avx2::score_avx24(codes4, &pq.rotated[..self.padded], &pq.lut) };
386 for (v, &r) in raw.iter().enumerate() {
387 let si = idx + v;
388 if self.alive[si] {
389 consider(r * self.scales[si], si, &mut heap);
390 }
391 }
392 idx += 4;
393 }
394 }
395 }
396 while idx < self.n {
397 if self.alive[idx] {
398 consider(self.score(&pq, idx), idx, &mut heap);
399 }
400 idx += 1;
401 }
402 let key_undo = |k: u32| -> u32 {
404 if k & 0x8000_0000 != 0 {
405 k ^ 0x8000_0000 } else {
407 !k }
409 };
410 let mut out: Vec<(usize, f32)> = heap
411 .into_iter()
412 .map(|r| (r.0 .1, f32::from_bits(key_undo(r.0 .0))))
413 .collect();
414 out.sort_by(|a, b| b.1.partial_cmp(&a.1).expect("no NaN scores"));
415 out
416 }
417}
418
419#[cfg_attr(target_arch = "aarch64", cfg(test))]
422pub(crate) fn score_scalar(codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 {
423 let nb = codes.len();
424 let mut acc = [0f32; 8];
425 let mut i = 0;
426 while i + 8 <= nb {
427 for j in 0..8 {
428 let b = codes[i + j];
429 let c = (i + j) * 2;
430 acc[j] += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
431 }
432 i += 8;
433 }
434 let mut tail = 0f32;
435 while i < nb {
436 let b = codes[i];
437 let c = i * 2;
438 tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
439 i += 1;
440 }
441 let s01 = acc[0] + acc[1];
442 let s23 = acc[2] + acc[3];
443 let s45 = acc[4] + acc[5];
444 let s67 = acc[6] + acc[7];
445 (s01 + s23) + (s45 + s67) + tail
446}
447
448#[cfg(target_arch = "aarch64")]
449mod neon {
450 use std::arch::aarch64::*;
463
464 #[inline]
470 unsafe fn gather16(tbl: uint8x16x4_t, nibbles: uint8x16_t) -> [float32x4_t; 4] {
471 let idx = vmulq_u8(nibbles, vdupq_n_u8(4)); let one = vdupq_n_u8(1);
473 let two = vdupq_n_u8(2);
474 let b0 = vqtbl4q_u8(tbl, idx);
475 let b1 = vqtbl4q_u8(tbl, vaddq_u8(idx, one));
476 let b2 = vqtbl4q_u8(tbl, vaddq_u8(idx, two));
477 let b3 = vqtbl4q_u8(tbl, vaddq_u8(idx, vdupq_n_u8(3)));
478 let z01 = vzip1q_u8(b0, b1); let z23 = vzip1q_u8(b2, b3); let z01b = vzip2q_u8(b0, b1);
482 let z23b = vzip2q_u8(b2, b3);
483 let lo16 = vreinterpretq_u16_u8(z01);
484 let hi16 = vreinterpretq_u16_u8(z23);
485 let lo16b = vreinterpretq_u16_u8(z01b);
486 let hi16b = vreinterpretq_u16_u8(z23b);
487 [
488 vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip1q_u16(lo16, hi16))),
489 vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip2q_u16(lo16, hi16))),
490 vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip1q_u16(lo16b, hi16b))),
491 vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip2q_u16(lo16b, hi16b))),
492 ]
493 }
494
495 #[inline]
497 pub unsafe fn score_neon(codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 {
498 let mut bytes = [0u8; 64];
500 for (c, &v) in lut.iter().enumerate() {
501 bytes[c * 4..c * 4 + 4].copy_from_slice(&v.to_le_bytes());
502 }
503 let tbl = uint8x16x4_t(
504 vld1q_u8(bytes[0..16].as_ptr()),
505 vld1q_u8(bytes[16..32].as_ptr()),
506 vld1q_u8(bytes[32..48].as_ptr()),
507 vld1q_u8(bytes[48..64].as_ptr()),
508 );
509 let mut acc_lo = vdupq_n_f32(0.0);
511 let mut acc_hi = vdupq_n_f32(0.0);
512 let nb = codes.len();
513 let mut i = 0;
514 while i + 8 <= nb {
515 let b8 = vld1_u8(codes.as_ptr().add(i)); let lo = vand_u8(b8, vdup_n_u8(0x0F));
519 let hi = vshr_n_u8(b8, 4);
520 let nibbles = vcombine_u8(lo, hi); let g = gather16(tbl, nibbles);
522 let q0 = vld1q_f32(q.as_ptr().add(i * 2));
526 let q1 = vld1q_f32(q.as_ptr().add(i * 2 + 4));
527 let q2 = vld1q_f32(q.as_ptr().add(i * 2 + 8));
528 let q3 = vld1q_f32(q.as_ptr().add(i * 2 + 12));
529 let q_even_lo = vuzp1q_f32(q0, q1); let q_even_hi = vuzp1q_f32(q2, q3);
531 let q_odd_lo = vuzp2q_f32(q0, q1);
532 let q_odd_hi = vuzp2q_f32(q2, q3);
533 let t_lo = vaddq_f32(vmulq_f32(q_even_lo, g[0]), vmulq_f32(q_odd_lo, g[2]));
535 let t_hi = vaddq_f32(vmulq_f32(q_even_hi, g[1]), vmulq_f32(q_odd_hi, g[3]));
536 acc_lo = vaddq_f32(acc_lo, t_lo);
537 acc_hi = vaddq_f32(acc_hi, t_hi);
538 i += 8;
539 }
540 let mut acc = [0f32; 8];
542 acc[0] = vgetq_lane_f32(acc_lo, 0);
543 acc[1] = vgetq_lane_f32(acc_lo, 1);
544 acc[2] = vgetq_lane_f32(acc_lo, 2);
545 acc[3] = vgetq_lane_f32(acc_lo, 3);
546 acc[4] = vgetq_lane_f32(acc_hi, 0);
547 acc[5] = vgetq_lane_f32(acc_hi, 1);
548 acc[6] = vgetq_lane_f32(acc_hi, 2);
549 acc[7] = vgetq_lane_f32(acc_hi, 3);
550 let mut tail = 0f32;
553 while i < nb {
554 let b = codes[i];
555 let c = i * 2;
556 tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
557 i += 1;
558 }
559 let s01 = acc[0] + acc[1];
560 let s23 = acc[2] + acc[3];
561 let s45 = acc[4] + acc[5];
562 let s67 = acc[6] + acc[7];
563 (s01 + s23) + (s45 + s67) + tail
564 }
565
566 #[inline]
572 pub unsafe fn score_neon4(codes4: &[u8], q: &[f32], lut: &[f32; 16]) -> [f32; 4] {
573 let mut bytes = [0u8; 64];
574 for (c, &v) in lut.iter().enumerate() {
575 bytes[c * 4..c * 4 + 4].copy_from_slice(&v.to_le_bytes());
576 }
577 let tbl = uint8x16x4_t(
578 vld1q_u8(bytes[0..16].as_ptr()),
579 vld1q_u8(bytes[16..32].as_ptr()),
580 vld1q_u8(bytes[32..48].as_ptr()),
581 vld1q_u8(bytes[48..64].as_ptr()),
582 );
583 let nb = codes4.len() / 4; let mut acc_lo = [vdupq_n_f32(0.0); 4];
585 let mut acc_hi = [vdupq_n_f32(0.0); 4];
586 let mut i = 0;
587 while i + 8 <= nb {
588 let q0 = vld1q_f32(q.as_ptr().add(i * 2));
590 let q1 = vld1q_f32(q.as_ptr().add(i * 2 + 4));
591 let q2 = vld1q_f32(q.as_ptr().add(i * 2 + 8));
592 let q3 = vld1q_f32(q.as_ptr().add(i * 2 + 12));
593 let q_even_lo = vuzp1q_f32(q0, q1);
594 let q_even_hi = vuzp1q_f32(q2, q3);
595 let q_odd_lo = vuzp2q_f32(q0, q1);
596 let q_odd_hi = vuzp2q_f32(q2, q3);
597 for v in 0..4 {
598 let b8 = vld1_u8(codes4.as_ptr().add(v * nb + i));
599 let lo = vand_u8(b8, vdup_n_u8(0x0F));
600 let hi = vshr_n_u8(b8, 4);
601 let nibbles = vcombine_u8(lo, hi);
602 let g = gather16(tbl, nibbles);
603 let t_lo = vaddq_f32(vmulq_f32(q_even_lo, g[0]), vmulq_f32(q_odd_lo, g[2]));
604 let t_hi = vaddq_f32(vmulq_f32(q_even_hi, g[1]), vmulq_f32(q_odd_hi, g[3]));
605 acc_lo[v] = vaddq_f32(acc_lo[v], t_lo);
606 acc_hi[v] = vaddq_f32(acc_hi[v], t_hi);
607 }
608 i += 8;
609 }
610 let mut out = [0f32; 4];
611 for v in 0..4 {
612 let mut a = [0f32; 8];
614 a[0] = vgetq_lane_f32(acc_lo[v], 0);
615 a[1] = vgetq_lane_f32(acc_lo[v], 1);
616 a[2] = vgetq_lane_f32(acc_lo[v], 2);
617 a[3] = vgetq_lane_f32(acc_lo[v], 3);
618 a[4] = vgetq_lane_f32(acc_hi[v], 0);
619 a[5] = vgetq_lane_f32(acc_hi[v], 1);
620 a[6] = vgetq_lane_f32(acc_hi[v], 2);
621 a[7] = vgetq_lane_f32(acc_hi[v], 3);
622 let mut tail = 0f32;
624 let mut j = i;
625 while j < nb {
626 let b = codes4[v * nb + j];
627 let c = j * 2;
628 tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
629 j += 1;
630 }
631 let s01 = a[0] + a[1];
632 let s23 = a[2] + a[3];
633 let s45 = a[4] + a[5];
634 let s67 = a[6] + a[7];
635 out[v] = (s01 + s23) + (s45 + s67) + tail;
636 }
637 out
638 }
639}
640
641#[cfg(target_arch = "x86_64")]
655mod avx2 {
656 use std::arch::x86_64::*;
657
658 pub fn available() -> bool {
660 std::is_x86_feature_detected!("avx2")
661 }
662
663 #[inline]
665 unsafe fn gather8(lut: &[f32; 16], nibbles: __m128i) -> __m256 {
666 let idx = _mm256_cvtepu8_epi32(nibbles);
669 _mm256_i32gather_ps(lut.as_ptr(), idx, 4)
670 }
671
672 #[inline]
675 unsafe fn deinterleave16(q: *const f32) -> (__m256, __m256) {
676 let qa = _mm256_loadu_ps(q);
677 let qb = _mm256_loadu_ps(q.add(8));
678 let fixup = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
682 let even = _mm256_permutevar8x32_ps(_mm256_shuffle_ps(qa, qb, 0x88), fixup);
683 let odd = _mm256_permutevar8x32_ps(_mm256_shuffle_ps(qa, qb, 0xDD), fixup);
684 (even, odd)
685 }
686
687 #[inline]
689 #[target_feature(enable = "avx2")]
690 pub unsafe fn score_avx2(codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 {
691 let mut acc = _mm256_setzero_ps();
692 let nb = codes.len();
693 let mut i = 0;
694 while i + 8 <= nb {
695 let b8 = _mm_loadl_epi64(codes.as_ptr().add(i) as *const __m128i);
696 let g_lo = gather8(lut, _mm_and_si128(b8, _mm_set1_epi8(0x0F)));
697 let g_hi = gather8(
698 lut,
699 _mm_and_si128(_mm_srli_epi16(b8, 4), _mm_set1_epi8(0x0F)),
700 );
701 let (even, odd) = deinterleave16(q.as_ptr().add(i * 2));
702 let term = _mm256_add_ps(_mm256_mul_ps(even, g_lo), _mm256_mul_ps(odd, g_hi));
704 acc = _mm256_add_ps(acc, term);
705 i += 8;
706 }
707 let mut a = [0f32; 8];
709 _mm256_storeu_ps(a.as_mut_ptr(), acc);
710 let mut tail = 0f32;
713 while i < nb {
714 let b = codes[i];
715 let c = i * 2;
716 tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
717 i += 1;
718 }
719 let s01 = a[0] + a[1];
720 let s23 = a[2] + a[3];
721 let s45 = a[4] + a[5];
722 let s67 = a[6] + a[7];
723 (s01 + s23) + (s45 + s67) + tail
724 }
725
726 #[inline]
731 #[target_feature(enable = "avx2")]
732 pub unsafe fn score_avx24(codes4: &[u8], q: &[f32], lut: &[f32; 16]) -> [f32; 4] {
733 let nb = codes4.len() / 4; let mut acc = [_mm256_setzero_ps(); 4];
735 let mut i = 0;
736 while i + 8 <= nb {
737 let (even, odd) = deinterleave16(q.as_ptr().add(i * 2));
739 for (v, acc_v) in acc.iter_mut().enumerate() {
740 let b8 = _mm_loadl_epi64(codes4.as_ptr().add(v * nb + i) as *const __m128i);
741 let g_lo = gather8(lut, _mm_and_si128(b8, _mm_set1_epi8(0x0F)));
742 let g_hi = gather8(
743 lut,
744 _mm_and_si128(_mm_srli_epi16(b8, 4), _mm_set1_epi8(0x0F)),
745 );
746 let term = _mm256_add_ps(_mm256_mul_ps(even, g_lo), _mm256_mul_ps(odd, g_hi));
747 *acc_v = _mm256_add_ps(*acc_v, term);
748 }
749 i += 8;
750 }
751 let mut out = [0f32; 4];
752 for v in 0..4 {
753 let mut a = [0f32; 8];
754 _mm256_storeu_ps(a.as_mut_ptr(), acc[v]);
755 let mut tail = 0f32;
757 let mut j = i;
758 while j < nb {
759 let b = codes4[v * nb + j];
760 let c = j * 2;
761 tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
762 j += 1;
763 }
764 let s01 = a[0] + a[1];
765 let s23 = a[2] + a[3];
766 let s45 = a[4] + a[5];
767 let s67 = a[6] + a[7];
768 out[v] = (s01 + s23) + (s45 + s67) + tail;
769 }
770 out
771 }
772}
773
774pub struct PreparedQuery {
776 rotated: Vec<f32>,
777 lut: [f32; 16],
778 #[allow(dead_code)]
779 norm: f32,
780}
781
782pub fn cosine_f32(a: &[f32], b: &[f32]) -> f32 {
784 let mut dot = 0f32;
785 let mut na = 0f32;
786 let mut nb = 0f32;
787 for i in 0..a.len() {
788 dot += a[i] * b[i];
789 na += a[i] * a[i];
790 nb += b[i] * b[i];
791 }
792 dot / (na.sqrt() * nb.sqrt())
793}
794
795#[cfg(test)]
796mod tests {
797 use super::*;
798
799 fn rand_unit(dim: usize, seed: u64) -> Vec<f32> {
800 let mut x = seed | 1;
802 let mut v = Vec::with_capacity(dim);
803 for _ in 0..dim {
804 x ^= x << 13;
805 x ^= x >> 7;
806 x ^= x << 17;
807 let u1 = ((x >> 11) as f64 / (1u64 << 53) as f64).max(1e-12);
808 x ^= x << 16;
809 let u2 = (x >> 11) as f64 / (1u64 << 53) as f64;
810 v.push(((-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()) as f32);
811 }
812 let norm: f32 = v.iter().map(|a| a * a).sum::<f32>().sqrt();
813 v.into_iter().map(|a| a / norm).collect()
814 }
815
816 #[test]
817 fn score_correlates_with_exact_cosine() {
818 let dim = 128;
819 let mut idx = VecqIndex::new(dim, 7);
820 let base: Vec<Vec<f32>> = (0..200).map(|i| rand_unit(dim, i + 1)).collect();
821 for v in &base {
822 idx.add(v);
823 }
824 let q = rand_unit(dim, 999);
825 let pq = idx.prepare_query(&q);
826 let exact: Vec<f32> = base.iter().map(|v| cosine_f32(&q, v)).collect();
827 let mut max_err = 0f32;
828 for (i, &e) in exact.iter().enumerate().take(200) {
829 let est = idx.score(&pq, i);
830 max_err = max_err.max((est - e).abs());
831 }
832 assert!(max_err < 0.2, "max score error {max_err}");
833 }
834
835 #[test]
836 fn score_reproducible_and_close_to_naive() {
837 let dim = 128;
838 let mut idx = VecqIndex::new(dim, 13);
839 for i in 0..50 {
840 idx.add(&rand_unit(dim, i + 21));
841 }
842 let q = rand_unit(dim, 321);
843 let pq = idx.prepare_query(&q);
844 for vi in 0..50 {
845 let base = vi * (idx.padded() / 2);
846 let mut naive = 0f32;
847 for i in 0..idx.padded() {
848 let b = idx.codes[base + i / 2];
849 let code = if i % 2 == 0 { b & 0x0F } else { b >> 4 };
850 naive += pq.rotated[i] * pq.lut[code as usize];
851 }
852 let s = idx.score(&pq, vi);
853 assert_eq!(s.to_bits(), idx.score(&pq, vi).to_bits());
854 assert!((s - naive * idx.scales[vi]).abs() < 1e-5, "vector {vi}");
855 }
856 }
857
858 #[test]
859 fn neon_matches_scalar_bitwise() {
860 let dim = 128;
861 let mut idx = VecqIndex::new(dim, 42);
862 for i in 0..30 {
863 idx.add(&rand_unit(dim, i + 500));
864 }
865 let q = rand_unit(dim, 777);
866 let pq = idx.prepare_query(&q);
867 for vi in 0..30 {
868 let base = vi * (idx.padded() / 2);
869 let codes = &idx.codes[base..base + idx.padded() / 2];
870 let qslice = &pq.rotated[..idx.padded()];
871 #[cfg(target_arch = "aarch64")]
872 {
873 let neon = unsafe { neon::score_neon(codes, qslice, &pq.lut) };
874 let scalar = score_scalar(codes, qslice, &pq.lut);
875 assert_eq!(
876 neon.to_bits(),
877 scalar.to_bits(),
878 "vector {vi}: NEON and scalar diverged"
879 );
880 }
881 #[cfg(not(target_arch = "aarch64"))]
882 {
883 let _ = (base, codes, qslice);
884 }
885 }
886 }
887
888 #[cfg(target_arch = "aarch64")]
889 #[test]
890 fn neon4_matches_neon_bitwise() {
891 let dim = 128;
892 let mut idx = VecqIndex::new(dim, 91);
893 for i in 0..12 {
894 idx.add(&rand_unit(dim, i + 90));
895 }
896 let q = rand_unit(dim, 1234);
897 let pq = idx.prepare_query(&q);
898 let bpv = idx.padded() / 2;
899 for chunk_start in (0..12).step_by(4) {
900 let codes4 = &idx.codes[chunk_start * bpv..(chunk_start + 4) * bpv];
901 let qslice = &pq.rotated[..idx.padded()];
902 {
903 let batched = unsafe { neon::score_neon4(codes4, qslice, &pq.lut) };
904 for v in 0..4 {
905 let single = unsafe {
906 neon::score_neon(&codes4[v * bpv..(v + 1) * bpv], qslice, &pq.lut)
907 };
908 assert_eq!(
909 batched[v].to_bits(),
910 single.to_bits(),
911 "chunk {chunk_start} vec {v}: neon4 diverged from neon"
912 );
913 }
914 }
915 }
916 }
917
918 #[cfg(target_arch = "x86_64")]
919 #[test]
920 fn avx2_matches_scalar_bitwise() {
921 if !avx2::available() {
922 return; }
924 for (dim, seed) in [(128, 42), (8, 43)] {
927 let mut idx = VecqIndex::new(dim, seed);
928 for i in 0..30 {
929 idx.add(&rand_unit(dim, i + 500));
930 }
931 let q = rand_unit(dim, 777);
932 let pq = idx.prepare_query(&q);
933 for vi in 0..30 {
934 let base = vi * (idx.padded() / 2);
935 let codes = &idx.codes[base..base + idx.padded() / 2];
936 let qslice = &pq.rotated[..idx.padded()];
937 let avx2raw = unsafe { avx2::score_avx2(codes, qslice, &pq.lut) };
938 let scalar = score_scalar(codes, qslice, &pq.lut);
939 assert_eq!(
940 avx2raw.to_bits(),
941 scalar.to_bits(),
942 "dim {dim} vector {vi}: AVX2 and scalar diverged"
943 );
944 }
945 }
946 }
947
948 #[cfg(target_arch = "x86_64")]
949 #[test]
950 fn avx24_matches_avx2_bitwise() {
951 if !avx2::available() {
952 return;
953 }
954 let dim = 128;
955 let mut idx = VecqIndex::new(dim, 91);
956 for i in 0..12 {
957 idx.add(&rand_unit(dim, i + 90));
958 }
959 let q = rand_unit(dim, 1234);
960 let pq = idx.prepare_query(&q);
961 let bpv = idx.padded() / 2;
962 for chunk_start in (0..12).step_by(4) {
963 let codes4 = &idx.codes[chunk_start * bpv..(chunk_start + 4) * bpv];
964 let qslice = &pq.rotated[..idx.padded()];
965 let batched = unsafe { avx2::score_avx24(codes4, qslice, &pq.lut) };
966 for v in 0..4 {
967 let single =
968 unsafe { avx2::score_avx2(&codes4[v * bpv..(v + 1) * bpv], qslice, &pq.lut) };
969 assert_eq!(
970 batched[v].to_bits(),
971 single.to_bits(),
972 "chunk {chunk_start} vec {v}: avx24 diverged from avx2"
973 );
974 }
975 }
976 }
977
978 #[cfg(target_arch = "x86_64")]
979 #[test]
980 fn search_dispatch_matches_scalar_on_avx2_hosts() {
981 let dim = 128;
984 let mut idx = VecqIndex::new(dim, 55);
985 for i in 0..30 {
986 idx.add(&rand_unit(dim, i + 800));
987 }
988 let q = rand_unit(dim, 888);
989 let pq = idx.prepare_query(&q);
990 for vi in 0..30 {
991 let base = vi * (idx.padded() / 2);
992 let codes = &idx.codes[base..base + idx.padded() / 2];
993 let qslice = &pq.rotated[..idx.padded()];
994 let scalar = score_scalar(codes, qslice, &pq.lut) * idx.scales[vi];
995 assert_eq!(idx.score(&pq, vi).to_bits(), scalar.to_bits());
996 }
997 }
998
999 #[test]
1000 fn search_returns_sorted_results() {
1001 let dim = 64;
1002 let mut idx = VecqIndex::new(dim, 3);
1003 for i in 0..50 {
1004 idx.add(&rand_unit(dim, i * 31 + 5));
1005 }
1006 let q = rand_unit(dim, 77);
1007 let res = idx.search(&q, 5);
1008 assert_eq!(res.len(), 5);
1009 for w in res.windows(2) {
1010 assert!(w[0].1 >= w[1].1);
1011 }
1012 }
1013
1014 #[test]
1015 fn quantized_size_is_one_eighth() {
1016 let dim = 384;
1017 let mut idx = VecqIndex::new(dim, 1);
1018 idx.add(&rand_unit(dim, 11));
1019 assert_eq!(idx.codes.len(), 512 / 2);
1020 }
1021
1022 #[test]
1023 fn keyed_add_search_remove() {
1024 let dim = 64;
1025 let mut idx = VecqIndex::new(dim, 5);
1026 for i in 0..50u64 {
1027 idx.add_keyed(1000 + i, &rand_unit(dim, i * 17 + 3));
1028 }
1029 assert_eq!(idx.len(), 50);
1030 assert!(idx.contains_key(1000));
1031 assert!(!idx.contains_key(999));
1032
1033 let q = rand_unit(dim, 77);
1034 let keyed = idx.search_keyed(&q, 5);
1035 assert_eq!(keyed.len(), 5);
1036 for w in keyed.windows(2) {
1037 assert!(w[0].1 >= w[1].1);
1038 }
1039 let positional = idx.search(&q, 5);
1041 for ((key, ks), (slot, ps)) in keyed.iter().zip(positional.iter()) {
1042 assert_eq!(key, &idx.key_of(*slot).unwrap());
1043 assert_eq!(ks.to_bits(), ps.to_bits());
1044 }
1045
1046 let top_key = keyed[0].0;
1048 assert!(idx.remove_keyed(top_key));
1049 assert!(!idx.remove_keyed(top_key), "second remove is a no-op");
1050 assert!(!idx.remove_keyed(12345), "unknown key returns false");
1051 assert_eq!(idx.len(), 49);
1052 assert_eq!(idx.tombstones(), 1);
1053 let keyed2 = idx.search_keyed(&q, 5);
1054 assert!(!keyed2.iter().any(|(k, _)| *k == top_key));
1055 for (k, s) in keyed2.iter() {
1056 let old = keyed.iter().find(|(ok, _)| ok == k).map(|(_, os)| *os);
1057 if let Some(os) = old {
1058 assert_eq!(s.to_bits(), os.to_bits(), "key {k} score changed");
1059 }
1060 }
1061 }
1062
1063 #[test]
1064 fn keyed_add_same_key_replaces() {
1065 let dim = 32;
1066 let mut idx = VecqIndex::new(dim, 9);
1067 idx.add_keyed(7, &rand_unit(dim, 101));
1068 idx.add_keyed(7, &rand_unit(dim, 202));
1069 assert_eq!(idx.len(), 1, "replace must not grow the index");
1070 assert_eq!(idx.tombstones(), 0);
1071 let q = rand_unit(dim, 202);
1073 let res = idx.search_keyed(&q, 1);
1074 assert_eq!(res[0].0, 7);
1075 }
1076
1077 #[test]
1078 fn keyed_slot_indices_stay_stable_across_remove_and_serialize() {
1079 let dim = 64;
1080 let mut idx = VecqIndex::new(dim, 15);
1081 for i in 0..20u64 {
1082 idx.add_keyed(i, &rand_unit(dim, i + 300));
1083 }
1084 let q = rand_unit(dim, 404);
1085 let before = idx.search(&q, 20);
1086 idx.remove_keyed(idx.key_of(before[0].0).unwrap());
1088 idx.remove_keyed(idx.key_of(before[5].0).unwrap());
1089 let after = idx.search(&q, 20);
1090 assert_eq!(after.len(), 18);
1091 for (slot, s) in &after {
1092 let old = before.iter().find(|(os, _)| os == slot);
1093 assert!(old.is_some(), "slot {slot} moved after remove");
1094 assert_eq!(old.unwrap().1.to_bits(), s.to_bits());
1095 }
1096 let bytes = idx.to_bytes();
1098 let disk = VecqIndex::from_bytes(&bytes).unwrap();
1099 assert_eq!(disk.len(), 18);
1100 assert_eq!(idx.search(&q, 20), after, "in-memory results unchanged");
1101 }
1102
1103 #[test]
1104 fn compact_drops_tombstones_and_preserves_results() {
1105 let dim = 64;
1106 let mut idx = VecqIndex::new(dim, 21);
1107 for i in 0..40u64 {
1108 idx.add_keyed(10 * i, &rand_unit(dim, i + 61));
1109 }
1110 for i in 0..20u64 {
1111 assert!(idx.remove_keyed(10 * i));
1112 }
1113 let q = rand_unit(dim, 123);
1114 let expected = idx.search_keyed(&q, 20);
1115 idx.compact();
1116 assert_eq!(idx.tombstones(), 0);
1117 assert_eq!(idx.len(), 20);
1118 assert_eq!(idx.search_keyed(&q, 20), expected);
1119 let bytes = idx.to_bytes();
1123 let back = VecqIndex::from_bytes(&bytes).unwrap();
1124 let reloaded = back.search(&q, 20);
1125 assert_eq!(reloaded.len(), 20);
1126 for ((slot, s), (key, ks)) in reloaded.iter().zip(expected.iter()) {
1127 assert_eq!(idx.key_of(*slot), Some(*key));
1128 assert!(
1129 (s - ks).abs() < 1e-3,
1130 "key {key} score drifted: {s} vs {ks}"
1131 );
1132 }
1133 }
1134
1135 #[test]
1136 fn keyed_search_on_empty_and_drained_index() {
1137 let dim = 32;
1138 let mut idx = VecqIndex::new(dim, 31);
1139 assert!(idx.search_keyed(&rand_unit(dim, 1), 3).is_empty());
1140 idx.add_keyed(1, &rand_unit(dim, 2));
1141 idx.add_keyed(2, &rand_unit(dim, 3));
1142 assert!(idx.remove_keyed(1));
1143 assert!(idx.remove_keyed(2));
1144 assert!(idx.is_empty(), "drained index reports empty");
1145 assert_eq!(idx.tombstones(), 2);
1146 assert!(idx.search_keyed(&rand_unit(dim, 4), 3).is_empty());
1147 }
1148
1149 #[test]
1150 fn keyed_index_from_file_supports_keyed_adds() {
1151 let dim = 64;
1152 let mut idx = VecqIndex::new(dim, 41);
1153 for i in 0..10u64 {
1154 idx.add(&rand_unit(dim, i + 700));
1155 }
1156 let bytes = idx.to_bytes();
1157 let mut back = VecqIndex::from_bytes(&bytes).unwrap();
1158 back.add_keyed(555, &rand_unit(dim, 999));
1159 assert!(back.contains_key(555));
1160 assert_eq!(back.len(), 11);
1161 let q = rand_unit(dim, 999);
1162 assert_eq!(back.search_keyed(&q, 1)[0].0, 555);
1163 }
1164
1165 #[test]
1166 #[should_panic(expected = "vector dim mismatch")]
1167 fn keyed_add_dim_mismatch_panics() {
1168 let mut idx = VecqIndex::new(32, 3);
1169 idx.add_keyed(1, &[0.5; 64]);
1170 }
1171}