分层结构
核心 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。而 kernel、driver、
special、dispatch 全都向下够到它们。pack 在 kernel 之下,因为是各家族的打包钩子
建立在打包原语之上,而不是反过来。
L0:词汇表,scalar.rs 与 simd.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:单位元常量,加上累加器类型。f16 和 bf16 在 f32 中累加。i8 在
i32 中累加。f32、f64 和复数类型都在自身中累加。Scalar 本身不带任何算术。真正的
运算要么向量化地发生在 SimdOps 里,要么在各家族的标量 epilogue 里,所以新增一个元素
类型永远不会连带出一套标量算术接口。细化 trait Float、NarrowFloat、ComplexFloat
按需叠加个别路径需要的额外能力。scalar.rs 刻意不知道的事:SIMD 的存在。它完全不知道
自己的常量将来会被广播进向量寄存器。
gemmkit/src/simd.rs 构成承重墙,和它一起的还有 simd/ 里的各后端:avx512.rs、
fma.rs、neon.rs、scalar.rs、wasm.rs,以及 complex.rs 里的复数胶水代码。三个
trait 分担这份职责。Simd 是一个零大小的 ISA 令牌。例子有 Avx512F、Fma、
Neon、ScalarTok、Simd128,外加具备点积能力的 Avx512Vnni 和 Avx512Bf16。它唯一
的方法是 vectorize,也就是 #[target_feature] 跳板。它把运行时选中的 intrinsic 放进
特性已启用的代码生成上下文。SimdOps<T> 是每个元素类型的厚词汇表:寄存器类型、
LANES、load/store/broadcast/mul/add/fma/reduce,以及可覆写的 accumulate_tile 调度。
再往上的 KernelSimd<L, R, A, O> 是拓宽/收窄接缝,让混合精度不需要驱动层的任何分支。这个
模块刻意不知道的事:它上面的一切。模块文档写明它只依赖 scalar 和 core,因此可以原封
不动拆成独立 crate。SimdOps 不知道什么是微面板、什么是缓存、什么是 GEMM。
L1:机械拷贝,pack.rs
gemmkit/src/pack.rs 存放两个共享的打包原语,把带步长的 A 或 B 区域拷成连续、微内核
尺寸的面板。这些就是 L4 各内核家族日后把打包钩子委托过来的那些拷贝。唯一的例外是复数
家族的实虚平面拆分打包,它随家族本身存放。
这两个原语是 pack_panels 和 pack_kgroup_panels。pack_panels 是微面板为主序的拷贝。
LHS 面板高 mr 行,RHS 面板宽 nr 列。两者用的是同一个例程,只是交换了“领先“和“深度“
两个步长的角色。尾部补零,带步长的源则走缓存分块转置。pack_kgroup_panels 是点积家族
使用的 k 组交错变体。
pack.rs 刻意不知道的事:自己的输出去了哪里。同一个例程既填充一次性调用的临时区,也
填充并行共享的打包缓冲,还填充调用方长期持有、贯穿整个进程生命周期的 PackedRhs。
pack.rs 从未见过 Workspace、工作线程或生命周期,只见过 dst、src 和步长。正是这份
漠不关心,让预打包路径与逐调用路径的字节完全一致。它只依赖 scalar,既不点名家族也不
点名缓存,因此才能坐得这么低。
L2:工作分发,parallel.rs
gemmkit/src/parallel.rs 拥有三样东西。
第一,Parallelism 枚举:Serial,或者 Rayon(n),其中 Rayon(0) 表示自动。
第二,感知工作量的线程数决策。总工作量低于某个门限时用一道串行闸门。显式指定的线程数
会被尊重,但有上限。自动档随总工作量 m*n*k 扩展,而不是直接跳到全部核心。访存受限的
矩阵乘向量形状另有一条独立的带宽规则。
第三,按需分发的机械装置。JobCursor 是一个无锁原子游标,工作线程从中拉取连续区块。
job_grain 和 packed_block_grain 这两个旋钮决定区块大小。for_each_worker 这个
fork-join 就是上层用作屏障的机制。parallel.rs 还提供 Ptr,让裸指针得以进入 rayon
闭包的 Send + Sync 垫片。
parallel.rs 只依赖 tuning,坐在一切调用它的东西之下。同一套工作线程词汇同样服务于
kernel、driver、special、dispatch。它刻意不知道的事:作业是什么。JobCursor 对
一个抽象计数发放索引区间。这个文件里没有任何地方提到 tile、矩阵或家族。这正是为什么日后
同一个游标既调度驱动的 tile,也不加区分地调度 B 打包面板、A 打包行块和 gemv 的行面板。
详见并行执行。
L3:机器模型,cache.rs
gemmkit/src/cache.rs 及其后端(cache/cpuid.rs、cache/sysfs.rs、cache/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.rs 与 kernel/(float.rs、mixed.rs、int.rs、complex.rs、
epilogue.rs)定义了 KernelFamily。它是把一种 GEMM 与另一种区分开的全部要素打成的
一捆。一个家族捆着 Lhs、Rhs、Acc、Out 这几个类型。它也捆着打包布局:pack_lhs
和 pack_rhs,两者都委托给 L1 的原语。它捆着微内核 microkernel_epi。它还捆着
OUT_IS_ACC、DEPTH_MULTIPLE 这类常量,告诉驱动该怎么为这个家族分块。
FloatGemm<T> 是基线。MixedGemm、IntGemm/IntGemmVnni、ComplexGemm 是原样复用
驱动的兄弟家族。这一层还拥有 Epilogue trait,带着它零成本的 Identity。它也拥有
AlphaStatus 和 BetaStatus 这两个枚举。驱动会预先算好这两者,好让微内核永远不必比较
浮点数。家族刻意不知道的事:自己的 tile 尺寸。MR_REG 和 NR 是微内核方法上的 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 产生完全相同的面板字节。
它的公开入口是 run、run_epilogue、run_packed_rhs、run_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.rs 与 special/(gemv.rs、small_k.rs、small_mn.rs、
batched.rs)存放寄存器分块驱动并不适合的形状的专用路径。它们是矩阵乘向量、低深度
GEMM、小 m,n 长 k 的内积形状,以及批量编排层。它们都藏在同样的公共入口背后,
特殊路径有完整介绍。特殊路径刻意不知道的事:自己为什么被选中。门限
gemv_threshold、small_k_threshold、small_mn_dim 住在上一层的分发逻辑和旁边的
tuning 模块里。small_k::run 甚至分辨不出自己在服务 gemm、gemm_fused 还是
gemm_map,因为 epilogue 是作为不透明的泛型参数到达的。批量路径是唯一一处某层向上
够回去的地方。batched.rs 把每个元素经由 dispatch::execute(L7)转发,于是它继承
上方的整条阶梯。这就是本页末尾讨论的那唯一一处带注解的例外。
L7:运行时 ISA 选择,dispatch.rs
gemmkit/src/dispatch.rs 与 dispatch/(isa.rs、float.rs、mixed.rs、int.rs、
complex.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.rs 与 api/(batched.rs、cplx.rs、fused.rs、int8.rs、
map.rs、packed.rs)定义了好几样东西。它们是带步长视图 MatRef/MatMut、各家族的
安全入口,以及 validate_gemm_views 的 panic 清单。安全入口分 *_with(调用方自备
工作区)和 *_unchecked(裸引擎)两种变体。这一层还负责把视图降解为 Task。
它刻意不知道的事:分发层以下的一切。API 层看不到将要运行哪个 ISA、会选出什么分块、
打包会不会发生。校验完成后,它把 Task 递给 dispatch::execute,任务就结束了。对称地,
MatRef 从不出现在这一层之下。crate 的其余部分只讲指针和步长。
为什么箭头只指向下方
依赖方向是这套架构唯一的硬规则。每一层由上层驱动,且对上层一无所知。simd 只依赖
scalar 和 core。驱动从不点名元素类型或 ISA。L7 以下没人知道分发层的存在。L8a 以下
没人听说过切片。
这条规则只有唯一一处刻意的、带注解的例外。special/batched.rs(L6)把每个批量元素
经由 dispatch::execute(L7)向上转发。这次再入让每个元素继承与其形状上独立 gemm
调用相同的 driver、small-k、small-mn、gemv 路由。它省去了要手工维护的第二套分发阶梯。
这是一支有意为之的向上箭头,也是 crate 里唯一的一支。
三个收益支撑着这份纪律。其一,扩展成本。因为知识只向下流动,新 ISA、新元素类型、新
家族都在自己那一层接入。其下的一切可证明地原封不动。扩展点描述的各条接缝
之所以成立,正是因为没有任何下层可能对上层做过特判。其二,审阅局部性:审计微内核只需读
kernel/float.rs 和 SimdOps 契约,别无其他。审计调度只需读 driver.rs 和
parallel.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 时,每次调用都使用
全新的工作区。详见打包与工作区。