动手实践：优化矩阵乘法

通用矩阵乘法 (GEMM) 是现代神经网络 (neural network)运行中占主导地位的工作负载。尽管高级框架将此操作抽象为 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$，而是遍历小块。这种分解的可视化表示了全局矩阵如何作为子矩阵网格进行处理。

将全局迭代空间分解为局部块以改善缓存局部性的图示。

阻塞因子 (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：向量 (vector)化与布局

现代处理器使用 SIMD（单指令多数据）单元执行指令，例如 Intel 上的 AVX-512 或 ARM 上的 NEON。这些单元可以在一个时钟周期内对多个数据点（例如，8 个浮点数）执行算术运算。

为利用 SIMD，最内层循环必须映射到连续的内存地址。如果我们的最内层循环是 j_inner 且数据布局是行主序，则内存访问是连续的。我们显式地将此轴标记 (token)为向量化 (quantization)。

s[C].vectorize(j_inner)

编译器现在将尝试用向量内联函数（例如 vaddps、vmulps）替换标量指令。如果内存访问不连续，或者存在阻止并行执行的数据依赖，编译器将无法进行向量化，更糟的是，会生成“收集/分散”指令，严重降低性能。

优化 4：循环展开

分支预测和循环索引维护会产生开销。循环展开将循环体在单次迭代中复制多次，减少了比较和跳转指令的数量。

我们对微核心内的归约轴 k_inner 应用展开。

s[C].unroll(k_inner)

这条指令告知编译器明确写出求和步骤。尽管这会增加二进制文件的大小，但它允许 CPU 更有效地进行指令管线化，并提供指令级并行化的机会。

性能分析

逐步应用这些转换会带来显著的性能提升。我们以 GFLOPS（每秒十亿次浮点运算）衡量吞吐量 (throughput)。基准实现通常由于内存带宽瓶颈，只能达到硬件理论峰值的一小部分。

通过优化调度，我们将瓶颈从内存带宽转移到计算能力。下图展示了在标准 CPU 架构上，随着我们应用每一层优化，性能是如何提升的。

不同优化阶段的吞吐量比较。从分块开始的性能提升意味着更好的缓存使用，而向量 (vector)化则利用了计算单元。

数据表明，虽然算法正确性足以满足功能需求，但精心设计循环调度对于性能来说是必不可少的。分块、重排和向量化 (quantization)的组合使软件逻辑与硬件的物理限制保持一致，与朴素实现相比，性能提升超过 40 倍。

在后续章节中，我们将讨论自动化搜索策略（自动调优）如何在无需人工干预的情况下，确定任意硬件目标的最佳拆分因子 (bn, bk)。