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

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

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

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

思路非常直接：

1. 建立一个足够大的搜索空间，保证可能的人工手写优化全部包含在这个搜索空间里面 
   
2. 快速地搜索这个这个空间，获取优化的实现
   

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

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

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

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