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

运行时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