内存层级优化：分块和预取

虽然多面体建模提供了一个强大的数学框架，用于变换循环嵌套以展现并行性并抽象地提升数据局部性，但性能的实现取决于对处理器内存层级的有效管理。现代处理器拥有多级缓存（L1、L2、L3），以弥补处理器核心与主内存（DRAM）之间明显的运行速度差异。张量运算处理可能非常大的数据集（权重、激活），经常会遇到内存瓶颈。即使是完全并行化的代码，如果持续等待来自慢速DRAM的数据，也会停滞。

本节审视两种基本的编译器和运行时技术，用于优化张量核内的内存访问模式：分块（也称作缓存阻塞）和软件预取。这些方法旨在最大化缓存内的数据重用，并隐藏无法避免的内存访问延迟。

分块（缓存阻塞）以增强数据局部性

分块的核心原则是增强数据局部性。许多张量运算，比如矩阵乘法（$C = A \times B$）或卷积，在数据重用方面展现出很大的潜力。例如，在矩阵乘法中，矩阵 A 和 B 的元素会被多次访问。然而，如果矩阵很大，标准的循环结构可能会在数据被重用之前就将其从缓存中逐出。

考虑一个标准的矩阵乘法：

// C = A * B 的标准三层嵌套循环
// 矩阵 A (M x K), B (K x N), C (M x N)
for (int i = 0; i < M; ++i) {
    for (int j = 0; j < N; ++j) {
        for (int k = 0; k < K; ++k) {
            C[i][j] += A[i][k] * B[k][j];
        }
    }
}

如果 M、N 和 K 很大，访问 B[k][j] 会为 C 的每个元素遍历 B 的整个列。如果 B 的一列（或在内层循环中访问的 A 的一行）不能完全放入缓存，元素将反复从主内存中获取，从而导致性能不佳。

分块通过增加外层循环来变换循环嵌套，这些外层循环遍历原始迭代空间的“块”或“单元”。内层循环随后只处理当前块内的数据。块的大小经过选择，使得工作集（块计算所需的数据）能够轻松放入特定的缓存层级（例如 L1 或 L2）。

分块矩阵乘法：

// 分块矩阵乘法 (块大小 T_M, T_N, T_K)
for (int i0 = 0; i0 < M; i0 += T_M) {
    for (int j0 = 0; j0 < N; j0 += T_N) {
        for (int k0 = 0; k0 < K; k0 += T_K) {
            // 内层循环处理一个块
            for (int i = i0; i < min(i0 + T_M, M); ++i) {
                for (int j = j0; j < min(j0 + T_N, N); ++j) {
                    for (int k = k0; k < min(k0 + T_K, K); ++k) {
                        C[i][j] += A[i][k] * B[k][j];
                    }
                }
            }
        }
    }
}

这里，T_M、T_N 和 T_K 是块的尺寸。内层循环现在使用 A 的 T_M x T_K 子块和 B 的 T_K x T_N 子块计算 C 的 T_M x T_N 子块。

从重复访问完整矩阵到访问具有高数据重用性的缓存友好块的转换。

分块的影响：

• 时间局部性： 通过在处理完一个 A 和 B 的块后再处理下一个，元素在它们可能仍驻留在缓存中时被多次重用。
• 空间局部性： 块内的访问通常是连续的或步长为1的，从而更有效地使用缓存行。

挑战和考量：

• 块大小选择： 这是一个重要的调整参数。最佳块大小很大程度上取决于缓存大小、关联性、缓存行大小、数据类型大小以及特定的处理器架构。选择略小于缓存容量的大小是很常见的，以避免冲突失效。多面体编译器通常包含分析模型或自动调整机制来确定合适的块大小。
• 多级分块： 对于具有多级缓存（L1、L2、L3）的架构，可以应用多级分块。循环被递归地分块，外层块用于 L3，中间块用于 L2，内层块（或寄存器阻塞）用于 L1/寄存器。
• 寄存器分块： 最内层的分块通常旨在将数据保持在处理器寄存器中，这提供了最快的访问。这有时被称为寄存器阻塞。
• 与并行化的关联： 分块可以与并行化策略（例如 OpenMP、线程化）关联。块本身可以成为并行工作的单元。

多面体框架擅长表示循环嵌套中涉及的迭代空间和依赖关系。基于多面体模型（如 ISL、Pluto、Polly）的工具能够系统地分析依赖关系，并应用复杂的分块变换，包括找到有效且可能最佳的块大小和循环置换，其可靠性远超针对复杂嵌套的手动启发式方法。

软件预取以隐藏延迟

即使进行了有效的分块，对当前不在缓存中的数据（例如加载下一个块）的访问仍然会产生主内存延迟。软件预取是一种技术，编译器会插入特殊指令（预取指令，例如 x86 上的 _mm_prefetch，ARM 上的 prfm），以便在加载或存储指令实际需要数据之前从内存中请求数据到缓存。此举旨在使数据传输时间与正在进行的计算重叠，从而有效隐藏延迟。

机制：

1. 识别未来访问： 编译器分析循环的内存访问模式，以预测在未来迭代中将需要的地址。对于像 A[i+10] 这样的简单仿射访问，这很直接。
2. 计算预取距离： 确定预取指令应提前多少次迭代。这个“预取距离”很重要：
   • 太短：数据在需要之前未到达；延迟未被隐藏。
   • 太长：预取的数据可能过早到达并在使用前被从缓存中逐出（缓存污染），或者可能取代其他有用数据。最佳距离取决于内存延迟和每次循环迭代的计算量。
3. 插入预取指令： 将适当的预取指令注入循环体，目标是计算出的未来地址。预取指令通常是对硬件的提示，如果地址无效，它们不会引起异常。

示例（简单循环中的预取）：

// 原始循环
for (int i = 0; i < N; ++i) {
    result += A[i] * B[i];
}

// 带有软件预取的循环
int prefetch_distance = 16; // 示例距离（调整参数）
for (int i = 0; i < N; ++i) {
    // 为迭代 i + prefetch_distance 预取数据
    if (i + prefetch_distance < N) {
         _mm_prefetch(&A[i + prefetch_distance], _MM_HINT_T0); // 将 A 元素预取到 L1/L2
         _mm_prefetch(&B[i + prefetch_distance], _MM_HINT_T0); // 将 B 元素预取到 L1/L2
    }

    // 迭代 i 的实际计算
    result += A[i] * B[i];
}

挑战和考量：

• 硬件预取器的关联： 现代 CPU 具有硬件预取器，可自动检测简单的访问模式（如顺序步长）并获取数据。软件预取需要补充而非冲突硬件预取器。有时软件预取针对硬件无法检测的复杂模式，或者使用不同的提示（例如预取到 L2 而不是 L1）。
• 地址计算开销： 计算未来地址会向循环体添加指令。此开销必须小于节省的延迟。
• 复杂访问模式： 预取间接访问（例如 A[index[i]]）明显更难，因为未来地址取决于一个可能也需要加载的值（index[i]）。
• 调优： 预取距离和特定的预取提示（表示要针对哪个缓存级别，例如 _MM_HINT_T0、_MM_HINT_T1、_MM_HINT_T2）是依赖于架构的调优参数。

分块与预取之间的关联

分块和预取通常结合使用以获得最大效益。分块通过确保工作集适合缓存来改进局部性，从而减少强制缓存失效（首次访问数据时的失效）的数量。预取有助于隐藏剩余失效的延迟，特别是那些在块之间移动或获取块的初始数据时发生的失效。分块变换所创建的块内部的常规访问模式，通常使得编译器（或硬件）更容易有效地预取该块内部后续迭代所需的数据。

通过分块和预取来优化内存访问，对于在具有复杂内存层级的硬件上执行的计算密集型 ML 核实现高性能是必不可少的。虽然多面体方法为循环变换提供了理论依据，但这些技术使这些变换基于缓存大小和内存延迟的物理现实，构成了抽象优化与加速之间的一个重要桥梁。高级 ML 编译器投入大量精力，根据目标硬件特性自动应用和调整这些内存层级优化。