Skip to main content

hermes_core/structures/
simd.rs

1//! Shared SIMD-accelerated functions for posting list compression
2//!
3//! This module provides platform-optimized implementations for common operations:
4//! - **Unpacking**: Convert packed 8/16/32-bit values to u32 arrays
5//! - **Delta decoding**: Prefix sum for converting deltas to absolute values
6//! - **Add one**: Increment all values in an array (for TF decoding)
7//!
8//! Supports:
9//! - **NEON** on aarch64 (Apple Silicon, ARM servers)
10//! - **SSE/SSE4.1** on x86_64 (Intel/AMD)
11//! - **Scalar fallback** for other architectures
12
13// ============================================================================
14// NEON intrinsics for aarch64 (Apple Silicon, ARM servers)
15// ============================================================================
16
17#[cfg(target_arch = "aarch64")]
18#[allow(unsafe_op_in_unsafe_fn)]
19mod neon {
20    use std::arch::aarch64::*;
21
22    /// SIMD unpack for 8-bit values using NEON
23    #[target_feature(enable = "neon")]
24    pub unsafe fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
25        let chunks = count / 16;
26        let remainder = count % 16;
27
28        for chunk in 0..chunks {
29            let base = chunk * 16;
30            let in_ptr = input.as_ptr().add(base);
31
32            // Load 16 bytes
33            let bytes = vld1q_u8(in_ptr);
34
35            // Widen u8 -> u16 -> u32
36            let low8 = vget_low_u8(bytes);
37            let high8 = vget_high_u8(bytes);
38
39            let low16 = vmovl_u8(low8);
40            let high16 = vmovl_u8(high8);
41
42            let v0 = vmovl_u16(vget_low_u16(low16));
43            let v1 = vmovl_u16(vget_high_u16(low16));
44            let v2 = vmovl_u16(vget_low_u16(high16));
45            let v3 = vmovl_u16(vget_high_u16(high16));
46
47            let out_ptr = output.as_mut_ptr().add(base);
48            vst1q_u32(out_ptr, v0);
49            vst1q_u32(out_ptr.add(4), v1);
50            vst1q_u32(out_ptr.add(8), v2);
51            vst1q_u32(out_ptr.add(12), v3);
52        }
53
54        // Handle remainder
55        let base = chunks * 16;
56        for i in 0..remainder {
57            output[base + i] = input[base + i] as u32;
58        }
59    }
60
61    /// SIMD unpack for 16-bit values using NEON
62    #[target_feature(enable = "neon")]
63    pub unsafe fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
64        let chunks = count / 8;
65        let remainder = count % 8;
66
67        for chunk in 0..chunks {
68            let base = chunk * 8;
69            let in_ptr = input.as_ptr().add(base * 2) as *const u16;
70
71            let vals = vld1q_u16(in_ptr);
72            let low = vmovl_u16(vget_low_u16(vals));
73            let high = vmovl_u16(vget_high_u16(vals));
74
75            let out_ptr = output.as_mut_ptr().add(base);
76            vst1q_u32(out_ptr, low);
77            vst1q_u32(out_ptr.add(4), high);
78        }
79
80        // Handle remainder
81        let base = chunks * 8;
82        for i in 0..remainder {
83            let idx = (base + i) * 2;
84            output[base + i] = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
85        }
86    }
87
88    /// SIMD unpack for 32-bit values using NEON (fast copy)
89    #[target_feature(enable = "neon")]
90    pub unsafe fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
91        let chunks = count / 4;
92        let remainder = count % 4;
93
94        let in_ptr = input.as_ptr() as *const u32;
95        let out_ptr = output.as_mut_ptr();
96
97        for chunk in 0..chunks {
98            let vals = vld1q_u32(in_ptr.add(chunk * 4));
99            vst1q_u32(out_ptr.add(chunk * 4), vals);
100        }
101
102        // Handle remainder
103        let base = chunks * 4;
104        for i in 0..remainder {
105            let idx = (base + i) * 4;
106            output[base + i] =
107                u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
108        }
109    }
110
111    /// SIMD prefix sum for 4 u32 values using NEON
112    /// Input:  [a, b, c, d]
113    /// Output: [a, a+b, a+b+c, a+b+c+d]
114    #[inline]
115    #[target_feature(enable = "neon")]
116    unsafe fn prefix_sum_4(v: uint32x4_t) -> uint32x4_t {
117        // Step 1: shift by 1 and add
118        // [a, b, c, d] + [0, a, b, c] = [a, a+b, b+c, c+d]
119        let shifted1 = vextq_u32(vdupq_n_u32(0), v, 3);
120        let sum1 = vaddq_u32(v, shifted1);
121
122        // Step 2: shift by 2 and add
123        // [a, a+b, b+c, c+d] + [0, 0, a, a+b] = [a, a+b, a+b+c, a+b+c+d]
124        let shifted2 = vextq_u32(vdupq_n_u32(0), sum1, 2);
125        vaddq_u32(sum1, shifted2)
126    }
127
128    /// SIMD delta decode: convert deltas to absolute doc IDs
129    /// deltas[i] stores (gap - 1), output[i] = first + sum(gaps[0..i])
130    /// Uses NEON SIMD prefix sum for high throughput
131    #[target_feature(enable = "neon")]
132    pub unsafe fn delta_decode(
133        output: &mut [u32],
134        deltas: &[u32],
135        first_doc_id: u32,
136        count: usize,
137    ) {
138        if count == 0 {
139            return;
140        }
141
142        output[0] = first_doc_id;
143        if count == 1 {
144            return;
145        }
146
147        let ones = vdupq_n_u32(1);
148        let mut carry = vdupq_n_u32(first_doc_id);
149
150        let full_groups = (count - 1) / 4;
151        let remainder = (count - 1) % 4;
152
153        for group in 0..full_groups {
154            let base = group * 4;
155
156            // Load 4 deltas and add 1 (since we store gap-1)
157            let d = vld1q_u32(deltas[base..].as_ptr());
158            let gaps = vaddq_u32(d, ones);
159
160            // Compute prefix sum within the 4 elements
161            let prefix = prefix_sum_4(gaps);
162
163            // Add carry (broadcast last element of previous group)
164            let result = vaddq_u32(prefix, carry);
165
166            // Store result
167            vst1q_u32(output[base + 1..].as_mut_ptr(), result);
168
169            // Update carry: broadcast the last element for next iteration
170            carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
171        }
172
173        // Handle remainder
174        let base = full_groups * 4;
175        let mut scalar_carry = vgetq_lane_u32(carry, 0);
176        for j in 0..remainder {
177            scalar_carry = scalar_carry.wrapping_add(deltas[base + j]).wrapping_add(1);
178            output[base + j + 1] = scalar_carry;
179        }
180    }
181
182    /// SIMD add 1 to all values (for TF decoding: stored as tf-1)
183    #[target_feature(enable = "neon")]
184    pub unsafe fn add_one(values: &mut [u32], count: usize) {
185        let ones = vdupq_n_u32(1);
186        let chunks = count / 4;
187        let remainder = count % 4;
188
189        for chunk in 0..chunks {
190            let base = chunk * 4;
191            let ptr = values.as_mut_ptr().add(base);
192            let v = vld1q_u32(ptr);
193            let result = vaddq_u32(v, ones);
194            vst1q_u32(ptr, result);
195        }
196
197        let base = chunks * 4;
198        for i in 0..remainder {
199            values[base + i] += 1;
200        }
201    }
202
203    /// Fused unpack 8-bit + delta decode using NEON
204    /// Processes 4 values at a time, fusing unpack and prefix sum
205    #[target_feature(enable = "neon")]
206    pub unsafe fn unpack_8bit_delta_decode(
207        input: &[u8],
208        output: &mut [u32],
209        first_value: u32,
210        count: usize,
211    ) {
212        output[0] = first_value;
213        if count <= 1 {
214            return;
215        }
216
217        let ones = vdupq_n_u32(1);
218        let mut carry = vdupq_n_u32(first_value);
219
220        let full_groups = (count - 1) / 4;
221        let remainder = (count - 1) % 4;
222
223        for group in 0..full_groups {
224            let base = group * 4;
225
226            // Load 4 bytes as a u32, then widen u8→u16→u32 via NEON
227            let raw = std::ptr::read_unaligned(input.as_ptr().add(base) as *const u32);
228            let bytes = vreinterpret_u8_u32(vdup_n_u32(raw));
229            let u16s = vmovl_u8(bytes); // 8×u8 → 8×u16 (only low 4 matter)
230            let d = vmovl_u16(vget_low_u16(u16s)); // 4×u16 → 4×u32
231
232            // Add 1 (since we store gap-1)
233            let gaps = vaddq_u32(d, ones);
234
235            // Compute prefix sum within the 4 elements
236            let prefix = prefix_sum_4(gaps);
237
238            // Add carry
239            let result = vaddq_u32(prefix, carry);
240
241            // Store result
242            vst1q_u32(output[base + 1..].as_mut_ptr(), result);
243
244            // Update carry
245            carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
246        }
247
248        // Handle remainder
249        let base = full_groups * 4;
250        let mut scalar_carry = vgetq_lane_u32(carry, 0);
251        for j in 0..remainder {
252            scalar_carry = scalar_carry
253                .wrapping_add(input[base + j] as u32)
254                .wrapping_add(1);
255            output[base + j + 1] = scalar_carry;
256        }
257    }
258
259    /// Fused unpack 16-bit + delta decode using NEON
260    #[target_feature(enable = "neon")]
261    pub unsafe fn unpack_16bit_delta_decode(
262        input: &[u8],
263        output: &mut [u32],
264        first_value: u32,
265        count: usize,
266    ) {
267        output[0] = first_value;
268        if count <= 1 {
269            return;
270        }
271
272        let ones = vdupq_n_u32(1);
273        let mut carry = vdupq_n_u32(first_value);
274
275        let full_groups = (count - 1) / 4;
276        let remainder = (count - 1) % 4;
277
278        for group in 0..full_groups {
279            let base = group * 4;
280            let in_ptr = input.as_ptr().add(base * 2) as *const u16;
281
282            // Load 4 u16 values and widen to u32
283            let vals = vld1_u16(in_ptr);
284            let d = vmovl_u16(vals);
285
286            // Add 1 (since we store gap-1)
287            let gaps = vaddq_u32(d, ones);
288
289            // Compute prefix sum within the 4 elements
290            let prefix = prefix_sum_4(gaps);
291
292            // Add carry
293            let result = vaddq_u32(prefix, carry);
294
295            // Store result
296            vst1q_u32(output[base + 1..].as_mut_ptr(), result);
297
298            // Update carry
299            carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
300        }
301
302        // Handle remainder
303        let base = full_groups * 4;
304        let mut scalar_carry = vgetq_lane_u32(carry, 0);
305        for j in 0..remainder {
306            let idx = (base + j) * 2;
307            let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
308            scalar_carry = scalar_carry.wrapping_add(delta).wrapping_add(1);
309            output[base + j + 1] = scalar_carry;
310        }
311    }
312
313    /// NEON Hamming distance: XOR + byte popcount + horizontal sum.
314    /// Processes 16 bytes per iteration (vs 8 for scalar u64 path).
315    #[target_feature(enable = "neon")]
316    pub unsafe fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
317        let len = a.len();
318        let chunks16 = len / 16;
319        let mut total = 0u32;
320
321        // Process 16 bytes at a time, flush u8 accumulators every 31 iters
322        // (vcntq_u8 returns 0-8 per lane; 31 * 8 = 248 ≤ 255, avoiding u8 overflow)
323        let mut i = 0;
324        while i < chunks16 {
325            let batch_end = (i + 31).min(chunks16);
326            let mut acc = vdupq_n_u8(0);
327            for j in i..batch_end {
328                let off = j * 16;
329                let va = vld1q_u8(a.as_ptr().add(off));
330                let vb = vld1q_u8(b.as_ptr().add(off));
331                let popcnt = vcntq_u8(veorq_u8(va, vb));
332                acc = vaddq_u8(acc, popcnt);
333            }
334            // Widen u8 -> u16 -> u32 -> u64 and horizontal sum
335            let sum64 = vpaddlq_u32(vpaddlq_u16(vpaddlq_u8(acc)));
336            total += vgetq_lane_u64(sum64, 0) as u32 + vgetq_lane_u64(sum64, 1) as u32;
337            i = batch_end;
338        }
339
340        // Remainder bytes (< 16)
341        let base = chunks16 * 16;
342        for k in base..len {
343            total += (a[k] ^ b[k]).count_ones();
344        }
345
346        total
347    }
348
349    /// Check if NEON is available (always true on aarch64)
350    #[inline]
351    pub fn is_available() -> bool {
352        true
353    }
354}
355
356// ============================================================================
357// SSE intrinsics for x86_64 (Intel/AMD)
358// ============================================================================
359
360#[cfg(target_arch = "x86_64")]
361#[allow(unsafe_op_in_unsafe_fn)]
362mod sse {
363    use std::arch::x86_64::*;
364
365    /// SIMD unpack for 8-bit values using SSE
366    #[target_feature(enable = "sse2", enable = "sse4.1")]
367    pub unsafe fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
368        let chunks = count / 16;
369        let remainder = count % 16;
370
371        for chunk in 0..chunks {
372            let base = chunk * 16;
373            let in_ptr = input.as_ptr().add(base);
374
375            let bytes = _mm_loadu_si128(in_ptr as *const __m128i);
376
377            // Zero extend u8 -> u32 using SSE4.1 pmovzx
378            let v0 = _mm_cvtepu8_epi32(bytes);
379            let v1 = _mm_cvtepu8_epi32(_mm_srli_si128(bytes, 4));
380            let v2 = _mm_cvtepu8_epi32(_mm_srli_si128(bytes, 8));
381            let v3 = _mm_cvtepu8_epi32(_mm_srli_si128(bytes, 12));
382
383            let out_ptr = output.as_mut_ptr().add(base);
384            _mm_storeu_si128(out_ptr as *mut __m128i, v0);
385            _mm_storeu_si128(out_ptr.add(4) as *mut __m128i, v1);
386            _mm_storeu_si128(out_ptr.add(8) as *mut __m128i, v2);
387            _mm_storeu_si128(out_ptr.add(12) as *mut __m128i, v3);
388        }
389
390        let base = chunks * 16;
391        for i in 0..remainder {
392            output[base + i] = input[base + i] as u32;
393        }
394    }
395
396    /// SIMD unpack for 16-bit values using SSE
397    #[target_feature(enable = "sse2", enable = "sse4.1")]
398    pub unsafe fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
399        let chunks = count / 8;
400        let remainder = count % 8;
401
402        for chunk in 0..chunks {
403            let base = chunk * 8;
404            let in_ptr = input.as_ptr().add(base * 2);
405
406            let vals = _mm_loadu_si128(in_ptr as *const __m128i);
407            let low = _mm_cvtepu16_epi32(vals);
408            let high = _mm_cvtepu16_epi32(_mm_srli_si128(vals, 8));
409
410            let out_ptr = output.as_mut_ptr().add(base);
411            _mm_storeu_si128(out_ptr as *mut __m128i, low);
412            _mm_storeu_si128(out_ptr.add(4) as *mut __m128i, high);
413        }
414
415        let base = chunks * 8;
416        for i in 0..remainder {
417            let idx = (base + i) * 2;
418            output[base + i] = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
419        }
420    }
421
422    /// SIMD unpack for 32-bit values using SSE (fast copy)
423    #[target_feature(enable = "sse2")]
424    pub unsafe fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
425        let chunks = count / 4;
426        let remainder = count % 4;
427
428        let in_ptr = input.as_ptr() as *const __m128i;
429        let out_ptr = output.as_mut_ptr() as *mut __m128i;
430
431        for chunk in 0..chunks {
432            let vals = _mm_loadu_si128(in_ptr.add(chunk));
433            _mm_storeu_si128(out_ptr.add(chunk), vals);
434        }
435
436        // Handle remainder
437        let base = chunks * 4;
438        for i in 0..remainder {
439            let idx = (base + i) * 4;
440            output[base + i] =
441                u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
442        }
443    }
444
445    /// SIMD prefix sum for 4 u32 values using SSE
446    /// Input:  [a, b, c, d]
447    /// Output: [a, a+b, a+b+c, a+b+c+d]
448    #[inline]
449    #[target_feature(enable = "sse2")]
450    unsafe fn prefix_sum_4(v: __m128i) -> __m128i {
451        // Step 1: shift by 1 element (4 bytes) and add
452        // [a, b, c, d] + [0, a, b, c] = [a, a+b, b+c, c+d]
453        let shifted1 = _mm_slli_si128(v, 4);
454        let sum1 = _mm_add_epi32(v, shifted1);
455
456        // Step 2: shift by 2 elements (8 bytes) and add
457        // [a, a+b, b+c, c+d] + [0, 0, a, a+b] = [a, a+b, a+b+c, a+b+c+d]
458        let shifted2 = _mm_slli_si128(sum1, 8);
459        _mm_add_epi32(sum1, shifted2)
460    }
461
462    /// SIMD delta decode using SSE with true SIMD prefix sum
463    #[target_feature(enable = "sse2", enable = "sse4.1")]
464    pub unsafe fn delta_decode(
465        output: &mut [u32],
466        deltas: &[u32],
467        first_doc_id: u32,
468        count: usize,
469    ) {
470        if count == 0 {
471            return;
472        }
473
474        output[0] = first_doc_id;
475        if count == 1 {
476            return;
477        }
478
479        let ones = _mm_set1_epi32(1);
480        let mut carry = _mm_set1_epi32(first_doc_id as i32);
481
482        let full_groups = (count - 1) / 4;
483        let remainder = (count - 1) % 4;
484
485        for group in 0..full_groups {
486            let base = group * 4;
487
488            // Load 4 deltas and add 1 (since we store gap-1)
489            let d = _mm_loadu_si128(deltas[base..].as_ptr() as *const __m128i);
490            let gaps = _mm_add_epi32(d, ones);
491
492            // Compute prefix sum within the 4 elements
493            let prefix = prefix_sum_4(gaps);
494
495            // Add carry (broadcast last element of previous group)
496            let result = _mm_add_epi32(prefix, carry);
497
498            // Store result
499            _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
500
501            // Update carry: broadcast the last element for next iteration
502            carry = _mm_shuffle_epi32(result, 0xFF); // broadcast lane 3
503        }
504
505        // Handle remainder
506        let base = full_groups * 4;
507        let mut scalar_carry = _mm_extract_epi32(carry, 0) as u32;
508        for j in 0..remainder {
509            scalar_carry = scalar_carry.wrapping_add(deltas[base + j]).wrapping_add(1);
510            output[base + j + 1] = scalar_carry;
511        }
512    }
513
514    /// SIMD add 1 to all values using SSE
515    #[target_feature(enable = "sse2")]
516    pub unsafe fn add_one(values: &mut [u32], count: usize) {
517        let ones = _mm_set1_epi32(1);
518        let chunks = count / 4;
519        let remainder = count % 4;
520
521        for chunk in 0..chunks {
522            let base = chunk * 4;
523            let ptr = values.as_mut_ptr().add(base) as *mut __m128i;
524            let v = _mm_loadu_si128(ptr);
525            let result = _mm_add_epi32(v, ones);
526            _mm_storeu_si128(ptr, result);
527        }
528
529        let base = chunks * 4;
530        for i in 0..remainder {
531            values[base + i] += 1;
532        }
533    }
534
535    /// Fused unpack 8-bit + delta decode using SSE
536    #[target_feature(enable = "sse2", enable = "sse4.1")]
537    pub unsafe fn unpack_8bit_delta_decode(
538        input: &[u8],
539        output: &mut [u32],
540        first_value: u32,
541        count: usize,
542    ) {
543        output[0] = first_value;
544        if count <= 1 {
545            return;
546        }
547
548        let ones = _mm_set1_epi32(1);
549        let mut carry = _mm_set1_epi32(first_value as i32);
550
551        let full_groups = (count - 1) / 4;
552        let remainder = (count - 1) % 4;
553
554        for group in 0..full_groups {
555            let base = group * 4;
556
557            // Load 4 bytes (unaligned) and zero-extend to u32
558            let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
559                input.as_ptr().add(base) as *const i32
560            ));
561            let d = _mm_cvtepu8_epi32(bytes);
562
563            // Add 1 (since we store gap-1)
564            let gaps = _mm_add_epi32(d, ones);
565
566            // Compute prefix sum within the 4 elements
567            let prefix = prefix_sum_4(gaps);
568
569            // Add carry
570            let result = _mm_add_epi32(prefix, carry);
571
572            // Store result
573            _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
574
575            // Update carry: broadcast the last element
576            carry = _mm_shuffle_epi32(result, 0xFF);
577        }
578
579        // Handle remainder
580        let base = full_groups * 4;
581        let mut scalar_carry = _mm_extract_epi32(carry, 0) as u32;
582        for j in 0..remainder {
583            scalar_carry = scalar_carry
584                .wrapping_add(input[base + j] as u32)
585                .wrapping_add(1);
586            output[base + j + 1] = scalar_carry;
587        }
588    }
589
590    /// Fused unpack 16-bit + delta decode using SSE
591    #[target_feature(enable = "sse2", enable = "sse4.1")]
592    pub unsafe fn unpack_16bit_delta_decode(
593        input: &[u8],
594        output: &mut [u32],
595        first_value: u32,
596        count: usize,
597    ) {
598        output[0] = first_value;
599        if count <= 1 {
600            return;
601        }
602
603        let ones = _mm_set1_epi32(1);
604        let mut carry = _mm_set1_epi32(first_value as i32);
605
606        let full_groups = (count - 1) / 4;
607        let remainder = (count - 1) % 4;
608
609        for group in 0..full_groups {
610            let base = group * 4;
611            let in_ptr = input.as_ptr().add(base * 2);
612
613            // Load 8 bytes (4 u16 values, unaligned) and zero-extend to u32
614            let vals = _mm_loadl_epi64(in_ptr as *const __m128i); // loadl_epi64 supports unaligned
615            let d = _mm_cvtepu16_epi32(vals);
616
617            // Add 1 (since we store gap-1)
618            let gaps = _mm_add_epi32(d, ones);
619
620            // Compute prefix sum within the 4 elements
621            let prefix = prefix_sum_4(gaps);
622
623            // Add carry
624            let result = _mm_add_epi32(prefix, carry);
625
626            // Store result
627            _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
628
629            // Update carry: broadcast the last element
630            carry = _mm_shuffle_epi32(result, 0xFF);
631        }
632
633        // Handle remainder
634        let base = full_groups * 4;
635        let mut scalar_carry = _mm_extract_epi32(carry, 0) as u32;
636        for j in 0..remainder {
637            let idx = (base + j) * 2;
638            let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
639            scalar_carry = scalar_carry.wrapping_add(delta).wrapping_add(1);
640            output[base + j + 1] = scalar_carry;
641        }
642    }
643
644    /// Check if SSE4.1 is available at runtime
645    #[inline]
646    pub fn is_available() -> bool {
647        is_x86_feature_detected!("sse4.1")
648    }
649}
650
651// ============================================================================
652// AVX2 intrinsics for x86_64 (Intel/AMD with 256-bit registers)
653// ============================================================================
654
655#[cfg(target_arch = "x86_64")]
656#[allow(unsafe_op_in_unsafe_fn)]
657mod avx2 {
658    use std::arch::x86_64::*;
659
660    /// AVX2 unpack for 8-bit values (processes 32 bytes at a time)
661    #[target_feature(enable = "avx2")]
662    pub unsafe fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
663        let chunks = count / 32;
664        let remainder = count % 32;
665
666        for chunk in 0..chunks {
667            let base = chunk * 32;
668            let in_ptr = input.as_ptr().add(base);
669
670            // Load 32 bytes (two 128-bit loads, then combine)
671            let bytes_lo = _mm_loadu_si128(in_ptr as *const __m128i);
672            let bytes_hi = _mm_loadu_si128(in_ptr.add(16) as *const __m128i);
673
674            // Zero extend first 16 bytes: u8 -> u32
675            let v0 = _mm256_cvtepu8_epi32(bytes_lo);
676            let v1 = _mm256_cvtepu8_epi32(_mm_srli_si128(bytes_lo, 8));
677            let v2 = _mm256_cvtepu8_epi32(bytes_hi);
678            let v3 = _mm256_cvtepu8_epi32(_mm_srli_si128(bytes_hi, 8));
679
680            let out_ptr = output.as_mut_ptr().add(base);
681            _mm256_storeu_si256(out_ptr as *mut __m256i, v0);
682            _mm256_storeu_si256(out_ptr.add(8) as *mut __m256i, v1);
683            _mm256_storeu_si256(out_ptr.add(16) as *mut __m256i, v2);
684            _mm256_storeu_si256(out_ptr.add(24) as *mut __m256i, v3);
685        }
686
687        // Handle remainder with SSE
688        let base = chunks * 32;
689        for i in 0..remainder {
690            output[base + i] = input[base + i] as u32;
691        }
692    }
693
694    /// AVX2 unpack for 16-bit values (processes 16 values at a time)
695    #[target_feature(enable = "avx2")]
696    pub unsafe fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
697        let chunks = count / 16;
698        let remainder = count % 16;
699
700        for chunk in 0..chunks {
701            let base = chunk * 16;
702            let in_ptr = input.as_ptr().add(base * 2);
703
704            // Load 32 bytes (16 u16 values)
705            let vals_lo = _mm_loadu_si128(in_ptr as *const __m128i);
706            let vals_hi = _mm_loadu_si128(in_ptr.add(16) as *const __m128i);
707
708            // Zero extend u16 -> u32
709            let v0 = _mm256_cvtepu16_epi32(vals_lo);
710            let v1 = _mm256_cvtepu16_epi32(vals_hi);
711
712            let out_ptr = output.as_mut_ptr().add(base);
713            _mm256_storeu_si256(out_ptr as *mut __m256i, v0);
714            _mm256_storeu_si256(out_ptr.add(8) as *mut __m256i, v1);
715        }
716
717        // Handle remainder
718        let base = chunks * 16;
719        for i in 0..remainder {
720            let idx = (base + i) * 2;
721            output[base + i] = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
722        }
723    }
724
725    /// AVX2 unpack for 32-bit values (fast copy, 8 values at a time)
726    #[target_feature(enable = "avx2")]
727    pub unsafe fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
728        let chunks = count / 8;
729        let remainder = count % 8;
730
731        let in_ptr = input.as_ptr() as *const __m256i;
732        let out_ptr = output.as_mut_ptr() as *mut __m256i;
733
734        for chunk in 0..chunks {
735            let vals = _mm256_loadu_si256(in_ptr.add(chunk));
736            _mm256_storeu_si256(out_ptr.add(chunk), vals);
737        }
738
739        // Handle remainder
740        let base = chunks * 8;
741        for i in 0..remainder {
742            let idx = (base + i) * 4;
743            output[base + i] =
744                u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
745        }
746    }
747
748    /// AVX2 add 1 to all values (8 values at a time)
749    #[target_feature(enable = "avx2")]
750    pub unsafe fn add_one(values: &mut [u32], count: usize) {
751        let ones = _mm256_set1_epi32(1);
752        let chunks = count / 8;
753        let remainder = count % 8;
754
755        for chunk in 0..chunks {
756            let base = chunk * 8;
757            let ptr = values.as_mut_ptr().add(base) as *mut __m256i;
758            let v = _mm256_loadu_si256(ptr);
759            let result = _mm256_add_epi32(v, ones);
760            _mm256_storeu_si256(ptr, result);
761        }
762
763        let base = chunks * 8;
764        for i in 0..remainder {
765            values[base + i] += 1;
766        }
767    }
768
769    /// AVX2 prefix sum for 8 u32 values (Hillis-Steele)
770    /// Input:  [a, b, c, d, e, f, g, h]
771    /// Output: [a, a+b, a+b+c, ..., a+b+c+d+e+f+g+h]
772    #[inline]
773    #[target_feature(enable = "avx2")]
774    unsafe fn prefix_sum_8(v: __m256i) -> __m256i {
775        // Step 1: intra-lane shift by 1 element (4 bytes) and add
776        let s1 = _mm256_slli_si256(v, 4);
777        let r1 = _mm256_add_epi32(v, s1);
778
779        // Step 2: intra-lane shift by 2 elements (8 bytes) and add
780        let s2 = _mm256_slli_si256(r1, 8);
781        let r2 = _mm256_add_epi32(r1, s2);
782
783        // Step 3: propagate lower lane sum to upper lane
784        // Broadcast element 3 (lower lane sum) within each lane
785        let lo_sum = _mm256_shuffle_epi32(r2, 0xFF);
786        // Duplicate lane 0 to both lanes
787        let carry = _mm256_permute2x128_si256(lo_sum, lo_sum, 0x00);
788        // Zero carry for lower lane, keep for upper
789        let carry_hi = _mm256_blend_epi32::<0xF0>(_mm256_setzero_si256(), carry);
790        _mm256_add_epi32(r2, carry_hi)
791    }
792
793    /// AVX2 fused unpack 8-bit + delta decode (processes 8 values at a time)
794    #[target_feature(enable = "avx2")]
795    pub unsafe fn unpack_8bit_delta_decode(
796        input: &[u8],
797        output: &mut [u32],
798        first_value: u32,
799        count: usize,
800    ) {
801        output[0] = first_value;
802        if count <= 1 {
803            return;
804        }
805
806        let ones = _mm256_set1_epi32(1);
807        let mut carry = _mm256_set1_epi32(first_value as i32);
808        let broadcast_idx = _mm256_set1_epi32(7);
809
810        let full_groups = (count - 1) / 8;
811        let remainder = (count - 1) % 8;
812
813        for group in 0..full_groups {
814            let base = group * 8;
815
816            // Load 8 bytes and zero-extend to 8×u32
817            let bytes = _mm_loadl_epi64(input.as_ptr().add(base) as *const __m128i);
818            let d = _mm256_cvtepu8_epi32(bytes);
819
820            // Add 1 (since we store gap-1)
821            let gaps = _mm256_add_epi32(d, ones);
822
823            // Compute prefix sum within 8 elements
824            let prefix = prefix_sum_8(gaps);
825
826            // Add carry from previous group
827            let result = _mm256_add_epi32(prefix, carry);
828
829            // Store 8 results
830            _mm256_storeu_si256(output[base + 1..].as_mut_ptr() as *mut __m256i, result);
831
832            // Update carry: broadcast element 7 to all positions
833            carry = _mm256_permutevar8x32_epi32(result, broadcast_idx);
834        }
835
836        // Handle remainder with scalar
837        let base = full_groups * 8;
838        let mut scalar_carry = _mm256_extract_epi32::<0>(carry) as u32;
839        for j in 0..remainder {
840            scalar_carry = scalar_carry
841                .wrapping_add(input[base + j] as u32)
842                .wrapping_add(1);
843            output[base + j + 1] = scalar_carry;
844        }
845    }
846
847    /// AVX2 fused unpack 16-bit + delta decode (processes 8 values at a time)
848    #[target_feature(enable = "avx2")]
849    pub unsafe fn unpack_16bit_delta_decode(
850        input: &[u8],
851        output: &mut [u32],
852        first_value: u32,
853        count: usize,
854    ) {
855        output[0] = first_value;
856        if count <= 1 {
857            return;
858        }
859
860        let ones = _mm256_set1_epi32(1);
861        let mut carry = _mm256_set1_epi32(first_value as i32);
862        let broadcast_idx = _mm256_set1_epi32(7);
863
864        let full_groups = (count - 1) / 8;
865        let remainder = (count - 1) % 8;
866
867        for group in 0..full_groups {
868            let base = group * 8;
869            let in_ptr = input.as_ptr().add(base * 2);
870
871            // Load 16 bytes (8 u16 values) and zero-extend to 8×u32
872            let vals = _mm_loadu_si128(in_ptr as *const __m128i);
873            let d = _mm256_cvtepu16_epi32(vals);
874
875            // Add 1 (since we store gap-1)
876            let gaps = _mm256_add_epi32(d, ones);
877
878            // Compute prefix sum within 8 elements
879            let prefix = prefix_sum_8(gaps);
880
881            // Add carry from previous group
882            let result = _mm256_add_epi32(prefix, carry);
883
884            // Store 8 results
885            _mm256_storeu_si256(output[base + 1..].as_mut_ptr() as *mut __m256i, result);
886
887            // Update carry: broadcast element 7 to all positions
888            carry = _mm256_permutevar8x32_epi32(result, broadcast_idx);
889        }
890
891        // Handle remainder with scalar
892        let base = full_groups * 8;
893        let mut scalar_carry = _mm256_extract_epi32::<0>(carry) as u32;
894        for j in 0..remainder {
895            let idx = (base + j) * 2;
896            let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
897            scalar_carry = scalar_carry.wrapping_add(delta).wrapping_add(1);
898            output[base + j + 1] = scalar_carry;
899        }
900    }
901
902    /// AVX2 Hamming distance using VPSHUFB-based popcount (Muła algorithm).
903    /// Processes 32 bytes per iteration with a nibble lookup table.
904    #[target_feature(enable = "avx2")]
905    pub unsafe fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
906        let len = a.len();
907        let chunks32 = len / 32;
908        let low_mask = _mm256_set1_epi8(0x0f);
909        // Nibble popcount lookup table: popcount(0..15)
910        let lookup = _mm256_setr_epi8(
911            0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2,
912            3, 3, 4,
913        );
914        let mut total = 0u64;
915
916        let mut i = 0;
917        while i < chunks32 {
918            // Accumulate in u8 lanes, flush every 31 iters to avoid overflow
919            // (nibble popcount gives 0-8 per lane; 31 * 8 = 248 ≤ 255)
920            let batch_end = (i + 31).min(chunks32);
921            let mut acc = _mm256_setzero_si256();
922            for j in i..batch_end {
923                let off = j * 32;
924                let va = _mm256_loadu_si256(a.as_ptr().add(off) as *const __m256i);
925                let vb = _mm256_loadu_si256(b.as_ptr().add(off) as *const __m256i);
926                let xored = _mm256_xor_si256(va, vb);
927                // VPSHUFB popcount: count bits per byte via nibble lookup
928                let lo = _mm256_and_si256(xored, low_mask);
929                let hi = _mm256_and_si256(_mm256_srli_epi16(xored, 4), low_mask);
930                let popcnt = _mm256_add_epi8(
931                    _mm256_shuffle_epi8(lookup, lo),
932                    _mm256_shuffle_epi8(lookup, hi),
933                );
934                acc = _mm256_add_epi8(acc, popcnt);
935            }
936            // Horizontal sum: u8 -> u64 via SAD against zero
937            let sad = _mm256_sad_epu8(acc, _mm256_setzero_si256());
938            total += _mm256_extract_epi64(sad, 0) as u64
939                + _mm256_extract_epi64(sad, 1) as u64
940                + _mm256_extract_epi64(sad, 2) as u64
941                + _mm256_extract_epi64(sad, 3) as u64;
942            i = batch_end;
943        }
944
945        // Remainder bytes (< 32)
946        let base = chunks32 * 32;
947        for k in base..len {
948            total += (a[k] ^ b[k]).count_ones() as u64;
949        }
950
951        total as u32
952    }
953
954    /// Check if AVX2 is available at runtime
955    #[inline]
956    pub fn is_available() -> bool {
957        is_x86_feature_detected!("avx2")
958    }
959}
960
961// ============================================================================
962// Scalar fallback implementations
963// ============================================================================
964
965#[allow(dead_code)]
966mod scalar {
967    /// Scalar unpack for 8-bit values
968    #[inline]
969    pub fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
970        for i in 0..count {
971            output[i] = input[i] as u32;
972        }
973    }
974
975    /// Scalar unpack for 16-bit values
976    #[inline]
977    pub fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
978        for (i, out) in output.iter_mut().enumerate().take(count) {
979            let idx = i * 2;
980            *out = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
981        }
982    }
983
984    /// Scalar unpack for 32-bit values
985    #[inline]
986    pub fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
987        for (i, out) in output.iter_mut().enumerate().take(count) {
988            let idx = i * 4;
989            *out = u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
990        }
991    }
992
993    /// Scalar delta decode
994    #[inline]
995    pub fn delta_decode(output: &mut [u32], deltas: &[u32], first_doc_id: u32, count: usize) {
996        if count == 0 {
997            return;
998        }
999
1000        output[0] = first_doc_id;
1001        let mut carry = first_doc_id;
1002
1003        for i in 0..count - 1 {
1004            carry = carry.wrapping_add(deltas[i]).wrapping_add(1);
1005            output[i + 1] = carry;
1006        }
1007    }
1008
1009    /// Scalar add 1 to all values
1010    #[inline]
1011    pub fn add_one(values: &mut [u32], count: usize) {
1012        for val in values.iter_mut().take(count) {
1013            *val += 1;
1014        }
1015    }
1016}
1017
1018// ============================================================================
1019// Public dispatch functions that select SIMD or scalar at runtime
1020// ============================================================================
1021
1022/// Unpack 8-bit packed values to u32 with SIMD acceleration
1023#[inline]
1024pub fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
1025    #[cfg(target_arch = "aarch64")]
1026    {
1027        if neon::is_available() {
1028            unsafe {
1029                neon::unpack_8bit(input, output, count);
1030            }
1031            return;
1032        }
1033    }
1034
1035    #[cfg(target_arch = "x86_64")]
1036    {
1037        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1038        if avx2::is_available() {
1039            unsafe {
1040                avx2::unpack_8bit(input, output, count);
1041            }
1042            return;
1043        }
1044        if sse::is_available() {
1045            unsafe {
1046                sse::unpack_8bit(input, output, count);
1047            }
1048            return;
1049        }
1050    }
1051
1052    scalar::unpack_8bit(input, output, count);
1053}
1054
1055/// Unpack 16-bit packed values to u32 with SIMD acceleration
1056#[inline]
1057pub fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
1058    #[cfg(target_arch = "aarch64")]
1059    {
1060        if neon::is_available() {
1061            unsafe {
1062                neon::unpack_16bit(input, output, count);
1063            }
1064            return;
1065        }
1066    }
1067
1068    #[cfg(target_arch = "x86_64")]
1069    {
1070        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1071        if avx2::is_available() {
1072            unsafe {
1073                avx2::unpack_16bit(input, output, count);
1074            }
1075            return;
1076        }
1077        if sse::is_available() {
1078            unsafe {
1079                sse::unpack_16bit(input, output, count);
1080            }
1081            return;
1082        }
1083    }
1084
1085    scalar::unpack_16bit(input, output, count);
1086}
1087
1088/// Unpack 32-bit packed values to u32 with SIMD acceleration
1089#[inline]
1090pub fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
1091    #[cfg(target_arch = "aarch64")]
1092    {
1093        if neon::is_available() {
1094            unsafe {
1095                neon::unpack_32bit(input, output, count);
1096            }
1097            return;
1098        }
1099    }
1100
1101    #[cfg(target_arch = "x86_64")]
1102    {
1103        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1104        if avx2::is_available() {
1105            unsafe {
1106                avx2::unpack_32bit(input, output, count);
1107            }
1108            return;
1109        }
1110        if sse::is_available() {
1111            unsafe {
1112                sse::unpack_32bit(input, output, count);
1113            }
1114            return;
1115        }
1116    }
1117
1118    scalar::unpack_32bit(input, output, count);
1119}
1120
1121/// Delta decode with SIMD acceleration
1122///
1123/// Converts delta-encoded values to absolute values.
1124/// Input: deltas[i] = value[i+1] - value[i] - 1 (gap minus one)
1125/// Output: absolute values starting from first_value
1126#[inline]
1127pub fn delta_decode(output: &mut [u32], deltas: &[u32], first_value: u32, count: usize) {
1128    #[cfg(target_arch = "aarch64")]
1129    {
1130        if neon::is_available() {
1131            unsafe {
1132                neon::delta_decode(output, deltas, first_value, count);
1133            }
1134            return;
1135        }
1136    }
1137
1138    #[cfg(target_arch = "x86_64")]
1139    {
1140        if sse::is_available() {
1141            unsafe {
1142                sse::delta_decode(output, deltas, first_value, count);
1143            }
1144            return;
1145        }
1146    }
1147
1148    scalar::delta_decode(output, deltas, first_value, count);
1149}
1150
1151/// Add 1 to all values with SIMD acceleration
1152///
1153/// Used for TF decoding where values are stored as (tf - 1)
1154#[inline]
1155pub fn add_one(values: &mut [u32], count: usize) {
1156    #[cfg(target_arch = "aarch64")]
1157    {
1158        if neon::is_available() {
1159            unsafe {
1160                neon::add_one(values, count);
1161            }
1162            return;
1163        }
1164    }
1165
1166    #[cfg(target_arch = "x86_64")]
1167    {
1168        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1169        if avx2::is_available() {
1170            unsafe {
1171                avx2::add_one(values, count);
1172            }
1173            return;
1174        }
1175        if sse::is_available() {
1176            unsafe {
1177                sse::add_one(values, count);
1178            }
1179            return;
1180        }
1181    }
1182
1183    scalar::add_one(values, count);
1184}
1185
1186/// Compute the number of bits needed to represent a value
1187#[inline]
1188pub fn bits_needed(val: u32) -> u8 {
1189    if val == 0 {
1190        0
1191    } else {
1192        32 - val.leading_zeros() as u8
1193    }
1194}
1195
1196// ============================================================================
1197// Rounded bitpacking for truly vectorized encoding/decoding
1198// ============================================================================
1199//
1200// Instead of using arbitrary bit widths (1-32), we round up to SIMD-friendly
1201// widths: 0, 8, 16, or 32 bits. This trades ~10-20% more space for much faster
1202// decoding since we can use direct SIMD widening instructions (pmovzx) without
1203// any bit-shifting or masking.
1204//
1205// Bit width mapping:
1206//   0      -> 0  (all zeros)
1207//   1-8    -> 8  (u8)
1208//   9-16   -> 16 (u16)
1209//   17-32  -> 32 (u32)
1210
1211/// Rounded bit width type for SIMD-friendly encoding
1212#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1213#[repr(u8)]
1214pub enum RoundedBitWidth {
1215    Zero = 0,
1216    Bits8 = 8,
1217    Bits16 = 16,
1218    Bits32 = 32,
1219}
1220
1221impl RoundedBitWidth {
1222    /// Round an exact bit width to the nearest SIMD-friendly width
1223    #[inline]
1224    pub fn from_exact(bits: u8) -> Self {
1225        match bits {
1226            0 => RoundedBitWidth::Zero,
1227            1..=8 => RoundedBitWidth::Bits8,
1228            9..=16 => RoundedBitWidth::Bits16,
1229            _ => RoundedBitWidth::Bits32,
1230        }
1231    }
1232
1233    /// Convert from stored u8 value (must be 0, 8, 16, or 32)
1234    #[inline]
1235    pub fn from_u8(bits: u8) -> Self {
1236        match bits {
1237            0 => RoundedBitWidth::Zero,
1238            8 => RoundedBitWidth::Bits8,
1239            16 => RoundedBitWidth::Bits16,
1240            32 => RoundedBitWidth::Bits32,
1241            _ => RoundedBitWidth::Bits32, // Fallback for invalid values
1242        }
1243    }
1244
1245    /// Get the byte size per value
1246    #[inline]
1247    pub fn bytes_per_value(self) -> usize {
1248        match self {
1249            RoundedBitWidth::Zero => 0,
1250            RoundedBitWidth::Bits8 => 1,
1251            RoundedBitWidth::Bits16 => 2,
1252            RoundedBitWidth::Bits32 => 4,
1253        }
1254    }
1255
1256    /// Get the raw bit width value
1257    #[inline]
1258    pub fn as_u8(self) -> u8 {
1259        self as u8
1260    }
1261}
1262
1263/// Round a bit width to the nearest SIMD-friendly width (0, 8, 16, or 32)
1264#[inline]
1265pub fn round_bit_width(bits: u8) -> u8 {
1266    RoundedBitWidth::from_exact(bits).as_u8()
1267}
1268
1269/// Pack values using rounded bit width (SIMD-friendly)
1270///
1271/// This is much simpler than arbitrary bitpacking since values are byte-aligned.
1272/// Returns the number of bytes written.
1273#[inline]
1274pub fn pack_rounded(values: &[u32], bit_width: RoundedBitWidth, output: &mut [u8]) -> usize {
1275    let count = values.len();
1276    match bit_width {
1277        RoundedBitWidth::Zero => 0,
1278        RoundedBitWidth::Bits8 => {
1279            for (i, &v) in values.iter().enumerate() {
1280                output[i] = v as u8;
1281            }
1282            count
1283        }
1284        RoundedBitWidth::Bits16 => {
1285            for (i, &v) in values.iter().enumerate() {
1286                let bytes = (v as u16).to_le_bytes();
1287                output[i * 2] = bytes[0];
1288                output[i * 2 + 1] = bytes[1];
1289            }
1290            count * 2
1291        }
1292        RoundedBitWidth::Bits32 => {
1293            for (i, &v) in values.iter().enumerate() {
1294                let bytes = v.to_le_bytes();
1295                output[i * 4] = bytes[0];
1296                output[i * 4 + 1] = bytes[1];
1297                output[i * 4 + 2] = bytes[2];
1298                output[i * 4 + 3] = bytes[3];
1299            }
1300            count * 4
1301        }
1302    }
1303}
1304
1305/// Unpack values using rounded bit width with SIMD acceleration
1306///
1307/// This is the fast path - no bit manipulation needed, just widening.
1308#[inline]
1309pub fn unpack_rounded(input: &[u8], bit_width: RoundedBitWidth, output: &mut [u32], count: usize) {
1310    match bit_width {
1311        RoundedBitWidth::Zero => {
1312            for out in output.iter_mut().take(count) {
1313                *out = 0;
1314            }
1315        }
1316        RoundedBitWidth::Bits8 => unpack_8bit(input, output, count),
1317        RoundedBitWidth::Bits16 => unpack_16bit(input, output, count),
1318        RoundedBitWidth::Bits32 => unpack_32bit(input, output, count),
1319    }
1320}
1321
1322/// Fused unpack + delta decode using rounded bit width
1323///
1324/// Combines unpacking and prefix sum in a single pass for better cache utilization.
1325#[inline]
1326pub fn unpack_rounded_delta_decode(
1327    input: &[u8],
1328    bit_width: RoundedBitWidth,
1329    output: &mut [u32],
1330    first_value: u32,
1331    count: usize,
1332) {
1333    match bit_width {
1334        RoundedBitWidth::Zero => {
1335            // All deltas are 0, meaning gaps of 1
1336            let mut val = first_value;
1337            for out in output.iter_mut().take(count) {
1338                *out = val;
1339                val = val.wrapping_add(1);
1340            }
1341        }
1342        RoundedBitWidth::Bits8 => unpack_8bit_delta_decode(input, output, first_value, count),
1343        RoundedBitWidth::Bits16 => unpack_16bit_delta_decode(input, output, first_value, count),
1344        RoundedBitWidth::Bits32 => {
1345            // Unpack count-1 deltas from input, then prefix sum to absolute values
1346            if count > 0 {
1347                output[0] = first_value;
1348                let mut carry = first_value;
1349                for i in 0..count - 1 {
1350                    let idx = i * 4;
1351                    let delta = u32::from_le_bytes([
1352                        input[idx],
1353                        input[idx + 1],
1354                        input[idx + 2],
1355                        input[idx + 3],
1356                    ]);
1357                    carry = carry.wrapping_add(delta).wrapping_add(1);
1358                    output[i + 1] = carry;
1359                }
1360            }
1361        }
1362    }
1363}
1364
1365// ============================================================================
1366// Fused operations for better cache utilization
1367// ============================================================================
1368
1369/// Fused unpack 8-bit + delta decode in a single pass
1370///
1371/// This avoids writing the intermediate unpacked values to memory,
1372/// improving cache utilization for large blocks.
1373#[inline]
1374pub fn unpack_8bit_delta_decode(input: &[u8], output: &mut [u32], first_value: u32, count: usize) {
1375    if count == 0 {
1376        return;
1377    }
1378
1379    output[0] = first_value;
1380    if count == 1 {
1381        return;
1382    }
1383
1384    #[cfg(target_arch = "aarch64")]
1385    {
1386        if neon::is_available() {
1387            unsafe {
1388                neon::unpack_8bit_delta_decode(input, output, first_value, count);
1389            }
1390            return;
1391        }
1392    }
1393
1394    #[cfg(target_arch = "x86_64")]
1395    {
1396        if avx2::is_available() {
1397            unsafe {
1398                avx2::unpack_8bit_delta_decode(input, output, first_value, count);
1399            }
1400            return;
1401        }
1402        if sse::is_available() {
1403            unsafe {
1404                sse::unpack_8bit_delta_decode(input, output, first_value, count);
1405            }
1406            return;
1407        }
1408    }
1409
1410    // Scalar fallback
1411    let mut carry = first_value;
1412    for i in 0..count - 1 {
1413        carry = carry.wrapping_add(input[i] as u32).wrapping_add(1);
1414        output[i + 1] = carry;
1415    }
1416}
1417
1418/// Fused unpack 16-bit + delta decode in a single pass
1419#[inline]
1420pub fn unpack_16bit_delta_decode(input: &[u8], output: &mut [u32], first_value: u32, count: usize) {
1421    if count == 0 {
1422        return;
1423    }
1424
1425    output[0] = first_value;
1426    if count == 1 {
1427        return;
1428    }
1429
1430    #[cfg(target_arch = "aarch64")]
1431    {
1432        if neon::is_available() {
1433            unsafe {
1434                neon::unpack_16bit_delta_decode(input, output, first_value, count);
1435            }
1436            return;
1437        }
1438    }
1439
1440    #[cfg(target_arch = "x86_64")]
1441    {
1442        if avx2::is_available() {
1443            unsafe {
1444                avx2::unpack_16bit_delta_decode(input, output, first_value, count);
1445            }
1446            return;
1447        }
1448        if sse::is_available() {
1449            unsafe {
1450                sse::unpack_16bit_delta_decode(input, output, first_value, count);
1451            }
1452            return;
1453        }
1454    }
1455
1456    // Scalar fallback
1457    let mut carry = first_value;
1458    for i in 0..count - 1 {
1459        let idx = i * 2;
1460        let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
1461        carry = carry.wrapping_add(delta).wrapping_add(1);
1462        output[i + 1] = carry;
1463    }
1464}
1465
1466/// Fused unpack + delta decode for arbitrary bit widths
1467///
1468/// Combines unpacking and prefix sum in a single pass, avoiding intermediate buffer.
1469/// Uses SIMD-accelerated paths for 8/16-bit widths, scalar for others.
1470#[inline]
1471pub fn unpack_delta_decode(
1472    input: &[u8],
1473    bit_width: u8,
1474    output: &mut [u32],
1475    first_value: u32,
1476    count: usize,
1477) {
1478    if count == 0 {
1479        return;
1480    }
1481
1482    output[0] = first_value;
1483    if count == 1 {
1484        return;
1485    }
1486
1487    // Fast paths for SIMD-friendly bit widths
1488    match bit_width {
1489        0 => {
1490            // All zeros = consecutive doc IDs (gap of 1)
1491            let mut val = first_value;
1492            for item in output.iter_mut().take(count).skip(1) {
1493                val = val.wrapping_add(1);
1494                *item = val;
1495            }
1496        }
1497        8 => unpack_8bit_delta_decode(input, output, first_value, count),
1498        16 => unpack_16bit_delta_decode(input, output, first_value, count),
1499        32 => {
1500            // 32-bit: unpack inline and delta decode
1501            let mut carry = first_value;
1502            for i in 0..count - 1 {
1503                let idx = i * 4;
1504                let delta = u32::from_le_bytes([
1505                    input[idx],
1506                    input[idx + 1],
1507                    input[idx + 2],
1508                    input[idx + 3],
1509                ]);
1510                carry = carry.wrapping_add(delta).wrapping_add(1);
1511                output[i + 1] = carry;
1512            }
1513        }
1514        _ => {
1515            // Generic bit width: fused unpack + delta decode
1516            let mask = (1u64 << bit_width) - 1;
1517            let bit_width_usize = bit_width as usize;
1518            let mut bit_pos = 0usize;
1519            let input_ptr = input.as_ptr();
1520            let mut carry = first_value;
1521
1522            for i in 0..count - 1 {
1523                let byte_idx = bit_pos >> 3;
1524                let bit_offset = bit_pos & 7;
1525
1526                // SAFETY: Caller guarantees input has enough data
1527                let word = unsafe { (input_ptr.add(byte_idx) as *const u64).read_unaligned() };
1528                let delta = ((word >> bit_offset) & mask) as u32;
1529
1530                carry = carry.wrapping_add(delta).wrapping_add(1);
1531                output[i + 1] = carry;
1532                bit_pos += bit_width_usize;
1533            }
1534        }
1535    }
1536}
1537
1538// ============================================================================
1539// Sparse Vector SIMD Functions
1540// ============================================================================
1541
1542/// Dequantize UInt8 weights to f32 with SIMD acceleration
1543///
1544/// Computes: output[i] = input[i] as f32 * scale + min_val
1545#[inline]
1546pub fn dequantize_uint8(input: &[u8], output: &mut [f32], scale: f32, min_val: f32, count: usize) {
1547    #[cfg(target_arch = "aarch64")]
1548    {
1549        if neon::is_available() {
1550            unsafe {
1551                dequantize_uint8_neon(input, output, scale, min_val, count);
1552            }
1553            return;
1554        }
1555    }
1556
1557    #[cfg(target_arch = "x86_64")]
1558    {
1559        if sse::is_available() {
1560            unsafe {
1561                dequantize_uint8_sse(input, output, scale, min_val, count);
1562            }
1563            return;
1564        }
1565    }
1566
1567    // Scalar fallback
1568    for i in 0..count {
1569        output[i] = input[i] as f32 * scale + min_val;
1570    }
1571}
1572
1573#[cfg(target_arch = "aarch64")]
1574#[target_feature(enable = "neon")]
1575#[allow(unsafe_op_in_unsafe_fn)]
1576unsafe fn dequantize_uint8_neon(
1577    input: &[u8],
1578    output: &mut [f32],
1579    scale: f32,
1580    min_val: f32,
1581    count: usize,
1582) {
1583    use std::arch::aarch64::*;
1584
1585    let scale_v = vdupq_n_f32(scale);
1586    let min_v = vdupq_n_f32(min_val);
1587
1588    let chunks = count / 16;
1589    let remainder = count % 16;
1590
1591    for chunk in 0..chunks {
1592        let base = chunk * 16;
1593        let in_ptr = input.as_ptr().add(base);
1594
1595        // Load 16 bytes
1596        let bytes = vld1q_u8(in_ptr);
1597
1598        // Widen u8 -> u16 -> u32 -> f32
1599        let low8 = vget_low_u8(bytes);
1600        let high8 = vget_high_u8(bytes);
1601
1602        let low16 = vmovl_u8(low8);
1603        let high16 = vmovl_u8(high8);
1604
1605        // Process 4 values at a time
1606        let u32_0 = vmovl_u16(vget_low_u16(low16));
1607        let u32_1 = vmovl_u16(vget_high_u16(low16));
1608        let u32_2 = vmovl_u16(vget_low_u16(high16));
1609        let u32_3 = vmovl_u16(vget_high_u16(high16));
1610
1611        // Convert to f32 and apply scale + min_val
1612        let f32_0 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_0), scale_v);
1613        let f32_1 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_1), scale_v);
1614        let f32_2 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_2), scale_v);
1615        let f32_3 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_3), scale_v);
1616
1617        let out_ptr = output.as_mut_ptr().add(base);
1618        vst1q_f32(out_ptr, f32_0);
1619        vst1q_f32(out_ptr.add(4), f32_1);
1620        vst1q_f32(out_ptr.add(8), f32_2);
1621        vst1q_f32(out_ptr.add(12), f32_3);
1622    }
1623
1624    // Handle remainder
1625    let base = chunks * 16;
1626    for i in 0..remainder {
1627        output[base + i] = input[base + i] as f32 * scale + min_val;
1628    }
1629}
1630
1631#[cfg(target_arch = "x86_64")]
1632#[target_feature(enable = "sse2", enable = "sse4.1")]
1633#[allow(unsafe_op_in_unsafe_fn)]
1634unsafe fn dequantize_uint8_sse(
1635    input: &[u8],
1636    output: &mut [f32],
1637    scale: f32,
1638    min_val: f32,
1639    count: usize,
1640) {
1641    use std::arch::x86_64::*;
1642
1643    let scale_v = _mm_set1_ps(scale);
1644    let min_v = _mm_set1_ps(min_val);
1645
1646    let chunks = count / 4;
1647    let remainder = count % 4;
1648
1649    for chunk in 0..chunks {
1650        let base = chunk * 4;
1651
1652        // Load 4 bytes as a single i32 and zero-extend u8→u32 via SSE4.1
1653        let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
1654            input.as_ptr().add(base) as *const i32
1655        ));
1656        let ints = _mm_cvtepu8_epi32(bytes);
1657        let floats = _mm_cvtepi32_ps(ints);
1658
1659        // Apply scale and min_val: result = floats * scale + min_val
1660        let scaled = _mm_add_ps(_mm_mul_ps(floats, scale_v), min_v);
1661
1662        _mm_storeu_ps(output.as_mut_ptr().add(base), scaled);
1663    }
1664
1665    // Handle remainder
1666    let base = chunks * 4;
1667    for i in 0..remainder {
1668        output[base + i] = input[base + i] as f32 * scale + min_val;
1669    }
1670}
1671
1672/// Compute dot product of two f32 arrays with SIMD acceleration
1673#[inline]
1674pub fn dot_product_f32(a: &[f32], b: &[f32], count: usize) -> f32 {
1675    #[cfg(target_arch = "aarch64")]
1676    {
1677        if neon::is_available() {
1678            return unsafe { dot_product_f32_neon(a, b, count) };
1679        }
1680    }
1681
1682    #[cfg(target_arch = "x86_64")]
1683    {
1684        if is_x86_feature_detected!("avx512f") {
1685            return unsafe { dot_product_f32_avx512(a, b, count) };
1686        }
1687        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
1688            return unsafe { dot_product_f32_avx2(a, b, count) };
1689        }
1690        if sse::is_available() {
1691            return unsafe { dot_product_f32_sse(a, b, count) };
1692        }
1693    }
1694
1695    // Scalar fallback
1696    let mut sum = 0.0f32;
1697    for i in 0..count {
1698        sum += a[i] * b[i];
1699    }
1700    sum
1701}
1702
1703#[cfg(target_arch = "aarch64")]
1704#[target_feature(enable = "neon")]
1705#[allow(unsafe_op_in_unsafe_fn)]
1706unsafe fn dot_product_f32_neon(a: &[f32], b: &[f32], count: usize) -> f32 {
1707    use std::arch::aarch64::*;
1708
1709    let chunks16 = count / 16;
1710    let remainder = count % 16;
1711
1712    let mut acc0 = vdupq_n_f32(0.0);
1713    let mut acc1 = vdupq_n_f32(0.0);
1714    let mut acc2 = vdupq_n_f32(0.0);
1715    let mut acc3 = vdupq_n_f32(0.0);
1716
1717    for c in 0..chunks16 {
1718        let base = c * 16;
1719        acc0 = vfmaq_f32(
1720            acc0,
1721            vld1q_f32(a.as_ptr().add(base)),
1722            vld1q_f32(b.as_ptr().add(base)),
1723        );
1724        acc1 = vfmaq_f32(
1725            acc1,
1726            vld1q_f32(a.as_ptr().add(base + 4)),
1727            vld1q_f32(b.as_ptr().add(base + 4)),
1728        );
1729        acc2 = vfmaq_f32(
1730            acc2,
1731            vld1q_f32(a.as_ptr().add(base + 8)),
1732            vld1q_f32(b.as_ptr().add(base + 8)),
1733        );
1734        acc3 = vfmaq_f32(
1735            acc3,
1736            vld1q_f32(a.as_ptr().add(base + 12)),
1737            vld1q_f32(b.as_ptr().add(base + 12)),
1738        );
1739    }
1740
1741    let acc = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
1742    let mut sum = vaddvq_f32(acc);
1743
1744    let base = chunks16 * 16;
1745    for i in 0..remainder {
1746        sum += a[base + i] * b[base + i];
1747    }
1748
1749    sum
1750}
1751
1752#[cfg(target_arch = "x86_64")]
1753#[target_feature(enable = "avx2", enable = "fma")]
1754#[allow(unsafe_op_in_unsafe_fn)]
1755unsafe fn dot_product_f32_avx2(a: &[f32], b: &[f32], count: usize) -> f32 {
1756    use std::arch::x86_64::*;
1757
1758    let chunks32 = count / 32;
1759    let remainder = count % 32;
1760
1761    let mut acc0 = _mm256_setzero_ps();
1762    let mut acc1 = _mm256_setzero_ps();
1763    let mut acc2 = _mm256_setzero_ps();
1764    let mut acc3 = _mm256_setzero_ps();
1765
1766    for c in 0..chunks32 {
1767        let base = c * 32;
1768        acc0 = _mm256_fmadd_ps(
1769            _mm256_loadu_ps(a.as_ptr().add(base)),
1770            _mm256_loadu_ps(b.as_ptr().add(base)),
1771            acc0,
1772        );
1773        acc1 = _mm256_fmadd_ps(
1774            _mm256_loadu_ps(a.as_ptr().add(base + 8)),
1775            _mm256_loadu_ps(b.as_ptr().add(base + 8)),
1776            acc1,
1777        );
1778        acc2 = _mm256_fmadd_ps(
1779            _mm256_loadu_ps(a.as_ptr().add(base + 16)),
1780            _mm256_loadu_ps(b.as_ptr().add(base + 16)),
1781            acc2,
1782        );
1783        acc3 = _mm256_fmadd_ps(
1784            _mm256_loadu_ps(a.as_ptr().add(base + 24)),
1785            _mm256_loadu_ps(b.as_ptr().add(base + 24)),
1786            acc3,
1787        );
1788    }
1789
1790    let acc = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
1791
1792    // Horizontal sum: 256-bit → 128-bit → scalar
1793    let hi = _mm256_extractf128_ps(acc, 1);
1794    let lo = _mm256_castps256_ps128(acc);
1795    let sum128 = _mm_add_ps(lo, hi);
1796    let shuf = _mm_shuffle_ps(sum128, sum128, 0b10_11_00_01);
1797    let sums = _mm_add_ps(sum128, shuf);
1798    let shuf2 = _mm_movehl_ps(sums, sums);
1799    let final_sum = _mm_add_ss(sums, shuf2);
1800
1801    let mut sum = _mm_cvtss_f32(final_sum);
1802
1803    let base = chunks32 * 32;
1804    for i in 0..remainder {
1805        sum += a[base + i] * b[base + i];
1806    }
1807
1808    sum
1809}
1810
1811#[cfg(target_arch = "x86_64")]
1812#[target_feature(enable = "sse")]
1813#[allow(unsafe_op_in_unsafe_fn)]
1814unsafe fn dot_product_f32_sse(a: &[f32], b: &[f32], count: usize) -> f32 {
1815    use std::arch::x86_64::*;
1816
1817    let chunks = count / 4;
1818    let remainder = count % 4;
1819
1820    let mut acc = _mm_setzero_ps();
1821
1822    for chunk in 0..chunks {
1823        let base = chunk * 4;
1824        let va = _mm_loadu_ps(a.as_ptr().add(base));
1825        let vb = _mm_loadu_ps(b.as_ptr().add(base));
1826        acc = _mm_add_ps(acc, _mm_mul_ps(va, vb));
1827    }
1828
1829    // Horizontal sum: [a, b, c, d] -> a + b + c + d
1830    let shuf = _mm_shuffle_ps(acc, acc, 0b10_11_00_01); // [b, a, d, c]
1831    let sums = _mm_add_ps(acc, shuf); // [a+b, a+b, c+d, c+d]
1832    let shuf2 = _mm_movehl_ps(sums, sums); // [c+d, c+d, ?, ?]
1833    let final_sum = _mm_add_ss(sums, shuf2); // [a+b+c+d, ?, ?, ?]
1834
1835    let mut sum = _mm_cvtss_f32(final_sum);
1836
1837    // Handle remainder
1838    let base = chunks * 4;
1839    for i in 0..remainder {
1840        sum += a[base + i] * b[base + i];
1841    }
1842
1843    sum
1844}
1845
1846#[cfg(target_arch = "x86_64")]
1847#[target_feature(enable = "avx512f")]
1848#[allow(unsafe_op_in_unsafe_fn)]
1849unsafe fn dot_product_f32_avx512(a: &[f32], b: &[f32], count: usize) -> f32 {
1850    use std::arch::x86_64::*;
1851
1852    let chunks64 = count / 64;
1853    let remainder = count % 64;
1854
1855    let mut acc0 = _mm512_setzero_ps();
1856    let mut acc1 = _mm512_setzero_ps();
1857    let mut acc2 = _mm512_setzero_ps();
1858    let mut acc3 = _mm512_setzero_ps();
1859
1860    for c in 0..chunks64 {
1861        let base = c * 64;
1862        acc0 = _mm512_fmadd_ps(
1863            _mm512_loadu_ps(a.as_ptr().add(base)),
1864            _mm512_loadu_ps(b.as_ptr().add(base)),
1865            acc0,
1866        );
1867        acc1 = _mm512_fmadd_ps(
1868            _mm512_loadu_ps(a.as_ptr().add(base + 16)),
1869            _mm512_loadu_ps(b.as_ptr().add(base + 16)),
1870            acc1,
1871        );
1872        acc2 = _mm512_fmadd_ps(
1873            _mm512_loadu_ps(a.as_ptr().add(base + 32)),
1874            _mm512_loadu_ps(b.as_ptr().add(base + 32)),
1875            acc2,
1876        );
1877        acc3 = _mm512_fmadd_ps(
1878            _mm512_loadu_ps(a.as_ptr().add(base + 48)),
1879            _mm512_loadu_ps(b.as_ptr().add(base + 48)),
1880            acc3,
1881        );
1882    }
1883
1884    let acc = _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3));
1885    let mut sum = _mm512_reduce_add_ps(acc);
1886
1887    let base = chunks64 * 64;
1888    for i in 0..remainder {
1889        sum += a[base + i] * b[base + i];
1890    }
1891
1892    sum
1893}
1894
1895#[cfg(target_arch = "x86_64")]
1896#[target_feature(enable = "avx512f")]
1897#[allow(unsafe_op_in_unsafe_fn)]
1898unsafe fn fused_dot_norm_avx512(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1899    use std::arch::x86_64::*;
1900
1901    let chunks64 = count / 64;
1902    let remainder = count % 64;
1903
1904    let mut d0 = _mm512_setzero_ps();
1905    let mut d1 = _mm512_setzero_ps();
1906    let mut d2 = _mm512_setzero_ps();
1907    let mut d3 = _mm512_setzero_ps();
1908    let mut n0 = _mm512_setzero_ps();
1909    let mut n1 = _mm512_setzero_ps();
1910    let mut n2 = _mm512_setzero_ps();
1911    let mut n3 = _mm512_setzero_ps();
1912
1913    for c in 0..chunks64 {
1914        let base = c * 64;
1915        let vb0 = _mm512_loadu_ps(b.as_ptr().add(base));
1916        d0 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base)), vb0, d0);
1917        n0 = _mm512_fmadd_ps(vb0, vb0, n0);
1918        let vb1 = _mm512_loadu_ps(b.as_ptr().add(base + 16));
1919        d1 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 16)), vb1, d1);
1920        n1 = _mm512_fmadd_ps(vb1, vb1, n1);
1921        let vb2 = _mm512_loadu_ps(b.as_ptr().add(base + 32));
1922        d2 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 32)), vb2, d2);
1923        n2 = _mm512_fmadd_ps(vb2, vb2, n2);
1924        let vb3 = _mm512_loadu_ps(b.as_ptr().add(base + 48));
1925        d3 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 48)), vb3, d3);
1926        n3 = _mm512_fmadd_ps(vb3, vb3, n3);
1927    }
1928
1929    let acc_dot = _mm512_add_ps(_mm512_add_ps(d0, d1), _mm512_add_ps(d2, d3));
1930    let acc_norm = _mm512_add_ps(_mm512_add_ps(n0, n1), _mm512_add_ps(n2, n3));
1931    let mut dot = _mm512_reduce_add_ps(acc_dot);
1932    let mut norm = _mm512_reduce_add_ps(acc_norm);
1933
1934    let base = chunks64 * 64;
1935    for i in 0..remainder {
1936        dot += a[base + i] * b[base + i];
1937        norm += b[base + i] * b[base + i];
1938    }
1939
1940    (dot, norm)
1941}
1942
1943// ============================================================================
1944// Batched Cosine Similarity for Dense Vector Search
1945// ============================================================================
1946
1947/// Fused dot-product + self-norm in a single pass (SIMD accelerated).
1948///
1949/// Returns (dot(a, b), dot(b, b)) — i.e. the dot product of a·b and ||b||².
1950/// Loads `b` only once (halves memory bandwidth vs two separate dot products).
1951#[inline]
1952fn fused_dot_norm(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1953    #[cfg(target_arch = "aarch64")]
1954    {
1955        if neon::is_available() {
1956            return unsafe { fused_dot_norm_neon(a, b, count) };
1957        }
1958    }
1959
1960    #[cfg(target_arch = "x86_64")]
1961    {
1962        if is_x86_feature_detected!("avx512f") {
1963            return unsafe { fused_dot_norm_avx512(a, b, count) };
1964        }
1965        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
1966            return unsafe { fused_dot_norm_avx2(a, b, count) };
1967        }
1968        if sse::is_available() {
1969            return unsafe { fused_dot_norm_sse(a, b, count) };
1970        }
1971    }
1972
1973    // Scalar fallback
1974    let mut dot = 0.0f32;
1975    let mut norm_b = 0.0f32;
1976    for i in 0..count {
1977        dot += a[i] * b[i];
1978        norm_b += b[i] * b[i];
1979    }
1980    (dot, norm_b)
1981}
1982
1983#[cfg(target_arch = "aarch64")]
1984#[target_feature(enable = "neon")]
1985#[allow(unsafe_op_in_unsafe_fn)]
1986unsafe fn fused_dot_norm_neon(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1987    use std::arch::aarch64::*;
1988
1989    let chunks16 = count / 16;
1990    let remainder = count % 16;
1991
1992    let mut d0 = vdupq_n_f32(0.0);
1993    let mut d1 = vdupq_n_f32(0.0);
1994    let mut d2 = vdupq_n_f32(0.0);
1995    let mut d3 = vdupq_n_f32(0.0);
1996    let mut n0 = vdupq_n_f32(0.0);
1997    let mut n1 = vdupq_n_f32(0.0);
1998    let mut n2 = vdupq_n_f32(0.0);
1999    let mut n3 = vdupq_n_f32(0.0);
2000
2001    for c in 0..chunks16 {
2002        let base = c * 16;
2003        let va0 = vld1q_f32(a.as_ptr().add(base));
2004        let vb0 = vld1q_f32(b.as_ptr().add(base));
2005        d0 = vfmaq_f32(d0, va0, vb0);
2006        n0 = vfmaq_f32(n0, vb0, vb0);
2007        let va1 = vld1q_f32(a.as_ptr().add(base + 4));
2008        let vb1 = vld1q_f32(b.as_ptr().add(base + 4));
2009        d1 = vfmaq_f32(d1, va1, vb1);
2010        n1 = vfmaq_f32(n1, vb1, vb1);
2011        let va2 = vld1q_f32(a.as_ptr().add(base + 8));
2012        let vb2 = vld1q_f32(b.as_ptr().add(base + 8));
2013        d2 = vfmaq_f32(d2, va2, vb2);
2014        n2 = vfmaq_f32(n2, vb2, vb2);
2015        let va3 = vld1q_f32(a.as_ptr().add(base + 12));
2016        let vb3 = vld1q_f32(b.as_ptr().add(base + 12));
2017        d3 = vfmaq_f32(d3, va3, vb3);
2018        n3 = vfmaq_f32(n3, vb3, vb3);
2019    }
2020
2021    let acc_dot = vaddq_f32(vaddq_f32(d0, d1), vaddq_f32(d2, d3));
2022    let acc_norm = vaddq_f32(vaddq_f32(n0, n1), vaddq_f32(n2, n3));
2023    let mut dot = vaddvq_f32(acc_dot);
2024    let mut norm = vaddvq_f32(acc_norm);
2025
2026    let base = chunks16 * 16;
2027    for i in 0..remainder {
2028        dot += a[base + i] * b[base + i];
2029        norm += b[base + i] * b[base + i];
2030    }
2031
2032    (dot, norm)
2033}
2034
2035#[cfg(target_arch = "x86_64")]
2036#[target_feature(enable = "avx2", enable = "fma")]
2037#[allow(unsafe_op_in_unsafe_fn)]
2038unsafe fn fused_dot_norm_avx2(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
2039    use std::arch::x86_64::*;
2040
2041    let chunks32 = count / 32;
2042    let remainder = count % 32;
2043
2044    let mut d0 = _mm256_setzero_ps();
2045    let mut d1 = _mm256_setzero_ps();
2046    let mut d2 = _mm256_setzero_ps();
2047    let mut d3 = _mm256_setzero_ps();
2048    let mut n0 = _mm256_setzero_ps();
2049    let mut n1 = _mm256_setzero_ps();
2050    let mut n2 = _mm256_setzero_ps();
2051    let mut n3 = _mm256_setzero_ps();
2052
2053    for c in 0..chunks32 {
2054        let base = c * 32;
2055        let vb0 = _mm256_loadu_ps(b.as_ptr().add(base));
2056        d0 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base)), vb0, d0);
2057        n0 = _mm256_fmadd_ps(vb0, vb0, n0);
2058        let vb1 = _mm256_loadu_ps(b.as_ptr().add(base + 8));
2059        d1 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 8)), vb1, d1);
2060        n1 = _mm256_fmadd_ps(vb1, vb1, n1);
2061        let vb2 = _mm256_loadu_ps(b.as_ptr().add(base + 16));
2062        d2 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 16)), vb2, d2);
2063        n2 = _mm256_fmadd_ps(vb2, vb2, n2);
2064        let vb3 = _mm256_loadu_ps(b.as_ptr().add(base + 24));
2065        d3 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 24)), vb3, d3);
2066        n3 = _mm256_fmadd_ps(vb3, vb3, n3);
2067    }
2068
2069    let acc_dot = _mm256_add_ps(_mm256_add_ps(d0, d1), _mm256_add_ps(d2, d3));
2070    let acc_norm = _mm256_add_ps(_mm256_add_ps(n0, n1), _mm256_add_ps(n2, n3));
2071
2072    // Horizontal sums: 256→128→scalar
2073    let hi_d = _mm256_extractf128_ps(acc_dot, 1);
2074    let lo_d = _mm256_castps256_ps128(acc_dot);
2075    let sum_d = _mm_add_ps(lo_d, hi_d);
2076    let shuf_d = _mm_shuffle_ps(sum_d, sum_d, 0b10_11_00_01);
2077    let sums_d = _mm_add_ps(sum_d, shuf_d);
2078    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2079    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2080
2081    let hi_n = _mm256_extractf128_ps(acc_norm, 1);
2082    let lo_n = _mm256_castps256_ps128(acc_norm);
2083    let sum_n = _mm_add_ps(lo_n, hi_n);
2084    let shuf_n = _mm_shuffle_ps(sum_n, sum_n, 0b10_11_00_01);
2085    let sums_n = _mm_add_ps(sum_n, shuf_n);
2086    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2087    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2088
2089    let base = chunks32 * 32;
2090    for i in 0..remainder {
2091        dot += a[base + i] * b[base + i];
2092        norm += b[base + i] * b[base + i];
2093    }
2094
2095    (dot, norm)
2096}
2097
2098#[cfg(target_arch = "x86_64")]
2099#[target_feature(enable = "sse")]
2100#[allow(unsafe_op_in_unsafe_fn)]
2101unsafe fn fused_dot_norm_sse(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
2102    use std::arch::x86_64::*;
2103
2104    let chunks = count / 4;
2105    let remainder = count % 4;
2106
2107    let mut acc_dot = _mm_setzero_ps();
2108    let mut acc_norm = _mm_setzero_ps();
2109
2110    for chunk in 0..chunks {
2111        let base = chunk * 4;
2112        let va = _mm_loadu_ps(a.as_ptr().add(base));
2113        let vb = _mm_loadu_ps(b.as_ptr().add(base));
2114        acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2115        acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2116    }
2117
2118    // Horizontal sums
2119    let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2120    let sums_d = _mm_add_ps(acc_dot, shuf_d);
2121    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2122    let final_d = _mm_add_ss(sums_d, shuf2_d);
2123    let mut dot = _mm_cvtss_f32(final_d);
2124
2125    let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2126    let sums_n = _mm_add_ps(acc_norm, shuf_n);
2127    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2128    let final_n = _mm_add_ss(sums_n, shuf2_n);
2129    let mut norm = _mm_cvtss_f32(final_n);
2130
2131    let base = chunks * 4;
2132    for i in 0..remainder {
2133        dot += a[base + i] * b[base + i];
2134        norm += b[base + i] * b[base + i];
2135    }
2136
2137    (dot, norm)
2138}
2139
2140/// Fast approximate reciprocal square root: 1/sqrt(x).
2141///
2142/// Uses the IEEE 754 bit trick (Quake III) + one Newton-Raphson iteration
2143/// for ~23-bit precision — sufficient for cosine similarity scoring.
2144/// ~3-5x faster than `1.0 / x.sqrt()` on most architectures.
2145#[inline]
2146pub fn fast_inv_sqrt(x: f32) -> f32 {
2147    let half = 0.5 * x;
2148    let i = 0x5F37_5A86_u32.wrapping_sub(x.to_bits() >> 1);
2149    let y = f32::from_bits(i);
2150    let y = y * (1.5 - half * y * y); // first Newton-Raphson step
2151    y * (1.5 - half * y * y) // second step: ~23-bit precision
2152}
2153
2154/// Batch cosine similarity: query vs N contiguous vectors.
2155///
2156/// `vectors` is a contiguous buffer of `n * dim` floats (row-major).
2157/// `scores` must have length >= n.
2158///
2159/// Optimizations over calling `cosine_similarity` N times:
2160/// 1. Query norm computed once (not N times)
2161/// 2. Fused dot+norm kernel — each vector loaded once (halves bandwidth)
2162/// 3. No per-call overhead (branch prediction, function calls)
2163/// 4. Fast reciprocal square root (~3-5x faster than 1/sqrt)
2164#[inline]
2165pub fn batch_cosine_scores(query: &[f32], vectors: &[f32], dim: usize, scores: &mut [f32]) {
2166    let n = scores.len();
2167    debug_assert!(vectors.len() >= n * dim);
2168    debug_assert_eq!(query.len(), dim);
2169
2170    if dim == 0 || n == 0 {
2171        return;
2172    }
2173
2174    // Pre-compute query inverse norm once
2175    let norm_q_sq = dot_product_f32(query, query, dim);
2176    if norm_q_sq < f32::EPSILON {
2177        for s in scores.iter_mut() {
2178            *s = 0.0;
2179        }
2180        return;
2181    }
2182    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2183
2184    for i in 0..n {
2185        let vec = &vectors[i * dim..(i + 1) * dim];
2186        let (dot, norm_v_sq) = fused_dot_norm(query, vec, dim);
2187        if norm_v_sq < f32::EPSILON {
2188            scores[i] = 0.0;
2189        } else {
2190            scores[i] = dot * inv_norm_q * fast_inv_sqrt(norm_v_sq);
2191        }
2192    }
2193}
2194
2195// ============================================================================
2196// f16 (IEEE 754 half-precision) conversion
2197// ============================================================================
2198
2199/// Convert f32 to f16 (IEEE 754 half-precision), stored as u16
2200#[inline]
2201pub fn f32_to_f16(value: f32) -> u16 {
2202    let bits = value.to_bits();
2203    let sign = (bits >> 16) & 0x8000;
2204    let exp = ((bits >> 23) & 0xFF) as i32;
2205    let mantissa = bits & 0x7F_FFFF;
2206
2207    if exp == 255 {
2208        // Inf/NaN
2209        return (sign | 0x7C00 | ((mantissa >> 13) & 0x3FF)) as u16;
2210    }
2211
2212    let exp16 = exp - 127 + 15;
2213
2214    if exp16 >= 31 {
2215        return (sign | 0x7C00) as u16; // overflow → infinity
2216    }
2217
2218    if exp16 <= 0 {
2219        if exp16 < -10 {
2220            return sign as u16; // too small → zero
2221        }
2222        let shift = (1 - exp16) as u32;
2223        let m = (mantissa | 0x80_0000) >> shift;
2224        // Round-to-nearest-even
2225        let round_bit = (m >> 12) & 1;
2226        let sticky = m & 0xFFF;
2227        let m13 = m >> 13;
2228        let rounded = m13 + (round_bit & (m13 | if sticky != 0 { 1 } else { 0 }));
2229        return (sign | rounded) as u16;
2230    }
2231
2232    // Round-to-nearest-even for normal numbers
2233    let round_bit = (mantissa >> 12) & 1;
2234    let sticky = mantissa & 0xFFF;
2235    let m13 = mantissa >> 13;
2236    let rounded = m13 + (round_bit & (m13 | if sticky != 0 { 1 } else { 0 }));
2237    // Check if rounding caused mantissa overflow (carry into exponent)
2238    if rounded > 0x3FF {
2239        let exp16_inc = exp16 as u32 + 1;
2240        if exp16_inc >= 31 {
2241            return (sign | 0x7C00) as u16; // overflow → infinity
2242        }
2243        (sign | (exp16_inc << 10)) as u16
2244    } else {
2245        (sign | ((exp16 as u32) << 10) | rounded) as u16
2246    }
2247}
2248
2249/// Convert f16 (stored as u16) to f32
2250#[inline]
2251pub fn f16_to_f32(half: u16) -> f32 {
2252    let sign = ((half & 0x8000) as u32) << 16;
2253    let exp = ((half >> 10) & 0x1F) as u32;
2254    let mantissa = (half & 0x3FF) as u32;
2255
2256    if exp == 0 {
2257        if mantissa == 0 {
2258            return f32::from_bits(sign);
2259        }
2260        // Subnormal: normalize
2261        let mut e = 0u32;
2262        let mut m = mantissa;
2263        while (m & 0x400) == 0 {
2264            m <<= 1;
2265            e += 1;
2266        }
2267        return f32::from_bits(sign | ((127 - 15 + 1 - e) << 23) | ((m & 0x3FF) << 13));
2268    }
2269
2270    if exp == 31 {
2271        return f32::from_bits(sign | 0x7F80_0000 | (mantissa << 13));
2272    }
2273
2274    f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mantissa << 13))
2275}
2276
2277// ============================================================================
2278// uint8 scalar quantization for [-1, 1] range
2279// ============================================================================
2280
2281const U8_SCALE: f32 = 127.5;
2282const U8_INV_SCALE: f32 = 1.0 / 127.5;
2283
2284/// Quantize f32 in [-1, 1] to u8 [0, 255]
2285#[inline]
2286pub fn f32_to_u8_saturating(value: f32) -> u8 {
2287    ((value.clamp(-1.0, 1.0) + 1.0) * U8_SCALE) as u8
2288}
2289
2290/// Dequantize u8 [0, 255] to f32 in [-1, 1]
2291#[inline]
2292pub fn u8_to_f32(byte: u8) -> f32 {
2293    byte as f32 * U8_INV_SCALE - 1.0
2294}
2295
2296// ============================================================================
2297// Batch conversion (used during builder write)
2298// ============================================================================
2299
2300/// Batch convert f32 slice to f16 (stored as u16)
2301pub fn batch_f32_to_f16(src: &[f32], dst: &mut [u16]) {
2302    debug_assert_eq!(src.len(), dst.len());
2303    for (s, d) in src.iter().zip(dst.iter_mut()) {
2304        *d = f32_to_f16(*s);
2305    }
2306}
2307
2308/// Batch convert f32 slice to u8 with [-1,1] → [0,255] mapping
2309pub fn batch_f32_to_u8(src: &[f32], dst: &mut [u8]) {
2310    debug_assert_eq!(src.len(), dst.len());
2311    for (s, d) in src.iter().zip(dst.iter_mut()) {
2312        *d = f32_to_u8_saturating(*s);
2313    }
2314}
2315
2316// ============================================================================
2317// NEON-accelerated fused dot+norm for quantized vectors
2318// ============================================================================
2319
2320#[cfg(target_arch = "aarch64")]
2321#[allow(unsafe_op_in_unsafe_fn)]
2322mod neon_quant {
2323    use std::arch::aarch64::*;
2324
2325    /// Fused dot(query_f16, vec_f16) + norm(vec_f16) for f16 vectors on NEON.
2326    ///
2327    /// Both query and vectors are f16 (stored as u16). Uses hardware `vcvt_f32_f16`
2328    /// for SIMD f16→f32 conversion (replaces scalar bit manipulation), processes
2329    /// 8 elements per iteration with f32 accumulation for precision.
2330    #[allow(clippy::incompatible_msrv)]
2331    #[target_feature(enable = "neon")]
2332    pub unsafe fn fused_dot_norm_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2333        let chunks16 = dim / 16;
2334        let remainder = dim % 16;
2335
2336        // 2 accumulator pairs to hide FMA latency (processes 16 f16 per iteration)
2337        let mut acc_dot0 = vdupq_n_f32(0.0);
2338        let mut acc_dot1 = vdupq_n_f32(0.0);
2339        let mut acc_norm0 = vdupq_n_f32(0.0);
2340        let mut acc_norm1 = vdupq_n_f32(0.0);
2341
2342        for c in 0..chunks16 {
2343            let base = c * 16;
2344
2345            // First 8 f16 elements
2346            let v_raw0 = vld1q_u16(vec_f16.as_ptr().add(base));
2347            let v_lo0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw0)));
2348            let v_hi0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw0)));
2349            let q_raw0 = vld1q_u16(query_f16.as_ptr().add(base));
2350            let q_lo0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw0)));
2351            let q_hi0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw0)));
2352
2353            acc_dot0 = vfmaq_f32(acc_dot0, q_lo0, v_lo0);
2354            acc_dot0 = vfmaq_f32(acc_dot0, q_hi0, v_hi0);
2355            acc_norm0 = vfmaq_f32(acc_norm0, v_lo0, v_lo0);
2356            acc_norm0 = vfmaq_f32(acc_norm0, v_hi0, v_hi0);
2357
2358            // Second 8 f16 elements (independent accumulator chain)
2359            let v_raw1 = vld1q_u16(vec_f16.as_ptr().add(base + 8));
2360            let v_lo1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw1)));
2361            let v_hi1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw1)));
2362            let q_raw1 = vld1q_u16(query_f16.as_ptr().add(base + 8));
2363            let q_lo1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw1)));
2364            let q_hi1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw1)));
2365
2366            acc_dot1 = vfmaq_f32(acc_dot1, q_lo1, v_lo1);
2367            acc_dot1 = vfmaq_f32(acc_dot1, q_hi1, v_hi1);
2368            acc_norm1 = vfmaq_f32(acc_norm1, v_lo1, v_lo1);
2369            acc_norm1 = vfmaq_f32(acc_norm1, v_hi1, v_hi1);
2370        }
2371
2372        // Combine accumulator pairs
2373        let mut dot = vaddvq_f32(vaddq_f32(acc_dot0, acc_dot1));
2374        let mut norm = vaddvq_f32(vaddq_f32(acc_norm0, acc_norm1));
2375
2376        // Handle remainder
2377        let base = chunks16 * 16;
2378        for i in 0..remainder {
2379            let v = super::f16_to_f32(*vec_f16.get_unchecked(base + i));
2380            let q = super::f16_to_f32(*query_f16.get_unchecked(base + i));
2381            dot += q * v;
2382            norm += v * v;
2383        }
2384
2385        (dot, norm)
2386    }
2387
2388    /// Fused dot(query, vec) + norm(vec) for u8 vectors on NEON.
2389    /// Processes 16 u8 values per iteration using NEON widening chain.
2390    #[target_feature(enable = "neon")]
2391    pub unsafe fn fused_dot_norm_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2392        let scale = vdupq_n_f32(super::U8_INV_SCALE);
2393        let offset = vdupq_n_f32(-1.0);
2394
2395        let chunks16 = dim / 16;
2396        let remainder = dim % 16;
2397
2398        let mut acc_dot = vdupq_n_f32(0.0);
2399        let mut acc_norm = vdupq_n_f32(0.0);
2400
2401        for c in 0..chunks16 {
2402            let base = c * 16;
2403
2404            // Load 16 u8 values
2405            let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2406
2407            // Widen: 16×u8 → 2×8×u16 → 4×4×u32 → 4×4×f32
2408            let lo8 = vget_low_u8(bytes);
2409            let hi8 = vget_high_u8(bytes);
2410            let lo16 = vmovl_u8(lo8);
2411            let hi16 = vmovl_u8(hi8);
2412
2413            let f0 = vaddq_f32(
2414                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo16))), scale),
2415                offset,
2416            );
2417            let f1 = vaddq_f32(
2418                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo16))), scale),
2419                offset,
2420            );
2421            let f2 = vaddq_f32(
2422                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi16))), scale),
2423                offset,
2424            );
2425            let f3 = vaddq_f32(
2426                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi16))), scale),
2427                offset,
2428            );
2429
2430            let q0 = vld1q_f32(query.as_ptr().add(base));
2431            let q1 = vld1q_f32(query.as_ptr().add(base + 4));
2432            let q2 = vld1q_f32(query.as_ptr().add(base + 8));
2433            let q3 = vld1q_f32(query.as_ptr().add(base + 12));
2434
2435            acc_dot = vfmaq_f32(acc_dot, q0, f0);
2436            acc_dot = vfmaq_f32(acc_dot, q1, f1);
2437            acc_dot = vfmaq_f32(acc_dot, q2, f2);
2438            acc_dot = vfmaq_f32(acc_dot, q3, f3);
2439
2440            acc_norm = vfmaq_f32(acc_norm, f0, f0);
2441            acc_norm = vfmaq_f32(acc_norm, f1, f1);
2442            acc_norm = vfmaq_f32(acc_norm, f2, f2);
2443            acc_norm = vfmaq_f32(acc_norm, f3, f3);
2444        }
2445
2446        let mut dot = vaddvq_f32(acc_dot);
2447        let mut norm = vaddvq_f32(acc_norm);
2448
2449        let base = chunks16 * 16;
2450        for i in 0..remainder {
2451            let v = super::u8_to_f32(*vec_u8.get_unchecked(base + i));
2452            dot += *query.get_unchecked(base + i) * v;
2453            norm += v * v;
2454        }
2455
2456        (dot, norm)
2457    }
2458
2459    /// Dot product only for f16 vectors on NEON (no norm — for unit_norm vectors).
2460    #[allow(clippy::incompatible_msrv)]
2461    #[target_feature(enable = "neon")]
2462    pub unsafe fn dot_product_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2463        let chunks8 = dim / 8;
2464        let remainder = dim % 8;
2465
2466        let mut acc = vdupq_n_f32(0.0);
2467
2468        for c in 0..chunks8 {
2469            let base = c * 8;
2470            let v_raw = vld1q_u16(vec_f16.as_ptr().add(base));
2471            let v_lo = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw)));
2472            let v_hi = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw)));
2473            let q_raw = vld1q_u16(query_f16.as_ptr().add(base));
2474            let q_lo = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw)));
2475            let q_hi = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw)));
2476            acc = vfmaq_f32(acc, q_lo, v_lo);
2477            acc = vfmaq_f32(acc, q_hi, v_hi);
2478        }
2479
2480        let mut dot = vaddvq_f32(acc);
2481        let base = chunks8 * 8;
2482        for i in 0..remainder {
2483            let v = super::f16_to_f32(*vec_f16.get_unchecked(base + i));
2484            let q = super::f16_to_f32(*query_f16.get_unchecked(base + i));
2485            dot += q * v;
2486        }
2487        dot
2488    }
2489
2490    /// Dot product only for u8 vectors on NEON (no norm — for unit_norm vectors).
2491    #[target_feature(enable = "neon")]
2492    pub unsafe fn dot_product_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2493        let scale = vdupq_n_f32(super::U8_INV_SCALE);
2494        let offset = vdupq_n_f32(-1.0);
2495        let chunks16 = dim / 16;
2496        let remainder = dim % 16;
2497
2498        let mut acc = vdupq_n_f32(0.0);
2499
2500        for c in 0..chunks16 {
2501            let base = c * 16;
2502            let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2503            let lo8 = vget_low_u8(bytes);
2504            let hi8 = vget_high_u8(bytes);
2505            let lo16 = vmovl_u8(lo8);
2506            let hi16 = vmovl_u8(hi8);
2507            let f0 = vaddq_f32(
2508                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo16))), scale),
2509                offset,
2510            );
2511            let f1 = vaddq_f32(
2512                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo16))), scale),
2513                offset,
2514            );
2515            let f2 = vaddq_f32(
2516                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi16))), scale),
2517                offset,
2518            );
2519            let f3 = vaddq_f32(
2520                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi16))), scale),
2521                offset,
2522            );
2523            let q0 = vld1q_f32(query.as_ptr().add(base));
2524            let q1 = vld1q_f32(query.as_ptr().add(base + 4));
2525            let q2 = vld1q_f32(query.as_ptr().add(base + 8));
2526            let q3 = vld1q_f32(query.as_ptr().add(base + 12));
2527            acc = vfmaq_f32(acc, q0, f0);
2528            acc = vfmaq_f32(acc, q1, f1);
2529            acc = vfmaq_f32(acc, q2, f2);
2530            acc = vfmaq_f32(acc, q3, f3);
2531        }
2532
2533        let mut dot = vaddvq_f32(acc);
2534        let base = chunks16 * 16;
2535        for i in 0..remainder {
2536            let v = super::u8_to_f32(*vec_u8.get_unchecked(base + i));
2537            dot += *query.get_unchecked(base + i) * v;
2538        }
2539        dot
2540    }
2541}
2542
2543// ============================================================================
2544// Scalar fallback for fused dot+norm on quantized vectors
2545// ============================================================================
2546
2547#[allow(dead_code)]
2548fn fused_dot_norm_f16_scalar(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2549    let mut dot = 0.0f32;
2550    let mut norm = 0.0f32;
2551    for i in 0..dim {
2552        let v = f16_to_f32(vec_f16[i]);
2553        let q = f16_to_f32(query_f16[i]);
2554        dot += q * v;
2555        norm += v * v;
2556    }
2557    (dot, norm)
2558}
2559
2560#[allow(dead_code)]
2561fn fused_dot_norm_u8_scalar(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2562    let mut dot = 0.0f32;
2563    let mut norm = 0.0f32;
2564    for i in 0..dim {
2565        let v = u8_to_f32(vec_u8[i]);
2566        dot += query[i] * v;
2567        norm += v * v;
2568    }
2569    (dot, norm)
2570}
2571
2572#[allow(dead_code)]
2573fn dot_product_f16_scalar(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2574    let mut dot = 0.0f32;
2575    for i in 0..dim {
2576        dot += f16_to_f32(query_f16[i]) * f16_to_f32(vec_f16[i]);
2577    }
2578    dot
2579}
2580
2581#[allow(dead_code)]
2582fn dot_product_u8_scalar(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2583    let mut dot = 0.0f32;
2584    for i in 0..dim {
2585        dot += query[i] * u8_to_f32(vec_u8[i]);
2586    }
2587    dot
2588}
2589
2590// ============================================================================
2591// x86_64 SSE4.1 quantized fused dot+norm
2592// ============================================================================
2593
2594#[cfg(target_arch = "x86_64")]
2595#[target_feature(enable = "sse2", enable = "sse4.1")]
2596#[allow(unsafe_op_in_unsafe_fn)]
2597unsafe fn fused_dot_norm_f16_sse(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2598    use std::arch::x86_64::*;
2599
2600    let chunks = dim / 4;
2601    let remainder = dim % 4;
2602
2603    let mut acc_dot = _mm_setzero_ps();
2604    let mut acc_norm = _mm_setzero_ps();
2605
2606    for chunk in 0..chunks {
2607        let base = chunk * 4;
2608        // Load 4 f16 values and convert to f32 using scalar conversion
2609        let v0 = f16_to_f32(*vec_f16.get_unchecked(base));
2610        let v1 = f16_to_f32(*vec_f16.get_unchecked(base + 1));
2611        let v2 = f16_to_f32(*vec_f16.get_unchecked(base + 2));
2612        let v3 = f16_to_f32(*vec_f16.get_unchecked(base + 3));
2613        let vb = _mm_set_ps(v3, v2, v1, v0);
2614
2615        let q0 = f16_to_f32(*query_f16.get_unchecked(base));
2616        let q1 = f16_to_f32(*query_f16.get_unchecked(base + 1));
2617        let q2 = f16_to_f32(*query_f16.get_unchecked(base + 2));
2618        let q3 = f16_to_f32(*query_f16.get_unchecked(base + 3));
2619        let va = _mm_set_ps(q3, q2, q1, q0);
2620
2621        acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2622        acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2623    }
2624
2625    // Horizontal sums
2626    let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2627    let sums_d = _mm_add_ps(acc_dot, shuf_d);
2628    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2629    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2630
2631    let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2632    let sums_n = _mm_add_ps(acc_norm, shuf_n);
2633    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2634    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2635
2636    let base = chunks * 4;
2637    for i in 0..remainder {
2638        let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2639        let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2640        dot += q * v;
2641        norm += v * v;
2642    }
2643
2644    (dot, norm)
2645}
2646
2647#[cfg(target_arch = "x86_64")]
2648#[target_feature(enable = "sse2", enable = "sse4.1")]
2649#[allow(unsafe_op_in_unsafe_fn)]
2650unsafe fn fused_dot_norm_u8_sse(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2651    use std::arch::x86_64::*;
2652
2653    let scale = _mm_set1_ps(U8_INV_SCALE);
2654    let offset = _mm_set1_ps(-1.0);
2655
2656    let chunks = dim / 4;
2657    let remainder = dim % 4;
2658
2659    let mut acc_dot = _mm_setzero_ps();
2660    let mut acc_norm = _mm_setzero_ps();
2661
2662    for chunk in 0..chunks {
2663        let base = chunk * 4;
2664
2665        // Load 4 bytes, zero-extend to i32, convert to f32, dequantize
2666        let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
2667            vec_u8.as_ptr().add(base) as *const i32
2668        ));
2669        let ints = _mm_cvtepu8_epi32(bytes);
2670        let floats = _mm_cvtepi32_ps(ints);
2671        let vb = _mm_add_ps(_mm_mul_ps(floats, scale), offset);
2672
2673        let va = _mm_loadu_ps(query.as_ptr().add(base));
2674
2675        acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2676        acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2677    }
2678
2679    // Horizontal sums
2680    let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2681    let sums_d = _mm_add_ps(acc_dot, shuf_d);
2682    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2683    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2684
2685    let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2686    let sums_n = _mm_add_ps(acc_norm, shuf_n);
2687    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2688    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2689
2690    let base = chunks * 4;
2691    for i in 0..remainder {
2692        let v = u8_to_f32(*vec_u8.get_unchecked(base + i));
2693        dot += *query.get_unchecked(base + i) * v;
2694        norm += v * v;
2695    }
2696
2697    (dot, norm)
2698}
2699
2700// ============================================================================
2701// x86_64 F16C + AVX + FMA accelerated f16 scoring
2702// ============================================================================
2703
2704#[cfg(target_arch = "x86_64")]
2705#[target_feature(enable = "avx", enable = "f16c", enable = "fma")]
2706#[allow(unsafe_op_in_unsafe_fn)]
2707unsafe fn fused_dot_norm_f16_f16c(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2708    use std::arch::x86_64::*;
2709
2710    let chunks16 = dim / 16;
2711    let remainder = dim % 16;
2712
2713    // 2 accumulator pairs to hide FMA latency (processes 16 f16 per iteration)
2714    let mut acc_dot0 = _mm256_setzero_ps();
2715    let mut acc_dot1 = _mm256_setzero_ps();
2716    let mut acc_norm0 = _mm256_setzero_ps();
2717    let mut acc_norm1 = _mm256_setzero_ps();
2718
2719    for c in 0..chunks16 {
2720        let base = c * 16;
2721
2722        // First 8 f16 elements
2723        let v_raw0 = _mm_loadu_si128(vec_f16.as_ptr().add(base) as *const __m128i);
2724        let vb0 = _mm256_cvtph_ps(v_raw0);
2725        let q_raw0 = _mm_loadu_si128(query_f16.as_ptr().add(base) as *const __m128i);
2726        let qa0 = _mm256_cvtph_ps(q_raw0);
2727        acc_dot0 = _mm256_fmadd_ps(qa0, vb0, acc_dot0);
2728        acc_norm0 = _mm256_fmadd_ps(vb0, vb0, acc_norm0);
2729
2730        // Second 8 f16 elements (independent accumulator chain)
2731        let v_raw1 = _mm_loadu_si128(vec_f16.as_ptr().add(base + 8) as *const __m128i);
2732        let vb1 = _mm256_cvtph_ps(v_raw1);
2733        let q_raw1 = _mm_loadu_si128(query_f16.as_ptr().add(base + 8) as *const __m128i);
2734        let qa1 = _mm256_cvtph_ps(q_raw1);
2735        acc_dot1 = _mm256_fmadd_ps(qa1, vb1, acc_dot1);
2736        acc_norm1 = _mm256_fmadd_ps(vb1, vb1, acc_norm1);
2737    }
2738
2739    // Combine accumulator pairs
2740    let acc_dot = _mm256_add_ps(acc_dot0, acc_dot1);
2741    let acc_norm = _mm256_add_ps(acc_norm0, acc_norm1);
2742
2743    // Horizontal sum 256→128→scalar
2744    let hi_d = _mm256_extractf128_ps(acc_dot, 1);
2745    let lo_d = _mm256_castps256_ps128(acc_dot);
2746    let sum_d = _mm_add_ps(lo_d, hi_d);
2747    let shuf_d = _mm_shuffle_ps(sum_d, sum_d, 0b10_11_00_01);
2748    let sums_d = _mm_add_ps(sum_d, shuf_d);
2749    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2750    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2751
2752    let hi_n = _mm256_extractf128_ps(acc_norm, 1);
2753    let lo_n = _mm256_castps256_ps128(acc_norm);
2754    let sum_n = _mm_add_ps(lo_n, hi_n);
2755    let shuf_n = _mm_shuffle_ps(sum_n, sum_n, 0b10_11_00_01);
2756    let sums_n = _mm_add_ps(sum_n, shuf_n);
2757    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2758    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2759
2760    let base = chunks16 * 16;
2761    for i in 0..remainder {
2762        let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2763        let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2764        dot += q * v;
2765        norm += v * v;
2766    }
2767
2768    (dot, norm)
2769}
2770
2771#[cfg(target_arch = "x86_64")]
2772#[target_feature(enable = "avx", enable = "f16c", enable = "fma")]
2773#[allow(unsafe_op_in_unsafe_fn)]
2774unsafe fn dot_product_f16_f16c(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2775    use std::arch::x86_64::*;
2776
2777    let chunks = dim / 8;
2778    let remainder = dim % 8;
2779    let mut acc = _mm256_setzero_ps();
2780
2781    for chunk in 0..chunks {
2782        let base = chunk * 8;
2783        let v_raw = _mm_loadu_si128(vec_f16.as_ptr().add(base) as *const __m128i);
2784        let vb = _mm256_cvtph_ps(v_raw);
2785        let q_raw = _mm_loadu_si128(query_f16.as_ptr().add(base) as *const __m128i);
2786        let qa = _mm256_cvtph_ps(q_raw);
2787        acc = _mm256_fmadd_ps(qa, vb, acc);
2788    }
2789
2790    let hi = _mm256_extractf128_ps(acc, 1);
2791    let lo = _mm256_castps256_ps128(acc);
2792    let sum = _mm_add_ps(lo, hi);
2793    let shuf = _mm_shuffle_ps(sum, sum, 0b10_11_00_01);
2794    let sums = _mm_add_ps(sum, shuf);
2795    let shuf2 = _mm_movehl_ps(sums, sums);
2796    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums, shuf2));
2797
2798    let base = chunks * 8;
2799    for i in 0..remainder {
2800        let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2801        let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2802        dot += q * v;
2803    }
2804    dot
2805}
2806
2807#[cfg(target_arch = "x86_64")]
2808#[target_feature(enable = "sse2", enable = "sse4.1")]
2809#[allow(unsafe_op_in_unsafe_fn)]
2810unsafe fn dot_product_u8_sse(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2811    use std::arch::x86_64::*;
2812
2813    let scale = _mm_set1_ps(U8_INV_SCALE);
2814    let offset = _mm_set1_ps(-1.0);
2815    let chunks = dim / 4;
2816    let remainder = dim % 4;
2817    let mut acc = _mm_setzero_ps();
2818
2819    for chunk in 0..chunks {
2820        let base = chunk * 4;
2821        let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
2822            vec_u8.as_ptr().add(base) as *const i32
2823        ));
2824        let ints = _mm_cvtepu8_epi32(bytes);
2825        let floats = _mm_cvtepi32_ps(ints);
2826        let vb = _mm_add_ps(_mm_mul_ps(floats, scale), offset);
2827        let va = _mm_loadu_ps(query.as_ptr().add(base));
2828        acc = _mm_add_ps(acc, _mm_mul_ps(va, vb));
2829    }
2830
2831    let shuf = _mm_shuffle_ps(acc, acc, 0b10_11_00_01);
2832    let sums = _mm_add_ps(acc, shuf);
2833    let shuf2 = _mm_movehl_ps(sums, sums);
2834    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums, shuf2));
2835
2836    let base = chunks * 4;
2837    for i in 0..remainder {
2838        dot += *query.get_unchecked(base + i) * u8_to_f32(*vec_u8.get_unchecked(base + i));
2839    }
2840    dot
2841}
2842
2843// ============================================================================
2844// Platform dispatch
2845// ============================================================================
2846
2847#[inline]
2848fn fused_dot_norm_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2849    #[cfg(target_arch = "aarch64")]
2850    {
2851        return unsafe { neon_quant::fused_dot_norm_f16(query_f16, vec_f16, dim) };
2852    }
2853
2854    #[cfg(target_arch = "x86_64")]
2855    {
2856        if is_x86_feature_detected!("f16c") && is_x86_feature_detected!("fma") {
2857            return unsafe { fused_dot_norm_f16_f16c(query_f16, vec_f16, dim) };
2858        }
2859        if sse::is_available() {
2860            return unsafe { fused_dot_norm_f16_sse(query_f16, vec_f16, dim) };
2861        }
2862    }
2863
2864    #[allow(unreachable_code)]
2865    fused_dot_norm_f16_scalar(query_f16, vec_f16, dim)
2866}
2867
2868#[inline]
2869fn fused_dot_norm_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2870    #[cfg(target_arch = "aarch64")]
2871    {
2872        return unsafe { neon_quant::fused_dot_norm_u8(query, vec_u8, dim) };
2873    }
2874
2875    #[cfg(target_arch = "x86_64")]
2876    {
2877        if sse::is_available() {
2878            return unsafe { fused_dot_norm_u8_sse(query, vec_u8, dim) };
2879        }
2880    }
2881
2882    #[allow(unreachable_code)]
2883    fused_dot_norm_u8_scalar(query, vec_u8, dim)
2884}
2885
2886// ── Dot-product-only dispatch (for unit_norm vectors) ─────────────────────
2887
2888#[inline]
2889fn dot_product_f16_quant(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2890    #[cfg(target_arch = "aarch64")]
2891    {
2892        return unsafe { neon_quant::dot_product_f16(query_f16, vec_f16, dim) };
2893    }
2894
2895    #[cfg(target_arch = "x86_64")]
2896    {
2897        if is_x86_feature_detected!("f16c") && is_x86_feature_detected!("fma") {
2898            return unsafe { dot_product_f16_f16c(query_f16, vec_f16, dim) };
2899        }
2900    }
2901
2902    #[allow(unreachable_code)]
2903    dot_product_f16_scalar(query_f16, vec_f16, dim)
2904}
2905
2906#[inline]
2907fn dot_product_u8_quant(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2908    #[cfg(target_arch = "aarch64")]
2909    {
2910        return unsafe { neon_quant::dot_product_u8(query, vec_u8, dim) };
2911    }
2912
2913    #[cfg(target_arch = "x86_64")]
2914    {
2915        if sse::is_available() {
2916            return unsafe { dot_product_u8_sse(query, vec_u8, dim) };
2917        }
2918    }
2919
2920    #[allow(unreachable_code)]
2921    dot_product_u8_scalar(query, vec_u8, dim)
2922}
2923
2924// ============================================================================
2925// Public batch cosine scoring for quantized vectors
2926// ============================================================================
2927
2928/// Batch cosine similarity: f32 query vs N contiguous f16 vectors.
2929///
2930/// `vectors_raw` is raw bytes: N vectors × dim × 2 bytes (f16 stored as u16).
2931/// Query is quantized to f16 once, then both query and vectors are scored in
2932/// f16 space using hardware SIMD conversion (8 elements/iteration on NEON).
2933/// Memory bandwidth is halved for both query and vector loads.
2934#[inline]
2935pub fn batch_cosine_scores_f16(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
2936    let n = scores.len();
2937    if dim == 0 || n == 0 {
2938        return;
2939    }
2940
2941    // Compute query inverse norm in f32 (full precision, before quantization)
2942    let norm_q_sq = dot_product_f32(query, query, dim);
2943    if norm_q_sq < f32::EPSILON {
2944        for s in scores.iter_mut() {
2945            *s = 0.0;
2946        }
2947        return;
2948    }
2949    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2950
2951    // Quantize query to f16 once (O(dim)), reused for all N vector scorings
2952    let query_f16: Vec<u16> = query.iter().map(|&v| f32_to_f16(v)).collect();
2953
2954    let vec_bytes = dim * 2;
2955    debug_assert!(vectors_raw.len() >= n * vec_bytes);
2956
2957    // Vectors file uses data-first layout with 8-byte padding between fields,
2958    // so mmap slices are always 2-byte aligned for u16 access.
2959    debug_assert!(
2960        (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
2961        "f16 vector data not 2-byte aligned"
2962    );
2963
2964    for i in 0..n {
2965        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
2966        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
2967
2968        let (dot, norm_v_sq) = fused_dot_norm_f16(&query_f16, f16_slice, dim);
2969        scores[i] = if norm_v_sq < f32::EPSILON {
2970            0.0
2971        } else {
2972            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
2973        };
2974    }
2975}
2976
2977/// Batch cosine similarity: f32 query vs N contiguous u8 vectors.
2978///
2979/// `vectors_raw` is raw bytes: N vectors × dim bytes (u8, mapping [-1,1]→[0,255]).
2980/// Converts u8→f32 using NEON widening chain (16 values/iteration), scores with FMA.
2981/// Memory bandwidth is quartered compared to f32 scoring.
2982#[inline]
2983pub fn batch_cosine_scores_u8(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
2984    let n = scores.len();
2985    if dim == 0 || n == 0 {
2986        return;
2987    }
2988
2989    let norm_q_sq = dot_product_f32(query, query, dim);
2990    if norm_q_sq < f32::EPSILON {
2991        for s in scores.iter_mut() {
2992            *s = 0.0;
2993        }
2994        return;
2995    }
2996    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2997
2998    debug_assert!(vectors_raw.len() >= n * dim);
2999
3000    for i in 0..n {
3001        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3002
3003        let (dot, norm_v_sq) = fused_dot_norm_u8(query, u8_slice, dim);
3004        scores[i] = if norm_v_sq < f32::EPSILON {
3005            0.0
3006        } else {
3007            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3008        };
3009    }
3010}
3011
3012// ============================================================================
3013// Batch dot-product scoring for unit-norm vectors
3014// ============================================================================
3015
3016/// Batch dot-product scoring: f32 query vs N contiguous f32 unit-norm vectors.
3017///
3018/// For pre-normalized vectors (||v|| = 1), cosine = dot(q, v) / ||q||.
3019/// Skips per-vector norm computation — ~40% less work than `batch_cosine_scores`.
3020#[inline]
3021pub fn batch_dot_scores(query: &[f32], vectors: &[f32], dim: usize, scores: &mut [f32]) {
3022    let n = scores.len();
3023    debug_assert!(vectors.len() >= n * dim);
3024    debug_assert_eq!(query.len(), dim);
3025
3026    if dim == 0 || n == 0 {
3027        return;
3028    }
3029
3030    let norm_q_sq = dot_product_f32(query, query, dim);
3031    if norm_q_sq < f32::EPSILON {
3032        for s in scores.iter_mut() {
3033            *s = 0.0;
3034        }
3035        return;
3036    }
3037    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3038
3039    for i in 0..n {
3040        let vec = &vectors[i * dim..(i + 1) * dim];
3041        let dot = dot_product_f32(query, vec, dim);
3042        scores[i] = dot * inv_norm_q;
3043    }
3044}
3045
3046/// Batch dot-product scoring: f32 query vs N contiguous f16 unit-norm vectors.
3047///
3048/// For pre-normalized vectors (||v|| = 1), cosine = dot(q, v) / ||q||.
3049/// Uses F16C/NEON hardware conversion + dot-only kernel.
3050#[inline]
3051pub fn batch_dot_scores_f16(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3052    let n = scores.len();
3053    if dim == 0 || n == 0 {
3054        return;
3055    }
3056
3057    let norm_q_sq = dot_product_f32(query, query, dim);
3058    if norm_q_sq < f32::EPSILON {
3059        for s in scores.iter_mut() {
3060            *s = 0.0;
3061        }
3062        return;
3063    }
3064    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3065
3066    let query_f16: Vec<u16> = query.iter().map(|&v| f32_to_f16(v)).collect();
3067    let vec_bytes = dim * 2;
3068    debug_assert!(vectors_raw.len() >= n * vec_bytes);
3069    debug_assert!(
3070        (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3071        "f16 vector data not 2-byte aligned"
3072    );
3073
3074    for i in 0..n {
3075        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3076        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3077        let dot = dot_product_f16_quant(&query_f16, f16_slice, dim);
3078        scores[i] = dot * inv_norm_q;
3079    }
3080}
3081
3082/// Batch dot-product scoring: f32 query vs N contiguous u8 unit-norm vectors.
3083///
3084/// For pre-normalized vectors (||v|| = 1), cosine = dot(q, v) / ||q||.
3085/// Uses NEON/SSE widening chain for u8→f32 conversion + dot-only kernel.
3086#[inline]
3087pub fn batch_dot_scores_u8(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3088    let n = scores.len();
3089    if dim == 0 || n == 0 {
3090        return;
3091    }
3092
3093    let norm_q_sq = dot_product_f32(query, query, dim);
3094    if norm_q_sq < f32::EPSILON {
3095        for s in scores.iter_mut() {
3096            *s = 0.0;
3097        }
3098        return;
3099    }
3100    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3101
3102    debug_assert!(vectors_raw.len() >= n * dim);
3103
3104    for i in 0..n {
3105        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3106        let dot = dot_product_u8_quant(query, u8_slice, dim);
3107        scores[i] = dot * inv_norm_q;
3108    }
3109}
3110
3111// ============================================================================
3112// Precomputed-norm batch scoring (avoids redundant query norm + f16 conversion)
3113// ============================================================================
3114
3115/// Batch cosine: f32 query vs N f32 vectors, with precomputed `inv_norm_q`.
3116#[inline]
3117pub fn batch_cosine_scores_precomp(
3118    query: &[f32],
3119    vectors: &[f32],
3120    dim: usize,
3121    scores: &mut [f32],
3122    inv_norm_q: f32,
3123) {
3124    let n = scores.len();
3125    debug_assert!(vectors.len() >= n * dim);
3126    for i in 0..n {
3127        let vec = &vectors[i * dim..(i + 1) * dim];
3128        let (dot, norm_v_sq) = fused_dot_norm(query, vec, dim);
3129        scores[i] = if norm_v_sq < f32::EPSILON {
3130            0.0
3131        } else {
3132            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3133        };
3134    }
3135}
3136
3137/// Batch cosine: precomputed `inv_norm_q` + `query_f16` vs N f16 vectors.
3138#[inline]
3139pub fn batch_cosine_scores_f16_precomp(
3140    query_f16: &[u16],
3141    vectors_raw: &[u8],
3142    dim: usize,
3143    scores: &mut [f32],
3144    inv_norm_q: f32,
3145) {
3146    let n = scores.len();
3147    let vec_bytes = dim * 2;
3148    debug_assert!(vectors_raw.len() >= n * vec_bytes);
3149    for i in 0..n {
3150        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3151        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3152        let (dot, norm_v_sq) = fused_dot_norm_f16(query_f16, f16_slice, dim);
3153        scores[i] = if norm_v_sq < f32::EPSILON {
3154            0.0
3155        } else {
3156            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3157        };
3158    }
3159}
3160
3161/// Batch cosine: precomputed `inv_norm_q` vs N u8 vectors.
3162#[inline]
3163pub fn batch_cosine_scores_u8_precomp(
3164    query: &[f32],
3165    vectors_raw: &[u8],
3166    dim: usize,
3167    scores: &mut [f32],
3168    inv_norm_q: f32,
3169) {
3170    let n = scores.len();
3171    debug_assert!(vectors_raw.len() >= n * dim);
3172    for i in 0..n {
3173        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3174        let (dot, norm_v_sq) = fused_dot_norm_u8(query, u8_slice, dim);
3175        scores[i] = if norm_v_sq < f32::EPSILON {
3176            0.0
3177        } else {
3178            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3179        };
3180    }
3181}
3182
3183/// Batch dot-product: precomputed `inv_norm_q` vs N f32 unit-norm vectors.
3184#[inline]
3185pub fn batch_dot_scores_precomp(
3186    query: &[f32],
3187    vectors: &[f32],
3188    dim: usize,
3189    scores: &mut [f32],
3190    inv_norm_q: f32,
3191) {
3192    let n = scores.len();
3193    debug_assert!(vectors.len() >= n * dim);
3194    for i in 0..n {
3195        let vec = &vectors[i * dim..(i + 1) * dim];
3196        scores[i] = dot_product_f32(query, vec, dim) * inv_norm_q;
3197    }
3198}
3199
3200/// Batch dot-product: precomputed `inv_norm_q` + `query_f16` vs N f16 unit-norm vectors.
3201#[inline]
3202pub fn batch_dot_scores_f16_precomp(
3203    query_f16: &[u16],
3204    vectors_raw: &[u8],
3205    dim: usize,
3206    scores: &mut [f32],
3207    inv_norm_q: f32,
3208) {
3209    let n = scores.len();
3210    let vec_bytes = dim * 2;
3211    debug_assert!(vectors_raw.len() >= n * vec_bytes);
3212    for i in 0..n {
3213        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3214        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3215        scores[i] = dot_product_f16_quant(query_f16, f16_slice, dim) * inv_norm_q;
3216    }
3217}
3218
3219/// Batch dot-product: precomputed `inv_norm_q` vs N u8 unit-norm vectors.
3220#[inline]
3221pub fn batch_dot_scores_u8_precomp(
3222    query: &[f32],
3223    vectors_raw: &[u8],
3224    dim: usize,
3225    scores: &mut [f32],
3226    inv_norm_q: f32,
3227) {
3228    let n = scores.len();
3229    debug_assert!(vectors_raw.len() >= n * dim);
3230    for i in 0..n {
3231        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3232        scores[i] = dot_product_u8_quant(query, u8_slice, dim) * inv_norm_q;
3233    }
3234}
3235
3236/// Compute cosine similarity between two f32 vectors with SIMD acceleration
3237///
3238/// Returns dot(a,b) / (||a|| * ||b||), range [-1, 1]
3239/// Returns 0.0 if either vector has zero norm.
3240#[inline]
3241pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
3242    debug_assert_eq!(a.len(), b.len());
3243    let count = a.len();
3244
3245    if count == 0 {
3246        return 0.0;
3247    }
3248
3249    let dot = dot_product_f32(a, b, count);
3250    let norm_a = dot_product_f32(a, a, count);
3251    let norm_b = dot_product_f32(b, b, count);
3252
3253    let denom = (norm_a * norm_b).sqrt();
3254    if denom < f32::EPSILON {
3255        return 0.0;
3256    }
3257
3258    dot / denom
3259}
3260
3261// ============================================================================
3262// Hamming distance for binary dense vectors
3263// ============================================================================
3264
3265/// Compute Hamming distance between two packed-bit vectors.
3266/// Returns the number of differing bits.
3267///
3268/// Uses NEON intrinsics on aarch64, POPCNT on x86_64, scalar fallback otherwise.
3269#[inline]
3270pub fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
3271    debug_assert_eq!(a.len(), b.len());
3272
3273    #[cfg(target_arch = "aarch64")]
3274    unsafe {
3275        neon::hamming_distance(a, b)
3276    }
3277
3278    #[cfg(target_arch = "x86_64")]
3279    {
3280        if avx2::is_available() {
3281            return unsafe { avx2::hamming_distance(a, b) };
3282        }
3283        hamming_distance_scalar(a, b)
3284    }
3285
3286    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3287    hamming_distance_scalar(a, b)
3288}
3289
3290/// Scalar Hamming distance using u64 chunks + count_ones().
3291/// On x86_64, count_ones() compiles to POPCNT when target-cpu supports it.
3292#[inline]
3293#[allow(dead_code)]
3294fn hamming_distance_scalar(a: &[u8], b: &[u8]) -> u32 {
3295    let len = a.len();
3296    let chunks = len / 8;
3297    let remainder = len % 8;
3298    let mut total = 0u32;
3299
3300    for i in 0..chunks {
3301        let off = i * 8;
3302        let va = unsafe { std::ptr::read_unaligned(a.as_ptr().add(off) as *const u64) };
3303        let vb = unsafe { std::ptr::read_unaligned(b.as_ptr().add(off) as *const u64) };
3304        total += (va ^ vb).count_ones();
3305    }
3306
3307    let base = chunks * 8;
3308    for i in 0..remainder {
3309        total += (a[base + i] ^ b[base + i]).count_ones();
3310    }
3311
3312    total
3313}
3314
3315/// Batch Hamming scoring: compute similarity scores for multiple binary vectors.
3316///
3317/// `query` and each vector in `db` are packed-bit vectors of `byte_len` bytes each.
3318/// `dim_bits` is the number of bits (dimensions) for normalization.
3319/// Score = 1.0 - hamming_distance / dim_bits (range [0.0, 1.0]).
3320pub fn batch_hamming_scores(
3321    query: &[u8],
3322    db: &[u8],
3323    byte_len: usize,
3324    dim_bits: usize,
3325    scores: &mut [f32],
3326) {
3327    let n = scores.len();
3328    debug_assert_eq!(query.len(), byte_len);
3329    debug_assert!(db.len() >= n * byte_len);
3330
3331    if byte_len == 0 || n == 0 || dim_bits == 0 {
3332        return;
3333    }
3334
3335    let inv_dim = 1.0 / dim_bits as f32;
3336
3337    for i in 0..n {
3338        let vec = &db[i * byte_len..(i + 1) * byte_len];
3339        let dist = hamming_distance(query, vec);
3340        scores[i] = 1.0 - dist as f32 * inv_dim;
3341    }
3342}
3343
3344#[cfg(test)]
3345mod tests {
3346    use super::*;
3347
3348    #[test]
3349    fn test_unpack_8bit() {
3350        let input: Vec<u8> = (0..128).collect();
3351        let mut output = vec![0u32; 128];
3352        unpack_8bit(&input, &mut output, 128);
3353
3354        for (i, &v) in output.iter().enumerate() {
3355            assert_eq!(v, i as u32);
3356        }
3357    }
3358
3359    #[test]
3360    fn test_unpack_16bit() {
3361        let mut input = vec![0u8; 256];
3362        for i in 0..128 {
3363            let val = (i * 100) as u16;
3364            input[i * 2] = val as u8;
3365            input[i * 2 + 1] = (val >> 8) as u8;
3366        }
3367
3368        let mut output = vec![0u32; 128];
3369        unpack_16bit(&input, &mut output, 128);
3370
3371        for (i, &v) in output.iter().enumerate() {
3372            assert_eq!(v, (i * 100) as u32);
3373        }
3374    }
3375
3376    #[test]
3377    fn test_unpack_32bit() {
3378        let mut input = vec![0u8; 512];
3379        for i in 0..128 {
3380            let val = (i * 1000) as u32;
3381            let bytes = val.to_le_bytes();
3382            input[i * 4..i * 4 + 4].copy_from_slice(&bytes);
3383        }
3384
3385        let mut output = vec![0u32; 128];
3386        unpack_32bit(&input, &mut output, 128);
3387
3388        for (i, &v) in output.iter().enumerate() {
3389            assert_eq!(v, (i * 1000) as u32);
3390        }
3391    }
3392
3393    #[test]
3394    fn test_delta_decode() {
3395        // doc_ids: [10, 15, 20, 30, 50]
3396        // gaps: [5, 5, 10, 20]
3397        // deltas (gap-1): [4, 4, 9, 19]
3398        let deltas = vec![4u32, 4, 9, 19];
3399        let mut output = vec![0u32; 5];
3400
3401        delta_decode(&mut output, &deltas, 10, 5);
3402
3403        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3404    }
3405
3406    #[test]
3407    fn test_add_one() {
3408        let mut values = vec![0u32, 1, 2, 3, 4, 5, 6, 7];
3409        add_one(&mut values, 8);
3410
3411        assert_eq!(values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
3412    }
3413
3414    #[test]
3415    fn test_bits_needed() {
3416        assert_eq!(bits_needed(0), 0);
3417        assert_eq!(bits_needed(1), 1);
3418        assert_eq!(bits_needed(2), 2);
3419        assert_eq!(bits_needed(3), 2);
3420        assert_eq!(bits_needed(4), 3);
3421        assert_eq!(bits_needed(255), 8);
3422        assert_eq!(bits_needed(256), 9);
3423        assert_eq!(bits_needed(u32::MAX), 32);
3424    }
3425
3426    #[test]
3427    fn test_unpack_8bit_delta_decode() {
3428        // doc_ids: [10, 15, 20, 30, 50]
3429        // gaps: [5, 5, 10, 20]
3430        // deltas (gap-1): [4, 4, 9, 19] stored as u8
3431        let input: Vec<u8> = vec![4, 4, 9, 19];
3432        let mut output = vec![0u32; 5];
3433
3434        unpack_8bit_delta_decode(&input, &mut output, 10, 5);
3435
3436        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3437    }
3438
3439    #[test]
3440    fn test_unpack_16bit_delta_decode() {
3441        // doc_ids: [100, 600, 1100, 2100, 4100]
3442        // gaps: [500, 500, 1000, 2000]
3443        // deltas (gap-1): [499, 499, 999, 1999] stored as u16
3444        let mut input = vec![0u8; 8];
3445        for (i, &delta) in [499u16, 499, 999, 1999].iter().enumerate() {
3446            input[i * 2] = delta as u8;
3447            input[i * 2 + 1] = (delta >> 8) as u8;
3448        }
3449        let mut output = vec![0u32; 5];
3450
3451        unpack_16bit_delta_decode(&input, &mut output, 100, 5);
3452
3453        assert_eq!(output, vec![100, 600, 1100, 2100, 4100]);
3454    }
3455
3456    #[test]
3457    fn test_fused_vs_separate_8bit() {
3458        // Test that fused and separate operations produce the same result
3459        let input: Vec<u8> = (0..127).collect();
3460        let first_value = 1000u32;
3461        let count = 128;
3462
3463        // Separate: unpack then delta_decode
3464        let mut unpacked = vec![0u32; 128];
3465        unpack_8bit(&input, &mut unpacked, 127);
3466        let mut separate_output = vec![0u32; 128];
3467        delta_decode(&mut separate_output, &unpacked, first_value, count);
3468
3469        // Fused
3470        let mut fused_output = vec![0u32; 128];
3471        unpack_8bit_delta_decode(&input, &mut fused_output, first_value, count);
3472
3473        assert_eq!(separate_output, fused_output);
3474    }
3475
3476    #[test]
3477    fn test_round_bit_width() {
3478        assert_eq!(round_bit_width(0), 0);
3479        assert_eq!(round_bit_width(1), 8);
3480        assert_eq!(round_bit_width(5), 8);
3481        assert_eq!(round_bit_width(8), 8);
3482        assert_eq!(round_bit_width(9), 16);
3483        assert_eq!(round_bit_width(12), 16);
3484        assert_eq!(round_bit_width(16), 16);
3485        assert_eq!(round_bit_width(17), 32);
3486        assert_eq!(round_bit_width(24), 32);
3487        assert_eq!(round_bit_width(32), 32);
3488    }
3489
3490    #[test]
3491    fn test_rounded_bitwidth_from_exact() {
3492        assert_eq!(RoundedBitWidth::from_exact(0), RoundedBitWidth::Zero);
3493        assert_eq!(RoundedBitWidth::from_exact(1), RoundedBitWidth::Bits8);
3494        assert_eq!(RoundedBitWidth::from_exact(8), RoundedBitWidth::Bits8);
3495        assert_eq!(RoundedBitWidth::from_exact(9), RoundedBitWidth::Bits16);
3496        assert_eq!(RoundedBitWidth::from_exact(16), RoundedBitWidth::Bits16);
3497        assert_eq!(RoundedBitWidth::from_exact(17), RoundedBitWidth::Bits32);
3498        assert_eq!(RoundedBitWidth::from_exact(32), RoundedBitWidth::Bits32);
3499    }
3500
3501    #[test]
3502    fn test_pack_unpack_rounded_8bit() {
3503        let values: Vec<u32> = (0..128).map(|i| i % 256).collect();
3504        let mut packed = vec![0u8; 128];
3505
3506        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits8, &mut packed);
3507        assert_eq!(bytes_written, 128);
3508
3509        let mut unpacked = vec![0u32; 128];
3510        unpack_rounded(&packed, RoundedBitWidth::Bits8, &mut unpacked, 128);
3511
3512        assert_eq!(values, unpacked);
3513    }
3514
3515    #[test]
3516    fn test_pack_unpack_rounded_16bit() {
3517        let values: Vec<u32> = (0..128).map(|i| i * 100).collect();
3518        let mut packed = vec![0u8; 256];
3519
3520        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits16, &mut packed);
3521        assert_eq!(bytes_written, 256);
3522
3523        let mut unpacked = vec![0u32; 128];
3524        unpack_rounded(&packed, RoundedBitWidth::Bits16, &mut unpacked, 128);
3525
3526        assert_eq!(values, unpacked);
3527    }
3528
3529    #[test]
3530    fn test_pack_unpack_rounded_32bit() {
3531        let values: Vec<u32> = (0..128).map(|i| i * 100000).collect();
3532        let mut packed = vec![0u8; 512];
3533
3534        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits32, &mut packed);
3535        assert_eq!(bytes_written, 512);
3536
3537        let mut unpacked = vec![0u32; 128];
3538        unpack_rounded(&packed, RoundedBitWidth::Bits32, &mut unpacked, 128);
3539
3540        assert_eq!(values, unpacked);
3541    }
3542
3543    #[test]
3544    fn test_unpack_rounded_delta_decode() {
3545        // Test 8-bit rounded delta decode
3546        // doc_ids: [10, 15, 20, 30, 50]
3547        // gaps: [5, 5, 10, 20]
3548        // deltas (gap-1): [4, 4, 9, 19] stored as u8
3549        let input: Vec<u8> = vec![4, 4, 9, 19];
3550        let mut output = vec![0u32; 5];
3551
3552        unpack_rounded_delta_decode(&input, RoundedBitWidth::Bits8, &mut output, 10, 5);
3553
3554        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3555    }
3556
3557    #[test]
3558    fn test_unpack_rounded_delta_decode_zero() {
3559        // All zeros means gaps of 1 (consecutive doc IDs)
3560        let input: Vec<u8> = vec![];
3561        let mut output = vec![0u32; 5];
3562
3563        unpack_rounded_delta_decode(&input, RoundedBitWidth::Zero, &mut output, 100, 5);
3564
3565        assert_eq!(output, vec![100, 101, 102, 103, 104]);
3566    }
3567
3568    // ========================================================================
3569    // Sparse Vector SIMD Tests
3570    // ========================================================================
3571
3572    #[test]
3573    fn test_dequantize_uint8() {
3574        let input: Vec<u8> = vec![0, 128, 255, 64, 192];
3575        let mut output = vec![0.0f32; 5];
3576        let scale = 0.1;
3577        let min_val = 1.0;
3578
3579        dequantize_uint8(&input, &mut output, scale, min_val, 5);
3580
3581        // Expected: input[i] * scale + min_val
3582        assert!((output[0] - 1.0).abs() < 1e-6); // 0 * 0.1 + 1.0 = 1.0
3583        assert!((output[1] - 13.8).abs() < 1e-6); // 128 * 0.1 + 1.0 = 13.8
3584        assert!((output[2] - 26.5).abs() < 1e-6); // 255 * 0.1 + 1.0 = 26.5
3585        assert!((output[3] - 7.4).abs() < 1e-6); // 64 * 0.1 + 1.0 = 7.4
3586        assert!((output[4] - 20.2).abs() < 1e-6); // 192 * 0.1 + 1.0 = 20.2
3587    }
3588
3589    #[test]
3590    fn test_dequantize_uint8_large() {
3591        // Test with 128 values (full SIMD block)
3592        let input: Vec<u8> = (0..128).collect();
3593        let mut output = vec![0.0f32; 128];
3594        let scale = 2.0;
3595        let min_val = -10.0;
3596
3597        dequantize_uint8(&input, &mut output, scale, min_val, 128);
3598
3599        for (i, &out) in output.iter().enumerate().take(128) {
3600            let expected = i as f32 * scale + min_val;
3601            assert!(
3602                (out - expected).abs() < 1e-5,
3603                "Mismatch at {}: expected {}, got {}",
3604                i,
3605                expected,
3606                out
3607            );
3608        }
3609    }
3610
3611    #[test]
3612    fn test_dot_product_f32() {
3613        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
3614        let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0];
3615
3616        let result = dot_product_f32(&a, &b, 5);
3617
3618        // Expected: 1*2 + 2*3 + 3*4 + 4*5 + 5*6 = 2 + 6 + 12 + 20 + 30 = 70
3619        assert!((result - 70.0).abs() < 1e-5);
3620    }
3621
3622    #[test]
3623    fn test_dot_product_f32_large() {
3624        // Test with 128 values
3625        let a: Vec<f32> = (0..128).map(|i| i as f32).collect();
3626        let b: Vec<f32> = (0..128).map(|i| (i + 1) as f32).collect();
3627
3628        let result = dot_product_f32(&a, &b, 128);
3629
3630        // Compute expected
3631        let expected: f32 = (0..128).map(|i| (i as f32) * ((i + 1) as f32)).sum();
3632        assert!(
3633            (result - expected).abs() < 1e-3,
3634            "Expected {}, got {}",
3635            expected,
3636            result
3637        );
3638    }
3639
3640    #[test]
3641    fn test_fused_dot_norm() {
3642        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
3643        let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
3644        let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3645
3646        let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3647        let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3648        assert!(
3649            (dot - expected_dot).abs() < 1e-5,
3650            "dot: expected {}, got {}",
3651            expected_dot,
3652            dot
3653        );
3654        assert!(
3655            (norm_b - expected_norm).abs() < 1e-5,
3656            "norm: expected {}, got {}",
3657            expected_norm,
3658            norm_b
3659        );
3660    }
3661
3662    #[test]
3663    fn test_fused_dot_norm_large() {
3664        let a: Vec<f32> = (0..768).map(|i| (i as f32) * 0.01).collect();
3665        let b: Vec<f32> = (0..768).map(|i| (i as f32) * 0.02 + 0.5).collect();
3666        let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3667
3668        let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3669        let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3670        assert!(
3671            (dot - expected_dot).abs() < 1.0,
3672            "dot: expected {}, got {}",
3673            expected_dot,
3674            dot
3675        );
3676        assert!(
3677            (norm_b - expected_norm).abs() < 1.0,
3678            "norm: expected {}, got {}",
3679            expected_norm,
3680            norm_b
3681        );
3682    }
3683
3684    #[test]
3685    fn test_batch_cosine_scores() {
3686        // 4 vectors of dim 3
3687        let query = vec![1.0f32, 0.0, 0.0];
3688        let vectors = vec![
3689            1.0, 0.0, 0.0, // identical to query
3690            0.0, 1.0, 0.0, // orthogonal
3691            -1.0, 0.0, 0.0, // opposite
3692            0.5, 0.5, 0.0, // 45 degrees
3693        ];
3694        let mut scores = vec![0f32; 4];
3695        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3696
3697        assert!((scores[0] - 1.0).abs() < 1e-5, "identical: {}", scores[0]);
3698        assert!(scores[1].abs() < 1e-5, "orthogonal: {}", scores[1]);
3699        assert!((scores[2] - (-1.0)).abs() < 1e-5, "opposite: {}", scores[2]);
3700        let expected_45 = 0.5f32 / (0.5f32.powi(2) + 0.5f32.powi(2)).sqrt();
3701        assert!(
3702            (scores[3] - expected_45).abs() < 1e-5,
3703            "45deg: expected {}, got {}",
3704            expected_45,
3705            scores[3]
3706        );
3707    }
3708
3709    #[test]
3710    fn test_batch_cosine_scores_matches_individual() {
3711        let query: Vec<f32> = (0..128).map(|i| (i as f32) * 0.1).collect();
3712        let n = 50;
3713        let dim = 128;
3714        let vectors: Vec<f32> = (0..n * dim).map(|i| ((i * 7 + 3) as f32) * 0.01).collect();
3715
3716        let mut batch_scores = vec![0f32; n];
3717        batch_cosine_scores(&query, &vectors, dim, &mut batch_scores);
3718
3719        for i in 0..n {
3720            let vec_i = &vectors[i * dim..(i + 1) * dim];
3721            let individual = cosine_similarity(&query, vec_i);
3722            assert!(
3723                (batch_scores[i] - individual).abs() < 1e-5,
3724                "vec {}: batch={}, individual={}",
3725                i,
3726                batch_scores[i],
3727                individual
3728            );
3729        }
3730    }
3731
3732    #[test]
3733    fn test_batch_cosine_scores_empty() {
3734        let query = vec![1.0f32, 2.0, 3.0];
3735        let vectors: Vec<f32> = vec![];
3736        let mut scores: Vec<f32> = vec![];
3737        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3738        assert!(scores.is_empty());
3739    }
3740
3741    #[test]
3742    fn test_batch_cosine_scores_zero_query() {
3743        let query = vec![0.0f32, 0.0, 0.0];
3744        let vectors = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
3745        let mut scores = vec![0f32; 2];
3746        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3747        assert_eq!(scores[0], 0.0);
3748        assert_eq!(scores[1], 0.0);
3749    }
3750
3751    // ================================================================
3752    // f16 conversion tests
3753    // ================================================================
3754
3755    #[test]
3756    fn test_f16_roundtrip_normal() {
3757        for &v in &[0.0f32, 1.0, -1.0, 0.5, -0.5, 0.333, 65504.0] {
3758            let h = f32_to_f16(v);
3759            let back = f16_to_f32(h);
3760            let err = (back - v).abs() / v.abs().max(1e-6);
3761            assert!(
3762                err < 0.002,
3763                "f16 roundtrip {v} → {h:#06x} → {back}, rel err {err}"
3764            );
3765        }
3766    }
3767
3768    #[test]
3769    fn test_f16_special() {
3770        // Zero
3771        assert_eq!(f16_to_f32(f32_to_f16(0.0)), 0.0);
3772        // Negative zero
3773        assert_eq!(f32_to_f16(-0.0), 0x8000);
3774        // Infinity
3775        assert!(f16_to_f32(f32_to_f16(f32::INFINITY)).is_infinite());
3776        // NaN
3777        assert!(f16_to_f32(f32_to_f16(f32::NAN)).is_nan());
3778    }
3779
3780    #[test]
3781    fn test_f16_embedding_range() {
3782        // Typical embedding values in [-1, 1]
3783        let values: Vec<f32> = (-100..=100).map(|i| i as f32 / 100.0).collect();
3784        for &v in &values {
3785            let back = f16_to_f32(f32_to_f16(v));
3786            assert!((back - v).abs() < 0.001, "f16 error for {v}: got {back}");
3787        }
3788    }
3789
3790    // ================================================================
3791    // u8 conversion tests
3792    // ================================================================
3793
3794    #[test]
3795    fn test_u8_roundtrip() {
3796        // Boundary values
3797        assert_eq!(f32_to_u8_saturating(-1.0), 0);
3798        assert_eq!(f32_to_u8_saturating(1.0), 255);
3799        assert_eq!(f32_to_u8_saturating(0.0), 127); // ~127.5 truncated
3800
3801        // Saturation
3802        assert_eq!(f32_to_u8_saturating(-2.0), 0);
3803        assert_eq!(f32_to_u8_saturating(2.0), 255);
3804    }
3805
3806    #[test]
3807    fn test_u8_dequantize() {
3808        assert!((u8_to_f32(0) - (-1.0)).abs() < 0.01);
3809        assert!((u8_to_f32(255) - 1.0).abs() < 0.01);
3810        assert!((u8_to_f32(127) - 0.0).abs() < 0.01);
3811    }
3812
3813    // ================================================================
3814    // Batch scoring tests for quantized vectors
3815    // ================================================================
3816
3817    #[test]
3818    fn test_batch_cosine_scores_f16() {
3819        let query = vec![0.6f32, 0.8, 0.0, 0.0];
3820        let dim = 4;
3821        let vecs_f32 = vec![
3822            0.6f32, 0.8, 0.0, 0.0, // identical to query
3823            0.0, 0.0, 0.6, 0.8, // orthogonal
3824        ];
3825
3826        // Quantize to f16
3827        let mut f16_buf = vec![0u16; 8];
3828        batch_f32_to_f16(&vecs_f32, &mut f16_buf);
3829        let raw: &[u8] =
3830            unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
3831
3832        let mut scores = vec![0f32; 2];
3833        batch_cosine_scores_f16(&query, raw, dim, &mut scores);
3834
3835        assert!(
3836            (scores[0] - 1.0).abs() < 0.01,
3837            "identical vectors: {}",
3838            scores[0]
3839        );
3840        assert!(scores[1].abs() < 0.01, "orthogonal vectors: {}", scores[1]);
3841    }
3842
3843    #[test]
3844    fn test_batch_cosine_scores_u8() {
3845        let query = vec![0.6f32, 0.8, 0.0, 0.0];
3846        let dim = 4;
3847        let vecs_f32 = vec![
3848            0.6f32, 0.8, 0.0, 0.0, // ~identical to query
3849            -0.6, -0.8, 0.0, 0.0, // opposite
3850        ];
3851
3852        // Quantize to u8
3853        let mut u8_buf = vec![0u8; 8];
3854        batch_f32_to_u8(&vecs_f32, &mut u8_buf);
3855
3856        let mut scores = vec![0f32; 2];
3857        batch_cosine_scores_u8(&query, &u8_buf, dim, &mut scores);
3858
3859        assert!(scores[0] > 0.95, "similar vectors: {}", scores[0]);
3860        assert!(scores[1] < -0.95, "opposite vectors: {}", scores[1]);
3861    }
3862
3863    #[test]
3864    fn test_batch_cosine_scores_f16_large_dim() {
3865        // Test with typical embedding dimension
3866        let dim = 768;
3867        let query: Vec<f32> = (0..dim).map(|i| (i as f32 / dim as f32) - 0.5).collect();
3868        let vec2: Vec<f32> = query.iter().map(|x| x * 0.9 + 0.01).collect();
3869
3870        let mut all_vecs = query.clone();
3871        all_vecs.extend_from_slice(&vec2);
3872
3873        let mut f16_buf = vec![0u16; all_vecs.len()];
3874        batch_f32_to_f16(&all_vecs, &mut f16_buf);
3875        let raw: &[u8] =
3876            unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
3877
3878        let mut scores = vec![0f32; 2];
3879        batch_cosine_scores_f16(&query, raw, dim, &mut scores);
3880
3881        // Self-similarity should be ~1.0
3882        assert!((scores[0] - 1.0).abs() < 0.01, "self-sim: {}", scores[0]);
3883        // High similarity with scaled version
3884        assert!(scores[1] > 0.99, "scaled-sim: {}", scores[1]);
3885    }
3886
3887    // ================================================================
3888    // Hamming distance tests
3889    // ================================================================
3890
3891    #[test]
3892    fn test_hamming_distance_identical() {
3893        let a = vec![0xAA; 64];
3894        assert_eq!(hamming_distance(&a, &a), 0);
3895    }
3896
3897    #[test]
3898    fn test_hamming_distance_opposite() {
3899        let a = vec![0xFF; 32];
3900        let b = vec![0x00; 32];
3901        assert_eq!(hamming_distance(&a, &b), 256);
3902    }
3903
3904    #[test]
3905    fn test_hamming_distance_known() {
3906        // Single byte: 0b10101010 vs 0b01010101 = 8 bits differ
3907        let a = vec![0xAA];
3908        let b = vec![0x55];
3909        assert_eq!(hamming_distance(&a, &b), 8);
3910
3911        // Two bytes
3912        let a = vec![0xFF, 0x00];
3913        let b = vec![0x00, 0x00];
3914        assert_eq!(hamming_distance(&a, &b), 8);
3915    }
3916
3917    #[test]
3918    fn test_hamming_distance_single_bit() {
3919        let a = vec![0x00; 16];
3920        let mut b = vec![0x00; 16];
3921        b[7] = 0x01; // flip one bit
3922        assert_eq!(hamming_distance(&a, &b), 1);
3923    }
3924
3925    #[test]
3926    fn test_hamming_distance_empty() {
3927        let a: Vec<u8> = vec![];
3928        assert_eq!(hamming_distance(&a, &a), 0);
3929    }
3930
3931    #[test]
3932    fn test_hamming_distance_remainder_path() {
3933        // 17 bytes: not aligned to 16 (NEON) or 32 (AVX2)
3934        let a = vec![0xFF; 17];
3935        let b = vec![0x00; 17];
3936        assert_eq!(hamming_distance(&a, &b), 136); // 17 * 8
3937
3938        // 33 bytes: tests 32-byte chunk + 1 remainder for AVX2
3939        let a = vec![0xFF; 33];
3940        let b = vec![0x00; 33];
3941        assert_eq!(hamming_distance(&a, &b), 264); // 33 * 8
3942    }
3943
3944    #[test]
3945    fn test_hamming_distance_large() {
3946        // 4096 bytes = 32768 bits, all differing
3947        let a = vec![0xFF; 4096];
3948        let b = vec![0x00; 4096];
3949        assert_eq!(hamming_distance(&a, &b), 32768);
3950    }
3951
3952    #[test]
3953    fn test_hamming_distance_scalar_matches() {
3954        // Verify SIMD path matches scalar for various sizes
3955        for size in [1, 7, 8, 15, 16, 31, 32, 63, 64, 100, 128, 255, 256] {
3956            let a: Vec<u8> = (0..size).map(|i| (i * 37 + 13) as u8).collect();
3957            let b: Vec<u8> = (0..size).map(|i| (i * 53 + 7) as u8).collect();
3958            let expected = hamming_distance_scalar(&a, &b);
3959            let got = hamming_distance(&a, &b);
3960            assert_eq!(got, expected, "mismatch at size {size}");
3961        }
3962    }
3963
3964    // ================================================================
3965    // Batch Hamming scoring tests
3966    // ================================================================
3967
3968    #[test]
3969    fn test_batch_hamming_scores_identical() {
3970        let query = vec![0xAA; 16];
3971        let db = vec![0xAA; 16]; // one vector, identical
3972        let mut scores = vec![0f32; 1];
3973        batch_hamming_scores(&query, &db, 16, 128, &mut scores);
3974        assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
3975    }
3976
3977    #[test]
3978    fn test_batch_hamming_scores_opposite() {
3979        let query = vec![0xFF; 16];
3980        let db = vec![0x00; 16];
3981        let mut scores = vec![0f32; 1];
3982        batch_hamming_scores(&query, &db, 16, 128, &mut scores);
3983        assert!((scores[0] - 0.0).abs() < 1e-6, "opposite: {}", scores[0]);
3984    }
3985
3986    #[test]
3987    fn test_batch_hamming_scores_multiple() {
3988        let byte_len = 8;
3989        let dim_bits = 64;
3990        let query = vec![0xFF; byte_len];
3991        let mut db = Vec::new();
3992        db.extend_from_slice(&vec![0xFF; byte_len]); // identical → 1.0
3993        db.extend_from_slice(&vec![0x00; byte_len]); // opposite → 0.0
3994        db.extend_from_slice(&vec![0x0F; byte_len]); // half bits differ → 0.5
3995
3996        let mut scores = vec![0f32; 3];
3997        batch_hamming_scores(&query, &db, byte_len, dim_bits, &mut scores);
3998
3999        assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4000        assert!((scores[1] - 0.0).abs() < 1e-6, "opposite: {}", scores[1]);
4001        assert!((scores[2] - 0.5).abs() < 1e-6, "half: {}", scores[2]);
4002    }
4003
4004    #[test]
4005    fn test_batch_hamming_scores_empty() {
4006        let query = vec![0xFF; 8];
4007        let db: Vec<u8> = vec![];
4008        let mut scores: Vec<f32> = vec![];
4009        batch_hamming_scores(&query, &db, 8, 64, &mut scores);
4010        assert!(scores.is_empty());
4011    }
4012
4013    #[test]
4014    fn test_batch_hamming_scores_zero_byte_len() {
4015        let query: Vec<u8> = vec![];
4016        let db: Vec<u8> = vec![];
4017        let mut scores = vec![0f32; 1];
4018        batch_hamming_scores(&query, &db, 0, 0, &mut scores);
4019        // Should return early without modifying scores
4020        assert_eq!(scores[0], 0.0);
4021    }
4022}
4023
4024// ============================================================================
4025// SIMD-accelerated linear scan for sorted u32 slices (within-block seek)
4026// ============================================================================
4027
4028/// Find index of first element >= `target` in a sorted `u32` slice.
4029///
4030/// Equivalent to `slice.partition_point(|&d| d < target)` but uses SIMD to
4031/// scan 4 elements per cycle. Faster than binary search for slices ≤ 256
4032/// elements because it avoids the data-dependency chain inherent in binary
4033/// search (~8-10 cycles/iteration vs ~1-2 cycles/iteration for SIMD scan).
4034///
4035/// Returns `slice.len()` if no element >= `target`.
4036#[inline]
4037pub fn find_first_ge_u32(slice: &[u32], target: u32) -> usize {
4038    #[cfg(target_arch = "aarch64")]
4039    {
4040        if neon::is_available() {
4041            return unsafe { find_first_ge_u32_neon(slice, target) };
4042        }
4043    }
4044
4045    #[cfg(target_arch = "x86_64")]
4046    {
4047        if sse::is_available() {
4048            return unsafe { find_first_ge_u32_sse(slice, target) };
4049        }
4050    }
4051
4052    // Scalar fallback (WASM, other architectures)
4053    slice.partition_point(|&d| d < target)
4054}
4055
4056#[cfg(target_arch = "aarch64")]
4057#[target_feature(enable = "neon")]
4058#[allow(unsafe_op_in_unsafe_fn)]
4059unsafe fn find_first_ge_u32_neon(slice: &[u32], target: u32) -> usize {
4060    use std::arch::aarch64::*;
4061
4062    let n = slice.len();
4063    let ptr = slice.as_ptr();
4064    let target_vec = vdupq_n_u32(target);
4065    // Bit positions for each lane: [1, 2, 4, 8]
4066    let bit_mask: uint32x4_t = core::mem::transmute([1u32, 2u32, 4u32, 8u32]);
4067
4068    let chunks = n / 16;
4069    let mut base = 0usize;
4070
4071    // Process 16 elements per iteration (4 × 4-wide NEON compares)
4072    for _ in 0..chunks {
4073        let v0 = vld1q_u32(ptr.add(base));
4074        let v1 = vld1q_u32(ptr.add(base + 4));
4075        let v2 = vld1q_u32(ptr.add(base + 8));
4076        let v3 = vld1q_u32(ptr.add(base + 12));
4077
4078        let c0 = vcgeq_u32(v0, target_vec);
4079        let c1 = vcgeq_u32(v1, target_vec);
4080        let c2 = vcgeq_u32(v2, target_vec);
4081        let c3 = vcgeq_u32(v3, target_vec);
4082
4083        let m0 = vaddvq_u32(vandq_u32(c0, bit_mask));
4084        if m0 != 0 {
4085            return base + m0.trailing_zeros() as usize;
4086        }
4087        let m1 = vaddvq_u32(vandq_u32(c1, bit_mask));
4088        if m1 != 0 {
4089            return base + 4 + m1.trailing_zeros() as usize;
4090        }
4091        let m2 = vaddvq_u32(vandq_u32(c2, bit_mask));
4092        if m2 != 0 {
4093            return base + 8 + m2.trailing_zeros() as usize;
4094        }
4095        let m3 = vaddvq_u32(vandq_u32(c3, bit_mask));
4096        if m3 != 0 {
4097            return base + 12 + m3.trailing_zeros() as usize;
4098        }
4099        base += 16;
4100    }
4101
4102    // Process remaining 4 elements at a time
4103    while base + 4 <= n {
4104        let vals = vld1q_u32(ptr.add(base));
4105        let cmp = vcgeq_u32(vals, target_vec);
4106        let mask = vaddvq_u32(vandq_u32(cmp, bit_mask));
4107        if mask != 0 {
4108            return base + mask.trailing_zeros() as usize;
4109        }
4110        base += 4;
4111    }
4112
4113    // Scalar remainder (0-3 elements)
4114    while base < n {
4115        if *slice.get_unchecked(base) >= target {
4116            return base;
4117        }
4118        base += 1;
4119    }
4120    n
4121}
4122
4123#[cfg(target_arch = "x86_64")]
4124#[target_feature(enable = "sse2")]
4125#[allow(unsafe_op_in_unsafe_fn)]
4126unsafe fn find_first_ge_u32_sse(slice: &[u32], target: u32) -> usize {
4127    use std::arch::x86_64::*;
4128
4129    let n = slice.len();
4130    let ptr = slice.as_ptr();
4131
4132    // For unsigned >= comparison: XOR with 0x80000000 converts to signed domain
4133    let sign_flip = _mm_set1_epi32(i32::MIN);
4134    let target_xor = _mm_xor_si128(_mm_set1_epi32(target as i32), sign_flip);
4135
4136    let chunks = n / 16;
4137    let mut base = 0usize;
4138
4139    // Process 16 elements per iteration (4 × 4-wide SSE compares)
4140    for _ in 0..chunks {
4141        let v0 = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4142        let v1 = _mm_xor_si128(
4143            _mm_loadu_si128(ptr.add(base + 4) as *const __m128i),
4144            sign_flip,
4145        );
4146        let v2 = _mm_xor_si128(
4147            _mm_loadu_si128(ptr.add(base + 8) as *const __m128i),
4148            sign_flip,
4149        );
4150        let v3 = _mm_xor_si128(
4151            _mm_loadu_si128(ptr.add(base + 12) as *const __m128i),
4152            sign_flip,
4153        );
4154
4155        // ge = eq | gt (in signed domain after XOR)
4156        let ge0 = _mm_or_si128(
4157            _mm_cmpeq_epi32(v0, target_xor),
4158            _mm_cmpgt_epi32(v0, target_xor),
4159        );
4160        let m0 = _mm_movemask_ps(_mm_castsi128_ps(ge0)) as u32;
4161        if m0 != 0 {
4162            return base + m0.trailing_zeros() as usize;
4163        }
4164
4165        let ge1 = _mm_or_si128(
4166            _mm_cmpeq_epi32(v1, target_xor),
4167            _mm_cmpgt_epi32(v1, target_xor),
4168        );
4169        let m1 = _mm_movemask_ps(_mm_castsi128_ps(ge1)) as u32;
4170        if m1 != 0 {
4171            return base + 4 + m1.trailing_zeros() as usize;
4172        }
4173
4174        let ge2 = _mm_or_si128(
4175            _mm_cmpeq_epi32(v2, target_xor),
4176            _mm_cmpgt_epi32(v2, target_xor),
4177        );
4178        let m2 = _mm_movemask_ps(_mm_castsi128_ps(ge2)) as u32;
4179        if m2 != 0 {
4180            return base + 8 + m2.trailing_zeros() as usize;
4181        }
4182
4183        let ge3 = _mm_or_si128(
4184            _mm_cmpeq_epi32(v3, target_xor),
4185            _mm_cmpgt_epi32(v3, target_xor),
4186        );
4187        let m3 = _mm_movemask_ps(_mm_castsi128_ps(ge3)) as u32;
4188        if m3 != 0 {
4189            return base + 12 + m3.trailing_zeros() as usize;
4190        }
4191        base += 16;
4192    }
4193
4194    // Process remaining 4 elements at a time
4195    while base + 4 <= n {
4196        let vals = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4197        let ge = _mm_or_si128(
4198            _mm_cmpeq_epi32(vals, target_xor),
4199            _mm_cmpgt_epi32(vals, target_xor),
4200        );
4201        let mask = _mm_movemask_ps(_mm_castsi128_ps(ge)) as u32;
4202        if mask != 0 {
4203            return base + mask.trailing_zeros() as usize;
4204        }
4205        base += 4;
4206    }
4207
4208    // Scalar remainder (0-3 elements)
4209    while base < n {
4210        if *slice.get_unchecked(base) >= target {
4211            return base;
4212        }
4213        base += 1;
4214    }
4215    n
4216}
4217
4218#[cfg(test)]
4219mod find_first_ge_tests {
4220    use super::find_first_ge_u32;
4221
4222    #[test]
4223    fn test_find_first_ge_basic() {
4224        let data: Vec<u32> = (0..128).map(|i| i * 3).collect(); // [0, 3, 6, ..., 381]
4225        assert_eq!(find_first_ge_u32(&data, 0), 0);
4226        assert_eq!(find_first_ge_u32(&data, 1), 1); // first >= 1 is 3 at idx 1
4227        assert_eq!(find_first_ge_u32(&data, 3), 1);
4228        assert_eq!(find_first_ge_u32(&data, 4), 2); // first >= 4 is 6 at idx 2
4229        assert_eq!(find_first_ge_u32(&data, 381), 127);
4230        assert_eq!(find_first_ge_u32(&data, 382), 128); // past end
4231    }
4232
4233    #[test]
4234    fn test_find_first_ge_matches_partition_point() {
4235        let data: Vec<u32> = vec![1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75];
4236        for target in 0..80 {
4237            let expected = data.partition_point(|&d| d < target);
4238            let actual = find_first_ge_u32(&data, target);
4239            assert_eq!(actual, expected, "target={}", target);
4240        }
4241    }
4242
4243    #[test]
4244    fn test_find_first_ge_small_slices() {
4245        // Empty
4246        assert_eq!(find_first_ge_u32(&[], 5), 0);
4247        // Single element
4248        assert_eq!(find_first_ge_u32(&[10], 5), 0);
4249        assert_eq!(find_first_ge_u32(&[10], 10), 0);
4250        assert_eq!(find_first_ge_u32(&[10], 11), 1);
4251        // Three elements (< SIMD width)
4252        assert_eq!(find_first_ge_u32(&[2, 4, 6], 5), 2);
4253    }
4254
4255    #[test]
4256    fn test_find_first_ge_full_block() {
4257        // Simulate a full 128-entry block
4258        let data: Vec<u32> = (100..228).collect();
4259        assert_eq!(find_first_ge_u32(&data, 100), 0);
4260        assert_eq!(find_first_ge_u32(&data, 150), 50);
4261        assert_eq!(find_first_ge_u32(&data, 227), 127);
4262        assert_eq!(find_first_ge_u32(&data, 228), 128);
4263        assert_eq!(find_first_ge_u32(&data, 99), 0);
4264    }
4265
4266    #[test]
4267    fn test_find_first_ge_u32_max() {
4268        // Test with large u32 values (unsigned correctness)
4269        let data = vec![u32::MAX - 10, u32::MAX - 5, u32::MAX - 1, u32::MAX];
4270        assert_eq!(find_first_ge_u32(&data, u32::MAX - 10), 0);
4271        assert_eq!(find_first_ge_u32(&data, u32::MAX - 7), 1);
4272        assert_eq!(find_first_ge_u32(&data, u32::MAX), 3);
4273    }
4274}