AutoTVM调优算法

DungDaj

出自论文《TVM: An Automated End-to-End Optimizing Compiler for Deep Learning》2018和《Learning to Optimize Tensor Programs》2018，前者描述了TVM的整体框架，后者则主要介绍了AutoTVM使用的调优算法。部分内容参考了博客Learning to Optimize Tensor Programs解读 - 知乎 (zhihu.com)，感谢知识分享。
1. 目标和背景

利用TVM的调度原语，建立一个调度搜索空间，TVM利用基于机器学习的优化器从搜索空间中找到一组高性能的优化方法，再将优化后的算子部署到后端设备，而这种基于机器学习的优化器就是AutoTVM
2. AutoTVM的框架和优化原理

2.1 框架和组成




AutoTVM的框架

e：计算高层描述

：计算底层执行方式（调度）
Exploration Module：核心模块
history data：用于对cost model更新
Cost Model：用于快速筛选调度
Objective Function：用于衡量cost model的优劣，从而更新cost model
Code Generator：用来生成底层代码，在实际硬件上运行
2.2 调度空间的确定

使用TVM调度原语来构建调度空间Se，包括：

• 各个循环层级上的tiling
  
• 循环的reordering
  
• GPU shared memory的caching
  
• unrolling和vectorization
  

【从Ansor论文】我们知道其实autoTVM的调度空间是通过人工来手写
2.3 算法流程



具体执行：

• 确定一个迭代上界
  
• 【cost model初筛】用Cost Model去衡量调度的优劣，并用模拟退火算法选择候选调度，加入集合Q
  
• 【$L(x)$二筛】用贪心算法选择$(1-ε)b$个使得$L(x)$最大的，加入集合S
  
• 【Diversity】随机选择调度$εb$个加入S
  
• 【实际硬件执行】对S中每一个s都在硬件上进行执行，得到运行时间c，然后将(e, s, c)加入数据集D
  
• 【完善cost model】利用数据集D Update cost model
  
• 按照batch来遍历整个搜索空间
  
• 最后得到最佳调度s*
  

2.4 cost model

有这两种cost model：基于XGBoost的Gradient Tree Boosting Model(GTBM)和TreeGRU（另外一篇论文中叫做TreeRNN）。GTB需要对AST进行处理，提取出程序中的一些feature，然后才能进行分析，但是它可以较快的运行在CPU上。而TreeGRU可以直接喂AST，不需要对程序做处理，但是它在训练和推断的速度较慢，需要GPU来加速。


2.5 Objective Function

regression loss function：


由于其实只需要cost function能够分辨出调度之间运行时间的快慢关系，因此可以用下面的rank loss function：


2.6 L(x)



3. 迁移学习

在前期缺少历史数据的条件下提高cost model的精确度。


为了提高GTBM在不同算子之间的泛化能力，解决不同算子的feature不同的问题，作者提出了Context Relation Features，即使之前学习的是其他类型的算子，也可以提供迁移学习。针对TreeGRU也提出了类似的编码方式来解决泛化问题。
4. 实验结果

基于Resnet18和ImageNet，
4.1 对比黑盒优化结果

GTBM和TreeGRU的效果相比GA和随机优化的黑盒优化算法效果都要更好


4.2 结合贝叶斯优化的效果

作者使用bootstrapping去获得cost model的uncertainty estimate，然后用expected improvement (EI) and upper confidence bound (UCB)方法来做，效果似乎也不明显（这一块不是很了解，但是由于之前有讨论过用贝叶斯优化来做，所以放一下结果。）


4.3 迁移学习对比不用迁移学习的效果



用其他算子的历史数据来做迁移学习，会比不用更快得到更好的结果。
5. 其他工作

文中提到了以下这些工作
ATLAS《Automatically tuned linear algebra software》（1998）和FFTW《Fftw: an adaptive software architecture for the fft》（1998）是两篇黑盒优化的工作，《Automatically scheduling halide image processing pipelines》（2016）和《A practical automatic polyhedral parallelizer and locality optimizer》（2008）是model-based的工作。多面体优化的工作有《Polyhedral parallel code generation for cuda》（2013）和TC这种