高效的内核执行,依赖于软件指令与硬件物理特性的一致性。循环分块虽能优化缓存容量,但它本身不能解决数据如何从缓存行中取出,或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的访问与跨步访问之间在内存遍历上的差异。digraph G { rankdir=TB; node [shape=box, style=filled, fontname="Arial"]; subgraph cluster_memory { label="线性内存布局(行主序 4x4)"; bgcolor="#f8f9fa"; color="#dee2e6"; struct1 [label="<f0> (0,0)|<f1> (0,1)|<f2> (0,2)|<f3> (0,3)|<f4> (1,0)|<f5> (1,1)|<f6> (1,2)|<f7> (1,3)|<f8> ...", fillcolor="#e9ecef"]; } subgraph cluster_access { label="访问模式"; bgcolor="#ffffff"; color="#ffffff"; node [shape=ellipse, style=filled]; stride1 [label="步长1(最佳)", fillcolor="#d0bfff"]; strided [label="跨步(次优)", fillcolor="#ffc9c9"]; } stride1 -> struct1:f0 [color="#7048e8", penwidth=2]; stride1 -> struct1:f1 [color="#7048e8", penwidth=2]; stride1 -> struct1:f2 [color="#7048e8", penwidth=2]; strided -> struct1:f0 [color="#fa5252", penwidth=2, style=dashed]; strided -> struct1:f4 [color="#fa5252", penwidth=2, style=dashed]; strided -> struct1:f8 [color="#fa5252", penwidth=2, style=dashed]; }内存访问模式决定了缓存效率。步长为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]; } } }在此配置中:$C[i][j]$ 在内层循环中不变(好)。$A[i][k]$ 以步长为1的方式访问(好)。$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$ 的置换中:$A[i][k]$ 在内层循环中不变。$B[k][j]$ 以步长为1的方式访问(好)。$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]; }循环展开提供以下三个特定好处:减少循环开销: 比较和分支指令的执行次数从 $N$ 次减少到 $N/4$ 次。指令级并行 (ILP): 现代CPU是超标量和乱序执行的。展开块内的独立指令可以同时分派到不同的算术逻辑单元(ALU)执行。向量化机会: 编译器可以轻松地将展开的标量操作映射到SIMD向量指令(例如AVX-512或NEON),在一个周期内处理整个块。优化空间与权衡循环展开并非单调优化,不是“越多越好”。它在流水线饱和度和寄存器压力之间存在权衡。寄存器压力: 每个展开的操作可能需要临时寄存器来保存中间值。如果编译器展开过度,CPU的体系结构寄存器会耗尽。这将强制进行“寄存器溢出”,即变量被移到栈(L1缓存/内存)中,这会引起严重的性能损失,抵消循环展开的好处。指令缓存(I-Cache): 展开操作会增加程序的二进制大小。如果内层循环变得过大,它可能无法再容纳在指令缓存中,导致流水线停顿,因为CPU需要从主内存中获取指令。以下图表显示了不同展开因子下的性能特征。请留意随着寄存器压力增加,收益递减以及最终的性能退步现象。{"layout": {"title": {"text": "性能与展开因子", "font": {"family": "Arial", "size": 18, "color": "#343a40"}}, "xaxis": {"title": {"text": "展开因子", "font": {"family": "Arial", "color": "#495057"}}, "tickvals": [1, 2, 4, 8, 16, 32], "gridcolor": "#e9ecef"}, "yaxis": {"title": {"text": "归一化吞吐量", "font": {"family": "Arial", "color": "#495057"}}, "gridcolor": "#e9ecef"}, "plot_bgcolor": "#ffffff", "paper_bgcolor": "#ffffff", "width": 600, "height": 400, "margin": {"l": 60, "r": 30, "t": 50, "b": 50}}, "data": [{"type": "bar", "x": [1, 2, 4, 8, 16, 32], "y": [1.0, 1.8, 2.5, 2.7, 2.2, 1.4], "marker": {"color": ["#adb5bd", "#a5d8ff", "#4dabf7", "#228be6", "#ff8787", "#fa5252"]}}]}随着循环开销的减少和ILP的提高,归一化吞吐量随展开因子增加。然而,当展开因子超过8后,由于寄存器溢出和指令缓存未命中,性能会下降。编译器实现策略在TVM或MLIR等框架中,这些转换方法在降低(lowering)阶段应用。编译器不会随意猜测;它会借助成本模型(在第6章讨论)来选择参数。密集型内核的典型转换方案包含一个组合策略:重排循环,以确保最内层循环以步长为1的方式访问最大的张量。拆分循环(分块),使工作集能放入L1缓存。展开最内层的分块循环,其因子由目标架构的寄存器数量决定(例如,x86上标准float32的展开因子为4或8)。通过结合重排以修正内存访问模式,以及展开以使指令流水线饱和,编译器生成了一个接近硬件理论峰值性能的内核。