Skip to main content

zeppelin_embed/kernels/
mod.rs

1//! Runtime-dispatched dot-product and packed-bit Hamming kernels.
2//!
3//! Callers validate dimensions once at the engine boundary. Every hot kernel
4//! therefore uses `debug_assert!` for the equal-length and row-shape contracts
5//! instead of repeating release-mode validation in the scan loop. Call
6//! [`crate::kernels::initialize`] during engine initialization to apply the optional
7//! `ZE_KERNEL=scalar|neon|avx2` override and surface unsupported requests as a
8//! typed [`crate::kernels::KernelInitError`]. If no explicit initialization occurs, the first
9//! kernel call safely caches the best runtime-detected arm.
10//!
11//! The scalar implementation is the behavioral oracle. [`crate::kernels::KernelVariant`]
12//! exposes every runtime-supported table to property tests, fuzzing, and the
13//! future Task 27 frontier harness without moving selection to build time.
14
15mod dispatch;
16/// Posting-block decode kernels: bit-unpack and delta prefix sum.
17pub 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
28/// Largest supported dimension for an i8 dot product.
29///
30/// The worst product is `(-128) * (-128) = 16_384`; at this dimension the
31/// worst sum is `2^30`, which remains representable by [`i32`].
32pub 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/// User-selectable runtime kernel families.
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum KernelArm {
40    /// Portable scalar oracle and fallback.
41    Scalar,
42    /// AArch64 NEON, using SDOT at runtime when available.
43    Neon,
44    /// x86-64 AVX2 plus scalar POPCNT.
45    Avx2,
46}
47
48/// Instruction tiers enumerated for the Task 27 B1 frontier campaign.
49#[derive(Clone, Copy, Debug, Eq, PartialEq)]
50pub enum InstructionTier {
51    /// Portable scalar instructions.
52    Scalar,
53    /// Baseline AArch64 widening multiply.
54    NeonWiden,
55    /// AArch64 FEAT_DotProd SDOT.
56    NeonDotprod,
57    /// x86-64 AVX2 plus POPCNT.
58    Avx2,
59    /// AArch64 FEAT_I8MM SMMLA batch tier.
60    NeonI8mmReserved,
61    /// Reserved, unimplemented FEAT_SME2 campaign tier.
62    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/// Stable identity of one concrete runtime kernel table.
88///
89/// This is test-support evidence rather than a user-selectable production
90/// arm. In particular, the experimental DotProd shapes are enumerable for
91/// parity campaigns but cannot be selected through `ZE_KERNEL`.
92#[cfg(any(test, feature = "benchmark-kernels"))]
93#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
94#[doc(hidden)]
95pub enum KernelBackendId {
96    /// Portable scalar table.
97    Scalar,
98    /// Baseline AArch64 widening table.
99    NeonWiden,
100    /// Shipped four-accumulator AArch64 DotProd table.
101    NeonDotprodU4,
102    /// AArch64 I8MM batch table.
103    NeonI8mm,
104    /// Two-accumulator AArch64 DotProd experiment.
105    NeonDotprodU2,
106    /// Six-accumulator AArch64 DotProd experiment.
107    NeonDotprodU6,
108    /// Eight-accumulator AArch64 DotProd experiment.
109    NeonDotprodU8,
110    /// Four-accumulator AArch64 DotProd table with batch prefetch.
111    NeonDotprodU4Prefetch,
112    /// x86-64 AVX2 table.
113    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    /// Returns the permanent canonical evidence label.
141    #[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/// Runtime CPU capabilities relevant to current and reserved kernel tiers.
173#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
174pub struct KernelFeatures {
175    /// AArch64 Advanced SIMD is executable.
176    pub neon: bool,
177    /// AArch64 FEAT_DotProd is executable.
178    pub dotprod: bool,
179    /// AArch64 half-precision conversion/arithmetic is executable.
180    pub fp16: bool,
181    /// AArch64 FEAT_I8MM is executable.
182    pub i8mm: bool,
183    /// AArch64 FEAT_SME2 was detected; implementation is deferred to B1.
184    pub sme2: bool,
185    /// x86-64 AVX2 is executable.
186    pub avx2: bool,
187    /// x86 POPCNT is executable.
188    pub popcnt: bool,
189}
190
191/// Typed failures produced while applying a runtime kernel override.
192#[derive(Clone, Debug, Eq, PartialEq)]
193pub enum KernelInitError {
194    /// `ZE_KERNEL` contained a Unicode value outside the supported set.
195    UnknownOverride {
196        /// Rejected environment value.
197        value: String,
198    },
199    /// `ZE_KERNEL` was not valid Unicode.
200    NonUnicodeOverride,
201    /// The requested family cannot execute on this CPU/architecture.
202    UnsupportedArm {
203        /// Requested runtime family.
204        requested: KernelArm,
205    },
206    /// A kernel call selected an arm before a conflicting explicit override.
207    AlreadyInitialized {
208        /// Arm already cached in the process.
209        selected: KernelArm,
210        /// Conflicting arm requested by `ZE_KERNEL`.
211        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/// Enumerable tuning declarations consumed directly by Task 27-H.
243///
244/// This is data only. Task 03 intentionally does not implement search,
245/// measurement ledgers, roofline calculation, or variant code generation.
246#[derive(Clone, Copy, Debug, Eq, PartialEq)]
247pub struct KnobSpace {
248    /// Candidate loop-unroll depths.
249    pub unroll: &'static [usize],
250    /// Candidate independent-accumulator counts.
251    pub accumulators: &'static [usize],
252    /// Candidate rows grouped per scan block.
253    pub rows_per_block: &'static [usize],
254    /// Candidate software-prefetch distances in rows; zero disables prefetch.
255    pub prefetch_dist: &'static [usize],
256    /// Current and reserved instruction tiers.
257    pub tier: &'static [InstructionTier],
258}
259
260/// Baseline micro-shape selected before the Task 27 B1 campaign.
261#[derive(Clone, Copy, Debug, Eq, PartialEq)]
262pub struct BaselineKernelConfig {
263    /// Loop-unroll depth.
264    pub unroll: usize,
265    /// Number of independent accumulators.
266    pub accumulators: usize,
267    /// Rows handled per row-major batch iteration.
268    pub rows_per_block: usize,
269    /// Software-prefetch distance in rows.
270    pub prefetch_dist: usize,
271}
272
273/// Shipped baseline configuration; B1 may replace every value with evidence.
274pub const BASELINE_KERNEL_CONFIG: BaselineKernelConfig = BaselineKernelConfig {
275    // Retained by tasks/evidence/opt-ledger/B1.md iterations 1–5.
276    unroll: 4,
277    accumulators: 4,
278    rows_per_block: 1,
279    prefetch_dist: 0,
280};
281
282// Search provenance: tasks/evidence/opt-ledger/B1.md.
283const 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]; // Search provenance: tasks/evidence/opt-ledger/B1.md.
295
296/// Registered Task 27-H kernel tuning space.
297pub 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/// One bounds-validated packed four-bit row borrowed from a mapped region.
345///
346/// `from_mapped_region` validates an in-bounds offset and length; the
347/// crate-only `from_validated_bytes` wraps a slice whose persisted range was
348/// already validated. Both retain a borrowed slice, so the stored address is
349/// readable for that slice's complete length while this handle is alive.
350#[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    /// Borrows one row at `offset` when the complete row lies in `region`.
366    ///
367    /// Returns `None` when `offset + row_bytes` overflows or exceeds the
368    /// mapped region. No alignment is required by any installed kernel.
369    #[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    /// Wraps bytes whose complete persisted row range was already validated.
376    #[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/// Four equally-sized packed rows whose addresses were validated before scoring.
391#[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    /// Validates four mapped-row handles once for the unchecked gather loop.
409    ///
410    /// # Errors
411    ///
412    /// Returns [`GatherShapeError::RowLength`] when any handle was constructed
413    /// for a different readable width. No alignment is required.
414    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/// Rejected shape at the safe four-row gather boundary.
433#[derive(Clone, Copy, Debug, Eq, PartialEq)]
434pub enum GatherShapeError {
435    /// The prepared query exceeds the exact i8 accumulator bound.
436    DimensionTooLarge {
437        /// Supplied coordinate count.
438        actual: usize,
439        /// Largest supported coordinate count.
440        maximum: usize,
441    },
442    /// `row_bytes` does not match the packed query dimension.
443    RowBytes {
444        /// Required packed bytes per row.
445        expected: usize,
446        /// Supplied packed bytes per row.
447        actual: usize,
448    },
449    /// A validated row handle was constructed for a different row width.
450    RowLength {
451        /// Zero-based pointer slot.
452        index: usize,
453        /// Required readable bytes.
454        expected: usize,
455        /// Bytes proven readable by the row handle.
456        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
485/// Prefetches caller-selected packed rows into L1 for a later gather group.
486///
487/// Traversal can call this only for degrees where M0 found a benefit (R >= 32),
488/// leaving small-degree scoring unchanged. Every validated row handle supplies
489/// a live readable address. This is a no-op on non-AArch64 targets.
490pub 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        // SAFETY: Bit4Row construction proves address belongs to a live readable
495        // region. PRFM is a non-mutating hint and has no alignment requirement.
496        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
508/// Prefetches one already-validated four-row gather group.
509pub 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            // SAFETY: Bit4Rows4 construction proves every address belongs to a
514            // live region readable for row_bytes. PRFM is a non-mutating hint.
515            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/// One concrete runtime dispatch table.
529///
530/// Variants are exposed so the scalar oracle, property suite, fuzzer, and
531/// future frontier harness can evaluate every arm independently of the global
532/// once-selected default.
533#[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    /// Returns the scalar oracle variant.
550    #[must_use]
551    pub fn scalar() -> Self {
552        Self {
553            table: scalar::table(),
554        }
555    }
556
557    /// Returns the once-selected runtime table for a repeated hot loop.
558    #[doc(hidden)]
559    #[must_use]
560    pub fn selected() -> Self {
561        Self {
562            table: *dispatch::active_table(),
563        }
564    }
565
566    /// Iterates over every arm executable on the current CPU.
567    ///
568    /// BL-072: `dispatch::variant_tables()` returns the SMMLA table plus the
569    /// u2/u6/u8/prefetch DotProd shape tables, none of which the engine ever
570    /// selects on its own. They exist so the scalar oracle, the property
571    /// suite, the fuzzer and the frontier harness can evaluate each arm
572    /// independently. That is a test and measurement surface, not an engine
573    /// surface, so it is gated: a consumer of the shipped staticlib does not
574    /// see it.
575    #[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    /// Returns this variant's runtime family.
584    #[must_use]
585    pub fn arm(self) -> KernelArm {
586        self.table.arm
587    }
588
589    /// Returns this variant's precise instruction tier.
590    #[must_use]
591    pub fn tier(self) -> InstructionTier {
592        self.table.tier
593    }
594
595    /// Returns this concrete table's stable evidence identity.
596    #[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    /// Computes an i8 dot product through this concrete table.
604    #[must_use]
605    pub fn dot_i8(self, a: &[i8], b: &[i8]) -> i32 {
606        (self.table.dot_i8)(a, b)
607    }
608
609    /// Computes an f32 dot product through this concrete table.
610    #[must_use]
611    pub fn dot_f32(self, a: &[f32], b: &[f32]) -> f32 {
612        (self.table.dot_f32)(a, b)
613    }
614
615    /// Computes an IEEE-f16-bits dot product through this concrete table.
616    #[must_use]
617    pub fn dot_f16(self, a: &[u16], b: &[u16]) -> f32 {
618        (self.table.dot_f16)(a, b)
619    }
620
621    /// Computes packed-byte Hamming distance through this concrete table.
622    #[must_use]
623    pub fn hamming_u1(self, a: &[u8], b: &[u8]) -> u32 {
624        (self.table.hamming_u1)(a, b)
625    }
626
627    /// Scores contiguous i8 rows through this concrete table.
628    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    /// Scores contiguous packed-bit rows through this concrete table.
633    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    /// Scores a signed-byte query against one packed four-bit row.
638    #[must_use]
639    pub fn dot_bit4(self, q: &[i8], codes: &[u8]) -> i32 {
640        (self.table.dot_bit4)(q, codes)
641    }
642
643    /// Scores a block-interleaved four-bit query against one packed row.
644    ///
645    /// Within each 32-coordinate query block, even coordinates precede odd
646    /// coordinates. `query_sum` is the sum of the original signed query codes.
647    #[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    /// Scores a signed-byte query against contiguous packed four-bit rows.
654    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    /// Scores a prepared four-bit query and row factors through this table.
659    #[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    /// Scores four arbitrary validated packed four-bit rows through this table.
681    ///
682    /// # Safety
683    ///
684    /// [`Bit4Rows4::from_rows`] proves that every row address is readable for a
685    /// common width and that the mapped regions outlive this call. This method
686    /// checks that the width equals `q.len().div_ceil(2)` and that the supported
687    /// dimension bound holds. No row alignment is required.
688    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    // SAFETY: Bit4Rows4::from_rows verified that every Bit4Row's readable
723    // slice is row_bytes long, and their borrow keeps each backing region alive
724    // for the dispatched call.
725    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
739/// Applies runtime feature detection and the optional `ZE_KERNEL` override.
740///
741/// Call this once at engine initialization, before any kernel call. Unknown or
742/// unsupported forced arms return a typed error and are never executed.
743pub fn initialize() -> Result<KernelArm, KernelInitError> {
744    dispatch::initialize()
745}
746
747/// Returns the cached arm, selecting the best safe runtime arm if necessary.
748#[must_use]
749pub fn selected_arm() -> KernelArm {
750    dispatch::active_table().arm
751}
752
753/// Reports whether a user-selectable arm is executable on this CPU.
754#[must_use]
755pub fn is_arm_supported(arm: KernelArm) -> bool {
756    dispatch::table_for_arm(arm).is_some()
757}
758
759/// Returns cached runtime feature detection, including reserved B1 features.
760#[must_use]
761pub fn detected_features() -> KernelFeatures {
762    dispatch::features()
763}
764
765/// Streams bytes through four independent wide-load accumulators for the
766/// platform-truth bandwidth harness.
767///
768/// This measurement-only hook returns `None` when an AArch64 NEON
769/// implementation is unavailable. It is not part of the distance-kernel API.
770#[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/// Computes the dot product of equal-length signed-byte vectors.
787///
788/// Callers must pre-validate equal lengths and `a.len() <=
789/// MAX_DOT_I8_DIMENSION`; the hot path performs `debug_assert!` checks only.
790#[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/// Computes the dot product of equal-length f32 vectors.
802///
803/// Callers must pre-validate equal lengths; the hot path performs a
804/// `debug_assert!` check only. SIMD accumulation may reassociate additions.
805#[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/// Computes an f32 dot product from equal-length IEEE f16 bit vectors.
817///
818/// Callers must pre-validate equal lengths; the hot path performs a
819/// `debug_assert!` check only. Subnormals, infinities, and NaNs follow IEEE
820/// conversion and f32 arithmetic behavior.
821#[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/// Computes Hamming distance over equal-length packed-byte vectors.
833///
834/// Callers must pre-validate equal lengths; the hot path performs a
835/// `debug_assert!` check only. Every bit of every supplied byte is significant,
836/// including unused high bits in a caller's final partial logical byte.
837#[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
848/// Scores one i8 query against contiguous row-major vectors.
849///
850/// Callers must pre-validate `q.len() == d`, `rows.len() == d * out.len()`,
851/// and `d <= MAX_DOT_I8_DIMENSION`; the hot path uses `debug_assert!` only.
852/// A zero dimension deterministically fills `out` with zero.
853pub 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
866/// Scores one packed-bit query against contiguous row-major byte vectors.
867///
868/// Callers must pre-validate `q.len() == d_bytes` and `rows.len() ==
869/// d_bytes * out.len()`; the hot path uses `debug_assert!` only. Every bit in
870/// the final byte is significant. A zero byte dimension fills `out` with zero.
871pub 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/// Scores a signed-byte query against one packed four-bit row.
885///
886/// Two unsigned fields are stored most-significant first in each byte and
887/// decoded to the doubled quantization grid `2 * code - 15`. Callers must
888/// pre-validate `codes.len() == q.len().div_ceil(2)` and `q.len() <=
889/// MAX_DOT_I8_DIMENSION`; an unused trailing field is ignored. The hot path
890/// allocates no memory.
891#[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
968/// Scores four arbitrary validated packed four-bit rows through runtime dispatch.
969///
970/// # Safety
971///
972/// [`Bit4Rows4::from_rows`] proves that every row pointer is readable for one
973/// common width and that each backing region outlives the call. This function
974/// checks that the width equals `q.len().div_ceil(2)` and that the supported
975/// dimension bound holds. No alignment is required.
976pub 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
1003/// Scores one signed-byte query against contiguous packed four-bit rows.
1004///
1005/// Callers must pre-validate `q.len() == d` and `rows.len() ==
1006/// d.div_ceil(2) * out.len()` and `d <= MAX_DOT_I8_DIMENSION`. A zero dimension
1007/// fills `out` with zero. No scratch storage is allocated.
1008pub 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}