ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

深度学习编译器论文研读指南:从TVM到MLIR的入门地图

深度学习编译器论文研读指南:从TVM到MLIR的入门地图 这个系列我想了很久今天总算动笔了。起因很简单组里要规划ML Compiler方向我从去年年底开始系统性读论文从TVM、XLA一路读到MLIR、Triton、Ansor前前后后几十篇踩了不少坑也渐渐捋出了一些门道。市面上讲ML Compiler的博客不算少但大多数要么只讲某个框架的源码要么只丢一篇论文列表出来不解释为什么值得读、不标注难度、也不说复现时容易栽在哪真正能让人按图索骥、从入门到能动手跑实验的整理非常少。所以我想用这个系列把自己读过的ML Compiler论文按主题分类每篇给出核心贡献、我的理解、关联参考和实际验证心得这篇就是整个系列的地图和开篇后续会持续更新。这个系列适合三类人刚接触ML Compiler、想知道模型为什么跑不快、怎么让它跑更快的算法工程师准备把自定义算子或推理引擎接入编译器体系的系统研发以及正在找方向、需要快速了解研究脉络的在校同学。内容上我会尽量讲人话涉及公式的地方也只解释物理含义和设计动机不做数学推导。读完这一篇你应该能对ML Compiler领域有一个整体认知并且知道下一步该按什么顺序去读哪些论文。1. 为什么我要整理这套ML Compiler论文1.1 ML Compiler到底解决什么问题ML Compiler机器学习编译器说白了就是把你在PyTorch、TensorFlow里写好的模型计算图经过一层层分析和变换最终变成能在GPU、CPU、NPU上高效运行的机器代码。过去这个活多半靠手写CUDA kernel或者依赖框架自带的算子库但模型结构越来越多样、硬件种类越来越多手动优化根本跟不上于是编译器这条路就被推到前台。一个直观的例子是你现在想在新的显卡上跑一个transformer模型如果只依赖cuDNN这类算子库很多融合优化是用不上的比如把LayerNorm里的多个reduce融合成一个kernel、把attention里的多次矩阵乘用一块内存复用掉这些全局优化必须跳出单个算子的视角站在整张计算图上去看。ML Compiler做的就是这件事把计算图拆解、分析依赖、做算子融合、设计内存复用方案、给每个算子生成适配目标硬件的高效代码。1.2 这份清单是给谁用的读前先判断自己的位置如果你是算法工程师平时只调框架不改底层那你可以重点看计算图优化和自动调优方向的论文这能帮你理解为什么同一个模型换一个batch size后性能波动那么大为什么有些融合开关打开后精度有细微变化。如果你是系统方向的工程师建议从编译器基础开始把中间表示这条线吃透再去看TVM、Triton这类上层编译框架会轻松很多。在校学生起步阶段别贪多先读我后面标了必读的几篇再顺着引用关系往外扩展比漫无目的地刷论文效率高得多。我还有一个很实在的建议读之前先给自己定一个目标。比如我要搞清楚TVM到底怎么生成GPU代码或者我要把torch模型通过MLIR跑到自研NPU上。带着具体问题去读远比漫无目的地一篇接一篇看更容易坚持也更容易有产出。2. ML Compiler系列论文的主流分类与研读地图2.1 计算图优化与端到端编译这个方向关注的是从整张计算图出发的联合优化包括算子融合、layout转换、常量折叠、内存复用、并行调度等。代表工作包括Google的XLA、TVM系列、以及后来出现的一些端到端编译栈。读这类论文时一定要记住它解决的问题是整图视角下的联合优化单看某个算子是没有意义的。以算子融合为例它不只是把两个kernel简单合并成一个而是要分析两个算子之间的数据依赖中间结果要不要写回全局内存能不能直接通过寄存器或者共享内存传递融合后循环结构怎么改才能不破坏依赖关系这些都是整图层面才能回答的问题。XLA的核心贡献之一就是把融合这件事系统化了它定义了一套HLO中间表示用HLO Pass来做跨算子优化这套思路后来被很多编译器沿用。2.2 算子生成与代码生成算子生成关注的是单个计算节点如何生成高性能实现核心难点是搜索空间表达和代码质量。经典代表是Halide它首次把算法Algorithm和调度Schedule解耦你可以在不改变计算逻辑的前提下调整tiling、向量化、并行化等调度策略。这个设计后来成为编译器前端设计的经典范式TVM的调度机制也明显受了它的影响。Triton则在另一条路上走得更远它用Tile块级编程模型让你像写普通Python一样写GPU算子编译器负责把Tile映射到GPU计算单元上。实测下来用Triton写一个fused kernel的开发速度比手写CUDA快很多性能也能达到甚至超过手写水平对非编译器背景的算法工程师尤其友好。算子生成方向还有Tiramisu、TACO等代表工作Tiramisu用多面体模型Polyhedral Model做循环变换TACO则专门解决稀疏张量的代码生成问题。2.3 自动调优与性能搜索自动调优要解决的问题很直接给定一个算子或一个子图编译器能从巨大的参数空间里自动找到最优的实现。AutoTVM用模板机器学习成本模型来做它把调度参数枚举成模板用XGBoost这类模型预测每种参数组合的性能再通过迭代采样逼近最优解。Ansor更进一步用程序采样进化搜索替代人工设计模板可以自动生成不同结构的调度在多个推理任务上都刷新了当时的SOTA。这几年自动调优的发展非常快。从模板搜索到基于成本模型的搜索再到用强化学习选择调度策略本质上都在做同一件事如何用尽量少的编译时间换尽量高的运行性能。如果你是做训练加速或者推理引擎的这个方向值得长期关注因为硬件迭代速度远快于手动优化算子库的更新速度能自动适配新硬件的编译方案会越来越重要。2.4 中间表示与可扩展基础设施MLIR是这十几年编译器基础设施领域最重要的一个项目。它提出了一套可分层、可扩展的中间表示体系核心概念是Dialect也就是一组operation、type和pass的集合。你可以为某个特定领域定义自己的Dialect比如Triton Dialect、ONNX Dialect也可以把高层次Dialect逐步lower到LLVM IR或目标硬件指令整个过程像搭积木一样灵活。很多人第一次看MLIR会看不进去因为概念实在太多Operation、Type、Trait、Pass、Dialect、Conversion……但坚持看完之后你对整个编译器架构的理解会上一个大台阶。AI硬件碎片化已经是既成事实各家NPU的指令集和内存架构都不同MLIR提供了一套一次编写、多端生成的基础设施Triton、torch-mlir、ONNX-MLIR等大量项目都建立在它之上所以这个方向几乎是所有做AI编译底层的人绕不开的。2.5 张量代数与稀疏计算TACO是张量代数编译器的代表作它解决的是稀疏张量高效计算的问题。推荐系统、科学计算里有大量稀疏数据但通用编译器很难针对稀疏格式做优化。TACO的贡献在于引入格式维度来统一描述稀疏和稠密计算让同一个张量表达式可以针对不同存储格式生成不同代码。稀疏计算的论文阅读门槛会比稠密方向稍高一点因为涉及格式选择CSR、COO、CSC等和循环结构的联合优化。建议先理解各种稀疏存储格式的数据结构再理解格式是如何影响计算循环的生成顺序的。你一旦理解了存储格式不是数据问题而是调度问题再读这类论文就会豁然开朗。2.6 异构调度与内存规划异构编译关注跨设备调度代表作有Rammer、Welder等。这类论文在GPU和CPU混合场景、多卡模型并行场景里很有价值。Rammer提出了rTask抽象把计算图中的节点按内存访问模式重新组织实现跨算子、跨设备的并行调度思路非常有启发性。内存规划是另一个大主题包括显存复用、换入换出、静态内存分配等。在做超大模型训练和推理时显存往往比算力更先成为瓶颈。读这类论文时要重点关注它的生命周期分析和池化策略是怎么做的——很多优化看起来复杂核心其实就是把不同张量的生命周期错开然后共享同一块显存。2.7 论文速查表一份可打印的研读清单下面是这个系列目前整理的核心论文速查表我会在后续更新中不断补充。难度用星级标注五颗星最困难推荐指数结合了我在实际科研和工程中的验证感受。论文/项目会议/年份方向难度推荐指数HalideSIGGRAPH 2012算子生成★★★五颗星TVMOSDI 2018端到端编译★★★★五颗星AutoTVMASPLOS 2018自动调优★★★四颗星MLIRCGO 2021编译基础设施★★★★★五颗星AnsorOSDI 2020自动调优★★★★五颗星TritonMAPL 2019算子生成★★★五颗星TACOOOPSLA 2017稀疏计算★★★★四颗星RammerOSDI 2020异构调度★★★★四颗星TiramisuCGO 2019多面体编译★★★★★三颗星XLAML for Systems 2017端到端编译★★★★四颗星这张表的排序并不完全代表阅读顺序只是帮你建立全局认知。接下来的章节我会重点拆解六篇最值得精读的论文并给出我实际复现和阅读时的经验。3. 绕不开的几篇经典论文我建议你这样读3.1 Halide先把算法与调度分离吃透Halide虽然最初是面向图像处理管线的但它的核心思想直接影响了整个深度学习编译领域。论文提出的算法与调度分离可以类比成做菜算法是菜谱告诉你需要哪些食材、按什么顺序加工调度就是烹饪顺序和火力安排你可以选择大火爆炒还是小火慢炖。菜谱不变最终菜品一样但时间和口感完全不同。具体到代码上Halide用Func、Expr这些概念描述计算逻辑调度则通过schedule指令控制循环顺序、分块大小、向量化宽度等。我强烈建议你动手跑一跑Halide的官方示例随便挑一个模糊算法先写默认调度再手动调成多级tiling向量化观察性能变化。这个实验做完你再去读TVM里的schedule代码会感觉异常亲切因为TVM的调度设计基本延续了同样的思路。3.2 TVM从AutoTVM到端到端编译栈TVM是深度学习编译领域绕不开的一座山。它的论文提出了一套包括图优化、张量表达式、调度、自动调优和后端代码生成的完整编译栈。读这篇论文我建议配合源码环境一起看光看文字特别容易绕晕。你不需要一开始就读完论文全文先跑通教程再回到论文看设计动机效果会好得多。实操上建议这样复现安装TVM之后跑一遍AutoTVM调优resnet-18的官方示例打开调优日志观察搜索过程。你会发现编译器在调优时不是直接搜机器码而是在搜索循环展开顺序、tiling大小、向量化方式、线程绑定这些调度参数的组合。一旦理解了调度与算法分离的意义你就能明白为什么TVM可以跨硬件算法是同一份调度参数是针对硬件搜索出来的。有个细节值得注意TVM的图优化Relay和算子优化TE/TIR是两层很多初学者搞混。Relay负责融合、layout转换这类跨算子的优化TE/TIR负责单个算子的代码生成。理解这两层的边界再去看TVM的pass流程会清晰很多。我自己当时卡了一周才搞明白图优化后的IR怎么变成TIR然后变成CUDA源码这条路理顺后整篇论文就通了。3.3 MLIR用Dialect重构编译基础设施MLIR的论文实际上是对一系列设计文档的总结所以单独读论文的效果远不如论文代码动手示例三者结合。我的建议是先用一个最简单的例子构建一个只包含addi的模块然后通过mlir-opt跑各种pass观察IR前后变化。再逐步增加toy语言教程里的操作看看一个Dialect从定义到lower的完整流程。动手操作过一遍之后你会发现MLIR最厉害的地方不是某个具体pass而是它的可扩展架构。你想做一个新的硬件后端不需要从零写编译器只要定义好这个硬件的Dialect然后把通用优化pass复用到新Dialect上就行了。我当时在一个自研后端上用MLIR做前端到后端的桥接做完才发现之前对PassManager的理解很多是错的不是pass越多越好而是要在正确的时间点用正确粒度的pass否则会破坏后续优化机会。MLIR的难点在于它概念很多但你不要被吓住。可以把Dialect理解成一套语法的方言Pass理解成对一个方言或跨方言做的变换PassPipeline就是按顺序执行的一组变换。只要你对编译器有过基本接触花两到三周时间是能摸熟这套体系的。3.4 TritonTile抽象如何降低GPU算子开发门槛Triton这篇论文提出的Tile抽象值得反复品味。它把GPU程序看成是对多个块Tile的并行处理编译器负责把块映射到GPU计算单元上同时处理shared memory分配和数据搬移。这样开发者不需要手动管理很多底层细节写起来比CUDA轻松太多。我实际用Triton写过fused MatMul和LayerNorm开发体验确实好直接用Python写了类似numpy的代码加一个triton.jit装饰器编译出来性能接近手写CUDA。读论文时建议重点看3个部分语言前端块级编程怎么表达、编译IRTritonIR的长相、后端怎么生成GPU代码。TritonIR是一种基于SSA的中间表示每个操作都作用在块上理解它之后再去看官方教程里的矩阵乘示例并观察生成的PTX代码你会对GPU编译有更直观的理解。这里有一个经验Triton注意到多使用但不要滥用如果遇到非常规shape或者动态shapeTriton的自动调优和边界处理有时候会不如手写CUDA灵活。它更适合那些结构规整、计算密集型的算子比如矩阵乘、卷积、归一化。下一次你用Triton没跑出手写性能时先别急着下结论检查一下tile size是不是偏小、有没有做足够多的自动调优轮数。3.5 Ansor与Rammer自动调优和异构调度的最新高水位Ansor的论文是AutoTVM之后自动调优方向的集大成者。它不再依赖人工设计的模板而是通过语法树枚举来自动生成不同结构的调度再用进化搜索加成本模型打分来选出最快实现。大量实验证明它在各种推理任务上都能稳定超过AutoTVM后来TVM默认的调优器也换成了Ansor。读Ansor时我建议关注它怎么处理搜索空间这是它和AutoTVM最大的区别。AutoTVM的搜索空间是模板内参数Ansor的搜索空间是程序结构加参数两层。它先通过random sketch生成一批粗粒度的调度模板再用进化搜索微调细粒度参数最后用学习到的成本模型剪枝。这个先粗后细、模型加速搜索的思路非常值得做系统优化的同学学习。Rammer则把目光放在多算子协作上它提出rTask抽象把计算图拆解成更细粒度的任务并按内存访问模式重新组织实现跨算子流水线并行。它解决的是传统编译器在一张很大计算图上做全局调度的困难。做过多卡推理或者芯片后端的人读Rammer会很有共鸣瓶颈往往不在单个算子而在于算子之间的数据依赖和同步开销。4. 从论文到实践研读路线与复现经验4.1 新手友好的研读顺序别一上来就啃MLIR我见过太多人一上来就读MLIR结果被Dialect、Pass、Conversion这些概念劝退。这里给一个我验证过的、相对友好的研读顺序第一步先读Halide或相关中文笔记理解算法与调度分离这个核心思想。第二步读AutoTVM/TVM知道一个完整编译栈大概有哪些模块。第三步读Ansor理解自动调优是怎么从模板走向程序生成的。第四步读MLIR此时你已经有了编译栈的全局认知再学基础设施会容易很多。第五步读Triton和Rammer看看当前热门方向在解决什么问题。每读一篇论文之前先花十五分钟搜一下背景资料或视频对问题建立感觉后再打开PDF效率会高很多。正式读的时候从摘要、图表和结论入手不要从Introduction第一句逐字读到末尾。等主线清楚了再回头看相关工作和原理细节这样不会迷失在文字里。4.2 环境配置与复现实验的真实经验复现ML Compiler相关论文最麻烦的一般是编译环境LLVM版本、Python版本、CUDA版本任何一个不对都会连环报错。我的建议是直接用官方提供的Docker镜像比如TVM社区的tlcpack镜像、MLIR官方的llvm-project镜像能省去大量配置时间。你没必要在本地环境死磕先把跑通的流程放在干净的容器里再根据需求添加自己的代码。跑实验时要把注意力放在编译时间、调优轮数、算子融合前后的耗时对比这些关键指标上而不是一上来就复现论文里的全部数据表。先把最小示例跑通再逐渐扩大模型。另外调优实验一定要固定随机种子否则结果会有抖动严重时甚至能颠倒结论。我踩过这个坑某次对比Ansor和AutoTVM时没有固定种子跑三次出来三种排名后来才发现是因为调优采样路径随机性太大。还有一个很实用的建议每次复现都记录环境版本和关键超参数哪怕只是一个Markdown文件也好。论文里往往不会写清楚某些细节比如调度搜索轮数、batch大小、是否用 warmup这些对结果影响很大。你不记录下次自己都复现不出自己的实验。5. 读论文和做ML Compiler工程最容易踩的坑5.1 读论文的三个典型误区第一个误区是把论文当小说读逐字逐句从头读到尾。正确做法是先抓住问题-方法-实验三段式再回填细节。第二个误区是只读论文不看代码ML Compiler方向尤其明显——尤其是MLIR相关的论文90%的内容都在代码里不实际跑一遍pass你很难理解它为什么要做某些设计。第三个误区是迷信论文里的性能数字论文里的环境、算子列表、硬件型号都是特定的复现时跑出不一样的结果很正常重点在于理解性能差异的来源比如是否做了算子级warmup、是否用了不同的CUDA版本而不是纠结数值本身。这三个误区我都经历过。最惨的一次是照着某篇论文的表格想复现它宣称的三倍加速折腾两周也没跑出来。后来仔细检查发现对方把baseline设成了一个很差的实现而我们的baseline已经做了不少手工优化。所以读论文时不妨先看它的对比实验是怎么设置的baseline是否合理这比关心标题里的XX倍加速重要得多。5.2 工程落地中的常见问题与排查思路接入TVM或MLIR做实际推理优化时我遇到最多的问题大概有这么几类。第一类是算子兼容性模型里有一些自定义算子编译器识别不了直接报错。解决思路是尽量用标准算子表达计算或者为编译器写自定义算子注册把新算子以外部函数或自定义调度的方式接进去。不要一上来就指望编译器能自动优化一切编译器的能力边界通常比你想象的小。第二类是编译时间过长在大模型上做自动调优动辄就是几小时甚至几天。解决方法是缩小搜索空间、用先验配置做冷启动、多用成本模型而不是全程在线搜索。你还可以先用一个小输入shape完成调优再用得到的schedule去跑正式shape虽然不一定最优但往往能省90%的调优时间。第三类是精度不一致图优化、低精度量化后结果和原模型有偏差。排查时先关闭所有pass再逐个开启用二分法定位是哪个pass引入的偏差。我见过最简单的一个case是某个fusion pass把浮点累加顺序改了导致极微小的数值差异结果被安全校验程序判定为模型不一致。遇到这类问题别慌先看pass日志确认融合了什么、改了哪层IR基本就能定位。这里我也把自己的排查思路整理成了一张速查表希望能帮到你现象常见原因排查思路编译报错算子不支持自定义算子未注册改用标准算子或注册外部函数调优时间过长搜索空间太大缩小搜索空间用小shape预调优推理结果精度偏差fusion或量化引入微小差异逐个关闭pass二分定位生成的kernel性能差tile size或绑定策略不合理查看调优日志对比不同schedule多卡或多线程结果抖动存在共享状态未同步加锁或固定随机种子重新验证6. 这个系列的后续计划与我的个人体会最后聊一点我自己比较深的感受。论文不是越新越好经典文章的价值在于它提出的抽象和概念框架能在很长一段时间里保持生命力。比如Halide的算法与调度分离思想从2012年影响到了今天TVM、MLIR、Triton都能看到它的影子。读ML Compiler系列论文这件事最大的收获不是掌握了某个框架的API而是学会了从计算图、调度、代码生成、后端抽象这条链路去思考性能问题。现在我看任何一个推理框架的性能瓶颈下意识就会去分析它是卡在算子实现上还是卡在图优化上还是卡在数据搬移上这种系统性的拆解能力就是读论文带给我的。这个系列会持续更新。每一篇论文我都会给出难度评级、阅读顺序、和复现参考尽量把我掉过的坑都写在里面。下一期我打算先从TVM的Relay和TIR源码开始拆解把整个编译流程从上到下走一遍如果读到这篇开篇的你也刚好在入门ML Compiler欢迎照着这个清单一篇篇读下去。踩坑了也别气馁这些都是学习的一部分。
返回列表