rlx-cpu 0.2.17

CPU backend for RLX — SIMD kernels, BLAS dispatch, thread pool, arena executor
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Apple matrix-coprocessor fast paths (AMX / SME).
//!
//! On Apple Silicon the fastest CPU matmul is not NEON — it's the matrix
//! coprocessor. Three independent, opt-in routes onto it live here, each behind
//! its own cargo feature, so the base build stays lean and the optimization is
//! something a consumer opts into rather than pays for by default:
//!
//! * `bnns`  (`amx-bnns`)  — low-precision (int8 W8A8) matmul via BNNS. Apple
//!   dispatches BNNS to AMX on M1–M3 and SME on M4+ for us, so we get the
//!   coprocessor for quantized inference without owning per-generation asm.
//! * `dense` (`amx-dense`) — verification + benchmark that Accelerate's sgemm
//!   already *is* the AMX path for dense f32/f64. Nothing to beat; we measure.
//! * `sme`   (`amx-sme`)   — a direct ARM SME2 `FMOPA` GEMM microkernel via
//!   `global_asm!`. The documented M4+ path; hand-written, runtime-gated.
//!
//! Whichever path is compiled in, `detect` provides the runtime probe that
//! decides whether the *current* chip actually has the unit before dispatch.

pub mod detect;

#[cfg(rlx_cpu_amx_bnns)]
pub mod bnns;

#[cfg(rlx_cpu_amx_dense)]
pub mod dense;

#[cfg(rlx_cpu_amx_sme)]
pub mod sme;