通用矩阵乘法 (GEMM) 是现代神经网络运行中占主导地位的工作负载。尽管高级框架将此操作抽象为 torch.matmul 或 np.dot 等单一函数调用,但底层编译器必须生成一系列机器指令,有效利用 CPU 或 GPU 管线。我们的目标是将一个朴素的、数学上正确的实现,转换为一个尊重硬件内存层次结构的调度。我们将使用类似于 TVM Schedule API 的高级 Python 语法来定义这些转换。这种方法使我们能够将算法定义与其执行策略分离。基准实现考虑两个矩阵,$A$ 的形状为 $(M, K)$,$B$ 的形状为 $(K, N)$,结果 $C$ 的形状为 $(M, N)$。数学定义描述了计算状态:$$C_{i,j} = \sum_{k=0}^{K-1} A_{i,k} \times B_{k,j}$$在像 C++ 这样的标准命令式语言中,直接转换会产生一个三层嵌套循环。// 朴素实现 for (int i = 0; i < M; ++i) { for (int j = 0; j < N; ++j) { for (int k = 0; k < K; ++k) { C[i][j] += A[i][k] * B[k][j]; } } }尽管功能正确,这段代码在大矩阵上的性能较差。核心问题在于矩阵 $B$ 的访问模式。假设按行主序存储,顺序访问 $B[k][j]$ 意味着每当 $k$ 增加时,在内存中会跳过 $N$ 个元素。这种跨步访问模式导致频繁的缓存未命中,因为 CPU 获取一个缓存行,但在请求远处的另一个内存地址之前只使用一个值。优化 1:循环分块 (Blocking)为减轻缓存抖动,我们应用循环分块。此技术将迭代空间划分为更小的块,这些块完全适合 L1 或 L2 缓存。通过一次计算 $C$ 的一个小子矩阵,我们确保 $A$ 和 $B$ 所需的元素保留在缓存中以供重复使用。在我们的调度中,我们将空间轴 $i$ 和 $j$ 拆分为外层和内层循环。# 将轴 i 拆分为大小为 bn 的块 i_outer, i_inner = s[C].split(axis=i, factor=bn) # 将轴 j 拆分为大小为 bn 的块 j_outer, j_inner = s[C].split(axis=j, factor=bn) # 将归约轴 k 拆分为大小为 bk 的块 k_outer, k_inner = s[C].split(axis=k, factor=bk)这种转换重写了执行顺序。处理器不再遍历整个行 $M$,而是遍历小块。这种分解的可视化表示了全局矩阵如何作为子矩阵网格进行处理。digraph G { rankdir=TB; node [shape=box, style="filled", color="#dee2e6", fontname="Helvetica"]; edge [fontname="Helvetica"]; subgraph cluster_global { label = "全局矩阵 (M x N)"; style=dashed; color="#adb5bd"; Global [label="完整迭代空间\ni: 0..M, j: 0..N"]; } subgraph cluster_tiled { label = "分块执行"; style=filled; color="#e9ecef"; OuterLoops [label="外层循环\n遍历块坐标", fillcolor="#a5d8ff"]; InnerLoops [label="内层循环\n计算微核心 (bn x bn)", fillcolor="#96f2d7"]; OuterLoops -> InnerLoops [label="实例化"]; } Global -> OuterLoops [label="拆分转换"]; }将全局迭代空间分解为局部块以改善缓存局部性的图示。阻塞因子 (bn, bk) 的选择很重要。如果块大小过小,循环开销会占主导地位。如果过大,工作集会超出缓存大小,导致性能回落到基准水平。优化 2:循环重排分块后,嵌套循环的顺序决定了数据局部性。拆分生成的默认顺序可能是 i_outer, i_inner, j_outer, j_inner, k_outer, k_inner。然而,我们希望最内层循环在瓦片上执行密集计算。我们通常会重排循环,将 outer 迭代器推到顶部,并保持 inner 迭代器相邻。一种常见策略是将循环排序为:i_outer, j_outer, k_outer, i_inner, j_inner, k_inner。这种排序确保对于 $C$ 的一个固定块,我们遍历 $A$ 和 $B$ 的必要块,从而最大限度地提高数据重用。s[C].reorder(i_outer, j_outer, k_outer, i_inner, j_inner, k_inner)优化 3:向量化与布局现代处理器使用 SIMD(单指令多数据)单元执行指令,例如 Intel 上的 AVX-512 或 ARM 上的 NEON。这些单元可以在一个时钟周期内对多个数据点(例如,8 个浮点数)执行算术运算。为利用 SIMD,最内层循环必须映射到连续的内存地址。如果我们的最内层循环是 j_inner 且数据布局是行主序,则内存访问是连续的。我们显式地将此轴标记为向量化。s[C].vectorize(j_inner)编译器现在将尝试用向量内联函数(例如 vaddps、vmulps)替换标量指令。如果内存访问不连续,或者存在阻止并行执行的数据依赖,编译器将无法进行向量化,更糟的是,会生成“收集/分散”指令,严重降低性能。优化 4:循环展开分支预测和循环索引维护会产生开销。循环展开将循环体在单次迭代中复制多次,减少了比较和跳转指令的数量。我们对微核心内的归约轴 k_inner 应用展开。s[C].unroll(k_inner)这条指令告知编译器明确写出求和步骤。尽管这会增加二进制文件的大小,但它允许 CPU 更有效地进行指令管线化,并提供指令级并行化的机会。性能分析逐步应用这些转换会带来显著的性能提升。我们以 GFLOPS(每秒十亿次浮点运算)衡量吞吐量。基准实现通常由于内存带宽瓶颈,只能达到硬件理论峰值的一小部分。通过优化调度,我们将瓶颈从内存带宽转移到计算能力。下图展示了在标准 CPU 架构上,随着我们应用每一层优化,性能是如何提升的。{ "layout": { "title": "GEMM 优化性能扩展", "xaxis": { "title": "优化阶段" }, "yaxis": { "title": "吞吐量 (GFLOPS)" }, "barmode": "group", "plot_bgcolor": "#f8f9fa", "paper_bgcolor": "#ffffff", "font": { "family": "Helvetica" } }, "data": [ { "type": "bar", "x": ["朴素循环", "分块", "重排", "向量化 + 展开"], "y": [2.5, 18.2, 45.6, 112.4], "marker": { "color": ["#adb5bd", "#74c0fc", "#63e6be", "#339af0"] }, "text": ["2.5", "18.2", "45.6", "112.4"], "textposition": "auto" } ] }不同优化阶段的吞吐量比较。从分块开始的性能提升意味着更好的缓存使用,而向量化则利用了计算单元。数据表明,虽然算法正确性足以满足功能需求,但精心设计循环调度对于性能来说是必不可少的。分块、重排和向量化的组合使软件逻辑与硬件的物理限制保持一致,与朴素实现相比,性能提升超过 40 倍。在后续章节中,我们将讨论自动化搜索策略(自动调优)如何在无需人工干预的情况下,确定任意硬件目标的最佳拆分因子 (bn, bk)。