循环优化空间

优化张量程序实际上是在有效程序变换的高维空间 (high-dimensional space)内进行搜索的问题。当深度学习 (deep learning)框架将计算图交给编译器时，算子的数学定义（例如卷积或矩阵乘法）是固定的。然而，用于计算该结果的指令具体顺序是灵活的。这种将计算（算什么）与调度（怎么算）分离的做法，是现代张量编译器（如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}$ 的内层循环。

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

3. 硬件映射（绑定）

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

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

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

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

制约与依赖

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

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

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

屋脊线模型与性能目标

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

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

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

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

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

组合爆炸

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

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

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

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

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

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