现代深度学习加速器提供大量的算术吞吐量。然而,实现峰值利用率很少受限于浮点运算单元的速度。相反,瓶颈几乎总是数据从内存移动到处理器的速率。从DRAM获取操作数的成本可能比执行乘法本身的成本高出几个数量级。这种差异通常被称为内存墙。为减少这种延迟,硬件架构师采用多级缓存(L1、L2、L3)。循环分块,也叫循环阻塞,是使软件执行与此硬件层级对齐的主要软件转换方法。它将循环嵌套的迭代空间划分为更小的块,这些块可以放入特定级别的快速内存中。这能保证数据块一旦载入,就会在被淘汰前尽可能多地重复使用。局部性机制考虑一个标准的稠密矩阵乘法 $C = A \times B$,其中所有矩阵的尺寸均为 $N \times N$。一个简单的实现通常包含三个嵌套循环:$$ \text{对于 } i \in [0, N): \text{ 对于 } j \in [0, N): \text{ 对于 } k \in [0, N): C_{i,j} \ += A_{i,k} \times B_{k,j} $$在行主序存储格式中,访问模式显示出性能风险:矩阵 A ($A_{i,k}$): 沿行顺序访问。这表现出良好的空间局部性。矩阵 B ($B_{k,j}$): 以步长 $N$ 访问。对于内部 $k$ 循环的每次迭代,内存地址会跳转 $N \times \text{sizeof(float)}$。这破坏了空间局部性,并且如果 $N$ 很大,通常会导致每次加载都发生缓存未命中。矩阵 C ($C_{i,j}$): 在内部循环中持续访问,表现出优秀的局部性。当 $N$ 足够大时,工作集大小会超出缓存容量。当程序完成 $B$ 的一行并需要重用 $A$ 的元素时,包含 $A$ 的缓存行很可能已被淘汰。循环分块转换分块通过引入新循环来按块遍历数据,以解决此问题。我们将原始迭代空间 $(i, j, k)$ 分为遍历块的外部循环 $(i_{outer}, j_{outer}, k_{outer})$ 和在块内迭代的内部循环 $(i_{inner}, j_{inner}, k_{inner})$。如果我们选择块大小为 $T$,转换后的逻辑如下:// 遍历块 for (int i_o = 0; i_o < N; i_o += T) { for (int j_o = 0; j_o < N; j_o += T) { for (int k_o = 0; k_o < N; k_o += T) { // 执行当前块(适应缓存) for (int i_i = i_o; i_i < min(i_o + T, N); i_i++) { for (int j_i = j_o; j_i < min(j_o + T, N); j_i++) { for (int k_i = k_o; k_i < min(k_o + T, N); k_i++) { C[i_i][j_i] += A[i_i][k_i] * B[k_i][j_i]; } } } } } }这种重构改变了操作顺序,但保持了数学上的等价性。处理器现在将 $A$ 和 $B$ 的一个 $T \times T$ 块载入缓存,并利用这些数据计算所有必需的部分积,然后再移至下一个块。下图说明了迭代空间如何被划分为块。灰色区域表示正在处理的活跃块,它必须完全适应目标缓存级别。digraph G { rankdir=TB; bgcolor="#ffffff"; node [style=filled, shape=rect, fontname="Arial", fontsize=12]; edge [fontname="Arial", fontsize=10, color="#adb5bd"]; subgraph cluster_matrix { label="矩阵迭代空间 (N x N)"; style=dashed; color="#adb5bd"; fontcolor="#495057"; // 块网格 node [color="#dee2e6", fillcolor="#f8f9fa", width=1.2, height=1.2, label=""]; t00 [pos="0,2!"]; t01 [pos="1.5,2!"]; t02 [pos="3,2!"]; t10 [pos="0,0.5!"]; t11 [fillcolor="#a5d8ff", label="活跃块\n(适应L1)"]; t12 [pos="3,0.5!"]; t20 [pos="0,-1!"]; t21 [pos="1.5,-1!"]; t22 [pos="3,-1!"]; edge [style=invis]; t00 -> t01 -> t02; t10 -> t11 -> t12; t20 -> t21 -> t22; t00 -> t10 -> t20; } memory [shape=cylinder, fillcolor="#e9ecef", label="主内存", width=1.5]; cache [shape=box, fillcolor="#b2f2bb", label="L1 缓存", width=1.5]; memory -> cache [label="载入块", color="#1c7ed6", penwidth=2]; cache -> t11 [label="高重用", color="#1c7ed6", penwidth=2]; }将大型矩阵迭代空间分解为更小的块。活跃块(蓝色)代表当前驻留在 L1 缓存中的工作集。分析成本模型我们可以通过计算数据重用因子来量化分块的好处。假设一个简化的缓存模型,其中缓存可以容纳内部循环所需的三个 $T \times T$ 块。不进行分块时,矩阵 $A$ 的内存传输总次数(假设严格执行的LRU淘汰策略和较大的 $N$)接近 $N^3$,因为 $A$ 的行会为 $B$ 的每一列重新获取。通过分块,要计算 $C$ 的一个 $T \times T$ 块,我们需要遍历 $A$ 和 $B$ 的 $N/T$ 个块。$C_{tile}$ 块的大小:$T^2$(保留在寄存器/L1中)。$A$ 的总载入量:我们载入整个矩阵 $A$。由于我们为 $C$ 的不同 $j$ 块重用 $A$ 块,因此 $A$ 的全局流量为 $(N/T) \times N^2$。$B$ 的总载入量:同理,$(N/T) \times N^2$。算术密度(每字节浮点运算数)随块大小 $T$ 线性提高:$$ \text{运算强度} \approx \frac{2N^3 \text{ (总运算量)}}{2 \frac{N^3}{T} \text{ (总字节数)}} = T $$较大的块意味着更好的带宽效率,前提是它们能适应缓存。如果选择 $T$ 使 $3T^2 \times \text{sizeof(float)} > \text{CacheSize}$,就会发生缓存颠簸,性能会急剧下降。下方图表显示了块大小与有效内存带宽利用率之间的理论关系。请注意,当块大小超出缓存容量时会出现“性能断崖”。{"layout": {"title": "有效带宽与块大小", "xaxis": {"title": "块大小(元素数量)", "showgrid": true}, "yaxis": {"title": "有效带宽(GB/秒)", "showgrid": true}, "plot_bgcolor": "#f8f9fa", "paper_bgcolor": "#ffffff", "font": {"family": "Arial, sans-serif"}}, "data": [{"x": [8, 16, 32, 64, 128, 256, 512], "y": [100, 180, 280, 350, 50, 40, 30], "type": "scatter", "mode": "lines+markers", "line": {"color": "#1c7ed6", "width": 3, "shape": "spline"}, "marker": {"size": 8, "color": "#228be6"}, "name": "L1 缓存性能"}, {"x": [64, 64], "y": [0, 350], "type": "line", "mode": "lines", "line": {"color": "#fa5252", "dash": "dash", "width": 2}, "name": "L1 容量限制"}]}块大小对内存带宽的影响。性能随块大小的增加而提高,因为重用效果更好,直到工作集超出 L1 缓存容量(红色虚线),导致因颠簸而下降。多级分块和寄存器阻塞现代CPU和GPU有多层内存(寄存器、L1、L2、VRAM/DRAM)。高级编译器不止步于单一级别的分块。它们执行多级分块以在层级结构的每一级实现最大化重用。L2/L3 分块: 一种粗粒度阻塞策略,将大块(例如 $128 \times 128$)保留在 L2 缓存中。L1 分块: 在 L2 块内部,我们遍历适应 L1 的更小子块(例如 $32 \times 32$)。寄存器分块(微内核): 在最内层,编译器展开循环,将一小块(例如 $4 \times 8$)累加器变量保存在 CPU 寄存器中。这可以防止 ALU 在等待 L1 缓存访问时停顿。TVM中的编译器实现在 Apache TVM 或 Halide 等深度学习编译器中,分块作为一种调度原语表达,与算法定义分离。这种分离使得寻找最佳块大小 ($T$) 的过程可以自动化。例如,在 TVM 的 te(张量表达式)语言中,分块调度应用于计算的轴:# 定义计算 A = te.placeholder((N, N), name='A') B = te.placeholder((N, N), name='B') k = te.reduce_axis((0, N), name='k') C = te.compute((N, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k)) # 创建调度 s = te.create_schedule(C.op) # 应用分块 # 将 i 和 j 轴按因子 bn 和 bm 分割 bn, bm = 32, 32 xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bm) # 重新排序,先外部循环,再规约,最后内部循环 k_outer, k_inner = s[C].split(k, factor=4) s[C].reorder(xo, yo, k_outer, xi, yi, k_inner)tile 原语有效地分割轴并重新排序它们。随后的 reorder 命令作用很大;它确保遍历缓存驻留块 (xi, yi) 的循环是最内层的空间循环。约束和填充当矩阵维度 $N$ 不能被块大小 $T$ 完全整除时,分块会带来一些复杂情况。这会导致“余数”循环,即矩阵边缘的小块计算。编译器通过两种方式处理这种情况:循环分割: 为边界条件生成一个单独的、未优化的循环体。这保持了主内核的代码紧凑性,但增加了二进制文件大小。填充: 分配略大的输入张量(向上取整到 $T$ 的最近倍数),并用零(或根据操作的单位值)填充额外空间。这使得优化后的内核可以在整个域上统一运行,通常通过避免内核内的分支发散来产生更好的性能。掌握循环分块,您就定义了硬件消耗数据的粒度。这是将数学定义转换为高性能二进制文件的主要步骤。随后的部分将说明如何对这些块进行矢量化以使用 SIMD 指令。