1#[cfg(target_arch = "aarch64")]
18#[allow(unsafe_op_in_unsafe_fn)]
19mod neon {
20 use std::arch::aarch64::*;
21
22 #[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 let bytes = vld1q_u8(in_ptr);
34
35 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 let base = chunks * 16;
56 for i in 0..remainder {
57 output[base + i] = input[base + i] as u32;
58 }
59 }
60
61 #[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 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 #[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 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 #[inline]
115 #[target_feature(enable = "neon")]
116 unsafe fn prefix_sum_4(v: uint32x4_t) -> uint32x4_t {
117 let shifted1 = vextq_u32(vdupq_n_u32(0), v, 3);
120 let sum1 = vaddq_u32(v, shifted1);
121
122 let shifted2 = vextq_u32(vdupq_n_u32(0), sum1, 2);
125 vaddq_u32(sum1, shifted2)
126 }
127
128 #[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 let d = vld1q_u32(deltas[base..].as_ptr());
158 let gaps = vaddq_u32(d, ones);
159
160 let prefix = prefix_sum_4(gaps);
162
163 let result = vaddq_u32(prefix, carry);
165
166 vst1q_u32(output[base + 1..].as_mut_ptr(), result);
168
169 carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
171 }
172
173 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 #[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 #[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 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); let d = vmovl_u16(vget_low_u16(u16s)); let gaps = vaddq_u32(d, ones);
234
235 let prefix = prefix_sum_4(gaps);
237
238 let result = vaddq_u32(prefix, carry);
240
241 vst1q_u32(output[base + 1..].as_mut_ptr(), result);
243
244 carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
246 }
247
248 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 #[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 let vals = vld1_u16(in_ptr);
284 let d = vmovl_u16(vals);
285
286 let gaps = vaddq_u32(d, ones);
288
289 let prefix = prefix_sum_4(gaps);
291
292 let result = vaddq_u32(prefix, carry);
294
295 vst1q_u32(output[base + 1..].as_mut_ptr(), result);
297
298 carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
300 }
301
302 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 #[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 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 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 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 #[inline]
351 pub fn is_available() -> bool {
352 true
353 }
354}
355
356#[cfg(target_arch = "x86_64")]
361#[allow(unsafe_op_in_unsafe_fn)]
362mod sse {
363 use std::arch::x86_64::*;
364
365 #[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 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 #[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 #[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 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 #[inline]
449 #[target_feature(enable = "sse2")]
450 unsafe fn prefix_sum_4(v: __m128i) -> __m128i {
451 let shifted1 = _mm_slli_si128(v, 4);
454 let sum1 = _mm_add_epi32(v, shifted1);
455
456 let shifted2 = _mm_slli_si128(sum1, 8);
459 _mm_add_epi32(sum1, shifted2)
460 }
461
462 #[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 let d = _mm_loadu_si128(deltas[base..].as_ptr() as *const __m128i);
490 let gaps = _mm_add_epi32(d, ones);
491
492 let prefix = prefix_sum_4(gaps);
494
495 let result = _mm_add_epi32(prefix, carry);
497
498 _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
500
501 carry = _mm_shuffle_epi32(result, 0xFF); }
504
505 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 #[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 #[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 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 let gaps = _mm_add_epi32(d, ones);
565
566 let prefix = prefix_sum_4(gaps);
568
569 let result = _mm_add_epi32(prefix, carry);
571
572 _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
574
575 carry = _mm_shuffle_epi32(result, 0xFF);
577 }
578
579 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 #[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 let vals = _mm_loadl_epi64(in_ptr as *const __m128i); let d = _mm_cvtepu16_epi32(vals);
616
617 let gaps = _mm_add_epi32(d, ones);
619
620 let prefix = prefix_sum_4(gaps);
622
623 let result = _mm_add_epi32(prefix, carry);
625
626 _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
628
629 carry = _mm_shuffle_epi32(result, 0xFF);
631 }
632
633 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 #[inline]
646 pub fn is_available() -> bool {
647 is_x86_feature_detected!("sse4.1")
648 }
649}
650
651#[cfg(target_arch = "x86_64")]
656#[allow(unsafe_op_in_unsafe_fn)]
657mod avx2 {
658 use std::arch::x86_64::*;
659
660 #[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 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 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 let base = chunks * 32;
689 for i in 0..remainder {
690 output[base + i] = input[base + i] as u32;
691 }
692 }
693
694 #[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 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 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 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 #[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 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 #[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 #[inline]
773 #[target_feature(enable = "avx2")]
774 unsafe fn prefix_sum_8(v: __m256i) -> __m256i {
775 let s1 = _mm256_slli_si256(v, 4);
777 let r1 = _mm256_add_epi32(v, s1);
778
779 let s2 = _mm256_slli_si256(r1, 8);
781 let r2 = _mm256_add_epi32(r1, s2);
782
783 let lo_sum = _mm256_shuffle_epi32(r2, 0xFF);
786 let carry = _mm256_permute2x128_si256(lo_sum, lo_sum, 0x00);
788 let carry_hi = _mm256_blend_epi32::<0xF0>(_mm256_setzero_si256(), carry);
790 _mm256_add_epi32(r2, carry_hi)
791 }
792
793 #[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 let bytes = _mm_loadl_epi64(input.as_ptr().add(base) as *const __m128i);
818 let d = _mm256_cvtepu8_epi32(bytes);
819
820 let gaps = _mm256_add_epi32(d, ones);
822
823 let prefix = prefix_sum_8(gaps);
825
826 let result = _mm256_add_epi32(prefix, carry);
828
829 _mm256_storeu_si256(output[base + 1..].as_mut_ptr() as *mut __m256i, result);
831
832 carry = _mm256_permutevar8x32_epi32(result, broadcast_idx);
834 }
835
836 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 #[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 let vals = _mm_loadu_si128(in_ptr as *const __m128i);
873 let d = _mm256_cvtepu16_epi32(vals);
874
875 let gaps = _mm256_add_epi32(d, ones);
877
878 let prefix = prefix_sum_8(gaps);
880
881 let result = _mm256_add_epi32(prefix, carry);
883
884 _mm256_storeu_si256(output[base + 1..].as_mut_ptr() as *mut __m256i, result);
886
887 carry = _mm256_permutevar8x32_epi32(result, broadcast_idx);
889 }
890
891 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 #[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 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 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 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 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 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 #[inline]
956 pub fn is_available() -> bool {
957 is_x86_feature_detected!("avx2")
958 }
959}
960
961#[allow(dead_code)]
966mod scalar {
967 #[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 #[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 #[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 #[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 #[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#[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 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#[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 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#[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 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#[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#[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 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#[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#[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 #[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 #[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, }
1243 }
1244
1245 #[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 #[inline]
1258 pub fn as_u8(self) -> u8 {
1259 self as u8
1260 }
1261}
1262
1263#[inline]
1265pub fn round_bit_width(bits: u8) -> u8 {
1266 RoundedBitWidth::from_exact(bits).as_u8()
1267}
1268
1269#[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#[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#[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 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 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#[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 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#[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 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#[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 match bit_width {
1489 0 => {
1490 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 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 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 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#[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 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 let bytes = vld1q_u8(in_ptr);
1597
1598 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 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 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 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 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 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 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#[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 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 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 let shuf = _mm_shuffle_ps(acc, acc, 0b10_11_00_01); let sums = _mm_add_ps(acc, shuf); let shuf2 = _mm_movehl_ps(sums, sums); let final_sum = _mm_add_ss(sums, shuf2); let mut sum = _mm_cvtss_f32(final_sum);
1836
1837 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#[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 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 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 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#[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); y * (1.5 - half * y * y) }
2153
2154#[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 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#[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 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; }
2217
2218 if exp16 <= 0 {
2219 if exp16 < -10 {
2220 return sign as u16; }
2222 let shift = (1 - exp16) as u32;
2223 let m = (mantissa | 0x80_0000) >> shift;
2224 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 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 if rounded > 0x3FF {
2239 let exp16_inc = exp16 as u32 + 1;
2240 if exp16_inc >= 31 {
2241 return (sign | 0x7C00) as u16; }
2243 (sign | (exp16_inc << 10)) as u16
2244 } else {
2245 (sign | ((exp16 as u32) << 10) | rounded) as u16
2246 }
2247}
2248
2249#[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 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
2277const U8_SCALE: f32 = 127.5;
2282const U8_INV_SCALE: f32 = 1.0 / 127.5;
2283
2284#[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#[inline]
2292pub fn u8_to_f32(byte: u8) -> f32 {
2293 byte as f32 * U8_INV_SCALE - 1.0
2294}
2295
2296pub 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
2308pub 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#[cfg(target_arch = "aarch64")]
2321#[allow(unsafe_op_in_unsafe_fn)]
2322mod neon_quant {
2323 use std::arch::aarch64::*;
2324
2325 #[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 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 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 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 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 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 #[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 let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2406
2407 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 #[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 #[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#[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#[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 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 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 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 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#[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 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 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 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 let acc_dot = _mm256_add_ps(acc_dot0, acc_dot1);
2741 let acc_norm = _mm256_add_ps(acc_norm0, acc_norm1);
2742
2743 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#[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#[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#[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 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 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 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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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
3315pub 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 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 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 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 let input: Vec<u8> = (0..127).collect();
3460 let first_value = 1000u32;
3461 let count = 128;
3462
3463 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 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 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 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 #[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 assert!((output[0] - 1.0).abs() < 1e-6); assert!((output[1] - 13.8).abs() < 1e-6); assert!((output[2] - 26.5).abs() < 1e-6); assert!((output[3] - 7.4).abs() < 1e-6); assert!((output[4] - 20.2).abs() < 1e-6); }
3588
3589 #[test]
3590 fn test_dequantize_uint8_large() {
3591 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 assert!((result - 70.0).abs() < 1e-5);
3620 }
3621
3622 #[test]
3623 fn test_dot_product_f32_large() {
3624 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 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 let query = vec![1.0f32, 0.0, 0.0];
3688 let vectors = vec![
3689 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, -1.0, 0.0, 0.0, 0.5, 0.5, 0.0, ];
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 #[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 assert_eq!(f16_to_f32(f32_to_f16(0.0)), 0.0);
3772 assert_eq!(f32_to_f16(-0.0), 0x8000);
3774 assert!(f16_to_f32(f32_to_f16(f32::INFINITY)).is_infinite());
3776 assert!(f16_to_f32(f32_to_f16(f32::NAN)).is_nan());
3778 }
3779
3780 #[test]
3781 fn test_f16_embedding_range() {
3782 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 #[test]
3795 fn test_u8_roundtrip() {
3796 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); 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 #[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, 0.0, 0.0, 0.6, 0.8, ];
3825
3826 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, -0.6, -0.8, 0.0, 0.0, ];
3851
3852 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 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 assert!((scores[0] - 1.0).abs() < 0.01, "self-sim: {}", scores[0]);
3883 assert!(scores[1] > 0.99, "scaled-sim: {}", scores[1]);
3885 }
3886
3887 #[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 let a = vec![0xAA];
3908 let b = vec![0x55];
3909 assert_eq!(hamming_distance(&a, &b), 8);
3910
3911 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; 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 let a = vec![0xFF; 17];
3935 let b = vec![0x00; 17];
3936 assert_eq!(hamming_distance(&a, &b), 136); let a = vec![0xFF; 33];
3940 let b = vec![0x00; 33];
3941 assert_eq!(hamming_distance(&a, &b), 264); }
3943
3944 #[test]
3945 fn test_hamming_distance_large() {
3946 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 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 #[test]
3969 fn test_batch_hamming_scores_identical() {
3970 let query = vec![0xAA; 16];
3971 let db = vec![0xAA; 16]; 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]); db.extend_from_slice(&vec![0x00; byte_len]); db.extend_from_slice(&vec![0x0F; byte_len]); 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 assert_eq!(scores[0], 0.0);
4021 }
4022}
4023
4024#[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 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 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 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 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 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 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 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 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 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 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(); assert_eq!(find_first_ge_u32(&data, 0), 0);
4226 assert_eq!(find_first_ge_u32(&data, 1), 1); assert_eq!(find_first_ge_u32(&data, 3), 1);
4228 assert_eq!(find_first_ge_u32(&data, 4), 2); assert_eq!(find_first_ge_u32(&data, 381), 127);
4230 assert_eq!(find_first_ge_u32(&data, 382), 128); }
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 assert_eq!(find_first_ge_u32(&[], 5), 0);
4247 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 assert_eq!(find_first_ge_u32(&[2, 4, 6], 5), 2);
4253 }
4254
4255 #[test]
4256 fn test_find_first_ge_full_block() {
4257 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 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}