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 #[cfg(any(test, feature = "test-seams"))]
157 const fn arm(self) -> KernelArm {
158 match self {
159 Self::Scalar => KernelArm::Scalar,
160 Self::NeonWiden
161 | Self::NeonDotprodU4
162 | Self::NeonI8mm
163 | Self::NeonDotprodU2
164 | Self::NeonDotprodU6
165 | Self::NeonDotprodU8
166 | Self::NeonDotprodU4Prefetch => KernelArm::Neon,
167 Self::Avx2 => KernelArm::Avx2,
168 }
169 }
170}
171
172#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
174pub struct KernelFeatures {
175 pub neon: bool,
177 pub dotprod: bool,
179 pub fp16: bool,
181 pub i8mm: bool,
183 pub sme2: bool,
185 pub avx2: bool,
187 pub popcnt: bool,
189}
190
191#[derive(Clone, Debug, Eq, PartialEq)]
193pub enum KernelInitError {
194 UnknownOverride {
196 value: String,
198 },
199 NonUnicodeOverride,
201 UnsupportedArm {
203 requested: KernelArm,
205 },
206 AlreadyInitialized {
208 selected: KernelArm,
210 requested: KernelArm,
212 },
213}
214
215impl std::fmt::Display for KernelInitError {
216 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
217 match self {
218 Self::UnknownOverride { value } => write!(
219 formatter,
220 "unknown ZE_KERNEL value {value:?}; expected scalar, neon, or avx2"
221 ),
222 Self::NonUnicodeOverride => formatter.write_str("ZE_KERNEL is not valid Unicode"),
223 Self::UnsupportedArm { requested } => {
224 write!(
225 formatter,
226 "forced kernel arm {requested:?} is unsupported on this CPU"
227 )
228 }
229 Self::AlreadyInitialized {
230 selected,
231 requested,
232 } => write!(
233 formatter,
234 "kernel dispatch already selected {selected:?}; cannot force {requested:?}"
235 ),
236 }
237 }
238}
239
240impl std::error::Error for KernelInitError {}
241
242#[derive(Clone, Copy, Debug, Eq, PartialEq)]
247pub struct KnobSpace {
248 pub unroll: &'static [usize],
250 pub accumulators: &'static [usize],
252 pub rows_per_block: &'static [usize],
254 pub prefetch_dist: &'static [usize],
256 pub tier: &'static [InstructionTier],
258}
259
260#[derive(Clone, Copy, Debug, Eq, PartialEq)]
262pub struct BaselineKernelConfig {
263 pub unroll: usize,
265 pub accumulators: usize,
267 pub rows_per_block: usize,
269 pub prefetch_dist: usize,
271}
272
273pub const BASELINE_KERNEL_CONFIG: BaselineKernelConfig = BaselineKernelConfig {
275 unroll: 4,
277 accumulators: 4,
278 rows_per_block: 1,
279 prefetch_dist: 0,
280};
281
282const UNROLL_KNOBS: [usize; 4] = [2, 4, 6, 8];
284const ACCUMULATOR_KNOBS: [usize; 4] = [2, 4, 6, 8];
285const ROW_BLOCK_KNOBS: [usize; 5] = [1, 2, 4, 8, 16];
286const PREFETCH_KNOBS: [usize; 5] = [0, 1, 2, 4, 8];
287const TIER_KNOBS: [InstructionTier; 6] = [
288 InstructionTier::Scalar,
289 InstructionTier::NeonWiden,
290 InstructionTier::NeonDotprod,
291 InstructionTier::Avx2,
292 InstructionTier::NeonI8mmReserved,
293 InstructionTier::Sme2Reserved,
294]; pub const KERNEL_KNOB_SPACE: KnobSpace = KnobSpace {
298 unroll: &UNROLL_KNOBS,
299 accumulators: &ACCUMULATOR_KNOBS,
300 rows_per_block: &ROW_BLOCK_KNOBS,
301 prefetch_dist: &PREFETCH_KNOBS,
302 tier: &TIER_KNOBS,
303};
304
305type DotI8Fn = fn(&[i8], &[i8]) -> i32;
306type DotF32Fn = fn(&[f32], &[f32]) -> f32;
307type DotF16Fn = fn(&[u16], &[u16]) -> f32;
308type HammingU1Fn = fn(&[u8], &[u8]) -> u32;
309type DotI8BatchFn = fn(&[i8], &[i8], usize, &mut [i32]);
310type HammingU1BatchFn = fn(&[u8], &[u8], usize, &mut [u32]);
311type DotPackedFn = fn(&[i8], &[u8]) -> i32;
312type DotBit4PreparedFn = fn(&[i8], i32, &[u8]) -> i32;
313type DotPackedBatchFn = fn(&[i8], &[u8], usize, &mut [i32]);
314type ScoreBit4PreparedBatchFn =
315 fn(&[i8], i32, f64, &[u8], usize, &[crate::quant::Bit4Factors], &mut [f32]);
316type ScoreBit4PtrsFn = unsafe fn(
317 &[i8],
318 i32,
319 f64,
320 &[*const u8; 4],
321 usize,
322 &[crate::quant::Bit4Factors; 4],
323 &mut [f32; 4],
324);
325#[derive(Clone, Copy)]
326struct KernelTable {
327 #[cfg_attr(not(any(test, feature = "benchmark-kernels")), allow(dead_code))]
328 backend: KernelBackendTag,
329 arm: KernelArm,
330 tier: InstructionTier,
331 dot_i8: DotI8Fn,
332 dot_f32: DotF32Fn,
333 dot_f16: DotF16Fn,
334 hamming_u1: HammingU1Fn,
335 dot_i8_batch: DotI8BatchFn,
336 hamming_u1_batch: HammingU1BatchFn,
337 dot_bit4: DotPackedFn,
338 dot_bit4_prepared: DotBit4PreparedFn,
339 dot_bit4_batch: DotPackedBatchFn,
340 score_bit4_prepared_batch: ScoreBit4PreparedBatchFn,
341 score_bit4_ptrs: ScoreBit4PtrsFn,
342}
343
344#[derive(Clone, Copy)]
351pub struct Bit4Row<'a> {
352 bytes: &'a [u8],
353}
354
355impl std::fmt::Debug for Bit4Row<'_> {
356 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
357 formatter
358 .debug_struct("Bit4Row")
359 .field("len", &self.bytes.len())
360 .finish()
361 }
362}
363
364impl<'a> Bit4Row<'a> {
365 #[must_use]
370 pub fn from_mapped_region(region: &'a [u8], offset: usize, row_bytes: usize) -> Option<Self> {
371 let end = offset.checked_add(row_bytes)?;
372 region.get(offset..end).map(|bytes| Self { bytes })
373 }
374
375 #[must_use]
377 pub(crate) const fn from_validated_bytes(bytes: &'a [u8]) -> Self {
378 Self { bytes }
379 }
380
381 const fn as_ptr(self) -> *const u8 {
382 self.bytes.as_ptr()
383 }
384
385 const fn len(self) -> usize {
386 self.bytes.len()
387 }
388}
389
390#[derive(Clone, Copy)]
392pub struct Bit4Rows4<'a> {
393 pointers: [*const u8; 4],
394 row_bytes: usize,
395 lifetime: std::marker::PhantomData<&'a [u8]>,
396}
397
398impl std::fmt::Debug for Bit4Rows4<'_> {
399 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
400 formatter
401 .debug_struct("Bit4Rows4")
402 .field("row_bytes", &self.row_bytes)
403 .finish()
404 }
405}
406
407impl<'a> Bit4Rows4<'a> {
408 pub fn from_rows(rows: [Bit4Row<'a>; 4], row_bytes: usize) -> Result<Self, GatherShapeError> {
415 for (index, row) in rows.iter().enumerate() {
416 if row.len() != row_bytes {
417 return Err(GatherShapeError::RowLength {
418 index,
419 expected: row_bytes,
420 actual: row.len(),
421 });
422 }
423 }
424 Ok(Self {
425 pointers: rows.map(Bit4Row::as_ptr),
426 row_bytes,
427 lifetime: std::marker::PhantomData,
428 })
429 }
430}
431
432#[derive(Clone, Copy, Debug, Eq, PartialEq)]
434pub enum GatherShapeError {
435 DimensionTooLarge {
437 actual: usize,
439 maximum: usize,
441 },
442 RowBytes {
444 expected: usize,
446 actual: usize,
448 },
449 RowLength {
451 index: usize,
453 expected: usize,
455 actual: usize,
457 },
458}
459
460impl std::fmt::Display for GatherShapeError {
461 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
462 match self {
463 Self::DimensionTooLarge { actual, maximum } => write!(
464 formatter,
465 "gather dimension {actual} exceeds supported maximum {maximum}"
466 ),
467 Self::RowBytes { expected, actual } => write!(
468 formatter,
469 "gather row width mismatch: expected {expected}, got {actual}"
470 ),
471 Self::RowLength {
472 index,
473 expected,
474 actual,
475 } => write!(
476 formatter,
477 "gather row {index} has {actual} readable bytes, expected {expected}"
478 ),
479 }
480 }
481}
482
483impl std::error::Error for GatherShapeError {}
484
485pub fn prefetch_bit4_rows(rows: &[Bit4Row<'_>]) {
491 #[cfg(target_arch = "aarch64")]
492 for row in rows.iter().filter(|row| row.len() != 0) {
493 let address = row.as_ptr();
494 unsafe {
497 std::arch::asm!(
498 "prfm pldl1keep, [{address}]",
499 address = in(reg) address,
500 options(readonly, nostack)
501 );
502 }
503 }
504 #[cfg(not(target_arch = "aarch64"))]
505 let _ = rows;
506}
507
508pub fn prefetch_bit4_row_group(rows: &Bit4Rows4<'_>) {
510 #[cfg(target_arch = "aarch64")]
511 if rows.row_bytes != 0 {
512 for &address in &rows.pointers {
513 unsafe {
516 std::arch::asm!(
517 "prfm pldl1keep, [{address}]",
518 address = in(reg) address,
519 options(readonly, nostack)
520 );
521 }
522 }
523 }
524 #[cfg(not(target_arch = "aarch64"))]
525 let _ = rows;
526}
527
528#[derive(Clone, Copy)]
534pub struct KernelVariant {
535 table: KernelTable,
536}
537
538impl std::fmt::Debug for KernelVariant {
539 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
540 formatter
541 .debug_struct("KernelVariant")
542 .field("arm", &self.table.arm)
543 .field("tier", &self.table.tier)
544 .finish()
545 }
546}
547
548impl KernelVariant {
549 #[must_use]
551 pub fn scalar() -> Self {
552 Self {
553 table: scalar::table(),
554 }
555 }
556
557 #[doc(hidden)]
559 #[must_use]
560 pub fn selected() -> Self {
561 Self {
562 table: *dispatch::active_table(),
563 }
564 }
565
566 #[cfg(any(test, feature = "benchmark-kernels"))]
576 pub fn available() -> impl Iterator<Item = Self> {
577 dispatch::variant_tables()
578 .into_iter()
579 .flatten()
580 .map(|table| Self { table })
581 }
582
583 #[must_use]
585 pub fn arm(self) -> KernelArm {
586 self.table.arm
587 }
588
589 #[must_use]
591 pub fn tier(self) -> InstructionTier {
592 self.table.tier
593 }
594
595 #[cfg(any(test, feature = "benchmark-kernels"))]
597 #[doc(hidden)]
598 #[must_use]
599 pub const fn backend_id(self) -> KernelBackendId {
600 KernelBackendId::from_tag(self.table.backend)
601 }
602
603 #[must_use]
605 pub fn dot_i8(self, a: &[i8], b: &[i8]) -> i32 {
606 (self.table.dot_i8)(a, b)
607 }
608
609 #[must_use]
611 pub fn dot_f32(self, a: &[f32], b: &[f32]) -> f32 {
612 (self.table.dot_f32)(a, b)
613 }
614
615 #[must_use]
617 pub fn dot_f16(self, a: &[u16], b: &[u16]) -> f32 {
618 (self.table.dot_f16)(a, b)
619 }
620
621 #[must_use]
623 pub fn hamming_u1(self, a: &[u8], b: &[u8]) -> u32 {
624 (self.table.hamming_u1)(a, b)
625 }
626
627 pub fn dot_i8_batch(self, q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
629 (self.table.dot_i8_batch)(q, rows, d, out);
630 }
631
632 pub fn hamming_u1_batch(self, q: &[u8], rows: &[u8], d_bytes: usize, out: &mut [u32]) {
634 (self.table.hamming_u1_batch)(q, rows, d_bytes, out);
635 }
636
637 #[must_use]
639 pub fn dot_bit4(self, q: &[i8], codes: &[u8]) -> i32 {
640 (self.table.dot_bit4)(q, codes)
641 }
642
643 #[doc(hidden)]
648 #[must_use]
649 pub fn dot_bit4_prepared(self, q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
650 (self.table.dot_bit4_prepared)(q, query_sum, codes)
651 }
652
653 pub fn dot_bit4_batch(self, q: &[i8], rows: &[u8], d: usize, out: &mut [i32]) {
655 (self.table.dot_bit4_batch)(q, rows, d, out);
656 }
657
658 #[doc(hidden)]
660 pub fn score_bit4_prepared_batch(
661 self,
662 query: (&[i8], i32, f64),
663 rows: &[u8],
664 d: usize,
665 factors: &[crate::quant::Bit4Factors],
666 out: &mut [f32],
667 ) {
668 let (q, query_sum, query_scale_half) = query;
669 (self.table.score_bit4_prepared_batch)(
670 q,
671 query_sum,
672 query_scale_half,
673 rows,
674 d,
675 factors,
676 out,
677 );
678 }
679
680 pub fn score_bit4_ptrs(
689 &self,
690 query: (&[i8], i32, f64),
691 rows: &Bit4Rows4<'_>,
692 factors: &[crate::quant::Bit4Factors; 4],
693 out: &mut [f32; 4],
694 ) -> Result<(), GatherShapeError> {
695 score_bit4_ptrs_with_table(&self.table, query, rows, factors, out)
696 }
697}
698
699fn score_bit4_ptrs_with_table(
700 table: &KernelTable,
701 query: (&[i8], i32, f64),
702 rows: &Bit4Rows4<'_>,
703 factors: &[crate::quant::Bit4Factors; 4],
704 out: &mut [f32; 4],
705) -> Result<(), GatherShapeError> {
706 let (q, query_sum, query_scale_half) = query;
707 if q.len() > MAX_DOT_I8_DIMENSION {
708 return Err(GatherShapeError::DimensionTooLarge {
709 actual: q.len(),
710 maximum: MAX_DOT_I8_DIMENSION,
711 });
712 }
713 let expected_row_bytes = q.len().div_ceil(2);
714 if rows.row_bytes != expected_row_bytes {
715 return Err(GatherShapeError::RowBytes {
716 expected: expected_row_bytes,
717 actual: rows.row_bytes,
718 });
719 }
720 debug_assert!(q.len() <= MAX_DOT_I8_DIMENSION);
721 debug_assert_eq!(rows.row_bytes, q.len().div_ceil(2));
722 unsafe {
726 (table.score_bit4_ptrs)(
727 q,
728 query_sum,
729 query_scale_half,
730 &rows.pointers,
731 rows.row_bytes,
732 factors,
733 out,
734 );
735 }
736 Ok(())
737}
738
739pub fn initialize() -> Result<KernelArm, KernelInitError> {
744 dispatch::initialize()
745}
746
747#[must_use]
749pub fn selected_arm() -> KernelArm {
750 dispatch::active_table().arm
751}
752
753#[must_use]
755pub fn is_arm_supported(arm: KernelArm) -> bool {
756 dispatch::table_for_arm(arm).is_some()
757}
758
759#[must_use]
761pub fn detected_features() -> KernelFeatures {
762 dispatch::features()
763}
764
765#[doc(hidden)]
771#[must_use]
772pub fn platform_wide_stream_checksum(bytes: &[u8]) -> Option<u64> {
773 #[cfg(target_arch = "aarch64")]
774 {
775 dispatch::features()
776 .neon
777 .then(|| neon::wide_stream_checksum(bytes))
778 }
779 #[cfg(not(target_arch = "aarch64"))]
780 {
781 let _ = bytes;
782 None
783 }
784}
785
786#[must_use]
791pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
792 let table = dispatch::active_table();
793 let result = (table.dot_i8)(a, b);
794 #[cfg(any(test, feature = "test-seams"))]
795 observe_selected_score(table, vector_fault_kernel_id_dot_i8(), a.len());
796 #[cfg(any(test, feature = "test-seams"))]
797 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32(result));
798 result
799}
800
801#[must_use]
806pub fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
807 let table = dispatch::active_table();
808 let result = (table.dot_f32)(a, b);
809 #[cfg(any(test, feature = "test-seams"))]
810 observe_selected_score(table, vector_fault_kernel_id_dot_f32(), a.len());
811 #[cfg(any(test, feature = "test-seams"))]
812 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32(result.to_bits()));
813 result
814}
815
816#[must_use]
822pub fn dot_f16(a: &[u16], b: &[u16]) -> f32 {
823 let table = dispatch::active_table();
824 let result = (table.dot_f16)(a, b);
825 #[cfg(any(test, feature = "test-seams"))]
826 observe_selected_score(table, vector_fault_kernel_id_dot_f16(), a.len());
827 #[cfg(any(test, feature = "test-seams"))]
828 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32(result.to_bits()));
829 result
830}
831
832#[must_use]
838pub fn hamming_u1(a: &[u8], b: &[u8]) -> u32 {
839 let table = dispatch::active_table();
840 let result = (table.hamming_u1)(a, b);
841 #[cfg(any(test, feature = "test-seams"))]
842 observe_selected_score(table, vector_fault_kernel_id_hamming(), a.len());
843 #[cfg(any(test, feature = "test-seams"))]
844 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::U32(result));
845 result
846}
847
848pub fn dot_i8_batch(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
854 let table = dispatch::active_table();
855 (table.dot_i8_batch)(q, rows, d, out);
856 #[cfg(any(test, feature = "test-seams"))]
857 observe_selected_score(
858 table,
859 vector_fault_kernel_id_dot_i8_batch(),
860 d.saturating_mul(out.len()),
861 );
862 #[cfg(any(test, feature = "test-seams"))]
863 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32s(out.to_vec()));
864}
865
866pub fn hamming_u1_batch(q: &[u8], rows: &[u8], d_bytes: usize, out: &mut [u32]) {
872 let table = dispatch::active_table();
873 (table.hamming_u1_batch)(q, rows, d_bytes, out);
874 #[cfg(any(test, feature = "test-seams"))]
875 observe_selected_score(
876 table,
877 vector_fault_kernel_id_hamming_batch(),
878 d_bytes.saturating_mul(out.len()),
879 );
880 #[cfg(any(test, feature = "test-seams"))]
881 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::U32s(out.to_vec()));
882}
883
884#[must_use]
892pub fn dot_bit4(q: &[i8], codes: &[u8]) -> i32 {
893 let table = dispatch::active_table();
894 let result = (table.dot_bit4)(q, codes);
895 #[cfg(any(test, feature = "test-seams"))]
896 observe_selected_score(table, vector_fault_kernel_id_dot_bit4(), q.len());
897 #[cfg(any(test, feature = "test-seams"))]
898 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32(result));
899 result
900}
901
902pub(crate) fn dot_bit4_prepared(q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
903 let table = dispatch::active_table();
904 let result = (table.dot_bit4_prepared)(q, query_sum, codes);
905 #[cfg(any(test, feature = "test-seams"))]
906 observe_selected_score(table, vector_fault_kernel_id_dot_bit4_prepared(), q.len());
907 #[cfg(any(test, feature = "test-seams"))]
908 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32(result));
909 result
910}
911
912pub(crate) fn score_bit4_prepared_batch(
913 q: &[i8],
914 query_sum: i32,
915 query_scale_half: f64,
916 rows: &[u8],
917 d: usize,
918 factors: &[crate::quant::Bit4Factors],
919 out: &mut [f32],
920) {
921 let table = dispatch::active_table();
922 (table.score_bit4_prepared_batch)(q, query_sum, query_scale_half, rows, d, factors, out);
923 #[cfg(any(test, feature = "test-seams"))]
924 observe_selected_score(
925 table,
926 vector_fault_kernel_id_score_bit4_batch(),
927 d.saturating_mul(out.len()),
928 );
929 #[cfg(any(test, feature = "test-seams"))]
930 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32s(
931 out.iter().map(|value| value.to_bits()).collect(),
932 ));
933}
934
935#[cfg(any(test, feature = "test-seams"))]
936fn observe_selected_score(
937 table: &KernelTable,
938 kernel: vector_fault::KernelOperationId,
939 work_items: usize,
940) {
941 vector_fault::observe_selected_score(table.backend, kernel, work_items);
942}
943
944macro_rules! vector_fault_kernel_id {
945 ($name:ident, $variant:ident) => {
946 #[cfg(any(test, feature = "test-seams"))]
947 const fn $name() -> vector_fault::KernelOperationId {
948 vector_fault::KernelOperationId::$variant
949 }
950 };
951}
952
953vector_fault_kernel_id!(vector_fault_kernel_id_dot_i8, DotI8);
954vector_fault_kernel_id!(vector_fault_kernel_id_hamming, HammingU1);
955vector_fault_kernel_id!(vector_fault_kernel_id_dot_f32, DotF32);
956vector_fault_kernel_id!(vector_fault_kernel_id_dot_f16, DotF16);
957vector_fault_kernel_id!(vector_fault_kernel_id_dot_i8_batch, DotI8Batch);
958vector_fault_kernel_id!(vector_fault_kernel_id_hamming_batch, HammingU1Batch);
959vector_fault_kernel_id!(vector_fault_kernel_id_dot_bit4, DotBit4);
960vector_fault_kernel_id!(vector_fault_kernel_id_dot_bit4_prepared, DotBit4Prepared);
961vector_fault_kernel_id!(vector_fault_kernel_id_dot_bit4_batch, DotBit4Batch);
962vector_fault_kernel_id!(
963 vector_fault_kernel_id_score_bit4_batch,
964 ScoreBit4PreparedBatch
965);
966vector_fault_kernel_id!(vector_fault_kernel_id_score_bit4_ptrs, ScoreBit4Ptrs);
967
968pub fn score_bit4_ptrs(
977 q: &[i8],
978 query_sum: i32,
979 query_scale_half: f64,
980 rows: &Bit4Rows4<'_>,
981 factors: &[crate::quant::Bit4Factors; 4],
982 out: &mut [f32; 4],
983) -> Result<(), GatherShapeError> {
984 let table = dispatch::active_table();
985 let result =
986 score_bit4_ptrs_with_table(table, (q, query_sum, query_scale_half), rows, factors, out);
987 #[cfg(any(test, feature = "test-seams"))]
988 if result.is_ok() {
989 #[cfg(any(test, feature = "test-seams"))]
990 observe_selected_score(
991 table,
992 vector_fault_kernel_id_score_bit4_ptrs(),
993 q.len().saturating_mul(out.len()),
994 );
995 #[cfg(any(test, feature = "test-seams"))]
996 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::F32s(
997 out.iter().map(|value| value.to_bits()).collect(),
998 ));
999 }
1000 result
1001}
1002
1003pub fn dot_bit4_batch(q: &[i8], rows: &[u8], d: usize, out: &mut [i32]) {
1009 let table = dispatch::active_table();
1010 (table.dot_bit4_batch)(q, rows, d, out);
1011 #[cfg(any(test, feature = "test-seams"))]
1012 observe_selected_score(
1013 table,
1014 vector_fault_kernel_id_dot_bit4_batch(),
1015 d.saturating_mul(out.len()),
1016 );
1017 #[cfg(any(test, feature = "test-seams"))]
1018 vector_fault::observe_selected_result(vector_fault::KernelScoreValue::I32s(out.to_vec()));
1019}