特殊路径
一次 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.rs 中 run_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 == 1 或 n == 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*C,kc = 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:小输出的水平点积
当 m 和 n 都远小于 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_batched 与 gemm_batched_fused(后者把同一份共享的 epilogue 穿过同一套骨架)共用一份调度实现,普通形态和融合形态因此不会出现分歧。
对于元素形状各不相同的批,指针数组形式 gemm_batched_ptr_unchecked 改为接受一个 GemmProblem 描述符切片,每个描述符携带自己的维度、步长和指针。这种形式采用更简单的 resolve_batch_flat 策略:把整个 GEMM 分给 worker,绝不做元素内部的拆分,因为这里没有统一的驻留性可供判断。gemm_batched_slice 是它经过校验的安全孪生版本。