循环展开与重排序

现代处理器使用深度指令流水线来获得高时钟频率。为了保持这些流水线充满，硬件必须提前准确预测下一条执行指令。循环结构对此机制构成挑战，因为它们引入频繁的分支。每次循环迭代时，处理器必须检查条件并跳回代码块的开头。

分块优化内存层级之间的数据移动，而循环展开与重排序则优化指令流本身。这些变换减少了控制逻辑的开销，并向处理器提供更多独立操作，从而实现更好的指令级并行（ILP）。

通过循环展开减少开销

循环展开是一种变换，它增加循环体的大小，同时减少迭代次数。其主要目的是减少每次迭代支付的“代价”。此代价包括三个主要组成部分：增加循环计数器、将计数器与边界比较以及执行条件跳转。

以一个简单的向量 (vector)加法操作为例，我们计算 $N$ 个元素的 $C[i] = A[i] + B[i]$。在标准实现中，控制逻辑执行 $N$ 次。

如果我们将循环展开 4 倍，我们重写代码，使其每次迭代处理四个元素。循环现在执行 $N/4$ 次。我们执行控制逻辑（递增、比较、跳转）的次数仅为原始代码的 $25\%$。

通过减少分支开销，展开为编译器提供更大的指令分析窗口。在原始循环中，迭代之间的依赖关系对硬件调度器来说可能不明显。在一个展开的代码块中，编译器明确列出独立操作。这使得 CPU 可以在同时添加 A[i] 和 B[i] 的同时加载 A[i+1] 的数据，有效地填充执行流水线的空闲。

然而，此方法并非没有代价。展开会增加总二进制文件大小，这可能会给指令缓存带来压力。此外，如果展开的循环体使用过多临时变量，它可能会超出可用物理寄存器的数量。这迫使编译器将寄存器“溢出”到内存中，这会带来显著的延迟开销，超过了减少分支所带来的好处。自动调优系统通常寻找理想的展开因子，通常是 4、8 或 16，以兼顾并行性与寄存器压力。

上方图表展示了循环展开所涉及的取舍。随着展开因子增加，吞吐量 (throughput)（蓝线）最初由于开销减少而提升。然而，一旦寄存器使用率（红条）接近硬件限制，性能会因寄存器溢出而下降。

循环重排序与空间局部性

展开优化指令流水线，而循环重排序（或循环置换）则优化内存访问模式。嵌套循环的迭代顺序决定了内存地址的访问顺序。因为硬件以连续的缓存行（通常为 64 字节）加载内存，所以顺序访问数据（步长为 1）比大步长访问数据快几个数量级。

此优化的典型例子是矩阵乘法。我们来分析计算 $C = AB$，其中 $C$、$A$ 和 $B$ 是 $N \times N$ 矩阵。数学定义表明了 $i, j, k$ 的顺序：

$C_{ij} = \sum_{k=0}^{N-1} A_{ik} \times B_{kj}$

在循环嵌套为 for i、for j、for k 的标准实现中，矩阵 $B$ 的内存访问模式效率低下。当内层循环 k 递增时，我们访问 $B[0][j]$，然后是 $B[1][j]$，接着是 $B[2][j]$。假设采用行主序存储格式（C/C++ 和 PyTorch 中的标准），这些元素在内存中相隔 $N$ 个位置存储。这导致了“步长为 N”的访问模式，这导致几乎每次读取操作都发生缓存未命中。

为了解决这个问题，编译器会重排序循环。通过将内层循环交换为 $i, k, j$ 顺序，访问模式会发生显著变化：

1. 循环 i（外层）： 选择 $A$ 和 $C$ 的行。
2. 循环 k（中层）： 选择特定元素 $A_{ik}$（内层循环的标量）以及 $B$ 的行。
3. 循环 j（内层）： 遍历列。

在这种配置下，内层循环使用 $B$ 的一行来更新 $C$ 的一行。$C$ 和 $B$ 都以步长为 1 的方式访问。处理器读取连续的内存块，使加载到缓存行中的每个字节都得到完全使用。这种简单的置换可以在现代 CPU 上将性能提高 5 到 10 倍，而无需更改任何算术操作。

此图比较了循环嵌套结构。在原始顺序中，内层循环对 'k' 进行迭代，导致从矩阵 B 进行非顺序读取。在重排序版本中，内层循环对 'j' 进行迭代，使内存读取与行主序存储对齐 (alignment)。

组合变换

实践中，编译器不会单独应用这些技术。它们组合使用以实现最大算术强度。由 TVM 或 XLA 等编译器生成的典型矩阵乘法核，可能首先涉及循环分块以使数据适合 L1 缓存，然后对这些分块内的循环进行重排序以确保步长为 1 的访问，最后展开最内层循环以最大化向量 (vector)化。

例如，编译器可能将一个简单的三重嵌套循环转换为一个具有六层或更多层嵌套（瓦片上的循环和瓦片内的循环）的复杂结构，其中最内层块完全展开。这种组合使得代码能够隐藏内存延迟：当流水线的一部分等待来自 L2 缓存的数据时（由于分块），另一部分则在寄存器上执行高吞吐量 (throughput)的向量算术（由于展开）。

编译器面临的难题是识别嵌套中哪些循环可以安全展开，以及哪些置换能保持程序的正确性。这通过分析数据依赖关系来确定。如果迭代 $k$ 依赖于迭代 $k-1$ 的结果，展开需要仔细处理，以防止竞态条件或错误的计算顺序。现代 ML 编译器构建循环变量的依赖图，以确保任何重排序或展开都符合原始数据流要求。