AutoTVM：让AI来编译优化AI系统底层算子

导言

Alexnet 横空出世的两年之前，我选择了深度学习作为本科毕设方向。当时的深度学习并没有现在那么火热，也没有现有的GPU卷积实现，我在实验室的 GTX 470上面第一次手写GPU卷积用以支持卷积RBM，当时针对CPU十多倍的加速比让我兴奋不已。虽然最终那段研究经历并没有结果，但是计算本身对于机器学习的推动作用却深深地印在了我的脑海里。

深度学习系统支是推动进步的引擎，而在引擎的核心，则是像cuDNN这样的针对硬件的高效算子实现。每一个高效的算子库的背后都是大量工程的工程努力。虽然我自己已经从事了几年深度学习系统的工作，如何实现高效的硬件代码始终对于我来说是一个头疼的问题。最近的深度学习系统趋势开始进行更多的高层优化：包括探索从不同的数据表示，精度和算子融合等各个方面，而这些优化势必会带来更多的算子版本需要去优化，这无疑是一个非常大的工程问题。我们还需要把这个问题乘上包括移动芯片，移动GPU和AI加速器等等后端硬件的种类。针对每一个平台都去实现一遍底层的优化代价无疑会非常大。

我们希望利用机器学习本身去解决这个问题，以深度学习编译器TVM为基础自动化地优化深度学习算子性能。 写到这里，第一个大家会问的问题是为什么机器有可能会比手工做的优化要好。其实具体原因也很简单，一个人的精力是有限的，只能针对几个特定的场景去做一些优化。而机器有无限的精力可以针对每一种类型，数据表示，算子融合和实际的输入尺寸做特定的调整。另外像算子融合这样的优化必须要利用自动代码生成才可以做到。

所以机器在一定程度上可以取巧，在一些方面虽然未必像手工那样做到十全十美，但是在其它的一些方面反而可以通过自动针对每个特定算子尺寸的优化来反超手工优化（而人一般没有精力专门针对某一层的某一个算子去专门优化）。从目前来看的经验来看，手工优化的优势一般在于最里层的micro kernel，而对于程序总体的内存访问等把控并没有机器做的好。在一些必要的情况下TVM也允许直接插入手工micro kernel，但是依然把外层循环的把控交给机器来做，从而结合手工和机器优化的优势。我们在系统论文中的一个超低精度（2bit，1bit）算子的结果就是通过这种结合得到的。

如何可以让机器匹配手写优化性能

思路非常直接：

1）建立一个足够大的搜索空间，保证可能的人工手写优化全部包含在这个搜索空间里面

2）快速地搜索这个这个空间，获取优化的实现

这两个想法很简单，其实在两年之前我们基本上就有了这一个解决问题的蓝图，但是执行都非常困难。第一个问题是如何建立足够大的搜索空间，使得其可以包含人可能达到的手写优化。这个在之前TVM的介绍里面我们提到过，我们通过参考改进Halide的schedule编程源语，加入了对于GPU，加速器的抽象支持，通过一年的努力使得搜索空间基本可以匹配手写优化的性能。值得指出的是手工的优化可以非常多样性，并不是所有的优化都可以被简单的一个固定pattern总结。这也是之前很多自动代码生成框架无法达到比较好的效果的一个原因。我们也需要不断地总结抽象的优化规律，加入到TVM的原语中。另外在必要的时候我们也可以通过tensorize这个抽象可以插入一些手工优化的micro kernel，但是依然使用框架来做外层调度，达到最好的效果。

有个足够好的搜索空间，剩下的问题是如何在几十亿的可能性里面去选择比较好的实现。这里有几种常见的做法。传统的高性能计算库会采用Auto Tuning，也就是把可能的参数都尝试一遍。这样做的潜在问题是空间太大之后枚举开销过高。另外一种常见的做法是类似于数据库的做法，针对程序建立一个代价估价函数，然后利用估价函数来搜索。这个做法可能碰到的主要问题是估价函数不一定可以估计准确，并且针对每个新的硬件特性必须要重新设计估价函数。

我们的利用机器学习来学习程序空间的代价估价函数。具体地说， 探索程序在一开始会随机地选取一些设定，直接到硬件上面去运行生成的代码，再通过得到的反馈数据来更新我们的程序代价估计函数。这里面比较有趣的一点是模型的可迁移性。因为真正的深度学习系统需要优化许多不一样输入类型，输入形状的算子。一个可迁移的模型可以通过学习已经看到过的算子优化记录来预测新的目标的代价，导致最后的搜索时间可以大幅降低。

其实在这个整个优化的过程中机器学习算法只是起到了重要但是只是一部分的作用。在一些如ARM我们的先验模版比较强的情况下，随机搜索就可以达到非常不错的效果。但是重要的是我们需要建立集群化，自动化，python可以对接的框架来做这些实验。这里特别提一下TVM最近引入的RPC tracker，支持了这一个功能，才可以使得我们可以在多种设备上进行实验。

论文链接

对于具体实验效果和方法介绍有兴趣的同学可以看论文里面的结果

- 机器学习：https://arxiv.org/abs/1805.08166

- 系统： https://arxiv.org/abs/1802.04799

未来展望

在可见的未来里面我们还是需要高性能优化的知识不断充实完善TVM。手写算子优化和自动优化可能会在比较长的一段时并存下去。但是我们个人相信更多的自动化和半自动化也会成为主流。文章中提到的搜索空间，RPC集群支持都已经正式加入到TVM中去。自动化模块也会在下一个版本和大家见面。

围绕这个问题，这个方向还留下包括模型表示，迁移等在内还有很多未知的问题等待着大家去探索，也欢迎更多的小伙伴加入TVM社区。

后记

后续我们也会继续在这一专栏介绍TVM深度学习自动优化等相关方向的内容，欢迎有兴趣的同学投稿。本文欢迎转载。包括本次工作在内的TVM工作是整个社区团队一起努力的结果（而并非第一作者作者个人），转载本专栏文章时请注明TVM社区。

本次机器学习相关工作主要成员来自华盛顿大学，上海交大和复旦大学。我本人十分感谢交大acm班学习经历对于这次工作的影响：十年之前的这个时候我在交大acm班正好在做编译原理大作业。本次论文的主力之一，第二作者郑怜悯也来自acm班，是今年编译原理的助教。我们当时编译原理课程主现在依然挂在李沐在交大的主页上 http://bcmi.sjtu.edu.cn/home/limu/tiger/