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

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

循环重排序与空间局部性

展开优化指令流水线，而循环重排序（或循环置换）则优化内存访问模式。嵌套循环的迭代顺序决定了内存地址的访问顺序。因为硬件以连续的缓存行（通常为 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$ 顺序，访问模式会发生显著变化：

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

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

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

组合变换

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

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

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

这部分内容有帮助吗？