tune_gemm 设计

设计目标

tune_gemm 是 tune 的参考领域：一个分块参数很重要、正确性可以精确检查、且每个候选都可以序列化的问题。它在矩阵乘法上展示了完整的调优规范（先验证再计时、统计工作区、测试尾部、记录配置和环境），而不依赖于特定硬件的内核。

数学背景

乘积及其布局

对每个批次 $b$ ， $C_b = A_b B_b$ ，其中 $c_{ij} = \sum_{l=0}^{k-1} a_{il}\, b_{lj}$ ，代价为 $2mnk$ 次浮点运算。因此每次操作 $t$ µs 的测量时间对应于

\text{GFLOP/s} = \frac{2\,m\,n\,k\cdot\text{batch}}{t \cdot 10^{3}} .

一个 $r \times c$ 矩阵把第 $b$ 个批次的元素 $(i, j)$ 存放在 $b\,rc + i\,c + j$ （行主序布局）或 $b\,rc + j\,r + i$ （列主序布局）；这两个映射都是从索引三元组到 $[0, \text{batch}\cdot rc)$ 的双射，因此每种布局组合描述的都是同一个数学乘积。

分块与工作区

gemm_blocked 把迭代空间划分为 $mc \times nc \times kc$ 块，再把每个块划分为 $mr \times nr$ 寄存器分片。一个块涉及 $A$ 的一个 $mc \times kc$ 块、 $B$ 的一个 $kc \times nc$ 面板和 $C$ 的一个 $mc \times nc$ 块：

W = 8\,(mc\cdot kc + kc\cdot nc + mc\cdot nc) \ \text{bytes}, \qquad I = \frac{2\, mc\, nc\, kc}{W} = \frac{mc\, nc\, kc}{4\,(mc\cdot kc + kc\cdot nc + mc\cdot nc)} \ \text{flop/byte}.

对于立方体 $mc = nc = kc = s$ ， $W = 24 s^2$ 且 $I = s/12$ ：块越大，每个加载的值被复用得越频繁，直到 $W$ 超出本应容纳它的缓存。调优搜索的正是这种权衡，而 workspace_bytes 就是 $W$ ，被 valid_candidate 用作硬性上限。默认候选中最大的 $128 \times 128 \times 64$ 需要 $W = 8(8192 + 8192 + 16384) = 262144$ 字节。¹1 K. Goto 和 R. A. van de Geijn，“Anatomy of high-performance matrix multiplication”，ACM TOMS 34(3)，2008。

分块内核的精确性

对每个元素，gemm_reference 计算

\hat c_{ij} = \mathrm{fl}\bigl(\cdots\mathrm{fl}(\mathrm{fl}(0 + a_{i0} b_{0j}) + a_{i1} b_{1j}) \cdots + a_{i,k-1} b_{k-1,j}\bigr).

gemm_blocked 在第一个 $k$ 块中从 $0.0$ 开始计算每个元素，把部分和存入 $C$ （双精度数，因此存储是精确的），在下一个 $k$ 块中重新加载并按 $l$ 递增的顺序继续。它以相同的顺序对相同的操作数执行相同的运算，因此只要编译器对两个循环一视同仁（特别是不把其中之一收缩为融合乘加），两个结果就逐位相同。打包只是复制值。因此对内置内核，正确性检查可以使用容差 $0$ 。

对求和重新排序的内核（向量化的微内核、不同的循环顺序）只能做到接近。设单位舍入误差 $u = 2^{-53}$ 、 $\gamma_k = ku/(1 - ku)$ ，任何求和顺序都满足²2 N. J. Higham，Accuracy and Stability of Numerical Algorithms，第 2 版，SIAM，2002，§3.1。

\lvert \hat c_{ij} - c_{ij}\rvert \le \gamma_k \sum_l \lvert a_{il}\rvert\,\lvert b_{lj}\rvert ,

因此两个这样的内核至多相差 $2\gamma_k \sum_l \lvert a_{il}\rvert\lvert b_{lj}\rvert$ 。对于来自 reproducible_matrix 的元素， $\lvert a\rvert, \lvert b\rvert \le 1$ ，界为 $2\gamma_k k \approx 2k^2 u$ ：当 $k = 64$ 时约为 $9.1\cdot 10^{-13}$ 。这就是在这些输入上为重排序内核传给 validate_gemm 的绝对容差。

可复现的输入

reproducible_matrix(seed, r, c, layout) 运行 64 位线性同余生成器

x_{j+1} = 6364136223846793005\, x_j + 1442695040888963407 \pmod{2^{64}}, \qquad x_0 = \text{seed} \oplus \mathtt{0x9E3779B97F4A7C15},

并把每个状态映射为 $v = \bigl((x \gg 32) \bmod 2001\bigr)/1000 - 1$ 。乘数满足 $\equiv 1 \pmod 4$ 且增量为奇数，因此由 Hull–Dobell 定理，该生成器具有满周期 $2^{64}$ 。32 位值 $x \gg 32$ 对 2001 取模：由于 $2^{32} = 2146410 \cdot 2001 + 886$ ，2001 个值中有 886 个出现的可能性略高，相对高出 $1/2146410 \approx 4.7\cdot 10^{-7}$ 。元素是 $[-1, 1]$ 中 $0.001$ 的倍数。逻辑矩阵按行主序生成，然后以请求的布局存储，因此布局不会改变值。

尾部

边缘分片处理不是块大小倍数的 $m$ 、 $n$ 或 $k$ 。boundary_shapes 返回 $m - 1$ 、 $m$ 和 $m + 1$ ：当 $m$ 是 $mr$ 的倍数时，它们分别考验短分片、恰好整除和单行尾部。

设计决策

可移植的标量内核

问题。 真实的 GEMM 调优依赖于因目标而异的 SIMD 微内核。选择。 一个标量分块内核，以 microkernel_id 作为标签。理由。 该包在每个后端上演示并测试调优机制；测试框架可以把该标签映射到自己的内核。

先验证再测量

gemm_is_correct 在候选被计时之前就运行它并与参考比较。对于无效的候选或形状，execute_candidate 返回 None，而不是返回一个看起来像快速空乘法的结果。

工作区作为硬约束

候选所需的内存由其参数算出，并在执行前与上限比较，因此搜索永远不会选中一个无法在目标内存预算内运行的配置。

显式的计时范围

TimingScope 和 AllocationMode 是候选及其 id 的一部分，因此“仅计算”和“端到端”的结果永远不会被当作测量同一事物来比较。

序列化的配置

config_json 写出重复一次调优运行所需的一切，64 位值写为字符串，模式版本为 mmkts_1。它与 GemmEnvironment 一起说明调优结果在何处有效。

正确性与不变量

在上述条件下，gemm_blocked 与 gemm_reference 逐位相等；测试检查了带尾部的全部八种布局组合。
对每个有效候选都有 workspace_bytes(c) <= limit。
enumerate_gemm_candidates 以相同的顺序返回相同的 486 个候选，它们在上限 262144 下全部有效。
reproducible_matrix 只取决于其参数；行主序和列主序的结果描述同一个逻辑矩阵。
validate_gemm 把长度差异中的每个元素都计为不匹配，因此被截断的结果永远不会有效。

被否决的方案

针对特定目标的微内核。 无法在 MoonBit 后端之间移植。
验证中使用相对容差。 $C$ 的元素可能接近零；对有界输入而言，上面的绝对界才是诚实的界。
用平台随机生成器生成输入。 无法跨目标复现。

边界

只有 m_n_k 循环顺序和标量代码；三个微内核 id 运行同一个循环。
pack_operand 把整个操作数复制为行主序；它不构建缓存大小的打包面板。
allocation_mode、timing_scope 和 reuse_count 是描述性的：内核在每次调用时都分配其结果；计时什么由测试框架决定。
没有搜索循环：由 tune 和应用驱动搜索。
reproducible_matrix 生成一个矩阵；批量操作数由调用者拼接。

K. Goto 和 R. A. van de Geijn，“Anatomy of high-performance matrix multiplication”，ACM TOMS 34(3)，2008。 ↩
N. J. Higham，Accuracy and Stability of Numerical Algorithms，第 2 版，SIAM，2002，§3.1。 ↩