• 文档
• 教程
• 人工智能

本篇技术指南聚焦 AISystem 项目中编译后端模块的算子循环优化（Loop Optimization）。在 AI 芯片实际执行计算时，算子本质上是大量嵌套 for 循环对存储数据的重复读写与重复指令（SIMT/SIMD），循环优化的核心目标是提升数据局部性与计算并行性，从而提升整体算子性能。读完本篇，你将掌握循环分块、循环展开、循环重排、循环融合、循环拆分五类变换的原理、Cache miss 推导方法与应用边界，并可对照仓库中的示例代码与 TVM 调度实践深入验证。

为什么算子需要循环优化

在具体硬件执行计算的时候，实际会大量地使用for等循环指令不断地去读取不同的数据执行重复的指令（SIMT/SIMD）。因此循环优化主要是为了提升数据的局部性或者计算的并行性，从而提升整体算子性能——当然，这二者都需要 AI 芯片硬件的支持。

从硬件视角看，AI 算子（如卷积、矩阵乘）天然具有高度规则化的多层嵌套循环结构，这为优化提供了丰富的技术手段。以逐元素操作为例，如 ReLU、加法（Add）和乘法（Mul）等，可以通过在所有循环轴上进行迭代来执行计算；即使是较为复杂的卷积（Conv），也可以通过七层嵌套循环实现。然而，如果仅仅采用直观的原生计算方法，往往会导致效率低下。在 算子手工优化 一节中，我们已经借助 Roofline 模型分析了算子的计算瓶颈（Compute-Bound）与访存瓶颈（Memory-Bound），而循环优化正是针对访存与并行性两大瓶颈的核心调度手段。

循环优化的两大挑战

数据局部性：Cache 容量有限，命中率决定访存开销

数据的局部性与计算机存储层次有关。计算机拥有速度由慢到快、容量由大到小的多层次存储器，以最优的控制调度算法和合理的成本构成存储系统：

• 寄存器位于最上层（一般在 CPU 芯片内），直接参与运算，速度最快但价格最高、容量最小；
• 主存存放运行中的程序与数据，但其速度与 CPU 差距较大；
• Cache（高速缓冲存储器）插入在主存与 CPU 之间，比主存速度更快、容量更小，用于匹配二者的速度差异。

主存与 Cache 之间的数据迁移由硬件自动完成，对程序员来说是透明的、无法直接编程控制的。在现代多核 CPU 架构中，Cache 分为 L1、L2、L3 三级，对计算机运算速度影响最大的是L3 Cache，因为该级 Cache 的不命中将导致片外访存。L3 Cache 的大小一般在几 MB 至几十 MB 之间，容量较小，因此无法把所有的数据都放到 Cache 里——Cache 里应该放 CPU 最有可能会对它进行处理的数据。

这里有两个著名的程序运行时局部性原理，诠释了什么样的数据很可能会被 CPU 处理：

• 时间局部性：CPU 处理了一个数据以后，它很有可能会对它进行第二次处理。
• 空间局部性：CPU 处理了内存中某一块的数据，它很可能会还对它附近的数据块进行读写操作。

这两个局部性原理十分匹配循环的特点。以两层嵌套循环为例：

for i in range(n): for j in range(m): A[i] += B[j] # 这个计算可能并没有什么意义，只是为了举例

• 时间局部性：A[i]在内层循环（J 循环）完整结束一次后被重用，每一个 A 中的元素会被重用 m 次；B[j]在每一次内层循环都改变，在 i 循环时重用，每一个 B 中的元素会被重用 n 次。
• 空间局部性：A 中的元素在 m 次后切换到下一个元素；B 中的元素每次都会切换到下一个。这两个数组具有不同的模式——A 的时间局部性更明显，B 的空间局部性更明显。

为什么要分析数据的局部性？因为 Cache 容量小，一次只能容纳一部分数据。对 Cache 而言，如果 CPU 要读写的数据在 Cache 中，则称为命中，可以直接在 Cache 这一级别返回数据；如果不在 Cache 中，则未命中，需要去下一级存储取数据，增加了 IO 的开销。

计算并行性：多核线程并行与 SIMD 向量化

现代 CPU 通常是多核结构，可以进行线程级并行。在单核 CPU 上运行多任务时，通常是模拟进行的并行，通过进程调度算法（如时间片轮转）来同时进行多任务，实际上 CPU 在一个时刻只进行了一个任务。

在多核 CPU 上，每一个核都可以进行计算。搭配超线程技术（Intel CPU）时，一个核可以执行两个线程。通过将多个计算线程分配到多个核，可以同时执行多线程计算实现并行加速，这是 CPU 上最有效的优化方式。在 Windows 中可以通过任务管理器查看内核与逻辑处理器数量——逻辑处理器就是用超线程技术模拟的，能让处理器在同一时间段内同时处理多个线程，从而实现更好的多任务处理能力。

向量化则是一种数据级并行优化。向量化即"批量操作"，在计算机中常见执行模型是单指令多数据（SIMD，Single Instruction Multiple Data）——通过对批量数据同时进行相同计算以提高效率。向量体系结构获取在存储器中散布的数据集，将多个数据元素放在大型的顺序寄存器堆（即向量寄存器）中，对整个寄存器进行操作，从而同时计算了多个数据元素。向量本身可以容纳不同大小的数据，因此如果一个向量寄存器可以容纳 64 个 64 bit 元素，那么也可以容纳 128 个 32 bit 元素或者 512 个 8 bit 元素。凭借这种硬件上的多样性，向量化特别适合用于多媒体应用和科学计算。

传统的执行方式为单指令单数据（SISD，Single Instruction Single Data），硬件不支持并行计算。现代 CPU 几乎都支持 SIMD 指令集，如 Intel 的 SSE（Streaming SIMD Extensions）和 AVX（Advanced Vector Extensions）系列指令集。

五类循环优化方案详解

循环的优化方案针对不同的数据局部性和计算并行性，主要有循环分块、循环展开、循环重排、循环融合、循环拆分等。仓库在 03Compiler/04Backend/code 目录下提供了 tiling、fusion、reorder、spliting 四份可直接对照的示例文件，下面逐一展开。

循环分块（Loop Blocking / Tiling）

循环分块是利用 Cache 的数据局部性进行优化的一种方法。现代 CPU 通常具有多级 Cache，Cache 是除 CPU 寄存器外最接近 CPU 的存储层次，相比主存速度更快但容量更小。Cache 中复制有 CPU 频繁使用的数据以进行快速访问。由于 Cache 容量有限，数据会在 Cache 中进行换入换出：当访问的数据在 Cache 中没有时，产生Cache miss，会向低一级存储层次发出访问请求，然后该数据存储进 Cache，访问时间大大提高；当访问数据就在 Cache 中时，会直接使用该数据以进行复用。

循环分块主要针对大型数据集进行优化——大数据集无法一次全部存入 Cache。当遍历该数据集时，循环按照顺序进行访问，会替换掉之前加载进 Cache 的数据，导致后面的指令对之前的数据无法复用，要重新加载数据，产生大量 Cache miss，数据复用性很差，程序执行时间变长，大量时间花费在载入数据上。

循环分块将大数据集分成多个小块以充分进行数据复用。数据块的内存访问是一个具有高内存局部性的小邻域，该数据块可以一次加载进 Cache，执行完所有或者尽可能多的计算任务后才被替换出。在实现中，将一层内层循环分成outer loop * inner loop，然后把 outer loop 移到更外层去，从而确保 inner loop 一定能满足 Cache。原始与分块后的数据存储访问模式如下：

以如下代码为例分析其 Cache 利用率：

for i in range(n): for j in range(m): A[i] += B[j]

假设 m 和 n 是很大的数（大数组），那么对于每一轮 I 循环，等到B[m-1]被访问时，B[1]、B[2]等已经被清出缓存了。假设每个 Cache line 可以容纳 b 个数组元素，以全相联的方式管理，则 A 的 Cache miss 是n/b，B 的 Cache miss 是n*m/b。

如果把 j 层循环进行 tile，分为j_o和j_i两层循环，并把j_o提到最外面，循环就变成了：

for j_o in range(0, m, T): for i in range(n): for j_i in range(j_o, min(j_o + T, m)): A[i] += B[j_i]

当 T 个 B 中的元素可以放进 Cache 时，在 i 层循环即使 i 切换了，但 B 的数据此时还在 Cache 中，只有j_o循环变换时 B 中的数据才会发生 Cache miss。对于 B 来说 Cache miss 为m/T * T/b = m/b；但对于 A 来说，Cache miss 变为了m/T * n/b = nm / Tb，反而增加了。那么再对 i 层循环进行 tile，最终变为：

for i_o in range(0, n, W): for j_o in range(0, m, T): for i_i in range(i_o, min(i_o + W, n)): for j_i in range(j_o, min(j_o + T, m)): A[i_i] += B[j_i]

假设 W 个 A 中的元素也能一次性放入 Cache：B 的 Cache miss 没变，A 的 Cache miss 变为n/W * W/b = n/b。仓库中的示例文件 tiling 正是这一变换的 C 风格伪代码实现：

for j = 0, n for i = 1, m A(i) += B(j) endfor endfor # After Tiling for j_o = 0, m, T: for i = 0, n: for j_i = j_o, j_o + T: A[i] += B[j_i] endfor endfor endfor

从公式上推导，似乎只要 T 和 W 满足能放进 Cache 的数量，那么 Cache miss 与 T、W 就无关，不同 T、W 的性能应该是一样的——但实际运行并不是这样。tile 大小的选择受到 Cache line 大小、Cache 关联度、数据替换策略、硬件存储架构等多个因素的共同影响：分块过大时数据尚未充分利用 Cache 的重用就被替换出去从而导致 miss；分块过小又会造成较大的成本开销，从而掩盖带来的性能优势。tile 大小的选择目前还没有一个确定的算法，目前流行的方法有：

• 基于分析的方法：早期有研究者对嵌套循环和硬件存储特征进行静态分析，为编译器选择合适的分块大小。这类方法主要用来解决容量失效、自干扰失效、交叉干扰失效导致的 Cache 不命中和局部性优化问题。其缺陷在于：1）理论分析不能完全反应实际存储的复杂过程，影响程序分块性能的因素很多，导致实际最优分块的性能和分析方法选择的分块性能差距较大；2）分析建模过程复杂、成本较高，对硬件和程序布局依赖严重，不具有通用性。
• 基于经验搜索的方法：将循环嵌套看作一个黑盒，根据经验选择一系列不同分块大小的组合，在目标机器上对这些分块组合进行自动调优，并从较好的分块大小组合中选取性能最优的分块大小。其缺陷在于对多层循环进行分块时要遍历庞大的搜索空间，导致时间成本过高。

值得说明的是，分块思想在现代 AI 编译器中已通过调度 API 高度工程化。在 TVM 开发算子 的示例中，cfg.define_split("tile_b", b, num_outputs=2)、cfg.define_split("tile_ci", c_i, num_outputs=2)等调度原语正是把"分块大小"定义为可搜索的配置项，交由自动调优器在目标机器上搜索最优组合——这正是上文"基于经验搜索方法"的 DSL 化落地。

循环展开（Loop Unrolling）

循环展开将一个循环中的多次迭代展开成多个单独的迭代，以减少程序执行的开销，提高代码的运行效率。在计算机执行程序的流水线中，每次跳转到循环体内部都需要进行额外的指令处理和跳转操作，这会增加程序的开销。而循环展开可以通过减少跳转次数、减少指令处理次数等方式，降低指令分支预测的开销来优化程序性能，提高程序执行的速度。

通常循环展开包含以下几个步骤：

1. 复制循环体 n 次，使展开后循环体包括原来循环体 n 个拷贝（这里的 n 一般被称为展开因子）；
2. 调整数组索引变量的增量以及内存读写操作的地址；
3. 删除除了最后一个循环体外的所有循环体中的循环计数变量的增加操作，并且修改最后一个循环体中循环计数变量的增量为原来的 n 倍；
4. 删除除了最后一个循环体外的所有循环中的循环条件判断语句。

例如原始循环：

for i in range(m): a[i] += b[i]

通过循环展开（展开因子 4），可以将其转换为以下形式：

for i in range(0, m-3, 4): a[i] += b[i] a[i+1] += b[i+1] a[i+2] += b[i+2] a[i+3] += b[i+3] for i in range(m-3, m): a[i] += b[i]

在展开后的循环中，原本执行了 n 次循环迭代，变成了执行n/4次循环展开。

从优化收益上分析，循环展开不仅可以减少循环开销（如循环变量测试及分支语句等），还提高了指令之间的并发度，并且因为减少了分支语句从而减少流水线停顿，提升了流水线效率。另一个角度是循环展开后可能会为其他优化（如指令级并行、向量化）提供更多机会。但循环展开也可能带来负面效果：

• 如果展开后循环体超过指令缓存容量，会引起缓存失效，造成程序性能下降；
• 循环展开会增加寄存器压力，可能导致生成更多的寄存器溢出处理操作，从而降低优化效果。

循环展开最关键的是确定展开因子，目前主要有三种方法：

1. 启发式方法：对循环体代码进行分析，然后使用静态模型计算展开因子。分析时需要考虑循环本身减少的循环开销、循环展开与其他优化的交互关系等，建立模型要充分考虑指令级并行度、流水线效率、数据局部性、指令缓存与寄存器的限制等。
2. 机器学习方法：根据循环的特征将循环分类，通过大量样本学习，使用分类器建立循环类型和展开因子之间的映射，在实际优化循环时根据循环类型确定最优展开因子。
3. 迭代编译：使用不同展开因子编译生成多个版本的程序并实际运行，选取运行时间最短的作为最优展开因子。

比较三个方法：启发式方法开销最小，展开因子的选择依赖于静态模型的准确性；机器学习开销次之，展开因子的选择不仅依赖于提取的循环特征，还需要大量样本进行训练；迭代编译开销最大，但在不考虑开销的情况下肯定可以找到最优展开因子。

循环重排（Loop Reorder）

循环重排序是矩阵乘法常见的优化方式，指的是对程序中的循环结构重新排列顺序，以优化数据访问模式，特别是在 CNN 中卷积层的应用。通过改变循环的嵌套顺序或者循环内部的迭代顺序，可以改善数据的局部性，减少缓存失效。

在矩阵乘法计算中，B 是逐列访问的，在行优先的存储模式下访问模式很不友好。切换内层的循环顺序可以使得所有元素按顺序读取和写入——一次计算输出的一行，得到的是中间结果，全部累加即可得到结果矩阵的一行最终结果，这种方式利用的是内存的空间局部性。

仓库中的示例文件 reorder 展示了最直观的重排变换——交换两层循环的顺序：

for i = 1, n for j = 1, m A(i,j) = B(i, j) * C(i, j) endfor endfor # After Reorder for j = 1, m for i = 1, n A(i, j) = B(i, j) * C(i, j) endfor endfor

需要注意的是，循环重排必须保证变换前后循环语义（迭代空间）等价。在 TVM 开发算子 的调度示例中，s[output].reorder(x_bo, x_co, x_bi, x_ci)正是将分块后的输出通道循环重排到更外层，配合compute_at将计算阶段锚定到合适的存储层级，实现数据局部性与并行性的协同优化。

循环融合（Loop Fusion）

循环融合用于将多个循环合并为一个更大的循环，将相邻或紧密间隔的循环融合在一起。通过合并多个循环，可以减少程序中的循环次数，从而减少循环开销；合并循环可以减少内存访问次数，提高数据局部性，减少缓存未命中的可能性，从而提高程序执行效率。

以下是一个简单的循环融合示例。两个独立的循环：

# 独立的循环 for i in range(len(a)): a[i] = b[i] + x for i in range(len(b)): d[i] = a[i] + y

在第一个循环中，a 的值被依次写入，在第二个循环中又被马上读取。当数组非常大时，在第二个循环要读取a[0]时，a[0]早已因为 Cache 容量的限制而被清除，需要从下一级存储中读取。

通过循环融合，可以将这两个循环合并为一个循环：

# 循环融合 for i in range(len(a)): a[i] = b[i] + x d[i] = a[i] + y

这样在第二个对 a 的读取语句执行的时候，a 的元素还在 Cache 中。除了这种数据局部性的收益，循环融合还减少了对分支跳转指令的生成。仓库中的示例文件 fusion 展示了更完整的形式——将两个循环体合并后，边界上的首尾迭代单独处理：

for i = 0, n A(i) = a(i) + b(i); c(i) = 2 * a(i); endfor for i = 1, n - 1 D(i) = c(i) + a(i); endfor # After fusion A(0) = a(0) + b(0); c(0) = 2 * a(0); A(n - 1) = a(n - 1) + b(n - 1); c(n - 1) = 2 * a(n - 1); for i = 1, n - 1 A(i) = a(i) + b(i) c(i) = 2 * a(i) D(i) = c(i) + a(i) endfor

循环融合并不总是具有正向收益的，有时反而会降低性能甚至导致错误的结果。当前后两个循环存在数据依赖关系时，将它们融合可能会导致错误的结果：

# 第一个循环 for i in range(N): A[i] = B[i] + C1 # 第二个循环 for i in range(N): D[i] = A[i+1] + C2 # 循环融合 for i in range(N): A[i] = B[i] + C1 D[i] = A[i+1] + C2

融合后，第二个循环本来要读取的是 A 改变之后的值，但现在读取的是改变之前的值，导致错误的结果。当然也可以通过修改源代码的方式进行对齐（peeling），把公共部分融合：

A[0] = B[0] + C1 for i in range(2, N-1): A[i] = B[i] + C1 D[i-1] = A[i] + C2 D[N-1] = A[N] + C2

循环拆分（Loop Splitting）

拆分主要是将循环分成多个循环，可以在有条件的循环中使用，分为无条件循环和含条件循环。以仓库示例 spliting 的代码为例：

for i in range(n): A[i] = a[i] + b[i] c[i]=2 * a[i] if(temp[i] > data): d[i] = a[i] # 循环拆分 for i in range(n): A[i] = a[i] + b[i] c[i]=2 * a[i] for i in range(n): if(temp[i] > data): d[i] = a[i]

通过拆分，将包含控制流的代码独立为一个循环：一部分代码只有计算，可以在加速器上计算，而加速器不支持的控制流部分就可以回退到 CPU 计算。

循环拆分一般可以创造出更多的优化机会，例如和循环融合结合：

# 第一个循环 for i in range(N): A[i] = B[i] + C1 E[i] = K[i] * 2 # 第二个循环 for i in range(N): D[i] = A[i+1] + E[i]

这两个循环无法直接合并，因为 A 数组在两个循环之间存在依赖关系。但是可以通过对第一个循环进行拆分，把 E 数组的部分拆分出来，再融合进第二个循环中：

# 拆分 for i in range(N): A[i] = B[i] + C1 # 融合 for i in range(N): E[i] = K[i] * 2 D[i] = A[i+1] + E[i]

这样做可以提升数组 E 的局部性，减少 Cache miss——这正体现了"拆分+融合"组合拳的价值：先拆分消除依赖阻碍，再融合重建局部性。

小结与思考

• 因为 Cache 容量小的特点，一次只能容纳一部分数据，因此需要分析计算的时间局部性和空间局部性。
• 循环的优化方案针对不同的数据局部性和计算并行性，有循环分块、循环展开、循环重排、循环融合、循环拆分等方案。
• 循环优化的目的是提升数据的局部性或计算的并行性，以提高整体算子性能。其挑战包括数据局部性和计算并行性。
• 在 AISystem 仓库中，五类变换均有可运行的对照示例（tiling、fusion、reorder、spliting），且 TVM/Triton 等 DSL 已把"分块、重排、融合"等变换封装为可组合、可自动调优的调度原语，开发者只需声明计算逻辑，编译器即可在目标硬件上搜索最优的循环变换组合。

• 文档
• 教程
• 人工智能