Skip to main content

KernelSimd

Trait KernelSimd 

Source
pub trait KernelSimd<L: Scalar, R: Scalar, A: Scalar, O: Scalar>: SimdOps<A> {
    const REQUANT_VECTOR: bool = false;

    // Required methods
    unsafe fn load_lhs(self, p: *const L) -> <Self as SimdOps<A>>::Reg;
    unsafe fn splat_rhs(self, v: R) -> <Self as SimdOps<A>>::Reg;
    unsafe fn load_out(self, p: *const O) -> <Self as SimdOps<A>>::Reg;
    unsafe fn store_out(self, p: *mut O, v: <Self as SimdOps<A>>::Reg);

    // Provided methods
    unsafe fn dot_accumulate<const MR_REG: usize, const NR: usize>(
        self,
        _kc: usize,
        _a: *const L,
        _b: *const R,
        _acc: &mut [[<Self as SimdOps<A>>::Reg; MR_REG]; NR],
    ) { ... }
    unsafe fn requant_store(
        self,
        _dst: *mut i8,
        _v: <Self as SimdOps<A>>::Reg,
        _scale: f64,
        _zp: i32,
        _lo: i32,
        _hi: i32,
    ) { ... }
}
Expand description

The capability an ISA token needs when a crate::kernel::KernelFamily has input types L/R, an accumulator A, and an output O that differ

The token accumulates in A (the SimdOps<A> supertrait) and moves family inputs and outputs into and out of A-typed registers. It widens on load and narrows on store wherever the element type is narrower than A

This is the seam that lets mixed precision (A != L) work without a per-type branch in the driver. The homogeneous case (L = R = A = O) is covered once by the blanket impl below. It forwards to plain SimdOps load, splat, and store. A narrow family (f16 or bf16 inputs, f32 accumulator) instead gets an ISA impl whose load_* widens and whose store_out narrows. The all-equal blanket and any mixed impl (L != A) can never overlap, because a mixed impl’s types are concrete and unequal

Provided Associated Constants§

Source

const REQUANT_VECTOR: bool = false

true only when Self::requant_store is a genuine vector implementation, rather than the default unreachable! stub. The requantizing epilogue’s vector store path is gated on this. A false token routes every element through the scalar map (KRequantize::apply) instead

Required Methods§

Source

unsafe fn load_lhs(self, p: *const L) -> <Self as SimdOps<A>>::Reg

Load LANES LHS values, widened to one A register (a plain load when L == A)

§Safety

p must be valid for LANES reads. Run this inside this token’s Simd::vectorize

Source

unsafe fn splat_rhs(self, v: R) -> <Self as SimdOps<A>>::Reg

Widen one RHS scalar and broadcast it to all A lanes (a plain splat when R == A)

§Safety

See the trait-level note

Source

unsafe fn load_out(self, p: *const O) -> <Self as SimdOps<A>>::Reg

Load LANES output values, widened to one A register, for the beta != 0 read of C (a plain load when O == A)

§Safety

p must be valid for LANES reads. Run this inside Simd::vectorize

Source

unsafe fn store_out(self, p: *mut O, v: <Self as SimdOps<A>>::Reg)

Narrow one A register to LANES output values and store them (a plain store when O == A). This rounds to nearest-even when it actually narrows

§Safety

p must be valid for LANES writes. Run this inside Simd::vectorize

Provided Methods§

Source

unsafe fn dot_accumulate<const MR_REG: usize, const NR: usize>( self, _kc: usize, _a: *const L, _b: *const R, _acc: &mut [[<Self as SimdOps<A>>::Reg; MR_REG]; NR], )

Accumulate one full MR_REG x NR microtile from dot-product-packed panels into the register-resident acc (pre-zeroed by the caller). This is the seam a dot-kernel family (crate::kernel::KernelFamily::DEPTH_MULTIPLE > 1) drives on instead of SimdOps::accumulate_tile. It folds DEPTH_MULTIPLE consecutive depth steps into one hardware instruction (vpdpbusd, vdpbf16ps), which reshapes the accumulation rounding in a way accumulate_tile’s contract forbids. a and b are the family’s interleaved panels, laid out by contract between the family’s packers and the overriding token. kc is the real, unpadded depth. The token reads ceil(kc / DEPTH_MULTIPLE) instruction-groups from the depth-padded panel. Any signedness or bias correction (VNNI’s +128) applies internally, so acc holds the true sum_k(A*B) on return

The default is unreachable. Only a dot-capable token (e.g. Avx512Vnni, Avx512Bf16) overrides it, and only a dot family ever calls it

§Safety

a and b must be valid for the family’s packed panel at this (MR_REG, NR, kc), and acc must be pre-initialized. Run this inside this token’s Simd::vectorize context

Source

unsafe fn requant_store( self, _dst: *mut i8, _v: <Self as SimdOps<A>>::Reg, _scale: f64, _zp: i32, _lo: i32, _hi: i32, )

Vectorized i32 -> i8 requantize store

This clamps each A-accumulator lane through the exact requant map and writes its low byte to LANES consecutive slots at dst. dst is a raw byte pointer regardless of O. A u8 output casts its pointer to call this. The result is bit-identical, because the low byte of a value clamped into [lo, hi] reads the same whether the type is i8 or u8. Only requant-vector-capable tokens (Self::REQUANT_VECTOR = true) override this. The default is unreachable!, the same seam pattern as Self::dot_accumulate

§Contract: bit-for-bit agreement with the scalar map

Each lane of an implementation does the following, in order:

  1. Widen i32 -> f64 (exact)
  2. Multiply by scale widened f32 -> f64 (exact widening, one IEEE multiply)
  3. Round to nearest-even in hardware. This agrees with the scalar round_ne_f64. That function’s 2^52 trick is roundTiesToEven below 2^52. Above 2^52 every f64 is already integral, so hardware rounding is the identity there too
  4. Add zp as f64
  5. Clamp with max(lo as f64) then min(hi as f64)
  6. Convert f64 -> i32, exact because the value is now integral and inside [lo, hi]
  7. Store the low byte by truncation, never a saturating pack

That sequence equals the scalar clamp(zp + round_ne(scale*v), lo, hi) case by case:

  • |t| < 2^52: t is integral and exact. The scalar t as i64 is exact, and its zp add cannot saturate. The f64 value t + zp is exact too, because both stay far below 2^53. The 2 paths feed identical values into an identical clamp
  • t >= 2^52: both clamp to hi (scalar through a saturating i64 + zp then clamp, vector through f64 + zp then min(hi)). By symmetry t <= -2^52 clamps both to lo
  • NaN cannot occur: the API validates that scale is finite and positive, and that v is a finite i32

The caller supplies v already bias-added. SIMD i32 add (paddd) wraps, matching the scalar wrapping_add it must agree with. lo and hi are parameters: -128/127 for the i8 output. The u8 output phase reuses the same machinery with (0, 255)

§Safety

dst must be valid for LANES byte writes. Run this inside this token’s Simd::vectorize

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§

Source§

impl KernelSimd<bf16, bf16, f32, bf16> for Avx512Bf16

Available on crate feature half and (x86 or x86-64) only.

bf16 -> f32 via vdpbf16ps. The widen-load/narrow-store seam delegates to Avx512F’s bf16 impl. That keeps one source of truth for the round-to-nearest-even bias plus the half-NaN narrowing, which must stay bit-identical to half::bf16::from_f32 and the scalar edge path. vectorize here enables a superset of avx512f, so the delegated conversions still land in a valid codegen context. splat_rhs is trait-required but unused, because the hot loop runs Self::dot_accumulate. load_out is used by the beta != 0 read of C

Source§

impl KernelSimd<bf16, bf16, f32, bf16> for Avx512F

Available on crate feature half and (x86 or x86-64) only.

bf16 via plain integer ops, all AVX-512F: widening is a 16-bit left shift into an f32. Narrowing uses the round-to-nearest-even bias trick followed by a truncate. The narrowing side is bit-identical to half::bf16::from_f32, NaN included (forced to (bits>>16) | 0x0040), so this conversion matches the scalar path. That keeps full and edge tiles of the same matrix consistent. Even so, the vdpbf16ps dot kernel’s fused 2-term MAC rounds differently from this widen-and-FMA path

Source§

impl KernelSimd<bf16, bf16, f32, bf16> for Fma

Available on crate feature half and (x86 or x86-64) only.

bf16 via plain integer ops (no dedicated bf16 hardware needed): widening is a 16-bit left shift into the top of an f32. Narrowing uses the round-to-nearest-even bias trick (add ((bits>>16)&1) + 0x7FFF, then shift right 16). The narrowing side is bit-identical to half::bf16::from_f32, NaN included (forced to (bits>>16) | 0x0040), so this conversion matches the scalar path exactly. That keeps full and edge tiles of the same matrix consistent. Even so, the vdpbf16ps dot kernel’s fused 2-term MAC rounds differently from this widen-and-FMA path

Source§

impl KernelSimd<bf16, bf16, f32, bf16> for ScalarTok

Available on crate feature half only.
Source§

impl KernelSimd<f16, f16, f32, f16> for Avx512F

Available on crate feature half and (x86 or x86-64) only.

f16 via AVX-512F’s vcvtph2ps/vcvtps2ph: round-to-nearest-even on store, matching half::f16::from_f32

Source§

impl KernelSimd<f16, f16, f32, f16> for Fma

Available on crate feature half and (x86 or x86-64) only.

f16 via F16C: widen all 8 lanes with vcvtph2ps, narrow with vcvtps2ph (round-to-nearest-even, matching half::f16::from_f32)

Source§

impl KernelSimd<f16, f16, f32, f16> for ScalarTok

Available on crate feature half only.
Source§

impl KernelSimd<i8, i8, i32, i32> for Avx512F

Available on crate feature int8 and (x86 or x86-64) only.

i8 -> i32 widen kernel: sign-extend 16 LHS bytes on load, broadcast a sign-extended RHS byte. Out == Acc == i32 here, so load_out/store_out are plain load/store

Source§

impl KernelSimd<i8, i8, i32, i32> for Avx512Vnni

Available on crate feature int8 and (x86 or x86-64) only.

i8 -> i32 via VNNI. The load/store seam matches Avx512F’s plain i32 epilogue. load_lhs and splat_rhs are required by the trait but unused, because the hot loop runs through Self::dot_accumulate, which reads the family’s k-quad-interleaved panels directly

Source§

impl KernelSimd<i8, i8, i32, i32> for Fma

Available on crate feature int8 and (x86 or x86-64) only.

i8 -> i32 widen kernel: sign-extend 8 LHS bytes on load, broadcast a sign-extended RHS byte. Out == Acc == i32 here, so load_out/store_out are plain load/store

Source§

impl KernelSimd<i8, i8, i32, i32> for ScalarTok

Available on crate feature int8 only.
Source§

impl<A: Scalar, S: SimdOps<A>> KernelSimd<A, A, A, A> for S

Homogeneous blanket: when every family type equals the accumulator type, there is nothing to widen or narrow. So load_lhs, splat_rhs, load_out, and store_out are plain SimdOps load, splat, and store. Any homogeneous family (e.g. FloatGemm<f32> or FloatGemm<f64>) needs zero per-ISA code to satisfy KernelSimd

Source§

impl<S: KernelSimd<bf16, bf16, f32, bf16>> KernelSimd<bf16, bf16, f32, f32> for S

Available on crate feature half only.

The bf16 head of the f32-output twin above: its dot_accumulate forward carries a token’s vdpbf16ps override through to crate::kernel::Bf16DotGemmF32

Source§

impl<S: KernelSimd<f16, f16, f32, f16>> KernelSimd<f16, f16, f32, f32> for S

Available on crate feature half only.

f32-output twin of the narrow mixed seam: f16 or bf16 inputs, f32 accumulator, and an f32 output (Out == Acc)

This is the seam the deep-contraction narrow twins (crate::kernel::MixedGemmF32 and crate::kernel::Bf16DotGemmF32) drive on, when a large-k narrow GEMM is re-blocked through an f32 scratch buffer. The accumulate-side ops (load_lhs/splat_rhs widen f16 -> f32, dot_accumulate folds pairs) forward verbatim to the narrow KernelSimd<f16, f16, f32, f16> impl. So a token’s override still applies, and the twin’s accumulation is bit-identical to the narrow family’s. load_out and store_out are a plain f32 load and store instead, because the C scratch is already f32 and needs no widen or narrow

This is written as 2 explicit impls, one per narrow type, rather than one blanket generic over N. A KernelSimd<N, N, f32, f32> blanket generic in N would collide with the homogeneous <A, A, A, A> blanket under the coherence check. The compiler cannot rule out N = f32. The concrete f16 and bf16 heads cannot unify with <A, A, A, A> (f16 != f32), so they are coherent. This is the same trick as the concrete-type impl_requant_blanket! heads above

Source§

impl<S: KernelSimd<i8, i8, i32, i32>> KernelSimd<i8, i8, i32, i8> for S

Source§

impl<S: KernelSimd<i8, i8, i32, i32>> KernelSimd<i8, i8, i32, u8> for S