Skip to content

Repository files navigation

简体中文 | English

gemmkit

CI crates.io docs.rs

gemmkit is a pure-Rust workspace for GEMM (general matrix multiply). It computes C <- alpha*A*B + beta*C over strided views or raw pointers. It picks the best available instruction set at runtime.

The core engine works on f32 and f64 by default. Behind Cargo features, it also works on f16/bf16 (mixed precision with f32 accumulation), on i8 -> i32 integers, and on c32/c64 complex data. Runtime ISA dispatch covers x86-64 FMA and AVX-512F (with AVX-512 VNNI for int8 and AVX-512 BF16 for bf16), aarch64 NEON, and wasm32 simd128. A portable scalar fallback backs every one of them, for a target with no vector support. The GEMMKIT_REQUIRE_ISA environment variable pins or forbids a backend.

Multithreading is optional and uses rayon. For a fixed input and a fixed configuration on one machine, results are reproducible. With default features off, the core builds under no_std. It then needs only core and alloc.

Beyond plain GEMM, gemmkit offers:

  • fused epilogues: bias, activation, i8/u8 requantization, and a user per-element map
  • prepacked-operand reuse for fixed-weight inner loops
  • batched GEMM
  • automatic bandwidth-bound paths for matrix-vector products and small shapes

Crates

The workspace ships 5 crates. They share version 0.1.2 and release in lockstep.

Crate Description
gemmkit Core GEMM engine: strided-view and raw-pointer entry points, runtime ISA dispatch, no_std support
gemmkit-ndarray Zero-copy adapter over ndarray matrix views
gemmkit-nalgebra Zero-copy adapter over nalgebra matrix views
gemmkit-faer Zero-copy adapter over faer matrix views
gemmkit-tune Install-time autotuner binary: sweeps the runtime knobs on the target machine and emits a GEMMKIT_* environment profile

The adapters wrap ndarray >= 0.17.1, nalgebra 0.35, and faer 0.24. Each one forwards its parallel, wasm_threads, half, complex, int8, and epilogue feature to the same-named gemmkit feature.

Quick start

[dependencies]
gemmkit = "0.1"
use gemmkit::{gemm, MatMut, MatRef, Parallelism};

fn main() {
    // 2x3 times 3x2 = 2x2, all 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]);
}

Strides express transposition. Use from_col_major, or give an explicit rs/cs in MatRef::new. A transposed operand needs no copy.

Element types and backends

Every element-type family below has a SIMD implementation on every backend, over the scalar fallback that runs anywhere.

Family Feature Accumulator
f32, f64 (built in) same type
f16, bf16 half f32
i8 -> i32 int8 i32
c32, c64 complex same type

Backends, selected at runtime (or pinned with GEMMKIT_REQUIRE_ISA):

  • Scalar: portable fallback, no target features required
  • x86-64 FMA
  • x86-64 AVX-512F, with AVX-512 VNNI (vpdpbusd) for int8 and AVX-512 BF16 (vdpbf16ps) for bf16
  • aarch64 NEON
  • wasm32 simd128 (compile-time feature detection)

The gemmkit Cargo features are:

  • std and parallel, both on by default
  • wasm_threads, for the wasm32-wasip1-threads target
  • complex, half, and int8
  • epilogue: fused bias and activation, i8/u8 requantization, and the user per-element map

With std off, the crate is no_std. parallel implies std.

Documentation

License

Licensed under either of MIT or Apache-2.0 at your option.

Minimum supported Rust version: 1.89 (edition 2024).

About

A clean, extensible, high-performance GEMM (general matrix multiply) engine

Resources

Code of conduct

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages