Ansor: 为深度学习生成高性能张量程序

1、背景

深度神经网络(DNN)的低延迟推理在自动驾驶、增强现实、语言翻译等应用中起着至关重要的作用。为了获得高性能，现有的深度学习框架（如TensorFlow、PyTorch、MXNet）基于硬件厂商提供的内核库（如cuDNN、MKL-DNN）来实现DNN中的算子。然而，这些内核库针对于不同硬件和算子均需要大量的工作来开发特定的优化代码，严重限制了新算子的开发。

为了以高效的方式在各种硬件平台上部署这些算子时都保持高性能，已经引入了多种编译器技术（如TVM、Halide、Tensor Comprehensions）。用户使用高级声明式语言以类似于数学表达式的形式定义计算，编译器根据定义生成优化的张量程序。

然而，从高级定义中自动生成高性能张量程序是极其困难的。根据目标平台的架构，编译器需要在包含优化组合选择（如tile structure, tile size, vectorization, parallelization）的极其庞大和复杂的空间中进行搜索。寻找高性能程序需要搜索策略覆盖一个全面的空间并进行有效地探索。其中，两种最新且有效的搜索方法为模板引导搜索和基于序列结构的搜索。

在模板引导搜索中，搜索空间由手动模板定义。如图1a所示，编译器（例如，TVM）要求用户手动为计算定义编写模板。该模板使用一些可调参数（如tile size和unrolling factor）定义了张量程序的结构。然后编译器为特定的输入形状和硬件目标搜索这些参数的最佳值。该方法在常用的深度学习算子上取得了很好的效果。但是，开发模板需要大量的工作，并且构造一个高质量的模板需要张量算子和硬件方面的专业知识。此外，手动模板仅涵盖有限的程序结构，因为手动枚举所有算子的所有优化选择是几乎不可能的。

基于序列结构的搜索通过将程序结构分解为一系列固定的决策来定义搜索空间。然后，编译器使用诸如beam search之类的算法来搜索好的决策（如Halide auto-scheduler）。在这种方法中，编译器通过依次展开计算图中的所有节点来构造张量程序。对于每个节点，编译器会对如何将其转换为低级张量程序做出一些决策（如computation location, storage location, tile size）。当所有节点展开时，构造一个完整的张量程序。这种方法为每个节点使用一组通用的展开规则，因此它可以自动搜索而不需要手动模板。由于每个决策的可能选择数较大，为了使序贯过程可行，该方法在每个决策后只保留前k个候选方案。编译器基于一个可学习的cost model评估和比较候选程序的性能，以选择前k个候选程序。在搜索过程中，候选程序是不完整的，因为只展开了部分计算图或做出了部分决策。图1b显示了这个过程。但是，在完整程序上训练的cost model不能准确地预测不完整程序的最终性能，并且顺序决策的固定顺序限制了搜索空间的设计。

Ansor是基于分层搜索空间构建的，如图1c所示，该空间将高级结构和低级细节解耦。Ansor自动构造计算图的搜索空间，无需手动开发模板。Ansor从空间中采样完整的程序，并对完整的程序进行微调，避免了对不完整程序的粗略估计。图1显示了Ansor方法和现有方法之间的主要区别。

2、设计概述

Ansor是一个自动生成张量程序的框架。图2显示了Ansor的总体架构。Ansor的输入是一组待优化的DNN模型。Ansor使用Relay的算子融合算法将DNN从流行的模型格式（如ONNX，TensorFlow PB）转换为小的子图。然后，Ansor为这些子图生成张量程序。Ansor有三个主要组成部分：（1）程序采样器，它构造了一个大的搜索空间并从中抽取不同的程序；(2）性能调试器，用于微调所采样程序的性能；(3）任务调度器，为优化DNN中的多个子图分配时间资源。

3、架构详解

3.1 Program Sampler

算法探索的搜索空间决定了它可以找到的最佳程序。现有方法中考虑的搜索空间受到以下因素的限制：（1）手动枚举（如TVM）。通过模板手动枚举所有可能的选择是不切实际的，因此现有的手动模板只能启发式地覆盖有限的搜索空间。（2）过度早期修剪（如Halide auto-scheduler）。基于评估不完整程序的过度早期修剪阻止了算法探索空间中的某些区域。

为了解决上述问题，Ansor通过递归地应用一组灵活的推导规则来自动扩展搜索空间，并在搜索空间中随机采样完整的程序。由于随机采样为每个要采样的点提供了均等的机会，因此搜索算法可以潜在地探索所考虑空间中的每个程序。不依赖随机采样来找到最佳程序，因为每个采样程序后还会进行微调。

为了采样可以覆盖大搜索空间的程序，Ansor定义了一个具有两个级别的分层搜索空间：sketch和annotation。将程序的高级结构定义为sketches，并将数十亿个低级选择（如tile size、unroll annotations）作为annotations。在顶层，通过递归应用一些推导规则来生成sketch，在底层，随机注释这些sketch以获得完整的程序。这种表示从数十亿个低级选择中总结了一些基本结构，实现了高级结构的灵活枚举和低级细节的有效采样。

3.1.1 草图生成

Ansor提出了一种基于推导的枚举方法，通过递归地应用几个基本规则来生成所有可能的sketches。此过程将DAG作为输入并返回草图列表。定义状态函数，其中是DAG当前部分生成的草图， 是当前工作节点的索引。DAG中的节点按从输出到输入的拓扑顺序排序。推导从初始naive程序和最后一个节点开始，将所有推导规则递归地应用于状态。对于每个规则，如果当前状态满足应用条件，将规则应用到并得到其中。这样索引（工作节点）单调递减，当 时，状态变为终止状态。在枚举过程中，可以对一个状态应用多个规则以生成多个后续状态。一个规则还可以生成多个可能的后续状态，所以维护一个队列来存储所有中间状态，当队列为空时该过程结束。所有处于终止状态的在草图生成结束时形成一个草图列表。一般来说，子图生成的草图数量小于10。

图3列出了用于CPU的推导规则。表示中的节点是否是一个可以内联的简单element-wise算子（如element-wise add，ReLU）。表示中的节点是否为计算密集型算子且具有大量算子内数据重用机会（如matmul、conv2d）。表示中的节点是否只有一个消费者且节点可以融合到节点i中（如matmul + bias_add、conv2d + relu）。表示中的节点是否在空间维度上无法并行，但在reduction维度上有足够的并行机会（如计算矩阵的乘法 ）。对计算定义执行静态分析以获得这些述语的值。分析是通过解析数学表达式中的读/写模式自动完成的。

3.1.2 随机注释

生成的草图是不完整的程序，因为它们只有tile结构，没有特定tile尺寸和循环注释，例如parallel, unroll, and vectorization。需要对草图进行注释，使它们成为用于微调和评估的完整程序。在生成的草图列表中，随机选择一个草图，填充tile sizes，并行化一些外循环，矢量化一些内循环，并展开一些内循环。此外，随机改变程序中一些节点的计算位置，对tile结构进行了微调。所有“随机”表示所有有效值的均匀分布。如果某些特殊算法需要自定义注释才能生效（如，特殊展开），允许用户在计算定义中给出简单的提示以调整注释策略。最后，由于更改常量张量的布局可以在编译时完成，并且不会带来运行时开销，因此根据多级tile结构重写常量张量的布局，使其尽可能缓存友好。这种优化十分有效，因为卷积或全连接层的权重张量是静态张量。

3.2 Performance Tuner

程序采样器采样的程序对搜索空间有很好的覆盖，但质量得不到保证。这是因为优化选择，例如tile结构和循环注释，都是随机采样的。性能调试器通过进化搜索和可学习的cost model微调采样程序的性能。

微调是迭代执行的。在每次迭代中，首先使用进化搜索根据学习到的cost model找到一小批性能较好的程序。然后在硬件上测量这些程序的实际执行时间。最后，使用测量得到的分析数据重新训练cost model，使其更加准确。

进化搜索使用随机抽样的程序以及上次评估中的高质量程序作为初始种群，并应用变异和交叉来生成下一代。可学习的cost model用于预测每个程序的性能，如吞吐量。运行固定代数的进化，并选择在搜索过程中的最佳程序。因为cost model可以给出相对准确的程序性能估计，同时比实际测量快几个数量级，它使我们能够在几秒钟内比较搜索空间中的数万个程序，并挑选出性能不错的程序进行评估。

其中，cost model使用加权平方误差作为损失函数。因为最关心的是从搜索空间中识别出性能良好的程序，所以重视运行速度更快的程序。具体来说，模型f在吞吐量为y的程序P上的损失函数为

其中是中最内层非循环语句的集合，使用吞吐量作为权重。训练一个梯度提升树作为基础模型，为来自所有DAG的所有张量程序训练了一个模型，将来自同一DAG的所有程序的吞吐量归一化在[0,1]的范围内。在优化DNN时，测量的程序数量通常小于30,000。在如此小的数据集上训练梯度提升决策树的速度非常快，因此每次都训练一个新模型，而不是进行增量更新。

3.3 Task Scheduler

一个DNN可以划分为许多独立的子图（例如，conv2d + relu）。对于某些子图，花时间调整它们并不会显着提高端到端DNN的性能。这是因为：（1）子图不是性能瓶颈，（2）调优对子图的性能带来较小的改进。为了避免在调整不重要的子图上浪费时间，Ansor会动态地为不同的子图分配不同的时间资源。Ansor的任务调度器以迭代的方式为任务分配时间资源。在每次迭代中，Ansor选择一个任务，为子图生成一批性能较好的程序，并在硬件上评估程序。为了最小化单个DNN的端到端延迟，定义目标函数为

其中为任务总数，为时间分配向量，是任务在DNN中的出现次数。Ansor采用基于梯度下降的调度算法来有效地优化目标函数，在给定当前分配的情况下，其思想是近似目标函数的梯度，以便选择任务使得。近似梯度为

其中是一个小的后向窗口大小，和可以从分配的历史中得知。是的相似任务的集合，是任务i中浮点运算的次数，是我们在任务中每秒可以实现的浮点运算次数。参数和控制权重。

4、评估

4.1单个算子的基准测试

在一组常见的深度学习算子上评估Ansor，包括1D、2D和3D卷积（分别为C1D、C2D和C3D）、矩阵乘法(GMM)、组卷积(GRP)、扩张卷积(DIL)、深度卷积(DEP)、转置2D卷积(T2D)、胶囊2D卷积(CAP)和矩阵2-范数(NRM)。对每个算子选择4个常见的形状，并使用两个批次大小（1和16）进行评估。如图4所示，Ansor在所有算子和Batchsize设置中，性能均是最优的。Ansor的性能比现有搜索框架高1.1-22.5倍。Ansor的性能提升来自其庞大的搜索空间和有效的探索策略。对于大多数算子，发现Ansor生成的最佳程序在现有搜索框架的搜索空间之外，因为Ansor能够探索更多优化组合。

4.2 子图的基准测试

在DNN中的两个常见子图上推理子图基准测试。“ConvLayer”是由2D卷积、批量归一化和ReLU激活组成的子图。“TBS”是一个由两个矩阵转置、一个批量矩阵乘法和一个softmax组成的子图。选择四种不同的形状和两种批量大小，图5显示Ansor的性能比手动库和其他搜索框架高1.1-14.2倍。Ansor可以在两个平台上为这些子图一致地生成高性能程序。FlexTensor对单个算子表现良好，但在子图上的优势较小，是因为它缺乏算子融合的支持。

4.3 端到端网络的基准测试

对多个DNN的端到端推理进行了基准测试，其中包括ResNet-50、MobileNet-V2、3D-ResNet18、DCGAN和BERT。在三个硬件平台上对这些DNN进行了基准测试，对于Intel CPU 和 NVIDIA GPU，测试了Batchsize为1和16的结果，对于ARM CPU，只测试了Batchsize为1的结果。图6显示了在Intel CPU、NVIDIA GPU和ARM CPU上的结果，Ansor在所有情况下的性能比基于搜索的AutoTVM高1.0-21.8倍。与最佳替代方案相比，Ansor将DNN在Intel CPU、ARM CPU和NVIDIA GPU上的推理性能分别提高了3.8倍、2.6倍和1.7倍。

4.4 搜索时间

Ansor搜索效率高，可以在更短的搜索时间内超越或匹配AutoTVM。如图7所示，Ansor可以在搜索时间减少一个数量级的情况下与AutoTVM的性能相匹配。在图7a中，搜索时间的节省来自任务调度程序、高效的微调和有效优化的全面覆盖。在图7b中，Ansor显示挂钟时间更省时，这是因为Ansor不会引入太多的搜索开销。