陈天奇：深度学习编译技术的现状和未来

来源：机器之心

传统的深度学习框架采用人工优化算子，然后建立运行时图解释器来解决内存分配调度等问题。深度学习编译器技术路线一般指在优化过程中采用了自动或者半自动的代码生成用以替代人工优化。深度学习编译器无疑是最近非常热门的话题。本文主要探讨深度学习编译技术的现状和未来。

为什么需要深度学习编译器

 深度学习编译器的部署目标传统的深度学习框架也可以做，一个非常自然的问题是为什么不直接沿用传统的框架。这是一个编译器研究者来往往会忽略的问题。深度学习编译器只有在各种场景超过人工优化的传统办法，才有机会真正被采用，到达这一目标之前之前深度学习编译只是玩具。

 从目前的现状来看，深度学习编译器TVM已经一定程度上到达了这一目标。在一些部署场景下，深度学习编译已经到达了可以和传统框架一拼高下或者超越传统框架的阶段。随着深度学习自动化编译研究的进展手写优化的经验会被融入到编译器中，从而把逐步替代传统的方案。为什么会如此呢？总结下来的核心是编译器可以带来的更多自动化。细化来说分以下几点：

无限的算力和有限的精力

如果给定一个特定的算子，一个工程师无疑可以做到非常精细的地步，通过精细地选择流水线，指令集，预读，寄存器分配来到达接近peak的效果。深度学习编译器需要一个大的搜索空间涵盖手写优化的方案。一旦搜索空间定义足够大，从而接近人工。很多人会觉得，一个足够精力的工程师一定可以超过编译，这一点无疑是对的。那么机器（自动编译）的优势在什么地方呢？机器的本质优势在于其强大的算力，可以针对目标网络的每一层，特定输入大小场景，进行专门的优化。而人的精力有限，一般会尝试优化常见的瓶颈算子，而这种有针对性的优化未必适用于网络的每一层，或者企业相关的应用场景。通过（接近无限）的算力去适配每一个应用场景看到的网络，这是深度学习编译器比人工路线强的地方。

当比较TVM和传统方法的时候的时候，我们往往会发现：在标准的benchmark（如imagenet resnet）上，编译带来的提升可能只在10%到20%，但是一旦模型相对不标准化（如最近的OctConv，Deformable, 甚至是同样的resnet不同的输入尺寸），编译技术可以带来的提升会非常巨大。原因也非常简单，有限的精力使得参与优化的人往往关注有限的公开标准benchmark，但是我们的部署目标往往并非这些benchmark，自动化可以无差别地对我们关心的场景进行特殊优化。接近无限的算力和有限的精力的差别正是为什么编译技术一定会越来越重要的原因。

编译器可以站在不一样的起跑线上

如果比拼优化一个固定某一层的3x3卷积，往往精心的手工优化是机器未必可以超越的。那么在一个标准的benchmark如resnet上，为什么编译技术还是可以超过人呢？其原因是往往编译器可以和人站在不一样的起跑线上。

举一个简单的例子，如果优化一个n层的MLP，我们要优化矩阵乘法。从人工优化的角度来说，往往我们会限制优化目标为，优化一个行优先（row-major）存储的矩阵乘法。但是对于神经网端到端络本身而言，内部到底是行优先或者是列优先，或者是（NCHWc4, NCHWc8）都是可以的，而且往往对于每一层最好的排布也不一定相同。同样的，我们往往可以把算子和后面几层的各种算子融合起来，或者针对每一层选择有利于其的量化方案。还是因为精力有限的缘故，人工优化的库往往会把全局问题限制在一些子问题（如行优先的矩阵乘法）上。如果编译器和工程师硬碰硬直接解决同样的子问题，或许并不能讨好。但是自动化后的编译器可以直接去考虑更大的解决空间，去自动选择更加高效的数据排布或者算子融合，从站在了不一样的起跑线上。这也是为什么即使在标准benchmark，只要编译可以做到人工的80%到90%，更好的起跑线带来的优势会掩盖单个部分微弱的劣势。

编译器和手工优化结合

最后，编译的目标并非替代手工优化，而是吸收手工优化的经验，使得优化更加自动化。一个笨编译器比不过聪明的脑子。怎么办呢，我们需要从实用主义的角度出发给编译技术提供定制，允许在需要的加入手工优化来助力编译器。

其中最简单的办法当然是直接把一些层offload给类似于cudnn这样的库。这也是XLA等在内的工具采取的技术路线。更进一步，TVM允许手工提供微内核（micro kernel）用于优化4x4外积等，但是依然采用自动优化的办法优化内存排布和loop。达到手工和自动化相结合的目标。类似的，用户可以通过构造特定的搜索空间模版来加入人工信息。

总结下来，深度学习编译器之所以可以到达今天的高度，正是得益于深度学习优化工程师的经验总结。接下来的一段时间，我们应该会看到越来越多优化工程师加入到深度学习编译器建设中去，使得两条路线逐渐融为一体。

深度学习编译和传统编译的技术路线差别

在优化需求上深度学习编译器和传统编译器有很大的差别。传统编译器注重于优化寄存器使用和指令集匹配，其优化往往偏向于局部。而深度学习编译器的优化往往需要涉及到全局的改写，包括之前提到的内存，算子融合等。目前深度学习框架的图优化或者高层优化（HLO）部分和传统编译的pass比较匹配，这些优化也会逐渐被标准的pass所替代。但是在高层还会有开放的问题，即高层的抽象如何可以做到容易分析又有足够的表达能力。TVM的Relay，XLA和Glow是三个在这个方向上的例子。

在自动代码生成上，传统编译器的目标是生成比较优化的通用代码。而深度学习编译器的目标是生成接近手写或者更加高效的特定代码（卷积，矩阵乘法等）。相对的，在一些情况下深度学习编译器可以花费更多的时间（去寻找这些解决方案。

整数集分析和Polyhedral Model

Polyhedral model是一个将近研究了十年的领域，其核心思想是采用整数集来表示循环迭代的范围，利用整数集之间的关系来表示迭代变量之间的依赖，从而达到程序分析变换的目的。传统的poly方法采用了线性约束来表示整数集和整数集之间的关系。Polyhedral方法也是被很多人觉得有希望用于优化深度学习算子的方法之一。Poly是真的比较重要呢？这个问题可以分两点来看：

从核心思想上来看，poly的核心思想是整数集抽象和分析。我们也可以称之为广义的poly思想。整数集抽象是一个非常值得所有深度学习编译器采纳的抽象。包括TVM，TC，MLIR等在内的各个方案都引入了整数集抽象。具体的技术路线差别会在于如何表示以及整合整数集来进行分析，在这一点上不同的框架的做法有所不同。

当然，如果直接研究传统的poly文献，比较狭义上来说的poly包含了一套基于线性约束分析整数集的方法和搜索空间。线性约束空间解法带来了一些效率上和对于整数集关系的限制。当然采用线性约束求解的好处是可以解决像三角形或者平行四边形约束的循坏，但是这一类循坏并不常见。反过来说，如何更好的优化规则循环成为更加紧要的问题。因为模型本身的限制，狭义的poly本身不能完全解决搜索空间的问题。

总的来说，直接采用狭义的polyhedral技术并不能完全解决深度学习编译问题，但是整数集分析（广义的poly）已经被广泛地采用于各个方案之中。

包容万象的搜索空间

搜索空间决定了深度学习编译器能力的上界，如何设计搜索空间是所有深度学习编译器需要仔细考虑的话题。很多目前的编译方案还是采用了有限的规则生成搜索空间。一般来说，这一搜索空间的定义需要大量吸收人工优化经验并且加以融入。具体的空间包括循环重排，映射部分计算到实际的加速器指令（张量化，tensorzation），流水线优化等。一般很难确定一个完整的解答，以TVM为例，这一搜索空间会在过去和接下来的几年里面通过社区不断迭代，到达越来越好的效果。

用机器学习优化机器学习

搜索空间本身确立之后剩下的问题是如何快速地找到比较好的算子。我们可以采用机器学习来自动优化算子。有兴趣的同学可以看 AutoTVM：让AI来编译优化AI系统底层算子

编译器之外的架构

需要指出的是，深度学习编译器和传统编译器的一大区别是出了编译器之外还有很多周边的架构用于支持编译和部署。这些架构包括快速的远程部署，不同前端的转化，灵活的运行时设计等。能否有比较好的周边架构也是会决定深度学习编译器易用性的关键。

AI芯片和编译器的协同设计

AI芯片需要编译器吗？这无疑是一个很有趣的问题。在一定程度上，AI芯片对于编译器的依赖取决于芯片本身的设计。一般而言，越灵活的芯片对于编译器的依赖会越大一些。在AI芯片设计之初，大有传统的CISC风格，把所有的优化在芯片本身解决。但是随着领域的演化，为了支持更加灵活的需求，AI芯片本身也会在保留张量指令集和特殊内存结构的前提下越来越灵活。相信未来的架构师需要芯片和系统协同设计的，自动化也会越来越多地被应用到专用芯片中去。有兴趣的同学可以阅读 VTA: 开源AI芯片栈