循环分块与块化

现代深度学习 (deep learning)加速器提供大量的算术吞吐量 (throughput)。然而，实现峰值利用率很少受限于浮点运算单元的速度。相反，瓶颈几乎总是数据从内存移动到处理器的速率。从DRAM获取操作数的成本可能比执行乘法本身的成本高出几个数量级。这种差异通常被称为内存墙。

为减少这种延迟，硬件架构师采用多级缓存（L1、L2、L3）。循环分块，也叫循环阻塞，是使软件执行与此硬件层级对齐 (alignment)的主要软件转换方法。它将循环嵌套的迭代空间划分为更小的块，这些块可以放入特定级别的快速内存中。这能保证数据块一旦载入，就会在被淘汰前尽可能多地重复使用。

局部性机制

考虑一个标准的稠密矩阵乘法 $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}$

在行主序存储格式中，访问模式显示出性能风险：

1. 矩阵 A ($A_{i,k}$): 沿行顺序访问。这表现出良好的空间局部性。
2. 矩阵 B ($B_{k,j}$): 以步长 $N$ 访问。对于内部 $k$ 循环的每次迭代，内存地址会跳转 $N \times \text{sizeof(float)}$。这破坏了空间局部性，并且如果 $N$ 很大，通常会导致每次加载都发生缓存未命中。
3. 矩阵 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$ 块载入缓存，并利用这些数据计算所有必需的部分积，然后再移至下一个块。

下图说明了迭代空间如何被划分为块。灰色区域表示正在处理的活跃块，它必须完全适应目标缓存级别。

将大型矩阵迭代空间分解为更小的块。活跃块（蓝色）代表当前驻留在 L1 缓存中的工作集。

分析成本模型

我们可以通过计算数据重用因子来量化 (quantization)分块的好处。假设一个简化的缓存模型，其中缓存可以容纳内部循环所需的三个 $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}$，就会发生缓存颠簸，性能会急剧下降。

下方图表显示了块大小与有效内存带宽利用率之间的理论关系。请注意，当块大小超出缓存容量时会出现“性能断崖”。

块大小对内存带宽的影响。性能随块大小的增加而提高，因为重用效果更好，直到工作集超出 L1 缓存容量（红色虚线），导致因颠簸而下降。

多级分块和寄存器阻塞

现代CPU和GPU有多层内存（寄存器、L1、L2、VRAM/DRAM）。高级编译器不止步于单一级别的分块。它们执行多级分块以在层级结构的每一级实现最大化重用。

1. L2/L3 分块： 一种粗粒度阻塞策略，将大块（例如 $128 \times 128$）保留在 L2 缓存中。
2. L1 分块： 在 L2 块内部，我们遍历适应 L1 的更小子块（例如 $32 \times 32$）。
3. 寄存器分块（微内核）： 在最内层，编译器展开循环，将一小块（例如 $4 \times 8$）累加器变量保存在 CPU 寄存器中。这可以防止 ALU 在等待 L1 缓存访问时停顿。

TVM中的编译器实现

在 Apache TVM 或 Halide 等深度学习 (deep learning)编译器中，分块作为一种调度原语表达，与算法定义分离。这种分离使得寻找最佳块大小 ($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$ 完全整除时，分块会带来一些复杂情况。这会导致“余数”循环，即矩阵边缘的小块计算。

编译器通过两种方式处理这种情况：

1. 循环分割： 为边界条件生成一个单独的、未优化的循环体。这保持了主内核的代码紧凑性，但增加了二进制文件大小。
2. 填充： 分配略大的输入张量（向上取整到 $T$ 的最近倍数），并用零（或根据操作的单位值）填充额外空间。这使得优化后的内核可以在整个域上统一运行，通常通过避免内核内的分支发散来产生更好的性能。

掌握循环分块，您就定义了硬件消耗数据的粒度。这是将数学定义转换为高性能二进制文件的主要步骤。随后的部分将说明如何对这些块进行矢量化 (quantization)以使用 SIMD 指令。