2021-12-22 05:20

2021年12月16日，第四次TVMCon如期在线举行。从大约四年多前行业的起点开始，深度学习编译技术也从初期萌芽阶段到了现在行业广泛参与的状况。我们也看到了许多关于编译和优化技术有趣的讨论。深度学习编译技术生态的蓬勃发展也使得我们有了许多新的思考。TVM作为最早打通的深度学习编译技术框架已经可以给大家提供不少价值。但是就好像深度学习框架本身会经历从从第一代（caffe）到下两代（TF, pytorch）的变化一样，我们非常清楚深度学习编译技术本身也需要经历几代的演化。因此从两年前开始我们就开始问这样一个问题：“目前深度学习编译技术的瓶颈在哪里，下一代技术是什么“。

比较幸运的是，作为最早打通编译流程的架构我们可以比较早地直接观察到只有尝试整合才可以看到的经验。而目前开始繁荣的生态本身 (MLIR的各种dialects，XLA, ONNX, TorchScript) 也开始给我们不少可以学习参考的目标。本文是对过去两年全面深度学习编译和硬件加速生态的总结思考，也是我们对于深度学习新一代编译技术的技术展望，希望对大家有一些参考价值。本文的内容大部分来自于今年我们在TVMCon的主题报告。

四类抽象

当然深度学习编译加速生态已经从萌芽阶段到开始繁荣生长的阶段。但是抛开具体实现而言，现在深度学习编译生态围绕着四类抽象展开:

计算图表示(computational graph)：计算图可以把深度学习程序表示成DAG，然后进行类似于算子融合，改写，并行等高级优化。Relay, XLA, Torch-MLIR，ONXX 等基本都在这一级别。

张量程序表示(tensor program): 在这个级别我们需要对子图进行循环优化，对于DSA支持还要包含张量化和内存搬移的优化。

算子库和运行环境(library and runtime): 算子库本身依然是我们快速引入专家输入优化性能的方式。同时运行环境快速支持数据结构运行库。

硬件专用指令 (hardware primitive) ：专用硬件和可编程深度学习加速器也引入了我们专用硬件张量指令的需求。

当然的深度学习编译生态基本上也是围绕着各种抽象的实现展开。上图是对于当前生态的一个不完全总结。值得指出的是，算子库和运行环境等本身未必直接和编译技术相关。但是因为我们的目标是新硬件部署和运行加速，运行环境抽象和硬件指令也是生态的重要一环。

我们需要多层抽象本身基本是深度学习编译和优化领域大家达成的一种共识。但是为了真正地支持机器学习，光依赖于各个组件本身是远远不够的。我们发现最大的问题是“如何设计各个层级的抽象，并且对它们进行有效的整合”。

目前的解决方案

如果我们仔细地研究目前的整个生态，包括深度学习框架，编译框架(包括基于MLIR，ONNX或者是TVM的解决方案）。大家都遵循一种叫做多层渐进优化(Multi-stage lowering)的方式。这种构建方式的大致思路是我们在每一层抽象中采用一个（有时多个）中间表示。我们会在每一个层级(dialect, abstraction）做一些内部优化，然后把问题丢给下一个层级继续进行优化。

因为生态本身设计和其他一些原因，当然的解决方案基本有以下特点：每一个层级抽象基本由一个比较独立的团体维护。层和层之间往往比较松耦合。最后解决方案本身往往是以一个黑盒工具的方式呈现给用户。

很多人的一个愿景是只要每一层的优化做的好，把东西拼起来，我们就可以组合成一个满足需求的解决方案。虽然这样的做法的确可以达到一定的效果，但是我们也需要反问，multi-staging lowering真的可足以解决深度学习优化的问题吗?

两种隔阂

我们大约在三年前完成了基于multi-staging lowering的解决方案。当我们把全栈解决方案搭建起来并且不断实践之后我们发现有两种隔阂阻碍整个行业的发展。

其中的第一种竖向隔阂阻隔了手工优化的方案和自动编译优化的方案。如果我们看当前的深度学习运行框架和编译框架。我们可以发现有两种流派：一类是以手工算子优化为主的算子库驱动方案。这一类方案一般可以比较容易地让更多的优化专家加入，但是本身也会带来比较多的工程开销。另一类方案是以自动优化为主的编译方案。编译核心方案往往可以带来更多的自动的效果，但是有时候也比较难以引入领域知识。大部分当前的框架基本都只为两者中的其中一个设计。而我们往往发现实际比较好的解决方案其实同时需要机器学习工程师的输入和自动化。并且随着领域的发展，最优的解决方案也会发生变化。怎么样打破这样的竖向墙，让手工优化，机器学习优化专家的知识和自动优化做有机整合，也是目前行业面临的一个大的问题。

除了竖向高墙之外，第二类隔阂则一定程度上和multi-stage lowering的生态直接相关。因为我们往往会把不同层级的抽象分开来设计，虽然在抽象内部可以做比较灵活的优化。但是在一个抽象到另外一个抽象的转换时候往往需要通过translator或者lowering批量转换。这样导致了我们很多困难都开始集中在一类抽象和另外一类抽象的边界上，这样导致了如果我们想要在边界上面做一些分步的优化（如把其中一部分子图交给一类编译逻辑，剩下的交给其他编译如逻辑）我们就必须在边界上面引入大量的工程。另外一个常见的现象是这类转换往往是单向的，我们一般会在高阶的抽象如计算图上面做一些优化，然后传递给张量计算层级。但是张量计算或者硬件层级的信息往往难以反馈给更高的层级。举个例子，其实很多时候张量程序的优化本身可以反过来指导计算图层级的算子融合和数据排布，但是当前的单向架构比较难自然地利用这一类反馈。

总结一下我们的经验。深度学习编译和优化本身不是一个一个层级可以全部完成优化的问题。解决相关问题需要各个层级抽象之间的联动。随着TVM和MLIR一类基础架构的出现，我们其实已经可以比较容易地搭建出某一个层级的抽象或者dialect并且让它们之间通过multi-stage lowering的方式从高到地级抽象进行逐层变换和转换。但是现在的困难点往往出现在抽象的转换边界上。不论是在边界引入更多的可模块化整合变换，或者是进行尝试反馈迭代，multi-stage lowering本身是远远不够的。不仅如此，随着抽象层级的增加，如果希望加入自定义算子，我们往往需要针对每个抽象层级进行进一步的架构，其中的开销变得更加庞大。

因为各个抽象之间的竖向和横向隔阂。不论我们如何做好一层抽象内部本身，我们依然难以做好端到端的整体优化。需要注意的是，这些隔阂和问题的存在和基础架构的选择无关，不论是基于MLIR，ONNX或者是TVM的方案，一旦采用了multi-stage lowering都会不可避免地面对这个问题。相信在领域里面努力的小伙伴不管采取什么基础架构，在把方案打通之后或多或少都会碰到这个本质的问题。