1mod dispatch;
16pub mod postings;
18mod scalar;
19#[cfg(any(test, feature = "test-seams"))]
20#[doc(hidden)]
21pub mod vector_fault;
22
23#[cfg(target_arch = "x86_64")]
24mod avx2;
25#[cfg(target_arch = "aarch64")]
26mod neon;
27
28pub const MAX_DOT_I8_DIMENSION: usize = 65_536;
33
34const _: [(); 1] = [(); (MAX_DOT_I8_DIMENSION * (i8::MIN as i16 * i8::MIN as i16) as usize
35 <= i32::MAX as usize) as usize];
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum KernelArm {
40 Scalar,
42 Neon,
44 Avx2,
46}
47
48#[derive(Clone, Copy, Debug, Eq, PartialEq)]
50pub enum InstructionTier {
51 Scalar,
53 NeonWiden,
55 NeonDotprod,
57 Avx2,
59 NeonI8mmReserved,
61 Sme2Reserved,
63}
64
65#[derive(Clone, Copy, Debug, Eq, PartialEq)]
66#[cfg_attr(not(any(test, feature = "benchmark-kernels")), allow(dead_code))]
67enum KernelBackendTag {
68 Scalar,
69 #[cfg(target_arch = "aarch64")]
70 NeonWiden,
71 #[cfg(target_arch = "aarch64")]
72 NeonDotprodU4,
73 #[cfg(target_arch = "aarch64")]
74 NeonI8mm,
75 #[cfg(target_arch = "aarch64")]
76 NeonDotprodU2,
77 #[cfg(target_arch = "aarch64")]
78 NeonDotprodU6,
79 #[cfg(target_arch = "aarch64")]
80 NeonDotprodU8,
81 #[cfg(target_arch = "aarch64")]
82 NeonDotprodU4Prefetch,
83 #[cfg(target_arch = "x86_64")]
84 Avx2,
85}
86
87#[cfg(any(test, feature = "benchmark-kernels"))]
93#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
94#[doc(hidden)]
95pub enum KernelBackendId {
96 Scalar,
98 NeonWiden,
100 NeonDotprodU4,
102 NeonI8mm,
104 NeonDotprodU2,
106 NeonDotprodU6,
108 NeonDotprodU8,
110 NeonDotprodU4Prefetch,
112 Avx2,
114}
115
116#[cfg(any(test, feature = "benchmark-kernels"))]
117impl KernelBackendId {
118 const fn from_tag(tag: KernelBackendTag) -> Self {
119 match tag {
120 KernelBackendTag::Scalar => Self::Scalar,
121 #[cfg(target_arch = "aarch64")]
122 KernelBackendTag::NeonWiden => Self::NeonWiden,
123 #[cfg(target_arch = "aarch64")]
124 KernelBackendTag::NeonDotprodU4 => Self::NeonDotprodU4,
125 #[cfg(target_arch = "aarch64")]
126 KernelBackendTag::NeonI8mm => Self::NeonI8mm,
127 #[cfg(target_arch = "aarch64")]
128 KernelBackendTag::NeonDotprodU2 => Self::NeonDotprodU2,
129 #[cfg(target_arch = "aarch64")]
130 KernelBackendTag::NeonDotprodU6 => Self::NeonDotprodU6,
131 #[cfg(target_arch = "aarch64")]
132 KernelBackendTag::NeonDotprodU8 => Self::NeonDotprodU8,
133 #[cfg(target_arch = "aarch64")]
134 KernelBackendTag::NeonDotprodU4Prefetch => Self::NeonDotprodU4Prefetch,
135 #[cfg(target_arch = "x86_64")]
136 KernelBackendTag::Avx2 => Self::Avx2,
137 }
138 }
139
140 #[must_use]
142 pub const fn as_str(self) -> &'static str {
143 match self {
144 Self::Scalar => "scalar",
145 Self::NeonWiden => "neon-widen",
146 Self::NeonDotprodU4 => "neon-dotprod-u4",
147 Self::NeonI8mm => "neon-i8mm",
148 Self::NeonDotprodU2 => "neon-dotprod-u2",
149 Self::NeonDotprodU6 => "neon-dotprod-u6",
150 Self::NeonDotprodU8 => "neon-dotprod-u8",
151 Self::NeonDotprodU4Prefetch => "neon-dotprod-u4-prefetch",
152 Self::Avx2 => "avx2",
153 }
154 }
155
156 const fn arm(self) -> KernelArm {
157 match self {
158 Self::Scalar => KernelArm::Scalar,
159 Self::NeonWiden
160 | Self::NeonDotprodU4
161 | Self::NeonI8mm
162 | Self::NeonDotprodU2
163 | Self::NeonDotprodU6
164 | Self::NeonDotprodU8
165 | Self::NeonDotprodU4Prefetch => KernelArm::Neon,
166 Self::Avx2 => KernelArm::Avx2,
167 }
168 }
169}
170
171#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
173pub struct KernelFeatures {
174 pub neon: bool,
176 pub dotprod: bool,
178 pub fp16: bool,
180 pub i8mm: bool,
182 pub sme2: bool,
184 pub avx2: bool,
186 pub popcnt: bool,
188}
189
190#[derive(Clone, Debug, Eq, PartialEq)]
192pub enum KernelInitError {
193 UnknownOverride {
195 value: String,
197 },
198 NonUnicodeOverride,
200 UnsupportedArm {
202 requested: KernelArm,
204 },
205 AlreadyInitialized {
207 selected: KernelArm,
209 requested: KernelArm,
211 },
212}
213
214impl std::fmt::Display for KernelInitError {
215 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
216 match self {
217 Self::UnknownOverride { value } => write!(
218 formatter,
219 "unknown ZE_KERNEL value {value:?}; expected scalar, neon, or avx2"
220 ),
221 Self::NonUnicodeOverride => formatter.write_str("ZE_KERNEL is not valid Unicode"),
222 Self::UnsupportedArm { requested } => {
223 write!(
224 formatter,
225 "forced kernel arm {requested:?} is unsupported on this CPU"
226 )
227 }
228 Self::AlreadyInitialized {
229 selected,
230 requested,
231 } => write!(
232 formatter,
233 "kernel dispatch already selected {selected:?}; cannot force {requested:?}"
234 ),
235 }
236 }
237}
238
239impl std::error::Error for KernelInitError {}
240
241#[derive(Clone, Copy, Debug, Eq, PartialEq)]
246pub struct KnobSpace {
247 pub unroll: &'static [usize],
249 pub accumulators: &'static [usize],
251 pub rows_per_block: &'static [usize],
253 pub prefetch_dist: &'static [usize],
255 pub tier: &'static [InstructionTier],
257}
258
259#[derive(Clone, Copy, Debug, Eq, PartialEq)]
261pub struct BaselineKernelConfig {
262 pub unroll: usize,
264 pub accumulators: usize,
266 pub rows_per_block: usize,
268 pub prefetch_dist: usize,
270}
271
272pub const BASELINE_KERNEL_CONFIG: BaselineKernelConfig = BaselineKernelConfig {
274 unroll: 4,
276 accumulators: 4,
277 rows_per_block: 1,
278 prefetch_dist: 0,
279};
280
281const UNROLL_KNOBS: [usize; 4] = [2, 4, 6, 8];
283const ACCUMULATOR_KNOBS: [usize; 4] = [2, 4, 6, 8];
284const ROW_BLOCK_KNOBS: [usize; 5] = [1, 2, 4, 8, 16];
285const PREFETCH_KNOBS: [usize; 5] = [0, 1, 2, 4, 8];
286const TIER_KNOBS: [InstructionTier; 6] = [
287 InstructionTier::Scalar,
288 InstructionTier::NeonWiden,
289 InstructionTier::NeonDotprod,
290 InstructionTier::Avx2,
291 InstructionTier::NeonI8mmReserved,
292 InstructionTier::Sme2Reserved,
293]; pub const KERNEL_KNOB_SPACE: KnobSpace = KnobSpace {
297 unroll: &UNROLL_KNOBS,
298 accumulators: &ACCUMULATOR_KNOBS,
299 rows_per_block: &ROW_BLOCK_KNOBS,
300 prefetch_dist: &PREFETCH_KNOBS,
301 tier: &TIER_KNOBS,
302};
303
304type DotI8Fn = fn(&[i8], &[i8]) -> i32;
305type DotF32Fn = fn(&[f32], &[f32]) -> f32;
306type DotF16Fn = fn(&[u16], &[u16]) -> f32;
307type HammingU1Fn = fn(&[u8], &[u8]) -> u32;
308type DotI8BatchFn = fn(&[i8], &[i8], usize, &mut [i32]);
309type HammingU1BatchFn = fn(&[u8], &[u8], usize, &mut [u32]);
310type DotPackedFn = fn(&[i8], &[u8]) -> i32;
311type DotBit4PreparedFn = fn(&[i8], i32, &[u8]) -> i32;
312type DotPackedBatchFn = fn(&[i8], &[u8], usize, &mut [i32]);
313type ScoreBit4PreparedBatchFn =
314 fn(&[i8], i32, f64, &[u8], usize, &[crate::quant::Bit4Factors], &mut [f32]);
315type ScoreBit4PtrsFn = unsafe fn(
316 &[i8],
317 i32,
318 f64,
319 &[*const u8; 4],
320 usize,
321 &[crate::quant::Bit4Factors; 4],
322 &mut [f32; 4],
323);
324#[derive(Clone, Copy)]
325struct KernelTable {
326 #[cfg_attr(not(any(test, feature = "benchmark-kernels")), allow(dead_code))]
327 backend: KernelBackendTag,
328 arm: KernelArm,
329 tier: InstructionTier,
330 dot_i8: DotI8Fn,
331 dot_f32: DotF32Fn,
332 dot_f16: DotF16Fn,
333 hamming_u1: HammingU1Fn,
334 dot_i8_batch: DotI8BatchFn,
335 hamming_u1_batch: HammingU1BatchFn,
336 dot_bit4: DotPackedFn,
337 dot_bit4_prepared: DotBit4PreparedFn,
338 dot_bit4_batch: DotPackedBatchFn,
339 score_bit4_prepared_batch: ScoreBit4PreparedBatchFn,
340 score_bit4_ptrs: ScoreBit4PtrsFn,
341}
342
343#[derive(Clone, Copy)]
350pub struct Bit4Row<'a> {
351 bytes: &'a [u8],
352}
353
354impl std::fmt::Debug for Bit4Row<'_> {
355 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
356 formatter
357 .debug_struct("Bit4Row")
358 .field("len", &self.bytes.len())
359 .finish()
360 }
361}
362
363impl<'a> Bit4Row<'a> {
364 #[must_use]
369 pub fn from_mapped_region(region: &'a [u8], offset: usize, row_bytes: usize) -> Option<Self> {
370 let end = offset.checked_add(row_bytes)?;
371 region.get(offset..end).map(|bytes| Self { bytes })
372 }
373
374 #[must_use]
376 pub(crate) const fn from_validated_bytes(bytes: &'a [u8]) -> Self {
377 Self { bytes }
378 }
379
380 const fn as_ptr(self) -> *const u8 {
381 self.bytes.as_ptr()
382 }
383
384 const fn len(self) -> usize {
385 self.bytes.len()
386 }
387}
388
389#[derive(Clone, Copy)]
391pub struct Bit4Rows4<'a> {
392 pointers: [*const u8; 4],
393 row_bytes: usize,
394 lifetime: std::marker::PhantomData<&'a [u8]>,
395}
396
397impl std::fmt::Debug for Bit4Rows4<'_> {
398 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
399 formatter
400 .debug_struct("Bit4Rows4")
401 .field("row_bytes", &self.row_bytes)
402 .finish()
403 }
404}
405
406impl<'a> Bit4Rows4<'a> {
407 pub fn from_rows(rows: [Bit4Row<'a>; 4], row_bytes: usize) -> Result<Self, GatherShapeError> {
414 for (index, row) in rows.iter().enumerate() {
415 if row.len() != row_bytes {
416 return Err(GatherShapeError::RowLength {
417 index,
418 expected: row_bytes,
419 actual: row.len(),
420 });
421 }
422 }
423 Ok(Self {
424 pointers: rows.map(Bit4Row::as_ptr),
425 row_bytes,
426 lifetime: std::marker::PhantomData,
427 })
428 }
429}
430
431#[derive(Clone, Copy, Debug, Eq, PartialEq)]
433pub enum GatherShapeError {
434 DimensionTooLarge {
436 actual: usize,
438 maximum: usize,
440 },
441 RowBytes {
443 expected: usize,
445 actual: usize,
447 },
448 RowLength {
450 index: usize,
452 expected: usize,
454 actual: usize,
456 },
457}
458
459impl std::fmt::Display for GatherShapeError {
460 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
461 match self {
462 Self::DimensionTooLarge { actual, maximum } => write!(
463 formatter,
464 "gather dimension {actual} exceeds supported maximum {maximum}"
465 ),
466 Self::RowBytes { expected, actual } => write!(
467 formatter,
468 "gather row width mismatch: expected {expected}, got {actual}"
469 ),
470 Self::RowLength {
471 index,
472 expected,
473 actual,
474 } => write!(
475 formatter,
476 "gather row {index} has {actual} readable bytes, expected {expected}"
477 ),
478 }
479 }
480}
481
482impl std::error::Error for GatherShapeError {}
483
484pub fn prefetch_bit4_rows(rows: &[Bit4Row<'_>]) {
490 #[cfg(target_arch = "aarch64")]
491 for row in rows.iter().filter(|row| row.len() != 0) {
492 let address = row.as_ptr();
493 unsafe {
496 std::arch::asm!(
497 "prfm pldl1keep, [{address}]",
498 address = in(reg) address,
499 options(readonly, nostack)
500 );
501 }
502 }
503 #[cfg(not(target_arch = "aarch64"))]
504 let _ = rows;
505}
506
507pub fn prefetch_bit4_row_group(rows: &Bit4Rows4<'_>) {
509 #[cfg(target_arch = "aarch64")]
510 if rows.row_bytes != 0 {
511 for &address in &rows.pointers {
512 unsafe {
515 std::arch::asm!(
516 "prfm pldl1keep, [{address}]",
517 address = in(reg) address,
518 options(readonly, nostack)
519 );
520 }
521 }
522 }
523 #[cfg(not(target_arch = "aarch64"))]
524 let _ = rows;
525}
526
527#[derive(Clone, Copy)]
533pub struct KernelVariant {
534 table: KernelTable,
535}
536
537impl std::fmt::Debug for KernelVariant {
538 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539 formatter
540 .debug_struct("KernelVariant")
541 .field("arm", &self.table.arm)
542 .field("tier", &self.table.tier)
543 .finish()
544 }
545}
546
547impl KernelVariant {
548 #[must_use]
550 pub fn scalar() -> Self {
551 Self {
552 table: scalar::table(),
553 }
554 }
555
556 #[doc(hidden)]
558 #[must_use]
559 pub fn selected() -> Self {
560 Self {
561 table: *dispatch::active_table(),
562 }
563 }
564
565 #[cfg(any(test, feature = "benchmark-kernels"))]
575 pub fn available() -> impl Iterator<Item = Self> {
576 dispatch::variant_tables()
577 .into_iter()
578 .flatten()
579 .map(|table| Self { table })
580 }
581
582 #[must_use]
584 pub fn arm(self) -> KernelArm {
585 self.table.arm
586 }
587
588 #[must_use]
590 pub fn tier(self) -> InstructionTier {
591 self.table.tier
592 }
593
594 #[cfg(any(test, feature = "benchmark-kernels"))]
596 #[doc(hidden)]
597 #[must_use]
598 pub const fn backend_id(self) -> KernelBackendId {
599 KernelBackendId::from_tag(self.table.backend)
600 }
601
602 #[must_use]
604 pub fn dot_i8(self, a: &[i8], b: &[i8]) -> i32 {
605 (self.table.dot_i8)(a, b)
606 }
607
608 #[must_use]
610 pub fn dot_f32(self, a: &[f32], b: &[f32]) -> f32 {
611 (self.table.dot_f32)(a, b)
612 }
613
614 #[must_use]
616 pub fn dot_f16(self, a: &[u16], b: &[u16]) -> f32 {
617 (self.table.dot_f16)(a, b)
618 }
619
620 #[must_use]
622 pub fn hamming_u1(self, a: &[u8], b: &[u8]) -> u32 {
623 (self.table.hamming_u1)(a, b)
624 }
625
626 pub fn dot_i8_batch(self, q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
628 (self.table.dot_i8_batch)(q, rows, d, out);
629 }
630
631 pub fn hamming_u1_batch(self, q: &[u8], rows: &[u8], d_bytes: usize, out: &mut [u32]) {
633 (self.table.hamming_u1_batch)(q, rows, d_bytes, out);
634 }
635
636 #[must_use]
638 pub fn dot_bit4(self, q: &[i8], codes: &[u8]) -> i32 {
639 (self.table.dot_bit4)(q, codes)
640 }
641
642 #[doc(hidden)]
647 #[must_use]
648 pub fn dot_bit4_prepared(self, q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
649 (self.table.dot_bit4_prepared)(q, query_sum, codes)
650 }
651
652 pub fn dot_bit4_batch(self, q: &[i8], rows: &[u8], d: usize, out: &mut [i32]) {
654 (self.table.dot_bit4_batch)(q, rows, d, out);
655 }
656
657 #[doc(hidden)]
659 pub fn score_bit4_prepared_batch(
660 self,
661 query: (&[i8], i32, f64),
662 rows: &[u8],
663 d: usize,
664 factors: &[crate::quant::Bit4Factors],
665 out: &mut [f32],
666 ) {
667 let (q, query_sum, query_scale_half) = query;
668 (self.table.score_bit4_prepared_batch)(
669 q,
670 query_sum,
671 query_scale_half,
672 rows,
673 d,
674 factors,
675 out,
676 );
677 }
678
679 pub fn score_bit4_ptrs(
688 &self,
689 query: (&[i8], i32, f64),
690 rows: &Bit4Rows4<'_>,
691 factors: &[crate::quant::Bit4Factors; 4],
692 out: &mut [f32; 4],
693 ) -> Result<(), GatherShapeError> {
694 score_bit4_ptrs_with_table(&self.table, query, rows, factors, out)
695 }
696}
697
698fn score_bit4_ptrs_with_table(
699 table: &KernelTable,
700 query: (&[i8], i32, f64),
701 rows: &Bit4Rows4<'_>,
702 factors: &[crate::quant::Bit4Factors; 4],
703 out: &mut [f32; 4],
704) -> Result<(), GatherShapeError> {
705 let (q, query_sum, query_scale_half) = query;
706 if q.len() > MAX_DOT_I8_DIMENSION {
707 return Err(GatherShapeError::DimensionTooLarge {
708 actual: q.len(),
709 maximum: MAX_DOT_I8_DIMENSION,
710 });
711 }
712 let expected_row_bytes = q.len().div_ceil(2);
713 if rows.row_bytes != expected_row_bytes {
714 return Err(GatherShapeError::RowBytes {
715 expected: expected_row_bytes,
716 actual: rows.row_bytes,
717 });
718 }
719 debug_assert!(q.len() <= MAX_DOT_I8_DIMENSION);
720 debug_assert_eq!(rows.row_bytes, q.len().div_ceil(2));
721 unsafe {
725 (table.score_bit4_ptrs)(
726 q,
727 query_sum,
728 query_scale_half,
729 &rows.pointers,
730 rows.row_bytes,
731 factors,
732 out,
733 );
734 }
735 Ok(())
736}
737
738pub fn initialize() -> Result<KernelArm, KernelInitError> {
743 dispatch::initialize()
744}
745
746#[must_use]
748pub fn selected_arm() -> KernelArm {
749 dispatch::active_table().arm
750}
751
752#[must_use]
754pub fn is_arm_supported(arm: KernelArm) -> bool {
755 dispatch::table_for_arm(arm).is_some()
756}
757
758#[must_use]
760pub fn detected_features() -> KernelFeatures {
761 dispatch::features()
762}
763
764#[doc(hidden)]
770#[must_use]
771pub fn platform_wide_stream_checksum(bytes: &[u8]) -> Option<u64> {
772 #[cfg(target_arch = "aarch64")]
773 {
774 dispatch::features()
775 .neon
776 .then(|| neon::wide_stream_checksum(bytes))
777 }
778 #[cfg(not(target_arch = "aarch64"))]
779 {
780 let _ = bytes;
781 None
782 }
783}
784
785#[must_use]
790pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
791 let table = dispatch::active_table();
792 let result = (table.dot_i8)(a, b);
793 #[cfg(any(test, feature = "test-seams"))]
794 observe_selected_score(table, vector_fault_kernel_id_dot_i8(), a.len());
795 #[cfg(any(test, feature = "test-seams"))]
796 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32(result));
797 result
798}
799
800#[must_use]
805pub fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
806 let table = dispatch::active_table();
807 let result = (table.dot_f32)(a, b);
808 #[cfg(any(test, feature = "test-seams"))]
809 observe_selected_score(table, vector_fault_kernel_id_dot_f32(), a.len());
810 #[cfg(any(test, feature = "test-seams"))]
811 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32(result.to_bits()));
812 result
813}
814
815#[must_use]
821pub fn dot_f16(a: &[u16], b: &[u16]) -> f32 {
822 let table = dispatch::active_table();
823 let result = (table.dot_f16)(a, b);
824 #[cfg(any(test, feature = "test-seams"))]
825 observe_selected_score(table, vector_fault_kernel_id_dot_f16(), a.len());
826 #[cfg(any(test, feature = "test-seams"))]
827 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32(result.to_bits()));
828 result
829}
830
831#[must_use]
837pub fn hamming_u1(a: &[u8], b: &[u8]) -> u32 {
838 let table = dispatch::active_table();
839 let result = (table.hamming_u1)(a, b);
840 #[cfg(any(test, feature = "test-seams"))]
841 observe_selected_score(table, vector_fault_kernel_id_hamming(), a.len());
842 #[cfg(any(test, feature = "test-seams"))]
843 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::U32(result));
844 result
845}
846
847pub fn dot_i8_batch(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
853 let table = dispatch::active_table();
854 (table.dot_i8_batch)(q, rows, d, out);
855 #[cfg(any(test, feature = "test-seams"))]
856 observe_selected_score(
857 table,
858 vector_fault_kernel_id_dot_i8_batch(),
859 d.saturating_mul(out.len()),
860 );
861 #[cfg(any(test, feature = "test-seams"))]
862 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32s(out.to_vec()));
863}
864
865pub fn hamming_u1_batch(q: &[u8], rows: &[u8], d_bytes: usize, out: &mut [u32]) {
871 let table = dispatch::active_table();
872 (table.hamming_u1_batch)(q, rows, d_bytes, out);
873 #[cfg(any(test, feature = "test-seams"))]
874 observe_selected_score(
875 table,
876 vector_fault_kernel_id_hamming_batch(),
877 d_bytes.saturating_mul(out.len()),
878 );
879 #[cfg(any(test, feature = "test-seams"))]
880 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::U32s(out.to_vec()));
881}
882
883#[must_use]
891pub fn dot_bit4(q: &[i8], codes: &[u8]) -> i32 {
892 let table = dispatch::active_table();
893 let result = (table.dot_bit4)(q, codes);
894 #[cfg(any(test, feature = "test-seams"))]
895 observe_selected_score(table, vector_fault_kernel_id_dot_bit4(), q.len());
896 #[cfg(any(test, feature = "test-seams"))]
897 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32(result));
898 result
899}
900
901pub(crate) fn dot_bit4_prepared(q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
902 let table = dispatch::active_table();
903 let result = (table.dot_bit4_prepared)(q, query_sum, codes);
904 #[cfg(any(test, feature = "test-seams"))]
905 observe_selected_score(table, vector_fault_kernel_id_dot_bit4_prepared(), q.len());
906 #[cfg(any(test, feature = "test-seams"))]
907 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32(result));
908 result
909}
910
911pub(crate) fn score_bit4_prepared_batch(
912 q: &[i8],
913 query_sum: i32,
914 query_scale_half: f64,
915 rows: &[u8],
916 d: usize,
917 factors: &[crate::quant::Bit4Factors],
918 out: &mut [f32],
919) {
920 let table = dispatch::active_table();
921 (table.score_bit4_prepared_batch)(q, query_sum, query_scale_half, rows, d, factors, out);
922 #[cfg(any(test, feature = "test-seams"))]
923 observe_selected_score(
924 table,
925 vector_fault_kernel_id_score_bit4_batch(),
926 d.saturating_mul(out.len()),
927 );
928 #[cfg(any(test, feature = "test-seams"))]
929 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32s(
930 out.iter().map(|value| value.to_bits()).collect(),
931 ));
932}
933
934#[cfg(any(test, feature = "test-seams"))]
935fn observe_selected_score(
936 table: &KernelTable,
937 kernel: vector_fault::KernelOperationId,
938 work_items: usize,
939) {
940 vector_fault::observe_selected_score(table.backend, kernel, work_items);
941}
942
943macro_rules! vector_fault_kernel_id {
944 ($name:ident, $variant:ident) => {
945 #[cfg(any(test, feature = "test-seams"))]
946 const fn $name() -> vector_fault::KernelOperationId {
947 vector_fault::KernelOperationId::$variant
948 }
949 };
950}
951
952vector_fault_kernel_id!(vector_fault_kernel_id_dot_i8, DotI8);
953vector_fault_kernel_id!(vector_fault_kernel_id_hamming, HammingU1);
954vector_fault_kernel_id!(vector_fault_kernel_id_dot_f32, DotF32);
955vector_fault_kernel_id!(vector_fault_kernel_id_dot_f16, DotF16);
956vector_fault_kernel_id!(vector_fault_kernel_id_dot_i8_batch, DotI8Batch);
957vector_fault_kernel_id!(vector_fault_kernel_id_hamming_batch, HammingU1Batch);
958vector_fault_kernel_id!(vector_fault_kernel_id_dot_bit4, DotBit4);
959vector_fault_kernel_id!(vector_fault_kernel_id_dot_bit4_prepared, DotBit4Prepared);
960vector_fault_kernel_id!(vector_fault_kernel_id_dot_bit4_batch, DotBit4Batch);
961vector_fault_kernel_id!(
962 vector_fault_kernel_id_score_bit4_batch,
963 ScoreBit4PreparedBatch
964);
965vector_fault_kernel_id!(vector_fault_kernel_id_score_bit4_ptrs, ScoreBit4Ptrs);
966
967pub fn score_bit4_ptrs(
976 q: &[i8],
977 query_sum: i32,
978 query_scale_half: f64,
979 rows: &Bit4Rows4<'_>,
980 factors: &[crate::quant::Bit4Factors; 4],
981 out: &mut [f32; 4],
982) -> Result<(), GatherShapeError> {
983 let table = dispatch::active_table();
984 let result =
985 score_bit4_ptrs_with_table(table, (q, query_sum, query_scale_half), rows, factors, out);
986 #[cfg(any(test, feature = "test-seams"))]
987 if result.is_ok() {
988 #[cfg(any(test, feature = "test-seams"))]
989 observe_selected_score(
990 table,
991 vector_fault_kernel_id_score_bit4_ptrs(),
992 q.len().saturating_mul(out.len()),
993 );
994 #[cfg(any(test, feature = "test-seams"))]
995 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32s(
996 out.iter().map(|value| value.to_bits()).collect(),
997 ));
998 }
999 result
1000}
1001
1002pub fn dot_bit4_batch(q: &[i8], rows: &[u8], d: usize, out: &mut [i32]) {
1008 let table = dispatch::active_table();
1009 (table.dot_bit4_batch)(q, rows, d, out);
1010 #[cfg(any(test, feature = "test-seams"))]
1011 observe_selected_score(
1012 table,
1013 vector_fault_kernel_id_dot_bit4_batch(),
1014 d.saturating_mul(out.len()),
1015 );
1016 #[cfg(any(test, feature = "test-seams"))]
1017 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32s(out.to_vec()));
1018}