gemmkit 0.1.2

//! Runtime ISA dispatch (layer L7)
//!
//! Every element type gemmkit supports gets its own `OnceLock`-memoized descriptor. Each
//! descriptor holds a small struct of monomorphized function pointers (plain, prepacked-RHS,
//! and, under `epilogue`, fused variants) plus the microtile its ISA probe chose
//!
//! Feature detection runs once per type, on that type's first dispatch. Every call after that is
//! a plain indirect call through a typed function pointer, with no `transmute` and no type
//! erasure
//!
//! Adding an ISA is small and mechanical. It takes a new arm in each affected type's `select_*`
//! ladder, plus a thin wrapper (and its packed/fused siblings). The wrapper delegates to the
//! shared generic driver entry, with no new logic. Adding an element type takes a new descriptor
//! plus a `GemmScalar` impl in its own file under this module, not a new crate
//!
//! ## Pinning the kernel: `GEMMKIT_REQUIRE_ISA`
//!
//! By default gemmkit selects the best available ISA at runtime. Setting the environment
//! variable `GEMMKIT_REQUIRE_ISA` to `scalar`, `fma`, `avx512f`, `avx512vnni`, `avx512bf16`,
//! `neon`, or `simd128` forces that kernel. `avx512vnni` selects the `i8` `vpdpbusd` dot kernel,
//! `avx512bf16` selects the `bf16` `vdpbf16ps` dot kernel, and every other type takes the plain
//! AVX-512F path
//!
//! The CPU (or an emulator such as Intel SDE) might not report the required feature. The
//! requested ISA might also not exist on this target architecture. Either case makes selection
//! panic instead of falling back. A CI job that means to exercise a given kernel then fails
//! loudly instead of silently testing a different one
//!
//! `neon` is valid only on aarch64, where it is the baseline ISA. `fma` and `avx512*` are valid
//! only on x86. `simd128` is valid only on a `wasm32` build compiled with
//! `-C target-feature=+simd128`. Pinning `simd128` makes the build assert the SIMD path is live.
//! It panics if that flag is missing, instead of silently falling back to the scalar kernel
//!
//! `auto`, unset, or empty selects the normal auto-selecting behavior. Each type reads the
//! variable at most once, the first time that type's descriptor is built. Its `select_*`
//! memoizes the choice in that type's `OnceLock`. Set the variable in the process environment
//! before the first GEMM call

// Without std, `x86_isa_detected!` collapses to a compile-time `cfg!(...)` constant. These
// clippy lints would otherwise fire on the resulting always-true/false checks in every select_*
// ladder in this module tree
#![cfg_attr(
    not(feature = "std"),
    allow(
        clippy::assertions_on_constants,
        clippy::nonminimal_bool,
        clippy::eq_op
    )
)]

// ForcedIsa (GEMMKIT_REQUIRE_ISA) parsing plus the OnceLock-memoized select-once plumbing
// shared by every family's select_* ladder
#[macro_use]
mod isa;

// c32/c64 complex GEMM dispatch: conj-aware orientation swap, per-ISA wrappers, ComplexScalar impls
#[cfg(feature = "complex")]
mod complex;
// f32/f64 homogeneous-float dispatch: driver entries, per-ISA wrappers, GemmScalar/FusedScalar
// impls
mod float;
// i8 -> i32 integer GEMM dispatch, plus the fused i8 requantizing path
#[cfg(feature = "int8")]
mod int;
// f16/bf16 mixed-precision dispatch (Acc = f32)
#[cfg(feature = "half")]
mod mixed;

#[cfg(feature = "complex")]
pub use complex::ComplexScalar;
#[cfg(feature = "complex")]
pub(crate) use complex::execute_complex;
#[cfg(all(feature = "complex", feature = "epilogue"))]
pub(crate) use complex::execute_complex_fused;
#[cfg(feature = "epilogue")]
pub use float::{FusedScalar, MapScalar};
#[cfg(feature = "epilogue")]
pub(crate) use float::{execute_fused, execute_map, execute_packed_fused};
#[cfg(feature = "int8")]
pub(crate) use int::{
    IntPackedConsume, IntTask, execute_int, execute_int_packed, i8_rhs_depth_multiple, i8_rhs_tile,
    pack_rhs_full_i8,
};
#[cfg(all(feature = "int8", feature = "epilogue"))]
pub(crate) use int::{RequantTask, execute_int_requant};

#[cfg(feature = "epilogue")]
use crate::kernel::epilogue::FusedEpi;
use crate::parallel::Parallelism;
use crate::scalar::{Float, Scalar};
use crate::tuning;
use crate::workspace::Workspace;

/// A fully described GEMM problem, `C <- alpha*A*B + beta*C`, as raw pointers and per-axis
/// `isize` element strides. The homogeneous-type dispatch boundary: A, B, and C all share
/// element type `T` (the accumulator, per [`GemmScalar`], may still differ)
#[derive(Copy, Clone)]
pub struct Task<T> {
    /// Row count of A and C
    pub m: usize,
    /// Contraction length: column count of A, row count of B
    pub k: usize,
    /// Column count of B and C
    pub n: usize,
    /// Scale applied to the A*B product
    pub alpha: T,
    /// LHS base pointer, element `(0,0)`
    pub a: *const T,
    /// LHS row stride
    pub rsa: isize,
    /// LHS column stride
    pub csa: isize,
    /// RHS base pointer
    pub b: *const T,
    /// RHS row stride
    pub rsb: isize,
    /// RHS column stride
    pub csb: isize,
    /// Scale applied to the incoming C before adding alpha*A*B
    pub beta: T,
    /// Output base pointer
    pub c: *mut T,
    /// Output row stride
    pub rsc: isize,
    /// Output column stride
    pub csc: isize,
}

/// One problem for the pointer-array batched API ([`crate::gemm_batched_ptr_unchecked`]): the
/// same `C <- alpha*A*B + beta*C` as `Task`, as raw pointers and `isize` element strides. Each
/// batch entry can have its own shape and live anywhere in memory. The strided
/// [`crate::gemm_batched`] instead shares one shape and steps every operand by a fixed batch
/// stride
#[derive(Copy, Clone)]
pub struct GemmProblem<T> {
    /// Row count of A and C
    pub m: usize,
    /// Contraction length: column count of A, row count of B
    pub k: usize,
    /// Column count of B and C
    pub n: usize,
    /// Scale applied to the A*B product
    pub alpha: T,
    /// LHS base pointer
    pub a: *const T,
    /// LHS row stride
    pub rsa: isize,
    /// LHS column stride
    pub csa: isize,
    /// RHS base pointer
    pub b: *const T,
    /// RHS row stride
    pub rsb: isize,
    /// RHS column stride
    pub csb: isize,
    /// Scale applied to the incoming C before adding alpha*A*B
    pub beta: T,
    /// Output base pointer
    pub c: *mut T,
    /// Output row stride
    pub rsc: isize,
    /// Output column stride
    pub csc: isize,
}

impl<T: Copy> GemmProblem<T> {
    /// Copy the fields into the internal [`Task`] the dispatch layer consumes, with no allocation
    #[inline]
    pub(crate) fn task(&self) -> Task<T> {
        Task {
            m: self.m,
            k: self.k,
            n: self.n,
            alpha: self.alpha,
            a: self.a,
            rsa: self.rsa,
            csa: self.csa,
            b: self.b,
            rsb: self.rsb,
            csb: self.csb,
            beta: self.beta,
            c: self.c,
            rsc: self.rsc,
            csc: self.csc,
        }
    }
}

/// A GEMM whose RHS is already packed into micropanels: `C <- alpha*A*(prepacked B) + beta*C`.
/// It carries the blocking geometry the buffer was packed for (`nr`, `kc`, `nc`). The consuming
/// driver call then reads panels with the exact tiling the pack step used, rather than
/// re-deriving it
///
/// This type is `pub`, like [`Task`], only so it can appear in the doc-hidden [`GemmScalar`]
/// methods. The `dispatch` module itself is private, so this type is not nameable from outside
/// the crate
pub struct PackedConsume<T> {
    /// Row count of A and C
    pub m: usize,
    /// Contraction length: column count of A, which must match the prepacked B's depth
    pub k: usize,
    /// Column count of the prepacked B, and of C
    pub n: usize,
    /// Scale applied to the A*B product
    pub alpha: T,
    /// LHS base pointer
    pub a: *const T,
    /// LHS row stride
    pub rsa: isize,
    /// LHS column stride
    pub csa: isize,
    /// Base of the prepacked RHS micropanel buffer (see [`crate::driver::pack_rhs_full`])
    pub packed: *const T,
    /// Micropanel width `packed` was built with, matching the kernel's `NR`
    pub nr: usize,
    /// Depth-slice size `packed` was built with
    pub kc: usize,
    /// Column-block size `packed` was built with
    pub nc: usize,
    /// Scale applied to the incoming C before adding alpha*A*B
    pub beta: T,
    /// Output base pointer
    pub c: *mut T,
    /// Output row stride
    pub rsc: isize,
    /// Output column stride
    pub csc: isize,
}

/// The dispatch layer runs `f32` and `f64` directly, as a homogeneous float. Under `half`, it
/// also runs `f16` and `bf16` as mixed precision, with `Acc = f32`
///
/// The bound is [`Scalar`], not `Float<Acc = Self>`, so the accumulator type may differ from the
/// element type (the mixed-precision seam). The methods below supply what a generic bound alone
/// cannot express: the degenerate `beta`-only scale, and which kernel family packs and dispatches
/// the value. This keeps the driver and the public API type-agnostic
pub trait GemmScalar: Scalar {
    /// Mirrors [`crate::kernel::KernelFamily::OUT_IS_ACC`] for this type: `true` for `f32`/`f64`,
    /// `false` for `f16`/`bf16`. The prepack constructor reads it so a prepacked buffer blocks
    /// with the same `kc` the consuming kernel will use
    const OUT_IS_ACC: bool;

    /// `C <- beta*C` over the strided output: the degenerate path taken when the `A*B` term
    /// vanishes (`k == 0` or `alpha == 0`). Narrow types scale in `f32` and round back once
    ///
    /// # Safety
    /// `c` valid for the `m x n` region at strides `rsc`/`csc`
    #[doc(hidden)]
    unsafe fn scale_c(beta: Self, c: *mut Self, m: usize, n: usize, rsc: isize, csc: isize);

    /// Pack a full RHS into the prepacked micropanel buffer, using this type's kernel family.
    /// The panel layout is the same for every element type, but the kernel family differs
    /// (`FloatGemm` vs `MixedGemm`). This method exists so [`crate::prepack_rhs`] can call it
    /// without hard-wiring either family
    ///
    /// # Safety
    /// As [`crate::driver::pack_rhs_full`]
    #[doc(hidden)]
    #[allow(clippy::too_many_arguments)]
    unsafe fn pack_rhs_full(
        dst: *mut Self,
        b: *const Self,
        rsb: isize,
        csb: isize,
        k: usize,
        n: usize,
        kc: usize,
        nc: usize,
        nr: usize,
    );

    /// Run the ISA-dispatched plain kernel for this type. Called by the API layer
    ///
    /// # Safety
    /// `task`'s pointers must be valid and `c` must not alias `a`/`b`
    #[doc(hidden)]
    unsafe fn dispatch(task: Task<Self>, par: Parallelism, ws: &mut Workspace);

    /// Run the ISA-dispatched prepacked-RHS kernel for this type
    ///
    /// # Safety
    /// `req`'s pointers must be valid, `c` must not alias `a`/`packed`, and `packed` must have
    /// been produced by [`GemmScalar::pack_rhs_full`] for the geometry recorded in `req`
    #[doc(hidden)]
    unsafe fn dispatch_packed(req: PackedConsume<Self>, par: Parallelism, ws: &mut Workspace);

    /// This type's dispatched kernel's microtile `(mr, nr)`, i.e. `(MR_REG*LANES, NR)`. The
    /// prepack constructor calls this to size the buffer's blocking geometry through the same
    /// ISA choice the consuming call will make
    #[doc(hidden)]
    fn rhs_tile() -> (usize, usize);

    /// This type's dispatched kernel family's [`crate::kernel::KernelFamily::DEPTH_MULTIPLE`].
    /// The prepack constructor rounds the packed depth up to it so the buffer's layout matches
    /// the consuming kernel's. `1` (the default) covers every family except the bf16
    /// `vdpbf16ps` dot kernel, which overrides it to `2`
    #[doc(hidden)]
    fn rhs_depth_multiple() -> usize {
        1
    }

    /// Run the ISA-dispatched fused-epilogue kernel for this type. Every type covered by
    /// [`FusedScalar`] provides one. The real floats (`f32`/`f64`) go through the `float` module.
    /// The narrow floats (`f16`/`bf16`, `Acc = f32`) go through `mixed`. This method is required
    /// rather than defaulted, because the [`FusedScalar`] bound on the public fused API admits
    /// exactly those 4 types
    ///
    /// # Safety
    /// `task`'s pointers are valid and `c` does not alias `a` or `b`. `epi`'s bias is valid and
    /// disjoint from `c` (validated by the API layer)
    #[doc(hidden)]
    #[cfg(feature = "epilogue")]
    unsafe fn dispatch_fused(
        task: Task<Self>,
        epi: FusedEpi<Self>,
        par: Parallelism,
        ws: &mut Workspace,
    );

    /// Run the ISA-dispatched prepacked-RHS fused-epilogue kernel for this type. It is the fused
    /// twin of [`GemmScalar::dispatch_packed`], and threads `epi` into the prepacked driver entry
    /// (`driver::run_packed_rhs_epilogue`). `epi` already lives in the prepacked buffer's
    /// oriented frame, because the packed path never re-orients inside the driver. The public
    /// `gemm_packed_a_fused` entry therefore pre-flips the bias axis before it builds the
    /// transposed consume
    ///
    /// # Safety
    /// As [`GemmScalar::dispatch_packed`], plus `epi`'s bias is valid for the problem's `m`/`n`
    /// and disjoint from `c` (validated by the API layer)
    #[doc(hidden)]
    #[cfg(feature = "epilogue")]
    unsafe fn dispatch_packed_fused(
        req: PackedConsume<Self>,
        epi: FusedEpi<Self>,
        par: Parallelism,
        ws: &mut Workspace,
    );
}

/// Top-level entry used by the API layer: handle the degenerate cases (here, where the element
/// type is concrete) and then run the ISA-dispatched kernel
///
/// # Safety
/// `task`'s pointers must be valid for the implied regions and `c` must not alias `a`/`b`
pub(crate) unsafe fn execute<T: GemmScalar>(task: Task<T>, par: Parallelism, ws: &mut Workspace) {
    unsafe {
        if task.m == 0 || task.n == 0 {
            return;
        }
        // k == 0 or alpha == 0: the A*B term vanishes, C <- beta*C only
        if task.k == 0 || task.alpha == T::ZERO {
            T::scale_c(task.beta, task.c, task.m, task.n, task.rsc, task.csc);
            return;
        }
        T::dispatch(task, par, ws);
    }
}

/// Top-level entry for the prepacked-RHS path
///
/// It handles the degenerate case, where the `A*B` term vanishes (`C <- beta*C` without ever
/// touching the packed buffer). It then runs the ISA-dispatched prepacked kernel
///
/// # Safety
/// As [`execute`], plus `req.packed` valid for the recorded geometry and not aliasing `c`
pub(crate) unsafe fn execute_packed<T: GemmScalar>(
    req: PackedConsume<T>,
    par: Parallelism,
    ws: &mut Workspace,
) {
    unsafe {
        if req.m == 0 || req.n == 0 {
            return;
        }
        if req.k == 0 || req.alpha == T::ZERO {
            T::scale_c(req.beta, req.c, req.m, req.n, req.rsc, req.csc);
            return;
        }
        T::dispatch_packed(req, par, ws);
    }
}

/// Normalizes orientation for every dispatch path: the float/mixed [`Task`], the integer
/// `IntTask`, and the requantizing `RequantTask`. When `C` is row-major-ish (`|csc| < |rsc|`),
/// it computes `C^T = B^T*A^T` instead, so the kernel still writes columns contiguously
/// (`rsc == 1` after the swap). It does this by exchanging `m` and `n`, the `A` and `B` pointers
/// and strides, and `rsc` and `csc`
///
/// It returns `true` when it swaps, so a caller can flip whatever policy co-varies with
/// orientation, such as the bias axis or conj flags. It takes the pointer type `L` as a generic
/// parameter, because the 3 task shapes differ only there. Each still pairs `A` and `B` under one
/// element type
#[inline]
#[allow(clippy::too_many_arguments)]
fn orient_swap<L>(
    m: &mut usize,
    n: &mut usize,
    a: &mut *const L,
    rsa: &mut isize,
    csa: &mut isize,
    b: &mut *const L,
    rsb: &mut isize,
    csb: &mut isize,
    rsc: &mut isize,
    csc: &mut isize,
) -> bool {
    if csc.unsigned_abs() < rsc.unsigned_abs() {
        core::mem::swap(m, n);
        core::mem::swap(a, b); // new A pointer = old B pointer
        core::mem::swap(rsa, csb); // new rsa = old csb, new csb = old rsa
        core::mem::swap(csa, rsb); // new csa = old rsb, new rsb = old csa
        core::mem::swap(rsc, csc);
        true
    } else {
        false
    }
}

/// [`orient_swap`] specialized to the homogeneous float / mixed [`Task`] shape
#[inline]
fn orient_transpose<T>(t: &mut Task<T>) -> bool {
    orient_swap(
        &mut t.m, &mut t.n, &mut t.a, &mut t.rsa, &mut t.csa, &mut t.b, &mut t.rsb, &mut t.csb,
        &mut t.rsc, &mut t.csc,
    )
}

/// `true` when a post-swap problem should take the horizontal `small_mn` path, from the raw
/// oriented dimensions and strides. This holds for small `m` and `n` with a long contraction,
/// when both operands already stream contiguously along `k`. A's columns are unit-stride
/// (`csa == 1`), and B's rows are unit-stride (`rsb == 1`). The horizontal kernel can then read
/// them in place
///
/// This is the field-level core. It lets the float/mixed [`Task`] path and the heterogeneous
/// integer `IntTask` path (which has no `Task<T>` to wrap) share one gate. This keeps the 2
/// paths from drifting apart into separate copies
#[inline]
fn small_mn_eligible_dims(m: usize, n: usize, k: usize, csa: isize, rsb: isize) -> bool {
    m <= tuning::small_mn_dim()
        && n <= tuning::small_mn_dim()
        && k > tuning::small_k_threshold()
        && csa == 1
        && rsb == 1
}

/// [`small_mn_eligible_dims`] over a [`Task`]'s fields, for the float and mixed dispatch entries
#[inline]
fn small_mn_eligible<T>(t: &Task<T>) -> bool {
    small_mn_eligible_dims(t.m, t.n, t.k, t.csa, t.rsb)
}

/// `true` when a post-swap problem clears the `small_mn` dims and `k` gates. At least one
/// operand still misses the unit-stride-along-`k` predicate (`csa != 1` or `rsb != 1`). The pack
/// tier then copies only the failing operand into `k`-contiguous scratch and runs the same
/// horizontal kernel (see [`crate::special::small_mn::prepack_operands`]). An all-row-major or
/// all-col-major shape hits this path with exactly one operand needing the copy
///
/// This shares its `m`/`n` bound with [`small_mn_eligible_dims`], so the 2 gates cannot drift
/// apart, but it has its own `k` floor ([`crate::tuning::small_mn_pack_min_k`], not
/// `small_k_threshold`). It also requires a failing stride where the zero-copy gate forbids one.
/// This makes the 2 gates mutually exclusive, so a small_mn-shaped call takes at most one of them
///
/// This is the field-level core, shared by the [`Task`] and `IntTask` paths exactly as
/// [`small_mn_eligible_dims`] is
#[inline]
fn small_mn_pack_eligible_dims(m: usize, n: usize, k: usize, csa: isize, rsb: isize) -> bool {
    m <= tuning::small_mn_dim()
        && n <= tuning::small_mn_dim()
        && k > tuning::small_mn_pack_min_k()
        && !(csa == 1 && rsb == 1)
}

/// [`small_mn_pack_eligible_dims`] over a [`Task`]'s fields, for the float and mixed dispatch
/// entries
#[inline]
fn small_mn_pack_eligible<T>(t: &Task<T>) -> bool {
    small_mn_pack_eligible_dims(t.m, t.n, t.k, t.csa, t.rsb)
}

/// `C <- beta*C` for a homogeneous float type (`f32`/`f64`). It scales in place.
/// `beta == 0` overwrites with zero rather than multiplying, so a NaN or inf already in `C`
/// does not poison the result. The float impl of [`GemmScalar::scale_c`] forwards here. Narrow
/// types use `scale_c_narrow` in the `mixed` module instead
unsafe fn scale_c_float<T: Float>(beta: T, c: *mut T, m: usize, n: usize, rsc: isize, csc: isize) {
    unsafe {
        for j in 0..n {
            for i in 0..m {
                let p = c.offset(i as isize * rsc + j as isize * csc);
                if beta == T::ZERO {
                    *p = T::ZERO;
                } else if beta != T::ONE {
                    *p = beta * *p;
                }
            }
        }
    }
}