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

机器之心转载
作者：陈天奇
陈天奇是机器学习领域著名的青年华人学者之一，本科毕业于上海交通大学ACM班，博士毕业于华盛顿大学计算机系，研究方向为大规模机器学习。在本文中，陈天奇回答了目前深度学习编译技术的瓶颈在哪里，下一代技术是什么？
2021年12月16日，第四次TVMCon如期在线举行。从大约四年多前行业的起点开始，深度学习编译技术也从初期萌芽阶段到了现在行业广泛参与的状况。我们也看到了许多关于编译和优化技术有趣的讨论。深度学习编译技术生态的蓬勃发展也使得我们有了许多新的思考。 TVM作为最早打通的深度学习编译技术框架已经可以给大家提供不少价值。但是就好像深度学习框架本身会经历从从第一代（caffe）到下两代（TF,pytorch）的变化一样，我们非常清楚深度学习编译技术本身也需要经历几代的演化。因此从两年前开始我们就开始问这样一个问题：“目前深度学习编译技术的瓶颈在哪里，下一代技术是什么“ 。
比较幸运的是，作为最早打通编译流程的架构我们可以比较早地直接观察到只有尝试整合才可以看到的经验。而目前开始繁荣的生态本身(MLIR的各种dialects ， XLA,ONNX,TorchScript)也开始给我们不少可以学习参考的目标。本文是对过去两年全面深度学习编译和硬件加速生态的总结思考，也是我们对于深度学习新一代编译技术的技术展望，希望对大家有一些参考价值。本文的内容大部分来自于今年我们在TVMCon的主题报告。
现状是什么：当前深度学习编译解决方案和瓶颈
四类抽象
当然深度学习编译加速生态已经从萌芽阶段到开始繁荣生长的阶段。但是抛开具体实现而言，现在深度学习编译生态围绕着四类抽象展开:
计算图表示(computationalgraph)：计算图可以把深度学习程序表示成DAG ，然后进行类似于算子融合，改写，并行等高级优化。 Relay,XLA,Torch-MLIR ， ONXX等基本都在这一级别。
张量程序表示(tensorprogram):在这个级别我们需要对子图进行循环优化，对于DSA支持还要包含张量化和内存搬移的优化。
算子库和运行环境(libraryandruntime):算子库本身依然是我们快速引入专家输入优化性能的方式。同时运行环境快速支持数据结构运行库。
硬件专用指令(hardwareprimitive)：专用硬件和可编程深度学习加速器也引入了我们专用硬件张量指令的需求。

文章图片
当然的深度学习编译生态基本上也是围绕着各种抽象的实现展开。上图是对于当前生态的一个不完全总结。值得指出的是，算子库和运行环境等本身未必直接和编译技术相关。但是因为我们的目标是新硬件部署和运行加速，运行环境抽象和硬件指令也是生态的重要一环。
我们需要多层抽象本身基本是深度学习编译和优化领域大家达成的一种共识。但是为了真正地支持机器学习，光依赖于各个组件本身是远远不够的。我们发现最大的问题是“如何设计各个层级的抽象，并且对它们进行有效的整合” 。
目前的解决方案
如果我们仔细地研究目前的整个生态，包括深度学习框架，编译框架(包括基于MLIR ， ONNX或者是TVM的解决方案）。大家都遵循一种叫做多层渐进优化(Multi-stagelowering)的方式。这种构建方式的大致思路是我们在每一层抽象中采用一个（有时多个）中间表示。我们会在每一个层级(dialect,abstraction）做一些内部优化，然后把问题丢给下一个层级继续进行优化。