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    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/// Runtime CPU capabilities relevant to current and reserved kernel tiers.
172#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
173pub struct KernelFeatures {
174    /// AArch64 Advanced SIMD is executable.
175    pub neon: bool,
176    /// AArch64 FEAT_DotProd is executable.
177    pub dotprod: bool,
178    /// AArch64 half-precision conversion/arithmetic is executable.
179    pub fp16: bool,
180    /// AArch64 FEAT_I8MM is executable.
181    pub i8mm: bool,
182    /// AArch64 FEAT_SME2 was detected; implementation is deferred to B1.
183    pub sme2: bool,
184    /// x86-64 AVX2 is executable.
185    pub avx2: bool,
186    /// x86 POPCNT is executable.
187    pub popcnt: bool,
188}
189
190/// Typed failures produced while applying a runtime kernel override.
191#[derive(Clone, Debug, Eq, PartialEq)]
192pub enum KernelInitError {
193    /// `ZE_KERNEL` contained a Unicode value outside the supported set.
194    UnknownOverride {
195        /// Rejected environment value.
196        value: String,
197    },
198    /// `ZE_KERNEL` was not valid Unicode.
199    NonUnicodeOverride,
200    /// The requested family cannot execute on this CPU/architecture.
201    UnsupportedArm {
202        /// Requested runtime family.
203        requested: KernelArm,
204    },
205    /// A kernel call selected an arm before a conflicting explicit override.
206    AlreadyInitialized {
207        /// Arm already cached in the process.
208        selected: KernelArm,
209        /// Conflicting arm requested by `ZE_KERNEL`.
210        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/// Enumerable tuning declarations consumed directly by Task 27-H.
242///
243/// This is data only. Task 03 intentionally does not implement search,
244/// measurement ledgers, roofline calculation, or variant code generation.
245#[derive(Clone, Copy, Debug, Eq, PartialEq)]
246pub struct KnobSpace {
247    /// Candidate loop-unroll depths.
248    pub unroll: &'static [usize],
249    /// Candidate independent-accumulator counts.
250    pub accumulators: &'static [usize],
251    /// Candidate rows grouped per scan block.
252    pub rows_per_block: &'static [usize],
253    /// Candidate software-prefetch distances in rows; zero disables prefetch.
254    pub prefetch_dist: &'static [usize],
255    /// Current and reserved instruction tiers.
256    pub tier: &'static [InstructionTier],
257}
258
259/// Baseline micro-shape selected before the Task 27 B1 campaign.
260#[derive(Clone, Copy, Debug, Eq, PartialEq)]
261pub struct BaselineKernelConfig {
262    /// Loop-unroll depth.
263    pub unroll: usize,
264    /// Number of independent accumulators.
265    pub accumulators: usize,
266    /// Rows handled per row-major batch iteration.
267    pub rows_per_block: usize,
268    /// Software-prefetch distance in rows.
269    pub prefetch_dist: usize,
270}
271
272/// Shipped baseline configuration; B1 may replace every value with evidence.
273pub const BASELINE_KERNEL_CONFIG: BaselineKernelConfig = BaselineKernelConfig {
274    // Retained by tasks/evidence/opt-ledger/B1.md iterations 1–5.
275    unroll: 4,
276    accumulators: 4,
277    rows_per_block: 1,
278    prefetch_dist: 0,
279};
280
281// Search provenance: tasks/evidence/opt-ledger/B1.md.
282const 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]; // Search provenance: tasks/evidence/opt-ledger/B1.md.
294
295/// Registered Task 27-H kernel tuning space.
296pub 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/// One bounds-validated packed four-bit row borrowed from a mapped region.
344///
345/// `from_mapped_region` validates an in-bounds offset and length; the
346/// crate-only `from_validated_bytes` wraps a slice whose persisted range was
347/// already validated. Both retain a borrowed slice, so the stored address is
348/// readable for that slice's complete length while this handle is alive.
349#[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    /// Borrows one row at `offset` when the complete row lies in `region`.
365    ///
366    /// Returns `None` when `offset + row_bytes` overflows or exceeds the
367    /// mapped region. No alignment is required by any installed kernel.
368    #[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    /// Wraps bytes whose complete persisted row range was already validated.
375    #[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/// Four equally-sized packed rows whose addresses were validated before scoring.
390#[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    /// Validates four mapped-row handles once for the unchecked gather loop.
408    ///
409    /// # Errors
410    ///
411    /// Returns [`GatherShapeError::RowLength`] when any handle was constructed
412    /// for a different readable width. No alignment is required.
413    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/// Rejected shape at the safe four-row gather boundary.
432#[derive(Clone, Copy, Debug, Eq, PartialEq)]
433pub enum GatherShapeError {
434    /// The prepared query exceeds the exact i8 accumulator bound.
435    DimensionTooLarge {
436        /// Supplied coordinate count.
437        actual: usize,
438        /// Largest supported coordinate count.
439        maximum: usize,
440    },
441    /// `row_bytes` does not match the packed query dimension.
442    RowBytes {
443        /// Required packed bytes per row.
444        expected: usize,
445        /// Supplied packed bytes per row.
446        actual: usize,
447    },
448    /// A validated row handle was constructed for a different row width.
449    RowLength {
450        /// Zero-based pointer slot.
451        index: usize,
452        /// Required readable bytes.
453        expected: usize,
454        /// Bytes proven readable by the row handle.
455        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
484/// Prefetches caller-selected packed rows into L1 for a later gather group.
485///
486/// Traversal can call this only for degrees where M0 found a benefit (R >= 32),
487/// leaving small-degree scoring unchanged. Every validated row handle supplies
488/// a live readable address. This is a no-op on non-AArch64 targets.
489pub 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        // SAFETY: Bit4Row construction proves address belongs to a live readable
494        // region. PRFM is a non-mutating hint and has no alignment requirement.
495        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
507/// Prefetches one already-validated four-row gather group.
508pub 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            // SAFETY: Bit4Rows4 construction proves every address belongs to a
513            // live region readable for row_bytes. PRFM is a non-mutating hint.
514            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/// One concrete runtime dispatch table.
528///
529/// Variants are exposed so the scalar oracle, property suite, fuzzer, and
530/// future frontier harness can evaluate every arm independently of the global
531/// once-selected default.
532#[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    /// Returns the scalar oracle variant.
549    #[must_use]
550    pub fn scalar() -> Self {
551        Self {
552            table: scalar::table(),
553        }
554    }
555
556    /// Returns the once-selected runtime table for a repeated hot loop.
557    #[doc(hidden)]
558    #[must_use]
559    pub fn selected() -> Self {
560        Self {
561            table: *dispatch::active_table(),
562        }
563    }
564
565    /// Iterates over every arm executable on the current CPU.
566    ///
567    /// BL-072: `dispatch::variant_tables()` returns the SMMLA table plus the
568    /// u2/u6/u8/prefetch DotProd shape tables, none of which the engine ever
569    /// selects on its own. They exist so the scalar oracle, the property
570    /// suite, the fuzzer and the frontier harness can evaluate each arm
571    /// independently. That is a test and measurement surface, not an engine
572    /// surface, so it is gated: a consumer of the shipped staticlib does not
573    /// see it.
574    #[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    /// Returns this variant's runtime family.
583    #[must_use]
584    pub fn arm(self) -> KernelArm {
585        self.table.arm
586    }
587
588    /// Returns this variant's precise instruction tier.
589    #[must_use]
590    pub fn tier(self) -> InstructionTier {
591        self.table.tier
592    }
593
594    /// Returns this concrete table's stable evidence identity.
595    #[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    /// Computes an i8 dot product through this concrete table.
603    #[must_use]
604    pub fn dot_i8(self, a: &[i8], b: &[i8]) -> i32 {
605        (self.table.dot_i8)(a, b)
606    }
607
608    /// Computes an f32 dot product through this concrete table.
609    #[must_use]
610    pub fn dot_f32(self, a: &[f32], b: &[f32]) -> f32 {
611        (self.table.dot_f32)(a, b)
612    }
613
614    /// Computes an IEEE-f16-bits dot product through this concrete table.
615    #[must_use]
616    pub fn dot_f16(self, a: &[u16], b: &[u16]) -> f32 {
617        (self.table.dot_f16)(a, b)
618    }
619
620    /// Computes packed-byte Hamming distance through this concrete table.
621    #[must_use]
622    pub fn hamming_u1(self, a: &[u8], b: &[u8]) -> u32 {
623        (self.table.hamming_u1)(a, b)
624    }
625
626    /// Scores contiguous i8 rows through this concrete table.
627    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    /// Scores contiguous packed-bit rows through this concrete table.
632    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    /// Scores a signed-byte query against one packed four-bit row.
637    #[must_use]
638    pub fn dot_bit4(self, q: &[i8], codes: &[u8]) -> i32 {
639        (self.table.dot_bit4)(q, codes)
640    }
641
642    /// Scores a block-interleaved four-bit query against one packed row.
643    ///
644    /// Within each 32-coordinate query block, even coordinates precede odd
645    /// coordinates. `query_sum` is the sum of the original signed query codes.
646    #[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    /// Scores a signed-byte query against contiguous packed four-bit rows.
653    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    /// Scores a prepared four-bit query and row factors through this table.
658    #[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    /// Scores four arbitrary validated packed four-bit rows through this table.
680    ///
681    /// # Safety
682    ///
683    /// [`Bit4Rows4::from_rows`] proves that every row address is readable for a
684    /// common width and that the mapped regions outlive this call. This method
685    /// checks that the width equals `q.len().div_ceil(2)` and that the supported
686    /// dimension bound holds. No row alignment is required.
687    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    // SAFETY: Bit4Rows4::from_rows verified that every Bit4Row's readable
722    // slice is row_bytes long, and their borrow keeps each backing region alive
723    // for the dispatched call.
724    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
738/// Applies runtime feature detection and the optional `ZE_KERNEL` override.
739///
740/// Call this once at engine initialization, before any kernel call. Unknown or
741/// unsupported forced arms return a typed error and are never executed.
742pub fn initialize() -> Result<KernelArm, KernelInitError> {
743    dispatch::initialize()
744}
745
746/// Returns the cached arm, selecting the best safe runtime arm if necessary.
747#[must_use]
748pub fn selected_arm() -> KernelArm {
749    dispatch::active_table().arm
750}
751
752/// Reports whether a user-selectable arm is executable on this CPU.
753#[must_use]
754pub fn is_arm_supported(arm: KernelArm) -> bool {
755    dispatch::table_for_arm(arm).is_some()
756}
757
758/// Returns cached runtime feature detection, including reserved B1 features.
759#[must_use]
760pub fn detected_features() -> KernelFeatures {
761    dispatch::features()
762}
763
764/// Streams bytes through four independent wide-load accumulators for the
765/// platform-truth bandwidth harness.
766///
767/// This measurement-only hook returns `None` when an AArch64 NEON
768/// implementation is unavailable. It is not part of the distance-kernel API.
769#[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/// Computes the dot product of equal-length signed-byte vectors.
786///
787/// Callers must pre-validate equal lengths and `a.len() <=
788/// MAX_DOT_I8_DIMENSION`; the hot path performs `debug_assert!` checks only.
789#[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/// Computes the dot product of equal-length f32 vectors.
801///
802/// Callers must pre-validate equal lengths; the hot path performs a
803/// `debug_assert!` check only. SIMD accumulation may reassociate additions.
804#[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/// Computes an f32 dot product from equal-length IEEE f16 bit vectors.
816///
817/// Callers must pre-validate equal lengths; the hot path performs a
818/// `debug_assert!` check only. Subnormals, infinities, and NaNs follow IEEE
819/// conversion and f32 arithmetic behavior.
820#[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/// Computes Hamming distance over equal-length packed-byte vectors.
832///
833/// Callers must pre-validate equal lengths; the hot path performs a
834/// `debug_assert!` check only. Every bit of every supplied byte is significant,
835/// including unused high bits in a caller's final partial logical byte.
836#[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
847/// Scores one i8 query against contiguous row-major vectors.
848///
849/// Callers must pre-validate `q.len() == d`, `rows.len() == d * out.len()`,
850/// and `d <= MAX_DOT_I8_DIMENSION`; the hot path uses `debug_assert!` only.
851/// A zero dimension deterministically fills `out` with zero.
852pub 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
865/// Scores one packed-bit query against contiguous row-major byte vectors.
866///
867/// Callers must pre-validate `q.len() == d_bytes` and `rows.len() ==
868/// d_bytes * out.len()`; the hot path uses `debug_assert!` only. Every bit in
869/// the final byte is significant. A zero byte dimension fills `out` with zero.
870pub 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/// Scores a signed-byte query against one packed four-bit row.
884///
885/// Two unsigned fields are stored most-significant first in each byte and
886/// decoded to the doubled quantization grid `2 * code - 15`. Callers must
887/// pre-validate `codes.len() == q.len().div_ceil(2)` and `q.len() <=
888/// MAX_DOT_I8_DIMENSION`; an unused trailing field is ignored. The hot path
889/// allocates no memory.
890#[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
967/// Scores four arbitrary validated packed four-bit rows through runtime dispatch.
968///
969/// # Safety
970///
971/// [`Bit4Rows4::from_rows`] proves that every row pointer is readable for one
972/// common width and that each backing region outlives the call. This function
973/// checks that the width equals `q.len().div_ceil(2)` and that the supported
974/// dimension bound holds. No alignment is required.
975pub 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
1002/// Scores one signed-byte query against contiguous packed four-bit rows.
1003///
1004/// Callers must pre-validate `q.len() == d` and `rows.len() ==
1005/// d.div_ceil(2) * out.len()` and `d <= MAX_DOT_I8_DIMENSION`. A zero dimension
1006/// fills `out` with zero. No scratch storage is allocated.
1007pub 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}