循环优化空间

优化张量程序实际上是在有效程序变换的高维空间内进行搜索的问题。当深度学习框架将计算图交给编译器时，算子的数学定义（例如卷积或矩阵乘法）是固定的。然而，用于计算该结果的指令具体顺序是灵活的。这种将计算（算什么）与调度（怎么算）分离的做法，是现代张量编译器（如TVM和Halide）的基本准则。

循环优化空间代表了所有可能的调度集合，这些调度要保持原始程序的功能正确性，同时改变其性能特性。在这个空间中进行操作需要平衡相互冲突的硬件制约，主要包括内存带宽、缓存局部性和指令级并行。

将计算与调度分离

为理解优化空间，我们首先明确区分算法及其执行方式。考虑一个标准矩阵乘法 $C = A \times B$，其中 $A \in \mathbb{R}^{M \times K}$ 且 $B \in \mathbb{R}^{K \times N}$。

其逻辑定义包含一个求和归约：

$C_{i,j} = \sum_{k=0}^{K-1} A_{i,k} \times B_{k,j}$

在一个朴素的C++实现中，这个逻辑直接对应于一个三层嵌套循环。

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

虽然在语义上是正确的，但这个顺序（$i, j, k$）仅仅是优化空间中的一个点。它决定了一种特定的内存遍历顺序。如果矩阵$A$和$B$很大，这种特定顺序可能导致$B$频繁的缓存未命中，因为内层循环以非连续方式遍历$B$（假设是行主序存储）。

一个调度定义了应用于这些循环的变换。我们可以将顺序更改为 $i, k, j$，将循环拆分为更小的块，或展开它们。每个组合构成搜索空间中的一个唯一坐标。

搜索空间的维度

优化空间由可组合的调度原语构成。我们可以将这些原语分为三个主要方面，它们定义了生成内核的形态和性能。

1. 迭代顺序（置换）

嵌套循环的顺序决定了内存访问模式。对于深度为$D$的循环嵌套，有 $D!$ 种可能的循环索引置换方式，前提是没有数据依赖阻止重新排序。将顺序从 $i, j, k$ 更改为 $i, k, j$ 会将矩阵$B$的内存访问方式从列向（跨步）变为行向（顺序），显著提高了空间局部性。

2. 循环结构（分块和拆分）

分块将一个单独的循环拆分为两个嵌套循环：一个遍历块的外层循环和一个在块内遍历的内层循环。这种变换增加了优化空间中的维度。一个从 $0$ 到 $N$ 的 $i$ 循环可以由因子 $T_{factor}$ 拆分，从而创建一个范围为 $\lceil N/T_{factor} \rceil$ 的外层循环和一个范围为 $T_{factor}$ 的内层循环。

拆分因子的选择是连续的（受限于循环范围）且组合性的。如果我们将矩阵乘法的所有三个循环进行分块，我们会引入三个新的可调参数。这些参数必须被选择，以使内层循环的工作集能够放入L1或L2缓存。

3. 硬件映射（绑定）

一旦循环被结构化，它们必须映射到硬件执行单元。这包括：

• 向量化： 标记一个内层循环使用SIMD指令（例如AVX-512, NEON）执行。
• 并行化： 将外层循环映射到并行线程（例如CPU上的OpenMP线程，GPU上的线程块）。
• 循环展开： 在编译时完全展开一个循环，以增加指令密度并允许CPU流水线提前查看。

下图描绘了根据这些决策，一个单独的循环定义如何扩展成一个可能的底层实现树。

从抽象逻辑到机器码的过程包含在结构（分块）、排序和硬件映射层面的分支决策。

制约与依赖

优化空间中的每个点并非都有效。编译器必须遵守数据依赖关系，以保持程序的正确性。

如果迭代 $k+1$ 依赖于迭代 $k$ 中产生的数据（写后读依赖），我们不能并行化 $k$ 循环，也不能随意重新排序它。多面体分析通常将这些依赖关系表示为一组线性不等式。一个有效的调度会保留由这些依赖关系决定的事件的部分顺序。

对于深度学习，许多操作，如卷积和矩阵乘法，在某些轴上（例如批次或输出通道维度）是“极易并行”的，但在其他轴上（例如输入通道）是归约操作。确定哪些轴允许并行映射以及哪些需要同步，是循环分析阶段的主要任务。

屋脊线模型与性能目标

我们在这个优化空间中进行搜索，以最大化吞吐量，通常以GFLOPS（每秒十亿次浮点运算）衡量。可达性能的上限由屋脊线模型描述，该模型关联算术强度与硬件限制。

算术强度是执行的浮点运算次数与从DRAM传输的数据字节数之比。

$I = \frac{\text{浮点运算次数}}{\text{传输字节数}}$

如果一个内核具有低算术强度，它就是内存限制型的；其性能受限于内存带宽。分块等变换旨在通过将数据保留在缓存中来增加$I$，有效避免从DRAM重复获取数据。如果一个内核具有高算术强度，它就是计算限制型的；性能受限于处理器的峰值计算能力。向量化和循环展开等变换旨在使计算单元饱和。

内核的性能受限于带宽（倾斜线）或峰值计算（水平线）。循环优化将执行点从内存限制区域（左侧，红色）移向计算限制区域（右侧，绿色）。

组合爆炸

优化空间的庞大程度使得对于每个硬件后端上的每个算子而言，手动调优都不切实际。

考虑一个简单的三层循环嵌套。

1. 分块： 我们可以对每个循环级别进行分块。如果选择块大小为 $\{2, 4, 8, 16, 32\}$，这将产生 $5^3 = 125$ 种组合。
2. 重新排序： 有 $3! = 6$ 种顺序。
3. 循环展开： 我们可能会以内层循环的因子 $\{0, 4, 16, 64\}$ 进行展开。

即使在这个简单的例子中，我们也有数千种潜在调度。在拥有几十个算子以及更深层循环嵌套（例如，带有批次和通道的卷积有7层循环嵌套）的生产网络中，这个空间包含数十亿个候选方案。

此外，性能函数是非凸且不连续的。块大小为32可能完美适配L1缓存，而块大小为33则会引起冲突未命中，导致性能下降一个数量级。这在优化空间中产生了“断崖”。

有效处理这个空间需要自动化策略。尽管我们将在第6章介绍自动调优算法（如遗传算法或模拟退火），但本章的目标是理解变换本身的机制。你必须学习如何手动构建一个有效且高性能的调度，才能理解自动调优器在寻找什么。