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 assert!(
1676 count <= a.len() && count <= b.len(),
1677 "dot_product_f32 count {count} exceeds input lengths ({}, {})",
1678 a.len(),
1679 b.len()
1680 );
1681 #[cfg(target_arch = "aarch64")]
1682 {
1683 if neon::is_available() {
1684 return unsafe { dot_product_f32_neon(a, b, count) };
1685 }
1686 }
1687
1688 #[cfg(target_arch = "x86_64")]
1689 {
1690 if is_x86_feature_detected!("avx512f") {
1691 return unsafe { dot_product_f32_avx512(a, b, count) };
1692 }
1693 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
1694 return unsafe { dot_product_f32_avx2(a, b, count) };
1695 }
1696 if sse::is_available() {
1697 return unsafe { dot_product_f32_sse(a, b, count) };
1698 }
1699 }
1700
1701 let mut sum = 0.0f32;
1703 for i in 0..count {
1704 sum += a[i] * b[i];
1705 }
1706 sum
1707}
1708
1709#[cfg(target_arch = "aarch64")]
1710#[target_feature(enable = "neon")]
1711#[allow(unsafe_op_in_unsafe_fn)]
1712unsafe fn dot_product_f32_neon(a: &[f32], b: &[f32], count: usize) -> f32 {
1713 use std::arch::aarch64::*;
1714
1715 let chunks16 = count / 16;
1716 let remainder = count % 16;
1717
1718 let mut acc0 = vdupq_n_f32(0.0);
1719 let mut acc1 = vdupq_n_f32(0.0);
1720 let mut acc2 = vdupq_n_f32(0.0);
1721 let mut acc3 = vdupq_n_f32(0.0);
1722
1723 for c in 0..chunks16 {
1724 let base = c * 16;
1725 acc0 = vfmaq_f32(
1726 acc0,
1727 vld1q_f32(a.as_ptr().add(base)),
1728 vld1q_f32(b.as_ptr().add(base)),
1729 );
1730 acc1 = vfmaq_f32(
1731 acc1,
1732 vld1q_f32(a.as_ptr().add(base + 4)),
1733 vld1q_f32(b.as_ptr().add(base + 4)),
1734 );
1735 acc2 = vfmaq_f32(
1736 acc2,
1737 vld1q_f32(a.as_ptr().add(base + 8)),
1738 vld1q_f32(b.as_ptr().add(base + 8)),
1739 );
1740 acc3 = vfmaq_f32(
1741 acc3,
1742 vld1q_f32(a.as_ptr().add(base + 12)),
1743 vld1q_f32(b.as_ptr().add(base + 12)),
1744 );
1745 }
1746
1747 let acc = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
1748 let mut sum = vaddvq_f32(acc);
1749
1750 let base = chunks16 * 16;
1751 for i in 0..remainder {
1752 sum += a[base + i] * b[base + i];
1753 }
1754
1755 sum
1756}
1757
1758#[cfg(target_arch = "x86_64")]
1759#[target_feature(enable = "avx2", enable = "fma")]
1760#[allow(unsafe_op_in_unsafe_fn)]
1761unsafe fn dot_product_f32_avx2(a: &[f32], b: &[f32], count: usize) -> f32 {
1762 use std::arch::x86_64::*;
1763
1764 let chunks32 = count / 32;
1765 let remainder = count % 32;
1766
1767 let mut acc0 = _mm256_setzero_ps();
1768 let mut acc1 = _mm256_setzero_ps();
1769 let mut acc2 = _mm256_setzero_ps();
1770 let mut acc3 = _mm256_setzero_ps();
1771
1772 for c in 0..chunks32 {
1773 let base = c * 32;
1774 acc0 = _mm256_fmadd_ps(
1775 _mm256_loadu_ps(a.as_ptr().add(base)),
1776 _mm256_loadu_ps(b.as_ptr().add(base)),
1777 acc0,
1778 );
1779 acc1 = _mm256_fmadd_ps(
1780 _mm256_loadu_ps(a.as_ptr().add(base + 8)),
1781 _mm256_loadu_ps(b.as_ptr().add(base + 8)),
1782 acc1,
1783 );
1784 acc2 = _mm256_fmadd_ps(
1785 _mm256_loadu_ps(a.as_ptr().add(base + 16)),
1786 _mm256_loadu_ps(b.as_ptr().add(base + 16)),
1787 acc2,
1788 );
1789 acc3 = _mm256_fmadd_ps(
1790 _mm256_loadu_ps(a.as_ptr().add(base + 24)),
1791 _mm256_loadu_ps(b.as_ptr().add(base + 24)),
1792 acc3,
1793 );
1794 }
1795
1796 let acc = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
1797
1798 let hi = _mm256_extractf128_ps(acc, 1);
1800 let lo = _mm256_castps256_ps128(acc);
1801 let sum128 = _mm_add_ps(lo, hi);
1802 let shuf = _mm_shuffle_ps(sum128, sum128, 0b10_11_00_01);
1803 let sums = _mm_add_ps(sum128, shuf);
1804 let shuf2 = _mm_movehl_ps(sums, sums);
1805 let final_sum = _mm_add_ss(sums, shuf2);
1806
1807 let mut sum = _mm_cvtss_f32(final_sum);
1808
1809 let base = chunks32 * 32;
1810 for i in 0..remainder {
1811 sum += a[base + i] * b[base + i];
1812 }
1813
1814 sum
1815}
1816
1817#[cfg(target_arch = "x86_64")]
1818#[target_feature(enable = "sse")]
1819#[allow(unsafe_op_in_unsafe_fn)]
1820unsafe fn dot_product_f32_sse(a: &[f32], b: &[f32], count: usize) -> f32 {
1821 use std::arch::x86_64::*;
1822
1823 let chunks = count / 4;
1824 let remainder = count % 4;
1825
1826 let mut acc = _mm_setzero_ps();
1827
1828 for chunk in 0..chunks {
1829 let base = chunk * 4;
1830 let va = _mm_loadu_ps(a.as_ptr().add(base));
1831 let vb = _mm_loadu_ps(b.as_ptr().add(base));
1832 acc = _mm_add_ps(acc, _mm_mul_ps(va, vb));
1833 }
1834
1835 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);
1842
1843 let base = chunks * 4;
1845 for i in 0..remainder {
1846 sum += a[base + i] * b[base + i];
1847 }
1848
1849 sum
1850}
1851
1852#[cfg(target_arch = "x86_64")]
1853#[target_feature(enable = "avx512f")]
1854#[allow(unsafe_op_in_unsafe_fn)]
1855unsafe fn dot_product_f32_avx512(a: &[f32], b: &[f32], count: usize) -> f32 {
1856 use std::arch::x86_64::*;
1857
1858 let chunks64 = count / 64;
1859 let remainder = count % 64;
1860
1861 let mut acc0 = _mm512_setzero_ps();
1862 let mut acc1 = _mm512_setzero_ps();
1863 let mut acc2 = _mm512_setzero_ps();
1864 let mut acc3 = _mm512_setzero_ps();
1865
1866 for c in 0..chunks64 {
1867 let base = c * 64;
1868 acc0 = _mm512_fmadd_ps(
1869 _mm512_loadu_ps(a.as_ptr().add(base)),
1870 _mm512_loadu_ps(b.as_ptr().add(base)),
1871 acc0,
1872 );
1873 acc1 = _mm512_fmadd_ps(
1874 _mm512_loadu_ps(a.as_ptr().add(base + 16)),
1875 _mm512_loadu_ps(b.as_ptr().add(base + 16)),
1876 acc1,
1877 );
1878 acc2 = _mm512_fmadd_ps(
1879 _mm512_loadu_ps(a.as_ptr().add(base + 32)),
1880 _mm512_loadu_ps(b.as_ptr().add(base + 32)),
1881 acc2,
1882 );
1883 acc3 = _mm512_fmadd_ps(
1884 _mm512_loadu_ps(a.as_ptr().add(base + 48)),
1885 _mm512_loadu_ps(b.as_ptr().add(base + 48)),
1886 acc3,
1887 );
1888 }
1889
1890 let acc = _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3));
1891 let mut sum = _mm512_reduce_add_ps(acc);
1892
1893 let base = chunks64 * 64;
1894 for i in 0..remainder {
1895 sum += a[base + i] * b[base + i];
1896 }
1897
1898 sum
1899}
1900
1901#[cfg(target_arch = "x86_64")]
1902#[target_feature(enable = "avx512f")]
1903#[allow(unsafe_op_in_unsafe_fn)]
1904unsafe fn fused_dot_norm_avx512(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1905 use std::arch::x86_64::*;
1906
1907 let chunks64 = count / 64;
1908 let remainder = count % 64;
1909
1910 let mut d0 = _mm512_setzero_ps();
1911 let mut d1 = _mm512_setzero_ps();
1912 let mut d2 = _mm512_setzero_ps();
1913 let mut d3 = _mm512_setzero_ps();
1914 let mut n0 = _mm512_setzero_ps();
1915 let mut n1 = _mm512_setzero_ps();
1916 let mut n2 = _mm512_setzero_ps();
1917 let mut n3 = _mm512_setzero_ps();
1918
1919 for c in 0..chunks64 {
1920 let base = c * 64;
1921 let vb0 = _mm512_loadu_ps(b.as_ptr().add(base));
1922 d0 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base)), vb0, d0);
1923 n0 = _mm512_fmadd_ps(vb0, vb0, n0);
1924 let vb1 = _mm512_loadu_ps(b.as_ptr().add(base + 16));
1925 d1 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 16)), vb1, d1);
1926 n1 = _mm512_fmadd_ps(vb1, vb1, n1);
1927 let vb2 = _mm512_loadu_ps(b.as_ptr().add(base + 32));
1928 d2 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 32)), vb2, d2);
1929 n2 = _mm512_fmadd_ps(vb2, vb2, n2);
1930 let vb3 = _mm512_loadu_ps(b.as_ptr().add(base + 48));
1931 d3 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 48)), vb3, d3);
1932 n3 = _mm512_fmadd_ps(vb3, vb3, n3);
1933 }
1934
1935 let acc_dot = _mm512_add_ps(_mm512_add_ps(d0, d1), _mm512_add_ps(d2, d3));
1936 let acc_norm = _mm512_add_ps(_mm512_add_ps(n0, n1), _mm512_add_ps(n2, n3));
1937 let mut dot = _mm512_reduce_add_ps(acc_dot);
1938 let mut norm = _mm512_reduce_add_ps(acc_norm);
1939
1940 let base = chunks64 * 64;
1941 for i in 0..remainder {
1942 dot += a[base + i] * b[base + i];
1943 norm += b[base + i] * b[base + i];
1944 }
1945
1946 (dot, norm)
1947}
1948
1949#[inline]
1958fn fused_dot_norm(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1959 #[cfg(target_arch = "aarch64")]
1960 {
1961 if neon::is_available() {
1962 return unsafe { fused_dot_norm_neon(a, b, count) };
1963 }
1964 }
1965
1966 #[cfg(target_arch = "x86_64")]
1967 {
1968 if is_x86_feature_detected!("avx512f") {
1969 return unsafe { fused_dot_norm_avx512(a, b, count) };
1970 }
1971 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
1972 return unsafe { fused_dot_norm_avx2(a, b, count) };
1973 }
1974 if sse::is_available() {
1975 return unsafe { fused_dot_norm_sse(a, b, count) };
1976 }
1977 }
1978
1979 let mut dot = 0.0f32;
1981 let mut norm_b = 0.0f32;
1982 for i in 0..count {
1983 dot += a[i] * b[i];
1984 norm_b += b[i] * b[i];
1985 }
1986 (dot, norm_b)
1987}
1988
1989#[cfg(target_arch = "aarch64")]
1990#[target_feature(enable = "neon")]
1991#[allow(unsafe_op_in_unsafe_fn)]
1992unsafe fn fused_dot_norm_neon(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1993 use std::arch::aarch64::*;
1994
1995 let chunks16 = count / 16;
1996 let remainder = count % 16;
1997
1998 let mut d0 = vdupq_n_f32(0.0);
1999 let mut d1 = vdupq_n_f32(0.0);
2000 let mut d2 = vdupq_n_f32(0.0);
2001 let mut d3 = vdupq_n_f32(0.0);
2002 let mut n0 = vdupq_n_f32(0.0);
2003 let mut n1 = vdupq_n_f32(0.0);
2004 let mut n2 = vdupq_n_f32(0.0);
2005 let mut n3 = vdupq_n_f32(0.0);
2006
2007 for c in 0..chunks16 {
2008 let base = c * 16;
2009 let va0 = vld1q_f32(a.as_ptr().add(base));
2010 let vb0 = vld1q_f32(b.as_ptr().add(base));
2011 d0 = vfmaq_f32(d0, va0, vb0);
2012 n0 = vfmaq_f32(n0, vb0, vb0);
2013 let va1 = vld1q_f32(a.as_ptr().add(base + 4));
2014 let vb1 = vld1q_f32(b.as_ptr().add(base + 4));
2015 d1 = vfmaq_f32(d1, va1, vb1);
2016 n1 = vfmaq_f32(n1, vb1, vb1);
2017 let va2 = vld1q_f32(a.as_ptr().add(base + 8));
2018 let vb2 = vld1q_f32(b.as_ptr().add(base + 8));
2019 d2 = vfmaq_f32(d2, va2, vb2);
2020 n2 = vfmaq_f32(n2, vb2, vb2);
2021 let va3 = vld1q_f32(a.as_ptr().add(base + 12));
2022 let vb3 = vld1q_f32(b.as_ptr().add(base + 12));
2023 d3 = vfmaq_f32(d3, va3, vb3);
2024 n3 = vfmaq_f32(n3, vb3, vb3);
2025 }
2026
2027 let acc_dot = vaddq_f32(vaddq_f32(d0, d1), vaddq_f32(d2, d3));
2028 let acc_norm = vaddq_f32(vaddq_f32(n0, n1), vaddq_f32(n2, n3));
2029 let mut dot = vaddvq_f32(acc_dot);
2030 let mut norm = vaddvq_f32(acc_norm);
2031
2032 let base = chunks16 * 16;
2033 for i in 0..remainder {
2034 dot += a[base + i] * b[base + i];
2035 norm += b[base + i] * b[base + i];
2036 }
2037
2038 (dot, norm)
2039}
2040
2041#[cfg(target_arch = "x86_64")]
2042#[target_feature(enable = "avx2", enable = "fma")]
2043#[allow(unsafe_op_in_unsafe_fn)]
2044unsafe fn fused_dot_norm_avx2(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
2045 use std::arch::x86_64::*;
2046
2047 let chunks32 = count / 32;
2048 let remainder = count % 32;
2049
2050 let mut d0 = _mm256_setzero_ps();
2051 let mut d1 = _mm256_setzero_ps();
2052 let mut d2 = _mm256_setzero_ps();
2053 let mut d3 = _mm256_setzero_ps();
2054 let mut n0 = _mm256_setzero_ps();
2055 let mut n1 = _mm256_setzero_ps();
2056 let mut n2 = _mm256_setzero_ps();
2057 let mut n3 = _mm256_setzero_ps();
2058
2059 for c in 0..chunks32 {
2060 let base = c * 32;
2061 let vb0 = _mm256_loadu_ps(b.as_ptr().add(base));
2062 d0 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base)), vb0, d0);
2063 n0 = _mm256_fmadd_ps(vb0, vb0, n0);
2064 let vb1 = _mm256_loadu_ps(b.as_ptr().add(base + 8));
2065 d1 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 8)), vb1, d1);
2066 n1 = _mm256_fmadd_ps(vb1, vb1, n1);
2067 let vb2 = _mm256_loadu_ps(b.as_ptr().add(base + 16));
2068 d2 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 16)), vb2, d2);
2069 n2 = _mm256_fmadd_ps(vb2, vb2, n2);
2070 let vb3 = _mm256_loadu_ps(b.as_ptr().add(base + 24));
2071 d3 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 24)), vb3, d3);
2072 n3 = _mm256_fmadd_ps(vb3, vb3, n3);
2073 }
2074
2075 let acc_dot = _mm256_add_ps(_mm256_add_ps(d0, d1), _mm256_add_ps(d2, d3));
2076 let acc_norm = _mm256_add_ps(_mm256_add_ps(n0, n1), _mm256_add_ps(n2, n3));
2077
2078 let hi_d = _mm256_extractf128_ps(acc_dot, 1);
2080 let lo_d = _mm256_castps256_ps128(acc_dot);
2081 let sum_d = _mm_add_ps(lo_d, hi_d);
2082 let shuf_d = _mm_shuffle_ps(sum_d, sum_d, 0b10_11_00_01);
2083 let sums_d = _mm_add_ps(sum_d, shuf_d);
2084 let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2085 let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2086
2087 let hi_n = _mm256_extractf128_ps(acc_norm, 1);
2088 let lo_n = _mm256_castps256_ps128(acc_norm);
2089 let sum_n = _mm_add_ps(lo_n, hi_n);
2090 let shuf_n = _mm_shuffle_ps(sum_n, sum_n, 0b10_11_00_01);
2091 let sums_n = _mm_add_ps(sum_n, shuf_n);
2092 let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2093 let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2094
2095 let base = chunks32 * 32;
2096 for i in 0..remainder {
2097 dot += a[base + i] * b[base + i];
2098 norm += b[base + i] * b[base + i];
2099 }
2100
2101 (dot, norm)
2102}
2103
2104#[cfg(target_arch = "x86_64")]
2105#[target_feature(enable = "sse")]
2106#[allow(unsafe_op_in_unsafe_fn)]
2107unsafe fn fused_dot_norm_sse(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
2108 use std::arch::x86_64::*;
2109
2110 let chunks = count / 4;
2111 let remainder = count % 4;
2112
2113 let mut acc_dot = _mm_setzero_ps();
2114 let mut acc_norm = _mm_setzero_ps();
2115
2116 for chunk in 0..chunks {
2117 let base = chunk * 4;
2118 let va = _mm_loadu_ps(a.as_ptr().add(base));
2119 let vb = _mm_loadu_ps(b.as_ptr().add(base));
2120 acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2121 acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2122 }
2123
2124 let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2126 let sums_d = _mm_add_ps(acc_dot, shuf_d);
2127 let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2128 let final_d = _mm_add_ss(sums_d, shuf2_d);
2129 let mut dot = _mm_cvtss_f32(final_d);
2130
2131 let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2132 let sums_n = _mm_add_ps(acc_norm, shuf_n);
2133 let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2134 let final_n = _mm_add_ss(sums_n, shuf2_n);
2135 let mut norm = _mm_cvtss_f32(final_n);
2136
2137 let base = chunks * 4;
2138 for i in 0..remainder {
2139 dot += a[base + i] * b[base + i];
2140 norm += b[base + i] * b[base + i];
2141 }
2142
2143 (dot, norm)
2144}
2145
2146#[inline]
2152pub fn fast_inv_sqrt(x: f32) -> f32 {
2153 let half = 0.5 * x;
2154 let i = 0x5F37_5A86_u32.wrapping_sub(x.to_bits() >> 1);
2155 let y = f32::from_bits(i);
2156 let y = y * (1.5 - half * y * y); y * (1.5 - half * y * y) }
2159
2160#[inline]
2171pub fn batch_cosine_scores(query: &[f32], vectors: &[f32], dim: usize, scores: &mut [f32]) {
2172 let n = scores.len();
2173 let required = n
2174 .checked_mul(dim)
2175 .expect("batch cosine vector length overflow");
2176 assert_eq!(query.len(), dim, "batch cosine query dimension mismatch");
2177 assert!(
2178 vectors.len() >= required,
2179 "batch cosine vectors are truncated: need {required}, got {}",
2180 vectors.len()
2181 );
2182
2183 if dim == 0 || n == 0 {
2184 return;
2185 }
2186
2187 let norm_q_sq = dot_product_f32(query, query, dim);
2189 if norm_q_sq < f32::EPSILON {
2190 for s in scores.iter_mut() {
2191 *s = 0.0;
2192 }
2193 return;
2194 }
2195 let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2196
2197 for i in 0..n {
2198 let vec = &vectors[i * dim..(i + 1) * dim];
2199 let (dot, norm_v_sq) = fused_dot_norm(query, vec, dim);
2200 if norm_v_sq < f32::EPSILON {
2201 scores[i] = 0.0;
2202 } else {
2203 scores[i] = dot * inv_norm_q * fast_inv_sqrt(norm_v_sq);
2204 }
2205 }
2206}
2207
2208#[inline]
2214pub fn f32_to_f16(value: f32) -> u16 {
2215 let bits = value.to_bits();
2216 let sign = (bits >> 16) & 0x8000;
2217 let exp = ((bits >> 23) & 0xFF) as i32;
2218 let mantissa = bits & 0x7F_FFFF;
2219
2220 if exp == 255 {
2221 return (sign | 0x7C00 | ((mantissa >> 13) & 0x3FF)) as u16;
2223 }
2224
2225 let exp16 = exp - 127 + 15;
2226
2227 if exp16 >= 31 {
2228 return (sign | 0x7C00) as u16; }
2230
2231 if exp16 <= 0 {
2232 if exp16 < -10 {
2233 return sign as u16; }
2235 let shift = (1 - exp16) as u32;
2236 let m = (mantissa | 0x80_0000) >> shift;
2237 let round_bit = (m >> 12) & 1;
2239 let sticky = m & 0xFFF;
2240 let m13 = m >> 13;
2241 let rounded = m13 + (round_bit & (m13 | if sticky != 0 { 1 } else { 0 }));
2242 return (sign | rounded) as u16;
2243 }
2244
2245 let round_bit = (mantissa >> 12) & 1;
2247 let sticky = mantissa & 0xFFF;
2248 let m13 = mantissa >> 13;
2249 let rounded = m13 + (round_bit & (m13 | if sticky != 0 { 1 } else { 0 }));
2250 if rounded > 0x3FF {
2252 let exp16_inc = exp16 as u32 + 1;
2253 if exp16_inc >= 31 {
2254 return (sign | 0x7C00) as u16; }
2256 (sign | (exp16_inc << 10)) as u16
2257 } else {
2258 (sign | ((exp16 as u32) << 10) | rounded) as u16
2259 }
2260}
2261
2262#[inline]
2264pub fn f16_to_f32(half: u16) -> f32 {
2265 let sign = ((half & 0x8000) as u32) << 16;
2266 let exp = ((half >> 10) & 0x1F) as u32;
2267 let mantissa = (half & 0x3FF) as u32;
2268
2269 if exp == 0 {
2270 if mantissa == 0 {
2271 return f32::from_bits(sign);
2272 }
2273 let mut e = 0u32;
2275 let mut m = mantissa;
2276 while (m & 0x400) == 0 {
2277 m <<= 1;
2278 e += 1;
2279 }
2280 return f32::from_bits(sign | ((127 - 15 + 1 - e) << 23) | ((m & 0x3FF) << 13));
2281 }
2282
2283 if exp == 31 {
2284 return f32::from_bits(sign | 0x7F80_0000 | (mantissa << 13));
2285 }
2286
2287 f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mantissa << 13))
2288}
2289
2290const U8_SCALE: f32 = 127.5;
2295const U8_INV_SCALE: f32 = 1.0 / 127.5;
2296
2297#[inline]
2299pub fn f32_to_u8_saturating(value: f32) -> u8 {
2300 ((value.clamp(-1.0, 1.0) + 1.0) * U8_SCALE) as u8
2301}
2302
2303#[inline]
2305pub fn u8_to_f32(byte: u8) -> f32 {
2306 byte as f32 * U8_INV_SCALE - 1.0
2307}
2308
2309pub fn batch_f32_to_f16(src: &[f32], dst: &mut [u16]) {
2315 debug_assert_eq!(src.len(), dst.len());
2316 for (s, d) in src.iter().zip(dst.iter_mut()) {
2317 *d = f32_to_f16(*s);
2318 }
2319}
2320
2321pub fn batch_f32_to_u8(src: &[f32], dst: &mut [u8]) {
2323 debug_assert_eq!(src.len(), dst.len());
2324 for (s, d) in src.iter().zip(dst.iter_mut()) {
2325 *d = f32_to_u8_saturating(*s);
2326 }
2327}
2328
2329#[cfg(target_arch = "aarch64")]
2334#[allow(unsafe_op_in_unsafe_fn)]
2335mod neon_quant {
2336 use std::arch::aarch64::*;
2337
2338 #[allow(clippy::incompatible_msrv)]
2344 #[target_feature(enable = "neon")]
2345 pub unsafe fn fused_dot_norm_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2346 let chunks16 = dim / 16;
2347 let remainder = dim % 16;
2348
2349 let mut acc_dot0 = vdupq_n_f32(0.0);
2351 let mut acc_dot1 = vdupq_n_f32(0.0);
2352 let mut acc_norm0 = vdupq_n_f32(0.0);
2353 let mut acc_norm1 = vdupq_n_f32(0.0);
2354
2355 for c in 0..chunks16 {
2356 let base = c * 16;
2357
2358 let v_raw0 = vld1q_u16(vec_f16.as_ptr().add(base));
2360 let v_lo0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw0)));
2361 let v_hi0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw0)));
2362 let q_raw0 = vld1q_u16(query_f16.as_ptr().add(base));
2363 let q_lo0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw0)));
2364 let q_hi0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw0)));
2365
2366 acc_dot0 = vfmaq_f32(acc_dot0, q_lo0, v_lo0);
2367 acc_dot0 = vfmaq_f32(acc_dot0, q_hi0, v_hi0);
2368 acc_norm0 = vfmaq_f32(acc_norm0, v_lo0, v_lo0);
2369 acc_norm0 = vfmaq_f32(acc_norm0, v_hi0, v_hi0);
2370
2371 let v_raw1 = vld1q_u16(vec_f16.as_ptr().add(base + 8));
2373 let v_lo1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw1)));
2374 let v_hi1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw1)));
2375 let q_raw1 = vld1q_u16(query_f16.as_ptr().add(base + 8));
2376 let q_lo1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw1)));
2377 let q_hi1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw1)));
2378
2379 acc_dot1 = vfmaq_f32(acc_dot1, q_lo1, v_lo1);
2380 acc_dot1 = vfmaq_f32(acc_dot1, q_hi1, v_hi1);
2381 acc_norm1 = vfmaq_f32(acc_norm1, v_lo1, v_lo1);
2382 acc_norm1 = vfmaq_f32(acc_norm1, v_hi1, v_hi1);
2383 }
2384
2385 let mut dot = vaddvq_f32(vaddq_f32(acc_dot0, acc_dot1));
2387 let mut norm = vaddvq_f32(vaddq_f32(acc_norm0, acc_norm1));
2388
2389 let base = chunks16 * 16;
2391 for i in 0..remainder {
2392 let v = super::f16_to_f32(*vec_f16.get_unchecked(base + i));
2393 let q = super::f16_to_f32(*query_f16.get_unchecked(base + i));
2394 dot += q * v;
2395 norm += v * v;
2396 }
2397
2398 (dot, norm)
2399 }
2400
2401 #[target_feature(enable = "neon")]
2404 pub unsafe fn fused_dot_norm_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2405 let scale = vdupq_n_f32(super::U8_INV_SCALE);
2406 let offset = vdupq_n_f32(-1.0);
2407
2408 let chunks16 = dim / 16;
2409 let remainder = dim % 16;
2410
2411 let mut acc_dot = vdupq_n_f32(0.0);
2412 let mut acc_norm = vdupq_n_f32(0.0);
2413
2414 for c in 0..chunks16 {
2415 let base = c * 16;
2416
2417 let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2419
2420 let lo8 = vget_low_u8(bytes);
2422 let hi8 = vget_high_u8(bytes);
2423 let lo16 = vmovl_u8(lo8);
2424 let hi16 = vmovl_u8(hi8);
2425
2426 let f0 = vaddq_f32(
2427 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo16))), scale),
2428 offset,
2429 );
2430 let f1 = vaddq_f32(
2431 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo16))), scale),
2432 offset,
2433 );
2434 let f2 = vaddq_f32(
2435 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi16))), scale),
2436 offset,
2437 );
2438 let f3 = vaddq_f32(
2439 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi16))), scale),
2440 offset,
2441 );
2442
2443 let q0 = vld1q_f32(query.as_ptr().add(base));
2444 let q1 = vld1q_f32(query.as_ptr().add(base + 4));
2445 let q2 = vld1q_f32(query.as_ptr().add(base + 8));
2446 let q3 = vld1q_f32(query.as_ptr().add(base + 12));
2447
2448 acc_dot = vfmaq_f32(acc_dot, q0, f0);
2449 acc_dot = vfmaq_f32(acc_dot, q1, f1);
2450 acc_dot = vfmaq_f32(acc_dot, q2, f2);
2451 acc_dot = vfmaq_f32(acc_dot, q3, f3);
2452
2453 acc_norm = vfmaq_f32(acc_norm, f0, f0);
2454 acc_norm = vfmaq_f32(acc_norm, f1, f1);
2455 acc_norm = vfmaq_f32(acc_norm, f2, f2);
2456 acc_norm = vfmaq_f32(acc_norm, f3, f3);
2457 }
2458
2459 let mut dot = vaddvq_f32(acc_dot);
2460 let mut norm = vaddvq_f32(acc_norm);
2461
2462 let base = chunks16 * 16;
2463 for i in 0..remainder {
2464 let v = super::u8_to_f32(*vec_u8.get_unchecked(base + i));
2465 dot += *query.get_unchecked(base + i) * v;
2466 norm += v * v;
2467 }
2468
2469 (dot, norm)
2470 }
2471
2472 #[allow(clippy::incompatible_msrv)]
2474 #[target_feature(enable = "neon")]
2475 pub unsafe fn dot_product_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2476 let chunks8 = dim / 8;
2477 let remainder = dim % 8;
2478
2479 let mut acc = vdupq_n_f32(0.0);
2480
2481 for c in 0..chunks8 {
2482 let base = c * 8;
2483 let v_raw = vld1q_u16(vec_f16.as_ptr().add(base));
2484 let v_lo = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw)));
2485 let v_hi = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw)));
2486 let q_raw = vld1q_u16(query_f16.as_ptr().add(base));
2487 let q_lo = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw)));
2488 let q_hi = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw)));
2489 acc = vfmaq_f32(acc, q_lo, v_lo);
2490 acc = vfmaq_f32(acc, q_hi, v_hi);
2491 }
2492
2493 let mut dot = vaddvq_f32(acc);
2494 let base = chunks8 * 8;
2495 for i in 0..remainder {
2496 let v = super::f16_to_f32(*vec_f16.get_unchecked(base + i));
2497 let q = super::f16_to_f32(*query_f16.get_unchecked(base + i));
2498 dot += q * v;
2499 }
2500 dot
2501 }
2502
2503 #[target_feature(enable = "neon")]
2505 pub unsafe fn dot_product_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2506 let scale = vdupq_n_f32(super::U8_INV_SCALE);
2507 let offset = vdupq_n_f32(-1.0);
2508 let chunks16 = dim / 16;
2509 let remainder = dim % 16;
2510
2511 let mut acc = vdupq_n_f32(0.0);
2512
2513 for c in 0..chunks16 {
2514 let base = c * 16;
2515 let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2516 let lo8 = vget_low_u8(bytes);
2517 let hi8 = vget_high_u8(bytes);
2518 let lo16 = vmovl_u8(lo8);
2519 let hi16 = vmovl_u8(hi8);
2520 let f0 = vaddq_f32(
2521 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo16))), scale),
2522 offset,
2523 );
2524 let f1 = vaddq_f32(
2525 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo16))), scale),
2526 offset,
2527 );
2528 let f2 = vaddq_f32(
2529 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi16))), scale),
2530 offset,
2531 );
2532 let f3 = vaddq_f32(
2533 vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi16))), scale),
2534 offset,
2535 );
2536 let q0 = vld1q_f32(query.as_ptr().add(base));
2537 let q1 = vld1q_f32(query.as_ptr().add(base + 4));
2538 let q2 = vld1q_f32(query.as_ptr().add(base + 8));
2539 let q3 = vld1q_f32(query.as_ptr().add(base + 12));
2540 acc = vfmaq_f32(acc, q0, f0);
2541 acc = vfmaq_f32(acc, q1, f1);
2542 acc = vfmaq_f32(acc, q2, f2);
2543 acc = vfmaq_f32(acc, q3, f3);
2544 }
2545
2546 let mut dot = vaddvq_f32(acc);
2547 let base = chunks16 * 16;
2548 for i in 0..remainder {
2549 let v = super::u8_to_f32(*vec_u8.get_unchecked(base + i));
2550 dot += *query.get_unchecked(base + i) * v;
2551 }
2552 dot
2553 }
2554}
2555
2556#[allow(dead_code)]
2561fn fused_dot_norm_f16_scalar(query_f16: &[u16], vec_f16: &[u16], 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 = f16_to_f32(vec_f16[i]);
2566 let q = f16_to_f32(query_f16[i]);
2567 dot += q * v;
2568 norm += v * v;
2569 }
2570 (dot, norm)
2571}
2572
2573#[allow(dead_code)]
2574fn fused_dot_norm_u8_scalar(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2575 let mut dot = 0.0f32;
2576 let mut norm = 0.0f32;
2577 for i in 0..dim {
2578 let v = u8_to_f32(vec_u8[i]);
2579 dot += query[i] * v;
2580 norm += v * v;
2581 }
2582 (dot, norm)
2583}
2584
2585#[allow(dead_code)]
2586fn dot_product_f16_scalar(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2587 let mut dot = 0.0f32;
2588 for i in 0..dim {
2589 dot += f16_to_f32(query_f16[i]) * f16_to_f32(vec_f16[i]);
2590 }
2591 dot
2592}
2593
2594#[allow(dead_code)]
2595fn dot_product_u8_scalar(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2596 let mut dot = 0.0f32;
2597 for i in 0..dim {
2598 dot += query[i] * u8_to_f32(vec_u8[i]);
2599 }
2600 dot
2601}
2602
2603#[cfg(target_arch = "x86_64")]
2608#[target_feature(enable = "sse2", enable = "sse4.1")]
2609#[allow(unsafe_op_in_unsafe_fn)]
2610unsafe fn fused_dot_norm_f16_sse(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2611 use std::arch::x86_64::*;
2612
2613 let chunks = dim / 4;
2614 let remainder = dim % 4;
2615
2616 let mut acc_dot = _mm_setzero_ps();
2617 let mut acc_norm = _mm_setzero_ps();
2618
2619 for chunk in 0..chunks {
2620 let base = chunk * 4;
2621 let v0 = f16_to_f32(*vec_f16.get_unchecked(base));
2623 let v1 = f16_to_f32(*vec_f16.get_unchecked(base + 1));
2624 let v2 = f16_to_f32(*vec_f16.get_unchecked(base + 2));
2625 let v3 = f16_to_f32(*vec_f16.get_unchecked(base + 3));
2626 let vb = _mm_set_ps(v3, v2, v1, v0);
2627
2628 let q0 = f16_to_f32(*query_f16.get_unchecked(base));
2629 let q1 = f16_to_f32(*query_f16.get_unchecked(base + 1));
2630 let q2 = f16_to_f32(*query_f16.get_unchecked(base + 2));
2631 let q3 = f16_to_f32(*query_f16.get_unchecked(base + 3));
2632 let va = _mm_set_ps(q3, q2, q1, q0);
2633
2634 acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2635 acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2636 }
2637
2638 let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2640 let sums_d = _mm_add_ps(acc_dot, shuf_d);
2641 let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2642 let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2643
2644 let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2645 let sums_n = _mm_add_ps(acc_norm, shuf_n);
2646 let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2647 let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2648
2649 let base = chunks * 4;
2650 for i in 0..remainder {
2651 let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2652 let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2653 dot += q * v;
2654 norm += v * v;
2655 }
2656
2657 (dot, norm)
2658}
2659
2660#[cfg(target_arch = "x86_64")]
2661#[target_feature(enable = "sse2", enable = "sse4.1")]
2662#[allow(unsafe_op_in_unsafe_fn)]
2663unsafe fn fused_dot_norm_u8_sse(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2664 use std::arch::x86_64::*;
2665
2666 let scale = _mm_set1_ps(U8_INV_SCALE);
2667 let offset = _mm_set1_ps(-1.0);
2668
2669 let chunks = dim / 4;
2670 let remainder = dim % 4;
2671
2672 let mut acc_dot = _mm_setzero_ps();
2673 let mut acc_norm = _mm_setzero_ps();
2674
2675 for chunk in 0..chunks {
2676 let base = chunk * 4;
2677
2678 let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
2680 vec_u8.as_ptr().add(base) as *const i32
2681 ));
2682 let ints = _mm_cvtepu8_epi32(bytes);
2683 let floats = _mm_cvtepi32_ps(ints);
2684 let vb = _mm_add_ps(_mm_mul_ps(floats, scale), offset);
2685
2686 let va = _mm_loadu_ps(query.as_ptr().add(base));
2687
2688 acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2689 acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2690 }
2691
2692 let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2694 let sums_d = _mm_add_ps(acc_dot, shuf_d);
2695 let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2696 let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2697
2698 let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2699 let sums_n = _mm_add_ps(acc_norm, shuf_n);
2700 let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2701 let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2702
2703 let base = chunks * 4;
2704 for i in 0..remainder {
2705 let v = u8_to_f32(*vec_u8.get_unchecked(base + i));
2706 dot += *query.get_unchecked(base + i) * v;
2707 norm += v * v;
2708 }
2709
2710 (dot, norm)
2711}
2712
2713#[cfg(target_arch = "x86_64")]
2718#[target_feature(enable = "avx", enable = "f16c", enable = "fma")]
2719#[allow(unsafe_op_in_unsafe_fn)]
2720unsafe fn fused_dot_norm_f16_f16c(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2721 use std::arch::x86_64::*;
2722
2723 let chunks16 = dim / 16;
2724 let remainder = dim % 16;
2725
2726 let mut acc_dot0 = _mm256_setzero_ps();
2728 let mut acc_dot1 = _mm256_setzero_ps();
2729 let mut acc_norm0 = _mm256_setzero_ps();
2730 let mut acc_norm1 = _mm256_setzero_ps();
2731
2732 for c in 0..chunks16 {
2733 let base = c * 16;
2734
2735 let v_raw0 = _mm_loadu_si128(vec_f16.as_ptr().add(base) as *const __m128i);
2737 let vb0 = _mm256_cvtph_ps(v_raw0);
2738 let q_raw0 = _mm_loadu_si128(query_f16.as_ptr().add(base) as *const __m128i);
2739 let qa0 = _mm256_cvtph_ps(q_raw0);
2740 acc_dot0 = _mm256_fmadd_ps(qa0, vb0, acc_dot0);
2741 acc_norm0 = _mm256_fmadd_ps(vb0, vb0, acc_norm0);
2742
2743 let v_raw1 = _mm_loadu_si128(vec_f16.as_ptr().add(base + 8) as *const __m128i);
2745 let vb1 = _mm256_cvtph_ps(v_raw1);
2746 let q_raw1 = _mm_loadu_si128(query_f16.as_ptr().add(base + 8) as *const __m128i);
2747 let qa1 = _mm256_cvtph_ps(q_raw1);
2748 acc_dot1 = _mm256_fmadd_ps(qa1, vb1, acc_dot1);
2749 acc_norm1 = _mm256_fmadd_ps(vb1, vb1, acc_norm1);
2750 }
2751
2752 let acc_dot = _mm256_add_ps(acc_dot0, acc_dot1);
2754 let acc_norm = _mm256_add_ps(acc_norm0, acc_norm1);
2755
2756 let hi_d = _mm256_extractf128_ps(acc_dot, 1);
2758 let lo_d = _mm256_castps256_ps128(acc_dot);
2759 let sum_d = _mm_add_ps(lo_d, hi_d);
2760 let shuf_d = _mm_shuffle_ps(sum_d, sum_d, 0b10_11_00_01);
2761 let sums_d = _mm_add_ps(sum_d, shuf_d);
2762 let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2763 let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2764
2765 let hi_n = _mm256_extractf128_ps(acc_norm, 1);
2766 let lo_n = _mm256_castps256_ps128(acc_norm);
2767 let sum_n = _mm_add_ps(lo_n, hi_n);
2768 let shuf_n = _mm_shuffle_ps(sum_n, sum_n, 0b10_11_00_01);
2769 let sums_n = _mm_add_ps(sum_n, shuf_n);
2770 let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2771 let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2772
2773 let base = chunks16 * 16;
2774 for i in 0..remainder {
2775 let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2776 let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2777 dot += q * v;
2778 norm += v * v;
2779 }
2780
2781 (dot, norm)
2782}
2783
2784#[cfg(target_arch = "x86_64")]
2785#[target_feature(enable = "avx", enable = "f16c", enable = "fma")]
2786#[allow(unsafe_op_in_unsafe_fn)]
2787unsafe fn dot_product_f16_f16c(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2788 use std::arch::x86_64::*;
2789
2790 let chunks = dim / 8;
2791 let remainder = dim % 8;
2792 let mut acc = _mm256_setzero_ps();
2793
2794 for chunk in 0..chunks {
2795 let base = chunk * 8;
2796 let v_raw = _mm_loadu_si128(vec_f16.as_ptr().add(base) as *const __m128i);
2797 let vb = _mm256_cvtph_ps(v_raw);
2798 let q_raw = _mm_loadu_si128(query_f16.as_ptr().add(base) as *const __m128i);
2799 let qa = _mm256_cvtph_ps(q_raw);
2800 acc = _mm256_fmadd_ps(qa, vb, acc);
2801 }
2802
2803 let hi = _mm256_extractf128_ps(acc, 1);
2804 let lo = _mm256_castps256_ps128(acc);
2805 let sum = _mm_add_ps(lo, hi);
2806 let shuf = _mm_shuffle_ps(sum, sum, 0b10_11_00_01);
2807 let sums = _mm_add_ps(sum, shuf);
2808 let shuf2 = _mm_movehl_ps(sums, sums);
2809 let mut dot = _mm_cvtss_f32(_mm_add_ss(sums, shuf2));
2810
2811 let base = chunks * 8;
2812 for i in 0..remainder {
2813 let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2814 let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2815 dot += q * v;
2816 }
2817 dot
2818}
2819
2820#[cfg(target_arch = "x86_64")]
2821#[target_feature(enable = "sse2", enable = "sse4.1")]
2822#[allow(unsafe_op_in_unsafe_fn)]
2823unsafe fn dot_product_u8_sse(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2824 use std::arch::x86_64::*;
2825
2826 let scale = _mm_set1_ps(U8_INV_SCALE);
2827 let offset = _mm_set1_ps(-1.0);
2828 let chunks = dim / 4;
2829 let remainder = dim % 4;
2830 let mut acc = _mm_setzero_ps();
2831
2832 for chunk in 0..chunks {
2833 let base = chunk * 4;
2834 let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
2835 vec_u8.as_ptr().add(base) as *const i32
2836 ));
2837 let ints = _mm_cvtepu8_epi32(bytes);
2838 let floats = _mm_cvtepi32_ps(ints);
2839 let vb = _mm_add_ps(_mm_mul_ps(floats, scale), offset);
2840 let va = _mm_loadu_ps(query.as_ptr().add(base));
2841 acc = _mm_add_ps(acc, _mm_mul_ps(va, vb));
2842 }
2843
2844 let shuf = _mm_shuffle_ps(acc, acc, 0b10_11_00_01);
2845 let sums = _mm_add_ps(acc, shuf);
2846 let shuf2 = _mm_movehl_ps(sums, sums);
2847 let mut dot = _mm_cvtss_f32(_mm_add_ss(sums, shuf2));
2848
2849 let base = chunks * 4;
2850 for i in 0..remainder {
2851 dot += *query.get_unchecked(base + i) * u8_to_f32(*vec_u8.get_unchecked(base + i));
2852 }
2853 dot
2854}
2855
2856#[inline]
2861fn fused_dot_norm_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2862 #[cfg(target_arch = "aarch64")]
2863 {
2864 return unsafe { neon_quant::fused_dot_norm_f16(query_f16, vec_f16, dim) };
2865 }
2866
2867 #[cfg(target_arch = "x86_64")]
2868 {
2869 if is_x86_feature_detected!("f16c") && is_x86_feature_detected!("fma") {
2870 return unsafe { fused_dot_norm_f16_f16c(query_f16, vec_f16, dim) };
2871 }
2872 if sse::is_available() {
2873 return unsafe { fused_dot_norm_f16_sse(query_f16, vec_f16, dim) };
2874 }
2875 }
2876
2877 #[allow(unreachable_code)]
2878 fused_dot_norm_f16_scalar(query_f16, vec_f16, dim)
2879}
2880
2881#[inline]
2882fn fused_dot_norm_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2883 #[cfg(target_arch = "aarch64")]
2884 {
2885 return unsafe { neon_quant::fused_dot_norm_u8(query, vec_u8, dim) };
2886 }
2887
2888 #[cfg(target_arch = "x86_64")]
2889 {
2890 if sse::is_available() {
2891 return unsafe { fused_dot_norm_u8_sse(query, vec_u8, dim) };
2892 }
2893 }
2894
2895 #[allow(unreachable_code)]
2896 fused_dot_norm_u8_scalar(query, vec_u8, dim)
2897}
2898
2899#[inline]
2902fn dot_product_f16_quant(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2903 #[cfg(target_arch = "aarch64")]
2904 {
2905 return unsafe { neon_quant::dot_product_f16(query_f16, vec_f16, dim) };
2906 }
2907
2908 #[cfg(target_arch = "x86_64")]
2909 {
2910 if is_x86_feature_detected!("f16c") && is_x86_feature_detected!("fma") {
2911 return unsafe { dot_product_f16_f16c(query_f16, vec_f16, dim) };
2912 }
2913 }
2914
2915 #[allow(unreachable_code)]
2916 dot_product_f16_scalar(query_f16, vec_f16, dim)
2917}
2918
2919#[inline]
2920fn dot_product_u8_quant(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2921 #[cfg(target_arch = "aarch64")]
2922 {
2923 return unsafe { neon_quant::dot_product_u8(query, vec_u8, dim) };
2924 }
2925
2926 #[cfg(target_arch = "x86_64")]
2927 {
2928 if sse::is_available() {
2929 return unsafe { dot_product_u8_sse(query, vec_u8, dim) };
2930 }
2931 }
2932
2933 #[allow(unreachable_code)]
2934 dot_product_u8_scalar(query, vec_u8, dim)
2935}
2936
2937#[inline]
2948pub fn batch_cosine_scores_f16(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
2949 let n = scores.len();
2950 let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
2951 let required = n
2952 .checked_mul(vec_bytes)
2953 .expect("f16 batch byte length overflow");
2954 assert_eq!(
2955 query.len(),
2956 dim,
2957 "f16 batch cosine query dimension mismatch"
2958 );
2959 assert!(
2960 vectors_raw.len() >= required,
2961 "f16 batch cosine vectors are truncated: need {required} bytes, got {}",
2962 vectors_raw.len()
2963 );
2964 if required > 0 {
2965 assert!(
2966 (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
2967 "f16 batch cosine vectors are not 2-byte aligned"
2968 );
2969 }
2970 if dim == 0 || n == 0 {
2971 return;
2972 }
2973
2974 let norm_q_sq = dot_product_f32(query, query, dim);
2976 if norm_q_sq < f32::EPSILON {
2977 for s in scores.iter_mut() {
2978 *s = 0.0;
2979 }
2980 return;
2981 }
2982 let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2983
2984 let query_f16: Vec<u16> = query.iter().map(|&v| f32_to_f16(v)).collect();
2986
2987 for i in 0..n {
2988 let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
2989 let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
2990
2991 let (dot, norm_v_sq) = fused_dot_norm_f16(&query_f16, f16_slice, dim);
2992 scores[i] = if norm_v_sq < f32::EPSILON {
2993 0.0
2994 } else {
2995 dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
2996 };
2997 }
2998}
2999
3000#[inline]
3006pub fn batch_cosine_scores_u8(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3007 let n = scores.len();
3008 let required = n.checked_mul(dim).expect("u8 batch byte length overflow");
3009 assert_eq!(query.len(), dim, "u8 batch cosine query dimension mismatch");
3010 assert!(
3011 vectors_raw.len() >= required,
3012 "u8 batch cosine vectors are truncated: need {required} bytes, got {}",
3013 vectors_raw.len()
3014 );
3015 if dim == 0 || n == 0 {
3016 return;
3017 }
3018
3019 let norm_q_sq = dot_product_f32(query, query, dim);
3020 if norm_q_sq < f32::EPSILON {
3021 for s in scores.iter_mut() {
3022 *s = 0.0;
3023 }
3024 return;
3025 }
3026 let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3027
3028 for i in 0..n {
3029 let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3030
3031 let (dot, norm_v_sq) = fused_dot_norm_u8(query, u8_slice, dim);
3032 scores[i] = if norm_v_sq < f32::EPSILON {
3033 0.0
3034 } else {
3035 dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3036 };
3037 }
3038}
3039
3040#[inline]
3049pub fn batch_dot_scores(query: &[f32], vectors: &[f32], dim: usize, scores: &mut [f32]) {
3050 let n = scores.len();
3051 let required = n
3052 .checked_mul(dim)
3053 .expect("batch dot vector length overflow");
3054 assert_eq!(query.len(), dim, "batch dot query dimension mismatch");
3055 assert!(
3056 vectors.len() >= required,
3057 "batch dot vectors are truncated: need {required}, got {}",
3058 vectors.len()
3059 );
3060
3061 if dim == 0 || n == 0 {
3062 return;
3063 }
3064
3065 let norm_q_sq = dot_product_f32(query, query, dim);
3066 if norm_q_sq < f32::EPSILON {
3067 for s in scores.iter_mut() {
3068 *s = 0.0;
3069 }
3070 return;
3071 }
3072 let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3073
3074 for i in 0..n {
3075 let vec = &vectors[i * dim..(i + 1) * dim];
3076 let dot = dot_product_f32(query, vec, dim);
3077 scores[i] = dot * inv_norm_q;
3078 }
3079}
3080
3081#[inline]
3086pub fn batch_dot_scores_f16(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3087 let n = scores.len();
3088 let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
3089 let required = n
3090 .checked_mul(vec_bytes)
3091 .expect("f16 batch byte length overflow");
3092 assert_eq!(query.len(), dim, "f16 batch dot query dimension mismatch");
3093 assert!(
3094 vectors_raw.len() >= required,
3095 "f16 batch dot vectors are truncated: need {required} bytes, got {}",
3096 vectors_raw.len()
3097 );
3098 if required > 0 {
3099 assert!(
3100 (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3101 "f16 batch dot vectors are not 2-byte aligned"
3102 );
3103 }
3104 if dim == 0 || n == 0 {
3105 return;
3106 }
3107
3108 let norm_q_sq = dot_product_f32(query, query, dim);
3109 if norm_q_sq < f32::EPSILON {
3110 for s in scores.iter_mut() {
3111 *s = 0.0;
3112 }
3113 return;
3114 }
3115 let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3116
3117 let query_f16: Vec<u16> = query.iter().map(|&v| f32_to_f16(v)).collect();
3118 for i in 0..n {
3119 let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3120 let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3121 let dot = dot_product_f16_quant(&query_f16, f16_slice, dim);
3122 scores[i] = dot * inv_norm_q;
3123 }
3124}
3125
3126#[inline]
3131pub fn batch_dot_scores_u8(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3132 let n = scores.len();
3133 let required = n.checked_mul(dim).expect("u8 batch byte length overflow");
3134 assert_eq!(query.len(), dim, "u8 batch dot query dimension mismatch");
3135 assert!(
3136 vectors_raw.len() >= required,
3137 "u8 batch dot vectors are truncated: need {required} bytes, got {}",
3138 vectors_raw.len()
3139 );
3140 if dim == 0 || n == 0 {
3141 return;
3142 }
3143
3144 let norm_q_sq = dot_product_f32(query, query, dim);
3145 if norm_q_sq < f32::EPSILON {
3146 for s in scores.iter_mut() {
3147 *s = 0.0;
3148 }
3149 return;
3150 }
3151 let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3152
3153 for i in 0..n {
3154 let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3155 let dot = dot_product_u8_quant(query, u8_slice, dim);
3156 scores[i] = dot * inv_norm_q;
3157 }
3158}
3159
3160#[inline]
3166pub fn batch_cosine_scores_precomp(
3167 query: &[f32],
3168 vectors: &[f32],
3169 dim: usize,
3170 scores: &mut [f32],
3171 inv_norm_q: f32,
3172) {
3173 let n = scores.len();
3174 let required = n
3175 .checked_mul(dim)
3176 .expect("precomputed cosine vector length overflow");
3177 assert_eq!(
3178 query.len(),
3179 dim,
3180 "precomputed cosine query dimension mismatch"
3181 );
3182 assert!(
3183 vectors.len() >= required,
3184 "precomputed cosine vectors are truncated: need {required}, got {}",
3185 vectors.len()
3186 );
3187 for i in 0..n {
3188 let vec = &vectors[i * dim..(i + 1) * dim];
3189 let (dot, norm_v_sq) = fused_dot_norm(query, vec, dim);
3190 scores[i] = if norm_v_sq < f32::EPSILON {
3191 0.0
3192 } else {
3193 dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3194 };
3195 }
3196}
3197
3198#[inline]
3200pub fn batch_cosine_scores_f16_precomp(
3201 query_f16: &[u16],
3202 vectors_raw: &[u8],
3203 dim: usize,
3204 scores: &mut [f32],
3205 inv_norm_q: f32,
3206) {
3207 let n = scores.len();
3208 let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
3209 let required = n
3210 .checked_mul(vec_bytes)
3211 .expect("precomputed f16 cosine batch byte length overflow");
3212 assert_eq!(
3213 query_f16.len(),
3214 dim,
3215 "precomputed f16 cosine query dimension mismatch"
3216 );
3217 assert!(
3218 vectors_raw.len() >= required,
3219 "precomputed f16 cosine vectors are truncated: need {required} bytes, got {}",
3220 vectors_raw.len()
3221 );
3222 if required > 0 {
3223 assert!(
3224 (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3225 "precomputed f16 cosine vectors are not 2-byte aligned"
3226 );
3227 }
3228 for i in 0..n {
3229 let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3230 let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3231 let (dot, norm_v_sq) = fused_dot_norm_f16(query_f16, f16_slice, dim);
3232 scores[i] = if norm_v_sq < f32::EPSILON {
3233 0.0
3234 } else {
3235 dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3236 };
3237 }
3238}
3239
3240#[inline]
3242pub fn batch_cosine_scores_u8_precomp(
3243 query: &[f32],
3244 vectors_raw: &[u8],
3245 dim: usize,
3246 scores: &mut [f32],
3247 inv_norm_q: f32,
3248) {
3249 let n = scores.len();
3250 let required = n
3251 .checked_mul(dim)
3252 .expect("precomputed u8 cosine batch byte length overflow");
3253 assert_eq!(
3254 query.len(),
3255 dim,
3256 "precomputed u8 cosine query dimension mismatch"
3257 );
3258 assert!(
3259 vectors_raw.len() >= required,
3260 "precomputed u8 cosine vectors are truncated: need {required} bytes, got {}",
3261 vectors_raw.len()
3262 );
3263 for i in 0..n {
3264 let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3265 let (dot, norm_v_sq) = fused_dot_norm_u8(query, u8_slice, dim);
3266 scores[i] = if norm_v_sq < f32::EPSILON {
3267 0.0
3268 } else {
3269 dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3270 };
3271 }
3272}
3273
3274#[inline]
3276pub fn batch_dot_scores_precomp(
3277 query: &[f32],
3278 vectors: &[f32],
3279 dim: usize,
3280 scores: &mut [f32],
3281 inv_norm_q: f32,
3282) {
3283 let n = scores.len();
3284 let required = n
3285 .checked_mul(dim)
3286 .expect("precomputed dot vector length overflow");
3287 assert_eq!(query.len(), dim, "precomputed dot query dimension mismatch");
3288 assert!(
3289 vectors.len() >= required,
3290 "precomputed dot vectors are truncated: need {required}, got {}",
3291 vectors.len()
3292 );
3293 for i in 0..n {
3294 let vec = &vectors[i * dim..(i + 1) * dim];
3295 scores[i] = dot_product_f32(query, vec, dim) * inv_norm_q;
3296 }
3297}
3298
3299#[inline]
3301pub fn batch_dot_scores_f16_precomp(
3302 query_f16: &[u16],
3303 vectors_raw: &[u8],
3304 dim: usize,
3305 scores: &mut [f32],
3306 inv_norm_q: f32,
3307) {
3308 let n = scores.len();
3309 let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
3310 let required = n
3311 .checked_mul(vec_bytes)
3312 .expect("precomputed f16 dot batch byte length overflow");
3313 assert_eq!(
3314 query_f16.len(),
3315 dim,
3316 "precomputed f16 dot query dimension mismatch"
3317 );
3318 assert!(
3319 vectors_raw.len() >= required,
3320 "precomputed f16 dot vectors are truncated: need {required} bytes, got {}",
3321 vectors_raw.len()
3322 );
3323 if required > 0 {
3324 assert!(
3325 (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3326 "precomputed f16 dot vectors are not 2-byte aligned"
3327 );
3328 }
3329 for i in 0..n {
3330 let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3331 let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3332 scores[i] = dot_product_f16_quant(query_f16, f16_slice, dim) * inv_norm_q;
3333 }
3334}
3335
3336#[inline]
3338pub fn batch_dot_scores_u8_precomp(
3339 query: &[f32],
3340 vectors_raw: &[u8],
3341 dim: usize,
3342 scores: &mut [f32],
3343 inv_norm_q: f32,
3344) {
3345 let n = scores.len();
3346 let required = n
3347 .checked_mul(dim)
3348 .expect("precomputed u8 dot batch byte length overflow");
3349 assert_eq!(
3350 query.len(),
3351 dim,
3352 "precomputed u8 dot query dimension mismatch"
3353 );
3354 assert!(
3355 vectors_raw.len() >= required,
3356 "precomputed u8 dot vectors are truncated: need {required} bytes, got {}",
3357 vectors_raw.len()
3358 );
3359 for i in 0..n {
3360 let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3361 scores[i] = dot_product_u8_quant(query, u8_slice, dim) * inv_norm_q;
3362 }
3363}
3364
3365#[inline]
3370pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
3371 assert_eq!(a.len(), b.len(), "cosine vector dimension mismatch");
3372 let count = a.len();
3373
3374 if count == 0 {
3375 return 0.0;
3376 }
3377
3378 let dot = dot_product_f32(a, b, count);
3379 let norm_a = dot_product_f32(a, a, count);
3380 let norm_b = dot_product_f32(b, b, count);
3381
3382 let denom = (norm_a * norm_b).sqrt();
3383 if denom < f32::EPSILON {
3384 return 0.0;
3385 }
3386
3387 dot / denom
3388}
3389
3390#[inline]
3399pub fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
3400 assert_eq!(a.len(), b.len(), "Hamming vector byte length mismatch");
3401
3402 #[cfg(target_arch = "aarch64")]
3403 unsafe {
3404 neon::hamming_distance(a, b)
3405 }
3406
3407 #[cfg(target_arch = "x86_64")]
3408 {
3409 if avx2::is_available() {
3410 return unsafe { avx2::hamming_distance(a, b) };
3411 }
3412 hamming_distance_scalar(a, b)
3413 }
3414
3415 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3416 hamming_distance_scalar(a, b)
3417}
3418
3419#[inline]
3422#[allow(dead_code)]
3423fn hamming_distance_scalar(a: &[u8], b: &[u8]) -> u32 {
3424 let len = a.len();
3425 let chunks = len / 8;
3426 let remainder = len % 8;
3427 let mut total = 0u32;
3428
3429 for i in 0..chunks {
3430 let off = i * 8;
3431 let va = unsafe { std::ptr::read_unaligned(a.as_ptr().add(off) as *const u64) };
3432 let vb = unsafe { std::ptr::read_unaligned(b.as_ptr().add(off) as *const u64) };
3433 total += (va ^ vb).count_ones();
3434 }
3435
3436 let base = chunks * 8;
3437 for i in 0..remainder {
3438 total += (a[base + i] ^ b[base + i]).count_ones();
3439 }
3440
3441 total
3442}
3443
3444pub fn batch_hamming_scores(
3450 query: &[u8],
3451 db: &[u8],
3452 byte_len: usize,
3453 dim_bits: usize,
3454 scores: &mut [f32],
3455) {
3456 let n = scores.len();
3457 let required = n
3458 .checked_mul(byte_len)
3459 .expect("Hamming batch byte length overflow");
3460 assert_eq!(query.len(), byte_len, "Hamming query byte length mismatch");
3461 assert!(
3462 db.len() >= required,
3463 "Hamming batch is truncated: need {required} bytes, got {}",
3464 db.len()
3465 );
3466
3467 if byte_len == 0 || n == 0 || dim_bits == 0 {
3468 return;
3469 }
3470
3471 let inv_dim = 1.0 / dim_bits as f32;
3472
3473 for i in 0..n {
3474 let vec = &db[i * byte_len..(i + 1) * byte_len];
3475 let dist = hamming_distance(query, vec);
3476 scores[i] = 1.0 - dist as f32 * inv_dim;
3477 }
3478}
3479
3480#[cfg(test)]
3481mod tests {
3482 use super::*;
3483
3484 #[test]
3485 fn vector_simd_boundaries_reject_dimension_mismatches() {
3486 let vectors = vec![1.0f32; 6];
3487 let raw_f16 = vec![0u8; 12];
3488 let raw_u8 = vec![0u8; 6];
3489 let mut scores = vec![0.0f32; 2];
3490
3491 for invalid_query in [vec![1.0, 2.0], vec![1.0, 2.0, 3.0, 4.0]] {
3492 assert!(
3493 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3494 batch_cosine_scores(&invalid_query, &vectors, 3, &mut scores)
3495 }))
3496 .is_err()
3497 );
3498 assert!(
3499 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3500 batch_dot_scores_f16(&invalid_query, &raw_f16, 3, &mut scores)
3501 }))
3502 .is_err()
3503 );
3504 assert!(
3505 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3506 batch_cosine_scores_u8(&invalid_query, &raw_u8, 3, &mut scores)
3507 }))
3508 .is_err()
3509 );
3510 }
3511 }
3512
3513 #[test]
3514 fn vector_simd_boundaries_reject_truncated_storage() {
3515 let query = [1.0f32, 2.0, 3.0];
3516 let mut scores = [0.0f32; 2];
3517
3518 assert!(
3519 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3520 batch_dot_scores(&query, &[0.0; 5], 3, &mut scores)
3521 }))
3522 .is_err()
3523 );
3524 assert!(
3525 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3526 batch_cosine_scores_f16(&query, &[0u8; 11], 3, &mut scores)
3527 }))
3528 .is_err()
3529 );
3530 assert!(
3531 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3532 dot_product_f32(&query, &query, 4)
3533 }))
3534 .is_err()
3535 );
3536 }
3537
3538 #[test]
3539 fn test_unpack_8bit() {
3540 let input: Vec<u8> = (0..128).collect();
3541 let mut output = vec![0u32; 128];
3542 unpack_8bit(&input, &mut output, 128);
3543
3544 for (i, &v) in output.iter().enumerate() {
3545 assert_eq!(v, i as u32);
3546 }
3547 }
3548
3549 #[test]
3550 fn test_unpack_16bit() {
3551 let mut input = vec![0u8; 256];
3552 for i in 0..128 {
3553 let val = (i * 100) as u16;
3554 input[i * 2] = val as u8;
3555 input[i * 2 + 1] = (val >> 8) as u8;
3556 }
3557
3558 let mut output = vec![0u32; 128];
3559 unpack_16bit(&input, &mut output, 128);
3560
3561 for (i, &v) in output.iter().enumerate() {
3562 assert_eq!(v, (i * 100) as u32);
3563 }
3564 }
3565
3566 #[test]
3567 fn test_unpack_32bit() {
3568 let mut input = vec![0u8; 512];
3569 for i in 0..128 {
3570 let val = (i * 1000) as u32;
3571 let bytes = val.to_le_bytes();
3572 input[i * 4..i * 4 + 4].copy_from_slice(&bytes);
3573 }
3574
3575 let mut output = vec![0u32; 128];
3576 unpack_32bit(&input, &mut output, 128);
3577
3578 for (i, &v) in output.iter().enumerate() {
3579 assert_eq!(v, (i * 1000) as u32);
3580 }
3581 }
3582
3583 #[test]
3584 fn test_delta_decode() {
3585 let deltas = vec![4u32, 4, 9, 19];
3589 let mut output = vec![0u32; 5];
3590
3591 delta_decode(&mut output, &deltas, 10, 5);
3592
3593 assert_eq!(output, vec![10, 15, 20, 30, 50]);
3594 }
3595
3596 #[test]
3597 fn test_add_one() {
3598 let mut values = vec![0u32, 1, 2, 3, 4, 5, 6, 7];
3599 add_one(&mut values, 8);
3600
3601 assert_eq!(values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
3602 }
3603
3604 #[test]
3605 fn test_bits_needed() {
3606 assert_eq!(bits_needed(0), 0);
3607 assert_eq!(bits_needed(1), 1);
3608 assert_eq!(bits_needed(2), 2);
3609 assert_eq!(bits_needed(3), 2);
3610 assert_eq!(bits_needed(4), 3);
3611 assert_eq!(bits_needed(255), 8);
3612 assert_eq!(bits_needed(256), 9);
3613 assert_eq!(bits_needed(u32::MAX), 32);
3614 }
3615
3616 #[test]
3617 fn test_unpack_8bit_delta_decode() {
3618 let input: Vec<u8> = vec![4, 4, 9, 19];
3622 let mut output = vec![0u32; 5];
3623
3624 unpack_8bit_delta_decode(&input, &mut output, 10, 5);
3625
3626 assert_eq!(output, vec![10, 15, 20, 30, 50]);
3627 }
3628
3629 #[test]
3630 fn test_unpack_16bit_delta_decode() {
3631 let mut input = vec![0u8; 8];
3635 for (i, &delta) in [499u16, 499, 999, 1999].iter().enumerate() {
3636 input[i * 2] = delta as u8;
3637 input[i * 2 + 1] = (delta >> 8) as u8;
3638 }
3639 let mut output = vec![0u32; 5];
3640
3641 unpack_16bit_delta_decode(&input, &mut output, 100, 5);
3642
3643 assert_eq!(output, vec![100, 600, 1100, 2100, 4100]);
3644 }
3645
3646 #[test]
3647 fn test_fused_vs_separate_8bit() {
3648 let input: Vec<u8> = (0..127).collect();
3650 let first_value = 1000u32;
3651 let count = 128;
3652
3653 let mut unpacked = vec![0u32; 128];
3655 unpack_8bit(&input, &mut unpacked, 127);
3656 let mut separate_output = vec![0u32; 128];
3657 delta_decode(&mut separate_output, &unpacked, first_value, count);
3658
3659 let mut fused_output = vec![0u32; 128];
3661 unpack_8bit_delta_decode(&input, &mut fused_output, first_value, count);
3662
3663 assert_eq!(separate_output, fused_output);
3664 }
3665
3666 #[test]
3667 fn test_round_bit_width() {
3668 assert_eq!(round_bit_width(0), 0);
3669 assert_eq!(round_bit_width(1), 8);
3670 assert_eq!(round_bit_width(5), 8);
3671 assert_eq!(round_bit_width(8), 8);
3672 assert_eq!(round_bit_width(9), 16);
3673 assert_eq!(round_bit_width(12), 16);
3674 assert_eq!(round_bit_width(16), 16);
3675 assert_eq!(round_bit_width(17), 32);
3676 assert_eq!(round_bit_width(24), 32);
3677 assert_eq!(round_bit_width(32), 32);
3678 }
3679
3680 #[test]
3681 fn test_rounded_bitwidth_from_exact() {
3682 assert_eq!(RoundedBitWidth::from_exact(0), RoundedBitWidth::Zero);
3683 assert_eq!(RoundedBitWidth::from_exact(1), RoundedBitWidth::Bits8);
3684 assert_eq!(RoundedBitWidth::from_exact(8), RoundedBitWidth::Bits8);
3685 assert_eq!(RoundedBitWidth::from_exact(9), RoundedBitWidth::Bits16);
3686 assert_eq!(RoundedBitWidth::from_exact(16), RoundedBitWidth::Bits16);
3687 assert_eq!(RoundedBitWidth::from_exact(17), RoundedBitWidth::Bits32);
3688 assert_eq!(RoundedBitWidth::from_exact(32), RoundedBitWidth::Bits32);
3689 }
3690
3691 #[test]
3692 fn test_pack_unpack_rounded_8bit() {
3693 let values: Vec<u32> = (0..128).map(|i| i % 256).collect();
3694 let mut packed = vec![0u8; 128];
3695
3696 let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits8, &mut packed);
3697 assert_eq!(bytes_written, 128);
3698
3699 let mut unpacked = vec![0u32; 128];
3700 unpack_rounded(&packed, RoundedBitWidth::Bits8, &mut unpacked, 128);
3701
3702 assert_eq!(values, unpacked);
3703 }
3704
3705 #[test]
3706 fn test_pack_unpack_rounded_16bit() {
3707 let values: Vec<u32> = (0..128).map(|i| i * 100).collect();
3708 let mut packed = vec![0u8; 256];
3709
3710 let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits16, &mut packed);
3711 assert_eq!(bytes_written, 256);
3712
3713 let mut unpacked = vec![0u32; 128];
3714 unpack_rounded(&packed, RoundedBitWidth::Bits16, &mut unpacked, 128);
3715
3716 assert_eq!(values, unpacked);
3717 }
3718
3719 #[test]
3720 fn test_pack_unpack_rounded_32bit() {
3721 let values: Vec<u32> = (0..128).map(|i| i * 100000).collect();
3722 let mut packed = vec![0u8; 512];
3723
3724 let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits32, &mut packed);
3725 assert_eq!(bytes_written, 512);
3726
3727 let mut unpacked = vec![0u32; 128];
3728 unpack_rounded(&packed, RoundedBitWidth::Bits32, &mut unpacked, 128);
3729
3730 assert_eq!(values, unpacked);
3731 }
3732
3733 #[test]
3734 fn test_unpack_rounded_delta_decode() {
3735 let input: Vec<u8> = vec![4, 4, 9, 19];
3740 let mut output = vec![0u32; 5];
3741
3742 unpack_rounded_delta_decode(&input, RoundedBitWidth::Bits8, &mut output, 10, 5);
3743
3744 assert_eq!(output, vec![10, 15, 20, 30, 50]);
3745 }
3746
3747 #[test]
3748 fn test_unpack_rounded_delta_decode_zero() {
3749 let input: Vec<u8> = vec![];
3751 let mut output = vec![0u32; 5];
3752
3753 unpack_rounded_delta_decode(&input, RoundedBitWidth::Zero, &mut output, 100, 5);
3754
3755 assert_eq!(output, vec![100, 101, 102, 103, 104]);
3756 }
3757
3758 #[test]
3763 fn test_dequantize_uint8() {
3764 let input: Vec<u8> = vec![0, 128, 255, 64, 192];
3765 let mut output = vec![0.0f32; 5];
3766 let scale = 0.1;
3767 let min_val = 1.0;
3768
3769 dequantize_uint8(&input, &mut output, scale, min_val, 5);
3770
3771 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); }
3778
3779 #[test]
3780 fn test_dequantize_uint8_large() {
3781 let input: Vec<u8> = (0..128).collect();
3783 let mut output = vec![0.0f32; 128];
3784 let scale = 2.0;
3785 let min_val = -10.0;
3786
3787 dequantize_uint8(&input, &mut output, scale, min_val, 128);
3788
3789 for (i, &out) in output.iter().enumerate().take(128) {
3790 let expected = i as f32 * scale + min_val;
3791 assert!(
3792 (out - expected).abs() < 1e-5,
3793 "Mismatch at {}: expected {}, got {}",
3794 i,
3795 expected,
3796 out
3797 );
3798 }
3799 }
3800
3801 #[test]
3802 fn test_dot_product_f32() {
3803 let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
3804 let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0];
3805
3806 let result = dot_product_f32(&a, &b, 5);
3807
3808 assert!((result - 70.0).abs() < 1e-5);
3810 }
3811
3812 #[test]
3813 fn test_dot_product_f32_large() {
3814 let a: Vec<f32> = (0..128).map(|i| i as f32).collect();
3816 let b: Vec<f32> = (0..128).map(|i| (i + 1) as f32).collect();
3817
3818 let result = dot_product_f32(&a, &b, 128);
3819
3820 let expected: f32 = (0..128).map(|i| (i as f32) * ((i + 1) as f32)).sum();
3822 assert!(
3823 (result - expected).abs() < 1e-3,
3824 "Expected {}, got {}",
3825 expected,
3826 result
3827 );
3828 }
3829
3830 #[test]
3831 fn test_fused_dot_norm() {
3832 let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
3833 let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
3834 let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3835
3836 let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3837 let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3838 assert!(
3839 (dot - expected_dot).abs() < 1e-5,
3840 "dot: expected {}, got {}",
3841 expected_dot,
3842 dot
3843 );
3844 assert!(
3845 (norm_b - expected_norm).abs() < 1e-5,
3846 "norm: expected {}, got {}",
3847 expected_norm,
3848 norm_b
3849 );
3850 }
3851
3852 #[test]
3853 fn test_fused_dot_norm_large() {
3854 let a: Vec<f32> = (0..768).map(|i| (i as f32) * 0.01).collect();
3855 let b: Vec<f32> = (0..768).map(|i| (i as f32) * 0.02 + 0.5).collect();
3856 let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3857
3858 let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3859 let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3860 assert!(
3861 (dot - expected_dot).abs() < 1.0,
3862 "dot: expected {}, got {}",
3863 expected_dot,
3864 dot
3865 );
3866 assert!(
3867 (norm_b - expected_norm).abs() < 1.0,
3868 "norm: expected {}, got {}",
3869 expected_norm,
3870 norm_b
3871 );
3872 }
3873
3874 #[test]
3875 fn test_batch_cosine_scores() {
3876 let query = vec![1.0f32, 0.0, 0.0];
3878 let vectors = vec![
3879 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, -1.0, 0.0, 0.0, 0.5, 0.5, 0.0, ];
3884 let mut scores = vec![0f32; 4];
3885 batch_cosine_scores(&query, &vectors, 3, &mut scores);
3886
3887 assert!((scores[0] - 1.0).abs() < 1e-5, "identical: {}", scores[0]);
3888 assert!(scores[1].abs() < 1e-5, "orthogonal: {}", scores[1]);
3889 assert!((scores[2] - (-1.0)).abs() < 1e-5, "opposite: {}", scores[2]);
3890 let expected_45 = 0.5f32 / (0.5f32.powi(2) + 0.5f32.powi(2)).sqrt();
3891 assert!(
3892 (scores[3] - expected_45).abs() < 1e-5,
3893 "45deg: expected {}, got {}",
3894 expected_45,
3895 scores[3]
3896 );
3897 }
3898
3899 #[test]
3900 fn test_batch_cosine_scores_matches_individual() {
3901 let query: Vec<f32> = (0..128).map(|i| (i as f32) * 0.1).collect();
3902 let n = 50;
3903 let dim = 128;
3904 let vectors: Vec<f32> = (0..n * dim).map(|i| ((i * 7 + 3) as f32) * 0.01).collect();
3905
3906 let mut batch_scores = vec![0f32; n];
3907 batch_cosine_scores(&query, &vectors, dim, &mut batch_scores);
3908
3909 for i in 0..n {
3910 let vec_i = &vectors[i * dim..(i + 1) * dim];
3911 let individual = cosine_similarity(&query, vec_i);
3912 assert!(
3913 (batch_scores[i] - individual).abs() < 1e-5,
3914 "vec {}: batch={}, individual={}",
3915 i,
3916 batch_scores[i],
3917 individual
3918 );
3919 }
3920 }
3921
3922 #[test]
3923 fn test_batch_cosine_scores_empty() {
3924 let query = vec![1.0f32, 2.0, 3.0];
3925 let vectors: Vec<f32> = vec![];
3926 let mut scores: Vec<f32> = vec![];
3927 batch_cosine_scores(&query, &vectors, 3, &mut scores);
3928 assert!(scores.is_empty());
3929 }
3930
3931 #[test]
3932 fn test_batch_cosine_scores_zero_query() {
3933 let query = vec![0.0f32, 0.0, 0.0];
3934 let vectors = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
3935 let mut scores = vec![0f32; 2];
3936 batch_cosine_scores(&query, &vectors, 3, &mut scores);
3937 assert_eq!(scores[0], 0.0);
3938 assert_eq!(scores[1], 0.0);
3939 }
3940
3941 #[test]
3946 fn test_f16_roundtrip_normal() {
3947 for &v in &[0.0f32, 1.0, -1.0, 0.5, -0.5, 0.333, 65504.0] {
3948 let h = f32_to_f16(v);
3949 let back = f16_to_f32(h);
3950 let err = (back - v).abs() / v.abs().max(1e-6);
3951 assert!(
3952 err < 0.002,
3953 "f16 roundtrip {v} → {h:#06x} → {back}, rel err {err}"
3954 );
3955 }
3956 }
3957
3958 #[test]
3959 fn test_f16_special() {
3960 assert_eq!(f16_to_f32(f32_to_f16(0.0)), 0.0);
3962 assert_eq!(f32_to_f16(-0.0), 0x8000);
3964 assert!(f16_to_f32(f32_to_f16(f32::INFINITY)).is_infinite());
3966 assert!(f16_to_f32(f32_to_f16(f32::NAN)).is_nan());
3968 }
3969
3970 #[test]
3971 fn test_f16_embedding_range() {
3972 let values: Vec<f32> = (-100..=100).map(|i| i as f32 / 100.0).collect();
3974 for &v in &values {
3975 let back = f16_to_f32(f32_to_f16(v));
3976 assert!((back - v).abs() < 0.001, "f16 error for {v}: got {back}");
3977 }
3978 }
3979
3980 #[test]
3985 fn test_u8_roundtrip() {
3986 assert_eq!(f32_to_u8_saturating(-1.0), 0);
3988 assert_eq!(f32_to_u8_saturating(1.0), 255);
3989 assert_eq!(f32_to_u8_saturating(0.0), 127); assert_eq!(f32_to_u8_saturating(-2.0), 0);
3993 assert_eq!(f32_to_u8_saturating(2.0), 255);
3994 }
3995
3996 #[test]
3997 fn test_u8_dequantize() {
3998 assert!((u8_to_f32(0) - (-1.0)).abs() < 0.01);
3999 assert!((u8_to_f32(255) - 1.0).abs() < 0.01);
4000 assert!((u8_to_f32(127) - 0.0).abs() < 0.01);
4001 }
4002
4003 #[test]
4008 fn test_batch_cosine_scores_f16() {
4009 let query = vec![0.6f32, 0.8, 0.0, 0.0];
4010 let dim = 4;
4011 let vecs_f32 = vec![
4012 0.6f32, 0.8, 0.0, 0.0, 0.0, 0.0, 0.6, 0.8, ];
4015
4016 let mut f16_buf = vec![0u16; 8];
4018 batch_f32_to_f16(&vecs_f32, &mut f16_buf);
4019 let raw: &[u8] =
4020 unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
4021
4022 let mut scores = vec![0f32; 2];
4023 batch_cosine_scores_f16(&query, raw, dim, &mut scores);
4024
4025 assert!(
4026 (scores[0] - 1.0).abs() < 0.01,
4027 "identical vectors: {}",
4028 scores[0]
4029 );
4030 assert!(scores[1].abs() < 0.01, "orthogonal vectors: {}", scores[1]);
4031 }
4032
4033 #[test]
4034 fn test_batch_cosine_scores_u8() {
4035 let query = vec![0.6f32, 0.8, 0.0, 0.0];
4036 let dim = 4;
4037 let vecs_f32 = vec![
4038 0.6f32, 0.8, 0.0, 0.0, -0.6, -0.8, 0.0, 0.0, ];
4041
4042 let mut u8_buf = vec![0u8; 8];
4044 batch_f32_to_u8(&vecs_f32, &mut u8_buf);
4045
4046 let mut scores = vec![0f32; 2];
4047 batch_cosine_scores_u8(&query, &u8_buf, dim, &mut scores);
4048
4049 assert!(scores[0] > 0.95, "similar vectors: {}", scores[0]);
4050 assert!(scores[1] < -0.95, "opposite vectors: {}", scores[1]);
4051 }
4052
4053 #[test]
4054 fn test_batch_cosine_scores_f16_large_dim() {
4055 let dim = 768;
4057 let query: Vec<f32> = (0..dim).map(|i| (i as f32 / dim as f32) - 0.5).collect();
4058 let vec2: Vec<f32> = query.iter().map(|x| x * 0.9 + 0.01).collect();
4059
4060 let mut all_vecs = query.clone();
4061 all_vecs.extend_from_slice(&vec2);
4062
4063 let mut f16_buf = vec![0u16; all_vecs.len()];
4064 batch_f32_to_f16(&all_vecs, &mut f16_buf);
4065 let raw: &[u8] =
4066 unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
4067
4068 let mut scores = vec![0f32; 2];
4069 batch_cosine_scores_f16(&query, raw, dim, &mut scores);
4070
4071 assert!((scores[0] - 1.0).abs() < 0.01, "self-sim: {}", scores[0]);
4073 assert!(scores[1] > 0.99, "scaled-sim: {}", scores[1]);
4075 }
4076
4077 #[test]
4082 fn test_hamming_distance_identical() {
4083 let a = vec![0xAA; 64];
4084 assert_eq!(hamming_distance(&a, &a), 0);
4085 }
4086
4087 #[test]
4088 fn test_hamming_distance_opposite() {
4089 let a = vec![0xFF; 32];
4090 let b = vec![0x00; 32];
4091 assert_eq!(hamming_distance(&a, &b), 256);
4092 }
4093
4094 #[test]
4095 fn test_hamming_distance_known() {
4096 let a = vec![0xAA];
4098 let b = vec![0x55];
4099 assert_eq!(hamming_distance(&a, &b), 8);
4100
4101 let a = vec![0xFF, 0x00];
4103 let b = vec![0x00, 0x00];
4104 assert_eq!(hamming_distance(&a, &b), 8);
4105 }
4106
4107 #[test]
4108 fn test_hamming_distance_single_bit() {
4109 let a = vec![0x00; 16];
4110 let mut b = vec![0x00; 16];
4111 b[7] = 0x01; assert_eq!(hamming_distance(&a, &b), 1);
4113 }
4114
4115 #[test]
4116 fn test_hamming_distance_empty() {
4117 let a: Vec<u8> = vec![];
4118 assert_eq!(hamming_distance(&a, &a), 0);
4119 }
4120
4121 #[test]
4122 fn test_hamming_distance_remainder_path() {
4123 let a = vec![0xFF; 17];
4125 let b = vec![0x00; 17];
4126 assert_eq!(hamming_distance(&a, &b), 136); let a = vec![0xFF; 33];
4130 let b = vec![0x00; 33];
4131 assert_eq!(hamming_distance(&a, &b), 264); }
4133
4134 #[test]
4135 fn test_hamming_distance_large() {
4136 let a = vec![0xFF; 4096];
4138 let b = vec![0x00; 4096];
4139 assert_eq!(hamming_distance(&a, &b), 32768);
4140 }
4141
4142 #[test]
4143 fn test_hamming_distance_scalar_matches() {
4144 for size in [1, 7, 8, 15, 16, 31, 32, 63, 64, 100, 128, 255, 256] {
4146 let a: Vec<u8> = (0..size).map(|i| (i * 37 + 13) as u8).collect();
4147 let b: Vec<u8> = (0..size).map(|i| (i * 53 + 7) as u8).collect();
4148 let expected = hamming_distance_scalar(&a, &b);
4149 let got = hamming_distance(&a, &b);
4150 assert_eq!(got, expected, "mismatch at size {size}");
4151 }
4152 }
4153
4154 #[test]
4159 fn test_batch_hamming_scores_identical() {
4160 let query = vec![0xAA; 16];
4161 let db = vec![0xAA; 16]; let mut scores = vec![0f32; 1];
4163 batch_hamming_scores(&query, &db, 16, 128, &mut scores);
4164 assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4165 }
4166
4167 #[test]
4168 fn test_batch_hamming_scores_opposite() {
4169 let query = vec![0xFF; 16];
4170 let db = vec![0x00; 16];
4171 let mut scores = vec![0f32; 1];
4172 batch_hamming_scores(&query, &db, 16, 128, &mut scores);
4173 assert!((scores[0] - 0.0).abs() < 1e-6, "opposite: {}", scores[0]);
4174 }
4175
4176 #[test]
4177 fn test_batch_hamming_scores_multiple() {
4178 let byte_len = 8;
4179 let dim_bits = 64;
4180 let query = vec![0xFF; byte_len];
4181 let mut db = Vec::new();
4182 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];
4187 batch_hamming_scores(&query, &db, byte_len, dim_bits, &mut scores);
4188
4189 assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4190 assert!((scores[1] - 0.0).abs() < 1e-6, "opposite: {}", scores[1]);
4191 assert!((scores[2] - 0.5).abs() < 1e-6, "half: {}", scores[2]);
4192 }
4193
4194 #[test]
4195 fn test_batch_hamming_scores_empty() {
4196 let query = vec![0xFF; 8];
4197 let db: Vec<u8> = vec![];
4198 let mut scores: Vec<f32> = vec![];
4199 batch_hamming_scores(&query, &db, 8, 64, &mut scores);
4200 assert!(scores.is_empty());
4201 }
4202
4203 #[test]
4204 fn test_batch_hamming_scores_zero_byte_len() {
4205 let query: Vec<u8> = vec![];
4206 let db: Vec<u8> = vec![];
4207 let mut scores = vec![0f32; 1];
4208 batch_hamming_scores(&query, &db, 0, 0, &mut scores);
4209 assert_eq!(scores[0], 0.0);
4211 }
4212}
4213
4214#[inline]
4227pub fn find_first_ge_u32(slice: &[u32], target: u32) -> usize {
4228 #[cfg(target_arch = "aarch64")]
4229 {
4230 if neon::is_available() {
4231 return unsafe { find_first_ge_u32_neon(slice, target) };
4232 }
4233 }
4234
4235 #[cfg(target_arch = "x86_64")]
4236 {
4237 if sse::is_available() {
4238 return unsafe { find_first_ge_u32_sse(slice, target) };
4239 }
4240 }
4241
4242 slice.partition_point(|&d| d < target)
4244}
4245
4246#[cfg(target_arch = "aarch64")]
4247#[target_feature(enable = "neon")]
4248#[allow(unsafe_op_in_unsafe_fn)]
4249unsafe fn find_first_ge_u32_neon(slice: &[u32], target: u32) -> usize {
4250 use std::arch::aarch64::*;
4251
4252 let n = slice.len();
4253 let ptr = slice.as_ptr();
4254 let target_vec = vdupq_n_u32(target);
4255 let bit_mask: uint32x4_t = core::mem::transmute([1u32, 2u32, 4u32, 8u32]);
4257
4258 let chunks = n / 16;
4259 let mut base = 0usize;
4260
4261 for _ in 0..chunks {
4263 let v0 = vld1q_u32(ptr.add(base));
4264 let v1 = vld1q_u32(ptr.add(base + 4));
4265 let v2 = vld1q_u32(ptr.add(base + 8));
4266 let v3 = vld1q_u32(ptr.add(base + 12));
4267
4268 let c0 = vcgeq_u32(v0, target_vec);
4269 let c1 = vcgeq_u32(v1, target_vec);
4270 let c2 = vcgeq_u32(v2, target_vec);
4271 let c3 = vcgeq_u32(v3, target_vec);
4272
4273 let m0 = vaddvq_u32(vandq_u32(c0, bit_mask));
4274 if m0 != 0 {
4275 return base + m0.trailing_zeros() as usize;
4276 }
4277 let m1 = vaddvq_u32(vandq_u32(c1, bit_mask));
4278 if m1 != 0 {
4279 return base + 4 + m1.trailing_zeros() as usize;
4280 }
4281 let m2 = vaddvq_u32(vandq_u32(c2, bit_mask));
4282 if m2 != 0 {
4283 return base + 8 + m2.trailing_zeros() as usize;
4284 }
4285 let m3 = vaddvq_u32(vandq_u32(c3, bit_mask));
4286 if m3 != 0 {
4287 return base + 12 + m3.trailing_zeros() as usize;
4288 }
4289 base += 16;
4290 }
4291
4292 while base + 4 <= n {
4294 let vals = vld1q_u32(ptr.add(base));
4295 let cmp = vcgeq_u32(vals, target_vec);
4296 let mask = vaddvq_u32(vandq_u32(cmp, bit_mask));
4297 if mask != 0 {
4298 return base + mask.trailing_zeros() as usize;
4299 }
4300 base += 4;
4301 }
4302
4303 while base < n {
4305 if *slice.get_unchecked(base) >= target {
4306 return base;
4307 }
4308 base += 1;
4309 }
4310 n
4311}
4312
4313#[cfg(target_arch = "x86_64")]
4314#[target_feature(enable = "sse2")]
4315#[allow(unsafe_op_in_unsafe_fn)]
4316unsafe fn find_first_ge_u32_sse(slice: &[u32], target: u32) -> usize {
4317 use std::arch::x86_64::*;
4318
4319 let n = slice.len();
4320 let ptr = slice.as_ptr();
4321
4322 let sign_flip = _mm_set1_epi32(i32::MIN);
4324 let target_xor = _mm_xor_si128(_mm_set1_epi32(target as i32), sign_flip);
4325
4326 let chunks = n / 16;
4327 let mut base = 0usize;
4328
4329 for _ in 0..chunks {
4331 let v0 = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4332 let v1 = _mm_xor_si128(
4333 _mm_loadu_si128(ptr.add(base + 4) as *const __m128i),
4334 sign_flip,
4335 );
4336 let v2 = _mm_xor_si128(
4337 _mm_loadu_si128(ptr.add(base + 8) as *const __m128i),
4338 sign_flip,
4339 );
4340 let v3 = _mm_xor_si128(
4341 _mm_loadu_si128(ptr.add(base + 12) as *const __m128i),
4342 sign_flip,
4343 );
4344
4345 let ge0 = _mm_or_si128(
4347 _mm_cmpeq_epi32(v0, target_xor),
4348 _mm_cmpgt_epi32(v0, target_xor),
4349 );
4350 let m0 = _mm_movemask_ps(_mm_castsi128_ps(ge0)) as u32;
4351 if m0 != 0 {
4352 return base + m0.trailing_zeros() as usize;
4353 }
4354
4355 let ge1 = _mm_or_si128(
4356 _mm_cmpeq_epi32(v1, target_xor),
4357 _mm_cmpgt_epi32(v1, target_xor),
4358 );
4359 let m1 = _mm_movemask_ps(_mm_castsi128_ps(ge1)) as u32;
4360 if m1 != 0 {
4361 return base + 4 + m1.trailing_zeros() as usize;
4362 }
4363
4364 let ge2 = _mm_or_si128(
4365 _mm_cmpeq_epi32(v2, target_xor),
4366 _mm_cmpgt_epi32(v2, target_xor),
4367 );
4368 let m2 = _mm_movemask_ps(_mm_castsi128_ps(ge2)) as u32;
4369 if m2 != 0 {
4370 return base + 8 + m2.trailing_zeros() as usize;
4371 }
4372
4373 let ge3 = _mm_or_si128(
4374 _mm_cmpeq_epi32(v3, target_xor),
4375 _mm_cmpgt_epi32(v3, target_xor),
4376 );
4377 let m3 = _mm_movemask_ps(_mm_castsi128_ps(ge3)) as u32;
4378 if m3 != 0 {
4379 return base + 12 + m3.trailing_zeros() as usize;
4380 }
4381 base += 16;
4382 }
4383
4384 while base + 4 <= n {
4386 let vals = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4387 let ge = _mm_or_si128(
4388 _mm_cmpeq_epi32(vals, target_xor),
4389 _mm_cmpgt_epi32(vals, target_xor),
4390 );
4391 let mask = _mm_movemask_ps(_mm_castsi128_ps(ge)) as u32;
4392 if mask != 0 {
4393 return base + mask.trailing_zeros() as usize;
4394 }
4395 base += 4;
4396 }
4397
4398 while base < n {
4400 if *slice.get_unchecked(base) >= target {
4401 return base;
4402 }
4403 base += 1;
4404 }
4405 n
4406}
4407
4408#[cfg(test)]
4409mod find_first_ge_tests {
4410 use super::find_first_ge_u32;
4411
4412 #[test]
4413 fn test_find_first_ge_basic() {
4414 let data: Vec<u32> = (0..128).map(|i| i * 3).collect(); assert_eq!(find_first_ge_u32(&data, 0), 0);
4416 assert_eq!(find_first_ge_u32(&data, 1), 1); assert_eq!(find_first_ge_u32(&data, 3), 1);
4418 assert_eq!(find_first_ge_u32(&data, 4), 2); assert_eq!(find_first_ge_u32(&data, 381), 127);
4420 assert_eq!(find_first_ge_u32(&data, 382), 128); }
4422
4423 #[test]
4424 fn test_find_first_ge_matches_partition_point() {
4425 let data: Vec<u32> = vec![1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75];
4426 for target in 0..80 {
4427 let expected = data.partition_point(|&d| d < target);
4428 let actual = find_first_ge_u32(&data, target);
4429 assert_eq!(actual, expected, "target={}", target);
4430 }
4431 }
4432
4433 #[test]
4434 fn test_find_first_ge_small_slices() {
4435 assert_eq!(find_first_ge_u32(&[], 5), 0);
4437 assert_eq!(find_first_ge_u32(&[10], 5), 0);
4439 assert_eq!(find_first_ge_u32(&[10], 10), 0);
4440 assert_eq!(find_first_ge_u32(&[10], 11), 1);
4441 assert_eq!(find_first_ge_u32(&[2, 4, 6], 5), 2);
4443 }
4444
4445 #[test]
4446 fn test_find_first_ge_full_block() {
4447 let data: Vec<u32> = (100..228).collect();
4449 assert_eq!(find_first_ge_u32(&data, 100), 0);
4450 assert_eq!(find_first_ge_u32(&data, 150), 50);
4451 assert_eq!(find_first_ge_u32(&data, 227), 127);
4452 assert_eq!(find_first_ge_u32(&data, 228), 128);
4453 assert_eq!(find_first_ge_u32(&data, 99), 0);
4454 }
4455
4456 #[test]
4457 fn test_find_first_ge_u32_max() {
4458 let data = vec![u32::MAX - 10, u32::MAX - 5, u32::MAX - 1, u32::MAX];
4460 assert_eq!(find_first_ge_u32(&data, u32::MAX - 10), 0);
4461 assert_eq!(find_first_ge_u32(&data, u32::MAX - 7), 1);
4462 assert_eq!(find_first_ge_u32(&data, u32::MAX), 3);
4463 }
4464}