gemmkit 0.1.2

A clean, extensible, high-performance GEMM (general matrix multiply) engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

//! # 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

#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![warn(missing_docs)]
#![allow(clippy::missing_safety_doc)] // documented at module or contract level, not per function

// Some modules, such as the batched and packed API buffers, need heap `Vec`s even in a
// no_std build. The crate pulls in `alloc` unconditionally instead of gating it on `std`
extern crate alloc;

// Cache topology detection and BLIS-model analytical blocking (layer L3)
pub mod cache;
// Generic 5-loop GEMM driver, generic over the kernel family and the ISA (layer L5)
pub mod driver;
// Kernel families: the operation-family seam (layer L4)
pub mod kernel;
// Numeric scalar abstraction: identity constants and the accumulator type (layer L0)
pub mod scalar;
// SIMD abstraction: ISA tokens and per-type vector ops (layer L0)
pub mod simd;
// Tuning surface: heuristic thresholds with env-var and setter overrides
pub mod tuning;

// Public API: safe slice/stride entry points plus the raw unchecked engine (layer L8a)
mod api;
// Shared validation and lowering surface for the view adapters, support for L8a and not a
// layer of its own. Hidden from the public docs and versioned in lockstep with the adapters
#[doc(hidden)]
pub mod adapter;
// Runtime ISA dispatch, memoized per element type (layer L7)
mod dispatch;
// Packing primitives shared by every kernel family (layer L1)
mod pack;
// Parallelism control and job splitting across workers (layer L2)
mod parallel;
// Special-case paths that bypass the driver for shapes it fits poorly (layer L6)
mod special;
// Packing-buffer workspace: a thread-local pool plus explicit reuse
mod workspace;

#[cfg(feature = "epilogue")]
pub use api::{
    Activation, Bias, gemm_batched_fused, gemm_batched_fused_unchecked,
    gemm_batched_fused_unchecked_with, gemm_batched_fused_with, gemm_fused, gemm_fused_unchecked,
    gemm_fused_unchecked_with, gemm_fused_with, gemm_map, gemm_map_unchecked,
    gemm_map_unchecked_with, gemm_map_with, gemm_packed_a_fused, gemm_packed_a_fused_unchecked,
    gemm_packed_a_fused_unchecked_with, gemm_packed_a_fused_with, gemm_packed_b_fused,
    gemm_packed_b_fused_unchecked, gemm_packed_b_fused_unchecked_with, gemm_packed_b_fused_with,
};
pub use api::{
    BatchProblem, MatMut, MatRef, PackedLhs, PackedRhs, gemm, gemm_batched,
    gemm_batched_ptr_unchecked, gemm_batched_slice, gemm_batched_unchecked,
    gemm_batched_unchecked_with, gemm_batched_with, gemm_packed_a, gemm_packed_a_unchecked,
    gemm_packed_a_unchecked_with, gemm_packed_a_with, gemm_packed_b, gemm_packed_b_unchecked,
    gemm_packed_b_unchecked_with, gemm_packed_b_with, gemm_unchecked, gemm_unchecked_with,
    gemm_with, prepack_lhs, prepack_lhs_unchecked, prepack_rhs, prepack_rhs_unchecked,
};
#[cfg(all(feature = "int8", feature = "epilogue"))]
pub use api::{
    RequantScale, Requantize, gemm_i8_requant, gemm_i8_requant_u8, gemm_i8_requant_u8_unchecked,
    gemm_i8_requant_u8_unchecked_with, gemm_i8_requant_u8_with, gemm_i8_requant_unchecked,
    gemm_i8_requant_unchecked_with, gemm_i8_requant_with,
};
#[cfg(feature = "complex")]
pub use api::{gemm_cplx, gemm_cplx_unchecked, gemm_cplx_unchecked_with, gemm_cplx_with};
#[cfg(all(feature = "complex", feature = "epilogue"))]
pub use api::{
    gemm_cplx_fused, gemm_cplx_fused_unchecked, gemm_cplx_fused_unchecked_with,
    gemm_cplx_fused_with,
};
#[cfg(feature = "int8")]
pub use api::{
    gemm_i8, gemm_i8_packed_b, gemm_i8_packed_b_unchecked, gemm_i8_packed_b_unchecked_with,
    gemm_i8_packed_b_with, gemm_i8_unchecked, gemm_i8_unchecked_with, gemm_i8_with, prepack_rhs_i8,
    prepack_rhs_i8_unchecked,
};
#[cfg(feature = "complex")]
pub use dispatch::ComplexScalar;
pub use dispatch::GemmProblem;
pub use dispatch::GemmScalar;
#[cfg(feature = "epilogue")]
pub use dispatch::{FusedScalar, MapScalar};
#[cfg(feature = "epilogue")]
pub use kernel::epilogue::BiasDim;
pub use parallel::Parallelism;
#[cfg(feature = "complex")]
pub use scalar::ComplexFloat;
#[cfg(feature = "half")]
pub use scalar::NarrowFloat;
pub use scalar::{Float, Scalar};
pub use workspace::Workspace;

#[doc(no_inline)]
pub use cache::{CacheTopology, Machine, topology};

/// Re-exported [`half`] narrow float types (`f16`, `bf16`), so you do not need to
/// depend on `half` directly to call the `half`-gated GEMM entry points. Both
/// accumulate in `f32` (their [`Scalar::Acc`])
#[cfg(feature = "half")]
#[doc(no_inline)]
pub use half::{bf16, f16};

/// Re-exported [`num_complex::Complex`] (and the `c32`/`c64` aliases below), so you
/// do not need to depend on `num-complex` directly. The element type for [`gemm_cplx`]
#[cfg(feature = "complex")]
#[doc(no_inline)]
pub use num_complex::Complex;
/// `Complex<f32>`: the single-precision complex element type
#[cfg(feature = "complex")]
#[allow(non_camel_case_types)]
pub type c32 = num_complex::Complex<f32>;
/// `Complex<f64>`: the double-precision complex element type
#[cfg(feature = "complex")]
#[allow(non_camel_case_types)]
pub type c64 = num_complex::Complex<f64>;