TVM Unity与新一代机器学习编译技术
2023年03月16日，TVMCon在线举行。在上一次TVMCon我们展望了如何突破当前机器学习编译解决方案碰到的瓶颈，并且提出了TVM Unity的技术展望。过去一年我们通过MLC课程也普及了一些相关的概念，有不少同学也给出了宝贵的建议，希望看到更多机器学习编译技术在实践样例。随着TVM社区unity branch的建立和开发，Unity也从概念开始走向实践。感谢各位小伙伴的努力，TVM Unity也从去年的概念宣传逐渐变成了真正可以执行的代码。本文以TVM Unity为主题来介绍我们对于机器学习编译技术变革和发展的一些讨论。本文的内容基于今年我们在TVMCon的主题报告。
为什么做TVM Unity
深度学习和机器学习编译技术一直在快速发展。虽然TVM目前的multi-staging lowering解决方案可以帮助我们解决不少问题，但是我们也碰到了在各种抽象边界上面的瓶颈。我们也一直非常清楚深度学习编译技术本身也需要经历几代的演化。因此从三年前开始我们就开始问这样一个问题：“目前深度学习编译技术的瓶颈在哪里，下一代技术是什么”。如何演化更新架构， 吸收了社区几年来的实践经验和大家的反馈是我们构建Unity的第一个原因。
随着机器学习框架和硬件需求的发展。机器学习编译本身也从一开始的萌芽到了现在行业广泛参与的状况。新硬件，模型和环境的落地需求也必然让机器学习编译成为行业不可或缺的一环。但是机器学习编译本身还是有比较高的门槛。如何让大家更加简单的入门机器学习编译， 如何让更多的同学可以使用和改进机器学习编译技术。者是推动Unity，并且基于它构建MLC课程系列的第二个原因。
最后一个原因来自于推动机器学习编译领域本身的发展。机器学习编译还是一个非常年轻的领域。每年我们都会发表不少优秀的工作。比较可惜的是大部分工作都需要重新拼接，深度改动或者从头构建一个从前端到后端的整个链条，虽然创新的内容可能是其中的一部分。导致从研究到落地有比较大的距离。而且很多时候因为我们自己魔改，几种优秀的工作也不一定能够很有机地结合在一起。相比之下，深度学习算法开发就显得更加模块化的多。在我构建新的residual module的时候可以直接拼接后续的detection loss和optimizer快速看到每个工作在垂直领域上面的效果。而深度学习本身的发展真是得益于各种简单基于python的快速迭代和创新组合。如何可以让机器学习编译领域本身可以这样可组合一起发展，并且让更多的同学可以快速迭代机器学习模型一样来迭代机器学习编译解决方案，如何让我们的每个工作可以更有机地结合在一起。这是我们推动TVM Unity的第三个原因。
三个元素
熟悉机器学习编译的小伙伴可能知道，整个机器学习编译需要经历非常多的环节。而在构建Unity的过程中的主要目标是做减法，把核心元素抽取出来。
整个TVM Unity包含了三个元素：一个统一核心抽象，开发者效率优先的机器学习编译开发，全后端部署优化。
一个统一核心抽象
TVM unity的核心抽象IRModule和TVMScript
TVM unity的所有操作都围绕一个核心抽象展开 – IRModule。一个IRModule包含了我们在执行过程中的（张量）计算函数。比如对于diffusion model而言。每一个函数可以包含我们的执行模块（比如VAE，CLIP，UNet）。随着优化的进行它们也可以对应于我们的算子模块，比如attention或者是conv和gemm。
机器学习编译的四大抽象
这里其实并没有什么创新的地方，就好像Tensor在深度学习存在之前就已经广泛存在一样，IRModule的概念在编译领域基本也是一个常见的概念。而我们集中打磨的有三点
[*]保证编译过程中四大抽象， 计算图，循环程序，算子库和硬件指令都可以统一地被表示在一个IRModule中，并且可以让它们有机地混合交互。
[*]所有变换中间的任意一个IRModule可以被打印成TVMScript直接嵌入在python AST中。TVMScript从而成为编译过程中的一个载体，可以用于结合手工构造，变换和修改调试。
[*]引入原生的dynamic shape的支持，并且进一步引入编译领域可能用到的结构化信息（first class struct info），可以方便地进行进一步的扩展。
TVM Unity对动态shape的支持
上图展示了TVM Unity对动态shape的支持。我们可以看到对于程序中的标记不光包含静态的shape（如4， 64，64）， 它也包含了符号 n贯穿程序始终。在程序中我们还支持符号计算，支持类似R.Tensor((n*2, 4, 64, 64)) 等情况。这些符号计算可以用于编译过程中的证明（如证明两个动态shape相等），帮助编译优化
开发者效率优先的机器学习编译开发
基于以上三点，IRModule和TVMScript成为了迭代的载体。我们的大量创新其实都是在自动地去完成程序之间的变换。我们可以提供各种各样的模块来和它们交互。首先我们可以从各个机器学习框架来导入模型到IRModule。同时我们提供了可交互的Schedule API来变换循环程序，并且利用TVMScript来进行检查和调试。
我们同时提供了各种pass来针对IRModule做各种模块化的变换。在任意一个阶段，我们都可以通过build把一个模块到处到一个可执行的模块。以上的所有过程都都提供了python API，使得我们可以很容易地把各个步骤组合起来。大部分的优化和变换都从原来的编译过程被抽取出来成为了IRModule到IRModule的变换，而build只完成最后最基本的步骤。
为什么选择python呢？是因为我们觉得有必要在生态中提供快速迭代尝试的手段，并且可以让更多的同学参与进来。比如各个流程的自动调优就可以变成在python里面调用各个解决方案的一个循环测试再调用的过程，我们可以在一天就可以构建尝试迭代这类的idea。就好像python推动了深度学习的变革一样，我们相信更加快速地迭代也可以推动深度学习编译的变革，我们也希望为整个生态提供一个这样的支持。当然每一个pass的内部实现和优化还是大量采用了C++的代码，来保证整个工程的产品化特性，只是每一个pass本身都可以在python端进行有机地组合调用。
全场景部署
当然我们也认识到快速的迭代开发需要最后在各个场景的落地。所以这就包含了第三个元素--全场景部署。
这得益于TVM runtime本身的FFI抽象。比起编译端，我们提供了比较灵活的部署端。我们目前支持基本大家可以接触到的运行环境，并且可以利用PackedFunc和Object系统灵活地整合手写代码，编译生成代码和各个厂商本身的driver。其核心概念也非常简单，每一个IRModule里面的函数都会变成一个各个部署端原生运行的函数，可以是python，javascript，c++或者其它任意tvm支持的语言。
TVM Unity可以做什么
上一个小节讲到了TVM unity的三个元素。这三个元素的有机整合可以让整个系统更加模块化，也可以让大家比较灵活地在工作中使用整合大家需要的部分。因为整个系统的核心重新围绕IRModule进行可组合构建，我们不再依赖于几个比较重的模块来完成多层优化。大家也可以更加有机地选择自己需要的部分。比如tuning搜索和搜索之后的最优选择都可以成为Pass。我们可以在我们的实验过程中有机地选择各种自动化或者启发式的方案。由此tuning将成为众多优化中的可选武器之一。大家也可以根据自己的需求来组合对应方案。同时社区也会推动一些常见的方案组合。我们来举两个具体的例子。
算子库和自动编译的有机整合
算子模版库和自动编译一直是机器学习框架的两大流派。而每个方案其实都有各自的优势。比如最近AITemplate的效果很大程度上得益于对于算子模版库的整合。是否可以让两种方案更加有机地整合在一起呢？
答案是肯定的。我们只需要引入了DispatchToLibrary的pass，专门针对整个过程中的几个关键模块替换成`call_dps_packed(“cutlass_attention”..` 等调用，用于直接调用对于模版的生成。而这类对于模版库的调用也可以直接和其他张量算子一起表达在IRModule中。于是我们可以进一步组合TIRLegalize把剩下的算子变换成TensorIR并且进行自动优化。这样的改动基本上只需要一行python代码。我们就可以完成部分计算到cutlass的切换。我们把这样的编译流程在stable diffusion上面进行验证，发现在GPU上面可以快速达到和AITemplate一样的效率。并且这些变换可以和后续的内存分配优化等继续一起联动，发挥更好地迭代效果。
TVM的统一标识可以同时支持算子调用和编译策略
类似的，假设我们要设计新的fusion策略（如完成多层算子融合），但是又不希望解决所有的case。我们可以在fusion之前引入一个自己的customize fusion，并且把它和后续的fuse ops和fuse tir组合起来，我们就可以快速验证新的fusion策略对于整个编译流程的效果。每当我们的迭代经验得到验证，我们可以把它们加入到已有的pipeline里面去。或者对于业务比较重要的模型进行专门的编译迭代优化。
我们后续会推出相关的python教程。希望可以通过这个方式可以让算子编写和编译进行有机的整合。并且可以让每一个人都可以快速地撬动整个编译的流程，看到整合自己的代码生成，手写算子或者编译流程对于整个端到端过程的推动。
快速支持优化新的后端：Web Stable Diffusion
我们的第二个具体场景来自于对于新后端的支持。我们知道diffusion模型本身取得了巨大的成功。而到目前为止大部分的情况下我们都需要GPU服务器来支持相关的计算。我们能否通过机器学习编译来快速把模型迁移到Web端呢？
web stable diffusion
TVM社区最近的项目，web stable diffusion给出了一个肯定的答案。在技术上，我们采用TVM Unity把stable diffusion的clip，scheduler，unet和vae四大模块全部表示在IRModule中，并且结合了内存优化和MetaSchedule技术来自动生成WebGPU的代码。最后通过TVM web runtime无缝运行在了chrome中。 大家可以阅读下面这篇文章来获取更多技术细节。
值得一提的是在迭代上面这个技术方案的过程中，我们用到了TVM unity里面几乎所有的元素，包括通过tvmscript来进行迭代开发，调试， 填坑，和持续迭代在新的后端中可能碰到的问题。而整个解决方案从无到有花费了不到一个月的时间。其中还包括了各种基础架构的迭代。这也得益于整个自动优化的过程和利用python的开发流程进行迭代优化。同时这也得益于我们最近的几个工作在TVM Unity中的有机整合（包括了TensorIR 用于支撑代码后端生成和MetaSchedule来自动生成代码）。
总结和未来展望
在最近的一个月TVM unity从基础架构建设到垂直应用打通走出了最初的一步。相比于从Halide演化而来的TVM而言，TVM Unity相当于一个全新的开始，TVM原有的特性会更有机模块化地体现在TVM Unity中，更加方便大家做有机扩展整合， 我们也会看到像dynamic shape和跨层级表示和包括算子库，自动优化和机器学习优化的新的可能。我们只是走出了小小的一步。相信可以帮助更多的机器学习编译研究者和机器学习系统工程的同学走出更大的一步。这也是我们对于机器学习编译领域的一个尝试，希望可以对于整个生态贡献一部分力量。社区后面也会继续完善和pytorch，onnx,jax等生态的对接，来更好地支持各种模型。我们帮助大家一起持续推进整个领域，帮助支持分布式优化，变长模型等实际各种应用场景。
我们也会基于大家的反馈继续推出MLC相关的课程材料，讲解机器学习编译技术在实际场景中的应用，让大家不光知道我们可以完成有趣的事情，并且可以知道这些有趣的事情背后发生了什么。感谢大家的支持和反馈。希望让更多的人可以使用机器学习编译技术来完成我们自己的工作，让大家根据自己的需求来利用相关的组件来提高整个机器学习系统优化落地的效率。TVM unity也在对应的branch快速开发中，非常欢迎大家参与到开源社区当中来。 dynamic shape的支持代码发布出来了吗 已经在unity branch中有IR的支持，文章更新了动态shape的额外讨论 TVM Unity 有实例吗？看起来是单层的统一IR，这样更好做优化 可以看mlc.ai ,以及 web stable diffusion的实例
页:
[1]