循环分块与缓存局部性

现代硬件加速器执行算术操作的速度比从主内存获取数据快好几个数量级。在优化机器学习内核时，特别是像矩阵乘法或卷积这样的密集操作，主要瓶颈很少是原始计算速度。相反，性能受限于内存带宽。为此，我们必须重新组织循环以优先考虑数据复用，这是一种以缓存局部性为核心的策略。

内存层级结构带来的挑战

处理器采用内存层级结构来减少从DRAM（动态随机存取内存）获取数据的延迟。最靠近算术单元的是寄存器和L1缓存，接着是L2和L3缓存。从L1缓存访问数据只需几个时钟周期，而从DRAM获取数据可能需要数百个周期。

循环优化的目标是尽可能长时间地将活跃工作数据集保持在最快的缓存级别内。这基于两个原理：

1. 时间局部性： 如果一个数据位置被引用，它很可能很快再次被引用。
2. 空间局部性： 如果一个数据位置被引用，具有相邻地址的数据位置很可能很快被引用。

深度学习算子通常展现出很高的复用潜力，但朴素的循环实现未能加以利用。

分析朴素矩阵乘法

考虑一个标准矩阵乘法 $C = A \times B$，其中所有矩阵都是 $N \times N$ 大小。其数学定义是：

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

在典型的行主序布局中（在C++和PyTorch中常见），一行的连续元素在内存中是连续存储的。一个使用三个嵌套循环的朴素实现如下所示：

for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j++) {
        for (int k = 0; k < N; k++) {
            C[i][j] += A[i][k] * B[k][j];
        }
    }
}

尽管这段代码在数学上是正确的，但矩阵 $B$ 的内存访问模式效率低下。随着内层循环 k 的增加，我们访问 B[0][j]、B[1][j]、B[2][j] 等。在行主序布局中，这些元素间隔 $N$ 个位置存储。如果 $N$ 很大，每次访问 $B$ 都会跳跃内存，很可能导致缓存未命中。等到循环再次访问之前加载的缓存行时，该缓存行可能已被逐出，以便为其他数据腾出空间。这就造成了所谓的缓存颠簸。

循环分块原理

循环分块，也称为阻塞，通过将迭代空间划分为可以完全放入缓存的更小块（瓦片）来解决此问题。我们不是遍历 $i, j, k$ 的整个范围，而是遍历矩阵的小块。

这种转换需要将原始循环拆分为“外层”循环（遍历瓦片）和“内层”循环（在瓦片内部迭代）。

这是矩阵乘法示例的结构变化。我们引入一个瓦片大小因子，$T_{size}$。

// 外层循环遍历瓦片
for (int i_outer = 0; i_outer < N; i_outer += T_size) {
    for (int j_outer = 0; j_outer < N; j_outer += T_size) {
        for (int k_outer = 0; k_outer < N; k_outer += T_size) {
            
            // 内层循环在当前瓦片内迭代
            for (int i_inner = i_outer; i_inner < i_outer + T_size; i_inner++) {
                for (int j_inner = j_outer; j_inner < j_outer + T_size; j_inner++) {
                    for (int k_inner = k_outer; k_inner < k_outer + T_size; k_inner++) {
                        C[i_inner][j_inner] += A[i_inner][k_inner] * B[k_inner][j_inner];
                    }
                }
            }
        }
    }
}

在这个分块版本中，处理器将矩阵 $A$ 和矩阵 $B$ 的一个小的 $T_{size} \times T_{size}$ 块加载到缓存中。内层循环执行所有涉及这些块的必要计算，然后再处理下一组。由于这些块足够小，可以放入缓存中，我们消除了重复的、开销大的主内存访问。

下面的图示说明了矩阵如何逻辑上划分为更小的处理块。

矩阵被划分为 4x4 个独立区域。突出显示的蓝色区域表示当前加载到缓存中的“活动瓦片”，这使得本地计算能够在数据被逐出之前完成。

选择瓦片大小

选择合适的 $T_{size}$ 非常关键，并且取决于硬件。

1. 过小： 循环开销（计数器递增、分支）会主导执行时间。您未能充分使用可用的缓存行。
2. 过大： 内层循环所需的数据超过缓存大小，导致数据被逐出，并将性能恢复到朴素实现的水平。

在像TVM或MLIR这样的自动化机器学习编译栈中，寻找最优瓦片大小通常由自动调优组件处理。调优器会经验性地测试不同的瓦片配置（例如，$16 \times 16$, $32 \times 32$, $64 \times 64$），以找到目标硬件特定L1/L2缓存大小的最佳点。

多级分块

现代CPU和GPU拥有多级缓存。为了最大化性能，编译器通常分层应用分块。我们可能会使用适合L2缓存的较大瓦片大小，然后将其细分为适合L1缓存或寄存器文件的小瓦片。

下面的转换图描绘了从主内存到计算单元的数据流，说明了为什么不同的瓦片大小对应不同的硬件内存级别。

数据移动层级结构显示了分块策略如何与硬件内存级别对齐。较大的瓦片面向L2缓存，而子瓦片面向L1和寄存器。

对算术强度的影响

我们最终试图改进的指标是算术强度，即浮点运算次数 (FLOPs) 与内存传输字节数的比率。

在朴素的矩阵乘法中，每次内存访问，我们执行2次操作（乘法和加法）。采用分块后，我们加载一个 $T_{size} \times T_{size}$ 大小的块（这需要为 A 和 B 传输 $2 \times T_{size}^2$ 个元素），并对这些数据执行 $T_{size}^3$ 次操作。随着 $T_{size}$ 的增加（达到缓存限制），计算与内存访问的比率显著提高。

通过确保算术单元不断地从缓存中获取数据，而不是等待主内存，循环分块将内存受限操作转换为计算受限操作，使硬件能够接近其理论峰值性能。在下一节中，我们将分块与向量化结合起来，以进一步发挥现代处理器的能力。