陈天奇高赞文章：新一代深度学习编译技术变革和展望( 二 )

文章图片
因为生态本身设计和其他一些原因，当然的解决方案基本有以下特点：每一个层级抽象基本由一个比较独立的团体维护。层和层之间往往比较松耦合。最后解决方案本身往往是以一个黑盒工具的方式呈现给用户。
很多人的一个愿景是只要每一层的优化做的好，把东西拼起来，我们就可以组合成一个满足需求的解决方案。虽然这样的做法的确可以达到一定的效果，但是我们也需要反问， multi-staginglowering真的可足以解决深度学习优化的问题吗?
两种隔阂
我们大约在三年前完成了基于multi-staginglowering的解决方案。当我们把全栈解决方案搭建起来并且不断实践之后我们发现有两种隔阂阻碍整个行业的发展。
其中的第一种竖向隔阂阻隔了手工优化的方案和自动编译优化的方案。如果我们看当前的深度学习运行框架和编译框架。我们可以发现有两种流派：一类是以手工算子优化为主的算子库驱动方案。这一类方案一般可以比较容易地让更多的优化专家加入，但是本身也会带来比较多的工程开销。另一类方案是以自动优化为主的编译方案。编译核心方案往往可以带来更多的自动的效果，但是有时候也比较难以引入领域知识。大部分当前的框架基本都只为两者中的其中一个设计。而我们往往发现实际比较好的解决方案其实同时需要机器学习工程师的输入和自动化。并且随着领域的发展，最优的解决方案也会发生变化。怎么样打破这样的竖向墙，让手工优化，机器学习优化专家的知识和自动优化做有机整合，也是目前行业面临的一个大的问题。

文章图片
除了竖向高墙之外，第二类隔阂则一定程度上和multi-stagelowering的生态直接相关。因为我们往往会把不同层级的抽象分开来设计，虽然在抽象内部可以做比较灵活的优化。但是在一个抽象到另外一个抽象的转换时候往往需要通过translator或者lowering批量转换。这样导致了我们很多困难都开始集中在一类抽象和另外一类抽象的边界上，这样导致了如果我们想要在边界上面做一些分步的优化（如把其中一部分子图交给一类编译逻辑，剩下的交给其他编译如逻辑）我们就必须在边界上面引入大量的工程。另外一个常见的现象是这类转换往往是单向的，我们一般会在高阶的抽象如计算图上面做一些优化，然后传递给张量计算层级。但是张量计算或者硬件层级的信息往往难以反馈给更高的层级。举个例子，其实很多时候张量程序的优化本身可以反过来指导计算图层级的算子融合和数据排布，但是当前的单向架构比较难自然地利用这一类反馈。
总结一下我们的经验。深度学习编译和优化本身不是一个一个层级可以全部完成优化的问题。解决相关问题需要各个层级抽象之间的联动。随着TVM和MLIR一类基础架构的出现，我们其实已经可以比较容易地搭建出某一个层级的抽象或者dialect并且让它们之间通过multi-stagelowering的方式从高到地级抽象进行逐层变换和转换。但是现在的困难点往往出现在抽象的转换边界上。不论是在边界引入更多的可模块化整合变换，或者是进行尝试反馈迭代， multi-stagelowering本身是远远不够的。不仅如此，随着抽象层级的增加，如果希望加入自定义算子，我们往往需要针对每个抽象层级进行进一步的架构，其中的开销变得更加庞大。
因为各个抽象之间的竖向和横向隔阂。不论我们如何做好一层抽象内部本身，我们依然难以做好端到端的整体优化。需要注意的是，这些隔阂和问题的存在和基础架构的选择无关，不论是基于MLIR ， ONNX或者是TVM的方案，一旦采用了multi-stagelowering都会不可避免地面对这个问题。相信在领域里面努力的小伙伴不管采取什么基础架构，在把方案打通之后或多或少都会碰到这个本质的问题。