Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

简介

gemmkit 是一个纯 Rust 编写的 GEMM 工作区,用于计算通用矩阵乘法 C <- alpha*A*B + beta*C。核心 crate 可以处理带步长的视图或裸指针。它会在运行时选择当前机器上最快的指令集。在固定的输入和配置下,它保证每次运行的结果都可复现。核心之外还有 3 个零拷贝适配器,分别对接 ndarraynalgebrafaer。此外还有一个安装时自动调优器,用来为实际部署的机器校准引擎。

这本书是整个工作区的叙述式文档。docs.rs 上的 API 参考仍然是函数签名和条目级细节的权威来源。本书负责把各个部分放在上下文里讲清楚:设计思路、取舍,以及那些参考页面无法深入展开的 API 细节。

本书内容

gemmkit 使用指南覆盖核心 crate。它从第一次矩阵乘法讲起,一直讲到大多数用户用不到的部分:

  • 矩阵视图与内存布局
  • 可选的元素类型(f16/bf16i8、复数)
  • 并行执行
  • 预打包操作数
  • 融合 epilogue
  • 批量 GEMM
  • 小形状与 GEMV
  • 指令集固定
  • 调优旋钮
  • no_std 与 WebAssembly 构建
  • unchecked 裸指针层

适配器指南展示如何用 ndarraynalgebrafaer 的原生类型直接驱动引擎,全程零拷贝。每个适配器都有自己的章节,包含一篇快速上手和一篇进阶用法。进阶用法覆盖完整的 API 面:融合运算、整数与复数 GEMM、批量运算、预打包,全部使用宿主库自己的类型。

gemmkit-tune 指南讲解自动调优器。它说明如何在部署机器上运行自动调优器、生成的配置文件里有什么,以及背后的扫描过程是怎么工作的。

架构说明逐层拆解引擎内部。它说明一次调用如何从公共 API 一路走到微内核,以及指令集和元素类型如何在不用宏的前提下保持可插拔。它还说明分块参数如何从缓存层级推导出来,以及整个引擎是如何被测试的。它是仓库里那份紧凑的 ARCHITECTURE.md 的详细版和易读版,适合从头到尾通读。

阅读建议

如果你只是想在应用里用上快速的矩阵乘法,从快速上手开始。按需求往后读使用指南即可。

如果你的矩阵本来就存在 ndarraynalgebrafaer 里,直接跳到对应适配器的章节。需要理解底层概念时,再回头查使用指南。适配器转发到同一个引擎,语义完全一致。

如果你好奇引擎内部是怎么工作的,或者打算参与开发,架构说明就是为你准备的路线。它假设你大致翻过使用指南,但不假设你了解 BLIS。它解释的是设计决策本身,而不只是代码。

约定与资源

代码示例基于 Rust edition 2024,工作区 MSRV 为 1.89。需要开启可选 Cargo feature 的示例会在出现处注明。像 gemmkit/src/driver.rs 这样的仓库路径均相对于仓库根目录

相关资源包括:API 参考CHANGELOG,以及 crates.io 上的各个 crate:gemmkitgemmkit-ndarraygemmkit-nalgebragemmkit-faergemmkit-tune

This book is also available in English.

快速上手

gemmkit 在普通 Rust 切片的带步长视图上计算 C <- alpha*A*B + beta*C,并在运行时挑选 CPU 实际支持的最快指令集。你不需要在编译期选择 ISA,也不需要链接任何 BLAS 库。你只需添加 1 个依赖,把 3 个矩阵交给它,然后调用 gemm

添加依赖

核心 crate 是 gemmkit。在普通(std)目标上做 f32/f64 计算时,下面这一行就够了:

[dependencies]
gemmkit = "0.1"

这会引入 2 个默认 feature:stdparallelstd 提供运行时的缓存与 CPU 特性检测,还提供 GEMMKIT_REQUIRE_ISAGEMMKIT_* 调优旋钮,以及一个线程本地的工作区池。这个池让相同尺寸的重复调用不再分配内存。parallel 加上基于 rayon 的多线程,并隐含开启 std

可选的元素类型族(halfcomplexint8)和 epilogue 能力默认关闭。纯浮点构建不会为它们的代码生成或依赖付出任何代价。若要以 no_std(仅 core + alloc)构建 crate,用 default-features = false 关掉默认 feature。参见no_std 与 WebAssembly了解具体做法。

第一个完整示例

这个示例计算一个 2x3 矩阵乘以一个 3x2 矩阵,全部行主序,单线程运行:

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]);
}

参数恰好就是 C <- alpha*A*B + beta*C 的各项:标量 alpha、2 个输入视图、标量 beta、输出视图,以及一个 Parallelism 选择器。

MatRef::from_row_major(&a, 2, 3)a 当作 2 行 3 列的行主序矩阵读取。各维必须对得上:A.cols 必须等于 B.rows,且 C 必须是 A.rowsB.cols 列。如果对不上,调用会在触碰内存之前 panic。

转置从不需要拷贝,因为它只是步长的改变,而非数据的搬移。MatRef::from_col_major(&b, 3, 2) 把同一段缓冲当作列主序矩阵读取。MatRef::new 让你直接设置行步长和列步长。

底层发生了什么

gemm 入口在做任何算术之前,会先做一小段准备工作。

首先它做校验:检查内部维度是否相符、每个视图是否都落在其切片之内、C 是否把每个 (i, j) 映射到互不相同的偏移,以及 C 的存储是否与 AB 的存储重叠。任何一项检查失败都会带着具体的信息 panic,发生在任何一次 unsafe 操作之前。只有到这一步之后,它才把这 3 个视图降低成裸指针与步长,交给分发层。

分发层决定运行哪个内核。对某一元素类型的第一次 GEMM 调用会执行一次 CPU 特性检测,把胜出的入口点记录进一个 OnceLock,然后返回。此后每次调用都只是经由那个缓存指针的一次普通间接调用,不再重复检测。因此运行时 ISA 选择是一次性成本,摊薄到整个进程。

你可以用 GEMMKIT_REQUIRE_ISA 环境变量覆盖自动选择,也可以用它为测试锁定某个后端。gemmkit 只读取该变量一次,并以同样的方式记忆结果。参见运行时 ISA 分发了解细节。一次 GEMM 调用的生命周期走了一遍从调用到微内核的完整路径。

精确理解 alpha 与 beta

alpha 缩放乘积 A*Bbeta 缩放 C 的原有内容。唯一值得记牢的细节是边界情形下会发生什么。

beta == 0 时,引擎完全不读取 C,而是直接用 alpha*A*B 覆写它。正是这条规则,让上面的 let mut c = [0.0_f32; 4] 即便缓冲原本未初始化也依然正确。具体来说,beta == 0 的输出切片可以持有垃圾值。经由 unchecked 层,它甚至可以是真正未初始化的内存,结果仍然是良定义的。

beta == 1 时,引擎让 C 的原有内容保持不动,并把乘积累加到它上面。其他任何 beta 值都会先把 C 乘一遍。

此外还有一条退化的快速路径。若 k == 0(空收缩)或 alpha == 0(乘积消失),调用就退化成 C <- beta*C。它完全不触碰 AB,只在原地缩放输出。结合 beta == 0 的规则,alpha == 0, beta == 0 会把 C 清零,而 k == 0, beta == 1 是一次空操作。窄类型在 f32 中缩放,并在写回时舍入,因此退化路径的舍入与完整内核完全一致。

Cargo feature

Feature默认解锁引入依赖
std运行时缓存/CPU 检测、环境旋钮、线程本地工作区池。关闭即 no_stdcore + allocraw-cpuid(仅 x86)
parallel基于 rayon 的多线程(Parallelism::Rayon)。隐含 stdrayon
wasm_threadswasm32-wasip1-threads 提供定尺寸的 rayon 池。隐含 parallel(经由 parallel
halff16/bf16 混合精度 GEMM,f32 累加half
complex带共轭的 c32/c64 GEMM(gemm_cplxnum-complex
int8i8 -> i32 整数 GEMM(gemm_i8(无)
epilogue融合偏置/激活、i8/u8 重量化、逐元素映射(无)

元素类型 feature 与能力 feature 可以组合:half + epilogue 给出融合的 f16 GEMM,int8 + epilogue 给出重量化入口,以此类推。每种组合的详情见元素类型融合 Epilogue

版本要求

gemmkit 面向 Rust 1.89edition 2024。采用 MIT OR Apache-2.0 双许可。API 参考见 docs.rs/gemmkit,本书是那份参考的长文版伴读。

接下来读什么

矩阵视图与内存布局

gemmkit 的每一次调用都以视图的形式接收操作数:一个切片、一个形状,加上 2 个步长。MatRef<'a, T> 是不可变的输入视图,MatMut<'a, T> 是可变的输出视图。两者都不拥有自己的数据,都只是借用你已经持有的切片。

库里整套布局词汇都体现在这 2 个步长数字上:行主序、列主序、转置、子矩阵、广播。因此同一段缓冲可以被十几种方式读取,而从不发生拷贝。

两个步长

视图中元素 (i, j) 位于切片偏移 i*rs + j*cs 处,其中 rs 是行步长,cs 是列步长。步长以元素计,而非字节。 rs 为 4 表示下一行在切片中再往后 4 个元素。那一条偏移公式就是整个模型,其余一切都只是 rscs 的取值选择。

3 个构造器覆盖了常见情形,MatRefMatMut 上都有:

#![allow(unused)]
fn main() {
use gemmkit::MatRef;

let data = [0.0_f32; 12];
let row_major = MatRef::from_row_major(&data, 3, 4); // rs = cols = 4, cs = 1
let col_major = MatRef::from_col_major(&data, 3, 4); // rs = 1, cs = rows = 3
let general   = MatRef::new(&data, 3, 4, 4, 1);      // explicit rs, cs (here == row-major)
}

from_row_major(data, rows, cols) 设定 rs = cols, cs = 1:行是连续的,即经典的 C 序。from_col_major(data, rows, cols) 设定 rs = 1, cs = rows:列是连续的,即 Fortran 序。

new(data, rows, cols, rs, cs) 原样接收步长。当两种规范布局都不匹配时就用它,比如一个子矩阵,或者前导维(leading dimension)不等于逻辑宽度的视图。MatRefMatMut 还提供 .rows().cols()

转置就是交换步长

由于 (i, j)i*rs + j*cs 映射,交换两个步长(以及两个维度)的角色,就在原地转置了视图。设 a 以行主序持有一个 m x k 矩阵(rs = k, cs = 1)。它的转置是那个 k x m 矩阵,其 (i, j) 即原矩阵的 (j, i),位于偏移 j*k + i。这个偏移恰好就是同一切片上的 rs = 1, cs = k

#![allow(unused)]
fn main() {
use gemmkit::MatRef;

// `a` is m x k row-major
let (m, k) = (2, 3);
let a = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];

let a_rowmajor = MatRef::from_row_major(&a, m, k); // m x k
let a_transposed = MatRef::from_col_major(&a, k, m); // k x m, same bytes, no copy
}

于是在行主序缓冲上的 from_col_major转置,反过来也一样成立。把 rs/cs 对调的 new 对任意布局做同样的事。因此转置操作数在 API 层面毫无代价:内核就按你给出的步长去走。这正是无需实体化转置就能喂入 A^T * BA * B^T 的方式。

子矩阵与带步长视图

子矩阵是这样一种视图:它的前导维(相邻行或列之间的距离)大于其逻辑范围。构造方法是对缓冲做切片,让子块的左上角元素落在切片起点,再交出父矩阵的步长。下面取一个 4 x 4 行主序矩阵中、从第 1 行第 1 列开始的左上 2 x 2 块:

#![allow(unused)]
fn main() {
use gemmkit::MatRef;

let parent = [0.0_f32; 16]; // 4x4 row-major, leading dimension 4
let block = MatRef::new(&parent[1 * 4 + 1..], 2, 2, 4, 1); // rs stays 4, cs stays 1
}

行步长仍是 4,也就是父矩阵的宽度。于是子块的相邻行会跳过你排除掉的那些列。切片从偏移 5 开始,也就是子块的 (0, 0)。安全 API 会核验这段尾切片足够长,能触及子块的远端角落。

同一机制也能表达广播输入:步长为 0 会让某一维重复同一个元素。一个 1 x n 的行向下广播 m 行,就是 MatRef::new(row, m, n, 0, 1)。每一逻辑行都读取同一段存储。gemmkit 允许只读输入 AB 使用广播,但对输出永不允许,下一节会解释原因。

安全 API 接受什么、拒绝什么

安全入口(gemmgemm_i8gemm_cplx,以及各融合变体)只接受非负步长,其中包括广播输入用的 0。负步长超出了 &[T] 视图所能安全描述的范围。一个落在缓冲内部、而非指向元素 (0, 0) 的基指针同样如此。这些情形都属于Unchecked 层,也就是适配器用来表达任意布局的裸指针引擎。

在任何算术之前,安全入口对 (A, B, C) 三元组跑一段校验前奏,每一处失败都是一次 panic,抛出在首个 unsafe 操作之前:

  • 形状相符。 A.cols == B.rowsA.rows == C.rowsB.cols == C.cols。不符会带着出错的那一对 panic,例如 gemmkit: A.cols (3) != B.rows (4)
  • 视图在界内。 对每个视图,引擎会计算它将触及的最高切片偏移,并与切片长度比对。切片太小会 panic:gemmkit: A view of 3x4 (strides 4,1) needs 12 elements but slice has 8。若某视图步长为负,或大到寻址溢出 usize,则会 panic:... has negative strides or is too large to address; use gemm_unchecked
  • C 唯一寻址每个元素。 gemmkit 要写入输出,因此 2 个不同的 (i, j) 绝不能落在同一偏移。一个自别名的 C,比如零行或零列步长,或者相互碰撞的步长,在并行模式下会变成数据竞争。这种情形从完全安全的代码即可触及,因此会 panic:gemmkit: C view aliases itself (...); C must address each (i,j) uniquely。正因如此,广播步长对 A/B(只读)没问题,但对 C 禁止。
  • C 不与 AB 重叠。 输出的字节范围必须与每个输入的不相交。gemmkit 按字节范围比较,而非按元素个数,这样异构整数 API(i8 输入、i32 输出)也能精确无误。重叠会 panic:gemmkit: C aliases A or B。在纯安全 Rust 里,借用检查器本就禁止重叠的 &mut/& 对,这道检查是同时覆盖裸指针降级路径的防御性兜底。

这些信息是稳定的,正确性测试套件会断言它们的措辞,因此你可以在测试中依赖它们。

零尺寸维度

带有零维度的视图是合法的,并且能顺利通过校验。0 x km x 0m x n x (k = 0) 这些形状,gemmkit 都接受。任何切片,哪怕是空的,都能满足在界检查,因为这样的视图不寻址任何东西。

m == 0n == 0,调用是一次空操作:没有输出要写。若仅 k == 0,收缩为空,调用退化为 C <- beta*C。这也就是 alpha == 0 所走的那条只缩放路径。那条退化规则见快速上手

接下来读什么

  • 元素类型:把同样的视图用在 f16/bf16i8 与复数数据上。
  • Unchecked 层:负步长、内部基指针,以及裸指针引擎。
  • 适配器(ndarraynalgebrafaer)会从各库的原生矩阵类型替你构造这些视图。

元素类型

gemmkit 相乘的不止是 f32。同一套引擎、驱动与分块模型服务于 4 个元素类型族。每一个都是一个 Cargo feature,每一个也都在每个后端上有 SIMD 实现,并以可移植的标量回退路径兜底。

各族之间变化的是输入类型、累加器类型和输出类型,随之而来的还有你应当预期的精度变化。本页是一张地图:有哪些可用,以及各自有多精确。

内置的实数浮点

f32f64 不需要任何 feature 开关。它们经由通用的 gemm(以及 gemm_with 和各 unchecked 入口),以自身类型累加,是其余每个族赖以对照的基准:

#![allow(unused)]
fn main() {
use gemmkit::{gemm, MatMut, MatRef, Parallelism};

let a = [1.0_f64, 2.0, 3.0, 4.0];
let b = [5.0_f64, 6.0, 7.0, 8.0];
let mut c = [0.0_f64; 4];
gemm(2.0, MatRef::from_row_major(&a, 2, 2), MatRef::from_row_major(&b, 2, 2),
     0.0, MatMut::from_row_major(&mut c, 2, 2), Parallelism::Serial);
}

精度遵循教科书里的 GEMM 故事:相对误差大致随收缩深度 k 和该类型的机器 epsilon 增长。正确性测试套件把结果约束在相对 Frobenius 门限 8*k*eps 之内,对照一个独立的 f64 参考核验。因此 f64 对任何现实的 k 都近乎精确,而 f32 保有其惯常的每元素约 1e-7 相对精度。

窄浮点:half feature

开启 half 后,f16bf16 成为元素类型。gemmkit 把它们重导出为 gemmkit::f16gemmkit::bf16,因此你无需直接依赖 half。它们共用通用的 gemm 表面:输入 MatRef<'_, f16>,输出 MatMut<'_, f16>,因为它们实现了与实数浮点相同的标量 trait。

其决定性特征是混合精度。引擎在加载时把输入加宽到 f32。整个收缩都在 f32 中累加。引擎只在写回时把结果舍入回窄类型一次。 k 循环内部没有反复的窄舍入,正是这唯一一次舍入让精度可用。

#![allow(unused)]
fn main() {
use gemmkit::{f16, gemm, MatMut, MatRef, Parallelism};

let a: Vec<f16> = (0..6).map(|i| f16::from_f32(i as f32)).collect();
let b: Vec<f16> = (0..6).map(|i| f16::from_f32(i as f32)).collect();
let mut c = vec![f16::ZERO; 4];
gemm(f16::ONE, MatRef::from_row_major(&a, 2, 3), MatRef::from_row_major(&b, 3, 2),
     f16::ZERO, MatMut::from_row_major(&mut c, 2, 2), Parallelism::Serial);
}

由于累加在 f32 中进行,主导误差是那唯一一次最终舍入,而非求和本身。f16 保有约 9.8e-4(2^-10)的相对精度,bf167.8e-3(2^-7),两者基本都与 k 无关。因此窄精度 GEMM 接近于用 f32 计算、只舍入一次,远比在 16 位中累加要精确。

只舍入一次带来一个后果:在大 k 时,单个深度面板会流式吞吐一个大到 L2 缓存放不下的中间结果。引擎会自行处理这个问题。越过一个自动导出的字节门限后,它会切换到一个 f32 输出的内部孪生体。这个孪生体把收缩重新分块以保持驻留缓存,并在最后才把结果舍入回窄类型。对常见的 beta in {0, 1} 情形,孪生体与单面板逐字节相同,其他情形则落在容差之内。这个切换是自动的,无需任何配置。机制细节见点积内核与深 K 孪生。在 AVX-512 BF16 硬件上,bf16 还会使用 vdpbf16ps 点积内核,参见运行时 ISA 分发了解该机制。

整数:int8 feature

int8 加入 gemm_i8,一个独立的入口。它的输入与输出类型不同,i8 进、i32 出,同构的 gemm<T> 表面无法表达这一点。alphabetaC 全都是 i32

#![allow(unused)]
fn main() {
use gemmkit::{gemm_i8, MatMut, MatRef, Parallelism};

let a = [1_i8, 2, 3, 4, 5, 6];
let b = [7_i8, 8, 9, 10, 11, 12];
let mut c = [0_i32; 4];
gemm_i8(1, MatRef::from_row_major(&a, 2, 3), MatRef::from_row_major(&b, 3, 2),
        0, MatMut::from_row_major(&mut c, 2, 2), Parallelism::Serial);
}

整数 GEMM 是精确的:它是溢出时回绕i32 环算术,即整数 GEMM 的惯例语义。谈不上什么容差,因为根本没有舍入。结果在每个 ISA 上逐位相同(标量、FMA、AVX-512F,以及 AVX-512 VNNI 的 vpdpbusd 点积内核)。串行与并行运行之间也逐位相同,因为环上的整数加法与顺序无关。

若你喂入的乘积会超出 i32,回绕是有定义且可复现的,而非未定义行为。int8 feature 不引入任何额外依赖。在其之上再加 epilogue 可解锁重量化入口,一趟即可得到 i8u8 输出。参见融合 Epilogue了解这些入口。

复数:complex feature

complexnum-complex 值上加入 gemm_cplx。gemmkit 把它们重导出为 gemmkit::c32Complex<f32>)和 gemmkit::c64Complex<f64>)。它的签名为每个操作数带一个共轭标志:

#![allow(unused)]
fn main() {
use gemmkit::{c32, gemm_cplx, Complex, MatMut, MatRef, Parallelism};

let a = [Complex::new(1.0_f32, 1.0), Complex::new(2.0, 0.0)];
let b = [Complex::new(0.0_f32, 1.0), Complex::new(1.0, 0.0)];
let mut c = [c32::default(); 1];
gemm_cplx(
    Complex::new(1.0, 0.0),
    MatRef::from_row_major(&a, 1, 2), false, // conj_a
    MatRef::from_row_major(&b, 2, 1), false, // conj_b
    Complex::new(0.0, 0.0),
    MatMut::from_row_major(&mut c, 1, 1),
    Parallelism::Serial,
);
}

计算为 C <- alpha*op(A)*op(B) + beta*C:当 conj_a 置位时 op(A)conj(A)conj_b 置位时 op(B) 同理为 conj(B)。传入 false, false 就得到普通乘积 A*B。这些标志只共轭操作数本身。

复数以自身类型累加,被约束在相对 Frobenius 门限 16*k*eps 之内,其中 eps 取实部的 epsilon。因此 c32 GEMM 的精度约与 f32 相当,c64 GEMM 约与 f64 相当。complex 引入 num-complex

在内部,复数并不搭乘浮点内核,而是使用一个专用的拆分(结构体数组,SoA)内核,这也是它成为独立入口的原因。该设计见复数拆分内核

如何选择类型

类型族Feature输入 / 累加 / 输出精度确定性
f32, f64(内置)同 / 同 / 同教科书级,约 8*k*eps可复现,如今在驱动路径上串行与并行逐位相同
f16, bf16half窄 / f32 / 窄只一次最终舍入,约 1e-3(f16)、8e-3(bf16)可复现,深 K 孪生对 beta in {0,1} 逐位相同
i8int8i8 / i32 / i32精确,回绕 i32跨每个 ISA 与每个工作线程数逐位相同
c32, c64complex同 / 同 / 同16*k*eps可复现,如今串行与并行逐位相同

若你需要速度、又能容忍约 1e-3 的精度,选 bf16f16。二者每个元素占用的字节都只有 f32 的一半,同时仍在 f32 中累加。若你需要精确,int8 给你精确。若你需要动态范围与精度,就留在 f32/f64

它们的可复现契约都是同一份。“可复现”承诺什么、不承诺什么,见并行实践

接下来读什么

并行实践

每个 GEMM 入口都把 Parallelism 作为最后一个参数。它是一个只有三种实用模式的小枚举。用好它其实只在于两件事:弄清楚 auto 模式替你做了哪些决定,以及知道什么时候该自己接管。

三种模式

#![allow(unused)]
fn main() {
pub enum Parallelism {
    Serial,       // single-threaded
    Rayon(usize), // rayon with at most n threads; Rayon(0) auto-detects
}
}

Serial 把整次调用跑在调用线程上。Rayon(n) 请求至多 n 个工作线程。Rayon(0) 是自动模式,也是 Default,因此 Parallelism::default() 就是自动模式。Rayon(n) 里的 n 是分区数的上限,不是把它们全部用满的承诺。工作量不足 n 个块、或核心数少于 n 的问题,会分得更少。

自动模式究竟做了什么

自动模式不等于“用满所有核心”,而是依据问题规模做出两个决定。

首先是一道工作量闸门。当 m*n*k 总工作量低于某个阈值(GEMMKIT_PARALLEL_THRESHOLD 旋钮,默认 48*48*256)时,调用无论如何都保持串行。在这么小的矩阵上,fork/join 的开销会淹没任何收益。这道闸门先于其余一切执行,因此哪怕显式传入 Rayon(n) 也逃不过它:低于闸门时,Rayon(8) 仍然跑在单线程上。

越过闸门后,自动模式会让工作线程数随总工作量扩展,而不是直接跳到满核心数。它以 m*n*k 除以 GEMMKIT_PAR_MNK_PER_WORKER(默认 2_000_000,即每这么多工作量配一个工作线程)为目标。然后它用机器的核心数和可用作业块数给结果封顶,下限为一。这个数量按工作量而非维度计算,因为最优工作线程数跟随的是总 flops,而不是线性尺寸,线性维度上的任何单一步长都覆盖不了这么大的跨度。小乘积只用少数几个工作线程,大乘积则用很多个。把 GEMMKIT_PAR_MNK_PER_WORKER 设为 0(其行为等同于 1)会让任何越过串行闸门的问题都强制用满全部宽度。

显式数量

n > 0Rayon(n) 会绕开爬升启发式,直接请求恰好 n 个分区。出于安全考虑,这仍然会被机器的核心数(available_parallelism)和问题实际切分出的作业块数封顶。因此在一台 16 核机器上计算一个小乘积时,Rayon(1000) 不会导致过度订阅,它会坍缩到机器与工作量所能吸收的程度。正是这种精确性,让测试套件与扩展诊断都使用显式数量:Rayon(4) 在有那么多工作量、那么多核心时会给你四路分区,而不是一个启发式猜测。当你已经测量过自己的工作负载、知道甜点区在哪里时,就用显式数量。想要跨运行可复现的分区以做基准测试时,也用它。

gemmkit 如何使用 rayon 池

gemmkit 不需要你交给它一个 rayon 池。如果你把一次调用包在自己那个池的 install 里,这次 GEMM 的工作线程就会跑在那个池上,不会去别处。

#![allow(unused)]
fn main() {
let pool = rayon::ThreadPoolBuilder::new().num_threads(4).build().unwrap();
pool.install(|| {
    gemm(1.0, a, b, 0.0, c, Parallelism::Rayon(0)); // runs on `pool`
});
}

gemmkit 选择的工作线程数量仍然受 available_parallelism(整台机器)约束。rayon 的工作窃取调度器会把这些分区分布到当前池所拥有的线程上,所以一个更小的自定义池只是把同样的分区跑在更少的线程上。一次调用内部的工作分配是按需驱动的:工作线程从一个共享的无锁游标上领取连续的块。在一个异构部件上(比如混有 P 核与 E 核的布局),更快的核心会比别的核心多领一些块,而不是所有人都等最慢的那个。

如果一次调用没有跑在你自己安装的池里,gemmkit 会转而使用它自己的某个池。这一行为在原生目标上默认开启,也就是 x86_64 与 aarch64。gemmkit 最多会保留 GEMMKIT_POOL_CLASSES(x86_64 默认 2,aarch64 默认 1)个私有的、持久化的池,每一个都恰好卡在机器宽度的某个减半档位上。在一台 32 线程的机器上,那就是 16 线程和 8 线程两档。在一台 14 核的 M4 Max 上,则是单独一个 7 宽的档位。每个池都在首次用到时才惰性构建,此后不再重建。

自动模式会把它的工作线程数精确卡在这些档位之一,而不是去 fork 那个满宽度的全局池。一次 fork 的开销跟的是池的空闲余量:也就是池所拥有、却没有实际参与工作的那些线程。一个小规模的 GEMM 一旦落进满宽度的池,就会被这份余量拖累,所以把池的尺寸卡准工作量能避免这个拖累。

以上这些都不会改变你已经知道的东西。一次 install 过的调用依然被完全尊重,绝不会被改道到某个档位池。显式的 Rayon(n) 仍然恰好拿到 n 个工作线程,只是会被路由进能容纳它的最小档位池。真正会变的是闲置内存:默认情况下,一个 x86_64 进程现在会额外常驻大约 24 个线程(16 宽与 8 宽这两个档位池),一个 aarch64 M4 Max 进程会常驻 7 个(它唯一的半宽档位)。这些线程都在休眠,直到某个小规模 GEMM 需要它们。把 GEMMKIT_POOL_CLASSES 设为 0 即可彻底关闭档位池,此后每次调用都回落到环境池。

线程化 wasm 的情况有所不同:那里 gemmkit 总是会为自己定制一个专用池。这种情况见no_std 与 WebAssembly

精确表述可复现承诺

对固定的输入、环境与配置,输出与工作线程数无关地保持一致。这就是这份契约。

它成立有两个原因。第一,kcnc 与固定的深度面板顺序,是唯一塑造每个输出元素求和方式的东西,而 gemmkit 计算这三者时与将由多少线程来运行它们无关。第二,每个输出元素都由单个工作线程沿完整的收缩深度从头归约到尾,因此不存在任何归约拆分,其顺序也就谈不上依赖调度。

扁平作业列表本身在不同工作线程数下并非严格一致:更宽的工作线程数可能会缩小 mc,以保持列表足够深。但 mc 始终是 mr 的倍数,所以微 tile 的集合及其数值都不变。打包出的字节同样不取决于是谁打包的。哪个工作线程计算某个 tile 会因运行而异,但数值结果不会。

被承诺的是 SerialRayon(n) 之间的逐位相同。这一点今天在驱动路径上恰好成立,因为串行与并行跑的是同一个内核。但你应当依据的保证是固定配置下的可复现性,而不是串行对并行的逐位相等。你在这里得不到跨机器或跨配置的逐位相等:浮点 GEMM 对顺序敏感,而配置(ISA、分块、线程上限)正是固定输入的一部分。整数 gemm_i8 是例外:它跨 ISA、跨工作线程数逐位相同,因为 i32 加法与顺序无关。

何时该选 Serial

在三种情形下伸手去用 Serial

  1. 小问题。 低于工作量闸门时,自动模式本来就是串行的。显式传入 Serial 还会彻底跳过 available_parallelism 探测和 fork 机制,这在一串微小 GEMM 的紧循环里更省成本。
  2. 当外层并行由你掌控时。 假设你已经在一个 rayon 池上运行许多独立的 GEMM,或者自己在并行化一个批处理循环。不要再让每次内层调用也跟着扇出,那会过度订阅机器,通常还会拖累性能,而不是帮上忙。让内层调用跑 Serial,把并行度留在外层。对一批乘积,优先用内置的批量 GEMM 入口,它把整批当作一个整体来调度。
  3. 对确定性敏感的调试。 用单线程路径排除掉调度这个变量。

带宽受限的形状有自己的策略

矩阵-向量乘(m == 1n == 1)以及其他内存受限的形状不是计算受限的,所以上面那套按工作量算的工作线程数对它们是错的模型,因此这些路径用一条独立的规则。

在一个由缓存导出的字节下限以下,矩阵装得进单核的私有缓存。该核已经独自打满了缓存,所以拆分只会增加争用,这个路径也就保持串行。越过这个下限后,路径会直接跨到一个按触碰字节选出的宽度。这个宽度的顶端是逻辑核心数的一半,因为 gemv 远在最后一个核心加入之前就已经打满了自身带宽。

这个宽度是一级一级往上跳的,不是平滑增长:每一级都是上文那些精确适配的线程池之一,所以带宽受限的调用也能拿到一个尺寸刚好合适的池。少数几个工作线程是带宽扩展曲线上最糟的一点,所以这个策略会直接跳过那一点,而不是从中爬过去。

整套策略都是自动的。字节下限、级间距,以及一个固定宽度的覆盖值都可以通过 GEMMKIT_GEMV_PARALLEL_BYTESGEMMKIT_GEMV_TIER_STEPGEMMKIT_GEMV_THREAD_CAP 来调。完整论述见小形状与 GEMV

接下来读什么

预打包操作数

在微内核触碰 AB 之前,引擎会先把它们各自拷贝成一种对缓存友好的微面板(micropanel)布局。这种布局由一块块连续的 tile 组成,微内核以单位步长逐块扫过。对于一次性的乘积,这次拷贝纯粹是准备工作,引擎付出一次代价之后就再也用不上它了。

很多工作负载会反复乘同一个矩阵。一个线性层会用同一个固定的权重矩阵去乘一串激活批次,一个求解器会用同一个固定的算子去乘许多右端项。如果每次调用都重新打包这个固定操作数,就等于把引擎已经做过的工作白白扔掉。预打包操作数 API 让你把打包的代价只付一次,然后在所有共享该操作数的乘积之间复用这个结果。

打包右操作数

最常见的情形是固定 B(权重),再配上一串大小各异的 A(激活)。调用一次 prepack_rhs,把 k x nB 变成一个 PackedRhs 句柄,之后每次乘积都把这个句柄交给 gemm_packed_b

#![allow(unused)]
fn main() {
use gemmkit::{prepack_rhs, gemm_packed_b, MatRef, MatMut, Parallelism};

// 固定权重:一个 k x n 矩阵,被许多激活批次复用
let (k, n) = (512, 256);
let weights = vec![0.0f32; k * n];
let packed = prepack_rhs(MatRef::from_col_major(&weights, k, n));

// 每个激活批次:一个 m x k 输入,共享已打包的权重
let mut c = vec![0.0f32; m * n];
gemm_packed_b(
    1.0,
    MatRef::from_row_major(&input, m, k),
    &packed,
    0.0,
    MatMut::from_col_major(&mut c, m, n),
    Parallelism::Rayon(0),
);
}

prepack_rhs 接受任意布局的 B,并通过其步长来读取它。行主序、列主序或转置视图打包出来的结果都一样。打包只在 prepack_rhs 内部单线程地跑一次,之后每次调用 gemm_packed_b 都会跳过它。

这块缓冲区记录了它构建时所用的分块几何:nrkcnc。之后每次消费调用都会原样读回这套几何,所以一个面板永远是按它自己的 tiling 来解读的。该缓冲区在整个 GEMM 期间只读,引擎在打包之后不会再写它。因此一个 PackedRhs 可以安全地跨线程、跨并发调用共享,无需任何额外同步。PackedRhs::rows() 返回原始的 kPackedRhs::cols() 返回原始的 n

这个句柄对任何 (k, n) 与所打包 B 相符的乘积都有效,前提是输出 C 是列主序倾向的(|csc| >= |rsc|)。这条约束是这套 API 里唯一的意外之处。一个行主序的 C 会迫使引擎在内部交换 AB,以保持存储连续。而预打包的 B 无法挪到 A 的角色上,所以 gemm_packed_b 遇到行主序的 C 会 panic,并提示你对那种布局改用普通的 gemm。被这样钉住的只有 CA 的布局不受约束。

在固定配置下,gemm_packed_b 会复现普通的 gemm,且结果对 worker 数是确定的。这里有一个很窄的例外。它出现在小乘积上,即 mn 都在 small_mn_dim 旋钮(默认 16,aarch64 上为 32)及以下的情形。它也出现在 gemv 形状的乘积上,即 m == 1n == 1 的情形。在这两种情形下,两次调用的结果可能在最后一个 ULP 上不同。原因在于路由,而不是误差:普通 gemm 会把这些形状改道到某条特殊路径,而预打包入口始终驱动通用的打包内核。两个答案都是正确的,只是恰好在那些本会被特殊路径接管的形状上,求和的次序略有不同。

左操作数的镜像

对称的情形是固定 A,再配上一串变化的 Bprepack_lhs 产生一个 PackedLhs 句柄,由 gemm_packed_a 消费。它与 RHS 那一对完全镜像,只是把轴换了标签:PackedLhs::rows() 是原始的 mPackedLhs::cols() 是共享的 k

在内部,LHS 打包并不是一条独立的代码路径。由引擎的 A/B 对称性可知,一个预打包的 A 恰好就是转置乘积 C^T = B^T A^T 的预打包 B。所以 prepack_lhs 铺下的是完全相同的微面板缓冲区,只是把维度记成了 LHS 的说法。

这带来一个可见的后果:C 的布局约束翻转了。gemm_packed_a 要求行主序倾向的 C|csc| <= |rsc|),与 RHS 入口正好相反。列主序的 C 会让 A 保持在真正的 LHS 角色上,而一块转置 RHS 缓冲区填不了这个角色。C 是行主序时选打包 A 的入口,C 是列主序时选打包 B 的入口。两者合起来覆盖了这两种取向。

融合变体

每个打包入口都有一个融合孪生,位于 epilogue feature 之后。gemm_packed_b_fusedgemm_packed_a_fused 会在打包内核本就要执行的那次存储里,加上逐行或逐列的偏置,外加一个可选的激活。偏置与激活的类型见融合 Epilogue

同一个 PackedRhsPackedLhs 句柄同时服务普通入口和融合入口。epilogue 只作用于存储那一步,从不触碰打包,所以你只需构建一次缓冲区,每次调用再决定是否融合。

有两点是打包路径特有的。其一,与普通 gemm_fused 不同,打包融合入口从不改道到 gemv、小 m,n 或小 k 内核,它们始终驱动通用打包内核,这与普通打包入口所记录的分歧一致。其二,gemmkit 始终以自然的用户坐标系给出逐行或逐列偏置:gemm_packed_a_fused 替你处理了内部转置,所以无论调用哪个入口,PerRow 偏置的长度都是 A.rows

预打包 i8 权重

int8 feature 之下,同样的模式延伸到量化推理。prepack_rhs_i8 把固定的 i8 权重矩阵打包成 PackedRhs<i8>,由 gemm_i8_packed_b 消费。它接受 i8 输入,产生 i32 输出。

预打包在这里比对浮点收益更大,原因是结构性的。AVX-512 VNNI 内核(vpdpbusd)从一种 k-四元组交错布局中读取它的 RHS。引擎无法就地生成这种布局,因此这个内核的 RHS 打包在每次调用时都是强制的。在小 m 下,这个逐调用的 O(k*n) 打包很容易压过 O(m*k*n) 的计算量。预打包把它彻底从热循环里删掉。

这块打包缓冲区还钉住了内核的选择:它是按进程分发所选定的那个整数内核铺设的,要么是 VNNI 交错布局,要么是加宽内核的普通面板。gemm_i8_packed_b 始终跑同一个家族,所以缓冲区绝不会被错读。

整数累加是精确的,且不依赖 ISA。因此对每个合法形状,打包路径与普通路径都逐位一致,完全没有小形状的例外。

#![allow(unused)]
fn main() {
use gemmkit::{prepack_rhs_i8, gemm_i8_packed_b, MatRef, MatMut, Parallelism};

let packed = prepack_rhs_i8(MatRef::from_col_major(&weights_i8, k, n));
let mut c = vec![0i32; m * n];
gemm_i8_packed_b(
    1,
    MatRef::from_row_major(&input_i8, m, k),
    &packed,
    0,
    MatMut::from_col_major(&mut c, m, n),
    Parallelism::Rayon(0),
);
}

什么时候值得预打包

预打包用一次前置的 O(k*n) 拷贝,换取此后针对该操作数的每次乘积都省下一次重打包。只有当这个操作数被复用得足够多、能摊销这次拷贝时,它才划算。对于一次性的乘积,或者每次调用都变化的操作数,预打包一无所获,此时打包纯属额外开销,普通 gemm 才是正确的选择。

还要留意,浮点路径本就不一定会先打包 B。在小 m 下,普通 gemm 会就地读取未打包的 B,这个选择由 rhs_pack_threshold 旋钮控制。所以给一个复用不多的浮点 B 做预打包,甚至可能反而更慢。

最清晰的收益出现在这套 API 得名于的固定权重推理循环上,尤其是 i8 VNNI 路径,它的 RHS 打包本来每一次调用都躲不掉。拿不准的时候,就把循环的两种写法都测一测。交叉点取决于你的复用次数和你的机器,而不是某条固定规则。

裸指针形式(prepack_rhs_uncheckedgemm_packed_b_unchecked,以及它们的 _with、LHS 与 i8 对应版本)是为自行校验输入的适配器和 FFI 而设的。见 Unchecked 层

融合Epilogue

GEMM 很少独自出场。它的输出通常紧接着就要做一次偏置加法、一个激活,或一步量化。如果按朴素的写法来实现,这就意味着要对 C 再扫描一遍:GEMM 先写出 m*n 个值,然后一个单独的循环把它们全部读回来、做变换、再写回去。

融合 epilogue 把这个变换直接折进 GEMM 自己的存储步骤里。每个输出元素在写出的那一刻就在寄存器里完成了变换,那趟额外的内存扫描也就不存在了。本页的所有内容都位于 epilogue 这个 Cargo feature 之后。

偏置与激活

gemm_fused 是向量化的主力入口,一趟就算出 C <- act(alpha*A*B + beta*C + bias)。偏置是一个 Bias 枚举:要么是 Bias::PerRow(&[T])(每个输出行一个值,长度 m),要么是 Bias::PerCol(&[T])(每列一个值,长度 n)。gemmkit 会在乘积算出之后,把这个值加到对应行或列的每个元素上。激活是一个 ActivationRelumax(v, 0))或 LeakyRelu(slope)。这两个参数都是 Option,两者都传 None 就直接委托给普通的 gemm

#![allow(unused)]
fn main() {
use gemmkit::{gemm_fused, Bias, Activation, MatRef, MatMut, Parallelism};

let bias = vec![0.0f32; m]; // 每个输出行一个值
gemm_fused(
    1.0,
    MatRef::from_row_major(&a, m, k),
    MatRef::from_col_major(&b, k, n),
    0.0,
    MatMut::from_col_major(&mut c, m, n),
    Some(Bias::PerRow(&bias)),
    Some(Activation::Relu),
    Parallelism::Rayon(0),
);
}

偏置、LeakyRelu 斜率和激活都在向量快路径上于寄存器内直接施加,所以这次融合相比裸 GEMM 几乎不多花代价。

任意的逐元素映射

当想要的变换既不是偏置也不是标准激活时,就用 gemm_map。它接受一个闭包 f(value, row, col) -> value,把它施加到每个输出元素的最终值上,恰好一次,并融合进存储那一步。它是 gemmkit 没有内置快路径的那些 epilogue 的通用扩展点,比如 GELU、sigmoid、clamp,或任何与位置相关的变换:

#![allow(unused)]
fn main() {
use gemmkit::{gemm_map, MatRef, MatMut, Parallelism};

let f = |v: f32, _r: usize, _c: usize| v.tanh();
gemm_map(
    1.0,
    MatRef::from_row_major(&a, m, k),
    MatRef::from_col_major(&b, k, n),
    0.0,
    MatMut::from_col_major(&mut c, m, n),
    &f,
    Parallelism::Rayon(0),
);
}

交给闭包的 (row, col)C 的用户坐标系。闭包可以按引用捕获它的环境,约束条件 + Sync 正是为了让这个引用能安全地在并行 worker 之间共享,比如借用一张查找表。gemm_map 只支持 f32/f64。它用每个输出元素一次间接调用换来完全的通用性,相对每个元素 O(k) 的工作量而言,这次间接调用很便宜。如果只是普通的偏置或激活,优先选择会把变换向量化的 gemm_fused

整数重量化

量化推理想要的恰是加宽 GEMM 的反面。它接受 i8 输入,累加进 i32,再把结果重新变回 i8(或 u8)输出,途中还要施加一个 scale 和一个 zero-point。gemm_i8_requantgemm_i8_requant_u8 一趟就做完整件事,省掉了单独一次 gemm_i8 调用再接一步重量化所需要的、对完整 m*ni32 的物化。这两个入口都接受一个 Requantize 结构体:

#![allow(unused)]
fn main() {
use gemmkit::{gemm_i8_requant_u8, Requantize, RequantScale, MatRef, MatMut, Parallelism};

let req = Requantize {
    scale: RequantScale::PerRow(&per_channel_scales), // 长度 m,逐通道
    zero_point: 128,
    bias: Some(&i32_bias),                             // 可选的逐行 i32 偏置,长度 m
};
gemm_i8_requant_u8(
    MatRef::from_row_major(&activations, m, k),
    MatRef::from_col_major(&weights, k, n),
    req,
    MatMut::from_col_major(&mut out_u8, m, n),
    Parallelism::Rayon(0),
);
}

输出为 C[i,j] = clamp(zero_point + round_ne(scale * (sum_k A*B + bias[i])), LO, HI),采用四舍六入五成双(round-half-to-even)。scale 要么是单个 RequantScale::PerTensor(f32),要么是逐行的 RequantScale::PerRow(&[f32])(逐通道约定)。钳位区间由具体入口决定:gemm_i8_requant[-128, 127]u8 孪生入口是 [0, 255]。这里没有 alpha,因为它已经并入了 scale。也没有 beta,因为往一个已经量化过的 C 里累加是没有良定义的。这个重量化映射在每一种 ISA(scalar、FMA、AVX-512F、VNNI)上、以及向量与标量两条存储路径之间都是逐位精确的,所以答案绝不取决于实际跑了哪个内核。

复数偏置

complex feature 之下,gemm_cplx_fused 给复数乘积加上逐行或逐列偏置:C <- alpha*op(A)*op(B) + beta*C + bias。它接受与 gemm_cplx 相同的可选操作数共轭。它按设计只支持偏置:像 ReLU 这样基于序的激活在复数上没有定义。conj_aconj_b 标志只共轭操作数本身,偏置是原样加上的,绝不会被共轭。

你可以依赖的保证

每个融合入口都把每个形状路由到普通 gemm 会选的同一个内核:通用 driver,或者某条特殊路径。它把 epilogue 融进那个内核的存储步骤,而不改变它的累加次序。所以融合调用不是另一种算法,它跑的是同一个 GEMM,只是在存储时施加了映射。具体的保证如下:

  • f32/f64,融合结果与普通 gemm 后接同一个标量映射逐位一致,这对每个形状、每种布局、每个 worker 数都成立。gemm_map 对逐元素的 f 给出同样的保证,复数偏置入口对“gemm_cplx 再加偏置“也给出同样的保证。
  • 对窄浮点 f16/bf16half feature)有一个明确记录的例外。gemmkit 把偏置和斜率精确加宽到 f32,在 f32 中施加 epilogue,只在存储时向输出做一次四舍五入取偶的收窄。这比 gemm 后接一次单独映射精确(后者会先舍入到窄类型、再加宽、再舍入一次),所以对窄类型,融合结果是有意地与这种两步式逐位相等的。可复现性与确定性不受影响。
  • 串行与并行运行在今天是逐位一致的。恒等融合的情形(None/None,或没有偏置)会常量折叠回严格的普通 gemm。可复现性契约只承诺同一个固定配置内的结果一致,而 worker 数正是该配置的一部分。

回报就是你不再需要做的那趟 C 扫描。在一个内存受限的 epilogue 上,那第二趟扫描的代价可能不亚于存储本身,所以在两步式并不便宜的场景下,把偏置或激活融进 GEMM 几乎是免费的。

融合 epilogue 也能和其它 API 档次组合使用。gemm_batched_fused 对一次批量 GEMM 的每个元素施加同一份共享的偏置和激活。gemm_packed_b_fusedgemm_packed_a_fused预打包操作数之上做融合。每个带检查的入口都有裸指针的 _unchecked 孪生版本,供适配器与 FFI 使用,它们用 (ptr, BiasDim) 对来携带偏置,而不是 Bias 枚举。见 Unchecked 层

批量GEMM

一个大 GEMM 单靠自己就能把现代 CPU 喂满,一堆小 GEMM 却做不到。注意力头、分组卷积、逐样本的线性层、块对角求解,都会产出许多彼此独立的小乘积。把它们当成一个普通的 gemm 调用循环来跑,会让机器大部分时间闲着:每次调用都小到无法有效并行,可循环却仍然要么每个元素付一次 fork/join,要么干脆串行。批量入口一次性接下整组乘积,把它当作一个整体来调度:把整个 GEMM 分派给 worker,于是一批小矩阵真正能把核心填满。

批量 GEMM 是一层编排,而不是一个新内核。每个元素都会重新经由完整的单 GEMM 引擎来分发,所以一次批量调用会自动与 driver、gemv 路径以及小形状路径组合起来:一批 1 x 1 x k 的乘积会在每个元素内部跑水平点积,一批普通形状则跑寄存器分块 driver。每个元素都是一个独立的 GEMM,整批在任意 worker 数下都可复现。

带步长形式

当各元素以规则的步长排布时(一个矩阵接一个矩阵地放在一块扁平缓冲区里),gemm_batched 只需接收一次单元素的形状和步长,外加 ABC 各自的一个批步长。元素 b 基于 A + b*a_batch_strideB + b*b_batch_strideC + b*c_batch_stride,所有元素共享同一个形状:

#![allow(unused)]
fn main() {
use gemmkit::{gemm_batched, MatRef, MatMut, Parallelism};

// batch 个独立的 m x k 乘 k x n 乘积,连续排布
gemm_batched(
    batch,
    1.0,
    MatRef::new(&a, m, k, 1, m as isize), (m * k) as isize, // A 单元素 + 批步长
    MatRef::new(&b, k, n, 1, k as isize), (k * n) as isize, // B 单元素 + 批步长
    0.0,
    MatMut::new(&mut c, m, n, 1, m as isize), (m * n) as isize, // C 单元素 + 批步长
    Parallelism::Rayon(0),
);
}

批步长为 0 会把一个操作数在整批上广播。这对只读的 AB 是合法的,比如让同一个共享权重矩阵去乘一批输入,但对 C 绝不合法,因为 C 的各元素是并发写入的,必须互不重叠。这样得到的结果精确复现了一个 gemm 调用循环。

epilogue feature 之下,gemm_batched_fused 对每个元素施加同一份共享偏置和同一个共享激活,也就是批量线性层的情形。它复现了一个 gemm_fused 调用循环。这份偏置向量是按单个元素定尺寸的,不是按整批。

切片形式:逐元素形状

当各元素形状不同,或者根本不落在固定步长上时,用 gemm_batched_slice。它接收一个 BatchProblem 切片,每个元素都携带自己的 alphaABbeta,以及一个独立的 &mut C 视图:

#![allow(unused)]
fn main() {
use gemmkit::{gemm_batched_slice, BatchProblem, MatRef, MatMut, Parallelism};

let mut problems: Vec<BatchProblem<'_, f32>> = /* 每个乘积一个,各有各的形状 */;
gemm_batched_slice(&mut problems, Parallelism::Rayon(0));
}

因为每个 C 都是一个独立的 &mut,各输出天然两两不相交,也不可能与输入产生别名。所以校验只需检查逐元素的形状是否一致、步长是否在界内。当你的矩阵本就以一个视图 Vec 的形式存在时,就该用这个形式。它的裸对应版本 gemm_batched_ptr_unchecked 接收一个 GemmProblem 切片,把同样的逐元素形状以裸指针的形式给出。它服务于自行校验输入、可能使用任意或负步长的 FFI 和适配器。两者都在 Unchecked 层中有详细说明。

一批是怎么调度的

真正有意思的决定,是这些工作怎么在核心间铺开。引擎会在每次调用时,根据共享形状和批大小做一次这个决定。一共有 3 种调度:

  • 批级并行(batch-parallel)。 各 worker 从一个共享游标里领取互不相交的元素区段,每个元素在一个 worker 上串行、缓存驻留地跑。这正是这套 API 的核心意义所在:对许多小矩阵,它为整批只付一次 fork/join,而不是每个元素一次,并且让每个核心都忙在完整的 GEMM 上。这种调度从不把某个元素拆到多个 worker 上,所以在任意 worker 数下都与串行运行逐位一致。
  • 串行(serial)。 整批在调用线程上跑,每个元素单线程执行。当总工作量太小、不值得付一次 fork/join 时,引擎会选用这种调度。
  • 顺序 + 内部并行(sequential with internal parallelism)。 对于元素少而大、受内存带宽支配的情形,引擎会把这批循环起来,逐个把完整的引擎并行度交给每个元素。当一个元素大到自己就能饱和内存带宽时,把它铺满所有核心,胜过同时跑好几个元素而互相冲刷缓存。这种调度只用于 m, n > 1 的形状,因为这些形状的路由本就会把每个输出的归约收在一个 worker 内完成,所以它仍然可复现。

对你而言,结论很简单:把整批交给引擎,让它去挑选调度方式。许多小而独立的乘积,正是批处理胜过手写循环之处。手写循环做不出「为所有元素只付一次 fork/join」这种选择,它要么并行每个小 GEMM(这几乎全是开销),要么干脆串行地跑。对于少数几个大乘积,批量调用会收敛到普通循环本就做得不错的行为,所以批处理在那里既帮不上多少忙,也不会拖后腿。

确定性贯穿这 3 种调度。每个元素都是独立的,所以整批在不同 worker 数下都可复现。串行与批级并行这两种调度更进一步:因为每个元素都只在一个 worker 上完整运行,它们在任意 worker 数下都逐位一致。元素少而大的那种调度,则继承了它所跑的那条路由自身的串行与并行行为。零长度的批是一个空操作。

和 API 的其余部分一样,每个批量入口都有一个 _with 变体,复用调用方持有的 Workspace,以避免每次调用都分配内存。有一个细节值得了解:在批级并行调度下,打包无法走单个共享的 Workspace,因为并发的 worker 会在它上面互相冲突。所以这种调度改为让每个 worker 走自己那份持久的线程局部池,其复用方式与你自己的 Workspace 是一样的。

小形状与GEMV

gemmkit 核心的寄存器分块 driver,需要每个输出 tile 有足够的工作量。它要靠这些工作量去摊销打包、缓存分块,以及一整个 MR x NR 累加器的开销。

有些形状彻底打破了这个前提。矩阵-向量乘完全没有 tile 复用。k = 4 的收缩在打包收回成本之前就算完了。一个 8 x 8 x 100000 的乘积会让 driver 把大部分精力都花在乘填充上。对这类形状,driver 是错的工具,于是引擎会悄悄绕开它。

这里没有什么需要你去开启。同一个 gemm 入口(以及 gemm_i8gemm_fusedgemm_map)会在分发的最顶端检查形状与步长。遇到合适的形状,它就改道到一个特殊内核;否则就落回通用 driver。每一次改道都藏在同一个公开入口之后。

每次改道都遵守库的可复现性契约:同一次调用,在同一台机器上、使用同一份配置,返回同一个结果。每次改道也都由一个调优旋钮门控,你无需重新编译就能移动或关闭它。你从不直接调用这些路径。你只是把形状写得自然,而不必手搓点积循环,就享受到了它们的好处。

gemv:受内存带宽支配的边界

m == 1n == 1 的形状是一个矩阵-向量乘。每个输出元素只需要 2k 次浮点运算,但读取它却要用到 k 个矩阵元素。算术本身是琐碎的,所以整个问题就变成了尽量减少 DRAM 流量。专用的 gemv 路径用同一个核心例程处理这两种情形。它把矩阵看作 rows x k 乘一个 k 向量,m == 1 时先转置矩阵。这条路径对每种布局都正确,并且对连续布局做了向量化。

这也覆盖了退化的 m == n == 1 点积。1 x k 的视图行跨度与列跨度都是 1,所以它同时符合列主序与行主序两种策略。列主序策略沿输出行向量化,可这里只有一个输出行,填不满一个 SIMD 寄存器。于是路由把任何短于一个寄存器的扫描都交给沿 k 向量化的那一支。这里没有什么需要开启,也没有什么需要知道。用 MatRef::from_col_major(a, 1, k) 构造的点积,也就是列主序库交给你的那种形状,现在走的是和行主序写法一样的快路径。

有两条性质对调用者要紧。其一,gemv 遵守库的一般可复现性契约:同一次调用,在同一台机器上、使用同一份配置,总是返回同一个结果。每个输出元素都由一个 worker 在一趟 k 扫描里归约完成。

把行分给多个 worker,改变的只是哪个 worker 做这份工作,而不是这份工作怎么做。库并不把 gemv 的这条保证延伸到不同 worker 数之间的逐位一致。worker 数是配置的一部分,这和 gemmkit 里其他地方一样。

其二,gemv 有自己的一套并行策略。它受带宽支配,所以 worker 数来自一个带宽模型,而不是通用 driver 用的那条计算爬坡。过了那几个能让 DRAM 饱和的核心之后,更多 worker 就不再有用了,只会增加 fork/join 开销和共享缓存争用。这个数遵循一道架在“本次调用触碰多少字节”之上的阶梯。在某个下限之下,矩阵还装得进单核的私有缓存,所以调用保持串行。过了那个下限,宽度就按台阶逐级上升,而且永远到不了整机宽度。

有四个旋钮暴露这套策略:

  • gemv_parallel_bytes 设定保持单线程的字节下限。
  • gemv_tier_step 设定阶梯的台阶之间隔多少字节。
  • gemv_thread_cap 设定一个固定宽度,直接取代整道阶梯。
  • gemv_threshold 与前面三个并列。由于 gemv 形状总满足 min(m, n) == 1,这个旋钮实际上起的是开关作用,而不是一个可分级的上限。

小 k 路径

收缩 k 也可能小到打包不划算。阈值是 small_k_threshold,x86 默认 16,aarch64 默认 8。在这个深度及以下,整个乘积就是单个深度面板,每个打包元素本来也只会被读一次。打包无从摊销,于是 driver 的打包步骤就变成了纯粹的开销。

k 路径覆盖这些瘦长、低深度的形状:gevv、rank-k 更新,以及高瘦乘积。它直接用家族的微内核计算 C就地读取 AB,不打包,一趟算完。这条路径免费继承了家族的加宽、偏置、共轭与舍入行为,并且对任意 worker 数都与串行运行逐位一致。它需要列主序的 A,也就是行单位步长(rsa == 1)。当 A 不是这种布局时,在这么小的 k 下打包本来也很少能摊销。这时路径就转而退回通用 driver,仍然算出正确的结果。

小 m,n 路径

镜像的情形是一个很小的输出配上一个很长的收缩。mn 都远低于微 tile,处在 small_mn_dim 及以下(x86 默认 16,aarch64 默认 32),而 k 很长。driver 会把很小的行 tile 和列 tile 填充到一整个微 tile。然后它会把大部分工作都花在这些填充上。

这条路径改为把每个输出算成一个水平点积,C[i,j] = alpha * <A[i,:], B[:,j]> + beta * C[i,j]。它沿收缩方向流式跑 SIMD,没有任何分块或取向机制。它还对这个小输出网格做寄存器分块,让好几条独立的 FMA 链同时在飞。

这个水平内核需要两个操作数都沿 k 单位步长。也就是说,A 的行必须连续(csa == 1,即行主序 A),B 的列必须连续(rsb == 1,即列主序 B)。两者都成立时,该路径就地零拷贝地读取 AB。这就是快路径。

两种最常见的布局各自恰好缺一边。全行主序缺 rsb,全列主序缺 csa。遇到这种情况,一次内部预打包会把仅那个不达标的操作数拷进 k 连续的暂存区一次,然后在其上跑同一个水平点积。这次拷贝读取大约 m*k 个元素,或者对另一个操作数是 n*k 个。

相对乘积本身 m*n*k 的工作量,这大约是 1/n(或 1/m)的一个零头。它的代价远小于水平路径省下的开销,所以一个带步长的小 m,n 形状仍然胜过落回 driver 的填充微 tile。

不过只按浮点运算量算会低估它:这次拷贝每搬一个字节做零次算术,而点积每字节约做两次,所以它能用来掩盖内存延迟的东西更少,占用的时间份额远大于它占用的工作量份额。因此它本身也被并行化了,而且与点积各自独立地决定宽度:小 m, n 留给点积去切的输出网格极小,而拷贝可切的深度却很长。这不需要任何配置。列主序、深收缩的小 m,n 形状在参考机上快了 1.1-3.1×(f32,自动宽度),其中收益较大的是被打包操作数仍能驻留缓存的那一段。

这个预打包档在 k 越过它自己的旋钮 small_mn_pack_min_k(默认 16)时启用,这个旋钮与零拷贝档所用的 small_k_threshold 是分开的。这条路径同样对任意 worker 数都与串行运行逐位一致。它把每个输出算成在一块不相交 tile 上的一次定序归约。

实用建议

优先用现成的入口,而不是手写点积循环。你的问题可能是一个矩阵-向量乘、一次 rank-k 更新,或者长收缩之上的一格小输出。不管哪种情况,gemm 都已经带着为那个形状调好的内核。它自带带宽感知的线程策略,以及手写循环得从头再造的那些可复现性保证。

布局是你手里唯一的杠杆。对水平的小 m,n 路径,行主序 A 加列主序 B 沿 k 单位步长流动,能命中零拷贝快路径。对小 k 路径,列主序 A 会留在就地路径上。其他任何布局仍然可用,只是要么付小 m,n 的预打包拷贝,要么落回通用 driver。

本章提到的每个阈值都是一个可调旋钮。每一个都按每次调用,依次从一个参数、一个程序化 setter、一个 GEMMKIT_* 环境变量,或一个校准过的编译期默认值中解析。如果某次改道对你的机器失准,或者你想把某个形状强制送上通用 driver,请参阅调优旋钮一章。它完整讲解了 gemv_thresholdsmall_k_thresholdsmall_mn_dimsmall_mn_pack_min_kgemv_parallel_bytesgemv_tier_stepgemv_thread_cap,以及 k_stream_max。至于特殊路径的内部机制,也就是每个内核为何长成那样,请参阅架构一章的特殊路径

运行时ISA分发

gemmkit 只发布一个引擎。它挑选自己运行所用的指令集,是在程序启动时决定的,而不是在编译时。设想一个在笔记本上构建、又拷到服务器上的构建产物:只要服务器支持 AVX-512,这个二进制就会用上它。同一个二进制放到较老的机器上运行,则会悄悄退回到较窄的内核。你不需要挑选后端,不需要靠 cfg 分门别类,也不需要为每台主机重新构建。第一次 GEMM 调用会探测 CPU 的特性、缓存下胜出的内核,之后每次调用都只是通过那个缓存下来的指针做一次普通的间接调用。

后端阵容

一个构建带有哪些内核,取决于目标架构。实际运行哪一个,取决于 CPU。从快到慢,候选者依次是:

  • x86-64 上的 AVX-512F,最宽的浮点内核。窄元素类型另有两个点积特化内核与它并列:面向 i8 -> i32AVX-512 VNNIvpdpbusd),以及面向 bf16AVX-512 BF16vdpbf16ps)。它们分别需要 int8half feature,并且 CPU 必须报告对应的特性位。
  • x86-64 上的 FMA / AVX2,面向没有 AVX-512 的机器的加宽 FMA 内核。
  • aarch64 上的 NEON,那里 SIMD 是基线(每颗 aarch64 CPU 都有),因此运行时无需探测。
  • wasm32 上的 simd128,它在编译期而非运行期被选定(见下文)。
  • scalar,可移植的兜底路径。它在每个目标上都存在,在没有更好的选择时接手。一个正确、即便未加速的结果,始终是可达的。

Tile 几何是唯一随 (元素类型, ISA) 组合而变的东西。微内核计算一个 MR x NR 的寄存器 tile,其尺寸按 ISA 的向量宽度定制。对 f32,随库发布的 tile 是:

ISAf32 tile(MR x NR)
AVX-512F32 x 12
FMA / AVX216 x 6
NEON16 x 4
simd1288 x 4
scalar4 x 4

这五个跑的是同一个通用浮点微内核,区别只在 tile 形状。MRMR_REG * LANES,所以向量越宽,tile 就越高。f64 的 lane 数减半,因此 MR 也随之减半(AVX-512F 上 f6416 x 12,以此类推)。VNNI 与 BF16 点积内核有各自按深度分组的几何结构,在元素类型一章里介绍。这张表只是背景知识,用来理解某个内核为什么这样打包、这样分块,而不是一个供你设置的旋钮。

自动选择

每种元素类型都独占一个分发槽位:一个持有类型化函数指针的 OnceLock。对该类型的第一次调用,会让选择阶梯运行一次特性探测,挑出当前可用的最优内核。接着它会存下那个内核的单态化入口(普通、预打包、融合),以及它的 tile 几何。不管哪次调用最先发生,都要承担这一次性开销:is_x86_feature_detected! 探测与 OnceLock 初始化。此后,分发就只是一次缓存指针加载加一次间接调用,不再对 ISA 做任何逐次调用的分支判断。这背后没有 transmute,也没有原子指针的腾挪,只是每种类型一个类型化槽位。

有一点值得说清楚:没有任何公开 API 会告诉你选中的是哪个 ISA。 这个选择是那个记忆化槽位的内部细节。如果你需要确认某个特定内核确实在跑,不要试图把它读回来,而应该把它钉住(见下一节),让不匹配的情况直接大声报错。

GEMMKIT_REQUIRE_ISA 钉住内核

设置环境变量 GEMMKIT_REQUIRE_ISA,会强制端到端只用某一个内核,而不再自动选择。可接受的取值(大小写不敏感,首尾空白会被去除)是:

取值强制使用同时接受
scalar可移植的标量内核
fmaFMA / AVX2 加宽内核avx2
avx512fAVX-512F 加宽内核
avx512vnnii8vpdpbusd 点积内核(其他类型走普通 AVX-512F)vnni
avx512bf16bf16vdpbf16ps 点积内核(其他类型走普通 AVX-512F)bf16
neonaarch64 的 NEON 内核
simd128wasm32 的 simd128 内核wasm
auto正常的自动选择(未设置或为空时也是它)

avx512vnniavx512bf16 这两个 pin,只为它们各自那一种窄类型选用点积内核。其余一切都走普通 AVX-512F 路径,所以在这两个 pin 之下,混合负载中的其他类型仍能正确运行。

契约是报错,而非回退。 只要所请求的 ISA 不可用,分发就会 panic,而不会悄悄换一个内核来跑。以下三种情况都算不可用:

  • CPU 没有报告该特性
  • 这个取值命名的 ISA,在当前目标架构上并不存在(x86 上的 neon、aarch64 上的 avx512f
  • 这个取值本身就是个拼写错误

这是有意为之的,也正是 CI 想要的效果。如果一次 CI 任务的全部目的就是要跑通 AVX-512 VNNI 路径,它就绝不能因为特性标志拼错了、或者模拟器配错了,而悄悄测试着标量兜底路径却蒙混过关。gemmkit 自己的 CI 就是这样钉住每个内核的:x86 的点积内核在 Intel SDE 下跑,NEON 在 aarch64 上跑,simd128 在 wasm 上跑。一个坏掉的 pin 会变成一次红色构建,而不是一次假的绿色构建。

这个值只读取一次,在首次分发之前,并与内核选择一起被记忆化。请在任何 GEMM 运行之前,把它设进进程环境。进程运行到一半再改它不会有任何效果,因为槽位早就填好了。一个无法识别的取值会是硬错误,就是为了不让它被误当成 auto 而蒙混过关。

WebAssembly 是编译期决定的

wasm32 没有运行时特性探测,所以 simd128 不是靠探测机器选出来的。它是由编译期 cfg 选定的,而构建必须真的用 -C target-feature=+simd128 把它打开。忘了这个标志,wasm 构建就会悄悄用标量兜底。钉住 GEMMKIT_REQUIRE_ISA=simd128,会把这种无声降级变成一个断言:如果 SIMD 路径没在跑,构建就会 panic。这正是 wasm CI 作业要钉住它的原因。完整的 wasm 构建流程,包括带线程的目标,见no_std 与 WebAssembly

调优旋钮

gemmkit 里的每一个启发式都是一个带随库默认值的具名阈值,而不是写死的常量。这涵盖了何时并行、何时打包某个操作数、一个形状从哪里起不再算“小”这类问题。少数几个旋钮按架构做了拆分,aarch64 需要和其余架构不同的值。这些默认值在大多数硬件上都表现良好。当某个默认值不适合你的机器时,你有三种途径能够到它,都不需要改动源码。

解析顺序

一个旋钮在被读取的那一刻解析,取以下第一个已设置的层次:

  1. 每次调用的参数。 凡是有调用点等价物的旋钮,它直接胜出。最清楚的例子是并行度:你传给 gemmParallelism 参数会覆盖任何全局线程策略。这一层位于 API 里,不在 tuning 里。
  2. 编程式 setter。 gemmkit::tuning::set_*(v) 会无条件地存下一个值。一旦设置,之后的读取就再也不查环境变量。这是给那种在代码里自我调优的应用准备的,它理应压过部署环境所提供的任何东西。
  3. 环境变量。 GEMMKIT_*。这是部署层。source 一份配置,比如 gemmkit-tune 生成的那种,就能为某台主机重调一个已经构建好的二进制,不用重新编译。
  4. 编译期默认值。 那个标定好的常量,需要时按架构拆分。

setter 压过环境变量是有意为之的。调用了 setter 的应用,就是选择不理会环境变量。想让部署配置生效的应用,只要不去调那些 setter 即可。

环境变量只在对该旋钮的首次访问时读取一次,随后缓存为一个原子量。在某个旋钮首次被读之后再设置的值会被忽略,所以要在进程启动之前就导出配置。一个已设置但无法解析为非负整数的 GEMMKIT_* 变量会被当作拼写错误,而不是无声的空操作。gemmkit 会在 stderr 上告警,并回退到默认值。由于回退结果随后会被缓存,这个告警每个旋钮只会触发一次。它绝不会 panic:一个性能旋钮的拼写错误不应该让进程崩溃。

旋钮清单

下表是覆盖所有 feature 与目标配置的完整清单。内部的 tuning::knob_env_names 注册表是这份清单的唯一真实来源。其中两个旋钮受 feature 或目标门控,只有编译进来时才存在。每个 getter 都有对应的 set_*。环境变量名就是 getter 名的大写形式,加上 GEMMKIT_ 前缀。

串行 / 并行门槛

环境变量Setter默认值控制什么
GEMMKIT_PARALLEL_THRESHOLDset_parallel_threshold4848256m*n*k 低于此值时,工作会被强制放到单线程上。这是串行到并行的盈亏平衡点。如果你的线程池 fork 代价高,就调高它。如果你的线程廉价,且有值得拆分的小乘积,就调低它。

打包门槛与步长

环境变量Setter默认值控制什么
GEMMKIT_RHS_PACK_THRESHOLDset_rhs_pack_threshold2048只有当 m(有多少行块会复用它)超过此值时,才打包 RHS 宏面板。低于它时,B 就地读取。
GEMMKIT_LHS_PACK_THRESHOLDset_lhs_pack_threshold1024(aarch64:256)只有当每个工作线程的列复用超过此值时,才打包 LHS。打包在 aarch64 上更便宜,所以那里从更低的复用起就划算。
GEMMKIT_LHS_PACK_STRIDEset_lhs_pack_stride0(自动)这是加在列主序深度步长 csa * sizeof(Lhs) 上的字节门槛。一旦步长达到这么多字节,A 就会被打包,以避开对 TLB 和缓存都不友好的跨步读取,这与复用无关。0 表示从操作系统页大小推导。这道门槛与下面的跨度门槛、复用门槛取“与”:只有步长、跨度、复用三者同时成立,强制打包才会触发。
GEMMKIT_LHS_PACK_SPANset_lhs_pack_span0(自动)这是上面那道步长门槛的地址跨度搭档。只有当整条深度切片的遍历(csa * sizeof(Lhs) * kc)也达到这么多字节时,页级步长才会强制打包列主序的 A。在这个跨度以下,遍历仍驻留在缓存里,只是重读温热的缓存行,所以就地读取比它本要付的打包更快。0 表示自动(4 MiB)。
GEMMKIT_LHS_PACK_REUSEset_lhs_pack_reuse128(aarch64:4)这是一道复用下限,为强制打包的收益定价,而不是为它的成本定价。上面的步长门槛与跨度门槛,只有在超过这道复用下限时才会触发。这道下限以复用每个打包面板的 nr 宽列 tile 数来衡量(min(n, nc) / nr,向上取整)。一个瘦高形状(m 远大于 n)跨度巨大,但列 tile 很少,把一次昂贵的打包摊到过少的复用上并不划算,这道下限就是用来拦住这种情况的。0 会去掉这道下限,只让步长加跨度这一对来决定。在 aarch64 上,这笔账几乎反过来了:打包在那里很便宜,而就地遍历要跨很小的页。所以 aarch64 的默认值从远低得多的复用量起就开始打包。
GEMMKIT_SHARED_LHS_MNKset_shared_lhs_mnk8e9(aarch64:6e6,32 位:禁用)这是并行打包路径上共享 A 预处理的 m*n*k 门槛。这道预处理会消除各工作线程的冗余打包,代价是一次 fork-join 屏障。这个临界点权衡的是屏障的代价与它省下的打包开销。打包在 aarch64 上代价相对更低,所以那里的门槛远低于 x86。与这道门槛无关,预处理还会在工作线程数达到 16 起开启,因为在那个宽度上,各自打包的冗余代价必定压过屏障。
GEMMKIT_PACK_TRANSPOSE_TILEset_pack_transpose_tile16这是打包某个跨步操作数时所用的缓存分块转置的条带长度,把逐元素的 gather 变成分块拷贝。实数与复数打包器都靠它支撑。

特殊路径阈值

环境变量Setter默认值控制什么
GEMMKIT_GEMV_THRESHOLDset_gemv_threshold无上限当另一维为 1 时,为专用 gemv 路径的 min(m, n) 设上限。触发 gemv 的是形状而不是规模,这个旋钮只是给它封顶。
GEMMKIT_SMALL_K_THRESHOLDset_small_k_threshold16(aarch64:8)k 在此值及以下时,形状会走通用的 small-k 路线(单个深度面板,不打包),而不是寄存器分块驱动器。
GEMMKIT_SMALL_MN_DIMset_small_mn_dim16(aarch64:32)mn 都在此值及以下(且 k 较长)时,会走横向内积路线,每个输出都是一次 SIMD 归约点积。0 会关闭这条路线。寄存器分块驱动器则会把小的行、列 tile 填充到一个完整的微 tile,把大量工作花在这些填充上。驱动器开始占优的那个点因机器而异,这就是为什么 aarch64 的上限高于 x86。
GEMMKIT_SMALL_MN_PACK_MIN_Kset_small_mn_pack_min_k16这是 small-m,n 打包层的 k 门槛:只有当 k 超过此值时,一个跨步的小形状才会把不满足条件的操作数拷贝进 k 连续的暂存区。
GEMMKIT_GEMV_PARALLEL_BYTESset_gemv_parallel_bytes0(自动)这是带宽受限的 gemv/gevv 保持单线程的字节下限。低于它时,矩阵装得进单核的私有缓存,该核已经独自打满了它,拆分只会白白吃亏。0 表示从探测到的缓存推导:在有 L3 的部件上,取单核私有 L2;在没有 L3 的 aarch64 部件上,取共享簇 L2 的八分之一。
GEMMKIT_GEMV_TIER_STEPset_gemv_tier_step0(自动)这是自动 gemv/gevv 工作线程阶梯各级之间的字节间距。从上面那道字节下限起,触碰字节每增长这么多倍,宽度就往上跳一个精确适配的池档位。0 表示取 8。1 会把整道阶梯压缩到它的顶档。活跃池档位少于 2 个时,这个旋钮不起作用,所以在单档位的 aarch64 默认设置下它什么也不做。
GEMMKIT_GEMV_AXPY_PAR_MIN_ROWSset_gemv_axpy_par_min_rows16384(x86)、1024(aarch64)这是输出行数的下限,低于它时,列主序 gemv 会保持串行,不再拆分它的行。对列主序矩阵而言,输出行这个轴是内层内存轴。一旦拆分,就会让每个工作线程都要跨步走遍整个矩阵,而串行路径只需一趟顺序扫描。只有当每个工作线程分到的那段跑得足够长时,拆分才划算,才配得上它放弃的那种连续性。0 会关闭这道下限。两个默认值相差一个数量级,因为 aarch64 上的临界点低得多。在该平台上,临界点跟随的是单列的字节数而非行数,所以以 f64 为主的负载应取默认值的一半。行主序 gemv 与 half 混合孪生版本永远不受此限制:两者在任何尺寸下拆分都能很好地扩展。
GEMMKIT_GEMV_THREAD_CAPset_gemv_thread_cap0(自动)这是带宽受限的 gemv/gevv 的一个固定工作线程数,直接取代上面那道阶梯。它就是“我实测过这台机器,就用这个数”的覆盖开关。0 会保留那道阶梯,它的顶端是逻辑核心数的一半,因为 gemv 远在用满全部宽度之前就已经打满了自身带宽。
GEMMKIT_K_STREAM_MAXset_k_stream_max32这是 k 的上限,axpy 形状的 gemv 在整个深度扫描过程中,会把输出面板一直留在寄存器里,直到超过这个上限。超过之后,普通的列外层形式更快。
GEMMKIT_SEQ_INTERNAL_BYTES_PER_WORKERset_seq_internal_bytes_per_worker128 KiB这是 aarch64 批量 GEMM 的一个临界点:一旦某个批元素分摊到每个工作线程的字节份额超过此值,它就会跨机器拆分,而不是一线程一个、缓存热运行。只在 aarch64 上被查询。
GEMMKIT_I8_VNNI_MIN_PAR_MNKset_i8_vnni_min_par_mnk768^3m*n*k 低于此值时,自动选中的 VNNI i8 内核会转而把一个多线程问题交给加宽兜底。VNNI 强制要求的 RHS 打包屏障,在一个小规模的并行问题上并不划算。它与 VNNI 逐位一致。需要 int8 feature。

调度粒度

环境变量Setter默认值控制什么
GEMMKIT_PARALLEL_OVERSAMPLEset_parallel_oversample8并行驱动器的目标,是让每个工作线程分到这么多个工作块,从一个共享游标上按需领取。调得更高,负载均衡更细、尾部更小,代价是原子领取次数更多。调得更低则更粗糙,但开销更小。
GEMMKIT_PAR_MNK_PER_WORKERset_par_mnk_per_worker2000000(线程化 wasm:262144)这是自动工作线程数的粒度。自动路径以 m*n*k 除以每个工作线程这么多工作量为目标,再用核心数与作业数给结果封顶,下限为 1。这样一来,数量就随总 flops 扩展,而不是随线性尺寸扩展。一个 wasm 工作线程的启用成本远低于一个原生线程,因此 wasm 的下限更低。0 的行为等同于 1(永远用满全部宽度)。
GEMMKIT_PACKED_OVERSAMPLEset_packed_oversample2这是打包 LHS 路径自己的拆分目标,和上面那个通用粒度不是一回事。拆得更狠,会让 A 被重新打包得太频繁,反而拖累性能,所以这个最优值更低。
GEMMKIT_POOL_CLASSESset_pool_classes2(aarch64:1,其余:0)这是满机器宽度以下的减半档位数:先是一半,再是四分之一。对每个活跃档位,gemmkit 都会保留一个私有的、持久化的 rayon 池,首次用到时才惰性构建,此后不再重建。自动路径会把它的工作线程数精确卡在某个档位上,让 fork/join 屏障处没有线程闲置。显式的 Rayon(n) 仍然恰好拿到 n 个工作线程,只是会跑在能容纳它的最小档位池里。0 会关闭档位池,让每次调用都留在环境池上。上限为 3。x86_64 默认 2 档,aarch64 默认 1 档,其余每个目标默认 0(关闭)。
GEMMKIT_FULL_WIDTH_MNKset_full_width_mnk0(自动)m*n*k 高于此值时,自动路径才会离开它最大的那个档位池,转向全部机器宽度。低于它时,即便存在更多核心,自动模式也会停留在它最大的档位上,因为这时那些额外的全宽工作线程,还赚不回它们额外付出的 fork/join 代价。0 表示按架构推导:x86 上是 110_000_000,aarch64 上是 14_000_000,在 aarch64 上,包含 E 核在内的全部宽度,在更小的问题规模下就已经划算。MAX 会让自动路径无条件钉死在最大档位上,全部宽度就永远不会介入。

分块上限

环境变量Setter默认值控制什么
GEMMKIT_MC_REG_PANELSset_mc_reg_panels8A 的宏面板被限制为这么多个微 tile 行(this * MR),遵循 BLIS 的规则:MC 始终是 MR 的一个小倍数。
GEMMKIT_NC_NO_L3_PANELSset_nc_no_l3_panels512没有 L3 时(比如 Apple Silicon 这类部件),列块是 min(this * NR, N)。存在 L3 时,这个旋钮不起作用。
GEMMKIT_TINY_BLOCK_DIMset_tiny_block_dim64mn 都在此值及以下的形状,会跳过完整的 BLIS 分块模型,只是把 A、B 面板留在 L2 里。
GEMMKIT_KCset_kc2048(aarch64:16384)这是微小矩阵捷径里的深度块:k 会被钳到此值。该计数以 4 字节元素为单位,更窄的元素会去除它,从而保住打包后的面板字节数。在 x86 上更宽的元素同样会去除它。在 aarch64 上更宽的元素保留完整深度,因为在那里做除法会让切片数成倍增加,而每个切片都要多付一次 worker fork。另外在 aarch64 上,更深的切片能一直领先到比 x86 更远,所以 aarch64 上的这条捷径几乎是单切片运行。
GEMMKIT_KC_MINset_kc_min512这是主模型的 kc 下限:那个贴合 L1 的深度估计值会被抬高到至少这个数,好让一个很小的 L1 缓存也不会让微内核的深度扫描挨饿。

深收缩与 wasm

环境变量Setter默认值控制什么
GEMMKIT_DEEP_KC_BYTESset_deep_kc_bytes0(自动)这是深收缩路径的启用门槛,以字节计。一个窄输出家族(f16/bf16)通常会把整个收缩当作单个深度面板来跑。一旦它的 RHS 微面板(nr * k * sizeof(N))超出此值,它就会切到一个 f32 输出的多切片孪生版本上。0 表示从探测到的 L2 的一半推导。
GEMMKIT_PREFETCH_MIN_BYTESset_prefetch_min_bytes0(自动)这是驱动器 C tile 软件预取的启用门槛,以字节计。一旦一次调用的工作集(A + B + C 字节)超过此值,输出微 tile 就会从 LLC 之外流入。驱动器接着会在每个微 tile 对应的微内核调用之前,为它发出一次 T0 预取,掩盖读-改-写的延迟。低于此值时,tile 都驻留在缓存里,这个提示反而纯属开销。0 表示从单核可达的 LLC 推导(有 L3 就取 L3,否则取 L2)。非零值就是字节门槛本身,所以 usize::MAX 会关闭预取,1 会强制开启它。这个旋钮只在 x86_64 上生效,在其他目标上是空操作,所以 aarch64 与 wasm 都不受影响。它对数值也不可见:开或关都逐位一致。
GEMMKIT_WASM_THREADSset_wasm_threads8这是带线程的 wasm 构建的工作线程数,因为 wasm 没有 available_parallelism 可查。它给 gemmkit 的 wasm rayon 池定尺寸。只在带 wasm_threads feature 的 wasm32 上存在。

关于 GEMMKIT_FAST_TEST

你可能会在测试框架里看到 GEMMKIT_FAST_TEST。它会把正确性扫描收窄,好让测试跑得更快,是一个仅测试套件使用的开关。库本身从不读取它,设置它对生产环境中的 GEMM 没有任何影响。

超越手工调参

手动设置旋钮,适用于你已经知道该动哪一个的场合。要为某台具体机器标定整套旋钮,就运行自动调优器。它会在一组探测形状上扫描每个旋钮,并写出一份 GEMMKIT_* 配置,你在运行前 source 它即可,不用重新编译。这是 gemmkit-tune 一章的主题。

no_std与WebAssembly

gemmkit 的核心不需要操作系统。把默认 feature 关掉,crate 就变成 #![no_std] 的。这时它只需要 corealloc,此外不依赖任何东西。这让它可以用在内核、嵌入式固件,以及 WebAssembly 里。同一条代码路径,也正是 wasm SIMD 后端的构建方式。本页讲的是 no_std 构建放弃了什么、保留了什么,以及 wasm 目标额外需要哪些步骤。

no_std 核心

std feature 默认开启,是 default = ["std", "parallel"] 的一部分。把默认关掉,你就走上了只带 alloc 的路径:

[dependencies]
gemmkit = { version = "0.1", default-features = false }

alloc 始终是必需的,因为两种构建里打包暂存都由堆支撑。除此之外,这种配置下 crate 什么也不拉。每个可选 feature 至多增加一个依赖:

  • stdraw-cpuid,仅 x86,用于 CPUID 缓存与特性探测。
  • parallelrayon
  • halfhalf
  • complexnum-complex

int8epilogue 完全不增加依赖。元素类型 feature 能和 no_std 自由组合。比如一个 default-features = false, features = ["half", "int8"] 的构建,就是一个合法、零依赖的 f16/bf16/i8 引擎。

注意 parallel 蕴含 std,因为 rayon 需要标准库。所以 no_std 构建总是单线程的。一切照常编译、照常运行。只是跑在一个线程上。

没有 std 时会发生什么

有三样东西会从运行期挪到编译期,或者从自动变成显式:

特性探测变成编译期的。std 时,x86 分发会调用 is_x86_feature_detected!。它会挑出运行中 CPU 所报告的最优内核。没有 std 时,就没有运行时 CPU 探测。这项能力藏在受 std 门控的 raw-cpuid 里。于是 ISA 阶梯改为回退到 cfg!(target_feature = ...),构建跑的是它编译期目标特性所保证的东西。

想从 no_std 构建里得到加速的 x86 内核,你就必须为之编译。比如传入 -C target-cpu=native,或者显式的 -C target-feature=+avx512f。不这样做,你拿到的就是标量兜底。在 aarch64 和 wasm 上,选择本来就是这样工作的,所以那里没有任何损失。

环境旋钮关闭了。 读取环境变量需要 std。没有它,GEMMKIT_REQUIRE_ISA 就从不被查询(分发总是自动选择),每一个 GEMMKIT_* 调优旋钮都直接解析到它的编译期默认值。编程式的 tuning::set_* setter 仍然有效,所以你可以在代码里、而不是通过环境变量,重调一个 no_std 构建。setter 这一层参见调优旋钮

每次调用的工作区取代了线程池。 默认的线程本地打包池是一个 std 构造。没有 std 就没有池。每次调用都会改为分配一个新的 Workspace 作为暂存,返回时释放。这样做是对的,但每次调用都要分配。要达到零分配的稳态,就创建一个 Workspace,把它穿过 *_with 入口:gemm_with,以及每个家族的 _with 变体。第一次足够大的调用之后,这些入口就会复用缓冲区,不再有进一步的堆流量。

为 WebAssembly 构建

wasm32 没有运行时特性探测,所以 simd128 后端由编译期 cfg 选定。构建必须显式打开那个目标特性。如果你忘了它,wasm 构建仍会正确编译、正确运行,但只能跑在标量兜底上,慢很多。通过 RUSTFLAGS 传入这个标志:

RUSTFLAGS="-C target-feature=+simd128" \
  cargo build --target wasm32-wasip1 --no-default-features --features std

要运行产物,你需要一个 wasm 运行时。gemmkit 的 CI 用 wasmtime,所以把 Cargo 的目标 runner 指向它。有时你想确认 SIMD 路径确实在跑,而不是悄悄地退回了标量。这时就把 ISA 钉住:GEMMKIT_REQUIRE_ISA=simd128 会把缺失的 +simd128 变成一次 panic,而不是无声的回退。这正是测试作业想要的。这个 pin 需要 std,而 wasm 构建里 std 已经是开着的。

RUSTFLAGS="-C target-feature=+simd128" \
CARGO_TARGET_WASM32_WASIP1_RUNNER="wasmtime --env GEMMKIT_REQUIRE_ISA=simd128" \
  cargo test --target wasm32-wasip1 --no-default-features --features std

基线的 wasm32-wasip1 没有线程。如果你为基线 wasm 目标开着 parallel 构建,gemmkit 不会 trap。一个内部守卫会让 rayon 在那里不可用,于是 Parallelism::Rayon(_) 会转而降级到串行循环。这样一来,一个可移植的 wasm 二进制就可以带着 parallel feature,单线程运行即可,无需为目标专门构建。

带线程的 wasm

wasm 上真正的多线程需要支持线程的目标,以及对应的 feature:

RUSTFLAGS="-C target-feature=+simd128" \
CARGO_TARGET_WASM32_WASIP1_THREADS_RUNNER="wasmtime -W threads=y -W shared-memory=y -S threads=y" \
  cargo test --target wasm32-wasip1-threads \
  --no-default-features --features std,parallel,wasm_threads

wasm_threads feature 蕴含 parallel,面向 wasm32-wasip1-threads。它会打开 gemmkit 专用的 wasm rayon 线程池。由于 wasm 运行时无法报告核数,池的宽度不是自动推导的。它改为来自 GEMMKIT_WASM_THREADS 旋钮,默认 8,这个旋钮既给自动线程数封顶,也给池定尺寸。把它设成与你的运行时实际配备的工作线程数相符。其余一切的表现都与原生带线程构建完全一致:分块、工作清单,以及可复现性。

Unchecked层

安全入口(gemmgemm_fused 等等)在触碰内存之前会先校验输入。形状必须相符。每个跨步视图都必须落在自己的切片之内。输出必须一次性寻址每个元素,并且不能与输入重叠。

在每一层这样的检查之下,坐着的都是同一个引擎。它通过一个裸指针、isize 步长、完全不做检查的接口触达。这就是 unchecked 层。它是为那些已经持有安全 API 本会重新推导出的那些不变量的调用者而存在的。

它面向谁

这里住着三类调用者。封装其他矩阵库的适配器,比如 ndarraynalgebrafaer,直接从宿主类型里就已经拿到了一个校验过的指针和步长。再检查一遍边界,就是在库已经保证过的数据上做冗余功。从 C 或其他语言过来的 FFI 调用者,手里只有一个指针和步长,根本没有 Rust 切片可供边界检查。代码库自有的自定义矩阵类型,可以降解为指针并直接调用引擎,而不必先拷进一个 MatRef。这几种情形里,知道内存有效的都是调用者自己,所以检查就搬到了那份知识所在之处。

如果以上都不是你的情况,就用安全 API。对单次调用而言,unchecked 层并不更快。相对于乘法本身,校验的代价很便宜。它存在的意义,是让已经持有不变量的调用者,不必再证明一遍。

目录清单

每个安全入口都有一个裸孪生,命名方式是加后缀 _unchecked,多数还提供一个接收调用者自有工作区的 _with 形式(见下一节)。按家族划分,完整的裸接口如下:

家族裸入口Feature
普通 GEMMgemm_uncheckedgemm_unchecked_with核心(f32/f64half 下另加 f16/bf16
复数gemm_cplx_uncheckedgemm_cplx_unchecked_withcomplex
整数gemm_i8_uncheckedgemm_i8_unchecked_withint8
融合偏置/激活gemm_fused_uncheckedgemm_fused_unchecked_withepilogue
Map(逐元素闭包)gemm_map_uncheckedgemm_map_unchecked_withepilogue
复数融合gemm_cplx_fused_uncheckedgemm_cplx_fused_unchecked_withcomplex + epilogue
重量化gemm_i8_requant_uncheckedgemm_i8_requant_u8_unchecked(及 _withint8 + epilogue
跨步批量gemm_batched_uncheckedgemm_batched_unchecked_with核心
指针数组批量gemm_batched_ptr_unchecked核心
批量融合gemm_batched_fused_uncheckedgemm_batched_fused_unchecked_withepilogue
预打包prepack_rhs_uncheckedprepack_lhs_uncheckedprepack_rhs_i8_unchecked核心 / int8
消费预打包gemm_packed_a_uncheckedgemm_packed_b_unchecked(及 _with_fused_核心 / epilogue
消费预打包(i8)gemm_i8_packed_b_unchecked(及 _withint8

指针数组批量形式值得单拎出来说。gemm_batched_ptr_unchecked 接收一个 GemmProblem<T> 的切片,每个元素都有自己的形状、自己的指针。于是一个批次既能混合不同尺寸,也能把操作数散落在内存任意处。它没有同样形状的安全对应物。表达“一组相互独立的裸问题”正是裸层的用途所在。nalgebra 和 faer 适配器就把它们的批量 GEMM 搭建在它之上。

安全契约

调入 unchecked 层,意味着为安全 API 本会检查的东西逐项签字:

  • 有效的指针与步长。 对由维度和步长隐含的每个 (i, j)ab 对读有效,c 对读写有效。没有任何东西会为此做边界检查。越界的步长是未定义行为,而非 panic。
  • 一个唯一寻址的输出。 C 的步长必须把每个不同的 (i, j) 映射到不同的位置。并行驱动器假定输出 tile 互不相交,并会并发写入它们。一个自别名的 C(例如 rsc == 0)就会因此成为一场数据竞争。输入则可以自由地自别名,因为它们只被读取。所以一个零步长的广播式 AB 没有问题。
  • CA/B 不重叠。 输出会被写入。如果它与某个输入重叠,结果就会是垃圾。

有一处放宽是随之而来的。当 beta == 0 时,输出不被读取,因此 C 无需初始化。这里还有一项安全 API 不给的能力:负步长,以及指向缓冲区中部的指针,二者都是允许的。一个反向视图(rs < 0),或者一个从最后一个元素往回寻址的操作数,正是安全的 MatRef 会拒绝的那类布局。裸引擎则接纳它。这也是为什么封装那些会产出反向步长的库的适配器要转发到这一层。

复用工作区

每个裸入口有两种分配风格。普通形式 gemm_unchecked 借用线程本地打包池。它每线程至多分配一次。_with 形式 gemm_unchecked_with 则改为接收一个你自有的 &mut Workspace

#![allow(unused)]
fn main() {
use gemmkit::{Workspace, Parallelism};

let mut ws = Workspace::new();
// each iteration reuses `ws`; after the first large call it does no heap work
for _ in 0..iters {
    // SAFETY: pointers/strides valid, c uniquely addressed, c disjoint from a/b
    unsafe {
        gemmkit::gemm_unchecked_with(
            &mut ws, m, k, n,
            1.0_f32, a, rsa, csa, b, rsb, csb, 0.0_f32, c, rsc, csc,
            Parallelism::Serial,
        );
    }
}
}

工作区会长到能容纳它服务过的最大问题。此后它复用那块分配,所以一个 GEMM 热循环能达到零稳态分配。这正是 no_std 构建赖以复用的机制,因为它们没有线程本地池。在 std 下,对于想把分配挪出热路径的实时或延迟敏感循环,它同样好用。

一个完整示例:自定义 tile 类型

假设你的代码已经带着自己的稠密行主序矩阵,你想把两个相乘,而不必先拷进一个 MatRef

#![allow(unused)]
fn main() {
use gemmkit::{gemm_unchecked, Parallelism};

// a dense row-major matrix the caller already owns
struct Tile {
    data: Vec<f32>,
    rows: usize,
    cols: usize,
}

// c = a * b for row-major tiles
fn matmul(a: &Tile, b: &Tile, c: &mut Tile) {
    assert_eq!(a.cols, b.rows);
    assert_eq!(a.rows, c.rows);
    assert_eq!(b.cols, c.cols);
    // row-major: row stride = cols, column stride = 1
    // SAFETY: shapes checked above; each tile owns a dense rows*cols buffer, so
    // every addressed element is in bounds; c is a distinct &mut, so it cannot
    // alias a or b, and a dense layout addresses each (i, j) once
    unsafe {
        gemm_unchecked(
            a.rows, a.cols, b.cols,
            1.0_f32,
            a.data.as_ptr(), a.cols as isize, 1,
            b.data.as_ptr(), b.cols as isize, 1,
            0.0_f32,
            c.data.as_mut_ptr(), c.cols as isize, 1,
            Parallelism::Serial,
        );
    }
}
}

assert_eq! 的形状检查与 &mut Tile 借用,两者合在一起就结清了整份契约。稠密存储让每个偏移都落在界内,也让每个 (i, j) 都各不相同。对 c 的独占借用则排除了它与 ab 的重叠。这就是该采用的范式:在你自己类型的边界处证明不变量,然后把裸指针交给引擎。

适配器就是参照

把这件事做好的最干净的例子,就是适配器 crate 自身。每一个都从原生视图里抠出指针和步长:C 序、F 序、一般步长,或反向步长,全都零拷贝。然后各自转发到 *_unchecked 引擎,并在每个调用点附上一段简短的安全论证。如果你在封装自己的矩阵类型,读一读某个适配器章节,照着它的结构来。nalgebra 一章是个不错的起点。至于上面目录里的预打包入口,它们服务的定权重复用范式,在预打包操作数中有介绍。

在ndarray中使用gemmkit

gemmkit-ndarrayndarray 二维数组与 gemmkit 引擎之间的一层薄桥接。它本身不做任何数值计算。每个入口接受一个 ArrayBase。它直接从中读出基指针和 2 个轴的步长,再把这些原始数据交给 gemmkit 的 unchecked 引擎。整个 crate 只是一层步长转接。gemmkit 会的一切,比如运行时 ISA 选择、缓存分块和可复现的并行,都原样适用。数组在传入时既不会被重排,也不会被拷贝。

这些入口接受任意存储类型 S: Data&ArrayBase<S, Ix2>。无论是拥有所有权的 &Array2<T>,还是借用的 ArrayView2<T>,都能用,ArcArrayCowArray 以及它们的切片同样可以。唯一的内部辅助函数值得一看。它就是适配器全部的数据提取逻辑:

gemmkit-ndarray/src/common.rs

#![allow(unused)]
fn main() {
pub(crate) fn dims_strides<T, S: Data<Elem = T>>(
    a: &ArrayBase<S, Ix2>,
) -> (usize, usize, isize, isize) {
    let (r, c) = a.dim();
    let s = a.strides();
    (r, c, s[0], s[1])
}
}

这个 (rows, cols, row_stride, col_stride) 元组,加上 a.as_ptr(),就是 gemmkit 需要的全部。步长是带符号的 isize,所以负(反转)步长和正步长一样直接转发。

加入项目

是 2 个 crate,不是 3 个。适配器把自己签名里出现的一切都重新导出了,所以常规配置中不需要直接依赖 gemmkit

[dependencies]
gemmkit-ndarray = "0.1"
ndarray = "0.17.1"

gemmkit_ndarray 把调用方需要的一切都重新导出了,所以很少需要再直接依赖 gemmkit

  • Parallelism 选择器,以及每个 _with 变体都要用到的 Workspace 类型。
  • fused 选择器 BiasActivation
  • 预打包句柄 PackedLhsPackedRhs
  • 重量化参数 RequantizeRequantScale
  • 元素类型约束 GemmScalarFusedScalarMapScalarComplexScalar。写一个对某个入口泛型的封装时,需要用到它们。
  • 各自受 feature 门控的元素类型 f16bf16Complexc32c64,这样 halfnum-complex 也不必进入你的 manifest。
  • tuning 模块。

请通过适配器去用 tuning,不要为此单独再加一个 gemmkit 依赖。这些调优旋钮是进程级全局原子量。第二份单独解析出来的 gemmkit 只会给你一组适配器根本不会读取的原子量。

gemmkit-ndarray 上的每个 feature 都直接转发到 gemmkit 中的同名 feature。在这里打开某项能力,对应的入口点就会随之启用:

  • parallel(默认):rayon 多线程。
  • wasm_threads:在 wasm32-wasip1-threads 上启用线程,隐含 parallel
  • halff16 / bf16 输入,以 f32 累加。
  • complexComplex<f32> / Complex<f64> 矩阵。
  • int8i8 输入,累加进 i32
  • epilogue:融合的偏置 / 激活、i8 / u8 重量化,以及用户自定义的逐元素映射。

默认 feature 是 ["parallel"]进阶页会介绍受 feature 门控的各个类型族,比如 gemm_cplxgemm_i8gemm_fused。支持的最低 ndarray 版本是 0.17.1

核心入口

3 个函数覆盖了实数的普通路径。dot 是便捷入口。它把 A * B 算进一个新分配的行主序 Array2,读起来和 ndarray 自带的 .dot() 一样。

#![allow(unused)]
fn main() {
use ndarray::array;

let a = array![[1.0_f32, 2.0], [3.0, 4.0]];
let b = array![[5.0_f32, 6.0], [7.0, 8.0]];
let c = gemmkit_ndarray::dot(&a, &b);
assert_eq!(c, array![[19.0, 22.0], [43.0, 50.0]]);
}

dotT: GemmScalar 泛型,也就是无条件支持的 f32f64,再加上开启 half feature 后的 f16bf16。它用 Parallelism::default() 并行,并自行分配输出。当你还没有目标矩阵、只想算一次乘积时,用它就好。

gemm 就地写入通用形式 C <- alpha*A*B + beta*Calphabeta、一个已有的累加器,以及显式的 Parallelism,都在这里出场。它的签名是:

#![allow(unused)]
fn main() {
pub fn gemm<T, S1, S2, SC>(
    alpha: T,
    a: &ArrayBase<S1, Ix2>,
    b: &ArrayBase<S2, Ix2>,
    beta: T,
    c: &mut ArrayBase<SC, Ix2>,
    par: Parallelism,
)
where
    T: GemmScalar,
    S1: Data<Elem = T>,
    S2: Data<Elem = T>,
    SC: DataMut<Elem = T>;
}

输出约束为 SC: DataMut,所以 C&mut Array2ArrayViewMut2,并且和输入一样,可以是任意布局。下面 A 是一块行主序缓冲区,被无拷贝地转置成列主序视图。乘法以单线程运行:

#![allow(unused)]
fn main() {
use gemmkit_ndarray::Parallelism;
use ndarray::{Array2, array};

// 行主序存储,无拷贝地转置成列主序视图
let a = Array2::from_shape_vec((2, 2), vec![1.0_f32, 2.0, 3.0, 4.0])
    .unwrap()
    .reversed_axes();
let b = Array2::from_elem((2, 2), 1.0_f32);
let mut c = Array2::zeros((2, 2));
gemmkit_ndarray::gemm(1.0, &a, &b, 0.0, &mut c, Parallelism::Serial);
assert_eq!(c, array![[4.0, 4.0], [6.0, 6.0]]);
}

零代价的布局

因为适配器只读步长,ndarray 能表达的任何二维视图都能无拷贝转发,包括:

  • 标准的 C 序(行主序)布局
  • 来自 .reversed_axes().t(),或用 .f() 构造的数组的 F 序(列主序)视图
  • 带非单位步长的 .slice(...) 窗口视图
  • 来自负步长切片(如 s![..;-1, ..])的反转视图,会产生一个负的行步长

目标 C 同样自由。Array2::zeros((m, n).f()) 给出一个列主序输出,gemm 会直接填充它。

这里的“零拷贝”指的是适配器从不为规整布局而拷贝。当微内核需要连续的 panel 时,gemmkit 引擎仍会把操作数打包进自己的临时缓冲区。这种内部打包是算法的一部分,不是对转置输入的物化。关键在于:无论你的数组是什么样子,你都不必为了满足调用而付出一次 to_owned() 或手动转置。

panic:只查形状,不查别名

适配器只校验形状,别的一概不查。每个入口都断言内维对齐,且 C 与乘积匹配。不满足时,它会以 gemmkit-ndarray: 开头的信息 panic,并指出出错的维度,例如 A.cols (k) != B.rows (kb)A.rows (m) != C.rows (cm),或 B.cols (n) != C.cols (cn)。维度不匹配是普通 gemmdot panic 的唯一原因。

适配器不在运行时检查别名,也无需检查。C&mut ArrayBase<SC, _> 传入,这是一个独占借用。类型系统已经保证它不会与 AB 的共享 & 借用重叠。这正是 gemmkit 的 _unchecked 引擎要求调用方维持的前提条件,而 &mut 签名免费地维持了它。fused 入口多一项运行时检查,即偏置切片是否与 C 重叠。这项检查见进阶页

选择并行度

Parallelism 由适配器重新导出。Parallelism::Serial 在调用线程上运行。Parallelism::Rayon(n) 使用一个至多 n 线程的 rayon 线程池,Rayon(0) 会自动探测机器的核数。Parallelism::default() 就是 Rayon(0),也是 dot 所用的值,因此 dot 开箱即并行。并行路径需要默认开启的 parallel feature。关掉这个 feature 后,把每次调用都当作串行处理。

分块方式和任务顺序都不依赖线程数,所以固定的输入和配置会给出可复现的结果。串行和并行运行如今逐位一致,因为两者走的是同一套分块和同一个内核。这种一致只是当前实现的一个特性,不是跨配置的承诺。线程数背后的取舍见并行实践

复用工作区

每个会分配的入口,在调用期间都从 gemmkit 内部的线程本地池借用临时空间。单次 gemm 调用绝不会把一次分配泄漏进你的稳态。当你要跑一个形状相近的热循环时,_with 变体让你转而自己持有这块临时空间。把一个 &mut Workspace 作为首个参数传入。它会一次性增长到循环所需的最大尺寸,之后被复用,不再分配。

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Parallelism, Workspace};
use ndarray::Array2;

let mut ws = Workspace::new();
let par = Parallelism::default();
for &(m, k, n) in &[(256, 256, 256), (512, 128, 512)] {
    let a = Array2::<f32>::zeros((m, k));
    let b = Array2::<f32>::zeros((k, n));
    let mut c = Array2::<f32>::zeros((m, n));
    // 跨迭代复用 ws,至多分配一次
    gemmkit_ndarray::gemm_with(&mut ws, 1.0, &a, &b, 0.0, &mut c, par);
}
}

除了首个工作区参数,gemm_withgemm 完全一致,给出相同的结果。适配器里的每个类型族都有对应的 _with 孪生入口。这个模式同样适用于接下来要介绍的整数、复数、fused、批量和预打包入口。如果你要用一个固定的权重矩阵去乘一串激活,工作区会与进阶页上的预打包操作数路径自然搭配。

ndarray适配器进阶用法

除了普通的实数乘积,适配器还镜像了 gemmkit 的整个表面:

  • 整数 GEMM
  • 面向量化推理的重量化输出
  • 带可选共轭的复数乘积
  • 融合的偏置与激活
  • 用户自定义的逐元素映射
  • 在三维数组上的批量乘法
  • 预打包操作数

适配器让每个类型族都由与之同名的 Cargo feature 门控。每个类型族也都保持快速上手页里那些普通入口的形态:直接从数组读步长,转发给 gemmkit,仅在维度不匹配时 panic。fused 入口再多一种 panic 情形:偏置切片与 C 重叠。每个入口也都有一个携带调用方自有 Workspace_with 孪生。

整数 GEMM(int8

gemm_i8i8 输入乘进一个 i32 累加器:C(i32) <- alpha*A(i8)*B(i8) + beta*C,其中 alphabetaC 都是 i32。它是独立于 gemm 的入口,因为输入与输出的元素类型不同。算术在溢出时回绕,这是整数 GEMM 的惯例语义。dot_i8 是它的便捷孪生。它返回一个新建的 Array2<i32>

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Parallelism, dot_i8, gemm_i8};
use ndarray::Array2;

let a = Array2::<i8>::zeros((16, 12));
let b = Array2::<i8>::zeros((12, 10));

// i8 输入,i32 累加器
let c: Array2<i32> = dot_i8(&a, &b);

// 带 i32 alpha/beta 的通用形式,累加进已有累加器
let mut acc = Array2::<i32>::zeros((16, 10));
gemm_i8(2, &a, &b, 1, &mut acc, Parallelism::Serial);
}

重量化输出(int8 + epilogue

量化推理很少想要原始的 i32 累加器,它要的是一个 8 位张量。gemm_i8_requant 把乘法和重量化融进 1 趟。它直接把 i32 累加器折叠成 i8 输出,全程不物化完整的 m*n 中间结果。它没有 alpha 参数,因为其折进了 scale。它也没有 beta 参数,因为累加进量化输出没有良好定义。参数装在一个 Requantize 里:

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Parallelism, RequantScale, Requantize, gemm_i8_requant, gemm_i8_requant_u8};
use ndarray::Array2;

let a = Array2::<i8>::zeros((16, 12));
let b = Array2::<i8>::zeros((12, 10));

// i8 输出,范围 [-128, 127],单个 per-tensor scale,per-row 偏置(长度 A.rows)
let bias: Vec<i32> = vec![0; 16];
let mut c = Array2::<i8>::zeros((16, 10));
let req = Requantize {
    scale: RequantScale::PerTensor(0.05),
    zero_point: -7,
    bias: Some(&bias),
};
gemm_i8_requant(&a, &b, req, &mut c, Parallelism::default());

// u8 输出,范围 [0, 255],per-channel scale,无偏置
let scales: Vec<f32> = vec![0.02; 16]; // 每个输出行 / 通道一个
let mut cu = Array2::<u8>::zeros((16, 10));
gemm_i8_requant_u8(
    &a,
    &b,
    Requantize { scale: RequantScale::PerRow(&scales), zero_point: 128, bias: None },
    &mut cu,
    Parallelism::default(),
);
}

输出为 clamp(zero_point + round_ne(scale * (accumulator + bias[i])), LO, HI),采用四舍六入五成双。其中 scale 是 per-tensor 的那个值,或 per-row 的 scale_iu8 变体是 ONNX-QLinearMatMul 风格的激活:除了输出域 [0, 255]zero_point 的取值范围外,与 gemm_i8_requant 完全相同。两者都会拒绝:

  • 非有限或非正的 scale,per-tensor 或 per-row
  • 长度不等于 A.rows 的 per-row scale 或偏置
  • C 重叠的切片
  • 超出该入口取值域的 zero_point

复数 GEMM(complex

复数乘积有自己的入口,因为这 2 个共轭标志放不进同构的实数签名。gemm_cplx 计算 C <- alpha*op(A)*op(B) + beta*C,元素类型为 Complex<f32>Complex<f64>。设置 conj_a 时,op(A)conj(A)。设置 conj_b 时,op(B)conj(B)dot_cplx 是不做共轭的便捷入口。

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Complex, Parallelism, dot_cplx, gemm_cplx};
use ndarray::Array2;

type C = Complex<f64>;
let a = Array2::<C>::from_elem((8, 6), Complex::new(0.0, 0.0));
let b = Array2::<C>::from_elem((6, 5), Complex::new(0.0, 0.0));

// 普通 A*B
let c = dot_cplx(&a, &b);

// 对 A 取共轭,累加进已有 C
let mut acc = Array2::<C>::from_elem((8, 5), Complex::new(0.0, 0.0));
gemm_cplx(
    Complex::new(1.0, 0.0),
    &a,
    true,  // conj_a
    &b,
    false, // conj_b
    Complex::new(0.0, 0.0),
    &mut acc,
    Parallelism::Serial,
);
}

同时开启 complexepilogue 后,gemm_cplx_fused 会在同一趟里加上一个可选的 Bias(原样相加,绝不共轭)。它不接受激活参数。像 ReLU 这样带序关系的激活,在复数上没有定义。

融合偏置、激活与映射(epilogue

gemm_fused 在 1 趟里算出 C <- act(alpha*A*B + beta*C + bias)。偏置是可选的 Bias::PerRow(长度 A.rows)或 Bias::PerCol(长度 B.cols)。激活是可选的 ReluLeakyRelu(slope),最后施加。两者都设为 None 时,gemm_fused 就是 gemm

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Activation, Bias, Parallelism, gemm_fused, gemm_map};
use ndarray::Array2;

let a = Array2::<f32>::zeros((12, 9));
let b = Array2::<f32>::zeros((9, 7));

// 一趟算出 C <- ReLU(A*B + bias);PerRow 偏置长度为 A.rows
let bias: Vec<f32> = vec![0.0; 12];
let mut c = Array2::<f32>::zeros((12, 7));
gemm_fused(
    1.0, &a, &b, 0.0, &mut c,
    Some(Bias::PerRow(&bias)),
    Some(Activation::Relu),
    Parallelism::default(),
);

// 任意逐元素闭包 f(value, row, col);这里是一个 relu6
let f = |v: f32, _r: usize, _c: usize| v.max(0.0).min(6.0);
let mut c2 = Array2::<f32>::zeros((12, 7));
gemm_map(1.0, &a, &b, 0.0, &mut c2, &f, Parallelism::default());
}

BiasActivationgemmkit_ndarray 重新导出,所以你无需为它们再点名 gemmkit。对 f32/f64,无论什么形状,gemm_fused 都与“先 gemm 再做同样的标量映射”逐位相同。对 f16/bf16,尾部运算会先以 f32 进行,再做一次窄化。这比单独的窄化映射更精确,所以对这些类型而言,结果与“先 gemm 再映射”并不逐位相等。

gemm_map 是通用的逐元素扩展点。闭包 f(value, row, col) 看到的是每个输出元素的最终值,(row, col) 处在 C 的用户坐标系里。gemmkit 对每个元素恰好调用它一次。它的代价是每个元素 1 次间接调用。普通的偏置或激活优先用 gemm_fused,因为它会向量化。GELU、sigmoid、clamp,或依赖位置的变换,才用 gemm_map。这里的 T 只能是 f32/f64

批量 GEMM

这是唯一没有普通 gemm 对应、在同类适配器里也没有对手的运算。它是一叠彼此独立的乘积,承载在三维 Array3 上,批次维在 0 轴。a(batch, m, k)b(batch, k, n)c(batch, m, n)。0 轴是各操作数的批次步长。1、2 轴是元素步长。gemm_batched 在批次上并行。每个元素都在 1 个 worker 上运行,因此结果精确复现一个 gemm 调用循环。

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Parallelism, dot_batched, gemm_batched};
use ndarray::Array3;

let a = Array3::<f32>::zeros((32, 8, 5)); // (batch, m, k)
let b = Array3::<f32>::zeros((32, 5, 6)); // (batch, k, n)

// 一叠乘积
let c = dot_batched(&a, &b); // (32, 8, 6)

// 累加进已有累加器的通用形式
let mut acc = Array3::<f32>::zeros((32, 8, 6));
gemm_batched(0.7, &a, &b, 1.3, &mut acc, Parallelism::default());
}

适配器只读步长,所以一个换轴的、或本就一般步长的三维视图,都能无拷贝转发。比如,a.view().permuted_axes([0, 2, 1]) 把一块 (batch, k, m) 缓冲区变成 (batch, m, k) 视图,并直接批量转发。

epilogue 下,gemm_batched_fused 对这叠里的每个元素都施加同一个共享的 Bias/Activation。这正是批量线性层的情形。偏置按单个元素定尺寸(PerRow 长度 mPerCol 长度 n),而不是整批。

预打包操作数

当一个操作数固定、另一个成流而来时,把固定那一侧打包一次并复用。这样就省掉了每次调用的重复打包。prepack_rhs 为复用的 B 返回一个 PackedRhs<T>,由 gemm_packed_b 消费。prepack_lhs 为复用的 A 返回一个 PackedLhs<T>,由 gemm_packed_a 消费。打包函数直接读步长,所以 BA 可以是任意布局。

两者各有 1 条朝向约束。gemm_packed_b 需要一个偏列主序的 C|col stride| >= |row stride|)。gemm_packed_a 需要一个偏行主序的 C|col stride| <= |row stride|)。另一种朝向会交换操作数,并使打包句柄失效,gemmkit 会拒绝这种情况。不合适的布局请用普通 gemm

融合孪生 gemm_packed_b_fusedgemm_packed_a_fused 接受同样的句柄,再加上偏置和激活。这正是固定权重的推理层。下面的例子把一个权重矩阵作为 LHS 打包一次,并在各推理步之间复用。它融进了一个 per-output-channel 的偏置和一个 ReLU:

#![allow(unused)]
fn main() {
use gemmkit_ndarray::{Activation, Bias, Parallelism, gemm_packed_a_fused, prepack_lhs};
use ndarray::Array2;

let (out, in_features) = (256usize, 512usize);

// 把固定权重 W: (out, in) 打包一次
let w = Array2::<f32>::zeros((out, in_features));
let packed = prepack_lhs(&w);
let bias: Vec<f32> = vec![0.0; out]; // per-output-channel,长度 C.rows

// 每个推理步:激活 x (in, batch) -> y (out, batch)
let batch = 32;
let x = Array2::<f32>::zeros((in_features, batch));
let mut y = Array2::<f32>::zeros((out, batch)); // 行主序(packed_a 朝向)
gemm_packed_a_fused(
    1.0,
    &packed,
    &x,
    0.0,
    &mut y,
    Some(Bias::PerRow(&bias)),
    Some(Activation::Relu),
    Parallelism::default(),
);
}

把打包句柄与 _with 工作区变体 gemm_packed_a_fused_with 搭配。这样一来,一个稳定的推理循环在首次调用后就不再分配。你在用户坐标系里指定偏置的轴。打包路径对 gemm_packed_b_fused 原样转发它,而对 gemm_packed_a_fused 则让核心去翻转它。无论你打包了哪个操作数,PerRow 始终表示“每个输出行 1 个值”。

本适配器与 ndarray 自带乘积的取舍

ndarray 本就能做矩阵乘法:.dot() 给出普通乘积,general_mat_mul 给出就地的 alpha/beta 形式。对一次没有额外需求的 f32/f64 乘积,这些函数就够用,还能少拉 1 个依赖。没必要出于习惯就走 gemmkit。

当你需要 ndarray 内建路径给不了的东西时,再选本适配器。gemmkit 会在运行的机器上,于运行时挑选最快的指令集。它不会在编译期就把某个选择写死(见运行时ISA分发)。它带来本页覆盖的更宽表面:融合偏置与激活、i8 与重量化推理、带共轭的复数、批量乘积,以及预打包。它还暴露调优旋钮,外加一个把分块校准到部署机器的安装期自动调优器(见调优旋钮)。

这些都不会改变你传入的数组或拿回的结果。适配器自始至终是同一套零拷贝步长转接。它只是拓宽了你能提出的请求。

在nalgebra中使用gemmkit

gemmkit-nalgebra 让你直接用 nalgebra 的矩阵驱动 gemmkit 引擎,不必先拷贝一份。它面向 nalgebra 0.35,接收 &Matrix<T, R, C, S>,其中存储类型满足 S: RawStorage<T, R, C> 即可。拥有所有权的 DMatrix、静态的 SMatrix,以及各种视图和切片类型都满足要求。适配器读出矩阵的数据指针和它的 2 个步长,交给 gemmkit 的底层引擎。引擎不会重排输入,不会把它转置进临时缓冲区,也不会复制一份。

nalgebra 的自然布局是列主序,这也正是 gemmkit 偏好的朝向,所以最常见的情形恰好就是最快的情形。行主序视图和一般步长视图同样能用,拷贝次数照样为零,因为引擎直接读取步长,而不是假定某种布局。

加入项目

Cargo.toml 里需要 2 个 crate。适配器把自己签名里出现的类型都重新导出了,所以常规配置不需要直接依赖 gemmkit

[dependencies]
gemmkit-nalgebra = "0.1"
nalgebra = "0.35"

gemmkit_nalgebra 重新导出了好几种类型,因此项目不需要为它们直接依赖 gemmkit

  • Parallelism 选择器,以及每个 _with 变体所需的 Workspace
  • 融合选择器 BiasActivation
  • 预打包句柄 PackedLhsPackedRhs
  • 重量化参数 RequantizeRequantScale
  • 元素类型约束 GemmScalarFusedScalarMapScalarComplexScalar。当你要写一个对某个入口泛型的封装时,需要用到它们
  • 对应 feature 下的元素类型 f16bf16Complexc32c64,因此 halfnum-complex 也不必进入你的 manifest
  • tuning 模块

请通过适配器去用 tuning,而不是自己再依赖一份 gemmkit。这些调优旋钮是进程级的全局原子量。第二份单独解析出来的 gemmkit 会给你一组适配器根本不会读取的原子量。

默认 feature 集启用了 parallel,它会打开引擎里基于 rayon 的线程化(gemmkit/parallel)。适配器的每个 feature 都只是对 gemmkit 上同名 feature 的一层转发:

  • half 增加 f16/bf16 输入
  • complex 增加 Complex<f32>/Complex<f64>
  • int8 增加 i8 -> i32 路径
  • epilogue 增加融合的偏置/激活以及逐元素映射
  • wasm_threadswasm32-wasip1-threads 上叠加 parallel

这些 feature 门控的入口在 nalgebra 适配器进阶用法 中讲解。本页只谈始终可用的实数标量接口。

3 个实数标量入口

基础接口是 3 个函数,都对 GemmScalar 泛型(它始终是 f32f64,在 half feature 下另加 f16bf16)。gemm 是做累加的乘法,gemm_with 是复用调用方持有的工作区的同一调用,dot 则是一个自行分配结果的便捷封装。

下面是 gemm 的原样签名,摘自 gemmkit-nalgebra/src/float.rs

#![allow(unused)]
fn main() {
pub fn gemm<T, R1, C1, S1, R2, C2, S2, RC, CC, SC>(
    alpha: T,
    a: &Matrix<T, R1, C1, S1>,
    b: &Matrix<T, R2, C2, S2>,
    beta: T,
    c: &mut Matrix<T, RC, CC, SC>,
    par: Parallelism,
) where
    T: GemmScalar,
    R1: Dim,
    C1: Dim,
    S1: RawStorage<T, R1, C1>,
    R2: Dim,
    C2: Dim,
    S2: RawStorage<T, R2, C2>,
    RC: Dim,
    CC: Dim,
    SC: RawStorageMut<T, RC, CC>,
{
    gemm_common(None, alpha, a, b, beta, c, par);
}
}

10 个泛型参数看着吓人,其实只说了一件简单的事。ABC 各自可以是任意 nalgebra 矩阵或视图,行维度、列维度和存储类型三者相互独立。AB 通过 RawStorage 读取。C 需要 RawStorageMut,因为 gemm 会就地写入它。运算是 C <- alpha*A*B + beta*C

gemm_with 的参数相同,只是最前面多一个 &mut Workspacedot(a, b) -> DMatrix<T>A*B 算进一块新分配的列主序矩阵。它内部以 beta == 0 调用 gemm,所以那块新缓冲区在被覆盖之前从不会被读取。

用 DMatrix 做第一次乘法

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::Parallelism;
use nalgebra::DMatrix;

let a = DMatrix::from_row_slice(2, 2, &[1.0_f32, 2.0, 3.0, 4.0]);
let b = DMatrix::from_row_slice(2, 2, &[5.0_f32, 6.0, 7.0, 8.0]);

// dot: A*B 写进一块新的列主序 DMatrix
let c = gemmkit_nalgebra::dot(&a, &b);
assert_eq!(c, DMatrix::from_row_slice(2, 2, &[19.0, 22.0, 43.0, 50.0]));

// gemm: 就地累加 C <- alpha*A*B + beta*C
let mut acc = DMatrix::<f32>::zeros(2, 2);
gemmkit_nalgebra::gemm(1.0, &a, &b, 0.0, &mut acc, Parallelism::default());
assert_eq!(acc, c);
}

当你只想要乘积、也乐意拿回一个 DMatrix 时,dot 是趁手的工具。当你已经拥有目标矩阵时,gemm 才是该用的工具。用它来缩放目标矩阵(beta)、缩放乘积(alpha),或者省掉 dot 那次分配。dot 始终返回 DMatrix<T>,即便输入是静态矩阵也一样,因为封装只有在值的层面才知道输出维度。

静态矩阵与混合形状

静态矩阵走同样的函数,没有任何特判。由于每个操作数的行、列、存储泛型相互独立,静态的 A 可以乘一个动态的 Bgemm 写入静态的 &mut SMatrix 输出,也和写入 DMatrix 一样自然。

#![allow(unused)]
fn main() {
use nalgebra::{DMatrix, Matrix2, SMatrix};

// 静态 x 静态
let a = Matrix2::new(1.0_f32, 2.0, 3.0, 4.0);
let b = Matrix2::new(5.0_f32, 6.0, 7.0, 8.0);
let c = gemmkit_nalgebra::dot(&a, &b); // -> DMatrix<f32>
assert_eq!(c[(0, 0)], 19.0);

// 静态 A x 动态 B:相互独立的 Dim 泛型使其成立
let a34 = SMatrix::<f64, 3, 4>::from_fn(|i, j| (i as f64) - 0.5 * (j as f64) + 1.0);
let b = DMatrix::<f64>::from_element(4, 2, 0.25);
let c = gemmkit_nalgebra::dot(&a34, &b); // -> DMatrix<f64>,形状 3x2
}

布局与零拷贝

适配器从不拷贝操作数。它从矩阵中取出 (rows, cols, row-stride, col-stride),把指针连同步长转交给引擎。源布局只决定引擎看到的是哪一组步长,不决定是否发生分配。一个列主序的 DMatrix 是就地读取的,一个用 from_slice_with_strides 构造的行主序切片也是。一个非连续的跳步视图,比如某个更大矩阵的每隔一行,同样是就地读取的。nalgebra 以非负的元素个数报告步长,适配器把它们拓宽成引擎所需的带符号步长。

引擎为喂饱其微内核所做的内部打包,与源布局无关,无论数据从哪来都会发生。这是 gemmkit 的性质,不是适配器引入的拷贝。如果你想消除对某个被反复使用的操作数的重复内部打包,就改用预打包操作数路径。它在进阶用法页中有说明。

何时 panic

这些入口会校验形状,遇到不一致就 panic,并在消息里带上出错的维度。gemm(及其同族)在触碰任何内存之前,先检查 3 个等式:A.cols == B.rowsA.rows == C.rowsB.cols == C.cols。比如内维不匹配时,会以 gemmkit-nalgebra: A.cols (k) != B.rows (kb) 中止,而不是越界读取。因此输出矩阵必须事先具备正确形状。gemm 会写入它,但不会调整其大小。而自行分配输出的 dot 只可能在内维检查上失败。

选择并行方式

每个调用都以一个 Parallelism 作为最后一个参数。它有 2 个变体:Parallelism::Serial 单线程运行,Parallelism::Rayon(n) 在 rayon 上以至多 n 个线程运行,其中 Rayon(0) 自动探测线程数。DefaultRayon(0),也正是 dot 内部所用的。

对于小矩阵,或者当你已经身处一个并行区域、想避免嵌套线程化时,就传 Parallelism::Serial。对于空闲机器上的大乘法,Parallelism::Rayon(0) 能让引擎把工作铺开。线程化策略,以及引擎如何挑选线程数,见并行实践

复用工作区

引擎需要暂存空间来打包 AB 的分块。默认情况下,它从一个线程局部的池子借用这块空间,因此在稳态下 gemmdot 每次调用都不会自行分配。当你在紧凑循环里做很多次乘法、想完全掌控那块缓冲区时,gemm_with 会接收一个你自己持有并复用的 &mut Workspace

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Parallelism, Workspace};
use nalgebra::DMatrix;

let mut ws = Workspace::new();
let a = DMatrix::<f64>::from_element(64, 64, 1.0);
let b = DMatrix::<f64>::from_element(64, 64, 2.0);

for _ in 0..1000 {
    let mut c = DMatrix::<f64>::zeros(64, 64);
    gemmkit_nalgebra::gemm_with(&mut ws, 1.0, &a, &b, 0.0, &mut c, Parallelism::default());
}
}

同一个 Workspace 可以在各次迭代里支撑不同形状的乘法。它会增长到能容纳所见过的最大那次,并保留该容量。Workspace::new() 从空开始,在第一次调用把它填满之前不花任何代价。_with 形式在所有做累加的入口上都有,包括 feature 门控的那些。同样的复用模式也适用于整数、复数和融合调用。

nalgebra适配器进阶用法

除了 在 nalgebra 中使用 gemmkit 里讲的实数标量 gemm/gemm_with/dot,适配器还在 Cargo feature 之后暴露了引擎的完整接口:

  • 整数 GEMM
  • 复数 GEMM
  • 融合 epilogue
  • 重量化输出
  • 逐元素映射
  • 预打包操作数
  • 批量

每个入口都恪守适配器的核心承诺:直接读取 nalgebra 的指针和步长,不做任何拷贝。它还会镜像一个同名的 gemmkit 核心函数。每个类型族都由 Cargo feature 门控,你只为开启的部分付出代价。

Feature增加
halff16/bf16 走同一套 gemm/gemm_fused 泛型
int8gemm_i8gemm_i8_withdot_i8i8 -> i32
complexgemm_cplxgemm_cplx_withdot_cplx
epiloguegemm_fusedgemm_map(及预打包融合孪生)
int8 + epiloguegemm_i8_requantgemm_i8_requant_u8
complex + epiloguegemm_cplx_fused

这些入口用到的每个辅助类型都来自适配器 crate 本身,所以你无需点名 gemmkit 即可使用它们。它们是 BiasActivationRequantScaleRequantizePackedLhsPackedRhsParallelismWorkspace,以及 Complex(及其 c32/c64 别名)。

整数 GEMM

int8 下,gemm_i8 把 2 个 i8 矩阵相乘成一个 i32 输出。输入是 i8alphabetaCi32,因为 i8*i8 的乘积需要更宽的累加器。算术在溢出时回绕,这是整数 GEMM 的惯例。它之所以是独立于 gemm 的入口,是因为输入和输出的元素类型不同。

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Parallelism, dot_i8, gemm_i8};
use nalgebra::DMatrix;

let a = DMatrix::from_row_slice(2, 3, &[1_i8, 2, 3, 4, 5, 6]);
let b = DMatrix::from_row_slice(3, 2, &[1_i8, 0, 0, 1, 1, 1]);

// dot_i8:A*B 写进一块新的 DMatrix<i32>
let c = dot_i8(&a, &b);

// gemm_i8:缩放并累加进 i32 输出
let mut acc = DMatrix::<i32>::zeros(2, 2);
gemm_i8(1, &a, &b, 0, &mut acc, Parallelism::Serial);
assert_eq!(acc, c);
}

dot_i8(a, b) -> DMatrix<i32> 是做分配的便捷形式,gemm_i8_with 则为定成本的量化推理循环复用调用方持有的 Workspace

重量化输出

重量化把量化推理里“反量化-缩放-取整-夹取”这一步折进 GEMM,于是那个 m*ni32 累加器从不必被完整物化。gemm_i8_requant 接收 i8 输入、写出 i8 输出。gemm_i8_requant_u8 则写出无符号的 u8 输出(ONNX QLinearMatMul 的约定)。两者都需要 int8 + epilogue。没有 alpha(它折进了 scale),也没有 beta(往一个已经量化的 C 上累加没有良好定义)。参数装在一个被重新导出的 Requantize 里:

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Parallelism, RequantScale, Requantize, gemm_i8_requant};
use nalgebra::DMatrix;

let a = DMatrix::from_row_slice(2, 3, &[10_i8, -4, 7, 3, 8, -2]);
let b = DMatrix::from_row_slice(3, 2, &[2_i8, 1, -1, 5, 4, 0]);
let bias = [100_i32, -50]; // 逐行,长度为 A.rows

let req = Requantize {
    scale: RequantScale::PerTensor(0.05),
    zero_point: -7,        // i8 输出取值域 [-128, 127]
    bias: Some(&bias),
};
let mut c = DMatrix::from_element(2, 2, 0_i8);
gemm_i8_requant(&a, &b, req, &mut c, Parallelism::Serial);
}

RequantScale::PerTensor(s) 对每个元素施加同一个 scale。RequantScale::PerRow(&[f32]) 给每个输出行(即每个输出通道,标准的逐通道约定)一个 scale,长度为 A.rows。每个 scale 都必须有限且 > 0。引擎会在取整之后把 zero_point 当作整数加入,它必须落在输出域内:gemm_i8_requant[-128, 127]gemm_i8_requant_u8[0, 255]。引擎还会把可选的逐行 i32 偏置(长度 A.rows)在缩放之前加到累加器上。

适配器会校验以上全部。它会在任何违规时以核心引擎的措辞 panic:

  • scale 非有限或非正
  • 逐行 scale 或 bias 长度不对
  • zero_point 越界
  • scale 或 bias 切片与 C 重叠

复数 GEMM

complex 下,gemm_cplxT = Complex<f32>Complex<f64> 计算 C <- alpha*op(A)*op(B) + beta*C。置起 conj_a 标志时 op(A) = conj(A),置起 conj_b 标志时 op(B) = conj(B)。这两个共轭标志正是复数需要独立入口的原因:它们塞不进同质的实数标量签名。dot_cplx(a, b) 是不做共轭的 A*B 便捷形式。要做共轭乘积,直接用 gemm_cplx

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Complex, Parallelism, dot_cplx, gemm_cplx};
use nalgebra::DMatrix;

type C = Complex<f64>;
let a = DMatrix::from_element(2, 2, C::new(1.0, 1.0));
let b = DMatrix::from_element(2, 2, C::new(0.0, -1.0));

// 普通乘积
let p = dot_cplx(&a, &b);

// 对 A 取共轭、B 不变,累加进已有的 C
let mut acc = DMatrix::from_element(2, 2, C::new(0.0, 0.0));
gemm_cplx(C::new(1.0, 0.0), &a, true, &b, false,
          C::new(0.0, 0.0), &mut acc, Parallelism::Serial);
}

带融合偏置的复数

gemm_cplx_fused(需要 complex + epilogue)在复数乘积的同一趟里加上偏置:C <- alpha*op(A)*op(B) + beta*C + bias。偏置是被重新导出的 Bias,要么是 Bias::PerRow(长度 A.rows),要么是 Bias::PerCol(长度 B.cols)。gemmkit 会原样加上它,不做共轭。这里刻意没有激活参数:像 ReLU 这样带次序的激活在复数上没有定义。当 bias == None 时,该调用就等同于 gemm_cplx

融合 epilogue

epilogue feature 增加了 gemm_fused,它在 f32/f64 上以单趟计算 C <- act(alpha*A*B + beta*C + bias)。开启 half 时,f16/bf16 也走同一条路径。它们的 epilogue 在 f32 中求值,最后只做一次收窄存储。可选的 BiasPerRow(长度 A.rows)或 PerCol(长度 B.cols)。可选的 ActivationReluLeakyRelu(slope),最后施加。两者都传 None 时,逐位等同于普通 gemm

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Activation, Bias, Parallelism, gemm_fused};
use nalgebra::DMatrix;

let a = DMatrix::<f32>::from_element(12, 9, 0.5);
let b = DMatrix::<f32>::from_element(9, 7, -0.25);
let bias: Vec<f32> = (0..12).map(|i| 0.5 * i as f32 - 2.0).collect();
let mut c = DMatrix::<f32>::zeros(12, 7);

gemm_fused(1.3, &a, &b, -0.7, &mut c,
           Some(Bias::PerRow(&bias)), Some(Activation::Relu), Parallelism::Serial);
}

融合这一趟不只是图方便。它省掉了对 C 的第二次扫描,也省掉了分开做偏置加法和激活时要付的那趟内存往返。对 f32/f64,其结果与先跑 gemm、再逐元素施加同样的偏置和激活是逐位一致的。你可以放心采用它,而不改变数值结果。它背后的设计见 融合 Epilogue

逐元素映射

gemm_map 适用于套不进“偏置加激活”形状的 epilogue。它把一个任意闭包施加到每个算完的输出元素上:C[r, c] <- f(alpha*A*B + beta*C, r, c)。这个闭包每个元素恰好触发一次,其中 (r, c) 处于 C 的用户坐标系。T 只能是 f32/f64。闭包类型是 &(dyn Fn(T, usize, usize) -> T + Sync)。它必须是 Sync 才能并行运行,并且可以按引用捕获数据。

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Parallelism, gemm_map};
use nalgebra::DMatrix;

let a = DMatrix::<f64>::from_element(8, 6, 0.3);
let b = DMatrix::<f64>::from_element(6, 5, 0.4);
let mut c = DMatrix::<f64>::zeros(8, 5);

// sigmoid,忽略位置
let sigmoid = |v: f64, _r: usize, _c: usize| 1.0 / (1.0 + (-v).exp());
gemm_map(1.0, &a, &b, 0.0, &mut c, &sigmoid, Parallelism::Serial);
}

对普通的偏置或 ReLU,优先用 gemm_fused,因为它会向量化。gemm_map 是通用的扩展点,覆盖 GELU、sigmoid、夹取、依赖位置的变换,代价是每个输出元素要付 1 次间接调用。和融合入口一样,它在 f32/f64 上的结果与先跑 gemm、再施加同一映射是逐位一致的。gemm_map_with 复用一个 Workspace

预打包操作数

举例来说,一个权重矩阵对着一串激活服务时,就有 1 个操作数在许多次乘法中固定不变。把它预打包一次,就能免掉每次调用的重打包。prepack_rhs(b) -> PackedRhs<T> 打包右操作数。gemm_packed_b 随后以该句柄代替 Bprepack_lhs/gemm_packed_a 对固定的左操作数做镜像。

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Parallelism, gemm_packed_b, prepack_rhs};
use nalgebra::DMatrix;

let weights = DMatrix::<f32>::from_fn(64, 32, |i, j| 0.01 * (i as f32 - j as f32));
let packed = prepack_rhs(&weights); // 把固定的 B 打包一次

for step in 0..100 {
    let x = DMatrix::<f32>::from_element(16, 64, step as f32); // 一批激活
    let mut y = DMatrix::<f32>::zeros(16, 32);                 // 列主序输出
    gemm_packed_b(1.0, &x, &packed, 0.0, &mut y, Parallelism::default());
}
}

这里有一个朝向约束。gemm_packed_b 需要偏列主序的 C|列步长| >= |行步长|)。行主序的 C 会迫使引擎在内部交换 AB。这会使预打包的 RHS 失效,因此 gemmkit 会拒绝行主序的 Cgemm_packed_a 则相反:它需要偏行主序的 C,并拒绝列主序的。对于朝向不对的 C,退回到普通 gemm

每个预打包入口都有做工作区复用的 _with 孪生。在 epilogue 下,融合孪生 gemm_packed_b_fusedgemm_packed_a_fused 在同一句柄上再加偏置和激活。PackedRhsPackedLhs 暴露 .rows().cols(),方便你复核维度。底层的复用模型见 预打包操作数

批量 GEMM

nalgebra 没有三维数组类型,所以批量 GEMM 不像 ndarray 适配器那样接收一个三维张量。gemm_batched 改为以一个逐元素 (&A, &B) 输入对的切片,按位置搭配一个 &mut C 输出的切片。它在 1 次调用里,对每个元素运行 C_e <- alpha*A_e*B_e + beta*C_e,走 gemmkit 的指针数组引擎。alphabetapar 由整个批次共享。

#![allow(unused)]
fn main() {
use gemmkit_nalgebra::{Parallelism, gemm_batched};
use nalgebra::DMatrix;

let a = DMatrix::from_row_slice(2, 2, &[1.0_f32, 2.0, 3.0, 4.0]);
let b = DMatrix::from_row_slice(2, 2, &[5.0_f32, 6.0, 7.0, 8.0]);
let mut c = vec![DMatrix::<f32>::zeros(2, 2), DMatrix::<f32>::zeros(2, 2)];

let ab = [(&a, &b), (&a, &b)];
gemm_batched(1.0, &ab, 0.0, &mut c, Parallelism::Serial);
assert_eq!(c[0], DMatrix::from_row_slice(2, 2, &[19.0, 22.0, 43.0, 50.0]));
}

各元素的形状可以不同,因此异构批次也没问题,只要每个元素自身的维度自洽:A_e.cols == B_e.rows 等等。共享的存储类型承载着变化的运行期维度。因此用 DMatrix 或动态步长视图,就能覆盖异构形状和混合布局,而且都是同一个类型。输入数量与输出数量必须相等:ab.len() == c.len()。数量不符,或任一元素维度不符,都会 panic。

批量是跨元素并行的,而不是在单个元素内部并行。分派器把整个逐元素 GEMM 分给各个 worker,每个 worker 都串行、缓存命中地跑完自己的 GEMM。因此,批量的结果重现了一个普通的 gemm 调用循环,在不同线程数下也保持可复现。串行调度和并行调度是逐位一致的,因为每个元素总是完整地跑在 1 个 worker 上。

C 切片是单一存储类型,单个共享融合 epilogue 没有对应的指针数组形式,所以不同于 ndarray 适配器,这里没有 gemm_batched_fused。批量模型的进一步讨论见 批量 GEMM

它与 nalgebra 自带乘法的分工

nalgebra 本来就会做矩阵乘法:&a * &ba.mul_to(&b, &mut c),以及其余的运算符接口。它们都返回类型规整的矩阵,并与它的常量泛型维度融为一体。对于一次普通的 f32/f64 乘积,尤其是小的静态矩阵,那些才是地道的选择,没有理由去动这个适配器。

适配器的价值在于 nalgebra 的运算符不提供的东西:

  • 引擎的运行时 SIMD 分发:在运行期而非编译期挑选机器上可用的最佳指令集
  • i8 -> i32 与重量化的整数路径
  • 单趟内的融合偏置/激活与逐元素 epilogue
  • 对许多小问题的批量乘法
  • 供固定权重跨调用复用的预打包操作数

当你的矩阵本就在 nalgebra 里、又需要以上任意一项时,就用这个适配器。它让你不离开 nalgebra 的类型、也不做拷贝,就拿到引擎的吞吐和功能。

在faer中使用gemmkit

gemmkit-faer 是一层很薄的零拷贝桥梁,把 faer 的视图类型接到 gemmkit 的 GEMM 引擎上。它对每个输入接受一个 MatRef<'_, T>,对输出接受一个 MatMut<'_, T>。它直接从视图里读出数据指针,以及以元素为单位的行、列步长,再把它们交给 gemmkit 的底层引擎。适配器在入口处不做任何转置、拷贝或重打包。

faer 存放步长的方式,正好就是 gemmkit 引擎所需要的。因此,一个 faer Mat、一个转置视图、一个带偏移的子矩阵,以及一个反转(负步长)视图,都会原样抵达内核。

本 crate 面向 faer 0.24,需要 Rust 1.89。

安装与 feature

gemmkit-faer 把自己签名里出现的所有类型都重新导出了。因此,常规配置里并不需要直接依赖 gemmkit

[dependencies]
gemmkit-faer = "0.1"
faer = "0.24"

gemmkit_faer 重新导出了:

  • Parallelism 选择器,以及每个 _with 变体都要用到的 Workspace 类型。
  • 融合选择器 BiasActivation
  • 预打包句柄 PackedLhsPackedRhs
  • 重量化参数 RequantizeRequantScale
  • 元素类型约束 GemmScalarFusedScalarMapScalarComplexScalar。当你要写一个对某个入口泛型的封装时,需要用到它们。
  • 对应 feature 下的元素类型 f16bf16Complexc32c64。这样 halfnum-complex 也不必进入你的 manifest。
  • tuning 模块。

请通过适配器来用 tuning,而不要自己再单独依赖一份 gemmkit。这些旋钮是进程级的全局原子量。另一份单独解析出来的 gemmkit 会给你一组不同的原子量,一组适配器根本不会读取的原子量。

每个 Cargo feature 都会转发到 gemmkit 中的同名 feature。所以你在这里启用某个元素类型族或某个融合入口时,底层核心也会一并启用它。

  • parallel(默认):基于 rayon 的并行。
  • wasm_threads:在 wasm32-wasip1-threads 上启用线程,同时也会启用 parallel
  • halff16bf16 元素类型,以 f32 累加。
  • complexc32c64 元素类型。
  • int8i8 输入进入 i32 输出。
  • epilogue:融合的偏置/激活、重量化,以及逐元素映射入口。

feature 门控的类型族与融合入口,在进阶用法页中介绍。本页只讲始终可用的 f32/f64(以及 half 下的 f16/bf16)这层接口。

这里的“零拷贝”指什么

每个入口都会经过同一个小小的辅助函数,从 MatRef 里取出原始的组成部分。faer 已经以 isize 按元素为单位报告步长,反转视图为负值,这正是 gemmkit 非检查引擎所期望的形状,所以完全不需要任何转换步骤。

#![allow(unused)]
fn main() {
// gemmkit-faer/src/common.rs
pub(crate) fn ref_parts<T>(a: MatRef<'_, T>) -> (usize, usize, isize, isize, *const T) {
    (a.nrows(), a.ncols(), a.row_stride(), a.col_stride(), a.as_ptr())
}
}

适配器自己校验三个共享维度,然后在一个 unsafe 块内,把指针和步长转发给 gemmkit 的 _unchecked 引擎。安全性论证很简短:faer 的视图类型保证指针加步长描述的是一个合法的、边界内的布局。输出是一个 MatMut,也就是独占借用,所以 C 不可能与 AB 存在别名。对于普通路径来说,这就是适配器的全部内容。

gemmkit 的缓存分块、ISA 分发、打包,以及并行调度,全都在核心里实现,核心那边也有相应的文档。内部机理见架构章节

gemm 与 dot

两个主力函数是 dotgemmdot 返回一个全新的乘积,gemm 则就地更新一个已有的输出。二者都对 GemmScalar 泛型:始终支持 f32f64,在启用 half feature 时还支持 f16bf16

#![allow(unused)]
fn main() {
use faer::Mat;

let a = Mat::from_fn(2, 2, |i, j| [[1.0_f64, 2.0], [3.0, 4.0]][i][j]);
let b = Mat::from_fn(2, 2, |i, j| [[5.0_f64, 6.0], [7.0, 8.0]][i][j]);
// A*B into a fresh column-major Mat
let c = gemmkit_faer::dot(a.as_dyn_stride(), b.as_dyn_stride());
assert_eq!(c[(0, 0)], 19.0);
assert_eq!(c[(1, 1)], 50.0);
}

dot(a, b)A*B 算进一个新分配的列主序 Mat 中。它以默认并行度运行,也就是 Parallelism::Rayon(0),会自动探测线程数。把 dot 当作一次性的便捷接口来用。当你自己持有输出缓冲区,或者想做通用更新时,改用 gemm

#![allow(unused)]
fn main() {
use faer::Mat;
use gemmkit_faer::{Parallelism, gemm};

let a = Mat::<f64>::from_fn(4, 3, |i, j| (i + j) as f64);
let b = Mat::<f64>::from_fn(3, 5, |i, j| (i as f64) * (j as f64));
let mut c = Mat::<f64>::zeros(4, 5);
// c <- 1.5 * a * b + 2.0 * c, single-threaded
gemm(1.5, a.as_dyn_stride(), b.as_dyn_stride(), 2.0, c.as_dyn_stride_mut(), Parallelism::Serial);
}

gemm(alpha, a, b, beta, c, par) 就地计算 C <- alpha*A*B + beta*C。当 beta == 0 时,gemm 会覆盖 C 原有的内容,且完全不读取它们。这正是 dot 内部所做的事。当 beta 非零时,调用会在 C 已有的值上累加。

签名就是上面看到的样子:输入是 MatRef<'_, T>,输出是 MatMut<'_, T>par 是一个 Parallelism.as_dyn_stride().as_dyn_stride_mut() 这两个转换,把 faer 静态类型化的步长变成适配器所接受的动态步长视图。它们在运行时没有任何开销。

无需拷贝即可直通的布局

适配器始终只读取一个指针和 2 个步长。正因如此,任何 faer 视图都无需拷贝、也无需退化路径即可工作。转置操作数是常见的“行主序 A”情形:把一个列主序矩阵转置,得到的视图行步长非单位,这个视图会直接送进内核。

#![allow(unused)]
fn main() {
// `at` is k x m column-major; `.transpose()` gives an m x k view with a non-unit
// row stride - read straight through, no copy
let a = at.as_dyn_stride().transpose();
let c = gemmkit_faer::dot(a, b.as_dyn_stride());
}

带偏移的子矩阵同理:submatrix(...) 会移动基指针,并保留非连续的列步长。反转视图也一样:reverse_rows()reverse_cols() 带有负步长。

gemmkit 的非检查路径直接处理负步长,所以一个反转的输入在 beta 下也会正确累加,和其他任何输入没有区别。关于引擎如何处理一般步长,见矩阵视图与内存布局

选择并行度

每个入口都接受一个 ParallelismParallelism::Serial 单线程运行。Parallelism::Rayon(n) 用 rayon 以至多 n 个线程运行。Rayon(0) 会自动探测线程数。

gemmkit 让线程数随负载渐进增长,而不是一上来就用满每一个核心。对于固定的机器和固定的配置,同一次调用会给出可复现的结果。今天,串行与并行的运行结果也恰好逐位一致。但这种一致并不是硬性保证:可复现性契约本身只覆盖固定配置这一种情形,而线程数正是配置的一部分。调度模型见并行实践指南。

跨调用复用工作区

gemm 从一个线程局部池中分配它的临时空间。每个入口也都有一个 _with 孪生版本。如果你在循环里驱动大量 GEMM,并想显式持有那块临时缓冲区,就改用 _with 版本。它把 &mut Workspace 作为第一个参数,并在多次调用之间复用这个工作区。

#![allow(unused)]
fn main() {
use gemmkit_faer::{Parallelism, Workspace, gemm_with};

let mut ws = Workspace::new();
for (a, b, mut c) in problems {
    // same result as `gemm`, but the scratch buffer is reused
    gemm_with(&mut ws, 1.0, a, b, 0.0, c.as_dyn_stride_mut(), Parallelism::Rayon(0));
}
}

一个 Workspace 会先增长到能容纳它见过的最大问题,此后 gemmkit 就直接复用它。这对一连串规模相近的中小型 GEMM 最有意义,否则分配开销原本会显现在性能剖析中。

panic 行为

适配器在分发之前会先校验三个共享维度,遇到不匹配就 panic:A.cols 必须等于 B.rowsA.rows 必须等于 C.rowsB.cols 必须等于 C.cols。适配器会给每条消息加上 gemmkit-faer: 前缀,并点名两个冲突的维度,例如 gemmkit-faer: A.cols (4) != B.rows (5)。这些是普通 gemm/dot 路径上仅有的 panic。

feature 门控的入口还会再加几种检查:偏置的长度与重叠、重量化参数,以及预打包 C 的朝向。这些检查复刻了 gemmkit 自身检查入口的措辞。进阶用法页上的每个入口都各自列出了自己的 panic。

faer适配器进阶用法

除了 gemmdot,faer 适配器还镜像了 gemmkit 其余的接口,包括额外的元素类型族、融合尾部运算、基于切片的批量 GEMM,以及预打包操作数。它们每一项都由 feature 门控。每一项也都直接从 faer 视图中读出原始指针与步长,所以转置、子矩阵与反转操作数的表现,同普通路径完全一致。本页逐一走过这些类型族,最后给出一段说明:在 faer 已自带 matmul 的前提下,这个适配器什么时候才值得使用。

入门页介绍了安装、零拷贝机制、gemm/gemm_with/dot、并行度,以及工作区模式。本页的内容都建立在它之上。与普通路径一样,每个入口也都有一个复用调用方持有的 Workspace_with 孪生版本。

整数 GEMM(int8

int8 feature 下,gemm_i8dot_i8 接受 i8 输入,并累加进 i32 输出。输入与输出的元素类型不同。这正是它作为独立入口、而非泛型的又一个实例的原因。faer 的视图类型对元素是泛型的,所以一个 i8MatRef 和一个 i32MatMut 无需任何特殊处理。

#![allow(unused)]
fn main() {
use faer::Mat;
use gemmkit_faer::{Parallelism, dot_i8, gemm_i8};

let a = Mat::<i8>::from_fn(16, 12, |i, j| ((i + j) as i8 % 7) - 3);
let b = Mat::<i8>::from_fn(12, 10, |i, j| ((i * 2 + j) as i8 % 5) - 2);
// i8 * i8 accumulated into a fresh Mat<i32>
let c = dot_i8(a.as_dyn_stride(), b.as_dyn_stride());

// Mat::zeros is ComplexField-only, so integer outputs use from_fn
let mut acc = Mat::<i32>::from_fn(16, 10, |_, _| 0);
// c <- 3 * a * b + (-2) * c, all of alpha/beta/C in i32
gemm_i8(3, a.as_dyn_stride(), b.as_dyn_stride(), -2, acc.as_dyn_stride_mut(), Parallelism::Serial);
}

alphabetaC 都是 i32。算术运算在溢出时回绕。这就是整数 GEMM 的常规语义。

重量化输出(int8 + epilogue

在同时启用 int8epilogue 时,gemm_i8_requant 把重量化这一步融进了内核的写回。i8 输入相乘后进入一个 i32 累加器。内核在一趟之内,把这个累加器缩放、加偏置、取整,并夹取为 i8 输出。它从不物化完整的 m*ni32gemm_i8_requant_u8 做的是同一件事,只是夹取到无符号的 u8 输出,也就是 ONNX QLinearMatMul 风格的激活值域。

这里没有 alpha,因为它已经折进了 scale。这里也没有 beta,因为在一个量化输出上累加是没有意义的。

参数装在一个 Requantize 里。crate 已经重新导出了这个类型,所以你无需为它单独依赖 gemmkitscale 是一个 RequantScale,可以是 PerTensor(f32),也可以是逐通道的 PerRow(&[f32])zero_point 在取整之后以整数形式并入。bias 是一个可选的逐行 i32 向量,在缩放之前加到累加器上。

#![allow(unused)]
fn main() {
use faer::Mat;
use gemmkit_faer::{Parallelism, RequantScale, Requantize, gemm_i8_requant};

let (m, n) = (17, 13);
let bias: Vec<i32> = (0..m as i32).map(|i| 40 * i - 200).collect();
let mut c = Mat::<i8>::from_fn(m, n, |_, _| 0);
let req = Requantize {
    scale: RequantScale::PerTensor(0.05),
    zero_point: -7,
    bias: Some(&bias),
};
gemm_i8_requant(a.as_dyn_stride(), b.as_dyn_stride(), req, c.as_dyn_stride_mut(), Parallelism::Serial);
}

输出为 C[i,j] = clamp(zero_point + round_ne(scale * (sum_k A*B + bias[i])), LO, HI),采用向偶数取整。[LO, HI]i8 入口是 [-128, 127],在 u8 入口是 [0, 255]

适配器在分发之前会校验重量化参数,复刻 gemmkit 自身检查入口的措辞。校验覆盖:

  • scale 非有限或非正。
  • 逐行 scale 切片长度不对,或者与 C 重叠。
  • zero_point 超出输出值域。
  • 偏置长度不对,或者与 C 重叠。

这套校验是针对 C 字节足迹的原始指针运算。适配器从不构造 C 切片。这正是它能把负步长视图安全转发给底层引擎的原因。

复数 GEMM(complex

complex feature 下,gemm_cplxgemm_cplx_withdot_cplx 作用于复数矩阵,并可对每个操作数分别选择是否取共轭。元素类型 TComplex<f32>Complex<f64>

这并非 faer 之外的另一套表示。faer 0.24 的 c32c64 就是 num_complex::Complex<f32>num_complex::Complex<f64> 的类型别名。本 crate 重新导出的正是这两个类型,命名为 Complex,并带有同名的 c32/c64 别名,ComplexScalar 约束也建立在它们之上。因此,一个 faer 复数 Mat 抵达适配器时不需要任何转换,就像实数一样。

gemm_cplx 之所以独立于 gemm,是因为共轭标志放不进同质的接口。它计算 C <- alpha*op(A)*op(B) + beta*C,其中当 conj_a 置位时 op(A) = conj(A),当 conj_b 置位时 op(B) = conj(B)

cplx.rs 里的实现取出与实数路径相同的原始部件,并把两个 bool 标志一路传给 gemm_cplx_unchecked。除此之外没有任何区别,所以转置、子矩阵和反转视图的表现完全一致。dot_cplx 是非共轭 A*B 的便捷接口。

#![allow(unused)]
fn main() {
use faer::Mat;
use gemmkit_faer::{Complex, Parallelism, gemm_cplx};

type C = Complex<f64>;
let a = Mat::<C>::from_fn(12, 9, |i, j| C::new(i as f64, j as f64));
let b = Mat::<C>::from_fn(9, 7, |i, j| C::new((i + j) as f64, 1.0));
let mut c = Mat::<C>::zeros(12, 7);
// C <- alpha * conj(A) * B + beta * C
gemm_cplx(
    C::new(1.3, -0.4),
    a.as_dyn_stride(), true,   // conjugate A
    b.as_dyn_stride(), false,  // leave B
    C::new(0.5, 0.7),
    c.as_dyn_stride_mut(),
    Parallelism::Serial,
);
}

complexepilogue 下,还有 gemm_cplx_fused。它在一趟之内加上一个可选偏置:C <- alpha*op(A)*op(B) + beta*C + bias。偏置是 Bias::PerRow(长度为 A.rows)或 Bias::PerCol(长度为 B.cols)。gemmkit 会把它原样加到该行或该列的每个元素上,绝不取共轭。

这里刻意没有激活参数。像 ReLU 这样基于序的激活,在复数上是没有定义的,所以融合的复数入口只带一个偏置。

融合偏置与激活(epilogue

epilogue 下,gemm_fused 在一趟之内计算 C <- act(alpha*A*B + beta*C + bias)。可选的 BiasPerRowPerCol。可选的 ActivationReluLeakyRelu(slope),最后应用。两者都传 None,就恰好等于 gemm。crate 把这两个选择子都重新导出了。

#![allow(unused)]
fn main() {
use gemmkit_faer::{Activation, Bias, Parallelism, gemm_fused};

let bias: Vec<f64> = (0..m).map(|i| 0.5 * i as f64 - 2.0).collect();
// C <- relu(1.3 * A*B - 0.7 * C + rowbias)
gemm_fused(
    1.3, a.as_dyn_stride(), b.as_dyn_stride(), -0.7,
    c.as_dyn_stride_mut(),
    Some(Bias::PerRow(&bias)),
    Some(Activation::Relu),
    Parallelism::Rayon(0),
);
}

f32/f64 来说,任何形状下,融合结果都与先做普通 gemm 再做同样的标量映射逐位相同。尾部运算折进了同一个内核的写回,不扰动累加的顺序。今天,串行与并行的运行结果也恰好逐位一致。但这种一致只是当前实现的一个特性,不是硬性保证:可复现性契约本身只覆盖固定配置这一种情形,而线程数正是配置的一部分。

f16/bf16(在 half 下)来说,融合结果更精确,而不是逐位相同。单独调用 gemm() 再做窄类型映射,会先取整到窄类型,再拓宽回来,然后再取整一次。融合路径省掉了这次多余的取整:偏置与斜率精确地拓宽到 f32,尾部运算在 f32 中进行,结果只对窄输出取整一次。所以对 f16/bf16,融合结果更精确,但它和先 gemm 再做窄类型映射并不逐位相同。上面对 f32/f64 的逐位保证,并不延伸到这些窄类型上。不过,串行与并行的运行结果,对这些类型今天也依然逐位一致,同样只在固定配置的可复现性契约之下成立。完整契约见融合 Epilogue指南。

对任意的逐元素函数,还有 gemm_map(仅 f32/f64):C[r,c] <- f(alpha*A*B + beta*C, r, c)。闭包在每个输出元素的最终值上恰好运行一次,(r, c) 处于 C 的用户坐标系中。

gemm_map 来做 GELU、sigmoid、夹取,或者与位置相关的变换。普通的偏置或 ReLU,则优先用 gemm_fused,因为它会向量化。gemm_map 每个元素都要付一次间接调用的代价。

批量 GEMM

faer 没有三维数组类型,所以 gemmkit-faer 改用切片来表达批量 GEMM。gemm_batched 取一个 &[(MatRef, MatRef)],也就是逐元素的 (A, B) 输入,与一个 &mut [MatMut]C 输出按位置配对。所有元素共享同一个 alphabetaParallelism

gemmkit 的指针数组引擎把这个批次跨元素并行。它的调度器把整个 GEMM 分配给工作线程。每个工作线程串行运行自己的 GEMM,并为它保持缓存热度。

#![allow(unused)]
fn main() {
use faer::Mat;
use gemmkit_faer::{Parallelism, gemm_batched};

let a = Mat::from_fn(2, 2, |i, j| [[1.0_f64, 2.0], [3.0, 4.0]][i][j]);
let b = Mat::from_fn(2, 2, |i, j| [[5.0_f64, 6.0], [7.0, 8.0]][i][j]);
let mut c0 = Mat::<f64>::zeros(2, 2);
let mut c1 = Mat::<f64>::zeros(2, 2);
let ab = [
    (a.as_dyn_stride(), b.as_dyn_stride()),
    (a.as_dyn_stride(), b.as_dyn_stride()),
];
let mut c = [c0.as_dyn_stride_mut(), c1.as_dyn_stride_mut()];
gemm_batched(1.0, &ab, 0.0, &mut c, Parallelism::Serial);
}

只要每个元素自身的维度自洽,各元素的形状可以不同,也就是一个异构批次。如果输入与输出的数量不一致,调用就会 panic。如果任何元素的维度不自洽,调用也会 panic,并点名出错的元素下标。

每个元素都会重新经过完整引擎分发。所以这个批次复现的是一次普通的 gemm 循环。它对线程数是确定的,因为每个元素都完整地运行在一个工作线程上。出于同样的原因,串行与批量并行的输出逐位相同。

这里没有融合的批量入口:ndarray 适配器提供了共享尾部运算的批量形式,但核心里没有它的指针数组对应物。调度策略见批量 GEMM

预打包操作数

当一个操作数在多次调用之间保持固定,比如权重面对一连串激活值时,把它预打包一次,就能省掉每次调用的重打包。prepack_rhs 把一个 B 变成可复用的 PackedRhs,由 gemm_packed_b 消费。prepack_lhs 把一个 A 变成 PackedLhs,由 gemm_packed_a 消费。crate 把这两个句柄都重新导出了。

#![allow(unused)]
fn main() {
use gemmkit_faer::{Parallelism, gemm_packed_b, prepack_rhs};

let packed = prepack_rhs(weights.as_dyn_stride()); // pack the fixed B once
for (act, mut out) in stream {
    // out must be column-major-ish (|col stride| >= |row stride|)
    gemm_packed_b(1.0, act.as_dyn_stride(), &packed, 0.0, out.as_dyn_stride_mut(), Parallelism::Rayon(0));
}
}

唯一的约束是输出的朝向。预打包的 B 固定了操作数的角色,所以 gemm_packed_b 需要一个偏列主序的 C|col stride| >= |row stride|)。行主序的 C 会交换 A/B 的角色,使打包好的 RHS 失效,因此 gemmkit 会拒绝它。对称地,gemm_packed_a 需要一个偏行主序的 C。如果输出布局不匹配,就退回去用普通的 gemm

epilogue 下,预打包入口还有融合的孪生版本:gemm_packed_b_fusedgemm_packed_a_fused。它们各自在同一个句柄上,接受与 gemm_fused 相同的 Bias/Activation。复用模型见预打包操作数指南。

何时该动用这个适配器

faer 自带了它自己的 matmul。对于两个 faer 矩阵的普通 f32/f64 乘积,直接用它就够了。这个适配器在你需要核心 faer 算子没有提供的能力时才值得使用,前提是要用在 faer 自己的类型上,并且不想离开 faer 生态:

  • 额外的元素类型族i8 -> i32 的整数 GEMM,以及一路融合到 i8u8 输出的重量化。
  • 融合尾部运算:内核在同一趟里计算偏置与激活,或任意的逐元素闭包,而不是对 C 再扫一遍。
  • 跨调用预打包:把一个固定权重矩阵打包一次,然后在一段长推理循环中复用它。
  • 共享的调优面:gemmkit 的三个适配器都坐在同一个引擎上,所以来自 gemmkit-tune 的一份 GEMMKIT_* 环境配置,对它们全部适用。旋钮面见调优旋钮

如果以上都不适用,就改用 faer 内置的 matmul。它是更简单的选择。这个适配器是它的补充,不是替代。

使用gemmkit-tune调优

gemmkit-tune 是一个小巧的命令行自动调优器。你在将要运行 gemmkit 工作负载的那台机器上运行它一次。它会扫描 gemmkit 的运行时 GEMMKIT_* 旋钮,为每个旋钮测量一组有代表性的矩阵形状,然后写出一个由 export GEMMKIT_*=... 行组成的 shell 配置文件。在启动 gemmkit 二进制之前 source 这个文件,已经编译好的二进制就会针对当前主机重新调优,不用重新编译,不用改代码,也不依赖二进制是怎么分发过来的。

它为什么存在

gemmkit 编译期内置的默认值并非凭空而来。gemmkit::tuning 里的每一个阈值都是依据真实测量、在同一台参考机器上手工校准的。在那台机器上最优的取值,编码的是它的缓存大小、核心数,以及内存带宽与算力之比。

换一颗 CPU,L2 不同、核心数不同、内存带宽也不同,各个转折点也就落在了别处,例如:

  • 打包开始划算的那个 k
  • 共享预打包胜过每个 worker 各自打包的那个问题规模
  • 带宽受限的 gemv 应当保持单线程的那个字节下限

gemmkit-tune 会在眼前这块芯片上重新找出这些转折点并把它们钉死。

由此有一个很有用的推论。在参考机器上运行本工具,它基本会重新选出出厂默认值,报告会显示每个旋钮都保持了默认。这是正确的结果,而不是令人失望的结果,它验证了工具本身。

真正的收益出现在与参考机器不同的主机上:一台笔记本、一台共享插槽的云实例、一颗宽核的 Graviton,或一颗带共享 cluster-L2 而没有 L3 的 Apple 芯片。部署主机离参考机器越远,可争取的空间就越大。

扫描过程不含任何随机性。给定同一台机器和同一组标志,它每次都产出相同的配置文件,因此一份配置文件就是一个可复现的产物,可以随部署一起提交进版本库。

为什么在部署主机上运行,而绝不放进 build.rs

这些旋钮是针对执行计算的那颗 CPU 校准的,所以工具必须在那里运行。由此有两点。

其一,不要放进 build.rs。构建主机通常并非部署主机:你在 CI runner 或开发者笔记本上编译,再把二进制发到完全不同的机器上。build.rs 里的自动调优测量的会是构建机,随后会把构建机的转折点烤进一个将在不同缓存层级上运行的二进制里。交叉编译的构建甚至无法执行目标机的代码。整件事的要点就是测量真实的芯片,所以工具必须在其上运行。

其二,这些旋钮就是进程启动时读取一次的普通运行时环境变量,应用一份配置文件并不需要重新构建。同一个发布出来的二进制会读取它所处环境里的任何 GEMMKIT_* 取值。调优主机、source 配置文件、然后启动,引擎便会自行重新配置。

安装与首次运行

安装该二进制并在目标机器上运行:

cargo install gemmkit-tune
gemmkit-tune

一次完整运行需要一两分钟,过程中会打印报告。默认情况下它会在当前目录写入 gemmkit-tune.env。在启动应用的 shell 中 source 它:

source gemmkit-tune.env
./your-gemmkit-app

整个流程就这些。下面都是细化:如何限定运行时长、如何配合你的部署方式,以及何时需要再做一次。

各个标志

gemmkit-tune 不接受位置参数,全部行为都由五个标志控制。

--threads <n>

按此工作线程数进行调优。每个并行探测都在 Parallelism::Rayon(n) 下运行,因此调度类旋钮(各个 oversample 因子、自动 worker 数爬坡)都是针对该宽度优化的。默认值是机器的可用并行度,并会被机器宽度封顶。你无法针对多于机器实有核心数的 worker 调优。标注的 worker 数始终真实。

经验法则:用你的应用实际会使用的 worker 数来跑扫描。如果你的应用固定用 Parallelism::Rayon(8),就传 --threads 8。针对 32 个 worker 调优的配置文件,其调度粒度可能与针对 8 个的不同,而这种错配是要付出代价的。

--time-budget <dur>

限定扫描时长并相应放粗。可写成 30s2m1h,或直接给出秒数。在预算之下,工具会为每次估算减少计时重复次数:无预算时为 7 次,90 秒以下为 5 次,30 秒以下为 3 次。一旦超过截止时刻,它就会停止扫描,把剩余旋钮列为因 “time budget exhausted” 而跳过。

当安装时长必须受限时用它,想要最可靠的配置文件就不要设。若预算小到连一个旋钮都测不完,报告会明说并提示你调大它。

--large-matrices <GiB>

选择性地探测两个内存开销大的旋钮 GEMMKIT_K_STREAM_MAXGEMMKIT_SHARED_LHS_MNK,以给定的 GiB 数作为大型 gemv 矩阵的预算。这两个旋钮只有在开销昂贵的区间才起作用:一个需要 gemv 输出溢出末级缓存,也就是数 GB 的矩阵;另一个需要越过共享预打包转折点的极高 FLOP 形状。两者默认都关闭。

如果你给的预算装不下所需探测,工具会干净地跳过该旋钮,并打印出应当重跑时使用的确切 GiB 数。GEMMKIT_K_STREAM_MAX 探测仅在 64 位下进行,GEMMKIT_SHARED_LHS_MNK 在 32 位上照常扫描。先从 48 起步,若它要求更多就照办。这两个探测究竟做什么,见深入扫描过程

--out <path>

把配置文件写到 ./gemmkit-tune.env 以外的位置。

--dry-run

跑完整扫描并打印报告,但不写出配置文件。适合在落盘之前预览一台机器会选出什么。-h / --help 打印用法。

生成的配置文件剖析

文件由一段头部注释、每个被扫描旋钮一行 export,以及一段列出未扫描项的尾注组成。它长这样,取值完全取决于主机:

# gemmkit-tune profile. Source this before you run a gemmkit app: `source <this file>`
# generated 2026-07-19 14:12:03 UTC by gemmkit-tune 0.1.2
# host: 16 logical cores; L1d 32 KiB, L2 1024 KiB, L3 32 MiB; page 4 KiB
# tuned for 16 worker(s)

export GEMMKIT_MC_REG_PANELS=8  # default (1.00x)
export GEMMKIT_LHS_PACK_THRESHOLD=256  # tuned (1.07x)
export GEMMKIT_PAR_MNK_PER_WORKER=4000000  # tuned (1.03x)

# not swept on this host:
#   GEMMKIT_PARALLEL_THRESHOLD: serial/parallel break-even is strongly shape-dependent ...
#   GEMMKIT_DEEP_KC_BYTES: narrow-only (f16/bf16 deep-contraction twin); no narrow probe here ...

头部记录了它所调优的主机:逻辑核心数、三级缓存、页大小、worker 数,以及工具版本和 UTC 时间戳。这段标注能告诉你,几个月后,一份配置文件是否还与你手上这台机器相符。

每一行 export 都带一段行尾注释:胜出值等于出厂默认时标为 default,胜出值发生了移动时标为 tuned,并附上相对默认的实测加速比。保持默认的旋钮也照样写出,所以配置文件是对每个决策的完整、自带说明的记录,而不只是差异项。

取值一律是原始整数。“无上界“的胜出值会写成它的数值,绝不写成 MAX 别名,因为 gemmkit 的环境变量解析器读的是普通十进制整数。格式错误的 GEMMKIT_* 取值不会致命:gemmkit 会在 stderr 上警告一次并回退到编译期默认,所以手工编辑出的笔误只会退化为默认,而不会崩溃。

部署这份配置文件

gemmkit 对每个 GEMMKIT_* 变量只在首次访问时读取一次,随后在进程整个生命周期内缓存。因此配置文件必须在第一次 GEMM 调用之前就位于环境中,在启动前 source 恰好能保证这一点。这通常有三种常见做法。

shell 配置或启动脚本。 最直接的情形:在运行二进制的那个 shell 里,或在服务的启动脚本里,运行 source gemmkit-tune.env。完整的 shell 语义适用,所以文件原样即可用。

容器 entrypoint。 把配置文件打进镜像,并在 entrypoint 里 exec 你的应用之前 source 它,这样每个容器一启动就已调优就绪。要在与容器运行时硬件相符的主机上调优,而不是在镜像构建机上。

systemd EnvironmentFile 这可行,但有一个注意点:systemd 的 EnvironmentFile 解析器要的是裸的 NAME=value 行,它不认识 export 关键字,也不认识行尾的 # tuned (...) 注释。先转换配置文件,例如运行 grep '^export' gemmkit-tune.env | sed -e 's/^export //' -e 's/[[:space:]]*#.*$//' > gemmkit.env,再让 EnvironmentFile= 指向结果。# 注释头部留着无妨,只有赋值行需要这一步转换。

关于优先级还有一点说明:GEMMKIT_* 环境变量会被应用中的 tuning::set_* 程序化调用覆盖。如果你的应用在代码里调优了某个旋钮,配置文件就不会改动该旋钮。这是刻意为之:自调优的代码胜过部署配置。希望配置文件生效的应用只要不调用这些 setter 即可。完整的优先级次序见调优旋钮

在干净的环境里运行

调优所在 shell 里已经设置的任何 GEMMKIT_* 变量都会扭曲扫描,因为 gemmkit 在测量基准时会读取它。工具会中和它所扫描的那些旋钮,并在发现有 GEMMKIT_* 变量被设置时向你警告。可靠的做法仍然是从一个没有任何这类变量的 shell 里调优。不要 source 一份旧的 gemmkit-tune.env 之后又在同一个 shell 里重跑工具,那正是警告所指的被污染基准。

何时需要重新调优

当配置文件所标注的对象发生变化时就要重调。这意味着换了部署机器:不同的 CPU 带来不同的缓存大小,这正是工具存在的全部理由。这也意味着换了 worker 数:针对 8 个 worker 调优的配置文件不适用于 32 个。

gemmkit 或 gemmkit-tune 的版本升级也可能新增旋钮,所以升级后要重新生成配置文件。一份不再与其头部标注相符的配置文件,就是一份该丢弃并重新生成的配置文件。

要理解扫描究竟测量什么、如何为候选打分,以及如何拿一份配置文件对照你自己的工作负载做校验,请读深入扫描过程

深入扫描过程

使用gemmkit-tune调优讲的是如何运行工具。这一页讲的是扫描背后的机制,包括:

  • 扫描测量什么
  • 它如何裁定胜者
  • 为什么它偏向出厂默认值
  • 如何检验一份配置文件是否真的对你有帮助

一次只扫一个旋钮

扫描是一组彼此独立的一维搜索,而不是联合优化。对每个旋钮,其余所有旋钮都保持默认。工具把该旋钮的候选取值背靠背地测量,选出胜者,然后在扫描下一个旋钮之前把它恢复为默认。因此每个旋钮都是在一个除它以外全为默认的引擎上被评估的。

这是刻意的简化。对约三十个旋钮做完整的联合搜索在组合上毫无希望,而且会被噪声淹没。这些旋钮所把守的转折点在设计上就是各自独立地有意义,因此一维搜索正合适。

代价是不考察旋钮之间的相互作用,而这是可以接受的权衡,因为默认值本就落在一个良好的联合工作点上。工具的任务只是把各个转折点移到主机所在之处。

工具按固定的顺序测量候选,其中不含任何随机性。默认值排第一,因为它是 tie-break 的在位者,随后是各个不相同的额外候选。它为每个形状都以相同的种子重建缓冲区。因此两个候选取值之间的 A/B 对比看到的是逐字节相同的输入,任何机器漂移都会相互抵消。

扫描表与引擎旋钮注册表严格同步

gemmkit 在唯一一处枚举它的旋钮:gemmkit::tuning::knob_env_names()。这个机器可读的注册表是每一个 GEMMKIT_* 名字的唯一真相来源。调优器把每个旋钮归类为 TUNED(有真实的扫描)或 NEVER_TUNED(附理由),并有一个测试断言这两张表恰好划分 knob_env_names():不缺一个,也无陈旧项。

这条实际保证很直接:加进 gemmkit 的旋钮无法悄悄逃出自动调优器。在有人为它写好扫描、或记下为什么刻意放过它之前,构建都会失败。因此你读到工具列出的被扫描旋钮时,读到的是一张由编译器对照引擎保持诚实的列表。

测量什么,以什么单位

每个候选的得分是一种吞吐率。GEMM、i8 或 batched 探测以 GFLOP/s 计分,算法是每次调用 2*m*k*n,batched 再乘以 batch 数。gemv 探测则以 GB/s 计分,因为矩阵乘向量是带宽受限的,此处搬运的字节数才是诚实的评判标准。

单形状的估算刻意做得稳健。工具先预热探测闭包几次,再自动定出迭代次数,使一批计时约运行 50 毫秒。它计时若干这样的批次,并报告中位数速率,同时给出观测到的最小值和最大值。最小值和最大值不是装饰:它们记录的是逐次运行的离散度,而胜者逻辑正是用这个离散度在噪声下保持诚实。

打分:在一组探测形状上取几何平均

旋钮从不以单一形状评判。每个旋钮都携带一小组探测形状。工具挑选这些形状,为的是让该旋钮确实起作用,并从两侧夹住它的转折点。候选的得分随后就是它各形状中位吞吐率的几何平均

几何平均不论绝对规模大小都给每个形状等权。因此单个大形状无法抬高一个只对它有利的取值。胜者必须是跨整组的普遍提升。最差形状的离散度会被带进几何平均,因此噪声闸门对整组保持保守,而非只信最平稳的那个形状。

探测按旋钮挑选,为的是让旋钮起作用。举几个例子:

旋钮探测形状族为什么是这些形状
MC_REG_PANELS方阵 f32,512 到 3072,并行3072 那档压测 A 宏面板在 L2 中的驻留
LHS_PACK_THRESHOLD列主序 A,候选 32..MAX同时夹住 aarch64 的低复用平台段与 x86 的默认值 1024
SMALL_K_THRESHOLD瘦长的大 m,nk,如 4096x16x4096k 跨越就地 / 打包驱动器的转折点
GEMV_PARALLEL_BYTESm 的 gemv,GB/s覆盖缓存驻留 / DRAM 受限的字节下限
GEMV_TIER_STEPGEMV_THREAD_CAP触碰约 2.4 到 134 MiB 的 gemv,GB/s跨越 gemv worker 阶梯的一级,因为整组都落在同一级里的探测集会把每个候选评成一样
SEQ_INTERNAL_BYTES_PER_WORKER(aarch64)让每 batch-worker 份额为 96/192/384/432 KiB 的 batched 形状从两侧夹住约 128 KiB 的默认,是一个双向校验器
I8_VNNI_MIN_PAR_MNK(x86)方阵 i8,384/512/640夹住 VNNI / 加宽回退的并行转折点

tie-break 偏向默认且感知噪声

取几何平均最高者是错的。在一台有噪声的机器上,1% 的领先通常只是运气。胜者逻辑改为从默认起步,只有当某个候选的几何平均超过当前最优、且超出两个候选中较大的那份实测离散度时,才会升级到它。逐次运行的噪声按构造无法越过这道门槛,因此永远无法改写一个旋钮。恰好打平则保留默认。

对默认为 0 的那些 “auto” 旋钮还有一道额外余量。这些旋钮从机器派生取值,例如 LLC 大小、核心数、页大小。固定候选必须在噪声之外再多赢 auto 5%。那些 auto 派生会适应探测集未覆盖的形状,所以一个在探测上仅以毫厘取胜的固定数,不值得拿这份适应性去换。

这种偏向默认在噪声下是正确取舍。默认是一个已知良好、经过刻意选择的取值,而工具常常在无人盯着的机器上无人值守地运行。这道不对称的门槛意味着最坏情况不过是工具重现了默认值,它绝不会把你退化进一个测量假象里。

扫描完全没有 RNG,因此一次运行值得信任:最坏什么都不做,而它移动某个旋钮时,一定是因为一个真实、可复现的提升越过了噪声。

时间预算如何封顶并放粗扫描

--time-budget 从两方面起作用。

其一,它预先放粗每次估算:无预算 7 次计时重复,90 秒以下 5 次,30 秒以下 3 次,以一点测量稳定性换速度。

其二,它强制一个硬截止。每个旋钮之前工具都查看时钟,一旦超过截止时刻,就不再开始新的扫描,把每个剩余旋钮记为因 “time budget exhausted” 而跳过。

因此紧预算既会模糊它确实做的那些测量,又会从尾部丢弃旋钮。无预算时扫描以全额重复次数跑到底。

哪些旋钮被跳过,以及为什么

有些旋钮从不扫描,报告和配置文件尾注会逐个说明原因:

  • PARALLEL_THRESHOLD:串行/并行的收支平衡点强烈依赖形状。单个 m*n*k 标量无法适配所有长宽比,因此工具保留经过校准的跨形状默认,而不去自动拟合它。对比 GEMV_THRESHOLD,它是干净的二元开/关决策,被扫描。
  • DEEP_KC_BYTES:它把守 f16/bf16 深度收缩孪生路径,而调优器不跑窄类型探测。其 auto 默认从 L2(一项机器属性)派生。若需重调窄类型的深 k 触发点,请直接覆盖它。
  • PREFETCH_MIN_BYTES:它把守驱动器的 C tile 预取。其 auto 默认从检测到的 LLC(一项机器属性)派生,而探测这个转折点需要每个候选都用超出 LLC 的工作集。请直接覆盖它以重调触发点(usize::MAX 关闭预取,1 强制开启)。

另一些旋钮在当前目标上是惰性的,因而以此为由跳过。SEQ_INTERNAL_BYTES_PER_WORKER 只被 aarch64 的 batched 拆分规划器读取:在那里扫描,在 x86 上惰性并跳过。I8_VNNI_MIN_PAR_MNK 把守的是 x86 VNNI 小并行回退,其他目标的 i8 内核并无此回退。NC_NO_L3_PANELS 只在没有 L3 的机器上被查阅:在那里扫描,在有 L3 的主机上惰性并跳过。

两个重型旋钮,除非你传入 --large-matrices,否则一律跳过。

–large-matrices 解锁什么

有两个旋钮只在一个复现起来昂贵的区间才重要,所以它们藏在内存预算之后、需显式开启。

K_STREAM_MAX 限定 axpy-gemv 的输出保持寄存器分块到多远。它只有当输出明确 DRAM 受限时才占优。因此它的探测把输出固定在约两倍末级缓存。1x-LLC 的输出正落在缓存边界上,测不出任何决定性结果,所以探测避开这个尺寸。探测随后在校准上限附近扫描 k

那个输出尺寸是固定的,不随预算缩放,因此达到它要用数 GB 的矩阵。如果你给的预算装不下最大的探测,工具会跳过该旋钮。它会打印出应当重跑时使用的 GiB 数(向上取整)。在 32 位目标上,它会直接跳过该旋钮,因为那些矩阵根本装不进地址空间。

SHARED_LHS_MNK 把守共享 LHS 预打包。这个预打包消除每个 worker 冗余的 A 打包,却增加一道 fork-join 屏障,因此只在越过一个较大的 m*n*k 取值(x86 上约 8e9)后才划算。它的探测用的是越过该转折点的瘦高、极高 FLOP 形状。

工具在普通扫描期间会中和这两个旋钮,无论自身是否在扫描它们。这样陈旧的环境变量取值就无法扭曲读取它们的基准。

读懂终端报告

报告开头是一张每旋钮一行的汇总表。它的列是旋钮、单位、形状数、默认、胜者、加速比,以及一个读作 keeps default-> <value> 的结果列,移动过的旋钮会高亮。

表格下方,一段候选明细为每个旋钮打印完整的扫描地形。它列出每个候选的几何平均中位数,用一个前置标记标出默认取值,用另一个标记标出胜者,让你看出最优点有多平坦或多陡峭。

再往下是带原因的跳过列表,然后是一段页脚。页脚计出扫描了多少旋钮、多少移离默认、多少被跳过。在参考机器上,页脚会注明所有旋钮都保持了默认,这是意料之中的,配置文件会重现它们。

校验一份配置文件

扫描测量的是合成的、大致方形的探测。这对找出一台机器的转折点是正确选择。但你的工作负载有自己的形状,所以在生产中信任一份配置文件之前,先确认收益能迁移过去。

有两种检验方式。直接的一种:在部署主机上,分别在 source 与不 source gemmkit-tune.env 的情况下对你自己的应用计时并比较。可复现的一种:运行 gemmkit 的 criterion 基准,它覆盖五个头部分组(sgemmdtypesgemvprepackedbatched),并以一个已保存的基线对比:

cargo bench -p gemmkit -- --save-baseline stock
source gemmkit-tune.env
cargo bench -p gemmkit -- --baseline stock

如果某个旋钮移动了而你在意的东西反倒退化,配置文件就是一个纯文本文件。删掉或注释掉那一行 export,其余照留。头部标注和每行的 tuned/default 标记让你很容易看出该动哪一行。

设计目标与总体图景

gemmkit 是一个纯 Rust 的 GEMM 引擎。它在带显式步长的 &[T] 切片上计算 C <- alpha*A*B + beta*C,也可以直接在带 isize 步长的裸指针上计算。运行时,它会选出 这台机器支持的最优指令集:x86-64 上是 AVX-512 或 FMA/AVX2,并配有专门的 VNNI 与 BF16 点积内核;aarch64 上是 NEON;wasm32 上是 simd128。其余平台一律有可移植的标量路径兜底。 整个工作空间使用 edition 2024、rust-version 1.89,许可证是 MIT OR Apache-2.0。

API 表面由三类调用者共同决定。应用代码走安全的切片入口,例如 gemmgemm_fusedgemm_i8,这些入口会在任何 unsafe 代码运行之前完成全部校验。线性代数库走 *_unchecked 层,这一类包括随仓库发布的 ndarraynalgebrafaer 适配器,也包括任何按同样方式 构建的库。*_unchecked 层信任调用方自己的不变量,因而能接受安全层无法表达的内存布局。 受限的部署场景则得到一个可以 #![no_std] 构建、零强制依赖的核心,一路可以支持到带编译期 SIMD 的 wasm32。本章其余内容都可以从四条设计信条推导出来。 ARCHITECTURE.md 在 “Goals and constraints” 一节里紧凑地列出了它们,本章则展开每一条背后的理由。

边界上的安全性

各个受检入口都会在触碰任何 unsafe 代码之前运行 validate_gemm_viewsgemmkit/src/api.rs)。它的 panic 清单刻意做到详尽:

  • 形状不匹配:gemmkit 会检查 A.cols != B.rowsA.rows != C.rowsB.cols != C.cols。每条 panic 消息都会写明两个不一致的数字。
  • 视图寻址越出切片:对 A、B、C 各自,gemmkit 都会算出步长能到达的最高偏移 (extent),并与切片长度比对。一个需要的元素数超过切片实际长度的视图,会带着确切的 差额 panic。
  • 负步长:安全层会直接拒绝,消息会指向 gemm_unchecked。带负步长的 &[T] 视图必然 要寻址到切片起点之前,安全层基于切片长度的推导没法为这样的地址担保。
  • 自混叠的输出C 上的某个步长可能把两个不同的 (i, j) 映射到同一个偏移,零步长 就是最常见的情形。这在 AB 上没问题,因为广播输入只会被读取。但在 C 上会 panic,因为并行驱动假设各个输出 tile 互不重叠。通过这样的视图写入会造成一个纯安全代码 就能触发的数据竞争。
  • CAB 重叠:gemmkit 按字节区间来检查,因此即便 C(i32)与 A/B (i8)的元素大小不同,检查依然精确。融合入口还会额外检查偏置长度(PerRow 为 m、 PerCol 为 n)以及偏置与 C 是否不相交。
  • 问题大到无法计量:广播步长允许逻辑维度逼近 isize::MAX,这会让内部打包缓冲区的 尺寸计算溢出 usize。所有这类乘积都会在元素数到字节数的必经关口(Workspace::regions) 以“失败即封闭“的方式 panic,而不是回绕后分配不足。

panic 的措辞本身也是一份被测试的契约。正确性测试套件会断言这些字符串的精确内容。因此 修改一条错误消息是一个刻意且可见的动作。

*_unchecked 层之所以存在,是因为这套校验只在一个边界上才有意义。适配器直接从 ndarraynalgebrafaer 的类型里取出指针和步长,这些类型自身的不变量已经保证了 有效性,重复检查纯粹是浪费。而且基于切片的检查根本无法表达适配器需要的东西:比如一个 反转的 ndarray 视图会有负步长,基指针也会落在分配区的中间,这两者对裸引擎来说都是合法 且健全的。于是安全成本恰好只需付一次,要么由 gemmkit 的校验器来付,要么由调用方的类型 系统来付,绝不会重复支付。

unchecked 入口都是带有明确文档契约的普通 unsafe fn。使用方式见 Unchecked 层

可复现、而非逐位相同的并行结果

gemmkit 承诺可复现的并行结果:在固定的输入、固定的环境和固定的配置下,输出不依赖工作 线程数。这个承诺由三个机制支撑。

第一个机制是分块尺寸。KCNC 只由缓存模型计算得出,从不依赖线程数。MC 至多只会 按 MR 对齐地重新分组。因此每次运行对每个输出元素的归约顺序都完全相同。

第二个机制是归约顺序。每个输出元素都由同一个工作线程沿完整深度从头归约到尾,引擎中不 存在被拆分到多个线程上的归约。

第三个机制是按需调度。打包出来的字节不依赖是谁做的打包,所以任何线程都可以接手任何 tile。哪个线程计算哪个 tile 因运行而异,结果却始终不变。

同样重要的是 gemmkit 没有承诺什么。串行与并行逐位相同并不在契约之内。它今天之所以 在驱动路径上成立,是因为串行和并行跑的是同一套内核、同一套分块,但没有任何东西把它钉死。 跨配置的逐位一致同样明确不在承诺内:改一个调优旋钮,分块就会跟着改变,浮点求和顺序也就 可能合法地随之改变。

同一类型不同内核之间的逐位一致也不在承诺内。bf16 的 vdpbf16ps 点积内核相对于拓宽再 FMA 的路径重塑了累加的舍入方式,gemmkit 对它只要求落在容差之内,而不要求精确相等。

为什么把线画在这里?若永久承诺串行与并行逐位相同,就会挡住真正有价值的工程手段:既挡住 把相邻深度对融合进一条指令的点积内核,也挡住会考虑并行度的分块选择。这样的承诺换不来任何 用户能跨机器、跨版本依赖的东西。

“感知并行度的分块“已经不再是设想。驱动器已经有一个作业深度下限,会随工作线程数收缩 MC,好让并行作业列表足够深。它之所以仍然逐位可复现,恰恰是因为这条较弱的契约给它留了 余地:MC 依然是 MR 的整数倍,所以微 tile 的集合以及每个元素由 KC 决定的累加顺序都 不受影响。

固定配置下的可复现性,才是测试能够断言、部署能够依赖、引擎又能在演进中一直守住的性质。 在能以低成本给出更强保证的地方,gemmkit 就会给出更强保证:gemv 按输出行划分工作,跨任意 线程数都逐位一致;i8 整数路径是精确算术,所以它的 VNNI 点积内核与拓宽内核逐位一致。

变化点上无宏、无 transmute

引擎沿三条轴变化:指令集、元素类型、运算家族。每条轴都是一个普通的 trait。SimdSimdOps 覆盖 ISA。Scalar 覆盖元素类型。KernelFamily 覆盖运算家族。

分发槽是缓存在 OnceLock 里的带类型函数指针。微 tile 的几何形状是一对在分发点选定的 const 泛型参数。一个“内核变体“实际长什么样,看 gemmkit/src/dispatch/float.rs 就知道:

#![allow(unused)]
fn main() {
unsafe fn gemm_f32_fma(t: Task<f32>, par: Parallelism, ws: &mut Workspace) {
    // MR = 2*8 = 16, NR = 6 -> 12 acc + 2 lhs + 1 rhs = 15 of 16 YMM
    unsafe { run_typed::<f32, Fma, 2, 6>(Fma, t, par, ws) }
}

unsafe fn gemm_f32_avx512f(t: Task<f32>, par: Parallelism, ws: &mut Workspace) {
    // MR = 2*16 = 32, NR = 12 -> 24 acc + 2 lhs + 1 rhs = 27 of 32 ZMM
    unsafe { run_typed::<f32, Avx512F, 2, 12>(Avx512F, t, par, ws) }
}
}

这就是每个(类型,ISA)组合的全部专属代码:一个点名令牌和 tile 的包装函数。另一条路, 也就是 C 语言 BLAS 传统里那种宏展开或手抄的每 ISA 内核,因为审阅成本和扩展成本太高而 被否决了。

trait 加 const 泛型,让全库只剩一个泛型微内核需要阅读、单步调试、修复。一次调度改进落地 一处,所有 ISA 就都跟着受益。编译器会对每个单态化实例做完整类型检查。OnceLock 槽里放的 是带类型的函数指针,不是类型擦除的指针,所以签名漂移是编译错误,不是潜伏的 transmute 缺陷。

扩展遵循同样的形状。新增一个 ISA,只需要一个零大小令牌、它的 SimdOps 实现,以及每个 选择阶梯里的一条分支。新增一个元素类型,只需要一个 Scalar 实现、一个家族(或者通过 拓宽/收窄接缝复用现有家族),再加一个分发槽。驱动、打包、分块永远不用改动。一个测试 (gemmkit/tests/open_closed.rs)通过用第二个平凡家族驱动驱动器来验证这一点。后续两页 SIMD 令牌与 ISA 分发标量与内核家族会 逐条展开每条接缝。

no_std 与零强制依赖的核心

关掉默认 feature 后,核心 crate 会以 #![no_std] 构建。它只需要 corealloc, 不再依赖任何其他东西。每个可选 feature 至多引入一个依赖:

Feature引入的依赖换来的能力
std(默认)raw-cpuid(仅 x86/x86-64 目标)运行时缓存与 CPU 特性探测、GEMMKIT_* 环境变量旋钮、线程本地工作区池
parallel(默认)rayonParallelism::Rayon 多线程
halfhalff16/bf16 混合精度 GEMM
complexnum-complexc32/c64 复数 GEMM
int8i8 -> i32 整数 GEMM
epilogue融合偏置/激活/映射 epilogue(重量化还需同时开启 int8
wasm_threadsparallel 外无线程化 wasm 上显式定容的 rayon 池

没有 std 时,编译期 target feature 会取代运行时 CPU 探测。环境变量旋钮会关闭,不过 程序化的 tuning::set_* setter 依然可用,因为它们只是普通的原子操作。线程本地池也会换成 按调用传入的工作区。

位于技术栈这么底层的一个内核,不应该把依赖策略强加给它的宿主。嵌入式或 wasm 部署得到的 是和桌面构建一样的驱动、一样的家族、一样的可复现性契约,只是少了那些真正需要操作系统的 机制。实操指南见 no_std 与 WebAssembly

工作空间地图

五个 crate 以 0.1.2 版本同步发布,另有一个模糊测试 crate,它被有意放在自己独立的工作空间 根目录下:

路径Crate角色
gemmkit/gemmkit核心 GEMM 引擎(本章描述的全部内容)
gemmkit-ndarray/gemmkit-ndarray面向 ndarray(>= 0.17.1)视图的零拷贝适配器
gemmkit-nalgebra/gemmkit-nalgebra面向 nalgebra 0.35 矩阵的零拷贝适配器
gemmkit-faer/gemmkit-faer面向 faer 0.24 矩阵的零拷贝适配器
gemmkit-tune/gemmkit-tune安装期自动调优程序,输出一份 GEMMKIT_* 环境变量配置
gemmkit/fuzz/gemmkit-fuzzcargo-fuzz 目标,仅限 nightly,被排除在稳定工作空间之外

适配器有意做得很薄。每个适配器都直接从宿主库的原生视图里取出矩阵指针和步长:C 序、 F 序、一般步长、反转步长皆可,全程零拷贝。取出之后,适配器就转发给 *_unchecked 引擎,靠宿主类型自身的不变量来完成安全层原本要做的校验。每个适配器还会把同名的 Cargo feature(parallelwasm_threadshalfcomplexint8epilogue)转发给 gemmkit,所以 feature 的行为在各处都保持一致。各适配器章节覆盖了它们完整的表面: ndarraynalgebrafaer

gemmkit-tune 是进程外的校准器。引擎里每个启发式阈值都是一个运行时旋钮(见 调优旋钮)。编译进去的默认值是在某一台机器上校准的, 调优器存在的意义就是在你自己的机器上重做一遍这次校准。

在部署主机上运行一次调优器二进制。它会对每个旋钮在一组探测形状上做扫描,然后写出一份由 export GEMMKIT_*=... 组成的 gemmkit-tune.env 配置文件。在启动应用之前 source 这个 文件即可。整个过程不需要重新编译,也没有构建期的耦合。调优器和库之间唯一的契约就是那份 文档化的环境变量表面。tuning::knob_env_names 注册表守着这份契约的诚实性,因为调优器的 扫描表会与它做断言比对。实操指南见 gemmkit-tune 章节

模糊测试 crate 被有意放在工作空间之外。cargo-fuzz 需要 nightly,因为要用到 build-std 和 AddressSanitizer。排除它可以让 cargo test --workspace 和 MSRV 构建始终留在 stable 上。

本章与 ARCHITECTURE.md 的关系

仓库里的 ARCHITECTURE.md 是紧凑的地图。它给出分层表、调用路径、各条接缝,以及每个 子系统一节的说明,写给那些开着另一个窗口看代码的读者。本书这一章是同一份素材的导览版。 它用的是同一套层号、同一批文件引用,但留出了篇幅讲清理由、被否决的备选方案,以及完整的 示例。两者若有出入,以代码为准,而且这意味着两份文档都有 bug。

接下来请读分层结构,了解静态骨架。再读 一次 GEMM 调用的生命周期,了解动态过程。

分层结构

核心 crate 里的每个模块,开篇都会声明自己在这个栈里的位置。api.rs 写着 “Public core API (layer L8a)”,driver.rs 写着 “The generic GEMM driver (layer L5)”,一路排到 L0 的 simd.rs。这些标签不是装饰,它们把 crate 的依赖纪律写在了读者不可能错过的地方。下面这张 图按依赖顺序列出各模块,正是这一点让“向下依赖“的主张可以被检验。本页自底向上走一遍这个 栈。等走到公共 API 时,它用到的每个词都已经有了定义。下一页 一次 GEMM 调用的生命周期会沿相反方向再穿越一次这个栈,跟随 一次真实的调用。

L8a  api        safe slice entries, *_with, *_unchecked; MatRef/MatMut
L7   dispatch   runtime ISA selection, one memoized fn pointer per type
L6   special    gemv, small-k, small-m,n, batched reroutes
L5   driver     the generic 5-loop blocked GEMM, one for all families
L4   kernel     KernelFamily seam (float/mixed/int/complex) + Epilogue
L3   cache      topology detection + BLIS analytical blocking
L2   parallel   worker-count resolution, JobCursor work distribution
L1   pack       micropanel packing primitives
L0   simd       ISA tokens + SimdOps vocabulary;  scalar: Scalar/Acc types
     ---        cross-cutting: tuning (GEMMKIT_* knobs), workspace (buffers)

有两处放置值得单独说说,因为正是它们让每支箭头都指向下方。parallel 的位置很低,在内核 家族、缓存模型和驱动之下,因为它是一套自足的工作线程词汇。Parallelism 策略枚举、Ptr 这个 Send 指针包装,以及 JobCursor,都只依赖 tuning。而 kerneldriverspecialdispatch 全都向下够到它们。packkernel 之下,因为是各家族的打包钩子 建立在打包原语之上,而不是反过来。

L0:词汇表,scalar.rssimd.rs

最底层定义了 crate 其余部分被允许谈论的东西。gemmkit/src/scalar.rs 是数据类型接缝, 而且刻意做得极小:

#![allow(unused)]
fn main() {
pub trait Scalar: Copy + Send + Sync + PartialEq + 'static {
    /// The type in which products are accumulated. `Self` for `f32`/`f64`
    type Acc: Scalar<Acc = Self::Acc>;
    /// The additive identity
    const ZERO: Self;
    /// The multiplicative identity
    const ONE: Self;
}
}

这就是整个 trait:单位元常量,加上累加器类型。f16bf16f32 中累加。i8i32 中累加。f32f64 和复数类型都在自身中累加。Scalar 本身不带任何算术。真正的 运算要么向量化地发生在 SimdOps 里,要么在各家族的标量 epilogue 里,所以新增一个元素 类型永远不会连带出一套标量算术接口。细化 trait FloatNarrowFloatComplexFloat 按需叠加个别路径需要的额外能力。scalar.rs 刻意不知道的事:SIMD 的存在。它完全不知道 自己的常量将来会被广播进向量寄存器。

gemmkit/src/simd.rs 构成承重墙,和它一起的还有 simd/ 里的各后端:avx512.rsfma.rsneon.rsscalar.rswasm.rs,以及 complex.rs 里的复数胶水代码。三个 trait 分担这份职责。Simd 是一个零大小的 ISA 令牌。例子有 Avx512FFmaNeonScalarTokSimd128,外加具备点积能力的 Avx512VnniAvx512Bf16。它唯一 的方法是 vectorize,也就是 #[target_feature] 跳板。它把运行时选中的 intrinsic 放进 特性已启用的代码生成上下文。SimdOps<T> 是每个元素类型的厚词汇表:寄存器类型、 LANES、load/store/broadcast/mul/add/fma/reduce,以及可覆写的 accumulate_tile 调度。 再往上的 KernelSimd<L, R, A, O> 是拓宽/收窄接缝,让混合精度不需要驱动层的任何分支。这个 模块刻意不知道的事:它上面的一切。模块文档写明它只依赖 scalarcore,因此可以原封 不动拆成独立 crate。SimdOps 不知道什么是微面板、什么是缓存、什么是 GEMM。

L1:机械拷贝,pack.rs

gemmkit/src/pack.rs 存放两个共享的打包原语,把带步长的 A 或 B 区域拷成连续、微内核 尺寸的面板。这些就是 L4 各内核家族日后把打包钩子委托过来的那些拷贝。唯一的例外是复数 家族的实虚平面拆分打包,它随家族本身存放。

这两个原语是 pack_panelspack_kgroup_panelspack_panels 是微面板为主序的拷贝。 LHS 面板高 mr 行,RHS 面板宽 nr 列。两者用的是同一个例程,只是交换了“领先“和“深度“ 两个步长的角色。尾部补零,带步长的源则走缓存分块转置。pack_kgroup_panels 是点积家族 使用的 k 组交错变体。

pack.rs 刻意不知道的事:自己的输出去了哪里。同一个例程既填充一次性调用的临时区,也 填充并行共享的打包缓冲,还填充调用方长期持有、贯穿整个进程生命周期的 PackedRhspack.rs 从未见过 Workspace、工作线程或生命周期,只见过 dstsrc 和步长。正是这份 漠不关心,让预打包路径与逐调用路径的字节完全一致。它只依赖 scalar,既不点名家族也不 点名缓存,因此才能坐得这么低。

L2:工作分发,parallel.rs

gemmkit/src/parallel.rs 拥有三样东西。

第一,Parallelism 枚举:Serial,或者 Rayon(n),其中 Rayon(0) 表示自动。

第二,感知工作量的线程数决策。总工作量低于某个门限时用一道串行闸门。显式指定的线程数 会被尊重,但有上限。自动档随总工作量 m*n*k 扩展,而不是直接跳到全部核心。访存受限的 矩阵乘向量形状另有一条独立的带宽规则。

第三,按需分发的机械装置。JobCursor 是一个无锁原子游标,工作线程从中拉取连续区块。 job_grainpacked_block_grain 这两个旋钮决定区块大小。for_each_worker 这个 fork-join 就是上层用作屏障的机制。parallel.rs 还提供 Ptr,让裸指针得以进入 rayon 闭包的 Send + Sync 垫片。

parallel.rs 只依赖 tuning,坐在一切调用它的东西之下。同一套工作线程词汇同样服务于 kerneldriverspecialdispatch。它刻意不知道的事:作业是什么。JobCursor 对 一个抽象计数发放索引区间。这个文件里没有任何地方提到 tile、矩阵或家族。这正是为什么日后 同一个游标既调度驱动的 tile,也不加区分地调度 B 打包面板、A 打包行块和 gemv 的行面板。 详见并行执行

L3:机器模型,cache.rs

gemmkit/src/cache.rs 及其后端(cache/cpuid.rscache/sysfs.rscache/sysctl.rs) 回答两个问题:缓存层级长什么样?由此推出什么分块?探测是一条不可能失败的尽力而为回退链。 它依次尝试 x86 上的 CPUID,然后 Linux sysfs,然后 macOS sysctl,最后是在 Zen5 上校准的 静态默认值。#[cfg] 只用来挑选探测方法,从不直接给值,结果在 Machine 中记忆化一次。

blocking() 随后按 BLIS 模型解析式地计算出 (MC, KC, NC)KC 的尺寸要让 A、B 微 面板在 L1 中共存。MC 的尺寸要让 A 宏面板装进 L2。NC 的尺寸要让 B 宏面板装进 L3。 关键类型是 Level(带着文档详尽的 shared_by 争用字段)、CacheTopology,以及分块 结果。这一层刻意不知道的事:线程数。blocking() 没有工作线程参数,而这个缺席是承重的。 与线程数无关的分块,正是设计目标里可复现性契约背后的机制。详见 分块与缓存模型

L4:运算家族接缝,kernel.rs

gemmkit/src/kernel.rskernel/float.rsmixed.rsint.rscomplex.rsepilogue.rs)定义了 KernelFamily。它是把一种 GEMM 与另一种区分开的全部要素打成的 一捆。一个家族捆着 LhsRhsAccOut 这几个类型。它也捆着打包布局:pack_lhspack_rhs,两者都委托给 L1 的原语。它捆着微内核 microkernel_epi。它还捆着 OUT_IS_ACCDEPTH_MULTIPLE 这类常量,告诉驱动该怎么为这个家族分块。

FloatGemm<T> 是基线。MixedGemmIntGemm/IntGemmVnniComplexGemm 是原样复用 驱动的兄弟家族。这一层还拥有 Epilogue trait,带着它零成本的 Identity。它也拥有 AlphaStatusBetaStatus 这两个枚举。驱动会预先算好这两者,好让微内核永远不必比较 浮点数。家族刻意不知道的事:自己的 tile 尺寸。MR_REGNR 是微内核方法上的 const 泛型,由高三层的分发点按(类型,ISA)选定。家族对任何几何都能编译,新 tile 是一次新的 实例化,永远不是新类型。

L5:引擎,driver.rs

gemmkit/src/driver.rs 是服务所有家族的那一个分块循环嵌套。它有 BLIS 顺序的 jc -> pc -> 扁平作业列表 结构。它做出自适应的打包决策。它可以按深度切片打包 B,也可以 就地读取 B。它可以按工作线程打包 A,通过共享预打包打包 A,或者干脆不打包 A。它还处理 预打包 RHS 的消费路径,走的是预打包 API 复用的 pack_rhs_full 布局。正是这份复用,让 预打包 GEMM 和普通 GEMM 产生完全相同的面板字节。

它的公开入口是 runrun_epiloguerun_packed_rhsrun_packed_rhs_epilogue, 全部汇入私有的 run_inner。它刻意不知道的事:任何具体的元素类型或 ISA。整个文件对 Fam: KernelFamily 和一个 KernelSimd 令牌泛型。它从不点名 f32,从不点名 AVX-512,从不按元素类型分支。这就是开放封闭性质。新增家族或 ISA 不会动这个文件一个字。 gemmkit/tests/open_closed.rs 证明了这一点。它用一个 crate 并不发布的第二个平凡家族去 驱动驱动器。

L6:改道路径,special.rs

gemmkit/src/special.rsspecial/gemv.rssmall_k.rssmall_mn.rsbatched.rs)存放寄存器分块驱动并不适合的形状的专用路径。它们是矩阵乘向量、低深度 GEMM、小 m,nk 的内积形状,以及批量编排层。它们都藏在同样的公共入口背后, 特殊路径有完整介绍。特殊路径刻意不知道的事:自己为什么被选中。门限 gemv_thresholdsmall_k_thresholdsmall_mn_dim 住在上一层的分发逻辑和旁边的 tuning 模块里。small_k::run 甚至分辨不出自己在服务 gemmgemm_fused 还是 gemm_map,因为 epilogue 是作为不透明的泛型参数到达的。批量路径是唯一一处某层向上 够回去的地方。batched.rs 把每个元素经由 dispatch::execute(L7)转发,于是它继承 上方的整条阶梯。这就是本页末尾讨论的那唯一一处带注解的例外。

L7:运行时 ISA 选择,dispatch.rs

gemmkit/src/dispatch.rsdispatch/isa.rsfloat.rsmixed.rsint.rscomplex.rs)把“这台机器该跑哪个内核“变成一锤子买卖。每个元素类型有一个 OnceLock<Dispatched<T>> 槽。特性探测只跑一次。胜出的单态化入口(普通、预打包、融合) 连同 tile 几何一起被缓存。之后每次调用都是经由带类型函数指针的一次普通间接调用, 没有 transmute,也没有 AtomicPtr<()>

这一层还拥有问题描述符 Task<T>execute 里的退化情形处理、方向归一化 orient_transpose,以及特殊路径门限。它也拥有 GEMMKIT_REQUIRE_ISA 这个钉选项,用来 强制指定或者响亮地拒绝某个特定内核。它刻意不知道的事:Task 里的指针从哪来。受检的 切片视图和 unchecked 的裸指针到这里长得一模一样。校验发生在上面,或者根本没发生,分发层 既不知道也不关心。参见 SIMD 令牌与 ISA 分发和面向用户的 运行时 ISA 分发

L8a:公共边界,api.rs

gemmkit/src/api.rsapi/batched.rscplx.rsfused.rsint8.rsmap.rspacked.rs)定义了好几样东西。它们是带步长视图 MatRef/MatMut、各家族的 安全入口,以及 validate_gemm_views 的 panic 清单。安全入口分 *_with(调用方自备 工作区)和 *_unchecked(裸引擎)两种变体。这一层还负责把视图降解为 Task

它刻意不知道的事:分发层以下的一切。API 层看不到将要运行哪个 ISA、会选出什么分块、 打包会不会发生。校验完成后,它把 Task 递给 dispatch::execute,任务就结束了。对称地, MatRef 从不出现在这一层之下。crate 的其余部分只讲指针和步长。

为什么箭头只指向下方

依赖方向是这套架构唯一的硬规则。每一层由上层驱动,且对上层一无所知。simd 只依赖 scalarcore。驱动从不点名元素类型或 ISA。L7 以下没人知道分发层的存在。L8a 以下 没人听说过切片。

这条规则只有唯一一处刻意的、带注解的例外。special/batched.rs(L6)把每个批量元素 经由 dispatch::execute(L7)向上转发。这次再入让每个元素继承与其形状上独立 gemm 调用相同的 driver、small-k、small-mn、gemv 路由。它省去了要手工维护的第二套分发阶梯。 这是一支有意为之的向上箭头,也是 crate 里唯一的一支。

三个收益支撑着这份纪律。其一,扩展成本。因为知识只向下流动,新 ISA、新元素类型、新 家族都在自己那一层接入。其下的一切可证明地原封不动。扩展点描述的各条接缝 之所以成立,正是因为没有任何下层可能对上层做过特判。其二,审阅局部性:审计微内核只需读 kernel/float.rsSimdOps 契约,别无其他。审计调度只需读 driver.rsparallel.rs。其三,可测试性:下层可以被隔离地检验。SIMD 一致性测试把每个令牌与标量 模型比对,开放封闭测试用外来家族驱动驱动器。这正是测试与验证里正确性 故事得以成立的原因。

两个横切模块

有两个模块立在栈的旁边而不是栈里,因为每一层都需要它们,而它们除 core/alloc 外不 依赖任何东西。

gemmkit/src/tuning.rs 是统一的旋钮表面。引擎里每个启发式阈值都住在这里:串行/并行 门限、打包门限与步长、特殊路径阈值、调度粒度、分块上限。每一个都按这个顺序解析:先看 按调用参数,再看程序化 setter(tuning::set_*),再看环境变量(GEMMKIT_*),最后看 编译期默认值。环境变量只读一次并缓存。无法解析的值会在 stderr 上告警后回退,而不是 panic,因为性能旋钮里的笔误绝不应该弄崩进程。全部 GEMMKIT_* 名字都在 tuning::knob_env_names 注册表中枚举。crate 外的消费者,也就是 gemmkit-tune 的扫描表、 旋钮属性测试、fuzz 的 setter,都与它做断言比对。所以新旋钮不可能悄悄逃出覆盖范围。面向 用户的导览见调优旋钮

gemmkit/src/workspace.rs 是临时内存的故事。Workspace 是可增长、64 字节对齐的缓冲区。 Workspace::regions 把它切成每工作线程(或每行块)的 LHS 区域,加一个共享的 RHS 区域, 并在元素数到字节数的必经关口做“失败即封闭“的溢出检查。在 std 下,一个可重入安全的 线程本地池提供默认工作区,所以普通 gemm 每线程至多分配一次。*_with 入口则改为穿入 调用方自备的工作区,从第一次足够大的调用之后堆分配为零。没有 std 时,每次调用都使用 全新的工作区。详见打包与工作区

一次GEMM调用的生命周期

上一页描述的是静止的栈。这一页跟随一次调用穿过它。样本是 crate 文档里的快速上手示例:

#![allow(unused)]
fn main() {
use gemmkit::{gemm, MatRef, MatMut, Parallelism};

// 2x3 * 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]);
}

这个玩具形状会在下文的某个早退口离开主路。整个走读同时记住两个问题:上面的 2x2x3, 以及一台 AVX-512 机器上的 2048x2048x2048 f32 乘积。后面这个更大的乘积会一路走到底, 穿过每一层。下面是压缩过的路线:

gemm(alpha, A, B, beta, C, par)
  |  validate_gemm_views: shapes, bounds, aliasing     [api.rs]
  v
Task<T>: raw pointers + isize strides
  |  m == 0 || n == 0        -> return                 [dispatch.rs]
  |  k == 0 || alpha == 0    -> C <- beta*C, done
  v
memoized per-type kernel (OnceLock fn pointer)
  |  gemv shape (m==1||n==1) -> special/gemv.rs        [dispatch/float.rs]
  |  orient: row-major-ish C -> compute C^T = B^T*A^T
  |  small m,n + long k      -> special/small_mn.rs
  |  k <= small_k_threshold  -> special/small_k.rs
  v
driver::run                                            [driver.rs]
  jc over NC -> pc over KC (never parallel)
    -> flat job list (ic row-block x jt column-tile),
       workers drain a shared JobCursor, pack A/B adaptively
  v
Fam::microkernel_epi: MR x NR tile in registers        [kernel/float.rs]
  alpha/beta epilogue store (vector fast path | scratch drain)

第一站:校验与降解

gemm 本身只有一行。它借出线程本地工作区,转发给 gemm_withgemmkit/src/api.rs)。 gemm_with 会运行 validate_gemm_views,也就是设计目标里的 完整 panic 清单。形状必须一致。每个视图都必须留在自己的切片以内。C 必须对每个 (i, j) 寻址唯一。C 不能与任何一个输入重叠。

然后视图就消解了。这一点以下的所有代码只讲 Task<T>:一个 Copy 结构体,装着 m, k, nalpha/beta,以及三个带 isize 行/列步长的裸指针。转置从来不是一个标志 位,转置视图不过是交换了步长。当 beta == 0 时,契约规定 C 永远不会被读取,所以它 可以是未初始化的。

unsafe 边界恰好在这里跨越,由刚刚跑完的校验背书。gemm_unchecked 则晚一步进场,改由 调用方自己扛起这份背书。

第二站:分发层的早退口

dispatch::executegemmkit/src/dispatch.rs)会趁元素类型还是具体的,先处理掉退化 的代数情形:

#![allow(unused)]
fn main() {
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);
}

输出为空意味着无事可做。A*B 项消失(k == 0alpha == 0)时,调用会退化成一次 C <- beta*C 缩放。这次缩放从不读取 AB。其中 beta == 0 会直接写零而不读 C,兑现未初始化 C 的契约。

只有真正的乘积才会到达 T::dispatch,它读取该类型的 OnceLock 槽。首次使用时,选择 阶梯会探测 CPU 特性,尊重 GEMMKIT_REQUIRE_ISA 的钉选(不满足时 panic 而非回退)。 随后它会缓存胜出的单态化入口,连同 tile 几何一起。此后每次调用都只是一次间接调用。在 那台 AVX-512 机器上,f32 解析为 run_typed::<f32, Avx512F, 2, 12>,也就是 32x12 的 tile。

第三站:run_typed 里的路由

run_typedgemmkit/src/dispatch/float.rs)是一串简短的闸门,每道闸门都把寄存器 分块驱动伺候不好的形状改道送走。

先是 gemv。若 m == 1 || n == 1,且该路径未被 GEMMKIT_GEMV_THRESHOLD 封顶关闭, 调用会直接去 special/gemv.rs。这发生在方向归一化之前,用的是用户的原始坐标系。 gemv 自己解决方向问题。它把 m == 1 的情形当作转置后的 rows x k 问题来处理,输出行 的划分也由它自己完成。

其余的都会经过 orient_transpose 做方向归一化。若 C 是行主序倾向 (|csc| < |rsc|),分发层会把问题改写成它的转置:C^T = B^T * A^T。这会交换 mnA/B 的指针与步长,以及 rsccsc

这个恒等式是免费的。数据一字节不动,只改描述符。它换来一条强不变量。此后,输出的 步长是小的那个(对完全连续的 C 即 rsc == 1)。输出的每一列因此在内存中连续, 内核得以沿连续的列往下走。

微内核的快速存储路径恰恰要求这一点:rsc == 1,这样它就能用向量存储写一列中 LANES 个连续行。下面的每一层都只需为一种方向优化,而不是两种。那个全行主序的 2048 立方 就命中了这次交换。引擎实际算的是 C^T,而分发层以下无人知晓。

接下来是归一化任务上的另外两道闸门。

m,n 形状会去 special/small_mn.rs。两个维度都必须不超过 small_mn_dim,且收缩 长度要超过 small_k_threshold。在那里,每个输出元素是一次水平 SIMD 点积。当两个操作 数都沿 k 单位步长流动时,这是零拷贝的。当某个操作数带步长时(k > small_mn_pack_min_k),会走一个打包档,只拷贝不合格的那个操作数。

k 形状(k <= small_k_threshold,默认 x86 为 16、aarch64 为 8)会去 special/small_k.rs。它把整个乘积当作一个就地读取的深度面板,直接过微内核,没有任何 分块和打包开销。

闯过全部闸门的(那个 2048 立方就是)会进入 driver::run。驱动声明的前置条件是: m, n, k > 0alpha != 0、方向已归一化。

第四站:驱动的循环嵌套

driver::run 带着零成本的 Identity epilogue 转发给 run_innergemmkit/src/driver.rs)。融合入口会带着真正的 epilogue 落进同一个函数。驱动对 家族和 ISA 令牌泛型。对这次调用,那就是 FloatGemm<f32>Avx512Fmr = MR_REG * LANES = 32nr = 12

分块要先算。cache::topology().blocking(mr, nr, sizeof_lhs, m, n, k) 按 BLIS 缓存 模型给出 (MC, KC, NC)。它们以打包输入元素计(sizeof(Lhs),而非累加器),所以 窄类型会得到更深的块。随后循环嵌套按 BLIS 顺序展开:

  • jc 遍历 NC:列块,尺寸保证打包后的 B 宏面板驻留 L3。
  • pc 遍历 KC:深度切片。这层循环永不并行。所有深度切片都累加进同一批 C tile。把深度并行化就意味着对 C 做同步的读-改-写,或者拆分归约。让深度保持串行,才 使每个输出元素由同一个工作线程从头归约到尾,这是可复现性契约的一半。beta 只在第 一个切片(pc == 0)参与。之后的切片以等效 beta 为一累加。混合精度家族 (OUT_IS_ACC = false)恰好只有一个切片,kc = k,运行中的部分和因此永远不会经过 窄输出类型的舍入。
  • 一份扁平的一维作业列表:每个深度切片内,剩下的工作是 n_mc 个行块乘 n_nt 个列 tile。它们压平成 n_jobs = n_mc * n_nt 个下标。工作线程按需从共享的无锁 JobCursor 拉取连续区块。没有静态划分,快核自然吸收更多工作。区块粒度对线程数过 采样(job_grain)。打包 LHS 路径改用与行块对齐的 packed_block_grain,区块因此 永不跨越打包边界。线程数本身来自 par.resolve(m*n*k, n_jobs)。这是以工作量为准: 随总工作量 m*n*k 除以每 worker 下限来扩展,而不是一步跳到全部核心。若这个线程数 会让作业列表浅到每个 worker 分不到几个块,驱动器会先缩小 mc。这只会切出更多、更 小的行块,因而不会移动任何结果比特。缩小 mc 会先把列表加深,再交给游标发放。

打包是自适应的,两侧各自独立决策。

B 会在 m 越过 rhs_pack_threshold 时每个深度切片打包一次。打包面板会被全部 n_mc 个行块复用,只有复用足够高,这次拷贝才划算。否则 B 按原始步长就地读取。真要打包时, 打包本身也是并行的。工作线程从游标拉取 nr 宽的列面板。for_each_worker 的汇合就是 先写后读的屏障。打包后的 B 是所有计算线程共享的唯一缓冲,这道屏障因此很关键。

A 有三种模式。每个工作线程都可以把手头的行块打包进自己的私有工作区区域。rsa != 1 或者行块不是 mr 的整倍数时,这是强制的。其余情况下,是否值得打包取决于每线程列复 用量,或者 TLB 不友好的列步长。大型并行问题上,一个共享预打包可以代替这一步:把每个 行块恰好打包一次。它打包进按块分配的区域,靠自己的屏障同步(shared_lhs_mnk 门限)。 这消除了各线程的重复打包。复用低到任何拷贝都摊不平时,A 就地读取。

这些区域的尺寸经由 Workspace::regions 预先切好,带着前文说过的“失败即封闭“溢出 检查。完全不打包的路线甚至不会碰工作区。

对每个作业,工作线程会解析出自己的 A 面板,已打包或就地。它会定位 B 面板:本次调用 打包的、预打包缓冲里的,或者就地的。随后它对块内每条 mr 行的条带调用微内核。这一切 都在 simd.vectorize 之内进行,整个条带因此都在 target-feature 代码生成上下文中 执行。

第五站:微内核与它的存储

Fam::microkernel_epigemmkit/src/kernel/float.rsmicrokernel_impl)计算 一个 MR x NR 的 tile。对这次调用,那是 32x12 个 f32 值,以 [[Reg; MR_REG]; NR] 数组的形式驻留在 24 个 ZMM 累加寄存器里。

满宽 tile 走 SimdOps::accumulate_tile,即升序 k 的融合乘加调度。这是一条接缝, 像 NEON 这样受载入约束的 ISA 会换上一个软件流水的变体,只重排载入,绝不重排算术。 列方向的边缘 tile 则走一个运行时定界的循环。这个循环恰好读 nr_eff 列,保证未打包的 B 永远不会被读过最后一个真实列。

然后 alpha 会折进累加器。alpha == 1 时这一步整段跳过,靠的是驱动预先算好的 AlphaStatus

存储是 beta 和 epilogue 的居所。它有两条路。

快路径在满 tile 且输出行步长为一时触发:mr_eff == mr && nr_eff == NR && rsc == 1。第三站的方向归一化正是让这个条件变得常见的原因。每个累加寄存器直接与 C 结合。 beta == 0 时原样存储(不读 C),beta == 1 时相加,一般 beta 时融合乘加。结果 再用向量存储写回。

边缘 tile 和带步长的输出走通用路径。所有累加器先倒进栈上的 scratch tile,即工作线程 栈帧里的 SCRATCH_LEN 数组,零分配。随后一个标量循环沿 C 的任意步长逐元素做同样的 beta 运算。

普通 gemmIdentity epilogue 穿过这一切。每个 epilogue 挂钩都由 !E::IS_IDENTITY 把守,这是一个关联 const。守卫在单态化时折叠殆尽,产出的内核与 无 epilogue 的代码逐字节相同。

融合调用,比如 gemm_fusedgemm_map、重量化,跑的是同一个引擎,带一个真正的 epilogue。这个 epilogue 只在 last_k 为真时点火,也就是在最后一个深度切片上,每个 输出元素恰好一次。这条线索在 Epilogue 融合里继续。

抄近路回家

那个 2x2x3 示例没见过上面大部分风景。它带着均为正的 m, n, kalpha == 1 进入 execute。它到达 run_typed,没过 gemv 闸门(n != 1m != 1)。它被方向交换, 随后没过小 m,n 闸门,因为 k = 3 算不上长收缩。以 k = 3 <= 16,它进了 special/small_k.rs。这条路线是同一个微内核上的一个就地深度面板,没有分块,没有 打包,工作区一次都没碰。

2048 立方走完了带并行 B 打包的完整驱动。在 Parallelism::Rayon(0) 下,它的线程数随 总工作量扩展。

同一个入口,同一份结果契约,两段截然不同的旅程。下面的层替调用者做了所有决定,调用者 从不需要操心。各站的深层机制见分块与缓存模型打包与工作区并行执行特殊路径

SIMD令牌与ISA分发

gemmkit 在运行时选择指令集,这个决定与 Rust 编译 SIMD 内建函数的方式相冲突。AVX 和 AVX-512 内建函数只有在 target feature 已启用的上下文中才能正确生成代码。这个上下文通常来自外层函数上的 #[target_feature(enable = "...")] 属性。程序只有在真正运行于某颗具体 CPU 上时,才知道哪些 feature 可以安全启用。微内核是所有指令集共用的同一个泛型函数,所以没有一个属性可以单独钉在它身上。

本页讲两件事。第一,L0 SIMD 层(gemmkit/src/simd.rsgemmkit/src/simd/)如何用零尺寸 ISA 令牌加一个蹦床函数化解这对矛盾。第二,L7 分发层(gemmkit/src/dispatch.rsgemmkit/src/dispatch/)如何选出并缓存获胜的内核。

ISA 令牌与 vectorize 蹦床

ISA 令牌是一个零尺寸类型,代表一种指令集选择。x86 上的令牌是 Fma(AVX2 + FMA)和 Avx512F,外加点积内核变体 Avx512VnniAvx512Bf16。aarch64 上的令牌是 Neon。wasm32 上的令牌是 Simd128ScalarTok 存在于每个平台,是可移植的兜底。

每个令牌都实现 Simd trait。它唯一的方法是 vectorize:在该令牌的 target feature 已启用的情况下运行一个闭包。下面的代码展示了整个机制,摘自 gemmkit/src/simd/fma.rs

#![allow(unused)]
fn main() {
/// AVX2 + FMA ISA token
#[derive(Copy, Clone, Default)]
pub struct Fma;

impl Simd for Fma {
    #[inline(always)]
    unsafe fn vectorize<R>(self, f: impl FnOnce() -> R) -> R {
        #[target_feature(enable = "avx2,fma,f16c")]
        unsafe fn inner<R>(f: impl FnOnce() -> R) -> R {
            f()
        }
        // SAFETY: the caller of `vectorize` (the runtime dispatcher) guarantees
        // the CPU supports avx2+fma(+f16c); `inner` then establishes the codegen
        // context, and `f` inlines into it
        unsafe { inner(f) }
    }
}
}

诀窍在于内联的方向。inner 是一个带 #[target_feature] 属性的极小函数。闭包 f 被内联进 innerf 装着打包循环和微内核调用,这些代码全部由 #[inline(always)] 原语构成。于是每一条内建指令都落在 feature 已启用的代码生成上下文里,泛型内核本身却从未被任何属性触碰。

unsafe 契约只有一条义务:调用者必须保证 CPU 确实支持该令牌的 feature。运行时分发器在每个进程里只确认一次这件事。

这与 pulp 和 faer 使用的模式相同。它对串行路径和 rayon 工作线程闭包同样适用。驱动层把每个列条带的微内核调用整体包进 simd.vectorize(|| ...),蹦床的开销就这样摊薄到许多个 tile 上。

ScalarTokvectorize 就是一句 f(),无需启用任何东西。这正是标量路径能在任何地方运行、包括在 Miri 下运行的原因。

SimdOps:按元素类型展开的指令词汇表

L0 一共定义了 3 个 trait,不是 2 个。Simd 就是上文的 ISA 令牌 trait。SimdOps<T> 是本节要讲的、按元素类型展开的词汇表。KernelSimd<L, R, A, O> 是第三个 trait:当一个家族的输入类型、累加器类型和输出类型并不完全相同时,它负责把加载值加宽、把存储值收窄。本节末尾会讲到它。

令牌本身对元素类型一无所知。所有实际运算都放在 SimdOps<T> 上,按 (ISA, T) 对分别实现一次。它给出寄存器类型 Reg、通道数 LANES,以及微内核需要的每一条原语。令牌与元素类型是解耦的,所以 LANES 随这个二元组变化。f32Fma 下是 8 通道,在 Avx512F 下是 16 通道。f64 的通道数是同一令牌下 f32 的一半。

这份词汇表刻意做得很厚。基础操作有 zerosplatloadustoreumuladd,以及融合乘加 mul_add。它的减法搭档是 fnma,计算 c - a*b。复数内核的某个累加项要靠 fnma 才能算。词汇表里还有水平求和 reduce_sum,供 gemv 与点积 epilogue 使用。

在这些之上还有几个原语。maxmin 只有实数浮点令牌才实现,供融合 ReLU 和 clip epilogue 使用。LANE_FMA 标志和它的 fma_bvec 方法给 NEON 提供了一条按通道索引的 FMA 路径:把一段 RHS 列作为一个向量整体加载,替代逐列发出 splat。accumulate_tile 就是 GEMM 的内层循环本体。它的可移植默认调度在任何乱序核心上都已经会被直接编译成教科书式的寄存器分块内核。

复数拆分内核在这里也有自己的接缝,叫 cplx_microkernel。点积内核也有自己的接缝,叫 dot_accumulate,它长在 KernelSimd 上,而不是 SimdOps 上。当一个家族的输入类型、累加器类型和输出类型并不全部相等时,KernelSimd<L, R, A, O> 就是它要驱动的那个接缝,比如 f16 输入配 f32 累加器。它把较窄的输入加载值加宽成累加器类型,再把累加器的值收窄后写出。同构家族的这 4 个类型全部相等,靠一个覆盖一切(blanket)实现就能免费获得 KernelSimd,完全不需要任何专属某个 ISA 的代码。这两个接缝的细节见标量与内核家族点积内核与深K孪生

厚,正是设计意图所在。matrixmultiply 的 per-ISA trait 很薄,逼着每个指令集都从头重写一遍内核。这里内核需要的每一条原语都在 SimdOps 背后,于是微内核成了横跨所有 ISA 的同一个泛型函数。新增一个指令集的成本是一个新令牌、它的 SimdOps 实现,加上每条分发阶梯里的一行。simd 模块只依赖 crate::scalarcore,对内核、驱动、缓存层没有任何反向依赖,所以整个抽象可以原封不动搬进一个独立的 crate。

分发层

分发层把“用哪个令牌“这个问题变成一次性的决定。每个可分发的元素类型都拥有一个 OnceLock 槽位,存放一个 Dispatched<T> 描述符。f32f64gemmkit/src/dispatch/float.rsf16bf16dispatch/mixed.rsi8dispatch/int.rs,它有自己的 IntDispatchedIntRequantDispatched 形态,因为这些类型是异构的。c32c64dispatch/complex.rs。下面这段代码摘自 dispatch/float.rs,略有删节。

#![allow(unused)]
fn main() {
#[derive(Copy, Clone)]
pub(super) struct Dispatched<T> {
    pub(super) run: GemmFn<T>,
    pub(super) run_packed: PackedFn<T>,
    #[cfg(feature = "epilogue")]
    pub(super) run_fused: FusedFn<T>,
    #[cfg(feature = "epilogue")]
    pub(super) run_packed_fused: PackedFusedFn<T>,
    pub(super) mr: usize,
    pub(super) nr: usize,
    pub(super) depth_multiple: usize,
}
}

这个槽位缓存了获胜的单态化入口:普通内核、预打包 RHS 内核,以及(在 epilogue feature 下)它们的融合孪生版本。它同时缓存微铺块几何 (mr, nr) 和该家族的 depth_multiple。缓存几何信息是为了让 prepack_rhs 能用与后续消费调用相同的 ISA 选择来确定缓冲区尺寸。depth_multiple 让 bf16 预打包路径把打包深度取整,匹配点积内核的布局。这里的一切都是带类型的函数指针,没有 transmute,也没有 AtomicPtr<()>

一次调用沿着固定的链路走:gemm 调用 dispatch::execute,它先处理退化情形。dispatch::execute 再调用 T::dispatch,读取记忆化的槽位。槽位解析为一次间接调用,进入像 gemm_f32_avx512f 这样的包装函数,这个包装函数把共享的泛型入口实例化为 run_typed::<f32, Avx512F, 2, 12>

选择只运行一次,就在 OnceLock 的初始化器里。它先处理 GEMMKIT_REQUIRE_ISA 锁定(见下文)。之后,自动阶梯在 x86 上先探测 avx512f,再探 avx2fma,最后落到标量。aarch64 上 NEON 是基线,架构规定 NEON 必备,所以这里无需探测。

wasm32 上根本没有运行时特性检测。simd128 在编译期由 cfg(target_feature = "simd128") 决定。构建必须传 -C target-feature=+simd128,否则就会拿到标量内核。标量在所有架构上都是地板。

各类型的阶梯在同一骨架上加自己的门槛。f16 的 FMA 分支还额外要求 f16c,因为 vcvtph2psvcvtps2ph 转换需要它。bf16 阶梯在普通 AVX-512F 之前先试 avx512bf16 点积内核。i8 阶梯在加宽内核之前先试 avx512vnni(连同 avx512bw)。

有两个构建模式的细节值得知道。有 std 时,特性检测走 is_x86_feature_detected!,结果记忆化在 OnceLock 里。没有 std 时不存在运行时 CPU 检测,因为 raw-cpuidstd 门控。探测宏退化为 cfg!(target_feature = ...)GEMMKIT_REQUIRE_ISA 的解析退化为 Auto,select 函数每次调用都会执行。不过其中每个分支此时都是编译期常量,所以会直接折叠成一个确定的选择。no_std 构建就只跑其编译期 target feature 保证的那条路径。参见 no_std 与 WebAssembly

作为 const 泛型的微铺块几何

除指令编码之外,真正随 (类型, ISA) 变化的只有微铺块形状。它表达为在分发点选定的一对 const 泛型 (MR_REG, NR),从来不是新类型、新 trait 或宏。MR_REG 是铺块的寄存器高度,行数即 MR = MR_REG * LANES。下表以 f32 为例。

ISA(MR_REG, NR)LANES铺块 MR x NR寄存器预算
AVX-512F(2, 12)1632 x 1224 累加 + 2 lhs + 1 rhs = 27 个 ZMM
FMA(AVX2)(2, 6)816 x 612 累加 + 2 lhs + 1 rhs = 15 个 YMM
NEON(4, 4)416 x 416 累加 + 4 lhs + 1 rhs = 21/32 个向量寄存器
simd128(2, 4)48 x 48 累加 + 2 lhs + 1 rhs = 11 个活跃 v128
标量(4, 4)14 x 4普通局部变量

f64 的通道数减半,同样的 (MR_REG, NR) 组合于是给出 16x12(AVX-512F)、8x6(FMA)、8x4(NEON)、4x4(simd128)。这些预算不是巧合。NEON 刻意留出约 11 个空闲寄存器,给宽乱序核心留出重命名余量,让它把下一步的加载与当前的 FMA 重叠起来。simd128 停在 11 个活跃向量,因为 LLVM 的 wasm 后端在大约 16 个之后就会开始溢出。这些注释就写在 dispatch/float.rs 里包装函数的旁边,所以上表就是代码本身,不是一个愿望。

用 GEMMKIT_REQUIRE_ISA 锁定内核

默认情况下,最优可用的 ISA 胜出。设置环境变量 GEMMKIT_REQUIRE_ISA 会强制锁定唯一一个内核,不再自动选择。它接受下面这些取值(不区分大小写):

  • scalar
  • fma(别名 avx2
  • avx512f
  • avx512vnni(别名 vnni
  • avx512bf16(别名 bf16
  • neon
  • simd128(别名 wasm
  • auto

未设置或空串等同于 auto。无法识别的值会直接 panic,这样 CI 配置里的拼写错误就不可能悄悄选中别的东西。avx512vnni 锁定 i8vpdpbusd 点积内核。avx512bf16 锁定 bf16vdpbf16ps 点积内核。对其余元素类型,两者都会解析为普通 AVX-512F 路径。

这里最核心的行为是:锁定永不回退。如果 CPU(或 Intel SDE 这类模拟器)没有报告所需的特性,选择就会 panic。如果所请求的 ISA 在目标架构上根本不存在,比如非 aarch64 上的 neon,或非 x86 上的 fma/avx512*,选择同样会 panic。panic 消息会指明缺失的特性。这么做是为了 CI 的诚实性:一个想要检验某个内核的任务,必须大声失败,而不是默默测了另一个内核。

simd128 的锁定在 wasm 上同样有用。那个 target feature 是一个极易被遗忘的编译期开关。锁定把“忘了传开关“从静默回退到标量,变成了拒绝运行的构建。

取值只读一次。选择结果记忆化在各类型的 OnceLock 里,所以必须在第一次 GEMM 调用之前,在进程环境中设好这个变量。之后再修改,在进程的生命周期内都不再生效。锁定的使用侧内容,包括 CI 配方以及它与调优旋钮的配合,见运行时 ISA 分发

标量与内核家族

gemmkit 能相乘好几种元素类型:f32f64f16bf16i8Complex<f32>Complex<f64>u8 也会出现,但只作为重量化输出。这些类型全部流经同一个驱动层、同一套打包框架、同一个缓存模型、同一个并行调度器。

全部差异由两个 trait 承载。L0 的 Scalargemmkit/src/scalar.rs)回答一个类型是什么、用什么类型累加。L4 的 KernelFamilygemmkit/src/kernel.rs)回答这一类 GEMM 与其他类有什么不同。驱动层对家族保持泛型,从不按元素类型分支。本页讲这道分割线是怎么划的,以及为什么划在这里。

Scalar:只有常量与累加器类型,别无其他

Scalar 刻意做得极小。整个 trait 如下,摘自 gemmkit/src/scalar.rs

#![allow(unused)]
fn main() {
pub trait Scalar: Copy + Send + Sync + PartialEq + 'static {
    /// The type in which products are accumulated. `Self` for `f32`/`f64`
    type Acc: Scalar<Acc = Self::Acc>;
    /// The additive identity
    const ZERO: Self;
    /// The multiplicative identity
    const ONE: Self;
}
}

Scalar 没有 Add,没有 Mul,也没有转换。它只有单位元常量和一个关联的累加器类型。这个省略是刻意的。所有向量化算术都在 SimdOps 里(见 SIMD 令牌与 ISA 分发)。epilogue 需要的标量算术改放在几个窄小的侧 trait 上,谁需要谁才实现。Float 覆盖 f32f64,复数类型也借 num-complex 的运算符实现它。NarrowFloat 覆盖 f16/bf16 的加宽与收窄转换。ComplexFloat 覆盖复数拆分内核的实部与虚部访问。

如果 Scalar 本身背着算术,每个新元素类型都要欠一整套它未必真正拥有的运算。i8 是最典型的例子。它完全不需要任何算术 trait,因为它的内核全部通过 SIMD 接缝和精确的 i32 整数运算完成。

Acc 是混合精度的接缝,表很短。

元素类型累加类型
f32f64自身
f16bf16f32
i8(以及仅作输出的 u8i32
Complex<f32>Complex<f64>自身

递归约束 Acc: Scalar<Acc = Self::Acc> 让这条链一步就封顶:f16 -> f32 -> f32 -> ...。泛型代码于是可以放心引用“累加器的累加器“,不必关心输入原本有多窄。对同质类型,Acc = Self 这条分支在编译期就会塌缩,零开销。

KernelFamily:一个 GEMM 家族的全部差异所在

KernelFamily 打包其余的一切。它携带四个元素类型(LhsRhsAccOut)、打包布局(pack_lhs/pack_rhs,负责写出微面板优先的面板),以及微内核。

三个关联常量决定驱动层如何对待一个家族。OUT_IS_ACC 说明一个运行中的部分和能否在深度面板之间经由 C 往返。这是最关键的一个常量,下文详述。FORCE_PACK_LHSFORCE_PACK_RHS 在打包执行了内核依赖的某种变换时置位,比如复数共轭或点积内核的交织。这种情况下,驱动层必须永不原地读该操作数。DEPTH_MULTIPLE 是点积内核的指令组深度。其余家族一律是 1

一个家族恰好覆写两个微内核方法中的一个。不做融合的家族覆写普通的 microkernel,继承默认的 microkernel_epi。这个默认实现会先断言 E::IS_IDENTITY,再转发到 microkernel。这道断言是一道失败即关闭的护栏。真正的 epilogue 落到不会融合的家族上,会直接 panic,而不是被悄悄丢弃。

会融合的家族转而覆写 microkernel_epi。浮点、混合精度、重量化家族都走这条路。每个家族都把 epilogue 穿进自己的存储路径。它们的普通 microkernel 方法就此成为死代码。它保留默认的 unreachable! 方法体。

铺块几何刻意不放在这个 trait 上。(MR_REG, NR) 是在分发点按 (家族, ISA) 选定的一对 const 泛型。新铺块因此是这对泛型的一次新实例化,不是新类型。

回报体现在驱动层的签名上。driver::run::<Fam, S, MR_REG, NR> 对家族和 ISA 令牌保持泛型。它调用 Fam::pack_lhsFam::pack_rhsFam::microkernel_epi,全程没有一个针对元素类型的 if。新增一类 GEMM 意味着写一个新家族,从来不需要碰驱动层。

家族清单

目前共有 10 个家族类型。它们按代际划分:同质、加宽、点积、重量化、复数。按这个顺序读下来,接缝一目了然。

家族类型(Lhs/Rhs -> Acc -> OutOUT_IS_ACCDEPTH_MULTIPLE说明
FloatGemm<T>T -> T -> Tf32/f64true1基线:一个泛型微内核覆盖所有 ISA
MixedGemm<N>N -> f32 -> Nf16/bf16false1KernelSimd 接缝的加宽 FMA
Bf16DotGemmbf16 -> f32 -> bf16false2vdpbf16ps 点积内核;两操作数强制打包、k 成对交织
MixedGemmF32<N> / Bf16DotGemmF32N -> f32 -> f32true1 / 2f32 输出的深 K 孪生:累加相同,存储为 f32
IntGemmi8 -> i32 -> i32true1精确、回绕;加载时符号扩展
IntGemmVnnii8 -> i32 -> i32true4vpdpbusd 点积内核,+128 符号修正,与 IntGemm 逐位一致
IntGemmQ<O> / IntGemmVnniQ<O>i8 -> i32 -> i8u8false1 / 4重量化变体(epilogue feature)
ComplexGemm<T, CONJ_A, CONJ_B>T -> T -> Tc32/c64true1拆分(SoA)内核;两操作数强制打包成平面布局;共轭是打包期的符号翻转

FloatGemm 是参照系。它是同质类型,一个泛型的 microkernel_impl 覆盖所有 ISA 和所有铺块。

混合精度与整数家族引入了 Acc != Lhs。它们完全倚仗下文讲的加宽/收窄接缝。点积家族 Bf16DotGemmIntGemmVnni 还要更进一步。它们各自换上交织的打包布局和一条硬件点积指令。f32 输出孪生的存在是为了让深收缩可以重新分块。这一切都在点积内核与深K孪生里展开。

重量化变体在整数累加之上,接了一段精确的 i32 -> i8/u8 重量化。它们是 Epilogue融合里融合故事的一部分。

ComplexGemm 保持 Acc = T,所以复数的 alpha/beta 能原样穿过驱动层。它的热循环转而经由一个专用接缝,跑在实数分量上。复数拆分内核讲的就是这个接缝。

本页只停留在清单层面,深入分析留给那几篇独立的页面。

KernelSimd:加宽与收窄的接缝

驱动层对 ISA 令牌的约束是 S: KernelSimd<Fam::Lhs, Fam::Rhs, Fam::Acc, Fam::Out>gemmkit/src/simd.rs)。KernelSimd<L, R, A, O> 扩展了 SimdOps<A>,所以它以 A 累加。它在类型边界上补齐了家族需要的四个动作。

load_lhs 加载 LANES 个 LHS 值,把它们加宽为一个 A 寄存器。splat_rhs 把一个 RHS 标量加宽后广播出去。load_outbeta != 0 读取 C 时把输出值加宽。store_out 把一个 A 寄存器收窄为 LANES 个输出值,真正收窄时向最近偶数舍入。

同质情形零成本。一条覆盖实现 KernelSimd<A, A, A, A> for S: SimdOps<A> 把全部四个方法转发到普通的 loadusplatstoreu。所以 FloatGemm<f32> 及其同类完全不需要任何 per-ISA 代码。

混合家族则要按 ISA 添加实现。它的加载会真正加宽,比如 f16 -> f32vcvtph2ps,或 i8 -> i32 走符号扩展。它的 store_out 会真正收窄。一致性检查在这里是白送的:全等的覆盖实现与 L != A 的混合实现,永远不可能描述同一组类型。

另有两组实现是推导出来的,而不是逐 ISA 手写的。重量化覆盖实现覆盖 Out = i8u8,把累加侧转发到 <i8, i8, i32, i32> 的实现。f32 输出孪生覆盖 N = f16bf16<N, N, f32, f32>,写成两个具体类型头,而不是对 N 泛型的一个覆盖实现。原因是一个泛型覆盖实现无法排除在 N = f32 处与同质覆盖实现相撞。

KernelSimd 还承载另外两条接缝。dot_accumulate 是点积接缝,只有具备点积能力的令牌才会覆写它,默认是 unreachable!requant_store 是向量化的重量化存储,遵循同一种模式。

把这道接缝与驱动层分块绑在一起的常量是 OUT_IS_ACC。驱动层通常把 k 切成若干 kc 面板来累加。部分和经由 C 往返,第一个面板之后 beta = 1。这种往返只有在 Out == Acc 时才是精确的。

当输出比累加器窄时,这种往返会在每个面板边界舍入到 16 位。于是窄类型家族声明 OUT_IS_ACC = false,驱动层随之回应以 kc = k:单个深度面板,整条收缩链在 f32 寄存器里累加完毕,最后只向窄输出舍入一次。

这条单次舍入的保证,正是混合精度结果站得住脚的根基。它也有代价:k 很大时,单面板意味着它的 RHS 微面板会撑破 L2 缓存。f32 输出孪生的存在,就是为了偿还这笔代价。点积内核与深K孪生讲了它们是怎么做到的。

开闭性质的证明

“家族接缝对扩展开放“这一论断不只是写在文字里,gemmkit/tests/open_closed.rs 强制执行了它。这是一个位于 crate 之外的集成测试,因此只能看见公开 API。

测试声明了 NaiveFloat,一个刻意朴素的第二浮点家族,与 FloatGemm 毫无共享。它从零重新实现微面板打包,因为 crate 内部的 pack 辅助函数对它不可见,这恰好就是第三方会遇到的处境。NaiveFloat 还提供了一个纯标量三重循环的 microkernel

测试接着驱动未经修改的泛型驱动层 driver::run::<NaiveFloat, ScalarTok, 4, 4>,跑一个 40x33x28 的问题,对照 f64 参考实现校验结果。

这个测试的主要价值在于它能编译通过。第二个家族在不改 driver.rs、不改 pack.rs 的情况下驱动了驱动层,用到的全是公开条目:gemmkit::kernel::KernelFamilygemmkit::simd::ScalarTokgemmkit::driver::runWorkspaceParallelism

任何封死这道接缝的重构都会先弄断这个文件,比如驱动层对某个具体家族的分支、新增的一个必需的私有辅助函数、或者 trait 签名里泄漏出来的内部类型。这样的重构会先伤到这个测试,然后才伤到下游用户。

更完整的测试图景,包括真实家族如何对照 oracle 交叉验证,见测试与验证。第三方能在这道接缝上搭建什么,见扩展点

点积内核与深K孪生

大多数内核家族每次只消费收缩维的一个深度步。每一步加载打包 A 面板的一列,广播打包 B 的一个元素,给每个累加器寄存器发出一条 FMA(或加宽乘加)。

两个 AVX-512 扩展打破了这个节奏。它们各自把好几个深度步折叠进一条指令。VNNI 的 vpdpbusd 把 4 个连续的 i8 深度步乘进 16 个 i32 通道中的每一个。AVX-512 BF16 的 vdpbf16ps 把 2 个连续的 bf16 深度步折叠进每个 f32 通道。

一条一次消费多个深度步的指令,要求这些深度步在内存中相邻。对浮点来说,这还会改变累加的舍入方式。所以 gemmkit 给点积内核单独建了内核家族和 ISA 令牌,而不是把它们藏在共享微内核的一个分支里。

本页依次讲三件事:承载点积内核的两道接缝、两个具体的点积内核本身,以及两个窄类型家族共享的深收缩路线。

为什么点积指令需要独立家族

点积内核与它的加宽(widen)兄弟恰好在两处不同。这两处差异分别落在引擎的两条不同扩展轴上。家族与令牌的划分见标量与内核家族

打包布局是家族的事。KernelFamily::pack_lhspack_rhs 不接受 ISA 参数,所以不同的交织方式只能由家族本身决定。这正是 Bf16DotGemmMixedGemm<bf16> 的兄弟家族、而不是它内部一个分支的原因。

内层循环是令牌的事。只有真正具备 vpdpbusdvdpbf16ps 的 CPU 才能执行它。所以这条指令藏在一个 KernelSimd 方法后面,只有具备点积能力的令牌才覆写这个方法。

具体的点积家族是:每指令折叠 4 个深度步的 IntGemmVnni(及其重量化变体 IntGemmVnniQ),以及每指令折叠 2 个深度步的 Bf16DotGemm。再加上本页末尾要讲的两个 f32 输出孪生家族,凑齐这一整套。

DEPTH_MULTIPLE 与 k 组打包

每指令折叠 Q 个深度步的家族声明 const DEPTH_MULTIPLE: usize = Q,默认值是 1gemmkit/src/kernel.rs 里写明的契约是这样运作的:家族的打包例程要写出 width * kc.next_multiple_of(Q) 个元素的面板,并把深度尾部补齐。驱动层按同样补齐后的深度为打包面板定步长,让两边保持锁步。对所有普通家族来说,DEPTH_MULTIPLE = 1,每一次 next_multiple_of 调用都会退化为恒等。

布局本身来自一个共享例程,gemmkit/src/pack.rs 里的 pack_kgroup_panels。它是交织下标算术的唯一权威来源。

普通的 pack_panels 按深度为主序存储面板。每个深度步存 width 个连续的先导元素:LHS 是 mr 行,RHS 是 nr 列。

pack_kgroup_panels 则把深度轴按 Q 个一组分组,让同一通道连续的 Q 个深度值变得连续。面板内第 g 组、第 i 个通道、组内第 t 个位置,落在偏移 g*width*Q + i*Q + t 处。

这正是一条点积指令要读的形状。一个 64 字节的 A 寄存器覆盖 LANES 行乘 Q 个连续深度元素。一个 B 组把某一列 Q 个连续深度值当作一次 32 位加载来广播。

这个共享打包器还有两个要点值得一提。它接受一个逐元素变换 xform:对 bf16 是恒等变换,对 VNNI 的 A 操作数是 +128 偏置。打包器把每个填充位置都填成 xform(0)。填充位置指超出块的先导位置,以及超出 kc 的深度位置。这让填充始终与活跃元素保持一致。

交织布局无法原地读取。所以每个点积家族都设置 FORCE_PACK_LHSFORCE_PACK_RHS,压过驱动层基于代价的打包决策。点积内核永远要付打包成本。这恰好是下文各个门槛所要对冲的开销。

pack.rs 测试里有一个字节级 oracle,用朴素写法重新实现该布局。它验证真实例程在宽度尾部、深度尾部、带步长源等各种情形下都逐位复现该布局。

在消费侧,KernelSimd::dot_accumulate 是各家族用来替代加宽 FMA 循环的接缝。它的默认实现是 unreachable!。只有具备点积能力的令牌会覆写它,也只有点积家族会调用它。

Avx512VnniAvx512Bf16 之所以是独立于 Avx512F 的令牌,是因为 #[target_feature] 按令牌生效。_mm512_dpbusd_epi32 需要 avx512vnni 的代码生成上下文。Avx512F::vectorize 只建立 avx512f 上下文,提供不了这个。

该方法接收真实、未补齐的 kc。它从补齐深度的面板中读取 ceil(kc / Q) 个指令组。一切符号或偏置修正都在内部完成,所以返回时累加器已经持有真实的 sum_k(A*B)

折叠之所以放在这个专用接缝上,而不是通用的 accumulate_tile 上,是有明文记录的原因的:折叠深度步会重塑累加的舍入方式,而 accumulate_tile 的契约不允许这样做。

i8 经 vpdpbusd

vpdpbusd 计算的是无符号乘有符号的点积。它的第一操作数取 u8,第二操作数取 i8。而 GEMM 要的是有符号乘有符号。

修正是代数式的,不是逐元素的。LHS 打包把每个字节偏移 +128,送进无符号域。这个变换就是 gemmkit/src/kernel/int.rs 里的 vnni_a_xform。它使用常量 VNNI_A_BIAS = 128,这个常量只在 gemmkit/src/simd.rs 里定义一次,让打包与修正永远不会失配。

sum_k((A+128)*B) = sum_k(A*B) + 128*sum_k(B)。所以内核通过减去逐列修正项 128 * sum_k(B[k][j]) 来还原真实乘积。Avx512Vnni::dot_accumulate 在向量循环之前,先用一小段标量代码对有符号的打包 B 面板求出这些列和,然后在最后把广播后的修正从每个累加器里减去。

填充在这套方案里也是配合好的。A 的填充是 xform(0) = 128,修正项恰好抵消它的贡献。B 的填充是 0,对乘积和列和都没有贡献。

i32 累加按环绕方式进行,而环绕加法在模 2^32 下满足结合律。所以把求和重新按四个一组分组,再加上偏置修正,与升序 k 的加宽求和逐位相等。IntGemmVnni 与加宽版的 IntGemm 在任何输入上都产生完全相同的输出。

ISA 的选择因此永远不会改变 i8 的结果。这比可复现性契约所要求的还要强。那份契约只承诺在固定机器和固定配置下结果可复现,并不承诺跨内核选择时逐位一致。

正是这种可以随时换内核的自由,被小规模并行回退门槛拿来利用。VNNI 的打包在两个操作数上都是强制的。在一个小型多线程问题上,这道打包壁垒有可能压过它本该节省下来的计算量。

触发这个回退需要同时满足三个条件。ISA 选择必须是自动的。并行度必须是 Rayon(n)n != 1。而且 m*n*k 必须低于 GEMMKIT_I8_VNNI_MIN_PAR_MNK,它的默认值是 768^3。三个条件同时成立时,dispatch/int.rs 才会把调用交给可原地读取的加宽内核。

串行运行和大规模并行运行都继续用 VNNI。强制设定 GEMMKIT_REQUIRE_ISA=avx512vnni 会完全关闭这道门槛,因为强制锁定就必须运行指定的那个内核。

预打包 RHS 路径同样绕过这道门槛,原因有二。其一,k 四元组交织的缓冲区只有 VNNI 家族能消费。其二,这道门槛本来要对冲的打包壁垒,已经在预打包阶段被摊销过一次了。VNNI 的 RHS 打包本来在每次调用中都是强制的,所以预打包在这里带来的收益,比任何一个能原地读取操作数的内核都大。

bf16 经 vdpbf16ps

Bf16DotGemm 是浮点这一侧的对应者。它的 DEPTH_MULTIPLE2。两个操作数都按 k 对交织打包,每一对存成一个 32 位的 __m512bh 元素。dot_accumulate 每个累加器、每个对步发出一条 vdpbf16ps

累加之后的一切都与 MixedGemm<bf16> 逐字共用,包括 alpha 折叠,以及加宽读、收窄写的 epilogue,都经由共享的 mixed_epilogue 辅助函数完成。这个家族保持 OUT_IS_ACC = false,让整个收缩在 f32 里累加、只向 bf16 舍入一次。

数值故事与 VNNI 有一处本质区别。vdpbf16ps 融合的二项点积,舍入方式与两条独立的加宽 FMA 不同。所以点积内核的结果与加宽路径只是容差相等,不是逐位相等。

这恰好就是引擎一致性标准所允许的:结果必须在固定的输入、环境与配置下可复现,不需要跨内核选择逐位一致。点积内核本身完全确定。串行、并行、预打包运行共用同一个内核和同一套打包布局,因此彼此逐位复现。

这条路径上没有任何规模门槛。只要 CPU 报告支持 avx512bf16,自动选择就会优先选 Bf16DotGemm,因为它相对普通加宽路径是结构性的胜出。和 VNNI 不同,这里没有小规模并行回退。

只有三条特殊路径改道是例外。gemv、small_mn、以及小 k 形状,都刻意留在 MixedGemm<bf16> 的加宽接缝上。微小或退化的输出折叠不了什么,点积打包的深度补齐在那里纯属亏损。i8 分发出于同样的理由,把自己的微小形状也改道到加宽内核。

深 K 问题

OUT_IS_ACC = false 用一种结构性代价换来单次舍入。驱动层运行 kc = k,用一个覆盖整个收缩的深度面板。这取代了每个同质家族原本都有的缓存模型 kc 切片(见分块与缓存模型)。

一次微块调用要读的 RHS 微面板于是有 nr * k * sizeof(N) 字节。这个微面板一旦超出 L2 缓存,一个列条带里 m/mr 次微块调用中的每一次,都要转而从 L3 或 DRAM 里把它流过来。更大的 mr * k LHS 微面板也一样要从那里流过来。这样形成的悬崖很陡峭。微面板还在 L2 里时吞吐接近峰值,一旦装不下就会下降。

gemmkit/src/dispatch/mixed.rs 里的接入门把这个微面板大小拿去和一个字节阈值比较:

#![allow(unused)]
fn main() {
let engage_deep_k = NR
    .checked_mul(t.k)
    .and_then(|x| x.checked_mul(core::mem::size_of::<N>()))
    .is_some_and(|bytes| bytes > crate::cache::deep_k_engage_bytes());
if engage_deep_k {
    run_deep_k_twin::<N, Fam::Twin, S, MR_REG, NR>(simd, &t, par, ws);
    return;
}
}

这个阈值就是 GEMMKIT_DEEP_KC_BYTES 旋钮,非零时原样采用。默认值 0 把阈值推导为每工作线程有效 L2 容量的一半。

用 L2 的一半,而不是整个 L2,是刻意的选择。如果把门槛设成整个 L2 大小,会接入得太晚,晚到微面板早已装不下了,也就错过了悬崖。用一半留出空间给工作集里的其他部分,能在还来得及避开悬崖的时候就接入孪生路线。

checked_mul 这条链条是失败即关闭的。一个广播操作数可能带着逻辑上荒谬的 k 通过校验。溢出的尺寸必须转而落到单面板上。单面板自己的打包尺寸检查会拒绝这个问题,而不是接入一个会对着这个 k 永远多切片下去的孪生。

f32 输出孪生

在门槛之上,分发根本不运行窄类型家族。一个小小的 DeepKTwin trait 把每个窄家族映射到它的 f32 输出孪生。MixedGemm<N> 映射到 MixedGemmF32<N>Bf16DotGemm 映射到 Bf16DotGemmF32。每个孪生唯一的类型变化是 Out = f32 = Acc

这一处变化让 OUT_IS_ACC 回到默认值 true。驱动层普通的多切片 K 分块随即原样适用,每个切片的面板重新变得 L2 常驻。这正是孪生存在的全部意义。

打包布局和累加循环都逐字沿用窄家族的。MixedGemmF32 复用 pack_panels 和共享的加宽 FMA 辅助函数。Bf16DotGemmF32 复用 pack_kgroup_panelsdot_accumulate。这些累加辅助函数只触及 KernelSimd 接缝的输入侧。所以它们产出的累加器,与窄家族本会算出的结果逐字节相同。

孪生以 alpha = 1beta = 0 运行,写入一块 m x n 列主序的 f32 暂存缓冲区。这块缓冲区取自一个专用的 Workspace。深 K 按定义就是大 k 场景,所以一块 m*n 的 f32 分配可以忽略不计。把它独立出来,还能让池化的打包工作区留给孪生驱动使用。

之后,一次向量化扫描计算 narrow(alpha*scratch + beta*widen(C))。这逐个操作复刻了 mixed_epilogue 的算术,包括同一个 store_out 收窄步骤。

让这条路线不止于近似的关键,在于切片是怎么衔接起来的。孪生的微内核通过第三个 KernelSimd<N, N, f32, f32> 接缝(gemmkit/src/kernel/mixed.rs 里的 twin_seed),从暂存缓冲区为累加器寄存器播种。在一个累加切片上,它把运行中的部分和加载进寄存器,继续升序 k 的累加链。它从不把某个切片单独从零求和、再事后相加。

f32 的存取往返是精确的。所以多切片之和恰好就是单面板之和,只是在切片边界处被切开而已。对 beta 属于 {0, 1} 的情形,深 K 结果与单面板结果逐字节相同。

对一般的 beta,结果就只能做到容差一致了。原因很平常。单面板在完整微块上用一条 FMA 融合 beta*C + AB,但在边缘微块上是不融合地组合这两项。没有哪一种扫描方式能同时匹配这两种情形。

串行运行与并行运行在任何情况下都保持逐位一致。孪生驱动的分块不依赖线程数,最后那次扫描又是逐元素的。

点积孪生还多一条对齐规则。驱动层把分块 kc 向上取整到 DEPTH_MULTIPLE 的倍数,这样内部的切片边界就永远不会切开一个 k 对。被切开的一对会在收缩中途补零,从而错误地重组融合点积。有了这条规则,只有最后的短尾会被补齐,与单面板情形完全一致。

有三条路线刻意保留单面板。它们是门槛之下的浅 k、融合 epilogue 的入口,以及预打包 RHS 路径。在预打包路径上,DEPTH_MULTIPLE > 1 的缓冲区要求整个收缩必须是同一个深度切片。驱动层用一个硬性 assert 强制这一点,因为违反它会静默地错位读取微面板。

这些等价性声明都有直接的测试。gemmkit/tests/deep_k_narrow.rs 拨动 GEMMKIT_DEEP_KC_BYTES:取值 1 会在任何 k 下强制走孪生路线,取值 usize::MAX 会强制走单面板。这个测试在宿主机选中的任意 ISA 上,检验 beta 属于 {0, 1} 时逐字节相等,一般 beta 时容差一致。调优旋钮与其余旋钮一起记录了这个旋钮。

复数拆分内核

复数 GEMM 是唯一不搭乘 FloatGemm 的同质类型家族。理论上它本可以搭乘。Complex<f32>Complex<f64> 都在自身类型里累加,所以 Lhs = Rhs = Acc = Out。这正是浮点家族能处理的形态。

问题出在内存布局上。num_complex 把一个复数存成相邻的 (re, im) 对。于是从复数切片加载的 SIMD 寄存器持有的是 re, im, re, im, ...。一次复数乘法需要跨通道组合:实部是 re*re - im*im,虚部是 re*im + im*re

在交错的通道上,这些组合迫使最内层循环里出现 shuffle 和 fmaddsub 一类指令。这个代价每个深度步都要重复一次,总共 O(mnk) 次。gemmkit 的做法是在打包时把布局改写一次,让热循环保持纯实数 FMA,循环内完全没有 shuffle。

本页依次讲五件事:拆分设计本身、共轭如何从中免费落出、内核经由的接缝、微块形状背后的寄存器预算算术,以及数值上的保证。

拆分布局

家族是 gemmkit/src/kernel/complex.rs 里的 ComplexGemm<T, CONJ_A, CONJ_B>。设计的核心就在它的打包例程里。

pack_planar 把每个微面板按结构数组(SoA)形式铺开。每个深度步,面板先存 width 个实部,紧接着存 width 个虚部。width 对 LHS 是 mr,对 RHS 是 nr,步长互换的方式与共享的 pack_panels 完全一致。

内核于是用普通的连续加载,取到一整个实部寄存器和一整个虚部寄存器。去交错的成本从 kc 内层循环移到了打包这一步。摊销后的成本变成 O(MK + KN),而不是 O(MNK)

内核只能消费这种平面布局。所以两个操作数永远都要打包。家族设置 FORCE_PACK_LHS = FORCE_PACK_RHS = true,压过驱动层原本基于代价、可能原地读取操作数的决策。

pack_planar 复刻了 pack_panels 的两条写入路径。一条是先导维连续时的直接遍历。另一条是带步长源时的缓存分块转置,这样行主序操作数打包时不会每个元素都错过一次缓存。两条路径写出的面板逐字节相同,只是写入顺序不同。共享框架见打包与工作区

共轭是打包时的符号翻转

共轭只对虚部取负。打包本来就单独写虚部平面。所以 conj(A)*BA*conj(B) 在热循环里零成本。

设置 CONJ_ACONJ_B 这个 const 泛型,会让打包器在拷贝时对虚部平面取负。这是真正的取负,+0.0 会映射到 -0.0,与 num_complex.conj() 一致。同一个实数 FMA 循环随后原样运行,任何地方都没有逐元素的共轭分支。

这也是强制打包标志存在的第二个理由。当打包做的不止是普通拷贝时,这个变换必须每次都执行。

运行时到编译时的桥梁在 gemmkit/src/dispatch/complex.rs 里。公开入口 gemm_cplxconj_aconj_b 当作普通 bool 接收。run_complex 对这一对布尔值只 match 一次,分发到四个 ComplexGemm 单态化里匹配的那一个。这个分支每次调用只发生一次,绝不会进到循环里面。

这里还带出一处细节。把行主序倾向的 C 规范化的方向交换,实际计算的是 C^T = B^T * A^T。由于 (conj(A)*B)^T = B^T * conj(A)^T,这次交换必须连同共轭标志一起交换。这个交换在 match 之前就完成了。

输出共轭(conjC)没有实现。在退化路径上,也就是 k == 0alpha == 0 时,这些标志根本无关紧要。没有 A*B 项,也就没有什么可共轭的。

热循环:每次复数乘加四条实数 FMA

循环真正跑起来之前,还有一个分层问题要解决。家族是同质的,所以驱动层的约束是 T 取复数类型的 KernelSimd<T, T, T, T>。这个约束只提供 SimdOps<Complex<..>>,而不是拆分内核真正需要的实数运算。

桥梁是 SimdOps::cplx_microkernel 这道接缝。家族的 microkernel 会转发给它。每个 ISA 令牌的覆写,由 gemmkit/src/simd/complex.rs 里的胶水宏 impl_complex_simd! 生成,再转发给唯一一个共享的、ISA 泛型的函数 soa_microkernel,它写在 S: SimdOps<C::Real> 之上。

累加器在家族接缝处保持复数类型。所以复数的 alphabeta 能原样穿过驱动层。但在接缝内部,累加器其实是两组实数寄存器。

薄薄的 SimdOps<Complex<..>> 胶水存在的唯一理由,是让驱动层能读到 LANES,也让同质 blanket 实现能够适用。它的元素运算全是 unreachable!,因为复数 GEMM 从不调用它们。LANES 被设成实数通道数,于是一个实数通道对应一个复数行,驱动层的 mr = MR_REG * LANES 数的正是复数行数。

循环本身,摘自 gemmkit/src/simd/complex.rs

#![allow(unused)]
fn main() {
for p in 0..kc {
    let are_p = a_re.add(p * 2 * mr); // re plane of this depth step
    let aim_p = are_p.add(mr); // im plane (offset by `mr`)
    let ar: [<S as SimdOps<C::Real>>::Reg; MR_REG] =
        core::array::from_fn(|i| simd.loadu(are_p.add(i * lanes)));
    let ai: [<S as SimdOps<C::Real>>::Reg; MR_REG] =
        core::array::from_fn(|i| simd.loadu(aim_p.add(i * lanes)));
    let bre_p = b_re.add(p * 2 * NR);
    let bim_p = bre_p.add(NR);
    for j in 0..NR {
        let br = simd.splat(*bre_p.add(j));
        let bi = simd.splat(*bim_p.add(j));
        for i in 0..MR_REG {
            acc_re[j][i] = simd.mul_add(ar[i], br, acc_re[j][i]); // += ar*br
            acc_re[j][i] = simd.fnma(ai[i], bi, acc_re[j][i]); //    -= ai*bi
            acc_im[j][i] = simd.mul_add(ar[i], bi, acc_im[j][i]); // += ar*bi
            acc_im[j][i] = simd.mul_add(ai[i], br, acc_im[j][i]); // += ai*br
        }
    }
}
}

一次复数乘加是四条融合的实数步骤,流入两组累加器。acc_re 拿到一条 mul_add 和一条 fnma(融合取负乘加,x86 上是 vfnmadd)。acc_im 拿到两条 mul_add

每条操作都是作用在连续加载和标量广播上的普通逐通道 FMA。在 epilogue 之前,没有任何操作会跨通道。固定的逐 p 顺序是刻意安排的。正是它让同一矩阵的完整微块与边缘微块舍入完全一致。

循环结束后,两组累加器排入平面暂存区。标量 epilogue 折叠复数 alphaalpha == 1 时跳过复数乘法),按情形合并 beta*C,并在写出时重新交错。这是一次摊销 O(MN) 的扫描,统一处理完整、边缘和带步长的输出微块。

打包中的标量去交错、以及 epilogue 中的标量重交错,都是刻意的选择,不是疏忽。内层循环本身就占了内核总成本的绝大部分。所以在每种 ISA 上,通用的标量路径对这两步来说都是下限。

寄存器压力与 NR 的选择

拆分设计让累加器数量翻倍。一个 MR_REG x NR 的复数微块需要 2*MR_REG*NR 个累加寄存器(一组实部、一组虚部),外加 2*MR_REG 个 A 平面寄存器,以及每个列步 2 个 B 广播。这份预算在 gemmkit/src/dispatch/complex.rs 里逐微块记录在案。它让复数微块比浮点微块更小,本页出现的这些铺块形状也都由它推出。

在 FMA 上,16 个 YMM 寄存器里,c32 取 MR_REG = 1(8 个实数通道对应 8 个复数行),NR = 5。这是 10 个累加器,加 2 个 A 寄存器,加 2 个 B 广播,占 16 个寄存器中的 14 个。空出的那两个很关键。若换成 NR = 6 的 16 占 16 满配微块,会把累加器溢出到栈上,所以 NR 被收缩到 5,代码注释里记下了原因。

AVX-512 的 32 个 ZMM 寄存器缓解了这份压力。c32 取 MR_REG = 2NR = 6,用掉 24 + 4 + 2 = 32 个中的 30 个。NEON 有 32 个向量寄存器,取 MR_REG = 2NR = 5,占 32 个中的 26 个,给在途的加载临时量留出空间。wasm 的 simd128MR_REG = 1NR = 4,共 12 个活跃的 v128 寄存器。

每个 c64 变体都沿用其 c32 兄弟的 MR_REGNR,只是通道数减半。这份预算算术本来就与通道数无关。

精度与可复现性

复数没有任何特殊路径。run_complex 把所有形状都送进 driver::run,没有 gemv、small_mn 或小 k 分支。特殊路径那套机制只服务实数浮点和整数。

这让数值契约很容易陈述。gemmkit/tests/correctness/complex.rs 对每一条都有直接测试。

确定性与线程无关性是按位成立的。分块与线程数无关,所以同一问题的串行与并行运行会产生逐位相同的输出。测试直接断言各线程数下 re/im 的原始位模式相等。

在单次运行内部,四条 FMA 固定的逐步顺序,让同一矩阵的完整微块和边缘微块舍入完全一致。所以结果不依赖于微块边界恰好落在哪里。

共轭完全不引入任何舍入。它只是对精确值的符号翻转。一个专门的小整数输入测试(其中每个乘积与和都精确可表示)用精确相等而不是容差,对照朴素参考检验全部四种共轭组合。

对外部 oracle 而言,标准必然要放宽一些。正确性套件把 gemm_cplx(包括每一种共轭组合)拿去和 gemm crate 比较,用的是套件惯常的 L2 型容差。它还单独检验了一个负行步长视图加共轭的情形,对照一个行反转的参考实现。之所以放宽标准,是因为分块的 SoA 收缩本来就有理由和另一个引擎的求和顺序算出不同的舍入结果。

这正是可复现性契约用在复数上的样子:固定的输入、环境与配置下,结果是同样的位。跨不同引擎或不同求和顺序时,结果只需容差一致。

融合偏置入口是拿到按位保证的例外。gemm_cplx_fused 支持逐行或逐列的复数偏置。它刻意不支持任何激活函数,因为 ReLU 一类基于序的激活在无序域上没有定义。

它的 epilogue 从不触碰内核本身的算术。SoA 内核存入的位,正是普通 gemm_cplx 会存入的那些位。一个局限于单个微块的后处理,只在最后一个深度面板上就地映射这些位。复数家族是 OUT_IS_ACC = true,所以中间面板必须保留原始的部分和。

结果与“先跑 gemm_cplx、再做同样的逐元素偏置加法“逐位相同,对每种形状、每种共轭组合都成立。Identity 单态化会把这个后处理整个常量折叠掉。所以非融合路径不为这个钩子的存在付出任何代价。通用机制见 Epilogue融合

分块与缓存模型

驱动器最底层的微内核负责计算 C 的一个 MR x NR 分块。它把这个分块保存在寄存器中,沿深度方向依次读入 A 的一条微面板和 B 的一条微面板。只有当这两路数据都来自邻近的缓存时,这个循环才能跑到机器峰值。一次 GEMM 触及的数据远远超过任何一级缓存的容量:A 的每个元素要在 n 个输出列上复用,B 的每个元素要在 m 个输出行上复用。

分块(blocking)就是对这种复用的安排。它把问题切分开,让每个操作数的分块只被读入某一级缓存一次,然后被反复读取,直到被逐出。KC 沿深度方向切片,让一次分块乘法所读的两条微面板,在整个分块计算期间都留在 L1 中。MC 决定打包后 A 宏面板的大小,使它在驱动器扫过当前列块的每个列瓦片时始终留在 L2 中。NC 决定打包后 B 宏面板的大小,使它在每个行块轮流扫过时始终留在 L3 中。让面板在整个循环嵌套中保持驻留,正是分块的全部意义所在。没有它,同样的字节就要从 DRAM 中反复读取 mnk 遍。

许多库按微架构把 (MC, KC, NC) 写死。gemmkit 则在每次调用时,在 CacheTopology::blockinggemmkit/src/cache.rs,L3 层)里解析地计算这三个值。这个函数遵循 BLIS 模型,根据运行时探测到的缓存几何来推导尺寸。它的输入是微块尺寸 (mr, nr)、单个打包输入元素的字节大小,以及问题形状 (m, n, k)。输出是驱动器循环嵌套所用的 Blocking { mc, kc, nc } 三元组。一次GEMM调用的生命周期展示了每个值落在嵌套的哪个位置。本页说明这些值是如何推导出来的,以及模型为何是这个样子。

三个约束,三个块尺寸

KC:两条微面板同驻 L1 且不自我逐出

每次微块调用都要走 kc 个深度步,每一步读取 mr 个打包后的 A 元素和 nr 个打包后的 B 元素。因此一条 mr x kc 的微面板和一条 nr x kc 的微面板必须在整个分块计算期间同时留在 L1d 中。这里微妙之处在于不自我逐出:缓存不是一个字节池,而是由若干组(set)划分出来的若干路(way)。一条面板如果把太多自己的缓存行映射到同一组上,还没等总大小逼近缓存容量,就已经把自己挤出去了。

模型因此以缓存行和组为单位计算,而不是以字节为单位。它先算出每条微面板每一个深度步占用多少条 L1 缓存行,再选出合并占用不超过 L1 相联度的最大 kc。然后把这个结果抬高到 GEMMKIT_KC_MIN 下限(默认 512),确保容量较小的 L1 也不会饿死微内核的深度扫描,最后把它钳制在 k 以内。最后还有一次再平衡:把 k 切成 ceil(k / kc) 份大小相近的面板,避免最后一个深度切片过于零碎。

MC:A 宏面板占 L2,扣除 B 的份额

在一个行块之内,驱动器要在当前列块的每个列瓦片上复用打包后的 mc x kc A 面板,所以这个面板理应填满 L2。但它不能占满整个 L2,因为同一深度切片的 nr x kc B 微面板也要在每次分块调用时流经 L2。模型先数出这条微面板占用多少 L2 路,预留这些路再加一路备用,剩下的容量才交给 A

它用这部分剩余容量除以 kc 得到 mc,向下取整到 mr 的倍数,再做一次再平衡使各行块大小均匀。最后,一个 BLIS 式的硬上限 GEMMKIT_MC_REG_PANELS * MR 行(默认 8 个微块行)钳住这个结果。这个上限是一个标定值,而不是从 L2 容量项严格推导出来的边界。实际情况中,这个上限往往先于 L2 容量项生效,所以 MC 通常就是 MR 的一个小倍数,L2 容量项大多只是留出的余量。

NC:B 宏面板占 L3,无 L3 时用面板数上限

有 L3 时,模型为途经的 A 流量预留一路,其余容量预算给打包后的 kc x nc B 宏面板。它用这部分容量除以 kc 得到 nc,向下取整到 nr 的倍数,再对 n 做一次再平衡。

有些机器根本不报告 L3。例如 Apple Silicon 的簇共享 L2 就是整个层级的顶端。这种情况下,模型改为跑满整个 N,直到一个面板数上限为止:ncGEMMKIT_NC_NO_L3_PANELS * nr(默认 512 个面板,也就是 nr = 4 时的 2048 列),并以 n 封顶。既然没有 L3 能让 B 驻留,B 无论如何都要从 DRAM 流入。这个上限只是限制共享的打包 B 缓冲区大小,而不是在为驻留建模。

按打包元素、而非累加器元素计大小

sizeof 参数是单个打包输入元素的大小。驱动器传入的是 size_of::<Fam::Lhs>(),而不是累加器的大小,因为模型预算的面板本来就是以打包后的 Lhs/Rhs 为单位存储的。对 f32f64 而言,这两个大小恰好相同,所以没有区别。

对窄类型来说,这个区分就很关键了。i8 每个元素打包 1 字节,对应的累加器 i32 却是 4 字节。f16bf16 每个元素打包 2 字节,累加器 f32 是 4 字节。如果按累加器计算大小,它们的 kcnc 会被砍到缓存实际可容纳量的四分之一或一半。窄类型因此得到成比例更深的块,这也是它们能在同样的硬件上跑得比 f32 更快的原因之一。预打包入口复用同一个模型,只是用一个哨兵行数,使预打包操作数的几何形状与最终的 m 无关。

把输出 tile 预取到 LLC 之外

这三个块尺寸让 AB 面板保持驻留,却完全没有提到 C。每次微 tile 调用都要对它的 mr x nr 输出 tile 做一次读取、修改、写回。一旦一次调用的工作集(ABC 三者字节数之和)超出单核可达的 LLC 容量,这个输出 tile 就不再驻留在缓存里,它的写回就要伸进 DRAM。

驱动器用一次软件预取来应对,这个决策每次调用只做一次。它把工作集与 cache::prefetch_ws_bytes(即 GEMMKIT_PREFETCH_MIN_BYTES 门槛,0 表示自动:取单核可达的 LLC,有 L3 就取 L3,否则取 L2)做比较。一旦超过门槛,驱动器就会在每个输出微 tile 对应的微内核调用之前,对它发出一次 T0 预取。这样可以在微内核仍在计算的时候,把即将被写入的缓存行提前拉进 L1。

这个预取沿着 tile 的单位步长维度,整条整条地拉取 64 字节的缓存行。两个维度都带步长的 tile 没有连续的缓存行可拉,会被跳过。驱动器只在 x86_64 上发出这个提示(prefetcht0,基线 SSE 指令,不需要额外的 feature 门控)。在其他目标上它会退化成空操作,因此 aarch64 和 wasm 都不受影响。这个预取只搬动缓存行,绝不涉及任何算术运算,所以无论门槛是打开、关闭还是被强制指定,结果都逐位一致。门槛以下,tile 本就驻留在缓存中,预取路径不会带来额外开销。

小矩阵捷径

mn 都不超过 GEMMKIT_TINY_BLOCK_DIM(默认 64)时,驱动器会跳过完整模型。它把 kc 设为 k,并钳制在 GEMMKIT_KC 上限内(默认 2048,aarch64 上为 16384)。它把 mc 设为在该深度下仍能让面板留在 L2 里的行数,并以 m 本身封顶。它把 nc 设为 n 向上取整到 nr 的倍数。

一个工作集整体就能装进 L2 的问题,从三级驻留分析中得不到任何好处。这个捷径把省下来的运算花在真正要紧的地方:小规模乘积中占主导地位的、每次调用固定不变的那部分开销。

这个上限以 4 字节元素为单位计数,窄元素会去除它。于是 f16 拿到 f32 两倍的深度,int8 拿到 4 倍的深度。保持不变的是字节预算,而硬件限制正是关于字节的。因此一个数就标定了所有元素族,在 f32 上跑出来的调优结果也能迁移到其余元素族。

这个上限决定深度切片的数量,而切片数量牵动两项方向相反的成本。每多一个切片,就要多重读、重写一遍 C,多进入驱动器一次,在并行路径上还要多 fork 一次 worker。这一侧要求上限尽量深。

而切片越深,打包后的 A、B 面板就越大。这些面板必须留在私有 L2 里,所以这一侧要求上限尽量浅。在 x86 上,默认值让面板占到约 1.1 MiB,正好是一块 Zen5 L2。两者之中并行侧的成本更大,所以默认值取在驻留上限处,而不是取在它下方。

元素正是这两项成本对“该不该做除法“意见不一的场合。字节预算不等于切片预算。在固定预算下,16 字节元素拿到的切片数是 4 字节元素的 4 倍,于是每切片成本它要多付 4 次。驻留一侧要求做除法,切片数一侧不要求。

哪一侧胜出是机器的属性,所以这个除数带一个按架构分叉的上限。在 x86 上私有 L2 只有 1 MiB,驻留先咬住,除法对所有元素大小都生效。在 aarch64 上有效 L2 大 4 倍,而且统一内存能很好地喂溢出的面板,驻留几乎咬不住,所以宽元素在那里保留完整的上限。在 M4 Max 上实测,对它做除法让 c64 的并行路径损失 21% 到 51%。

探测:一条不会失败的回退链

模型的好坏取决于喂给它的几何数据的质量,而查询缓存几何并没有一种可移植的办法。gemmkit 因此运行一条尽力而为的回退链,其中 #[cfg] 只用来挑选探测方法,从不用来决定数值本身。#[cfg(target_arch)] 分不清 Intel 部件和 AMD 部件,虚拟机或容器又可能屏蔽 CPUID、隐藏 /sys。所以每个后端都返回一个 Option,链条最终落在一个绝不会失败的常量上。

#![allow(unused)]
fn main() {
// gemmkit/src/cache.rs
#[cfg(feature = "std")]
fn detect() -> CacheTopology {
    // try the CPUID backend
    #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), not(miri)))]
    if let Some(t) = cpuid::detect().filter(plausible) {
        return t;
    }
    // try the sysfs backend
    #[cfg(all(target_os = "linux", not(miri)))]
    if let Some(t) = sysfs::detect().filter(plausible) {
        return t;
    }
    // try the sysctl backend
    #[cfg(all(target_os = "macos", not(miri)))]
    if let Some(t) = sysctl::detect().filter(plausible) {
        return t;
    }
    ZEN5_FALLBACK
}
}

各个后端按以下顺序运行。

CPUIDcache/cpuid.rs)是一条指令,因此无论操作系统如何,它在容器和大多数虚拟机里都能工作。CPUID 通过按缓存实例枚举的拓扑叶(Intel 的 04h、AMD 的 0x8000_001D)读取两家厂商的信息,这个叶描述的是执行核实际可达的那一块缓存。在多 die 的部件上,比如双 CCD 的 Ryzen,L3 报出的是一个核真正能命中的那个复合体(9950X 上是 32 MiB),而不是整片封装的总量。这个单核可达的数值,正是每个消费者都想要的语义。没有该叶的 AMD 部件或虚拟化环境会回退到传统的 L1 叶(0x8000_0005)和 L2/L3 叶(0x8000_0006)。在那里,L3 大小以 512 KiB 为单位、按整片总量报告,16 路相联度甚至无法编码。

Linux sysfscache/sysfs.rs)用纯 std::fs 解析 /sys/devices/system/cpu/cpu0/cache/index*/。它在 x86 Linux 上是一个回退项,用于应对屏蔽了 CPUID 的虚拟化环境;在没有 CPUID 指令的 aarch64 Linux 上,它则是主要的数据来源。

macOS sysctlcache/sysctl.rs)通过一段两行的 extern "C" 声明读取 sysctlbyname 键,不引入任何 libc 依赖。它优先读取 Apple Silicon 按性能级划分的键(hw.perflevel0.*,对应 P 核),把扁平的 Intel Mac 键作为回退。sysctl 不暴露相联度信息,所以这个后端采用保守的典型值。这样做是安全的,因为模型只需要相联度的近似值,并用 .max(2) 兜底。

链条的最末端是 ZEN5_FALLBACK,一个在 Ryzen 9950X 开发机上标定的静态默认值。L1d 是 48 KiB、12 路。L2 是 1 MiB、16 路、私有。L3 是 32 MiB、16 路。

两道守卫让这条链不只是“按顺序尝试“那么简单。plausible 会剔除半残的读数:只要有一级小于 4 KiB、缓存行小于 16 字节,或者相联度为零,整个后端就会被判定失败,被屏蔽的叶因此不可能用零值污染分块结果。探测本身每个进程最多只运行一次:Machine::current() 把拓扑数据连同操作系统页大小(getpagesize,校验为 4 KiB 到 2 MiB 之间的 2 的幂)一起用 OnceLock 记忆化。这个页大小驱动着打包与工作区里描述的 LHS 打包步长门槛。no_std 构建完全跳过探测,直接使用 Zen5 回退值和 4 KiB 页大小。

shared_by:对驱动器实际放置的数据的争用

每个 Level 携带 bytesassocline 三个字段,再加上一个推导字段 shared_by,它把该级缓存的容量除成模型实际预算所用的 effective_bytes。把硬件层面的核共享数直接填进这个字段,看起来很自然,但那是错的。shared_by 建模的其实是针对驱动器实际放在该级缓存中的数据的按 worker 争用。驱动器的放置方式是:每个 worker 的 A/B 微面板放在 L1d,每个 worker 私有的 A 宏面板放在 L2,唯一一份共享的 B 宏面板放在 L3。

这个放置方式就决定了取值。L1d 是按核私有的,它的全部容量都服务于一个 worker 的微面板,所以 shared_by = 1。L3 在硬件上被所有核共享,但驱动器放在那里的却是一条所有 worker 共同读取的面板:是同一份字节,而不是每个 worker 一份拷贝。因此整级容量都归这一条面板所有,shared_by 依然是 1。如果按原始核数去除,预算会被无谓地缩水很多倍,把 NC 白白拉低。

只有 L2 真正存放着按 worker 私有的数据,所以只有 L2 使用物理核层面的 L2 共享度。在私有 L2 的部件上,例如主流 x86 和 Neoverse,这个共享度是 1。在核簇共享一个 L2 的部件上,例如 Apple Silicon,它就是簇的大小,因为在那里,多个 worker 私有的 A 面板确实要争抢同一批路。每个后端都必须推导出这个值,而不能直接照抄一个原始计数:sysfs 用 L1d 共享列表读出的 SMT 度去除 L2 原始的 shared_cpu_list 计数,避免把超线程兄弟核重复计入;sysctl 读取 hw.perflevel0.cpusperl2;CPUID 后端则直接把它硬编码为 1,因为 x86 的 L2 本就按物理核私有。

在 x86 和 Graviton 上,整套机制因此都归约成全部为 1。它真正派上用场的是簇共享 L2 的部件:在那里,它决定了模型究竟是按一个 worker 实际独占的 L2 来分块,还是按一个必须与整个核簇共享的 L2 来分块。

线程数能挪动什么,不能挪动什么

blocking 不接收线程数参数。对 KCNC 来说,这个缺席意义重大:两者都只取决于机器本身和问题形状,所以一次串行运行和一次宽度很大的并行运行会推导出完全相同的 KCNC。这就让每次运行都得到相同的深度切片,以及每个输出元素相同的、顺序固定的深度累加链。

MC 是驱动器确实会为并行度调整的那一个分块维度。宽 worker 数有可能让扁平任务列表变得太浅:每个 worker 分不到几个块,运行的尾部就会退化成一群空等最后几块被人拿走的 worker。一旦出现这种情况,驱动器就会缩小 MC,切出更多行块,把列表加深。并行执行详细说明了这个并行作业深度下限。因此,面板边界和扁平任务列表已经不再严格独立于 worker 数了。

尽管如此,逐位一致性依然成立,因为这次缩小本身不涉及任何数值运算。MC 始终保持为 MR 的倍数,所以它产出的微 tile 集合(每个 MR 对齐的行偏移,加上唯一的那一条 m 尾巴)在任何切分方式下都相同;更宽的 worker 数只是把同一批 tile 重新分组成更多、更小的行块。而 KC,这个唯一决定 tile 累加顺序的分块维度,从不随线程数变化。因此,在固定配置下,仅仅改变 worker 数,也不会改变任何输出元素的累加顺序,这正是 gemmkit 可复现契约背后的机制。并行执行拼装了完整的契约内容,并说明了它确切的适用范围。

并行度在其他方面只影响打包决策本身。LHS 打包门槛看的是每个 worker 的列复用程度,共享 A 预打包也只在大型并行问题上才会启用。这些决策只决定驱动器把打包字节放在哪里、由谁写入,从不改变计算出的数值。任务列表如何切分、契约如何端到端成立,见并行执行;本页提到的这些旋钮,以及其他每一个 GEMMKIT_* 阈值,都收录在调优旋钮中。

打包与工作区

微内核对输入的要求只有一种形状:对每个深度步,它都要 mrA 元素在内存中连续排列,nrB 元素在内存中连续排列,一条面板接一条面板,中间不留任何空隙。用户矩阵几乎从来都不是这个样子。它们有着任意的行、列步长,尾部除不尽微块尺寸,深度方向的遍历甚至可能跨越内存页。

打包(packing)就是弥合这道缝隙的那一次拷贝。它把每个宏块一次性重排成微面板优先(micropanel-major)布局,让最内层循环每次都能从 64 字节对齐的暂存区里读到纯单位步长的数据流,读到完整的 mr/nr 向量。这次拷贝的开销是 O(mc*kc),而复用它的计算开销是 O(mc*kc*nc)。这就是它能摊销掉的原因,也是驱动器在复用程度低到不值得付出这个开销时会直接跳过它的原因。

一个例程,两个操作数

这次机械的拷贝集中在一个例程里:pack_panelsgemmkit/src/pack.rs,L1 层)。LHS 和 RHS 的布局其实是同一种布局,只是从两个不同的方向去看。LHS 宏块打包成若干条 mr 行高的面板,按列逐列存放:面板 0 存放第 0..mr 行,每个深度步的 mr 个元素连续排列;面板 1 存放第 mr..2*mr 行,以此类推。RHS 宏块打包成若干条 nr 列宽的面板,按行逐行存放。

两种布局都是“每个深度步 width 个连续的前导元素“。唯一的区别在于哪条矩阵轴充当前导轴。于是两个 KernelFamily 钩子调用的是同一个例程,只是把步长对调了一下:

#![allow(unused)]
fn main() {
// gemmkit/src/kernel/float.rs
#[inline]
unsafe fn pack_rhs(
    dst: *mut T,
    src: *const T,
    rs: isize,
    cs: isize,
    kc: usize,
    nc: usize,
    nr: usize,
) {
    // RHS panels are `nr` columns wide, stored row-by-row: the "leading"
    // direction is columns (stride `cs`) and the "depth" is rows (stride
    // `rs`), the transpose of the LHS case, handled by swapping strides
    unsafe {
        pack_panels(
            dst, src, /*lead*/ cs, /*depth*/ rs, /*n_lead*/ nc, kc, nr,
        )
    }
}
}

pack_lhs 是它的镜像:lead = rsdepth = cswidth = mr。当块的大小除不尽时,这个例程会把尾部面板里空出来的车道填零。这样内核就总能读到完整的 mr/nr 向量,乘法本身也就不需要对边缘块做任何掩码处理。

例程内部有两条路径,写出的结果逐字节一致。第一种情况是前导维连续(lead == 1),对应列主序的 A 或行主序的 B。这时每个深度步的 live 个元素在源数据里本来就相邻。这条面板于是就是一串直白的 copy_nonoverlapping 调用,加上尾部补零。

第二种情况是前导维带步长。这时朴素的逐元素收集每读一个元素都可能撞上一次缓存未命中(每个深度步要做 width 次跨步加载)。例程转而跑一趟缓存分块转置:它沿源数据的连续维,以 GEMMKIT_PACK_TRANSPOSE_TILE 个深度步(默认 16)为一条,把每一条散布进面板。这样产出的打包字节和一次纯粹的重排拷贝完全一样,但对跨步的数据源要便宜得多。这正是行主序 A 布局的代价并不比列主序高多少的原因。

点积家族(i8 的 VNNI、bf16vdpbf16ps)还有一个姊妹例程 pack_kgroup_panels。它在此基础上,把每条车道连续的 DEPTH_MULTIPLE 个深度步交织在一起,让一条点积指令能一次吞下整组数据。这种布局属于点积内核与深K孪生的内容。

打不打包,由驱动器决定,而这两个操作数并不对称。

微内核以 mr 宽的向量来读取 A,所以只要 A 的行不是单位步长,或者行面板不完整,驱动器就必须打包它。除此之外,当每个 worker 的列复用超过 GEMMKIT_LHS_PACK_THRESHOLD(aarch64 上默认 256 列,其他平台默认 1024 列)时,驱动器也会打包 A

对于列主序的 A,驱动器还会在它的深度遍历同时满足步长达到页级、跨度足够宽、并且被足够多列块复用、值得付出这个成本时打包它。以下三个条件必须同时成立:

  1. 每一步的步长达到半个内存页(GEMMKIT_LHS_PACK_STRIDE,从 Machine 记忆化的页大小自动推导)。
  2. 整条深度切片的遍历(csa * sizeof(Lhs) * kc)达到 GEMMKIT_LHS_PACK_SPAN 字节(自动值:4 MiB)。
  3. 至少有 GEMMKIT_LHS_PACK_REUSEnr 宽的列块复用每一条打包好的面板(min(n, nc) / nr,向上取整,x86 上默认 128,aarch64 上默认 4)。

每一道门槛排除的都是打包不划算的一种情形。一个页级的步长如果发生在仍然驻留缓存的跨度之内,那就只是在重新遍历本来就还热着的缓存行,就地读取 A 的代价反而比付出一次打包更低。跨度这道门槛让 A 保持就地,直到这趟遍历真的宽到足以打垮 TLB,无论后面复用多少次都是如此。复用门槛针对的是另一种失衡:一个瘦高的形状(m 远大于 n)只靠极少的列块就会堆出很大的跨度,把一次昂贵的拷贝摊到太少的复用上并不划算。

复用门槛在不同架构上取值不同,是因为打包和就地读取之间的权衡本身就因架构而异。在 x86 上,打包相对就地读取的代价更高,所以驱动器要等到复用足够多才愿意付出这个代价,默认门槛是 128 条列块。在 aarch64 上,打包相对就地跨步读取的代价更低,驱动器可以更早就选择打包,默认门槛是 4 条列块。

B 则不同,它永远只以广播单个元素的方式被读取,所以任何布局不打包也能用。驱动器打包 B 纯粹是为了复用:每个深度切片打包一次,条件是 m 超过 GEMMKIT_RHS_PACK_THRESHOLD(默认 2048),并且会有足够多的行块反复读取它。由谁来执行这些打包、打包和计算之间的屏障如何安排,属于调度层面的问题,并行执行详细讲述了这部分内容。

预打包操作数

当同一个操作数在一次次调用中反复出现时,每次调用都重新打包就是白费功夫,这正是推理场景的模式:固定的权重反复对上一串流动的激活值。gemmkit/src/api/packed.rs 里的预打包入口会把整个操作数一次性提前打包好。prepack_rhs 沿着任意布局的 B 的步长遍历它,返回一个 PackedRhs<T>gemm_packed_b 随后就用它来做乘法,完全跳过按调用的 RHS 打包。这套 API 的使用方式见预打包操作数。从架构角度看,有三条性质值得关注。

第一,缓冲区记录了它构建时所用的分块几何:nrkcnc。消费它的调用会原封不动地读回这份几何信息。驱动器用记录下来的 kcnc 顶替自己模型算出的结果,只有 mc 仍然按真实的 m 推导。因此,即便打包和消费之间某个调优旋钮发生了变化,面板地址也始终与打包时保持一致。几何本身是通过和普通调用相同的 blocking 模型求解出来的,只是用了一个 tiny_block_dim() + 1 的哨兵行数,让它永远走不到小矩阵分支,因而与最终的 m 无关。

第二,布局只有一个事实来源。prepack_rhs 通过 driver::pack_rhs_full 来填充缓冲区,这个函数铺设面板的顺序,和驱动器自己按片打包时写出的顺序完全一致:最外层是 jc 块,然后是深度切片,再然后是每个切片里 nr 宽的面板。预打包出的字节因此和按调用打包出的字节完全相等,所以在相同配置下,预打包 GEMM 会复现一次普通的 gemm 调用。文档中标注的例外是极小的乘积(mn 都不超过 tiny_block_dim)和 gemv 形状的乘积。普通的 gemm 会把它们改道到特殊路径上,所以它们可能在最后一个 ULP 上有所出入。

第三,缓冲区在整个 GEMM 期间都是只读的,所以每个 worker 都能无需任何同步地共享它。这一点和按调用打包的 B 不同,它不需要任何屏障。

PackedLhs 几乎不用额外的代码,靠的是这套引擎在 A、B 之间的对称性。一个 m x k 的 LHS,本身就是转置乘积 C^T = B^T*A^T 的 RHS。所以 prepack_lhs 只是把步长对调之后,委托给 prepack_rhs_uncheckedgemm_packed_a 也就通过这个转置后的问题来消费它。这种对称性也解释了取向方面的断言:预打包的 B 要求 C 近似列主序(|csc| >= |rsc|),预打包的 A 要求 C 近似行主序。换一种取向的话,分发层就会交换两个操作数的角色,烘焙好的布局就派不上用场了。

int8 feature 增加了一个异构的孪生入口,prepack_rhs_i8gemm_i8_packed_b,有三处刻意做出的不同。

第一,它的布局固定为本进程记忆化的分发所选中的那个整数内核所用的布局:要么是 VNNI 的 k 四元交织布局,要么是加宽内核的普通面板布局。消费入口永远运行同一个家族,所以缓冲区永远不可能被读错。

第二,它把缓冲区深度向上取整到点积内核的 DEPTH_MULTIPLE = 4,并把整个收缩打包成单独一个深度切片,满足驱动器对深度补齐家族的单切片约束。

第三,它刻意绕开了普通 gemm_i8 在低于 GEMMKIT_I8_VNNI_MIN_PAR_MNK 时才会启用的小规模并行加宽回退。vpdpbusd 的缓冲区是四元交织的,加宽内核根本无法消费它。由于整数累加是精确的,无论走哪条路径,结果都与普通的 gemm_i8 逐位一致。

预打包正是在这条路径上收益最大:VNNI 的 RHS 打包本来每次调用都是强制性的,所以在 m 较小时,按调用付出的 O(k*n) 打包开销会压过 O(m*k*n) 的计算开销。

工作区

所有这些打包都需要暂存内存,Workspacegemmkit/src/workspace.rs)就是这份内存的分配器:一块可以增长、64 字节对齐(足以满足 AVX-512 存储的要求)的缓冲区,按 2 的幂增长,并且从不收缩。每次调用中,Workspace::regions 都会把它切分成 a_regions 份大小相等的 LHS 区域,外加一份共享的 RHS 区域,每份区域都向上取整到对齐边界。

LHS 区域的数目,在按 worker 打包的路径上等于 worker 数,在共享 A 的路径上等于行块数,两种路径下切分方式完全一样。当两个操作数都不需要打包时,驱动器会干脆跳过这次预留,让一个完全就地计算的负载永远不会撑大这个池子。

在字节乘积处失败即拒绝

尺寸计算这部分藏着一个内存安全方面的微妙之处。gemmkit 接受广播(零步长)视图:它们只需要一小片后备存储就能通过边界校验,却呈现出逼近 isize::MAX逻辑维度。于是,用来计算打包缓冲区大小的乘积就真的有可能让 usize 溢出,而一旦尺寸回绕成一个偏小的值,就会导致缓冲区分配不足,随后打包时就会越界写入。

驱动器用 checked_mul 守住了元素计数的乘积,但只检查元素计数是不够的。以混合精度路径上的 k = 2^56 为例,那里 kc == k。一个 mc * kc 元素的 LHS 区域,比如说 32 * 2^56 = 2^61,完全能装进 usize,能顺利通过每一项元素级别的检查。可一旦拿这个元素数去乘元素大小、再向上取整到 64 字节对齐,数值就回绕了。溢出恰恰只在元素到字节的换算这一步才会显现,所以守卫也必须设在这里,设在每一份区域大小都要流经的这个咽喉位置:

#![allow(unused)]
fn main() {
// gemmkit/src/workspace.rs
fn region_bytes(elems: usize, esize: usize) -> usize {
    elems
        .checked_mul(esize)
        .and_then(|b| b.checked_next_multiple_of(ALIGN))
        .unwrap_or_else(|| workspace_too_large())
}
}

Workspace 会检查每一步:字节乘积、对齐取整、区域总和,以及最后 A + B 的总量。任何一步溢出,都会以和驱动器自身尺寸检查相同的“too large”契约触发 panic。这就是失败即拒绝:代码会大声拒绝一个荒谬的问题,而不是悄悄败坏内存。驱动器无条件地运行这些元素计数守卫,也是出于同样的理由,即便某条路线最终什么都不打包也不例外:一旦跳过这些检查,就等于同时跳过了中止,会把那个荒谬的 k 送进就地循环里,让它近乎无限期地空转下去。

池子、_withno_std

调用者很少会直接看到一个 Workspace,因为一个线程本地的池子会透明地提供一个。常规的 gemm 调用每个线程最多分配一次,之后每次调用都复用同一块缓冲区。

这个池子同时是可重入安全的。嵌套的 rayon 有可能在一个已经身处某次 GEMM 中的线程上再进入一次 GEMM:比如一个 worker 在自己的 for_each 里阻塞时,窃取了另一次 GEMM 的任务;又或者一个批量并行的 worker 内联地跑了其中一个元素。这种情况下,池子的 RefCell 已经被借出,于是 with_thread_pool 会为这一次调用单独发放一块全新的暂存工作区,而不是让程序 panic。打包缓冲区在调用之间不携带任何结果状态,所以这个回退是完全无感的,唯一被跳过的只是那一次的缓冲区复用。

如果需要显式控制,还有 *_with 这一层:每个入口都有一个变体(gemm_withgemm_packed_b_with 等等),可以传入一个调用者自己持有的 Workspace。从第二次足够大的调用开始,这样做能做到零堆分配,是热循环中大量小乘积、以及延迟敏感代码的合适工具,而 Workspace::with_capacity 甚至能免去首次调用时的分配尖峰。

没有 std 时不存在线程本地存储,所以 with_thread_pool 只会为每次调用简单地新建一个工作区。因为 parallel 本身就依赖 std,这种构建下也就没有线程需要重入。想要复用的调用者可以自己持有一个 Workspace,用 *_with 系列接口,这也是 no_std与WebAssembly 推荐的用法。

并行执行

gemmkit 的并行机制位于一个很小的层里,gemmkit/src/parallel.rs(L2 层),只做两件事。它决定一个问题值得多少个 worker,并把哪部分工作分给每一个 worker。

这两个决定都刻意偏保守,因为线程并不是免费的。这一层的设计出发点是一个观察:worker 数选错,比调度方式选错,损失的性能更多。这两个决定的具体形式也都保证了数值结果不会依赖于其中任何一个。

面向用户的接口只有一个枚举类型:要么是 Parallelism::Serial,要么是 Parallelism::Rayon(n)。默认值 Rayon(0) 表示自动。

感知负载的 worker 数解析

Parallelism::resolve 把请求转换成实际的分区数,而请求只是其中一个输入,负载本身是另一个。

首先是总工作量的串行门槛。当 m*n*k 低于 GEMMKIT_PARALLEL_THRESHOLD(默认 48*48*256)时,一切都保持串行,解析器甚至都不会去采样核心数:为一个只需要几微秒的乘积去 fork rayon,代价会超过乘积本身。这道门槛先于请求本身生效,所以即便显式指定了 Rayon(n),低于门槛时也照样保持串行。

过了这道门槛之后,解析器会尊重一个显式的计数,但会用核心数和可用任务数把它封顶。这样一来,Rayon(huge) 既不会让机器被过度订阅,也不会让按 worker 分配的打包区域超额分配。只有自动路径是启发式的,这让强制指定的宽度对测试和扩展性诊断而言始终精确。

自动路径以工作量为准。它把总工作量 m*n*k 除以 GEMMKIT_PAR_MNK_PER_WORKER(默认 2_000_000,即每 worker 下限,低于它 fork/join 开销就会抵消收益),得到 worker 数,下限为 1,并以核心数和任务数封顶。

之所以按工作量而非线性维度来判断,是因为最优的 worker 数跟随的是总浮点运算量,而不是线性尺寸。一个小规模的立方体乘积串行跑最快,一个中等规模的立方体乘积只需要几个 worker,一个大规模的立方体乘积则需要用满机器上所有的硬件线程。线性维度上的任何单一步幅都无法拟合这样的跨度。

要让中等规模的问题扩展到全宽,还有一件事很关键:避免每个 worker 都冗余地打包同一条 A 面板。打包与工作区里讲述的 LHS 就地读取门,正是用来阻止这种冗余发生的机制。GEMMKIT_PAR_MNK_PER_WORKER 这个旋钮,则是留给那些每 worker 下限和编译期默认值不同的机器的一个逃生口。

带宽受限的形状适用一条完全不同的规则。gemv 或 gevv 每字节只做 O(1) 次算术运算,计算量坡道那一套逻辑在这里不适用。resolve_bandwidth 改为按触碰的字节数设置门槛。

低于一个由缓存推导出来的字节下限时,矩阵能装进单核的私有缓存,该核就能独自把它跑满。这时再拆分,只会增加 fork/join 和共享缓存的争用,却换不来任何带宽收益。gemv_parallel_floor_bytes(在 cache.rs 里)根据拓扑推导这个下限:有 L3 的部件上取单核私有 L2 的大小,没有 L3 的部件上取整个共享簇 L2 的一个比例。GEMMKIT_GEMV_PARALLEL_BYTES 可以直接覆盖这个下限。

超过下限之后,矩阵就会溢出到共享的 L3,其带宽单核无法独自打满。于是自动计数会直接跳到与这些字节数相称的更宽宽度,而不是逐级爬升上去。这个宽度沿着下文所述的精确适配池档位组成一道阶梯往上爬:最小的档位就在下限处,触碰字节数每增加一个 GEMMKIT_GEMV_TIER_STEP(默认自动,8)倍,就上一个档位。这道阶梯止步于最大的档位,而不是整机宽度,因为 gemv 早在机器核心用完之前就已经把带宽打满了,超过这一点之后再加 worker,付出的代价可能比换来的收益还大。

阶梯的每一级都直接取自池档位本身,而不是另立一套比例,这意味着自动选出的 gemv 宽度总能对应一个精确适配的池,永远不必付出这些档位本就是为了消除而存在的那份空闲税。GEMMKIT_GEMV_THREAD_CAP 可以用一个固定宽度取代整道阶梯,供需要把宽度定死的部署使用。

档位之间不设渐进的坡道,因为在带宽受限的扩展曲线上,少量 worker 恰恰是最差的那个点:fork/join 和争用的成本已经付出了,聚合带宽却还没有到手,曲线在这里出现凹陷。任何试图穿过这个凹陷的坡道,最终都会输给凹陷两侧的端点。所以规则保持简单:下限以下串行,下限以上直接取某个档位的宽度,中间没有过渡。

批量 GEMM 有它自己的解析器 resolve_batch,在三种方案中做选择。Serial 在调用线程上依次跑完每一个元素。BatchParallel(n) 把整批、缓存热的 GEMM 整个地交给每个 worker 去跑,整批只需要付出一次 fork/join。因为没有任何一个元素会被拆开,这个方案在任何 worker 数下都逐位一致。SequentialInternal 则反过来,在调用线程上遍历整批,依次把完整的引擎并行度让给每一个体量大、受 DRAM 带宽限制的元素。

resolve_batch 只把 SequentialInternal 这种拆分方式用于 m, n > 1 的形状,这类形状的路由与 worker 数无关。gemv 形状的元素则始终整体留在一个 worker 上。特殊路径讲述了这部分路由,批量GEMM讲述了对应的 API。

按需拉取的工作分发

确定 worker 数之后,驱动器并不会搭建一棵嵌套的任务树。对每个列块和每个深度切片,它都把内部工作压平成一条一维任务列表:n_mc 个行块乘以 n_nt 个列瓦片。任务 q 解码为 (ic_idx, jt) = (q / n_nt, q % n_nt)。各个 worker 从一个共享的、无锁的游标里拉取连续的任务块,直到游标被取空为止:

#![allow(unused)]
fn main() {
// gemmkit/src/parallel.rs
impl JobCursor {
    /// Atomically claim the next `[start, end)` chunk, or `None` once the job space
    /// is exhausted
    #[inline]
    pub(crate) fn next_chunk(&self) -> Option<(usize, usize)> {
        let start = self.next.fetch_add(self.grain, Ordering::Relaxed);
        if start >= self.n_jobs {
            None
        } else {
            Some((start, (start + self.grain).min(self.n_jobs)))
        }
    }
}
}

每次认领只需要一次 fetch_add,没有锁,也没有按任务的队列。按需拉取正是让异构核心配置也能跑得好的关键:在大小核(big.LITTLE)机器上,处理任务更快的 P 核自然会按比例多拉一些任务,而一次静态的 n_jobs / n_threads 均分只会让所有核都去等最慢的那一个。在同构机器上,同样的性质也能吸收系统噪声和频率差异。

任务块的粒度要在两种代价之间取得平衡。粒度太粗,任务列表的尾部就会在汇合时让 worker 闲置。粒度太细,原子认领本身的开销就会显现出来,打包 LHS 的路径上还会在任务块边界处出现重复打包。

通用粒度对 worker 数做了过采样。job_grain 的目标是让每个 worker 拿到 GEMMKIT_PARALLEL_OVERSAMPLE 个任务块(默认 8),这样每个 worker 都预期要拉取好几次,失衡也就能自我纠正。

打包 LHS 的路径是个特例。它天然的任务块就是整个行块(n_nt 个连续任务),这样一个 worker 只需要打包一次该行块的 A 面板,就能在这个行块的所有列瓦片上复用它。这样一来就只有 n_mc 个任务块。所以当行块数较少时,packed_block_grain 会把每个行块切分成 2 的幂个列子块,直到任务块总数达到大约 GEMMKIT_PACKED_OVERSAMPLE * n_threads 个(默认目标为 2)为止,而且只按 n_nt 的因数来切分,这样一个任务块就永远不会跨越行块边界,也就不会在任务块中途重新打包 A。切分得比这个目标更狠,只会导致重复打包过于频繁,反而让性能变差。

每个深度切片里,计算区域之前都会先跑两个并行阶段,它们的边界也是驱动器里仅有的屏障。

B 打包时,各个 worker 从各自的游标里拉取 nr 宽的列面板。这个阶段的 fork/join,正是计算区域所依赖的那个“先写后读”屏障,因为打包好的 B 是唯一一个被所有 worker 非互斥共享的缓冲区。

共享 A 预打包会把每个行块的面板只打包一次,放进一个共享槽位,遵循同样的纪律。它在越过一个尺寸门槛之后开启,也会在 worker 数达到 16 之后无论尺寸大小都开启:在那样的宽度下,每多一个 worker,就多一份对它所触及面板的冗余拷贝,所以即便问题规模中等,去重也已经划算。

除此之外的一切在构造上都是互斥的。各个 worker 只写自己的输出 tile 和自己的打包区域,正是这条不变量,让 Ptr 这个垫片得以在一处经过审计的地方,把捕获到的裸指针声明为 Send + Sync

按规模分级的线程池

rayon 的 fork/join 代价并不随问题规模变化,它随的是池子的空闲余量:池子拥有的线程数,和某次调用实际用到的 worker 数之间的差距。向一个宽度为 w 的池子里 fork k 个 worker,被唤醒的是 w 个线程,而不是 k 个。那 w - k 个分不到活干的线程,仍然要为屏障、以及操作系统层面的唤醒/挂起往返各付一份代价。

对一次小规模的并行 GEMM 来说,全宽的全局池是最差的情形。一个中等规模的乘积往往只需要机器线程数中的一小部分,把它 fork 进全宽的全局池,就会把大部分宽度都白白浪费在每次调用的空闲税上。

gemmkit 的做法是准备一小撮私有的、持久化的池子,每一个都恰好对应自动路径实际会请求的某个 worker 数。在没到整机宽度之前,它最多维护 GEMMKIT_POOL_CLASSES 个减半档位(x86_64 上默认 2,aarch64 上默认 1,封顶 3,其余架构上默认 0):一个半宽档位、一个四分之一宽档位,依此类推,每一级都是对机器物理宽度的再一次减半。每个档位对应的池子都是惰性构建的,只在第一次用到时才建立,只需要一次很小的一次性开销,此后就不再重建。这些档位是机器宽度的固定减半,而不是针对某个具体形状调出来的值。

自动路径会把它选出的 worker 数精确地卡在某个档位上,按构造就没有任何空闲余量。它会一直停留在能覆盖自己的最大档位上,直到总工作量 m*n*k 越过 GEMMKIT_FULL_WIDTH_MNK(默认自动,按架构区分:x86_64 上为 110_000_000,aarch64 上为 14_000_000)。过了这一点之后,那些额外的全宽 worker 才终于能赚回它们带来的 fork/join 代价,全机宽度这才会真正接管。

有三条规则确保这套机制永远不会和调用方自己的调度打架。

  1. 一次已经运行在某个 rayon 池里的调用,无论是调用方自己的 install,还是嵌套在另一次 gemmkit 调用里,都不会被改道到某个档位池。环境池永远优先,就和档位池出现之前完全一样。
  2. 一次显式的 Rayon(n) 会保持它精确的语义,即恰好 n 个 worker。它只会挑一个能容纳 n 的最小档位池去 fork,而不是 fork 进全局池。worker 数本身不受影响,变化的只是它 fork 进哪一个池。
  3. 线程化的 wasm 保留它自己专属的池(见下文),完全不受这一切影响。档位池是原生、非 wasm 场景才有的机制。

wasm 的故事

wasm32-wasip1 上没有线程可以生成,rayon 一旦尝试就会 trap。编译期常量 RAYON_USABLE 记录了目标平台究竟能不能运行 worker。在没有开启线程选项的 wasm 构建上,每个解析器都返回 1,for_each_worker 直接跑普通的串行循环。parallel 就这样优雅降级,而不是直接 trap。

这个选项就是 wasm_threads feature,面向 wasm32-wasip1-threads,或者带 SharedArrayBuffer 的浏览器。由于 wasm 上不支持 available_parallelism,rayon 的全局池原本会悄悄把自己配成单线程。所以 gemmkit 转而自建一个池,大小由 GEMMKIT_WASM_THREADS 旋钮决定(默认 8),并把 worker 循环安装进这个池里执行。部署者只需要声明宽度,其余一切保持不变。wasm 构建的具体细节见 no_std与WebAssembly

拼装可复现契约

gemmkit 的可复现契约很简单:对同一台机器、同一份配置,引擎会产出可复现的结果。这不是对不同配置之间逐位一致的承诺,而 worker 数本身也算作这份配置的一部分。

不仅如此,gemmkit 今天还具备一条更强的性质,但这是一个工程事实,而不是另一条独立的承诺:只改变 worker 数、其余一切保持不变时,输出依然逐位一致。有三个机制共同保证了这一点,下面说明它们如何配合。

第一,数值本身不依赖于 worker 数。分块只从缓存模型推导 kcnc,从不看线程数,深度切片也始终按同一个固定的 pc 顺序运行。因此,无论是一个 worker 清空游标,还是很多个 worker 一起清空游标,每个输出元素的浮点归约都呈现完全相同的形状。worker 数唯一能挪动的分块维度是 mc:宽 worker 数会通过驱动器的并行作业深度下限把它缩小,好让扁平任务列表在每个 worker 上都还能保持几个块的深度。因此,任务列表本身在不同宽度下并不总是逐块完全相同。但 mc 始终保持为 mr 的倍数,所以微 tile 集合(每个 mr 对齐的行偏移,加上唯一那条 m 尾巴)在任何切分方式下都相同。kcpc 顺序都不曾改变,也就没有任何一个结果比特会移动。worker 数改变的是工作如何分组、如何分配,从不改变每一个 tile 究竟算出什么。

第二,任何一次归约都不会被拆分到多个 worker 上。在一个深度切片之内,一个输出 tile 的整次更新由一个 worker 独自完成,所以一个任务块永远是一组完整的 tile。深度切片本身则顺序执行,因为 pc 循环从不并行。beta 只作用于第一个切片,之后的切片只做累加。于是每个输出元素的浮点归约都按一个仅由分块决定的固定顺序进行。特殊路径也遵循同样的纪律:gemv 按寄存器面板边界来切分输出,每一行的 SIMD/标量拆分都不依赖于具体的分区方式,这使得 gemv 在任何 worker 数下都直接逐位一致。批量计划要么把完整的元素留在一个 worker 上,要么只拆分那些路由本就与 worker 数无关的形状。

第三,打包出来的字节不依赖于是谁打包的。pack_panels 是一次纯粹的重排,它的两条分支写出的字节完全相同。因此,无论一条面板是由某个 worker 打包的,还是由共享 A 预打包打的,又或者是更早之前由 prepack_rhs 打的,都是同一串字节。参见打包与工作区。内核完全无法察觉是谁准备好了它的输入。

究竟是哪个 worker 算出了某个 tile,确实每次运行都会不一样,因为游标把任务块交给最先来要的那个人。但由上面这三个机制可知,没有任何数值会依赖于这个选择。这种跨 worker 数的逐位一致,是 gemmkit 今天设计上的一个性质,而不是一份更宽泛的保证。它也因此给像 bf16 点积路径这样按容差保持一致的内核留出了空间,让它们可以改变自己的累加方式而不违背真正的契约。并行实践讲述了实践中应当如何选择 worker 数,以及这份契约对测试意味着什么。

特殊路径

一次 GEMM 调用的生命周期中描述的寄存器分块 driver 建立在一个前提之上:每个输出 tile 要有足够的工作量,才能摊销打包、分块和一整个 MR x NR 寄存器累加器的开销。有些形状会彻底打破这个前提。矩阵向量乘完全没有 tile 复用。k = 4 的乘积在打包成本收回之前就已经算完了。一个 8 x 8 x 100000 的收缩,会让 driver 把大部分精力都花在乘零填充上。

L6 层(gemmkit/src/special/)把这些形状统一改道到专用内核上。每条改道路径都保持三个性质。它藏在同样的公开入口之后,gemm 及其同族函数从不暴露自己走了哪条路。它由调优旋钮按阈值门控,校准失准的门可以不重新编译就移动或关闭。它同时维持库的可复现性契约:同一台机器、同一份配置,给出可复现的结果。下面大多数路径还更进一步,对固定形状而言,在任意 worker 数下都逐位一致。gemv 是唯一的例外,原因见它自己的小节。

各个门位于每个类型分发入口的顶部,顺序固定。下面是 gemmkit/src/dispatch/float.rsrun_typed 的删节版:

#![allow(unused)]
fn main() {
// gemmkit/src/dispatch/float.rs (run_typed, trimmed)
if (t.n == 1 || t.m == 1) && core::cmp::min(t.m, t.n) <= tuning::gemv_threshold() {
    gemv::run_typed_epi::<T, S, Identity>(/* user frame, before orientation */);
    return;
}
orient_transpose(&mut t);
if small_mn_eligible(&t) || small_mn_pack_eligible(&t) {
    small_mn::run_epi::<T, S, Identity>(/* horizontal dot kernel */);
    return;
}
if t.k <= tuning::small_k_threshold() {
    small_k::run::<FloatGemm<T>, S, MR_REG, NR>(/* one depth panel, in place */);
    return;
}
driver::run::<FloatGemm<T>, S, MR_REG, NR>(/* the general blocked driver */);
}

gemv 在方向归一化之前触发,用的是用户自己的坐标系。其余的门都作用在已经归一化的问题上。每条特殊路径也都有对应的融合 epilogue 形态,因此一次 gemm_fused 调用走的路线,与它对应的非融合调用完全相同。这条契约是 Epilogue 融合一章的主题。

gemv:受内存带宽支配的边界

m == 1n == 1 的形状(gemmkit/src/special/gemv.rs)每个输出元素只做 2k 次浮点运算,却要读一遍 k 个矩阵元素。这使它成为内存受限的操作。整个设计问题是如何压低 DRAM 流量,而不是如何调度 FMA。两种取向都归约到同一个核心例程:把矩阵(m == 1 时先转置)看作一个 rows x k 的块乘以一个 k 向量。

这个门看的是形状,不是大小。对任何 gemv 形状,min(m, n) 都是 1,所以与之比较的 GEMMKIT_GEMV_THRESHOLD 旋钮实际上起的是开关作用,而不是尺寸上限的作用。把它设为 0,就能把 gemv 形状强制送回通用 driver,结果依然正确。

并行度依据的是带宽模型,而不是计算量的爬升曲线。Parallelism::resolve_bandwidth 在一个由缓存推导出的字节下限以下保持串行。低于这个下限时,矩阵能装进单核的私有缓存,由这一个核独自打满即可。越过下限后,它会直接跳到与所触及字节数相称的宽度。这个宽度沿着精确适配的池子档位往上走,顶端封在逻辑核数的一半,因为只用少数几个 worker,恰恰是带宽扩展曲线上最差的一点。

输出行按 SIMD 宽度的整数倍为粒度切成 panel,分给各个 worker。每一行完整的 k 归约都留在同一个 worker 内部完成,没有任何 worker 会去合并另一个 worker 的部分结果。这让 gemv 在固定的 worker 数下保持可复现,这也是引擎其余部分共同遵守的底线。把这些行切给不同数目的 worker,本身并不承诺在这些不同 worker 数之间逐位一致,这一点和下文的 small-k、small-mn、batched 都不一样。

在一个 worker 自己的行区间内部,代码依据布局在四种策略中选择一种。列主序矩阵走 axpy 形态,它提供两个变体,两者刻意做到逐位一致,只在内存流量上有差别。

寄存器分块输出形态把一个输出行 panel 保持在 SIMD 寄存器里,扫完整个 k,因此矩阵和输出各自只读一遍。普通列外层形态每隔几列就重读一次输出,但矩阵是作为一条连续的流来读取的。二者之间的取舍依据是输出的缓存驻留性,由 output_register_block 计算得出。当输出(rows * sizeof)超出末级缓存的某个比例时(此时普通形态的重读本会落到 DRAM 上),并且 k 不超过 GEMMKIT_K_STREAM_MAX(默认 32)时,这条路径会选择寄存器分块。超过这个 k 之后,寄存器分块形态里那许多条并发的列流会开始挤爆硬件预取器。两个变体对每个元素执行同样的升序 k 融合累加,也用同样的按行 SIMD/标量划分。在两者之间切换,绝不会改变输出的任何一位,只会改变速度。

行主序矩阵改走 dot 形态。它把行按 4 个一组做寄存器分块,以重叠 FMA 的延迟链,而每一行仍然执行同一套共享的、顺序固定的 dot_contiguous 归约。完全跨步的操作数则退回标量循环。

有一种形状同时满足上述两种分类,而这个平局倒向哪一边,决定了这次计算究竟能不能向量化。axpy 形态沿输出行向量化,一次把 lanes 行放进一个寄存器。dot 形态则沿 k 向量化。只有一行的矩阵,其行跨度与列跨度恰好都是 1,所以把它描述成列主序或行主序同样准确,这正是纯粹的点积,m == n == 1

把这种形状交给 axpy 形态,它的向量循环(while i + lanes <= e,而 e == 1)根本进不去,整个归约就会落到标量尾巴上。axpy_yields_to_dot 避免了这种情况:只要行数不足一个 SIMD 寄存器、且 dot 形态自身的跨度条件也成立,它就把这趟扫描让给 dot 形态。dot 形态的累加树更宽,精度也更好。这个选择在实践中很重要:列主序的适配器库,例如 nalgebra 和 faer,正是用这种方式描述一个行向量,所以调用者随手写一个点积,原本会默认落在慢的那一类上。

行是否要切给多个 worker,是与由哪种策略计算它们相互独立的另一个决定。对列主序矩阵而言,通常的答案是不切。输出行正是这种矩阵内层、变化最快的内存轴。切开它,就等于让每个 worker 都要跨步走遍整个矩阵,却只消费每一列里属于自己的那一小段。串行路径反而只需一趟顺序扫描:row_sweep 会直接短路成一次 body(0, rows) 调用,完全不分块。这一趟扫描本就已经跑在单流可达速率附近,所以在某个行数以下,额外的 worker 能赢得的很少,要牺牲的顺序性却很多。

GEMMKIT_GEMV_AXPY_PAR_MIN_ROWS 就是那道行数下限。低于它时,无论请求了多少 worker,axpy 的切分都保持串行。有两条路径被特意排除在这道下限之外。行主序矩阵让每个 worker 拿到整段沿 k 连续的行,所以即使切分,它的数据流依然顺序,切分在任何规模下都值得。混合精度孪生的加宽 axpy 计算量足够大,足以在同一条列主序数据流上扩展,所以切分同样值得。

混合精度孪生 run_mixed(feature half)服务于 f16/bf16 的 gemv。它使用与上面浮点例程相同的行划分,也遵循相同的可复现性论证。每次加载都经由 KernelSimd<N, N, f32, N> 接缝拓宽到 f32,归约在 f32 中进行,结果只在存储那一刻舍入回窄类型一次。

这条单次舍入规则带来一处不对称:混合 axpy 总是采用寄存器分块输出。普通的列外层形态会按列组重读重写窄输出,这会变成每组舍入一次,而不是每个元素舍入一次。

混合的融合 gemv 干脆完全不走这条路,这是刻意的。浮点融合 gemv 的融合方式,是把已经存储的输出重读回来、就地映射,而这只有在浮点输出本身就是累加器时才逐位精确。窄类型的输出在存储时已经舍入过一次,再读回来映射一次就等于舍入了两次。混合融合入口没有把 epilogue 穿进每一次拓宽存储,而是干脆让 gemv 形状留在通用 driver 上,driver 本来就会在唯一一次收窄之前,以 f32 应用 epilogue(gemmkit/src/dispatch/mixed.rs)。

small-k:单个深度 panel,无可摊销

在很小的 k 下(gemmkit/src/special/small_k.rs),整个乘积就是一个深度 panel。driver 的缓存分块模型、工作区划分,尤其是 A/B 打包,全都会是纯粹的开销,因为每个打包元素只会被读一次。这条路径改为直接在 family 的 microkernel 上计算 C <- alpha*A*B + beta*Ckc = k。它就地读取 A 和 B:不打包、不分块、不产生工作区流量。因为它对 KernelFamily 保持泛型,还能免费继承 family 的拓宽和舍入语义。

门是 k <= GEMMKIT_SMALL_K_THRESHOLD。它的默认值按架构区分:x86 上是 16,aarch64 上是 8。更窄的 NEON microkernel tile 打包成本足够低,driver 在那里更早胜出,这就是默认值更低的原因。

就地读取需要三个前置条件。只要有一个不满足,这条路径就退回到 driver,结果依然正确,只是调度方式不同。第一,microkernel 要求 LHS 行单位步长,所以 A 必须是列主序(rsa == 1)。第二,像复数这样的 FORCE_PACK_* family,打包时会把数据变换成平面布局,从构造上就无法就地读取。第三,k 超过硬上限 SMALL_K_MAX = 32 会溢出这条路径唯一使用的那块栈缓冲区:一个为底部不完整行 tile 准备的零填充 panel。这个 panel 仍然需要打包,因为 microkernel 总是加载满 mr 行。

工作按输出 tile 划分,每个 tile 由一个 worker 完成完整的一遍 k。worker 数本身由带宽模型封顶,因为在小 k 下,m*n 的输出写入占主导。因为每个 tile 都是由单个 worker 完成的一次完整归约,所以结果在任意 worker 数下都保持逐位一致,这比 gemv 所遵守的性质更强。

small-mn:小输出的水平点积

mn 都远小于 microtile,而 k 很长时(gemmkit/src/special/small_mn.rs),driver 会把微小的行列 tile 填充成完整的 MR x NR microtile,算的大多是填充。这条路径改为把每个输出元素算成一次水平 SIMD 点积,C[i,j] = alpha*<A[i,:], B[:,j]> + beta*C[i,j],沿收缩方向流式计算。输出按 4 x 4 的累加器 tile 做寄存器分块,16 条独立的 FMA 链在整个 k 扫描中保持飞行状态,每个 A 行和 B 列每个 tile 只加载一次。这与 gemv 的 dot 形态是同一个延迟隐藏技巧,只是推广到了一个小网格上。

维度门是 m, n <= GEMMKIT_SMALL_MN_DIM(默认 16,aarch64 上为 32),再加上 k 高于 small-k 阈值。这个上限按架构区分,是因为 driver 的填充开销超过这条路径水平点积开销的那个临界点,因机器而异。两条小形状路径就是这样瓜分 k 轴的。

内核要求两个操作数都沿 k 单位步长:它需要 A 的行连续(csa == 1,行主序 A),B 的列连续(rsb == 1,列主序 B)。这就是零拷贝档。最常见的两种布局各自恰好缺一边:全行主序缺 B,全列主序缺 A。针对这两种情况,第二个与之互斥的门(k > GEMMKIT_SMALL_MN_PACK_MIN_K,默认 16)会启用一个预打包档:prepack_operands 只把不合格的那个操作数拷贝进一段 k 连续的工作区暂存缓冲,然后让同一个内核以单位步长在其上运行。

这次拷贝触及 m*k(或 n*k)个元素,相对于点积本身 m*n*k 的工作量而言,是一笔很小的开销,水平内核的收益远超于此,所以跨步布局的小 m,n 形状仍然胜过退回 driver 的填充 microtile。暂存缓冲把每条线的步长向上取整到奇数个缓存行(packed_line_stride)。若步长恰好取 k,一旦 k 是 2 的幂,就会让所有打包行落进同一个 L1 组,重读的收益也随之崩塌。取奇数行数正是为了避免这一点。

这笔开销在浮点运算量上确实很小,但浮点运算量对它是错误的度量:这次拷贝每搬一个字节做零次算术,而点积每字节大约做两次,所以拷贝能用来掩盖内存延迟的东西严格更少;在长 k 上它占用的时间份额远大于它占用的工作量份额。留在调用线程上做时,它反而支配了整条路线。因此拷贝本身也交给多个 worker 去做——前提是它的流量越过了带宽受限路径共用的那个由缓存推导的字节下限(GEMMKIT_GEMV_PARALLEL_BYTES;低于它则拷贝仍然串行)。它与随后的 tile 扫描分别决定宽度,因为两者能提供的并行度完全不同:tile 扫描受 MT x NT 输出网格封顶,而小 m, n 让这个网格极小;拷贝则受深度本身封顶,而长 k 让它很大。

切分沿深度进行,绝不沿 lead:连续的 t 区间让每个 worker 每步读入整条深度线、即 lead 个连续元素;若改为切分那寥寥几条 lead 线,则每个 worker 每步只能从每条线里挑出一个元素,而且要横跨整个操作数。Zen5 参考机实测(f32,自动宽度,列主序 A):8x8x524288 3.1×、16x16x262144 2.0×、4x4x1048576 1.8×、8x8x2097152 1.7×、16x16x1048576 1.1×。其中唯一贯穿始终的规律是占用空间——固定 m,n 时,被打包操作数越小收益越大;而固定占用空间、改变 m,n 时次序并不单调,所以这里除了“拷贝在串行时间中的占比“之外不主张任何单一机理。

预打包步骤是一次纯粹的重排:同样的值,同样的每行顺序。因此打包路线与本就合格的布局路线保持逐位一致;而且每个格子都恰好被写一次、写入的正是串行拷贝会写的值,所以把拷贝拆给多个 worker 同样不会移动任何一位。混合精度(f16/bf16,拓宽到 f32,每格恰好舍入一次)与整数(i8 -> i32,回绕运算,因此与 driver 逐位精确)两个同胞,共享同一套 tiling、同一个预打包助手,以及同一套可复现性论证。

batched:编排层,不是内核

批量 GEMM(gemmkit/src/special/batched.rs)刻意不实现新内核。每个批元素都重新经过完整的单 GEMM 引擎分发,所以批量调用自动能与 driver、gemv、small-k、small-mn 组合起来。这一层添加的是一份调度方案,由 Parallelism::resolve_batch 每次调用时选定一次:

  • BatchParallel:当元素数量足够填满所有 worker 时选用。每个 worker 串行、缓存热地跑完整个 GEMM,整个批只需付一次 fork/join,而不是每个元素一次。这正是“大量小矩阵“这一动机工作负载所对应的模型。
  • SequentialInternal:为数量少、体积大、受 DRAM 带宽限制的元素而选用。它在单线程上循环整个批,轮流把引擎的全部并行度交给每个元素。在 x86 上,一旦某个元素溢出单核的 L2,这个方案就会启用。在 aarch64 上,由于簇内共享 L2、且统一内存带宽较高,一旦每个批 worker 分到的份额 elem_bytes / batch 超过 GEMMKIT_SEQ_INTERNAL_BYTES_PER_WORKER(默认 128 KiB),它就会启用。这个方案会把单个元素自身的工作拆到多个 worker 上,因此只对 m, n > 1 的形状开放。driver、small-k、small-mn 这几条路线都是在单个 worker 内部归约每个输出,所以在当前这套与线程数无关的分块下,串行和并行结果逐位一致。gemv 只被要求满足基本的可复现性承诺,而不是这种逐位一致,因此这个方案把它排除在外。
  • Serial:当总工作量低于门限,或没有可用线程时选用。

由于每个元素都相互独立,串行方案和批级并行方案都不会拆分任何一个元素,所以在这两种方案下,批量结果在任意 worker 数下都逐位一致,与每个元素内部走的是哪条路线无关。带步长的批量入口 gemm_batchedgemm_batched_fused(后者把同一份共享的 epilogue 穿过同一套骨架)共用一份调度实现,普通形态和融合形态因此不会出现分歧。

对于元素形状各不相同的批,指针数组形式 gemm_batched_ptr_unchecked 改为接受一个 GemmProblem 描述符切片,每个描述符携带自己的维度、步长和指针。这种形式采用更简单的 resolve_batch_flat 策略:把整个 GEMM 分给 worker,绝不做元素内部的拆分,因为这里没有统一的驻留性可供判断。gemm_batched_slice 是它经过校验的安全孪生版本。

批量 GEMM小形状与 GEMV介绍了这一切面向用户的视角。并行执行描述了这些路径所依赖的 worker 数机制。

Epilogue融合

GEMM 的输出很少以原始形态离开例程。推理层要加偏置和激活,量化流水线要把 i32 累加器重量化成一个字节。朴素地做,这些操作每一个都是对 C 的第二次完整遍历:每个元素都要写入内存、被逐出、读回、变换、再写一次。

epilogue feature(gemmkit/src/kernel/epilogue.rs)改为把这个变换直接融合进 microkernel 的存储步骤。它在元素本就占据的寄存器(或暂存槽位)里完成变换,时机就是 microkernel 本来要存储它的那一刻,第二次遍历因此完全消失。对重量化而言,省下的开销还要更大:非融合流程必须先把整个 m x n 矩阵以 i32 物化出来,再收窄它。

接缝

这条接缝是 Epilogue trait。它经由 KernelFamily::microkernel_epi 穿入,让每个 family 的存储点都能应用它,而 driver 完全不需要知道它的存在:

#![allow(unused)]
fn main() {
// gemmkit/src/kernel/epilogue.rs (trimmed)
pub trait Epilogue<Fam: KernelFamily>: Copy + Send + Sync {
    /// true => every hook const-folds away; the kernel is bit-identical to non-fused
    const IS_IDENTITY: bool = false;
    /// true => apply_reg is implemented, enabling the fast vector store path
    const VECTOR: bool = false;
    /// true => apply_store is implemented (the Out != Acc requantize pattern)
    const VECTOR_STORE: bool = false;

    /// Scalar transform at absolute (row, col) in the oriented problem frame
    unsafe fn apply(&self, v: Fam::Acc, row: usize, col: usize) -> Fam::Out;
    /// Vector transform of LANES consecutive rows; MUST agree with apply bit-for-bit
    unsafe fn apply_reg<S>(&self, simd: S, v: ..., row: usize, col: usize) -> ...;
    /// Vector transform of a whole MR_REG x NR register tile; defaults to a loop over
    /// apply_reg. Overridden to hoist a runtime discriminant out of the unrolled pass
    unsafe fn apply_tile<S, const MR_REG: usize, const NR: usize>(&self, simd: S, acc: ..., row0: usize, col0: usize) -> ...;
    /// Vector store-transform from Acc scratch to Out; same bit-agreement contract
    unsafe fn apply_store<S>(&self, simd: S, src: *const Fam::Acc, dst: *mut Fam::Out, ...);
}
}

整个设计靠两条不变量撑起来。第一条是零开销恒等。普通 gemm 传入的是 Identity epilogue,它的 IS_IDENTITY = true 让每个钩子都在编译期折叠掉。单态化出来的非融合内核,和这条接缝存在之前完全逐位一致。调用者不用融合时,融合不花一分钱。

第二条是恰好触发一次的语义。driver 把一个 last_k 标志交给 microkernel,epilogue 只在最后一个深度 panel 上应用。更早的 panel 存储的是原始的 Acc 部分和,和非融合内核完全一样。OUT_IS_ACC = false 的 family(比如窄类型 f16/bf16 的输出)从构造上就以单个 kc = k panel 跑完整个收缩,所以 last_k 在那里天然为真。深 K 孪生会破坏这个单 panel 的保证,因此它在融合路径上刻意从不启用。特殊路径也天然只触发一次,因为它们的每个输出元素本身就是一次完整归约加一次存储。

内置 epilogue

随库发布三个 epilogue,各自有自己的公开入口(feature epilogue)。重量化入口还额外需要 int8。面向用户的视角见融合 Epilogue

FusedEpi 是运行期组合出来的“偏置加激活“epilogue。它先加上按行或按列的偏置(Bias::PerRow / Bias::PerCol),再应用 ReluLeakyRelu(slope)。一次单态化就覆盖了所有组合,所以融合内核的数量不会随 epilogue 种类的数目相乘。但这一点成立,仅仅是因为枚举分支每个 tile 只解码一次,在 apply_tile 的覆写里完成,而不是每个累加器解码一次。

这个区别不是什么微优化。tile 的 const 泛型把内核的存储遍展开。如果用逐寄存器的钩子,两个 match 都会在 tile 的每一个累加器槽位上被复制一遍。在一个宽 tile 上,由此形成的分支网会让编译器付出整个累加器 tile 的代价:它不再把 acc 留在寄存器里,而是从 kc 循环内部就把每个值写穿到栈上,而 epilogue 在那里根本不会执行。把解码提到循环外面、改成每个 tile 只做一次,就能完全避免这次溢出。tests/perf/fused.rs 钉住了融合速率与普通速率之间的比值,这样以后的改动就不会在无人察觉的情况下让这个问题重新出现。

FusedEpi 支撑着 gemm_fused 以及它的整个家族。这包括在整批之间共享同一份偏置和激活的 gemm_batched_fused、预打包孪生函数 gemm_packed_b_fusedgemm_packed_a_fused,以及复数入口 gemm_cplx_fused。复数入口只有偏置,因为基于大小比较的激活函数在复数上没有数学定义。FusedEpiVECTOR 设为 true。在快路径上,偏置加法和激活都以寄存器操作的形式运行,比如 max(v, 0)。它在 SIMD max/min 上的 NaN 约定经过精心选择,使向量形态与标量形态严格一致:两者都算出 ReLU(NaN) = 0

MapEpi 是逃生舱口。gemm_map 把一个任意的用户闭包 f(value, row, col) -> value 应用到每个输出元素的最终值上,(row, col) 使用用户坐标系。这个闭包是借用的 &dyn Fn + Sync,所以每个 (类型, ISA) 只需一次单态化,而不是每个闭包一次。它以标量方式运行,每个元素调用一次,由每个元素背后 O(k) 次浮点运算摊销这个开销。MapEpi 只支持 f32/f64。窄类型将不得不先舍入到 N、应用 N 域上的闭包、再舍入一次,这会破坏下文描述的逐位契约。

KRequantize 实现了量化推理的存储步骤:C[r,c] = clamp(zp + round_ne(scale*(acc + bias)), LO, HI)。它把 i32 累加器映射到 i8gemm_i8_requant,值域 [-128, 127])或 u8gemm_i8_requant_u8,值域 [0, 255],即 ONNX QLinearMatMul 的惯例)。scale 可以按张量取,也可以按行取(RequantScale)。zero point 在舍入步骤之后以整数形式并入。可选的 i32 偏置在唯一一次 f64 舍入步骤之前以整数形式并入。舍入本身是 round-half-to-even,通过一个 no_std 安全的 2^52 技巧实现(round_ne_f64)。KRequantize 没有 alpha,因为它已经折进了 scale。它也没有 beta,因为在一个已经量化过的 C 上继续累加没有清晰的含义。

正确性契约

这份契约之所以写得如此精确,正是因为 epilogue 的测试逐位钉住的就是它。三个要素合成了它:

  1. 相同的路由。 一次融合调用把每个形状都送进普通 gemm 会用的同一个内核。通用 driver、gemv、small-k、small-mn 各自都有对应的融合形态(见特殊路径),融合分发入口与普通的门一一镜像。没有哪个形状会仅仅因为要了个偏置就得多付 driver 的开销。唯一刻意的例外是混合 f16/bf16 的融合 gemv,出于特殊路径一章解释的舍入原因,它留在 driver 上。反正窄类型本来就在下文的逐位契约之外。
  2. 与 epilogue 无关的引擎。 分块、调度、打包和累加顺序都不依赖穿入的是哪个 epilogue,epilogue 只触碰存储这一步。
  3. 逐位一致的两条应用路径。 完整的列主序 tile 走向量路径存储,即 apply_tile(默认实现是 apply_reg),或者走 apply_store。边缘或跨步的 tile 则改为经暂存走标量 apply。同一个输出矩阵可以自由混用这两条路径。所以 trait 契约要求两者在同一个 token 下逐位一致。apply_tile 的覆写继承了这条义务,它必须逐元素留下与 apply_reg 完全相同的结果。

这三点合起来给出了头条保证:对 f32/f64 而言,gemm_fusedgemm_map,以及批量和预打包的融合入口,都等于 gemm() 后接同一个标量映射,逐位相等,对每一种形状都成立

MapEpi 展示了这份保证有多刻意。它把 VECTOR 设为 true,不是为了向量化闭包(做不到),而是为了让内核走上与普通 gemm 完全相同的路径选择。快路径融合的 beta*C + alpha*AB 存储,在一般的 beta 下,与标量路径的非融合算式相差 1 个 ULP。因此只走暂存的 epilogue,会把普通 gemm 实际上从未写出过的值交给闭包。apply_reg 转而把寄存器排空到一个栈缓冲区,再逐 lane 调用同一个标量 apply,这样 f 看到的永远是普通 gemm 产生的那些精确位。

文档化的例外是 f16/bf16。窄类型的 blanket 实现在唯一一次 round-to-nearest-even 收窄之前,就在 f32 累加器上应用了偏置和激活。这刻意比“gemm() 再映射“更精确,因为后者要先舍入到窄类型、再拓宽回来、再舍入一次。所以对窄类型而言,融合入口并不与“gemm 再映射“逐位相等,文档直接说明这一点,而不是削弱融合语义去迁就那个精度更差的替代方案。在同一次融合运行内部,向量路径和标量路径仍然逐位一致,因为两者都在 f32 中计算 act(bias(v)),并且都只舍入一次。跨 worker 数的可复现性也保持不变。

KRequantize 的向量路径是以另一种方式获得资格的。x86 token 实现了 KernelSimd::requant_store:向量化的拓宽到 f64、乘 scale、硬件 round-to-nearest-even、clamp,再写出低字节。它的文档带着逐情形的证明,说明每个 lane 都等于标量的 clamp(zp + round_ne(scale*v), lo, hi)i32 -> f64f32 -> f64 的拓宽都是精确的,2^52 技巧在 2^52 以下与硬件舍入一致,饱和行为在其上也一致,而 NaN 不可能出现,因为 API 已经校验过每个 scale 都是有限且为正的。按行取值的 scale 会逐 lane 变化,这种情况改走逐 lane 的标量映射。非 x86 token 保持 REQUANT_VECTOR = false,全程使用标量映射。一个模块内的一致性扫描,也就是 gemmkit/src/simd.rs 里的 requant_store 测试,会在每个具备这项能力的 token 上检查这种逐位相等。契约里说的“已证明“,因此是靠测试强制执行的,而不只是一个愿景。

最后还有一个角落。当 A*B 项消失时,也就是 k == 0alpha == 0,融合入口仍然欠着一笔 C <- act(beta*C + bias)。这个退化映射在用户坐标系里逐元素运行。gemmkit/src/dispatch/float.rs 里的 fused_degenerate 处理这种情况,它还有一个窄类型的同胞版本,在 f32 中合成结果、只收窄一次。因此,即便是乘积项为空的情形,也遵守与完整内核相同的语义。

扩展点

gemmkit 的变化点是 trait、const 泛型,以及存放在 OnceLock 槽位里的带类型函数指针,没有宏,也没有 transmute。这套纪律只为了一个目的:库预期会有四类增长,分别是新指令集、新元素类型、新点积指令、新融合变换。每一类都应该以纯增量代码落地,只带来一份简短、可核对的触点清单,并且都不应该改动 driver、打包例程和分块模型本身。

本页把这四份配方展开成实操走查,面向想要扩展这个 crate 本身的人。这些接缝的公开程度足够高,其中最关键的一条,也就是用你自己的内核 family 去驱动泛型 driver,在 crate 之外也同样成立,并且有测试证明这一点。

新的 ISA 后端

一个 ISA 后端是一个零大小的 token,加上一套词汇表实现。wasm simd128 后端(gemmkit/src/simd/wasm.rs)是最近的一个完整范例,值得从头到尾读一遍,因为它只有一个文件加上几行分发代码。

这个 token 唯一的固有行为是 Simd::vectorize,也就是 #[target_feature] 蹦床。运行期 CPU 检测无法与泛型内核上固定的 #[target_feature] 属性搭配使用,所以每次内核调用都跑在一个带注解的小函数里。#[inline(always)] 的原语会折叠进这个函数,让每个 intrinsic 都落在特性已启用的代码生成上下文里:

#![allow(unused)]
fn main() {
// gemmkit/src/simd/wasm.rs
impl Simd for Simd128 {
    #[inline(always)]
    unsafe fn vectorize<R>(self, f: impl FnOnce() -> R) -> R {
        #[target_feature(enable = "simd128")]
        fn inner<R>(f: impl FnOnce() -> R) -> R {
            f()
        }
        inner(f)
    }
}
}

清单如下:

  1. token。 在新的 gemmkit/src/simd/ 模块里加一个 Copy + Send + Sync + 'static 的零大小结构体,按架构做 cfg 门控,并给它加上前面展示的 vectorize 蹦床。

  2. SimdOps<T> 实现。 为该 ISA 加速的每个元素类型都加一份实现,各自需要一个寄存器类型、LANES,以及一套原语词汇:load、store、splat、mul、add、mul_addfnmareduce_sum,如果希望融合浮点 epilogue 能够向量化,还要加上 max/min。这套词汇表刻意做得很“厚“,这样 microkernel 才能始终保持为一个泛型函数。你实现的是原语,而不是内核本身。

    在这里一定要遵守文档化的契约。simd128 的实现用的是 f32x4_pmax,而不是 f32x4_max,因为 trait 里的 max 要求 a 为 NaN 时返回 b,这正是向量与标量 epilogue 之间“ReLU(NaN) = 0“的约定。它还把两个操作数反过来传,写成 f32x4_pmax(b, a),因为 pmax(x, y) 计算的是 x < y ? y : x。如果按自然顺序传参,a 为 NaN 时会返回 NaN,max(-0.0, +0.0) 会返回 -0.0,这两种情况都恰好与契约相反。它还把 mul_add 写成未融合的 mul 后接 add,因为 wasm 没有硬件 FMA,而 relaxed-SIMD 提供的替代方案在规范上是不确定的,会破坏可复现性。

  3. tile 几何形状。 为每个类型选定 (MR_REG, NR),并把它编码成分发模块里各 ISA 包装函数的 const 泛型。这是唯一一个按 (类型, ISA) 设置的旋钮。要明确地做寄存器预算:simd128 对 f32 用 2x4 的布局,也就是 8 个累加器、2 个 LHS 寄存器、1 个 RHS 寄存器,共 11 个活跃的 v128 值,因为 LLVM 的 wasm 后端在活跃向量数超过约 16 个之后就会开始溢出。NEON 则用 4x4 的布局,刻意留出富余的寄存器。

  4. 一个 Dispatched 描述符,以及每条 select_* 阶梯上的一条分支。 记忆化的选择阶梯位于 gemmkit/src/dispatch/ 下:浮点用 select_f32/select_f64,混合精度用 select_f16/select_bf16,整数用 select_i8,复数用 select_c32/select_c64,再加上 map-epilogue 的选择器。每条阶梯分支都把普通、预打包、融合三种入口点和 tile 几何形状捆在一起,所以新增一个 ISA,只需要一个描述符常量,加上每个受益类型一条 match 分支。

  5. 一个 GEMMKIT_REQUIRE_ISA 取值。gemmkit/src/dispatch/isa.rs 里加一个 ForcedIsa 变体及其解析字符串。目前的取值有 scalarfmaavx512favx512vnniavx512bf16neonsimd128auto。遵守“响亮失败“的规则:如果被钉住的 ISA 不受支持,分发必须直接 panic,而不是回退。这样一来,一个原本想测试你这个内核的 CI 任务,就不可能悄悄跑到别的内核上却通过了测试。

  6. 测试基本上是免费搭车的。 tests/simd_conformance.rs 直接构造各个 token,把每个原语拿去和标量参考实现逐一核对。再加上一个 env_isa_* 钉住二进制和一个 CI 任务,分发路线本身也就可测了(见测试与验证)。

你应该完全不需要碰 driver.rs、任何内核 family、pack.rscache.rs。simd128 后端一个都没改。

新的元素类型

元素类型沿着两个小 trait 变化(见标量与内核家族)。Scalargemmkit/src/scalar.rs)只声明恒等常量和累加器类型 Acc。选择 Acc 是这里影响最深远的一个决定,因为它决定了整套舍入方式。f16 选择了 Acc = f32i8 选择了 Acc = i32,这让整数 GEMM 严格精确。KernelFamilygemmkit/src/kernel.rs)则捆起了区分一种运算所需的其余一切:Lhs/Rhs/Acc/Out 类型、打包布局,以及 microkernel。

很多时候根本不需要新的 family。如果新类型只是既有累加器之上的一种窄输入,那就改为在有能力的 token 上实现 KernelSimd<L, R, A, O> 这条拓宽/收窄接缝:拓宽加载,再加一次收窄存储,然后复用泛型 microkernel,就像 MixedGemm<f16>MixedGemm<bf16> 那样。同质情形由一个 blanket 实现覆盖,混合实现不可能与它重叠。真正全新的运算形态,比如平面布局的复数内核,或者重量化的整数 family,才需要拥有自己的 KernelFamily

把新类型接入公开 API,意味着在 gemmkit/src/dispatch/ 下新增一个分发模块,并为该类型准备自己的 OnceLock 槽位。特性检测只跑一次,胜出的单态化入口点会被缓存下来,之后每次调用都只是一次间接调用。如果是两个类型、gemv/small-mn/small-k 的改道逻辑,再加上一点点点积内核选择上的细节,可以照抄 dispatch/mixed.rs 的模式。如果是异质的任务类型,则改为照抄 dispatch/int.rs

这里的开闭性质不是口口相传的说法,而是被 gemmkit/tests/open_closed.rs 直接强制执行的。那个测试定义了 NaiveFloat:一个独立编写、自带打包逻辑和朴素标量 microkernel、只使用公开条目构建的 family。测试用它去驱动完全未改动的公开函数 driver::run,再对照 f64 参考结果做检查。一旦 driver 的改动破坏了 family 这道接缝,这个测试就会连编译都通不过。它同时也是新建一个 family 时可以照着写的模板。

点积指令

vpdpbusdvdpbf16ps 这类指令把好几个深度步骤折进一条操作里,这会重塑累加的舍入方式,所以它们绝不能以“对可移植 tile 循环的巧妙覆写“这种形式出现。为了保持这个区分,这条接缝被拆成了几部分:

  • family 声明 DEPTH_MULTIPLE = Q(大于 1),并经由 pack_kgroup_panelsgemmkit/src/pack.rs)打包。这个函数把 Q 个连续的深度步骤按 lane 交织排列在一起。driver 会把 panel 的深度向上取整到 Q 的整数倍,并保证 k 组不会跨越切片边界。
  • 有能力的 token 覆写 KernelSimd::dot_accumulate,整组整组地消费这些 panel 里的指令组。打包出来的布局是 family 的打包器与覆写它的 token 之间的私有契约。任何符号修正,比如 VNNI 那个带列和补偿的 +128 技巧,都封装在覆写内部完成,这样累加器返回时持有的就是真实的和。
  • SimdOps::accumulate_tile 的覆写只保留给那些调度层面、并且不改变舍入形状的改动,比如需要显式软件流水线的顺序执行核心,或者长度不是编译期常量的可伸缩向量 ISA。它的文档说得很明确:会重塑舍入的指令不在这条接缝的适用范围内,那类指令应该改用带点积接缝的新 family。accumulate_tile 的覆写必须保持确定性,并且要和边缘路径的舍入方式一致。默认实现已经能在任何宽乱序核心上打满 FMA 管线,所以在保留一个覆写之前,先证明它确实值得。

IntGemmVnniBf16DotGemm 是两个现成的例子。IntGemmVnni 对拓宽路径逐位精确,因为整数算术满足结合律。Bf16DotGemm 则改为按容差把关,仍然落在可复现性契约之内。点积内核与深 K 孪生对两者都有深入讨论。

新的融合变换

一个融合变换就是一份 Epilogue 实现(gemmkit/src/kernel/epilogue.rs)。driver 的 last_k 管道、零开销的 Identity 默认值,以及穿过每条特殊路径的路由,全都是免费获得的(见 Epilogue 融合)。真正需要设计的是选定一条应用路径,并遵守一条硬规则:向量路径与标量路径必须逐位一致。完整 tile 走向量路径,边缘和跨步 tile 走标量路径,而同一个输出矩阵可以自由混用这两条路径。

  • 作用在 Acc 类型值上、有自然寄存器形态的变换,设 VECTOR = true 并实现 apply_reg。这是 FusedEpi 的模式。这里要留意 NaN 和带符号零的语义,比如 LeakyRelu 在两种形态下都写成完全相同的 max + slope*min 组合。
  • Acc 收窄成不同 Out 的变换,设 VECTOR_STORE = true 并实现 apply_store。这是 KRequantize 的模式,需要逐情形论证逐位相等,并用一致性扫描把它钉住。
  • 没有划算向量形态的变换,让两个标志都保持 false,一切都经暂存走标量 apply,这对任何 tile 形状都正确。但如果标量值可能与快路径的融合存储相差 1 个 ULP,就改为借用 MapEpi 的技巧:设 VECTOR = true,把 apply_reg 实现成“排空到栈、再逐 lane 应用“,这样这个变换看到的永远是普通 gemm 本会存储的那些精确位。

不管走哪条路径,都要在 gemmkit/tests/epilogue/ 里,紧挨着已有的测试,为这个新变换补上它自己的“gemm 再映射“等价测试。那套测试正是逐位契约真正被强制执行的地方。

测试与验证

一个把“这里逐位一致、那里按容差、处处可复现“当作头条承诺的库,成败全在于它的测试能多精确地钉住这些字眼。gemmkit 的测试套件都放在 gemmkit/tests/ 下。第一个结构性的决定是:什么不算测试。性能 harness 是测量工具,它们从不给 CI 把关。

tests/perf/ 是详尽的内部调查套件。它在一套取 9 次中位数的 harness 上跑 #[ignore] 基准测试,并且串行在一把共享锁之后,因为每一个基准都会打满所有核心。只有在某次改动需要拿到具体数字时,才会有人手动跑它。gemmkit/benches/gemm_bench.rs 则是精选出来的公开 cargo bench 界面。它包含五个头条 criterion 基准组:sgemmdtypesgemvprepackedbatched,用于配合 --save-baselinegemm crate 和 matrixmultiply 做回归追踪。这两套东西都不能让合并失败,因为在共享的 CI runner 上做性能断言,断言到的大多是噪声。

正确性、性质、一致性、模糊测试

正确性套件(tests/correctness/)针对形状、布局、alpha/beta 组合做扫描,对照一个独立的 f64 参考 GEMM 实现。这份参考实现,连同它周围的精度机制,只存在一份,放在 tests/oracle_common/ 里,包括元素 trait、确定性填充、f64 参考本身,以及每种元素类型各自的相对 Frobenius 精度门限。正确性套件和性质测试套件都用 #[path] 引入这个模块,所以值得信任的 oracle 只有一份。

在这些 oracle 扫描之上,还叠加了好几层检查。对外部 gemm crate 的交叉核对,能抓住仓库内参考实现抓不到的共同盲区错误,因为它是一份独立的实现。并行逐位一致性测试覆盖了做出这项承诺的那些路线。逐 ISA 的内核运行都经由泛型 driver 完成。还有容易被低估的一点:安全 API 精确的 panic 措辞,由 #[should_panic(expected = ...)] 里的子串(比如 "A.cols""aliases itself")把关,这样一条校验信息就不会在悄无声息中退化成没那么有用的版本。

性质测试把这些扫描进一步推广,三套测试都用 proptest 在形状、步长和旋钮值上做驱动。tests/props_api.rs 覆盖 oracle 精度、跨运行的位级确定性、串行与并行结果一致、beta == 0 的覆写语义、广播步长、批量调用,以及 panic 保证。tests/props_packed.rs 覆盖一般情形下预打包与普通调用的逐位一致性,以及在文档化的 tiny/gemv 例外集合上按容差校验。tests/props_knobs.rs 覆盖随机旋钮设置下的行为。

再往下一层,tests/simd_conformance.rs 检查的是 L0 词汇表本身。它直接构造宿主支持的每一个 ISA token,绕开分发逻辑,把每个 SimdOps 原语、同质的 KernelSimd blanket 实现,以及可移植的 fma_bvec 回退实现,逐 lane 拿去和标量参考对照。产品内核很少用到的原语,比如整数 reduce_sumfnma,以及拓宽接缝,正是在这里才得到覆盖。模块内部的扫描,比如 gemmkit/src/simd.rsrequant_store 的逐位相等测试,对向量化重量化契约做的是同一件事。这套测试不依赖 proptest,因此也能在 wasm 上运行,并对编译期的 simd128 token 做一致性测试。

模糊测试位于 gemmkit/fuzz/,这是一个仅限 nightly 的 cargo-fuzz 子工作空间,拥有自己的 workspace 根目录,被排除在稳定版构建之外。它包含六个 libFuzzer 目标。fuzz_gemm 生成构造即合法的问题实例,与朴素参考实现做差分检查,所以那里出现任何 panic 都是库本身的缺陷。fuzz_batched 对批量调用做同样的事。fuzz_prepackfuzz_prepack_i8 在预打包 API 上做往返测试,其中 i8 那个按位精确地把关。fuzz_api_validation 把对抗性的几何形状扔给带校验的入口,一个文档化的 "gemmkit:" panic 算作可接受的结果,其他任何情况都算作一处校验缺口。fuzz_knobs 在每次运行之前,把所有进程级的调优旋钮都设成对抗性的值,正是它以机械的方式找出了分块模型里那些算术溢出的问题类别。

隔离纪律

有两类全局状态,会让朴素的测试组织方式产生竞争。整套测试的结构正是围绕这两类状态来设计的。

调优旋钮是进程级的原子变量。每一个会改动旋钮的测试,都住在自己专属的二进制里:tests/tuning.rs 放 setter,tests/env.rs 负责环境变量解析,还有 tests/props_knobs.rs,以及负责拨动 GEMMKIT_DEEP_KC_BYTES、强制走各条深 K 路线的 tests/deep_k_narrow.rstests/env.rs 里只放了恰好一个测试,所以它对环境变量的访问从构造上就是单线程的。独立的二进制就是独立的进程,不可能和另一个二进制里的旋钮状态产生竞争。

不过同一个二进制内部,libtest 仍然会并发地跑测试。所以每一个会碰旋钮的测试,都要在该二进制专属的 KNOB_LOCK 互斥锁下串行执行,并且在释放这把锁之前,把自己改动过的值恢复原状。性质测试的二进制在此之上又加了一层 RAII 守卫,专门用来在 proptest 内部的 catch_unwind 中也能存活下来。

GEMMKIT_REQUIRE_ISA 更黏一些,因为分发逻辑每个进程只对它做一次记忆化。于是每个取值都有自己专属的钉住二进制:tests/env_isa_avx512f.rs_vnni_bf16_scalar_neon_wasm,再加上负责断言未知取值会 panic 的 env_isa_garbage.rs。每个二进制都让自己的所有测试,经由 tests/env_isa_common/ 里共享的一个 Once,在任何分发逻辑真正解析之前,完成那一次 set_var 调用。因为同一个二进制里的所有测试都钉住同一个值,所以究竟是哪个测试赢得了运行这个 Once 的先手,并不重要。这次写入会刻意覆盖继承来的 GEMMKIT_REQUIRE_ISA,正是这一点,让下面用 SDE 钉住的那些 CI 任务能够运行这些同样的二进制,同时依然走的是真实的、逐 ISA 的路线。

Miri 补上了模糊测试的 sanitizer 覆盖不到的那部分内存安全故事。CI 在 Miri 之下运行标量路径的正确性套件(miri_scalar_path),以及复数负步长的 unchecked 入口,由 Miri 直接解释打包和 microkernel 路径上那些真实的 unsafe 指针算术。cfg(miri) 的绕行分支,只在 Miri 无法解释某个硬件转换时才会出现,绝不是为了跳过逻辑本身。

CI 矩阵

.github/workflows/ci.yml 把这套钉住机制,变成了对 runner 本身并不具备的那些内核的覆盖:

任务覆盖的内容
test先用默认 feature,再用 --all-features,再关掉 parallel。此外还有关掉 std 的 4 种 feature 组合构建,呈现 no_std 风格。
kernel-scalar / kernel-fmaGEMMKIT_REQUIRE_ISA 钉在每个原生可用的内核上,跑完整套件。
avx512f_test / avx512vnni_test / avx512bf16_test在 Intel SDE(sde64 -spr)之下,钉住 AVX-512F、VNNI 点积、BF16 点积三种内核跑套件。SDE 仿真的是硅片,代码路径依然真实。
kernel-neon先在一台 arm64 macOS runner 上原生跑一遍整个工作空间,再钉住 neon 重跑一遍。
wasm_simd128 / wasm_simd128_threads在 wasmtime 之下、于 wasm32-wasip1 上钉住 simd128,跑正确性与一致性测试。threads 任务在 wasm32-wasip1-threads 上跑真实的 8 路并行。
no_stdx86_64-unknown-noneaarch64-unknown-nonewasm32-unknown-unknown 构建。
i686_check / msrv / lint / miri / coverage32 位检查、在 Rust 1.89.0(受支持的最低版本)上构建、fmtclippy -D warnings、上面提到的 Miri 任务,以及仅报告结果的 cargo-llvm-cov,它钉住了一份固定的 ISA 列表,这样报告出来的百分比就不会随 runner 池的变化而波动。

SDE 仿真跑起来比原生执行慢得多,这正是 GEMMKIT_FAST_TEST 派上用场的地方。它是一个只属于测试套件的开关,只在 tests/fast_test_common/ 里实现一次,再由各个 harness 引入,库本身从不读取它。这个开关把确定性的维度和系数扫描收缩成每个冗余组合只保留一个代表,同时依然会走到每一个分支和每一类路径。SDE 任务会打开这个开关,并配合 PROPTEST_CASES=16。原生任务则保留完整的扫描。把这个开关挡在库本身之外,就意味着一个为测试图方便而设的开关,永远不可能改变发布出去的行为。

最终效果把整章串了起来:前面几页做出的每一条断言,都有一个可以指认的测试把关,而这个测试所在二进制的隔离规则,让它的结果值得信任。这既包括特殊路径里的逐位一致性保证,也包括 Epilogue 融合里“gemm 再映射“的等价性,还有扩展点里的开闭性质,以及每个被钉住的内核,在项目自己并不拥有的硬件上的正确性。