矩阵乘法实践

矩阵乘法是几乎所有深度学习模型的计算核心。从多层感知器中的全连接层到Transformer中的查询-键-值投影，通用矩阵乘法 (GEMM) 操作的效率直接决定了神经网络的训练速度和推理延迟。了解如何优化这个单一内核为优化几乎任何张量操作提供了依据。

在本节中，我们将运用前面讨论的调度基本方法——重排、分块和向量化，将一个朴素的矩阵乘法实现转换为一个高性能内核。我们将观察每种转换如何影响内存访问模式和算术强度。

基准实现

让我们考虑两个矩阵 $A$（大小 $M \times K$）和 $B$（大小 $K \times N$）的乘法，以生成矩阵 $C$（大小 $M \times N$）。其数学定义如下：

$C_{i,j} = \sum_{k=0}^{K-1} A_{i,k} \times B_{k,j}$

在标准的行主序存储格式中，行中的元素在内存中是连续存储的。将数学定义直接转换为代码会产生一个三重嵌套循环。我们称之为“朴素”实现。

# 朴素实现
for i in range(M):
    for j in range(N):
        for k in range(K):
            C[i, j] += A[i, k] * B[k, j]

尽管正确，但此实现在现代硬件上效率不高。最内层循环迭代变量是 $k$。对于矩阵 $A$，我们访问 $A[i, k]$，随着 $k$ 的增加，它在内存中是顺序移动的。然而，对于矩阵 $B$，我们访问 $B[k, j]$。由于 $B$ 以行主序存储，增加 $k$ 意味着在内存中跳过 $N$ 个元素才能到达下一行。

这种跨步访问模式会导致频繁的缓存未命中。CPU 获取包含 $B[k, j]$ 的缓存行，只使用一个值，然后将其丢弃以获取下一行。这将性能限制在主内存的速度，而不是处理器的速度。

优化 1：循环重排

第一个也是通常最简单的优化是改变循环的顺序。我们希望确保最内层循环对正在读写的矩阵进行连续内存访问。通过交换 j 和 k 循环，我们改变了迭代模式。

# 重排实现
for i in range(M):
    for k in range(K):
        for j in range(N):
            C[i, j] += A[i, k] * B[k, j]

在这种配置中：

1. $A[i, k]$ 对于内层循环 j 来说是常量。编译器可以将其加载到寄存器中一次，并重用 $N$ 次。
2. $B[k, j]$ 随着 j 的增加而顺序访问。这能最大化空间局部性，使得 CPU 可以利用获取到的缓存行中的每个元素。
3. $C[i, j]$ 也被顺序访问。

这种简单的改变通常通过减少内存流量带来显著的速度提升，尽管算术操作的总数保持不变。

优化 2：循环分块

虽然重排提升了空间局部性，但它本身并未解决大型矩阵时间局部性的问题。如果矩阵 $A$、$B$ 和 $C$ 过大，无法完全放入 L1 或 L2 缓存，处理器将驱逐稍后需要的数据，导致“缓存颠簸”。

循环分块（也称阻塞）通过将执行划分为适合缓存的更小子矩阵（块）来处理这个问题。我们将迭代空间拆分为外层循环（遍历块）和内层循环（在块内迭代）。

分块将全局迭代空间划分为更小的块，以最大化缓存层次结构内的数据重用。

考虑将 i 和 j 循环按因子 bn 和 bm 进行分块。结构从三重循环转换为六层循环嵌套（或更少，取决于哪些维度进行了分块）。通过一次处理 $C$ 的一个小块，我们确保在计算过程中 $A$ 和 $B$ 的相应块保留在缓存中。

块大小的选择与硬件相关。这需要在缓存大小和寄存器压力之间取得平衡。如果块太小，我们会失去均摊的效益。如果块太大，数据会溢出缓存。

优化 3：向量化

现代处理器支持 SIMD（单指令多数据）扩展，例如 AVX-512 或 NEON。这些指令允许 CPU 同时对多个数据点执行相同的操作。例如，一个指令可能不是乘以两个浮点数，而是乘以两个包含 8 或 16 个浮点数的向量。

编译器可以自动向量化循环，但前提是内存访问模式可预测且依赖关系明确。重排后的循环结构（i, k, j）有助于实现这一点，因为最内层循环 j 对连续内存执行逐元素加法和乘法。

为了明确地针对 SIMD 单元，我们通常会展开最内层循环或使用向量化编译指示标记它。在像 TVM 或 MLIR 这样的机器学习编译器栈中，这是在调度阶段通过根据向量通道大小拆分内层循环并将内部分拆绑定到向量线程来完成的。

性能影响分析

结合这些技术，可以得到一个比朴素基准快几个数量级的内核。下面的图表说明了在标准 CPU 上优化密集矩阵乘法时通常观察到的累积性能提升。

对矩阵乘法内核应用顺序优化时观察到的累积加速因子。

总结

在实际的编译器工作流程中，您不会手动在 C++ 中编写这些循环。相反，您定义计算，然后应用调度。

1. 定义： 您使用高级 DSL 定义 $C = A \times B$。
2. 调度： 您创建一个调度对象。
3. 转换：
   • split：您将轴 i 和 j 拆分为外层和内层循环以定义块大小。
   • reorder：您重排轴，将分块循环置于顶部，将密集计算循环置于底部。
   • vectorize：您标记最内层迭代轴，以便用 SIMD 指令替换。

最终生成的汇编代码将大量使用向量寄存器，并最大程度地减少内存地址之间的跳转。这种检查循环嵌套、找出内存访问瓶颈并系统地应用调度基本方法的过程是内核优化的核心。尽管自动调优可以找出具体的块大小（例如，32x32 是否比 64x64 更好？），但理解结构转换使您能够有效地定义搜索空间。