Ansor: Generating High-Performance Tensor Programs for Deep Learning

文章标题: Ansor：为深度学习生成高性能张量程序
作者/机构: Lianmin Zheng (UC Berkeley), Chengfan Jia (Alibaba Group), Minmin Sun (Alibaba Group), Zhao Wu (Alibaba Group), Cody Hao Yu (Amazon Web Services, Inc), Ameer Haj-Ali (UC Berkeley), Yida Wang (Amazon Web Services), Jun Yang (Alibaba Group), Danyang Zhuo (UC Berkeley and Duke University), Koushik Sen (UC Berkeley), Joseph E. Gonzalez (UC Berkeley), Ion Stoica (UC Berkeley)

主要贡献

核心问题: 深度神经网络（DNN）的高效执行离不开高性能的张量程序。然而，为不同硬件平台上的各种算子获得高性能的张量程序极具挑战性。

现有方法及其局限:
1. 供应商提供的内核库 (如 cuDNN, MKL-DNN): 这类库需要投入巨大的工程努力为每个硬件平台和算子进行手动调优。这种巨大的手动工作量限制了新算子【7, Machine learning systems are stuck in a rut, 2019, HotOS】和专用加速器【35, A hardware–software blueprint for flexible deep learning specialization, 2019, IEEE Micro】的开发和创新。
2. 基于搜索的编译方法:
* 基于模板的方法 (如 TVM, FlexTensor): 依赖于预定义的、手动编写的模板，这限制了搜索空间的广度，无法覆盖许多有效的优化组合。
* 基于激进剪枝的方法 (如 Halide auto-scheduler): 通过评估不完整的程序来进行激进剪枝，这也阻碍了对全面搜索空间的探索。

研究目标: 本文旨在探索一种新颖的搜索策略来生成高性能的张量程序。该策略能够自动生成一个覆盖全面优化的大型搜索空间，并给予空间中每个张量程序被选择的机会，从而找到现有方法所遗漏的高性能程序。

面临的挑战:
1. 构建大型搜索空间: 如何为给定的计算定义自动构建一个足够大的搜索空间，以覆盖尽可能多的张量程序。
2. 高效搜索: 如何在比现有模板大几个数量级的搜索空间中高效搜索，同时避免比较不完整的程序。
3. 端到端性能优化: 在优化包含多个子图的整个DNN时，如何识别并优先处理对端到端性能至关重要的子图。

本文贡献:
* 大型分层搜索空间生成机制: 为计算图生成一个大型的分层张量程序搜索空间。
* 带有学习成本模型的进化策略: 使用进化策略和学习的成本模型来微调张量程序的性能。
* 基于梯度下降的调度算法: 在优化DNN的端到端性能时，优先处理重要的子图。
* Ansor系统实现与评估: 实现并全面评估了Ansor系统，证明其在多种DNN和硬件平台上均优于最先进的系统。

背景知识

深度学习生态与硬件多样性: 深度学习生态系统正面临着日益增长的硬件平台多样性，包括CPU、GPU、FPGA和ASIC。为了在这些平台上部署DNN，需要为DNN中使用的算子提供高性能的张量程序。这些算子集通常包含标准算子（如矩阵乘法、二维卷积）和机器学习研究人员发明的新型算子（如胶囊卷积【23, Matrix capsules with em routing, 2018】、空洞卷积【57, Multi-scale context aggregation by dilated convolutions, 2015, arXiv】）。

编译技术的作用: 为了在各种硬件平台上高效地提供可移植的算子性能，引入了多种编译器技术（如TVM【11, Tvm: an automated end-to-end optimizing compiler for deep learning, 2018, OSDI】、Halide【41, Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, SIGPLAN Notices】、Tensor Comprehensions【49, Tensor comprehensions: framework-agnostic high-performance machine learning abstractions, 2018, arXiv】）。用户使用高级声明性语言定义计算，编译器根据定义生成优化的张量程序。图1展示了在TVM张量表达式语言中矩阵乘法的计算定义，用户主要需要定义张量的形状以及输出张量中每个元素的计算方式。

自动生成高性能程序的挑战: 从高级定义自动生成高性能张量程序极其困难。编译器需要在包含优化组合选择（如分块结构、分块大小、向量化、并行化）的巨大且复杂的空间中进行搜索。找到高性能程序要求搜索策略能覆盖全面的空间并高效地探索。本节将描述两种近期有效的方法。

模板引导的搜索: 在模板引导的搜索中，搜索空间由手动模板定义。如图2a所示，编译器（如TVM）要求用户为计算定义手动编写模板。模板定义了张量程序的结构，并带有一些可调参数（如分块大小和展开因子）。然后，编译器针对特定的输入形状配置和硬件目标搜索这些参数的最佳值。这种方法在常见的深度学习算子上取得了良好性能。然而，开发模板需要大量精力。例如，TVM的代码库中已有超过15,000行用于这些模板的代码。随着新算子和新硬件平台的出现，这个数字还在持续增长。此外，构建高质量的模板需要算子和硬件两方面的专业知识。开发高质量模板需要相当大的研究努力【32, Optimizing cnn model inference on cpus, 2019, USENIX ATC】、【55, A unified optimization approach for cnn model inference on integrated gpus, 2019, ICPP】、【59, Flextensor: an automatic schedule exploration and optimization framework for tensor computation on heterogeneous system, 2020, ASPLOS】。尽管模板设计复杂，但手动模板仅覆盖了有限的程序结构，因为手动枚举所有算子的所有优化选择是不可行的。这种方法通常需要为每个算子定义一个模板。FlexTensor【59, Flextensor: an automatic schedule exploration and optimization framework for tensor computation on heterogeneous system, 2020, ASPLOS】提出了一个通用模板来覆盖多个算子，但其模板仍是为单个算子粒度设计的，无法包含涉及多个算子（如算子融合）的优化。优化具有多个算子的计算图的搜索空间应包含组合算子的不同方式。基于模板的方法无法实现这一点，因为它不能在搜索过程中分解其固定的模板并重新组合它们。
搜索策略比较。伪代码显示了带有循环嵌套的张量程序。橙色背景中的问号表示低级参数。

基于顺序构建的搜索: 这种方法通过将程序构建分解为一系列固定的决策来定义搜索空间。然后，编译器使用诸如束搜索（beam search）【34, Speech understanding systems: report of a steering committee, 1977, Artificial Intelligence】之类的算法来搜索好的决策（例如，Halide auto-scheduler【2, Learning to optimize halide with tree search and random programs, 2019, TOG】）。在这种方法中，编译器通过顺序展开计算图中的所有节点来构建张量程序。对于每个节点，编译器会做出一些决策，决定如何将其转换为低级张量程序（即决定计算位置、存储位置、分块大小等）。当所有节点都展开后，一个完整的张量程序就构建好了。这种方法对每个节点使用一组通用的展开规则，因此可以自动搜索而无需手动模板。由于每个决策的可能选择数量很大，为了使顺序过程可行，该方法在每次决策后仅保留前k个候选程序。编译器使用学习的成本模型来估计和比较候选程序的性能以选择前k个候选者；而其他候选者则被剪枝。在搜索过程中，候选程序是不完整的，因为只有部分计算图被展开或只做出了一部分决策。图2b展示了这一过程。

顺序构建搜索的局限性: 估计不完整程序的最终性能在几个方面都很困难：(1) 在完整程序上训练的成本模型无法准确预测不完整程序的最终性能。成本模型只能在完整程序上训练，因为我们需要编译程序并测量其执行时间以获得训练标签。直接使用此模型比较不完整程序的最终性能将导致准确性差。作为一个案例研究，我们在我们的搜索空间中的20,000个随机完整程序上训练了我们的成本模型（§5.2），并使用该模型来预测不完整程序的最终性能。不完整程序是通过仅应用完整程序的一部分循环转换获得的。我们使用两个排名指标进行评估：成对比较的准确性和top-k程序的召回率@k分数（k=10）。如图3所示，两条曲线分别从50%和0%开始，这意味着零信息的随机猜测给出50%的成对比较准确性和0%的top-k召回率。随着程序变得完整，两条曲线迅速增加，这意味着成本模型对于完整程序表现非常好，但无法准确预测不完整程序的最终性能。(2) 顺序决策的固定顺序限制了搜索空间的设计。例如，一些优化需要向计算图添加新节点（例如，添加缓存节点，使用rfactor【46, Parallel associative reductions in halide, 2017, CGO】）。不同程序的决策数量变得不同。很难对齐不完整程序进行公平比较。(3) 基于顺序构建的搜索不可扩展。扩大搜索空间需要增加更多的顺序构建步骤，然而，这会导致更差的累积误差。
随机部分程序的成对比较准确率和 top-k 召回率曲线。在两个子图中，值越高越好。

Ansor的分层方法: 如图2c所示，Ansor由一个分层搜索空间支持，该空间将高级结构与低级细节解耦。Ansor自动为计算图构建搜索空间，无需手动开发模板。然后，Ansor从该空间中采样完整的程序，并对完整的程序进行微调，避免了对不完整程序的不准确估计。

方法细节

设计总览

Ansor是一个自动化的张量程序生成框架。图4展示了Ansor的整体架构。Ansor的输入是一组待优化的DNN。Ansor使用Relay【42, Relay: a high-level compiler for deep learning, 2019, arXiv】中的算子融合算法，将来自流行模型格式（如ONNX【6, Onnx: open neural network exchange, 2019】、TensorFlow PB）的DNN转换为分区后的小子图。然后，Ansor为这些子图生成张量程序。Ansor有三个主要组成部分：(1) 一个程序采样器，用于构建大型搜索空间并从中采样多样化的程序；(2) 一个性能调优器，用于微调采样程序的性能；(3) 一个任务调度器，用于为优化DNN中的多个子图分配时间资源。

程序采样器: Ansor必须解决的一个关键挑战是为给定的计算图生成一个大型搜索空间。为了覆盖具有各种高级结构和低级细节的多样化张量程序，Ansor利用了一个具有两个层次的分层搜索空间表示：草图（sketch）和注解（annotation）（§4）。Ansor将程序的高级结构定义为草图，并将数十亿的低级选择（例如，分块大小、并行、展开注解）留作注解。这种表示使得Ansor能够灵活地枚举高级结构并高效地采样低级细节。Ansor包含一个程序采样器，该采样器从空间中随机采样程序，以提供对搜索空间的全面覆盖。

性能调优器: 随机采样程序的性能不一定好。下一个挑战是如何对它们进行微调。Ansor采用进化搜索和学习的成本模型来进行迭代式微调（§5）。在每次迭代中，Ansor使用重新采样的新程序以及来自先前迭代的优秀程序作为初始种群来开始进化搜索。进化搜索通过变异和交叉来微调程序，这些操作执行乱序重写并解决了顺序构建的局限性。查询学习的成本模型比实际测量快几个数量级，因此我们可以在几秒钟内评估数千个程序。

任务调度器: 使用程序采样和性能微调使得Ansor能够为计算图找到高性能的张量程序。直观地，将整个DNN视为单个计算图并为其生成完整的张量程序可能实现最佳性能。然而，这样做效率低下，因为它必须处理不必要的指数级搜索空间爆炸。通常，编译器将DNN的大计算图划分为几个小子图【11, Tvm: an automated end-to-end optimizing compiler for deep learning, 2018, OSDI】、【42, Relay: a high-level compiler for deep learning, 2019, arXiv】。由于DNN的逐层构建特性，这种划分对性能的影响可以忽略不计。这带来了Ansor的最后一个挑战：在为多个子图生成程序时如何分配时间资源。Ansor中的任务调度器（§6）使用一种基于梯度下降的调度算法，将资源分配给那些更有可能提高端到端DNN性能的子图。
系统总览。灰色箭头显示了从深度学习模型中提取子图并为其生成优化程序流程。绿色箭头表示测量器返回分析数据以更新系统中所有组件的状态。

程序采样

现有方法的搜索空间局限性: 一个算法能找到的最佳程序取决于其探索的搜索空间。现有方法中考虑的搜索空间受以下因素限制：(1) 手动枚举 (例如, TVM 【12, Learning to optimize tensor programs, 2018, NIPS】)。通过模板手动枚举所有可能的选择是不切实际的，因此现有的手动模板仅启发式地覆盖了有限的搜索空间。(2) 激进的早期剪枝 (例如, Halide auto-scheduler 【2, Learning to optimize halide with tree search and random programs, 2019, TOG】)。基于评估不完整程序的激进早期剪枝阻止了搜索算法探索空间中的某些区域。

Ansor的搜索空间扩展技术: 本节介绍了通过解决上述限制来扩展所考虑搜索空间边界的技术。为了解决问题(1)，我们通过递归地应用一组灵活的派生规则来自动扩展搜索空间。为了避免问题(2)，我们在搜索空间中随机采样完整的程序。由于随机采样为每个点提供了被采样的同等机会，我们的搜索算法可以潜在地探索所考虑空间中的每个程序。我们不依赖随机采样来找到最优程序，因为每个采样程序随后都会被微调（§5）。

分层搜索空间: 为了采样能够覆盖大型搜索空间的程序，我们定义了一个具有两个层次的分层搜索空间：草图（sketch）和注解（annotation）。我们将程序的高级结构定义为草图，并将数十亿的低级选择（例如，分块大小、并行、展开注解）作为注解。在顶层，我们通过递归应用一些派生规则来生成草图。在底层，我们随机注解这些草图以获得完整的程序。这种表示从数十亿的低级选择中总结出一些基本结构，从而能够灵活地枚举高级结构并高效地采样低级细节。

平台支持: 尽管Ansor同时支持CPU和GPU，我们以CPU为例，在§4.1和§4.2中解释采样过程。然后在§4.3中讨论GPU过程的不同之处。

草图生成

草图生成过程: 如图4所示，程序采样器接受分区后的子图作为输入。图5的第一列展示了两个输入示例。输入有三种等价形式：数学表达式、通过直接展开循环索引获得的相应朴素程序，以及相应的计算图（有向无环图，或DAG）。为了为具有多个节点的DAG生成草图，我们按拓扑顺序访问所有节点并迭代地构建结构。对于计算密集且有大量数据复用机会的计算节点（例如，conv2d, matmul），我们为它们构建基本的分块和融合结构作为草图。对于简单的逐元素节点（例如，ReLU，逐元素加法），我们可以安全地内联它们。注意，在草图生成过程中，新的节点（例如，缓存节点，布局转换节点）也可能被引入到DAG中。

基于派生的枚举方法: 我们提出一种基于派生的枚举方法，通过递归应用几个基本规则来生成所有可能的草图。这个过程以一个DAG作为输入，并返回一个草图列表。我们定义状态 $s = (S, i)$，其中 $S$ 是DAG的当前部分生成的草图，而 $i$ 是当前工作节点的索引。DAG中的节点按从输出到输入的拓扑顺序排序。派生从初始的朴素程序和最后一个节点开始，即初始状态 $s = (\text{naive program, index of the last node})$。然后我们尝试递归地将所有派生规则应用于这些状态。对于每个规则，如果当前状态满足应用条件，我们将规则应用于 $s = (S, i)$ 并得到 $s' = (S', i')$，其中 $i' \leq i$。这样，索引 $i$（工作节点）单调递减。当 $i = 0$ 时，一个状态成为终止状态。在枚举过程中，多个规则可以应用于一个状态以生成多个后继状态。一个规则也可以生成多个可能的后继状态。因此，我们维护一个队列来存储所有中间状态。当队列为空时，过程结束。所有终止状态中的 $s.S$ 构成了草图生成结束时的草图列表。对于一个典型的子图，草图的数量少于10个。
用于生成草图的派生规则。条件在当前状态 s = (S, i) 上运行。应用从当前状态 s 派生出下一个状态 s' = (S', i')。注意，某些函数（例如 AddRfactor, FuseConsumer）可以返回多个可能的 S' 值。在这种情况下，我们收集所有可能的 S'，并为单个输入状态 s 返回多个下一个状态 s'。

派生规则: 表1列出了我们用于CPU的派生规则。我们首先提供所用谓词的定义，然后描述每个规则的功能。
- IsStrictInliable(S, i): 指示S中的节点i是否是一个可以始终被内联的简单逐元素算子（例如，逐元素加法，ReLU）。
- HasDataReuse(S,i): 指示S中的节点i是否是计算密集型算子并且具有丰富的算子内数据复用机会（例如，matmul, conv2d）。
- HasFusibleConsumer(S,i): 指示S中的节点i是否只有一个消费者j，并且节点j可以被融合到节点i中（例如，matmul + bias_add, conv2d + relu）。
- HasMoreReductionParallel(S, i): 指示S中的节点i在空间维度上的并行性很小，但在归约维度上有充足的并行机会（例如，计算矩阵的2-范数，matmul $C_{2 \times 2} = A_{2 \times 512} \cdot B_{512 \times 2}$）。
我们对计算定义进行静态分析以获得这些谓词的值。该分析通过解析数学表达式中的读/写模式自动完成。接下来，我们介绍每个派生规则的功能。

规则1 (跳过): 如果一个节点不是严格可内联的，就简单地跳过它。
规则2 (总是内联): 总是内联严格可内联的节点。由于规则1和规则2的条件是互斥的，一个状态若 $i > 1$ 总能满足其中一个并继续派生。
规则3, 4, 和 5: 处理具有数据复用节点的多级分块和融合。
- 规则3 (多级分块): 对数据可复用节点执行多级分块。对于CPU，我们使用“SSRSRS”分块结构，其中“S”代表空间循环的一个分块级别，“R”代表归约循环的一个分块级别。例如，在矩阵乘法 $C(i, j) = \sum_{k} A[i, k] \times B[k, j]$ 中，i和j是空间循环，k是归约循环。“SSRSRS”分块结构将原始的3级循环(i, j, k)扩展为10级循环(i0, j0, i1, j1, k0, i2, j2, k1, i3, j3)。虽然我们不改变循环顺序，但这种多级分块也可以覆盖一些重排序的情况。例如，上述10级循环可以通过将其他循环的长度设置为1来特化为简单的重排序(k0, j2, i3)。“SSRSRS”分块结构对于深度学习中计算密集的稠密算子（例如，matmul, conv2d, conv3d）是通用的，因为它们都由空间循环和归约循环组成。
- 规则4 (带融合的多级分块): 执行多级分块并融合可融合的消费者。例如，我们将逐元素节点（如ReLU，偏置加法）融合到分块的节点（如conv2d，matmul）中。
- 规则5 (添加缓存): 如果当前数据可复用节点没有可融合的消费者，则添加一个缓存节点。例如，DAG中的最终输出节点没有任何消费者，默认情况下它直接将结果写入主内存，由于内存访问的高延迟，这是低效的。通过添加一个缓存节点，我们向DAG中引入一个新的可融合消费者，然后规则4可以应用于将这个新添加的缓存节点融合到最终输出节点中。融合缓存节点后，最终输出节点现在将其结果写入一个缓存块，当该块中的所有数据计算完毕后，缓存块将一次性写入主内存。
规则6 (归约因子分解): 可以使用 rfactor【46, Parallel associative reductions in halide, 2017, CGO】将一个归约循环分解为一个空间循环，以带来更多的并行性。

示例: 图5展示了生成的草图的三个例子。这些草图与TVM中的手动模板不同，因为手动模板同时指定了高级结构和低级细节，而草图只定义了高级结构。
- 对于输入示例1，DAG中四个节点的排序是(A, B, C, D)。为了派生DAG的草图，我们从输出节点D(i=4)开始，并逐个对节点应用规则。具体来说，生成草图1的派生过程是：
$$\begin{array}{l} \text { Input } 1 \rightarrow \sigma\left(S_{0}, i=4\right) \xrightarrow{\text { Rule } 1} \sigma\left(S_{1}, i=3\right) \xrightarrow{\text { Rule } 4} \\ \sigma\left(S_{2}, i=2\right) \xrightarrow{\text { Rule } 1} \sigma\left(S_{3}, i=1\right) \xrightarrow{\text { Rule } 1} \text { Sketch } 1 \end{array}$$
- 对于输入示例2，五个节点的排序是(A, B, C, D, E)。同样，我们从输出节点E(i=5)开始，并递归地应用规则。生成草图2的派生过程是：
$$ \begin{array}{l} \text{Input } 2 \to \sigma(S_0, i=5) \xrightarrow{\text{Rule 5}} \sigma(S_1, i=5) \xrightarrow{\text{Rule 4}} \\ \sigma(S_2, i=4) \xrightarrow{\text{Rule 1}} \sigma(S_3, i=3) \xrightarrow{\text{Rule 1}} \\ \sigma(S_4, i=2) \xrightarrow{\text{Rule 2}} \sigma(S_5, i=1) \xrightarrow{\text{Rule 1}} \text{Sketch 2} \end{array} $$
- 类似地，生成草图3的派生过程是：
$$ Input~2 \to \sigma(S_0, i=5) \xrightarrow{Rule~6} \sigma(S_1, i=4) \xrightarrow{Rule~1} \\ \sigma(S_2, i=3) \xrightarrow{Rule~1} \sigma(S_3, i=2) \xrightarrow{Rule~2} \\ \sigma(S_4, i=1) \xrightarrow{Rule~1} Sketch~3 $$

定制化: 尽管所提出的规则足够实用，可以覆盖大多数算子的结构，但总有例外。例如，一些特殊算法（如Winograd卷积【30, Fast algorithms for convolutional neural networks, 2016, CVPR】）和加速器内建函数（如TensorCore【37, Nvidia tensor cores, 2017】）需要特殊的tile结构才能有效。虽然模板引导的搜索方法（在TVM中）可以为每个新情况制作一个新的模板，但这需要大量的设计工作。另一方面，Ansor中基于派生的草图生成足够灵活，可以为新兴算法和硬件生成所需的结构，因为我们允许用户注册新的派生规则，并将其与现有规则无缝集成。

随机注解

将草图转换为完整程序: 前一小节生成的草图是不完整的程序，因为它们只有分块结构，没有具体的分块大小和循环注解，如并行、展开和向量化。在本小节中，我们对草图进行注解，使它们成为用于微调和评估的完整程序。

注解过程: 给定一个生成的草图列表，我们随机选择一个草图，随机填充分块大小，并行化一些外层循环，向量化一些内层循环，并展开一些内层循环。我们还随机更改程序中某些节点的计算位置，以对分块结构进行微调。本小节中所有“随机”均指在所有有效值上的均匀分布。如果某些特殊算法需要自定义注解才能有效（例如，特殊的展开），我们允许用户在计算定义中给出简单的提示来调整注解策略。最后，由于可以在编译时更改常量张量的布局且不带来运行时开销，我们根据多级分块结构重写常量张量的布局，使其尽可能缓存友好。这种优化是有效的，因为对于推理应用，卷积层或全连接层的权重张量是常量。

示例: 随机采样的示例如图5所示。采样出的程序可能比草图有更少的循环，因为长度为1的循环被简化了。

GPU支持

GPU的特定调整: 对于GPU，我们将多级分块结构从"SSRSRS"更改为"SSSRRSRS"，以匹配GPU的架构。前三个空间分块中的循环分别绑定到BlockIdx、虚拟线程（用于减少bank冲突）和ThreadIdx。我们添加了两个草图派生规则，一个用于通过插入缓存节点来利用共享内存（类似于规则5），另一个用于跨线程归约（类似于规则6）。

性能微调

微调的必要性: 程序采样器采样的程序对搜索空间有很好的覆盖，但其质量无法保证。这是因为优化选择，如分块结构和循环注解，都是随机采样的。本节介绍性能调优器，它通过进化搜索和学习的成本模型来微调采样程序的性能。
生成的草图和采样程序的示例。此图显示了两个示例输入、三个生成的草图和四个采样程序。代码示例是类似python语法的伪代码。

迭代微调过程: 微调是迭代进行的。在每次迭代中，我们首先使用进化搜索根据学习的成本模型找到一小批有前途的程序。然后，我们在硬件上测量这些程序以获得实际的执行时间成本。最后，从测量中得到的性能分析数据用于重新训练成本模型，使其更加准确。

进化搜索机制: 进化搜索使用随机采样的程序以及来自先前测量的高质量程序作为初始种群，并应用变异和交叉来生成下一代。学习的成本模型用于预测每个程序的适应度，在我们的案例中是程序的吞吐量。我们运行进化固定代数，并挑选出搜索过程中找到的最佳程序。我们利用学习的成本模型，因为成本模型可以对程序的适应度给出相对准确的估计，同时比实际测量快几个数量级。它使我们能够在几秒钟内比较搜索空间中的数万个程序，并挑选出有希望的程序进行实际测量。

进化搜索

进化搜索概述: 进化搜索【54, Evolutionary algorithms: a critical review and its future prospects, 2016, ICGTSPICC】是一种受生物进化启发的通用元启发式算法。通过迭代地变异高质量的程序，我们可以生成可能具有更高质量的新程序。进化从采样的初始代开始。为了生成下一代，我们首先根据特定概率从当前代中选择一些程序。选择一个程序的概率与其由学习的成本模型（§5.2）预测的适应度成正比，这意味着性能得分越高的程序被选择的概率越高。对于选定的程序，我们随机应用其中一种进化操作来生成新程序。基本上，对于我们在采样（§4.2）期间做出的决策，我们设计了相应的进化操作来重写和微调它们。

分块大小变异: 此操作扫描程序并随机选择一个分块循环。对于此分块循环，它将一个分块级别的分块大小除以一个随机因子，并将此因子乘到另一个级别。由于此操作保持分块大小的乘积等于原始循环长度，因此变异后的程序始终是有效的。

并行变异: 此操作扫描程序并随机选择一个已用并行注解的循环。对于此循环，此操作通过融合其相邻循环级别或按一个因子拆分它来更改并行粒度。

Pragma变异: 程序中的一些优化是由编译器特定的pragma指定的。此操作扫描程序并随机选择一个pragma。对于此pragma，此操作随机将其变异为另一个有效值。例如，我们的底层代码生成器通过提供一个auto_unroll_max_step=N pragma来支持最大步数的自动展开。我们随机调整数字N。

计算位置变异: 此操作扫描程序并随机选择一个未进行多级分块的灵活节点（例如，卷积层中的填充节点）。对于此节点，该操作随机将其计算位置更改为另一个有效的附加点。

基于节点的交叉: 交叉是一种通过组合两个或多个父代的基因来生成新后代的操作。在Ansor中，一个程序的基因是其重写步骤。Ansor生成的每个程序都是从其初始的朴素实现重写而来的。Ansor在草图生成和随机注解期间为每个程序保留了完整的重写历史。我们可以将重写步骤视为程序的基因，因为它们描述了该程序是如何从初始的朴素程序形成的。基于此，我们可以通过组合两个现有程序的重写步骤来生成一个新程序。然而，任意组合两个程序的重写步骤可能会破坏步骤中的依赖关系并创建一个无效的程序。因此，Ansor中交叉操作的粒度是基于DAG中的节点，因为不同节点之间的重写步骤通常依赖性较小。Ansor为每个节点随机选择一个父代，并合并所选节点的重写步骤。当节点之间存在依赖关系时，Ansor会尝试用简单的启发式方法分析和调整步骤。Ansor进一步验证合并后的程序以保证功能正确性。验证很简单，因为Ansor只使用一小组循环转换重写步骤，并且底层代码生成器可以通过依赖性分析来检查正确性。

进化搜索的优势: 进化搜索利用变异和交叉重复生成一组新的候选者，持续数轮，并输出一小组得分最高的程序。这些程序将在目标硬件上编译和测量，以获得真实的运行时间成本。收集到的测量数据随后用于更新成本模型。通过这种方式，学习的成本模型的准确性逐渐提高，以匹配目标硬件。因此，进化搜索逐渐为目标硬件平台生成更高质量的程序。与TVM和FlexTensor中只能在固定的网格状参数空间中工作的搜索算法不同，Ansor中的进化操作是专门为张量程序设计的。它们可以应用于通用的张量程序，并能处理具有复杂依赖关系的搜索空间。与Halide自动调度器中的展开规则不同，这些操作可以对程序执行乱序修改，解决了顺序限制。

学习的成本模型

成本模型的作用: 在搜索过程中，成本模型对于快速估计程序性能是必要的。我们采用了一个与相关工作【2, Learning to optimize halide with tree search and random programs, 2019, TOG】、【12, Learning to optimize tensor programs, 2018, NIPS】类似的 learned cost model，但设计了新的程序特征。一个基于学习成本模型的系统具有很好的可移植性，因为单个模型设计可以通过输入不同的训练数据来复用于不同的硬件后端。

模型设计: 由于我们的目标程序主要是数据并行的张量程序，它们由多个交错的循环嵌套构成，最内层是几个赋值语句，我们训练成本模型来预测循环嵌套中一个最内层非循环语句的得分。对于一个完整的程序，我们对每个最内层的非循环语句进行预测，并将预测值相加作为总分。我们通过在完整程序的上下文中提取特征来为最内层的非循环语句构建特征向量。提取的特征包括算术特征和内存访问特征。提取特征的详细列表见本文扩展版【58, Ansor: generating high-performance tensor programs for deep learning, 2020, https://arxiv.org/abs/2006.06762】的附录。

损失函数与训练: 我们使用加权平方误差作为损失函数。因为我们主要关心从搜索空间中识别出性能良好的程序，所以我们对运行速度更快的程序赋予更大的权重。具体来说，模型 $f$ 对吞吐量为 $y$ 的程序 $P$ 的损失函数是 $loss(f, P, y) = w_p(\sum_{s \in S(P)} f(s) - y)^2 = y(\sum_{s \in S(P)} f(s) - y)^2$，其中 $S(P)$ 是 $P$ 中最内层非循环语句的集合。我们直接使用吞吐量 $y$ 作为权重。我们训练一个梯度提升决策树【9, Xgboost: a scalable tree boosting system, 2016, KDD】作为底层模型 $f$。我们为所有来自所有DAG的张量程序训练一个单一模型，并将来自同一DAG的所有程序的吞吐量归一化到[0, 1]的范围内。在优化一个DNN时，测量的程序数量通常少于30,000个。在这样的小数据集上训练梯度提升决策树非常快，所以我们每次都训练一个新模型，而不是进行增量更新。

任务调度器

调度问题的背景: 一个DNN可以被划分为许多独立的子图（例如，conv2d + relu）。对于某些子图，花费时间调优它们并不能显著提高端到端的DNN性能。这有两个原因：(1) 该子图不是性能瓶颈，或者 (2) 调优对该子图性能的提升微乎其微。

Ansor的动态资源分配: 为了避免在调优不重要的子图上浪费时间，Ansor动态地为不同的子图分配不同数量的时间资源。以ResNet-50为例，图划分后它有29个独特的子图。这些子图大多是具有不同形状配置（输入大小、核大小、步长等）的卷积层。我们需要为不同的卷积层生成不同的程序，因为最佳的张量程序依赖于这些形状配置。在实际中，用户可能为其所有应用拥有多个DNN。这导致了更多的子图，也为减少总调优时间提供了更多机会，因为我们可以在子图之间共享和重用知识。一个子图也可能在一个DNN中或跨不同DNN多次出现。

任务定义与调度: 我们将为子图生成高性能程序的过程定义为一个任务。这意味着优化单个DNN需要完成数十个任务（例如，ResNet-50需要29个任务）。Ansor的任务调度器以迭代的方式为任务分配时间资源。在每次迭代中，Ansor选择一个任务，为该子图生成一批有前途的程序，并在硬件上测量这些程序。我们将这样一次迭代定义为一个时间资源单位。当我们为一个任务分配一个时间资源单位时，该任务就获得了一次生成和测量新程序的机会，这也意味着有机会找到更好的程序。接下来，我们介绍调度问题的形式化及其解决方案。

问题形式化

目标函数: 在调优一个DNN或一组DNN时，用户可以有各种类型的目标，例如，减少一个DNN的延迟，满足一组DNN的延迟要求，或者在调优不再显著提高DNN性能时最小化调优时间。因此，我们为用户提供了一组目标函数来表达他们的目标。用户也可以提供自己的目标函数。

数学表达: 假设总共有 $n$ 个任务。令 $t \in Z^n$ 为分配向量，其中 $t_i$ 是在任务 $i$ 上花费的时间单位数。令任务 $i$ 实现的最小子图延迟是分配向量的函数 $g_i(t)$。令DNN的端到端成本是子图延迟的函数 $f(g_1(t), g_2(t), ..., g_n(t))$。我们的目标是最小化端到端成本：
$$ \min_t f(g_1(t), g_2(t), ..., g_n(t)) \quad \text{s.t.} \sum_{i=1}^n t_i \leq T $$
其中 $T$ 是总时间预算。

单一DNN延迟最小化: 为了最小化单个DNN的端到端延迟，我们可以定义 $f(g_1, g_2, ..., g_n) = \sum_{i=1}^n w_i \times g_i$，其中 $w_i$ 是任务 $i$ 在DNN中出现的次数。这个公式是直接的，因为 $f$ 是端到端DNN延迟的一个近似。

多个DNN的目标函数: 当调优一组DNN时，有几个选项。表2显示了调优多个DNN的几个示例目标函数。设 $m$ 是DNN的数量，$S(j)$ 是属于DNN $j$ 的任务集合。
- $f_1$ 将每个DNN的延迟相加，这意味着优化一个顺序运行所有DNN一次的流水线的成本。
- 在 $f_2$ 中，我们将 $L_j$ 定义为DNN $j$ 的延迟要求，这意味着如果一个DNN的延迟已经满足要求，我们不想再花时间在其上。
- 在 $f_3$ 中，我们将 $B_j$ 定义为DNN $j$ 的参考延迟。因此，我们的目标是最大化相对于给定参考延迟的加速比的几何平均值。
- 最后，在 $f_4$ 中，我们定义了一个函数 $ES(g_i, t)$，它通过查看任务 $i$ 的延迟历史返回一个提前停止值。它可以实现每个任务提前停止的效果。
$$f_1 = \sum_{j=1}^{m} \sum_{i \in S(j)} w_i \times g_i(t)$$
$$f_2 = \sum_{j=1}^{m} \max\left(\sum_{i \in S(j)} w_i \times g_i(t), L_j\right)$$
$$f_3 = -\left(\prod_{j=1}^{m} \frac{B_j}{\sum_{i \in S(j)} w_i \times g_i(t)}\right)^{\frac{1}{m}}$$
$$f_4 = \sum_{j=1}^{m} \sum_{i \in S(j)} w_i \times \max(g_i(t), ES(g_i,t))$$

使用梯度下降进行优化

基于梯度的调度算法: 我们提出一种基于梯度下降的调度算法来有效优化目标函数。给定当前分配 $t$，其思想是近似目标函数的梯度 $\frac{\partial f}{\partial t_i}$，以便选择任务 $i$ 使得 $i = \arg\max_i |\frac{\partial f}{\partial t_i}|$。我们通过近似梯度来做出决策。
$$ \frac{\partial f}{\partial t_i} \approx \frac{\partial f}{\partial g_i} \left( \alpha \frac{g_i(t_i) - g_i(t_i - \Delta t)}{\Delta t} + (1 - \alpha) \left( \min \left( - \frac{g_i(t_i)}{t_i}, \beta \frac{C_i}{\max_{k \in \mathcal{N}(i)} V_k} - g_i(t_i) \right) \right) \right) $$
其中 $\Delta t$ 是一个小的反向窗口大小，$g_i(t_i)$ 和 $g_i(t_i - \Delta t)$ 从分配历史中可知。$N(i)$ 是 $i$ 的相似任务集合，$C_i$ 是任务 $i$ 中的浮点运算次数，$V_k$ 是我们在任务 $k$ 中每秒可以实现的浮点运算次数。参数 $\alpha$ 和 $\beta$ 控制信任某些预测的权重。

算法执行流程: Ansor从 $t = 0$ 开始，并通过一轮轮询（round-robin）进行热身，以获得初始分配向量 $t = (1, 1, ..., 1)$。热身之后，在每次迭代中，我们计算每个任务的梯度，并选择 $\arg\max_i |\frac{\partial f}{\partial t}|$。然后我们将资源单位分配给任务 $i$ 并更新分配向量 $t_i = t_i + 1$。优化过程一直持续到时间预算用完。为了鼓励探索，我们采用了一种 $\epsilon$-greedy策略【47, Reinforcement learning: an introduction, 2018, MIT press】，保留一个 $\epsilon$ 的概率随机选择一个任务。

调度策略示例: 以优化单个DNN的端到端延迟为例，Ansor会优先处理初始延迟高的子图，因为我们乐观地猜测可以迅速降低其延迟。之后，如果Ansor在其上花费了许多迭代而没有观察到延迟下降，Ansor会离开该子图，因为其 $|\frac{\partial f}{\partial t}|$ 减小了。

实验环境

硬件配置:
- Intel CPU: 18核 Platinum 8124M @ 3.0 GHz
- NVIDIA GPU: V100
- ARM CPU: 树莓派 3b+ 上的 4核 Cortex-A53 @ 1.4GHz
软件配置:
- Ansor实现: 核心用C++实现，约12K行代码。生成的程序先表示为Ansor自己的中间表示（IR），然后降级为TVM IR进行代码生成。
- 数据类型: 所有评估均使用float32。
- 基线框架:
  - PyTorch (v1.5)
  - TensorFlow (v2.0)
  - TensorRT (v6.0)
  - TensorFlow Lite (V2.0)
  - Halide auto-scheduler (commit: 1f875b0)
  - FlexTensor (commit: 7ac302c)
  - AutoTVM (commit: 69313a7)
数据集/工作负载:
- 单算子: 1D/2D/3D卷积(C1D, C2D, C3D), 矩阵乘法(GMM), 分组卷积(GRP), 空洞卷积(DIL), 深度可分离卷积(DEP), 转置卷积(T2D), 胶囊卷积(CAP), 矩阵2-范数(NRM)。每个算子有4种常见形状配置，批量大小为1和16。
- 子图:
  - ConvLayer: 2D卷积 + 批归一化 + ReLU。
  - TBS: 两个矩阵转置 + 批处理矩阵乘法 + softmax。
- 端到端网络:
  - 图像分类: ResNet-50, MobileNet-V2
  - 动作识别: 3D-ResNet18
  - 图像生成: DCGAN 生成器
  - 语言理解: BERT

实验结果

单算子基准测试

实验设置: 在Intel CPU上进行，每个测试案例最多允许1000次测量试验。由于Halide auto-scheduler没有预训练的AVX-512成本模型，因此在本节和子图基准测试中禁用了AVX-512。

结果: 如图6所示，Ansor在所有算子和批量大小设置中表现最佳或并列最佳。Ansor相比现有搜索框架性能提升1.1至22.5倍。性能提升源于其巨大的搜索空间和有效的探索策略。
- 对于大多数算子，Ansor找到的最佳程序都超出了现有搜索框架的搜索空间，因为它能探索更多的优化组合。
- 在NRM上的显著加速是因为Ansor可以并行化归约循环。
- 在T2D上的巨大加速是因为Ansor能使用正确的分块结构和展开策略，让代码生成器简化带步长的转置卷积中与零的乘法。
- 相比之下，其他框架在其搜索空间中未能捕获许多有效优化。例如，Halide的展开规则不拆分GMM中的归约循环；AutoTVM的模板有有限的分块结构；FlexTensor的模板不改变padding的计算位置，且无法运行NRM等归约算子。
单算子性能基准测试，在一台20核Intel-Platinum-8269CY上进行。y轴是每个算子相对于最佳吞吐量归一化的吞吐量。

消融研究: 在ResNet-50的最后一个卷积层上运行Ansor的四种变体（图7）。
- Ansor (ours): 使用所有提出的技术。
- Beam Search: 在采样过程中使用成本模型剪枝不完整的程序，不使用微调。
- No fine-tuning: 禁用微调，仅依赖随机采样。
- Limited space: 限制搜索空间，使其与现有手动模板相似。
结果表明，去掉大型搜索空间或高效的微调都会显著降低最终性能。“Beam Search”中的激进早期剪枝因不准确的估计而丢弃了最终性能良好的不完整程序。
在一个卷积算子上对Ansor的四个变体进行消融研究。y轴是相对于最佳程序吞吐量的吞吐量。

子图基准测试

实验设置: 在Intel CPU和NVIDIA GPU上对"ConvLayer"和"TBS"两个子图进行测试。每个测试案例最多1000次测量试验。FlexTensor无法在复杂的"TBS"子图上运行。

结果: 如图8所示，Ansor在两个平台上对这些子图的表现均优于手动库和其他搜索框架，性能提升1.1至14.2倍。这表明Ansor能为子图持续生成高性能程序。FlexTensor在单算子上表现良好，但在子图上优势减弱，因为它缺乏算子融合的支持。
在一台20核Intel-Platinum-8269CY和一台NVIDIA V100上的子图性能基准测试。“@C”表示CPU结果，“@G”表示GPU结果。y轴是每个子图相对于最佳吞吐量归一化的吞吐量。

端到端网络基准测试

实验设置: 在Intel CPU、NVIDIA GPU和ARM CPU上测试五个DNN的端到端推理时间。AutoTVM和Ansor的调优预算为每个DNN的子图数 * 1000次测量试验。本节在Intel CPU上启用了AVX-512。

结果: 如图9所示，Ansor在所有情况下表现最佳或并列最佳。
- 与AutoTVM相比: Ansor在所有情况下均持平或超越AutoTVM，加速比为1.0至21.8倍。
- 与最佳替代方案相比: Ansor在Intel CPU、ARM CPU和NVIDIA GPU上将DNN的执行性能分别提升高达3.8倍、2.6倍和1.7倍。
- DCGAN上的显著加速是因为其主要由转置2D卷积（T2D）构成，Ansor对此有很好的优化。
- AutoTVM在Intel CPU上的ResNet-50表现优异，得益于其高度优化的2D卷积模板和全局布局搜索。Ansor虽未进行全局布局搜索，但通过重写权重张量的布局（打包成更多层级），为ResNet-50带来了约40%的性能提升。
在三个硬件平台上的网络推理性能基准测试。y轴是每个网络相对于最佳吞吐量归一化的吞吐量。

消融研究: 如图10所示，对Ansor的变体进行测试。
- "Limited space" 最终性能最差，证明最佳程序不在受限空间内。
- "No fine-tuning" 扩大了搜索空间，性能有所提升，但在给定时间内仍无法超越AutoTVM。
- "No task scheduler" (使用轮询策略) 启用了微调后，性能超越了AutoTVM。
- "Ansor (ours)" 使用任务调度器优先处理性能瓶颈，因此在搜索效率和最终性能上都表现最佳。
网络性能自动调优曲线。y轴是相对于AutoTVM的加速比。

搜索时间

效率对比: Ansor的搜索效率很高，可以用更少的搜索时间达到或超过AutoTVM的性能。Ansor通过任务调度器同时优化所有子图，优先处理重要子图，避免了在不重要子图上浪费时间。

结果: 如表3所示，Ansor在Intel CPU网络基准测试中，可以用少一个数量级的搜索时间（无论是测量次数还是时钟时间）达到AutoTVM的性能。时钟时间节省更多，部分原因在于Ansor的搜索开销小，且在Intel CPU上的测量实现更优（通过显式刷新缓存减少重复次数）。
Ansor在Intel CPU上（批量大小=1）达到AutoTVM性能所用的测量次数和挂钟时间。

成本模型评估

评估设置: 使用调优ResNet-50时测量的25,000个程序作为数据集，其中20,000个用于训练，5,000个用于测试。

结果: 如图11所示，预测吞吐量与测量吞吐量高度相关。
- 图11a中，点分布在对角线周围，表明预测准确。
- 图11b显示了探索程序的性能分布。
- 模型在测试集上实现了0.079的RMSE，0.958的R2相关性，0.851的成对比较准确率，以及0.624的前30名程序召回率（recall@30）。
测量吞吐量与预测吞吐量对比。

补充细节

局限性与未来工作

动态形状: Ansor的一个局限性是它无法优化具有动态形状的图【45, Nimble: Efficiently compiling dynamic neural networks for model inference, 2020, arXiv】。Ansor要求计算图中的形状是静态且预先已知的，以便进行分析、构建搜索空间和执行测量。如何为符号或动态形状生成程序是一个有趣的未来方向。

稀疏算子: 另一个局限性是Ansor只支持稠密算子。为了支持稀疏神经网络【17, The state of sparsity in deep neural networks, 2019, arXiv】和图神经网络【25, Featgraph: A flexible and efficient backend for graph neural network systems, 2020, arXiv】中常用的稀疏算子（如SpMM），我们预计Ansor的大部分内容仍可复用，但需要重新设计搜索空间。

特殊指令利用: 最后，Ansor只在较高层次上执行程序优化，但依赖于其他代码生成器（如LLVM和NVCC）来进行平台相关的优化（如指令选择）。Ansor未能充分利用特殊指令，如Intel VNNI、NVIDIA Tensor Core和ARM Dot，用于混合精度和低精度算子，而这些目前现成的代码生成器处理得并不好。

结论

本文提出了Ansor，一个为深度神经网络生成高性能张量程序的自动化搜索框架。通过高效地探索一个大型搜索空间并优先处理性能瓶颈，Ansor找到了现有方法搜索空间之外的高性能程序。Ansor在多种神经网络和硬件平台上，其性能优于现有的手动库和基于搜索的框架，最高可达3.8倍。通过自动搜索更好的程序，我们希望Ansor能帮助弥合计算能力日益增长的需求与有限硬件性能之间的差距。Ansor已被集成到Apache TVM开源项目中。

PaperCache

Ansor: Generating High-Performance Tensor Programs for Deep Learning

Ansor: Generating High-Performance Tensor Programs for Deep Learning

主要贡献

背景知识

方法细节

设计总览

程序采样

草图生成

随机注解

GPU支持

性能微调

进化搜索

学习的成本模型

任务调度器

问题形式化

使用梯度下降进行优化

实验环境

实验结果

单算子基准测试

子图基准测试

端到端网络基准测试

搜索时间

成本模型评估

补充细节

相关工作

局限性与未来工作

结论

💬 评论讨论

Ansor: Generating High-Performance Tensor Programs for Deep Learning

主要贡献

背景知识

方法细节

设计总览

程序采样

草图生成

随机注解

GPU支持

性能微调

进化搜索

学习的成本模型

任务调度器

问题形式化

使用梯度下降进行优化

实验环境

实验结果

单算子基准测试

子图基准测试

端到端网络基准测试

搜索时间

成本模型评估

补充细节

相关工作

局限性与未来工作

结论

💬 评论讨论

登录

注册

忘记密码

重发验证邮件