论文精读：Ansor: Generating High-Performance Tensor Programs for Deep Learning

Posted 2023-03-09 夏小悠

文章目录

• • 1. Abstract
  • 2. Introduction
  • 3. Background
  • 4. Design Overview
  • 5. Program Sampling
  • • 5.1 Sketch Generation
    • 5.2 Random Annotation
  • 6. Performance Fine-tuning
  • • 6.1 Evolutionary Search
    • 6.2 Learned Cost Model
  • 7. Task Scheduler
  • • 7.1 Problem Formulation
    • 7.2 Optimizing with Gradient Descent

1. Abstract

  高性能张量程序是保证深度神经网络有效执行的关键。然而，在各种硬件平台上获得不同算子的性能较好的张量程序是非常具有挑战性的。目前，深度学习系统依赖于硬件供应商提供的内核库(kernel libraries)或各种各样的搜索策略来获取性能较好的张量程序。但这些方法有两个弊端：
  (1)需要巨大的工程量来开发平台特定的优化代码；
  (2)有限的搜索空间和无效的搜索策略导致很难发现高性能张量程序。
  作者基于上述弊端，提出了Ansor，一个用于深度学习应用的张量程序生成框架。对比现有的搜索策略，Ansor有以下特征：
  (1)通过从搜索空间的分层表示(hierarchical representation)中采样程序以探索更多的优化组合(optimization combinations)；
  (2)使用进化搜索(evolutionary search)和学习成本模型(cost model)对采样程序进行微调(fine-tune)，以确定最优的程序；
  (3)利用任务调度器(task scheduler)来同时优化深度神经网络的多个子图。
  作者的实验表明，Ansor能够找到现有最先进(state-of-the-art，SOAT)方法搜索空间之外的高性能程序：是Intel CPU的3.8倍，ARM CPU的2.6倍，NVIDIA GPU的1.7倍。

2. Introduction

  深度神经网络(DNN)的低延迟执行在自动驾驶(autonomous driving)、增强现实(augmented reality)、语言翻译(language translation)及其他的AI应用中发挥着至关重要的作用。DNN可以表示为一个有向无环计算图(directed acyclic graph, DAG)，结点表示算子(卷积，矩阵乘)，有向边表示算子之间的依赖关系。现有的深度学习框架(Tensorflow, PyTorch, MXNet)将DNN中的算子映射为供应商提供的内核库(cuDNN, MKL-DNN)以获取高性能。然而，这些内核库需要巨大的工程量为每个硬件平台和算子进行手动调优，为每个目标加速器产生有效的算子实现所需的大量手工工作限制了新算子和特定加速器的开发和创新。
  鉴于DNN性能的重要性，研究者和行业从业者已经转向基于编译器搜索(search-based compilation)来自动生成张量程序，比如张量算子的低级实现。对于一个算子或者多个算子的子图，用户需要用高级声明性语言来定义计算，然后编译器搜索针对不同硬件平台的定制程序。

3. Background

  深度学习生态系统正在拥抱快速增长的硬件平台多样性，包括CPU、GPU、FPGA和ASIC。为了在这些平台上部署DNN，需要为DNN中使用的算子提供高性能张量程序，所需的算子集通常包含标准算子(matmul, conv2d)和机器学习研究人员发明的新算子(capsule conv2d, dilated conv2d)。
  为了以高效的方式在广泛的硬件平台上提供这些算子的可移植性，多种编译器技术已经出现(TVM, Halide, Tensor Comprehensions)。用户使用高级声明性语言以类似数学表达式的形式定义计算，编译器根据定义生成优化的张量程序。下图显示了TVM张量表达式语言中矩阵乘法的计算定义，用户主要需要定义输入张量的形状，以及如何计算输出张量中的每个元素。

  然而，从高级定义自动生成高性能张量程序是极其困难的。根据目标平台的体系结构，编译器需要在一个包含优化组合选择的极其大而复杂的空间中进行搜索(例如，展开结构(tile structure)，展开大小(tile size)，向量化(vectorization)，并行化(parallelization))，寻找高性能的程序需要搜索策略覆盖一个全面的空间，并有效地探索它。

4. Design Overview

  Program sampler：Ansor必须解决的一个关键挑战是为给定的计算图生成大的搜索空间。为了覆盖具有各种高级结构和低级细节的各种张量程序，Ansor利用了具有两个级别的搜索空间的分层表示:草图(sketch)和注释(annotation)。Ansor将程序的高级结构定义为草图，并将数十亿个低级选择(例如，平铺大小(tile size)、并行(parallel)、展开注释(unroll annotations))作为注释，这种表示法允许Ansor灵活地枚举高级结构并有效地采样低级细节。
  Performance tuner：随机抽样程序的性能不一定好，下一个挑战是对它们进行微调。Ansor采用进化搜索和学习成本模型来迭代地执行微调，在每次迭代中，Ansor使用重新采样的新程序以及以前迭代中的好程序作为初始种群来开始进化搜索。进化搜索通过变异和交叉对程序进行微调，执行乱序重写并解决顺序构造的限制。查询学习到的成本模型比实际测量快几个数量级，因此我们可以在几秒钟内评估数千个程序。
  Task scheduler：使用程序采样和性能微调允许Ansor为计算图找到高性能张量程序。直观地说，处理一个完整DNN作为一个单一的计算图，并为其生成一个完整的张量程序，可以潜在地实现最佳性能。然而，这是低效的，因为它必须处理搜索空间不必要的指数爆炸。通常，编译器将DNN的大计算图划分为几个小的子图，由于DNN的逐层(layer-by-layer)构造特性，这种划分对性能的影响可以忽略不计，这就带来了Ansor的最后一个挑战：在为多个子图生成程序时如何分配时间资源？
  Ansor中的任务调度器使用基于梯度下降的调度算法将资源分配给更有可能提高端到端DNN性能的子图。

5. Program Sampling

  算法探索的搜索空间决定了它能找到的最佳程序。现有方法所考虑的搜索空间受到以下因素的限制:
  (1)手动枚举(TVM)：通过模板手动枚举所有可能的选择是不切实际的，因此现有的手动模板只能启发式地覆盖有限的搜索空间；
  (2)积极的早期剪枝(Halide auto-scheduler)：基于评估不完整程序的激进早期修剪阻止搜索算法探索空间中的某些区域。
  为了解决(1)，作者通过递归应用一组灵活的推导规则来自动扩展搜索空间；
  为了避免(2)，作者在搜索空间中随机抽样完整的程序。
  由于随机抽样给每个被抽样的点机会是相等的，作者提出的搜索算法可以潜在地探索所考虑空间中的每个程序，不依赖于随机抽样来找到最优程序，因为每个抽样程序后来都经过了微调。
  在顶层，通过递归地应用一些派生规则来生成草图。在底层，随机地注释这些草图以得到完整的程序。这种表示从数十亿个低级选择中总结了一些基本结构，从而实现了对高级结构的灵活枚举和对低级细节的高效采样。

5.1 Sketch Generation

  上图中的第一列显示了两个输入示例。输入有三种等效形式：数学表达式、直接展开循环指标得到的相应naive程序和相应的计算图(DAG)。

  在计算机编程领域中，"naive program"通常是指一种简单或者朴素的程序实现方式。这种程序可能没有考虑所有可能的情况，或者没有利用现有的优化技术。"naive"这个词通常用来描述某些程序员在编写代码时缺乏经验或技术水平较低的情况。在这种情况下，程序员可能会使用一些基本的算法或数据结构，而没有考虑到更复杂或高效的解决方案。这种程序通常会占用大量的计算资源，运行速度较慢。------by ChatGPT

  为了给具有多个节点的DAG生成草图，我们以拓扑顺序访问所有节点，并迭代地构建结构。对于计算密集型和有大量数据重用机会的计算节点(conv2d, matmul)，我们为它们构建基本的平铺和融合结构作为草图，对于简单的元素节点(ReLU, elementwise add)，我们可以安全地内联它们。注意，新节点(缓存节点(caching nodes)，布局转换节点(layout transform nodes))也可以在草图生成过程中引入DAG。
  作者提出了一种基于派生的枚举(derivation-based enumeration)方法，通过递归应用几个基本规则来生成所有可能的草图。这个过程以DAG作为输入，并返回草图列表。我们定义State $\sigma =(S;i)$，其中 $S$ 是DAG当前部分生成的草图，$i$ 是当前工作节点的索引，DAG中的节点按照从输出到输入的拓扑顺序进行排序。推导从初始naive程序和最后一个节点开始，或者初始状态$\sigma =(naiveprogram;indexofthelastnode)$，然后我们尝试递归地将所有推导规则应用于这些状态。对于每条规则，如果当前状态满足应用条件，我们应用这条规则$\sigma =(S;i)$得到$\sigma \prime =(S\prime ;i\prime ),i\prime <i$，这样，索引$i$(工作节点)单调地减小，当$i=0$时，一个状态就变成了终端状态。在枚举过程中，可以对一个状态应用多个规则，从而生成多个后续状态，一个规则还可以生成多个可能的后续状态。因此，我们维护一个队列来存储所有中间状态，当队列为空时，进程结束。所有处于终端状态的$\sigma .S$在草图生成结束时形成草图列表。对于典型的子图，草图的数量小于10。

// 递归应用几个基本规则来生成所有可能的sketch
// Derivation rule based enumeration
Array<State> out_states;
while (!pnow->empty()) 
  pnext->clear();
  for (const State& state : *pnow) 
    int stage_id = cur_stage_id_map[state];

    // Reaches to the terminal stage
    if (stage_id < 0) 
      out_states.push_back(state);
      continue;
    

    // Try all derivation rules
    for (const auto& rule : sketch_rules) 
      auto cond = rule->MeetCondition(*this, state, stage_id);
      if (cond != SketchGenerationRule::ConditionKind::kSkip) 
        for (const auto& pair : rule->Apply(*this, state, stage_id)) 
          cur_stage_id_map[pair.first] = pair.second;
          pnext->push_back(pair.first);
        
        // Skip the rest rules
        if (cond == SketchGenerationRule::ConditionKind::kApplyAndSkipRest) 
          break;
        
      
    
  
  std::swap(pnow, pnext);

// Conv2d(3, 64, kernel_size=(7, 7), stride=2, padding=1)有3个sketch生成

  Derivation rules：上述表格列出了用于CPU的派生规则。作者首先提供所使用谓词的定义，然后描述每个规则的功能，然后对计算定义执行静态分析，以获得这些谓词的值，分析是通过解析数学表达式中的读/写模式自动完成的。我对上述表格进行了整理：

Condition	Description
$IsStrictInliable(S,i)$	表示$S$中的节点$i$是一个简单的按元素(element-wise)计算的算子，比如element-wise add和ReLU
$HasDataReuse(S,i)$	表示$S$中的节点$i$是计算密集型(compute-intensive)算子，并且具有大量的算子内数据重用机会，比如matmul和 conv2d
$HasFusibleConsumer(S,i)$	表示$S$中的节点$i$只有一个消费者节点$j$，节点$j$可以融合到节点$i$中，比如matmul + bias_add和conv2d + relu
$HasMoreReductionParallel(S,i)$	表示$S$中的节点$i$在空间维度上并行性很小，但在降维上有足够的并行机会，比如计算矩阵的L2范数，矩乘$C_{2\times 2}=A_{2\times 512}\cdot B_{512\times 2}$

  在计算机编程中，"inline"通常指的是一种编译器优化技术，即在编译代码时将函数调用直接替换为函数体内的代码。这样可以避免函数调用时的额外开销，从而提高代码的执行效率。
  在C++中，我们可以使用关键字"inline"来告诉编译器，将某个函数作为inline函数来处理。在C++程序中使用inline函数的好处是可以减少函数调用的开销，从而提高程序的运行效率。此外，使用inline函数还可以减少代码的重复，因为每次调用该函数时都会将函数的代码嵌入到调用位置。
  需要注意的是，虽然使用inline函数可以提高程序的性能，但并不是所有函数都适合作为inline函数。一般来说，较小的、频繁调用的函数最适合作为inline函数，而较大的、复杂的函数则不适合作为inline函数。此外，inline函数可能会增加代码的体积，因此需要在代码大小和性能之间进行权衡。------by ChatGPT

  Rule 1只是简单地跳过一个节点，如果这个节点不是严格内联的；
  Rule 2始终是严格内联节点，由于Rule1和Rule2的条件是互斥的，$i>1$的状态总是可以满足其中一个条件并继续推导；
  Rule 3是为数据可重用节点执行多级平铺。对于CPU，我们使用"SSRSRS"平铺结构，其中"S"代表一个平铺级别的空间循环(space loop)，"R"代表一个平铺级别的缩减循环(reduction loop)。例如，在矩乘$C(i,j)={\sum }_{k}A[i,k]\times B[k,j]$， $i$和$j$是空间环，$k$是缩减环。矩乘的"SSRSRS"平铺结构将原来的3级循环$(i,j,k)$扩展为一个10级循环$(i_0,j_0,i_1,j_1,k_0,i_2,j_2,k_1,i_3,j_3)$，虽然没有打乱循环顺序，但这种多级平铺也可以覆盖一些重新排序的情况。例如，上面的10级循环可以专门用于简单的重排序$(k_0,j_2,j_3)$通过设置其他循环的长度为1。"SSRSRS"平铺结构一般用于深度学习中的计算密集型密集算子(matmul, conv2d, conv3d)，因为它们都由space loop和reduction loop组成；
  Rule 4是执行多级平铺，还融合了可融合的消费者。例如，可以将按元素划分的节点(ReLU，bias_add)融合到平铺节