Crate gemmkit

Expand description

§gemmkit

A GEMM (general matrix multiply) engine with no ndarray dependency. It computes C <- alpha*A*B + beta*C over data-type-agnostic &[T] and stride views, or raw pointers. It picks the fastest instruction set the running CPU supports

§Quick start

use gemmkit::{gemm, MatRef, MatMut, Parallelism};

// 2x3 * 3x2 = 2x2, row-major
let a = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
let b = [7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0];
let mut c = [0.0_f32; 4];
gemm(
    1.0,
    MatRef::from_row_major(&a, 2, 3),
    MatRef::from_row_major(&b, 3, 2),
    0.0,
    MatMut::from_row_major(&mut c, 2, 2),
    Parallelism::Serial,
);
assert_eq!(c, [58.0, 64.0, 139.0, 154.0]);

§Architecture in brief

Every axis of variation lives behind a trait. The instruction set is simd::SimdOps, the element type is scalar::Scalar, and the operation family is kernel::KernelFamily. One generic 5-loop driver and one generic microkernel cover every (type, ISA, tile) combination with no macros and no transmute. See ARCHITECTURE.md for the full tour

§Features

std (default): runtime cache and CPU-feature detection, the GEMMKIT_REQUIRE_ISA env var, tuning knobs, and the thread-local workspace pool. With std off, the crate is #![no_std] and needs only core and alloc
parallel (default, implies std): rayon multithreading. With it off, the crate still compiles and runs, single-threaded
wasm_threads (implies parallel): a stable-toolchain rayon thread pool for a threaded wasm target, sized from a tuning knob because wasm cannot report its own core count
complex: complex GEMM over c32/c64 with optional conjugation (gemm_cplx). Pulls in num-complex
half: f16/bf16 mixed-precision GEMM, accumulating in f32. Pulls in half
int8: i8 -> i32 integer GEMM (gemm_i8), with no extra dependency
epilogue: fused epilogues for bias and activation (gemm_fused, gemm_batched_fused*, the prepacked gemm_packed_a_fused* and gemm_packed_b_fused*, and, with complex, gemm_cplx_fused*). Also a user-defined per-element closure (gemm_map, f32/f64 only), and, with int8, requantized i8/u8 output (gemm_i8_requant*). No extra dependency

complex, half, and int8 are off by default, so a plain f32/f64 build pays for none of their codegen or dependencies. epilogue is off by default too, so a plain-GEMM build pays for none of its codegen

Re-exports§

pub use kernel::epilogue::BiasDim;epilogue
pub use scalar::ComplexFloat;complex
pub use scalar::NarrowFloat;half
pub use scalar::Float;
pub use scalar::Scalar;
pub use cache::CacheTopology;
pub use cache::Machine;
pub use cache::topology;
pub use half::bf16;half
pub use half::f16;half
pub use num_complex::Complex;complex

Modules§

cache: Cache topology and analytical blocking (layer L3)
driver: The generic GEMM driver (layer L5): one 5-loop nest, fully generic over the KernelFamily and the ISA token. It never names a concrete element type, a concrete ISA, or a macro, so adding a family or an ISA never touches this file. This is the open/closed property the architecture promises. A test proves it by declaring a 2nd family and driving it through this same run
kernel: Kernel families (layer L4): the seam between the driver and a concrete kind of GEMM
scalar: The element-type seam (layer L0): what a value must supply to be a GEMM operand
simd: SIMD abstraction layer (L0): the per-ISA register vocabulary every microkernel builds on
tuning: Unified tuning surface (cross-cutting)

Structs§

BatchProblem: 1 element of a checked pointer-array batched GEMM (gemm_batched_slice): C <- alpha*A*B + beta*C over safe views
GemmProblem: 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
MatMut: A mutable strided matrix view over a slice
MatRef: An immutable strided matrix view over a slice
PackedLhs: An A matrix packed once into gemmkit’s internal micropanel-major layout, for reuse across many products that share the same A. This fits a fixed weight matrix streamed against differently shaped B operands. prepack_lhs builds it. gemm_packed_a and its _with, _unchecked, and _fused siblings consume it and skip the per-call LHS pack
PackedRhs: A B matrix packed once into gemmkit’s internal micropanel-major layout, for reuse across many products that share the same B. This fits a fixed weight matrix streamed against many activation matrices. prepack_rhs builds it. gemm_packed_b and its _with, _unchecked, and _fused siblings consume it and skip the per-call RHS pack
Requantizeepilogue and int8: The quantization parameters for the requantizing entries: RequantScale, an integer zero_point, and an optional per-row i32 bias (length m, the standard qlinear layer bias). The output is C[i,j] = clamp(zero_point + round_ne(scale*(sum_k A*B + bias[i])), LO, HI) with round-half-to-even. scale is the per-tensor value or the per-row scale_i. [LO, HI] is set by the entry: [-128, 127] for gemm_i8_requant, [0, 255] for gemm_i8_requant_u8
Workspace: A growable, 64-byte-aligned scratch buffer for packing A and B

Enums§

Activationepilogue: The activation a gemm_fused call applies last, after the bias add
Biasepilogue: 1 bias value broadcast across a whole output row or column, combined into a gemm_fused call’s epilogue after the alpha*A*B + beta*C product and before the activation
Parallelism: Threading strategy for a GEMM call
RequantScaleepilogue and int8: The output scale for the requantizing entries: either 1 value applied to the whole tensor, or 1 value per output row/channel (the per-channel quantized-inference convention). Every scale must be finite and > 0

Traits§

ComplexScalarcomplex: Complex element types gemmkit can dispatch: Complex<f32> and Complex<f64>
FusedScalarepilogue: The sealed element-type bound for the fused-epilogue public API
GemmScalar: 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
MapScalarepilogue: The sealed element-type bound for the user-defined map-epilogue public API (crate::gemm_map): the real floats f32/f64 only

Functions§

gemm: C <- alpha*A*B + beta*C over safe slice views, using the thread-local workspace pool
gemm_batched: Strided-batched GEMM: C_b <- alpha*A_b*B_b + beta*C_b for b in 0..batch, in 1 call, parallelized across the batch rather than within each element. Every element shares the single-element shape and strides of a/b/c. Element b is based at a.data + b*a_batch_stride, and likewise for b/c. A *_batch_stride of 0 broadcasts 1 operand across the whole batch, valid for the read-only A/B but never for C. Uses the thread-local workspace pool
gemm_batched_fusedepilogue: Strided-batched GEMM with a fused epilogue shared by every element: C_b <- act(alpha*A_b*B_b + beta*C_b + bias) for b in 0..batch, in 1 call, parallelized across the batch. 1 bias vector and 1 activation apply to every element, the batched-linear-layer case of 1 layer applied to a batch of inputs. Shape, stride, and broadcast conventions match gemm_batched. Uses the thread-local workspace pool. When bias == None && act == None, this takes the plain gemm_batched path
gemm_batched_fused_unchecked^⚠epilogue: The raw strided-batched fused engine: C_e <- act(alpha*A_e*B_e + beta*C_e + bias) for e in 0..batch, over pointers and isize strides, with no bounds, alias, or shape checks. This is the raw-parts form of gemm_batched_fused, combining gemm_batched_unchecked’s per-element shape with the shared bias/activation of gemm_fused_unchecked. Element e is based at a + e*a_batch_stride / b + e*b_batch_stride / c + e*c_batch_stride, all sharing the single-element shape (m, k, n) and element strides. The 1 bias (a (ptr, dim) pair, read only when has_bias) and 1 act apply to every element. Uses the thread-local workspace pool
gemm_batched_fused_unchecked_with^⚠epilogue: As gemm_batched_fused_unchecked but with a caller-owned Workspace
gemm_batched_fused_withepilogue: Like gemm_batched_fused but reuses a caller-owned Workspace. This uses the same split as gemm_batched_with: the serial and few-but-large schedules pack through ws. The batch-parallel schedule packs through each worker’s own thread-local pool instead
gemm_batched_ptr_unchecked^⚠: Runs a pointer-array batched GEMM: every element in problems is an independent product with its own shape and pointers (GemmProblem). This parallelizes across the batch, with whole GEMMs assigned to workers, each run serially and cache-hot. The raw counterpart of gemm_batched_slice, for callers (FFI, adapters) that validate their own inputs and may use arbitrary pointers or negative strides. Deterministic across thread counts, since each element runs wholly on 1 worker, and takes the problems slice as-is with no per-call allocation
gemm_batched_slice: Runs a checked pointer-array batched GEMM: problems[i].c <- alpha*A*B + beta*C for each element, each an independent product over safe views, parallelized across the batch. This is the safe counterpart of gemm_batched_ptr_unchecked. Because every c is a distinct MatMut, the outputs are pairwise disjoint and cannot alias the inputs by construction. The borrow checker already forbids 2 overlapping &mut borrows, so validation only covers per-element shape agreement, in-bounds strides, and self-aliasing. Deterministic across thread counts
gemm_batched_unchecked^⚠: The raw strided-batched engine: gemm_batched over pointers and isize strides, with no bounds, alias, or shape checks. Element e is based at a + e*a_batch_stride / b + e*b_batch_stride / c + e*c_batch_stride, all sharing the single-element shape (m, k, n) and element strides. Adapter crates (e.g. an ndarray Array3 batched on axis 0) and FFI callers that supply their own pointers or arbitrary strides use this path. Uses the thread-local workspace pool
gemm_batched_unchecked_with^⚠: As gemm_batched_unchecked but with a caller-owned Workspace
gemm_batched_with: Like gemm_batched but reuses a caller-owned Workspace across calls. The serial and few-but-large schedules pack through ws. The batch-parallel schedule instead has each worker pack through its own thread-local pool, since 1 Workspace cannot back concurrent packing from several threads
gemm_cplxcomplex: Complex GEMM with optional conjugation: C <- alpha*op(A)*op(B) + beta*C, where op(A) = conj(A) when conj_a (independently, op(B) = conj(B) when conj_b). T is Complex<f32> or Complex<f64> (re-exported as crate::c32 / crate::c64). Uses the thread-local workspace pool
gemm_cplx_fusedcomplex and epilogue: Complex GEMM with an optional fused per-row / per-col bias: C <- alpha*op(A)*op(B) + beta*C + bias in 1 pass, where op conjugates an operand exactly as in gemm_cplx. T is Complex<f32> or Complex<f64>. Uses the thread-local workspace pool
gemm_cplx_fused_unchecked^⚠complex and epilogue: The raw complex fused engine: C <- alpha*op(A)*op(B) + beta*C + bias over pointers and isize strides, with no bounds, alias, or shape checks: the raw-parts form of gemm_cplx_fused. op conjugates an operand when its conj_* flag is set. bias is a (ptr, dim) pair, read only when has_bias, and added verbatim, never conjugated. There is no activation parameter, since it has no definition on complex numbers. Uses the thread-local workspace pool
gemm_cplx_fused_unchecked_with^⚠complex and epilogue: As gemm_cplx_fused_unchecked but with a caller-owned Workspace
gemm_cplx_fused_withcomplex and epilogue: Like gemm_cplx_fused but reuses a caller-owned Workspace
gemm_cplx_unchecked^⚠complex: The raw complex engine: C <- alpha*op(A)*op(B) + beta*C over pointers and isize strides, with no bounds, alias, or shape checks: the complex counterpart of gemm_unchecked, where op conjugates an operand when its conj_* flag is set. Adapter crates (e.g. ndarray) use this path to express transposed or negative strides that the checked API rejects. Uses the thread-local workspace pool
gemm_cplx_unchecked_with^⚠complex: As gemm_cplx_unchecked but with a caller-owned Workspace
gemm_cplx_withcomplex: Like gemm_cplx but reuses a caller-owned Workspace
gemm_fusedepilogue: C <- act(alpha*A*B + beta*C + bias) in 1 pass over safe slice views, using the thread-local workspace pool. bias == None && act == None delegates to plain gemm
gemm_fused_unchecked^⚠epilogue: The raw fused engine: C <- act(alpha*A*B + beta*C + bias) over pointers and isize strides, with no bounds, alias, or shape checks. bias is a (ptr, dim) pair, read only when has_bias. Uses the thread-local workspace pool
gemm_fused_unchecked_with^⚠epilogue: As gemm_fused_unchecked but with a caller-owned Workspace
gemm_fused_withepilogue: Like gemm_fused but reuses a caller-owned Workspace. Accepts the same f32/f64 and, under half, f16/bf16 types, with the same pre-narrow f32 epilogue precision described at gemm_fused
gemm_i8int8: Integer GEMM: C <- alpha*A*B + beta*C with i8 inputs accumulated into an i32 output (alpha, beta, C are i32). Wraps on overflow, the standard integer-GEMM convention. Uses the thread-local workspace pool
gemm_i8_packed_bint8: C(i32) <- alpha*A(i8)*(prepacked B) + beta*C, consuming a PackedRhs<i8> (B prepacked once) instead of B itself, using the thread-local workspace pool. This is the integer (i8 -> i32) twin of gemm_packed_b. It skips the RHS pack that, for the VNNI kernel, would otherwise run on every call
gemm_i8_packed_b_unchecked^⚠int8: As gemm_i8_packed_b but over raw A/C pointers and strides, with no bounds or alias checks. This is the heterogeneous (i8 -> i32) counterpart of gemm_packed_b_unchecked. The shared k and output n come from packed, and m is A’s row count, the same as C’s. Uses the thread-local workspace pool
gemm_i8_packed_b_unchecked_with^⚠int8: As gemm_i8_packed_b_unchecked but with a caller-owned Workspace
gemm_i8_packed_b_withint8: Like gemm_i8_packed_b but reuses a caller-owned Workspace
gemm_i8_requantepilogue and int8: Requantizing integer GEMM: i8 inputs multiply into an i32 accumulator, then requantize to an i8 output in 1 pass. This skips the full m x n i32 materialization that a separate gemm_i8 call followed by a requantize pass would need. No alpha (it folds into scale) and no beta (accumulating into an already-quantized C is not well-defined). Uses the thread-local workspace pool
gemm_i8_requant_u8epilogue and int8: Requantizing integer GEMM with an unsigned u8 output, the ONNX QLinearMatMul activation convention. i8 inputs multiply into an i32 accumulator, then requantize in 1 pass to C[i,j] = clamp(zero_point + round_ne(scale*(sum_k A*B + bias[i])), 0, 255) with round-half-to-even. scale is the per-tensor value or the per-row scale_i. This is the u8-output twin of gemm_i8_requant, differing only in the output domain ([0, 255] instead of [-128, 127]) and the accepted zero_point range. No alpha (it folds into scale) and no beta (accumulating into an already-quantized C is not well-defined). Uses the thread-local workspace pool
gemm_i8_requant_u8_unchecked^⚠epilogue and int8: C(u8) <- clamp(zp + round_ne(scale*(A*B + bias)), 0, 255) over raw pointers and isize element strides, with no bounds, alias, or shape checks: the unsigned twin of gemm_i8_requant_unchecked. bias is a per-row i32 pointer, read only when has_bias. The scale is the scalar scale unless has_row_scales is set, in which case row_scales supplies 1 f32 per output row (length m) instead. Uses the thread-local workspace pool
gemm_i8_requant_u8_unchecked_with^⚠epilogue and int8: Like gemm_i8_requant_u8_unchecked but reuses a caller-owned Workspace instead of the thread-local pool
gemm_i8_requant_u8_withepilogue and int8: Like gemm_i8_requant_u8 but reuses a caller-owned Workspace instead of the thread-local pool
gemm_i8_requant_unchecked^⚠epilogue and int8: C(i8) <- clamp(zp + round_ne(scale*(A*B + bias)), -128, 127) over raw pointers and isize element strides, with no bounds, alias, or shape checks. bias is a per-row i32 pointer, read only when has_bias. The scale is the scalar scale unless has_row_scales is set, in which case row_scales supplies 1 f32 per output row (length m) instead. Uses the thread-local workspace pool
gemm_i8_requant_unchecked_with^⚠epilogue and int8: Like gemm_i8_requant_unchecked but reuses a caller-owned Workspace instead of the thread-local pool
gemm_i8_requant_withepilogue and int8: Like gemm_i8_requant but reuses a caller-owned Workspace instead of the thread-local pool
gemm_i8_unchecked^⚠int8: C(i32) <- alpha*A(i8)*B(i8) + beta*C over raw pointers and isize element strides, with no bounds, alias, or shape checks. This is the i8 -> i32 escape hatch for negative strides or raw-pointer callers. gemm_unchecked is typed for the homogeneous surface and cannot express a differing output type. Uses the thread-local workspace pool
gemm_i8_unchecked_with^⚠int8: Like gemm_i8_unchecked but reuses a caller-owned Workspace instead of the thread-local pool
gemm_i8_withint8: Like gemm_i8 but reuses a caller-owned Workspace instead of the thread-local pool
gemm_mapepilogue: C[r, c] <- f(alpha*A*B + beta*C, r, c) in 1 fused pass. This is a plain GEMM with a caller closure applied to each output element at its final value. It runs over safe slice views, using the thread-local workspace pool. (r, c) is the user-frame coordinate of C: row r, column c. f fires exactly once per element, at the point the plain kernel would store it
gemm_map_unchecked^⚠epilogue: C[r, c] <- f(alpha*A*B + beta*C, r, c) over raw pointers and isize element strides, with no bounds, alias, or shape checks. (r, c) is the user-frame coordinate of C. Uses the thread-local workspace pool
gemm_map_unchecked_with^⚠epilogue: Like gemm_map_unchecked but reuses a caller-owned Workspace instead of the thread-local pool
gemm_map_withepilogue: Like gemm_map but reuses a caller-owned Workspace instead of the thread-local pool: zero heap allocation once the workspace has grown to fit the 1st sufficiently large call
gemm_packed_a: C <- alpha*A*B + beta*C, consuming a PackedLhs (A prepacked once) instead of A itself, using the thread-local workspace pool. This skips the per-call LHS pack that gemm would run
gemm_packed_a_fusedepilogue: C <- act(alpha*(prepacked A)*B + beta*C + bias) in one pass. This is a fused epilogue over a reused PackedLhs, using the thread-local workspace pool. It is the fused twin of gemm_packed_a. The bias is folded in with 1 IEEE add right after the final beta-scaled store. The activation applies next, fused into the same store the packed kernel already runs. bias == None && act == None reproduces gemm_packed_a bit-for-bit
gemm_packed_a_fused_unchecked^⚠epilogue: As gemm_packed_a_fused but over raw B/C pointers and strides, with no bounds or alias checks. bias is a (ptr, dim) pair, enabled by has_bias and ignored when has_bias == false, in the user frame, where a PerRow bias indexes A.rows, the same as C.rows. act applies last. Uses the thread-local workspace pool
gemm_packed_a_fused_unchecked_with^⚠epilogue: As gemm_packed_a_fused_unchecked but with a caller-owned Workspace
gemm_packed_a_fused_withepilogue: Like gemm_packed_a_fused but reuses a caller-owned Workspace
gemm_packed_a_unchecked^⚠: As gemm_packed_a but over raw B/C pointers and strides, with no bounds or alias checks. The shared k and output-row count m come from packed, and n is B’s column count, the same as C’s. Uses the thread-local workspace pool
gemm_packed_a_unchecked_with^⚠: As gemm_packed_a_unchecked but with a caller-owned Workspace
gemm_packed_a_with: Like gemm_packed_a but reuses a caller-owned Workspace
gemm_packed_b: C <- alpha*A*B + beta*C, consuming a PackedRhs (B prepacked once) instead of B itself, using the thread-local workspace pool. This skips the per-call RHS pack that gemm would run
gemm_packed_b_fusedepilogue: C <- act(alpha*A*(prepacked B) + beta*C + bias) in one pass. This is a fused epilogue over a reused PackedRhs, using the thread-local workspace pool. It is the fused twin of gemm_packed_b. The bias is folded in with 1 IEEE add right after the final beta-scaled store. The activation applies next, fused into the same store the packed kernel already runs. bias == None && act == None reproduces gemm_packed_b bit-for-bit
gemm_packed_b_fused_unchecked^⚠epilogue: As gemm_packed_b_fused but over raw A/C pointers and strides, with no bounds or alias checks. bias is a (ptr, dim) pair, enabled by has_bias and ignored when has_bias == false, in the user frame, because this path never swaps orientation. act applies last. Uses the thread-local workspace pool
gemm_packed_b_fused_unchecked_with^⚠epilogue: As gemm_packed_b_fused_unchecked but with a caller-owned Workspace
gemm_packed_b_fused_withepilogue: Like gemm_packed_b_fused but reuses a caller-owned Workspace
gemm_packed_b_unchecked^⚠: As gemm_packed_b but over raw A/C pointers and strides, with no bounds or alias checks. The shared k and output n come from packed, and m is A’s row count, the same as C’s. Uses the thread-local workspace pool
gemm_packed_b_unchecked_with^⚠: As gemm_packed_b_unchecked but with a caller-owned Workspace
gemm_packed_b_with: Like gemm_packed_b but reuses a caller-owned Workspace
gemm_unchecked^⚠: The raw engine: C <- alpha*A*B + beta*C over pointers and isize strides, with no bounds, alias, or shape checks. Uses the thread-local workspace pool
gemm_unchecked_with^⚠: Like gemm_unchecked but reuses a caller-owned Workspace
gemm_with: Like gemm but reuses a caller-owned Workspace: zero heap allocation once the workspace has grown to fit the 1st sufficiently large call
prepack_lhs: Pack an m x k A view into a PackedLhs for reuse across many gemm_packed_a calls. The pack runs once, single-threaded, right here, so every later call skips it
prepack_lhs_unchecked^⚠: As prepack_lhs but over a raw m x k A pointer and strides, with no bounds check. Use this raw form for an adapter or FFI caller that validates its own inputs
prepack_rhs: Pack a k x n B view into a PackedRhs for reuse across many gemm_packed_b calls. The pack runs once, single-threaded, right here, so every later call skips it
prepack_rhs_i8int8: Pack a k x n i8 RHS into a PackedRhs<i8> for reuse across many gemm_i8_packed_b calls. This fits the quantized-inference pattern of constant i8 weights against a stream of i8 activation batches. The pack runs once, single-threaded, right here, so later calls skip it
prepack_rhs_i8_unchecked^⚠int8: As prepack_rhs_i8 but over a raw k x n B pointer and strides, with no bounds check. Use this raw form for an adapter or FFI caller that validates its own inputs
prepack_rhs_unchecked^⚠: As prepack_rhs but over a raw k x n B pointer and strides, with no bounds check. Use this raw form for an adapter or FFI caller that validates its own inputs

Type Aliases§

c32complex: Complex<f32>: the single-precision complex element type
c64complex: Complex<f64>: the double-precision complex element type

Crate gemmkit

Crate gemmkit Copy item path

§gemmkit

§Quick start

§Architecture in brief

§Features

Re-exports§

Modules§

Structs§

Enums§

Traits§

Functions§

Type Aliases§

Crate gemmkit