循环重排与展开

高效的内核执行，依赖于软件指令与硬件物理特性的一致性。循环分块虽能优化缓存容量，但它本身不能解决数据如何从缓存行中取出，或CPU流水线使用效率如何的问题。循环重排和展开是针对空间局部性和指令级并行(ILP)的转换方法。

这些转换方法在不改变语义结果的前提下，修改执行顺序和结构。对于深度学习编译器来说，目的是重新排列迭代空间，使内存访问连续，并保持指令流水线饱和。

内存布局与空间局部性

为了明白循环顺序的重要性，我们必须查看多维张量如何被展平到线性内存中。在C/C++以及大多数深度学习框架（例如PyTorch或TVM默认设置）中，张量以行主序存储。这意味着内存中最后一个维度变化最快。

考虑一个形状为 $(N, M)$ 的张量 $A$。元素 $A_{i,j}$ 存储在地址 base + i*M + j 处。

• 先访问 $A_{i,j}$ 再访问 $A_{i,j+1}$（内部索引递增），会形成步长为1的访问模式。这种方式最佳，因为加载 $A_{i,j}$ 会将相邻元素带入缓存行中。
• 先访问 $A_{i,j}$ 再访问 $A_{i+1,j}$（外部索引递增），会形成步长为 $M$ 的访问模式。如果 $M$ 很大，每次访问都可能引起缓存未命中。

下图说明了步长为1的访问与跨步访问之间在内存遍历上的差异。

内存访问模式决定了缓存效率。步长为1的访问方式能充分使用整个缓存行，而跨步访问则会造成内存带宽的浪费。

矩阵乘法中的循环置换

循环重排的经典例子是矩阵乘法（MatMul）。其数学定义涉及三个索引：$i$、$j$ 和 $k$。

$C_{i,j} += A_{i,k} \times B_{k,j}$

一个朴素的实现通常采用 $i-j-k$ 的顺序：

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

在此配置中：

1. $C[i][j]$ 在内层循环中不变（好）。
2. $A[i][k]$ 以步长为1的方式访问（好）。
3. $B[k][j]$ 以步长为 $M$ 的方式访问（差）。

因为 $B$ 在存储时是行主序，但访问时是列主序，此调度会遇到严重的缓存抖动。通过将循环顺序改变为 $i-k-j$，我们修改了访问模式：

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

在 $i-k-j$ 的置换中：

1. $A[i][k]$ 在内层循环中不变。
2. $B[k][j]$ 以步长为1的方式访问（好）。
3. $C[i][j]$ 以步长为1的方式访问（好）。

这种简单的重排能带来显著的加速，在标准CPU上常能达到5到10倍，仅仅通过遵从硬件对线性访问的偏好。编译器会采用多面体分析来建立模型，以此确定无依赖的置换，使空间局部性达到最优。

循环展开

循环重排优化数据移动，而循环展开则优化指令流。展开操作包括在单次迭代中多次复制循环体，并降低循环计数器的更新频率。

考虑一个标准的点积循环：

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

CPU执行这段代码会产生开销：i 的递增、i 与 N 的比较以及分支指令。此外，标量操作通常无法充分利用处理器的多个执行单元。

一个展开后的版本可能如下所示：

for (int i = 0; i < N; i += 4) {
    sum += A[i] * B[i];
    sum += A[i+1] * B[i+1];
    sum += A[i+2] * B[i+2];
    sum += A[i+3] * B[i+3];
}

循环展开提供以下三个特定好处：

1. 减少循环开销： 比较和分支指令的执行次数从 $N$ 次减少到 $N/4$ 次。
2. 指令级并行 (ILP)： 现代CPU是超标量和乱序执行的。展开块内的独立指令可以同时分派到不同的算术逻辑单元（ALU）执行。
3. 向量化机会： 编译器可以轻松地将展开的标量操作映射到SIMD向量指令（例如AVX-512或NEON），在一个周期内处理整个块。

优化空间与权衡

循环展开并非单调优化，不是“越多越好”。它在流水线饱和度和寄存器压力之间存在权衡。

1. 寄存器压力： 每个展开的操作可能需要临时寄存器来保存中间值。如果编译器展开过度，CPU的体系结构寄存器会耗尽。这将强制进行“寄存器溢出”，即变量被移到栈（L1缓存/内存）中，这会引起严重的性能损失，抵消循环展开的好处。
2. 指令缓存（I-Cache）： 展开操作会增加程序的二进制大小。如果内层循环变得过大，它可能无法再容纳在指令缓存中，导致流水线停顿，因为CPU需要从主内存中获取指令。

以下图表显示了不同展开因子下的性能特征。请留意随着寄存器压力增加，收益递减以及最终的性能退步现象。

随着循环开销的减少和ILP的提高，归一化吞吐量随展开因子增加。然而，当展开因子超过8后，由于寄存器溢出和指令缓存未命中，性能会下降。

编译器实现策略

在TVM或MLIR等框架中，这些转换方法在降低（lowering）阶段应用。编译器不会随意猜测；它会借助成本模型（在第6章讨论）来选择参数。

密集型内核的典型转换方案包含一个组合策略：

1. 重排循环，以确保最内层循环以步长为1的方式访问最大的张量。
2. 拆分循环（分块），使工作集能放入L1缓存。
3. 展开最内层的分块循环，其因子由目标架构的寄存器数量决定（例如，x86上标准float32的展开因子为4或8）。

通过结合重排以修正内存访问模式，以及展开以使指令流水线饱和，编译器生成了一个接近硬件理论峰值性能的内核。