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 的并行机制位于一个很小的层里，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 数。分块只从缓存模型推导 kc 和 nc，从不看线程数，深度切片也始终按同一个固定的 pc 顺序运行。因此，无论是一个 worker 清空游标，还是很多个 worker 一起清空游标，每个输出元素的浮点归约都呈现完全相同的形状。worker 数唯一能挪动的分块维度是 mc：宽 worker 数会通过驱动器的并行作业深度下限把它缩小，好让扁平任务列表在每个 worker 上都还能保持几个块的深度。因此，任务列表本身在不同宽度下并不总是逐块完全相同。但 mc 始终保持为 mr 的倍数，所以微 tile 集合（每个 mr 对齐的行偏移，加上唯一那条 m 尾巴）在任何切分方式下都相同。kc 和 pc 顺序都不曾改变，也就没有任何一个结果比特会移动。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 数，以及这份契约对测试意味着什么。