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 能相乘好几种元素类型: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 交叉验证,见测试与验证。第三方能在这道接缝上搭建什么,见扩展点