资讯动态

ML Compiler 论文阅读路线图:从计算图优化到算子融合与自动调优

发布时间:2026/9/8 17:59:39 来源:尧图企业网站定制
追了两年的 ML Compiler 系列论文我最大的感受是这个方向看似小众实际上已经成了深度学习框架性能竞争的核心战场。从 TensorFlow 的 XLA、PyTorch 的 TorchInductor再到各种软硬件协同设计的新框架底层真正比拼的就是编译优化的功底。如果你也对这个方向感兴趣想要系统性地阅读和整理 ML Compiler 系列论文我这份持续更新的路线图和解剖式解读应该能帮你节省不少时间。它既涵盖计算图优化、算子融合、自动调优等经典话题也延伸到 MLIR 统一基础设施、搜索式编译器、AI for CodeGen 等前沿方向。适用于想做 AI 框架开发的工程师、从事高性能计算的科研人员以及刚入门的在校学生。1. 先搞清楚 ML Compiler 到底在做什么1.1 从一次性能调优说起我在一次线上推理服务优化中遇到了一个非常典型的问题同一个卷积模型卧槽换了 GPU 型号之后延迟反而变高了。最开始我以为是驱动或者 CUDA 版本不对折腾了半天才意识到问题出在算子的底层实现上——老 GPU 上有些 kernel 是手工调优过的新 GPU 架构变了手工 kernel 没跟上反倒不如编译器自动生成的代码。这件事给我留下的印象特别深。它让我明白一个深度学习模型部署到实际硬件上能不能发挥出性能早已不只是模型结构的事。计算图怎么切分、算子怎么融合、数据布局怎么转换、kernel 怎么生成、循环怎么展开这一整条链路就是 ML Compiler 要解决的问题。ML Compiler说白了就是把一个高层描述的神经网络模型转换成能在特定硬件上高效执行的底层代码。它和传统编译器有相似之处——都有前端、优化、后端的分层结构但又有本质差异。传统编译器处理的是for循环和指针ML Compiler 处理的是张量和计算图优化目标通常是访存开销、显存占用、并行度和算子启动开销。1.2 一张技术全景图我习惯把 ML Compiler 的技术栈切成四层方便整理论文时归类前端与图优化层负责把模型转换成计算图做算子融合Operation Fusion、常量折叠、死代码消除、Common Subexpression Elimination。代表性工作包括 TensorFlow 的 Grappler、PyTorch 的 TorchScript 与 Dynamo、XLA 的 HLO 优化 pass。中间表示与抽象层业界花了很大的力气去设计统一的中间表示IR最典型的就是 MLIR。MLIR 提供了一套可扩展的、多层级的 IR 基础设施让不同方言Dialect之间可以互相转换把前端框架和后端硬件解耦。算子实现与代码生成层把计算图里的算子映射到底层 kernel。这个方向有两个流派一个是手工/半自动写出高性能模板比如 cuDNN 里的很多 kernel另一个是自动代码生成比如 TVM 的算子模板与 AutoTVM、Ansor、Halide 的调度语言。近几年 Triton 凭借着更好的编程体验杀了出来走的是“块级编程 编译器自动优化”的路线。自动调优与搜索层解决“在巨大参数空间中寻找最优 kernel 配置”的问题。例如 TVM 的 x86/GPU 算子调优器、TensorRT 的 tactic 选择、cuDNN 的 heuristic 与 benchmark 模式。近年来的新方向则尝试把强化学习、贝叶斯优化甚至大语言模型引入代码生成与调优过程。读论文的时候先判断它落在哪一层再看它和上下层的接口如何衔接整体脉络就会清晰很多。很多时候一个循环优化技巧看起来简单但放到编译器上下文里难点在于如何让它对用户透明、可推广到多种硬件这些工程层面的考量往往才是论文的隐藏主线。1.3 框架、编译器与底层库的分工刚接触这个领域时我经常被一堆名词绕晕PyTorch 是编译器吗cuDNN 是编译器吗TensorRT 是不是编译器这里可以做一个通俗类比——如果把训练模型比作做一顿大餐那么框架是菜谱编译器是厨师基于菜谱切配、掌勺、摆盘的全流程管理底层库则像是已经加工好的调味料包。PyTorch 是动态图前端负责接收模型定义、记录计算过程并把算子调度到后端执行。它内部有 TorchScript、TorchInductor 等编译路径。CUDA / cuDNN / cuBLAS 这些底层库提供了一批高度优化的原语算子。编译器可以在做算子融合时把这些原语当作“不可再分”的基础块也可以绕过它们直接生成更深度的优化代码。XLA、MLIR、TVM、Triton 这类的真正“编译器”才是居中调度、生成、优化的核心。它们接收框架输出的计算图经过层层改写最终产出能在 GPU 上跑的 CUDA 代码或底层指令。理解了这个分工再看论文就不会再把不同层级的贡献混为一谈。我见过不少人看了 TVM 的论文后以为这是在和 PyTorch 竞争其实它俩完全不在同一层。TVM 更多是作为 PyTorch 的后端编译器存在或者作为一种独立的部署编译工具链。2. 论文阅读路线图从入门到进阶2.1 入门组计算图优化经典论文如果你刚起步我建议从这三篇开始它们基本覆盖了深度学习框架早期如何“编译”模型的底层逻辑TensorFlow XLA 对应的论文《XLA: Optimizing Compiler for TensorFlow》以及配套的 HLO IR 设计文档。XLA 的思路非常直观把整个计算图拿过来做子图捕获cluster然后对每个子图进行融合优化生成高效的设备代码。值得关注的是它引入了跨算子融合的核心理念——多个算子合并成一个 kernel减少 kernel 启动和中间张量读写开销。《TensorFlow Graph Optimization》或 Grappler 相关的系统论文讲的是如何在计算图级别做常量折叠、算术简化、布局优化。虽然这些内容在工程实现上偏琐碎但它能让你建立“编译器是 pass 管道”的心智模型。PyTorch 的 TorchScript 早期论文和 Dynamo 的系统描述重点看它如何把动态图转换成静态图。动态图灵活但逐个算子调度的开销高静态图可优化空间大但会牺牲一定的编程灵活性。Dynamo 用字节码分析加图捕获的方式在保持易用性的同时争取性能这条思路今天看仍然很先进。入门阶段不要贪多把这三篇的框架图和优化 pass 流程吃透基本上就能看懂后续大部分论文在讨论什么了。另外强烈建议一边读论文一边跑一下对应的框架代码光看术语是记不住的。2.2 核心中间表示从图 IR 到 MLIR 演进史ML Compiler 里最有意思的一段历史就是中间表示从各家自研逐步走向 MLIR 统一的过程。早期 Google 有 HLO、TensorFlow 有 Grappler GraphDefTVM 有 Relay / TIRPyTorch 有 TorchScript IR每个系统都搞了一套自己的表示互相之间又不能直接互转生态割裂非常严重。Chris Lattner 带着 MLIR 的论文和工作出现后情况发生了很大变化。MLIR 的核心思想是不要试图设计一个万能 IR而是定义一套开放、可扩展的 IR 基础设施让不同的方言Dialect共存并通过 dialect conversion 框架进行逐步降级。这种设计很像乐高积木——底层块是通用的但你可以自由组合出不同的形状。读 MLIR 相关论文时有几个概念必须搞清楚Operation、Value、Region、Block 这些基础结构以及它们如何构成一个嵌套的 IR 树。Dialect 机制为什么有了多种方言前端框架就不需要为每个后端重写一遍语法树。Dialect Conversion 与 Pattern Rewriter这是实现 pass 的核心机制理解了它就能看懂大部分优化 pass 的实现逻辑。MemRef 与张量类型系统张量如何在设备内存中分配、索引、读写直接影响算子的访存效率。此外还推荐读一下 TVM 的 Relay 论文和 XLA HLO 的设计文档用来对比不同 IR 的取舍这样你对“为什么需要 MLIR”会有更深的体会。MLIR 的价值更多在于生态整合和降低碎片化成本而不是在单一后端上做到极限性能。2.3 关键突破算子融合与内核生成算子融合是 ML Compiler 里最落地、收益最直接的技术之一。所谓算子融合就是把计算图中相邻的多个算子合并成一个算子减少中间结果的写回和读取。一个最常见的例子是 Conv-BN-ReLU 融合传统实现里卷积输出先写回显存再读出来做 BatchNorm然后再写回去再做 ReLU。融合之后这三个步骤可以在一个 kernel 里完成中间数据只留在寄存器或共享内存里。TVM 的论文和 Halide 的早期论文是理解算子融合与调度Schedule的最好教材。Halide 把算法Algorithm和调度Schedule分离的哲学后来被大量 ML 编译器继承。你可以先用纯函数式的方式描述“要算什么”再用调度原语去控制“怎么算”——包括循环划分、合并、向量化、展开、并行化等。Triton 的出现在这个领域是个大事件。它提出了一种以“块”为中心的编程模型你不需要手写线程级别的代码只需要描述每个 block 要处理的数据和计算逻辑。编译器负责把 block 映射到 GPU 的 warp 调度、shared memory 分配和向量化访存。对很多 GPU kernel 开发者来说Triton 大大降低了编写高性能 CUDA kernel 的门槛。用 TVM 的 Ansor 论文作为扩展阅读也很好它把算子生成和调优问题视为一个搜索问题通过模板自动生成 强化学习/进化算法搜索自动找到最优的循环结构和并行策略。这个方向直接催生了后续自动代码生成的一系列工作。2.4 搜索式编译自动调优的新浪潮搜索式编译是我个人认为最值得关注的研究方向之一。传统的编译器优化依赖人工设计的 pass 和启发式规则但面对不同模型和硬件很难保证每次都找到最优解。搜索式编译的思路是把“选什么调度、用什么 tile 大小、怎么排布共享内存”当作一个搜索问题。AutoTVM 用模板 成本模型 搜索算法自动为算子挑选最优配置。Ansor 则更进一步不依赖固定模板从底层开始生成候选程序再通过搜索寻找高效实现。这类系统的共同点是训练一个性能预测模型来缩减搜索空间避免在真实硬件上穷举所有配置因为真实评测一次 kernel 可能要几十毫秒到几十秒代价太高。最近让我眼前一亮的工作还包括把 Graph Neural Network 用于学习计算图表示以指导优化以及利用大语言模型生成或改写代码来提升 kernel 性能。比如有些论文用 LLM 对简单算子直接生成 UVM/共享内存版本的 CUDA 代码再把这些代码和模板搜索结合。这个方向还在快速变化中但“AI 来编编译器”的思路已经变成了现实。我建议读者在看搜索式编译论文时关注两个关键指标搜索时间成本和交出的性能收益。很多论文在几个 benchmark 模型上赢了但搜索过程耗时半天在实际迭代开发中很难用起来。反而是那些结合了启发式 成本模型、能在几分钟内收敛的方案落地可能性更大。3. 读论文的实操方法论3.1 三遍阅读法在 ML Compiler 论文上的变体读系统类论文和读算法类论文的方法不太一样。ML Compiler 的论文里系统设计和工程实现往往占了大头有的还伴随着巨大的代码库如 TVM、MLIR、Halide。完全从第一页读到最后一页很容易迷失。我的做法是“三遍法 代码验证”第一遍只读标题、摘要、引言和图表目标是搞清楚这篇论文解决了什么问题它和其他工作的核心差异是什么。这一步通常 20 分钟以内。第二遍细读方法论与系统设计部分不看公式推导只看关键流程。比如它的优化 pass 是怎么组织的搜索空间是怎么定义的评估指标是什么。边读边在纸上画数据流向。第三遍逐行研读公式和结果并且打开对应的开源代码去验证论文里的关键实现。如果论文没有开源那第三遍的收益会打折扣但仍可从实验设置里学到很多评估技巧。大部分时候你不需要对每篇论文都做到第三遍。可以根据自己的研究主题确定哪些需要精读、哪些泛读。我的经验是越偏系统的论文越需要配合代码精读越偏理论比如某些调度算法的复杂度分析的论文第三遍的收益越高。3.2 复现论文的五个标准步骤很多同学读到一篇让人兴奋的论文马上就想复现但常常踩坑。我总结了一套相对成熟的复现流程先跑通官方的 Hello World / 单元测试。无论论文吹了什么先把它的基础示例跑起来检查环境依赖和硬件兼容性。选择一个论文报告的 benchmark按照论文的设置跑一遍对比数字。如果连 baseline 数字都对不上说明环境差异较大先不要急着继续。把论文里的核心优化功能依次关闭/打开观察性能变化。这一步最有价值能帮你确认哪些优化对最终收益贡献最大。找一个新的模型或者算子测试它的泛化能力。论文里的结果往往是精心挑选的换个模型看看是否依然成立。修改关键参数如 tile size、搜索迭代次数、batch size画出敏感性曲线。这会极大加深你对系统行为模式的理解。在复现过程中GPU 型号、驱动版本、CUDA 版本、cuDNN 版本都会影响结果。建议在同一份环境里对比不同论文不要跨环境直接对比数据。3.3 追踪最新论文的工具链ML Compiler 论文更新速度很快光靠 Arxiv 邮件订阅很容易被淹没。我常用的追踪方式包括arXiv-Sanity 和 Papers with Code按关键词订阅既有定性。重点搜索这些关键词MLIR, TVM, Ansor, Halide, Triton, deep learning compiler, kernel generation, operator fusion, scheduling optimization。GitHub 上关注对应的仓库尤其是 TVM、MLIR、Triton、XLA、TorchInductor。很多真正有价值的技术细节会在 commit message 和 issue 讨论里出现比论文本身更及时。会议论文MLSys、OSDI、SOSP、ASPLOS、CGO、PACT 是主要阵地。其中 MLSys 和 CGO 特别适合 ML 编译器OSDI 和 SOSP 则偏系统整体设计。行业博客很多大厂的编译器团队会发技术博客比如 OpenAI 的 Triton 相关博客、Google 的 XLA/MLIR 博客。这些内容比论文更偏工程实践有时还包含一些论文不会细说的踩坑经验。我建议每周固定一个时间集中扫一遍新论文把看起来有价值的丢进“稍后读”列表集中到周末精读。不要逐篇在当时精读那样容易打断整块时间而且很多论文最终你并不会想细读。4. 一个值得深挖的方向以 Triton 为例的实操记录4.1 Triton 的工作机制拆解前面说到Triton 的核心贡献是把 GPU 编程从“线程级别”提升到了“块级别”。为了更具体地理解这一点我建议读者去跑一个简单的向量加法 kernel。对比 CUDA C 和 Triton 的写法差异你能立刻感受到抽象级别的不同。在 CUDA 里你可能要写明确声明 threadIdx、blockIdx、blockDim 等线程组织结构手动把数据索引换算成全局内存的线性偏移手动控制 shared memory 的声明和同步__syncthreads。在 Triton 里你只需要写一个triton.jit修饰的函数把输入张量按块切分然后用类似 numpy 的语义描述计算。编译器会推断出需要多少 block、每个 block 里的线程怎么排布、数据如何加载到寄存器/共享内存以及如何做内存合并访问。Triton 的底层会借用 MLIR 之类的 IR 来做优化包括布局推断、访问合并、L2 缓存复用等。理解它内部如何工作对读相关论文会很有帮助。4.2 一个 MatMul 融合实践假设你要实现一个既算矩阵乘法又做偏置相加和 ReLU 激活的融合 kernel。传统做法可能是先调用 cuBLAS 做矩阵乘把结果写回全局内存再启动一个 elementwise kernel 做偏置与 ReLU。这个流程中中间结果要往返一次显存带来了大量额外带宽开销。如果用 Triton写法大致是import torch import triton import triton.language as tl triton.jit def fused_matmul_bias_relu( a_ptr, b_ptr, bias_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr ): pid_m tl.program_id(0) pid_n tl.program_id(1) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) offs_k tl.arange(0, BLOCK_K) a_ptrs a_ptr (offs_m[:, None] * stride_am offs_k[None, :] * stride_ak) b_ptrs b_ptr (offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn) acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for k in range(0, K, BLOCK_K): a tl.load(a_ptrs, maskoffs_k[None, :] K - k, other0.0) b tl.load(b_ptrs, maskoffs_k[:, None] K - k, other0.0) acc tl.dot(a, b) a_ptrs BLOCK_K * stride_ak b_ptrs BLOCK_K * stride_bk bias tl.load(bias_ptr offs_n) acc acc bias[None, :] acc tl.maximum(acc, 0) c_ptrs c_ptr (offs_m[:, None] * stride_cm offs_n[None, :] * stride_cn) tl.store(c_ptrs, acc, mask(offs_m[:, None] M) (offs_n[None, :] N))这段代码很短但背后的信息量很大。tl.dot并不是简单调用 cuBLAS而是会尝试生成一个拆分 K 维并使用 tensor core 的循环BLOCK_M、BLOCK_N、BLOCK_K的选择影响着 shared memory 占用和计算访存比。在这个基础上你可以继续融合 LayerNorm、Residual 等操作进一步减少中间张量。我实际跑过这个 kernel相比单独调用torch.matmul加后续 elementwise 算子在小矩阵例如 128x128x128上能快 30% 左右在大矩阵上收益会略小因为此时访存不再是瓶颈。这个例子很好地向初学者展示了 ML Compiler “跨算子融合”的实际效益。4.3 调度搜索与编译器后端的联动Triton 这类系统并非万能。它的优化高度依赖编译器对 block 的排布和对 tensor core 的使用。像 persistent matmul常驻矩阵乘这种高级写法传统 TVM 模板里可能要花很大的功夫实现但在 Triton 中只需要在调度层面做少量修改例如控制 program_id 与循环的绑定方式就能让一个程序块反复使用已经加载的数据。这也引出我对“编译器后端”更深的理解调度和代码生成从来不是孤立的。你选择的 tile 大小、向量宽度、循环顺序要跟硬件的手性比如 GPU 的 memory coalescing、L1/L2 缓存策略、tensor core 的形状要求匹配。很多论文工作的价值正是帮你自动化寻找这种匹配关系而不用像早期 kernel 工程师那样逐手调参。如果你对 Triton 的下层感兴趣可以把它和 TVM 的 TIR、XLA 的 Gemm 类 lowering 做对比。你会发现不同系统之间其实有很多底层优化是类似的只是暴露给开发者的抽象层级不同。5. 常见问题与排查技巧实录5.1 论文实验结果总是复现不出来这是我遇到最多的问题也是初学者最容易崩溃的地方。复现失败很多时候不是你错了而是论文没有完整交代实验环境。GPU 型号、驱动、L2 缓存大小、编译器版本、依赖库版本这些都能显著影响最终性能。我的排查顺序是首先确认硬件和软件版本是不是和论文一致。如果 GPU 不同性能差异是正常的。检查 benchmark 是否包含了 warmup 时间。GPU kernel 第一次运行往往有初始化开销有的论文已经做了 warmup有的没有这会导致数字直接不可比。看论文有没有明确指定 batch size、shape 分布。很多编译器对特定 shape 会触发不同的 code path。尝试关闭所有其他后台任务保证显存和 CPU 资源充足。如果都排除了还是复现不了就直接给作者发邮件或提 GitHub issue。大部分做 ML 编译器的人还是很乐意交流的。5.2 显存与 kernel 启动开销的取舍很多刚开始接触算子融合的人觉得“融合越激进越好”实际工程中并不是。当算子数量特别多、中间张量特别大时激进的融合可能导致寄存器溢出spill反而让 kernel 变慢。TVM 的某些早期模板在融合策略上相对激进在某些图形形状上就出现过这个问题。另外一个感知不深、但非常重要的因素是 kernel launch 开销。GPU 上启动一个 kernel 大约有 5-10 微秒的开销如果你的模型里有几千个小算子仅启动时间就可能占掉总时长的一半。这也是为什么大子图捕获big cluster和层融合这么重要——把几十个小算子合成一两个大 kernel启动开销立刻被摊薄。实践中我通常会先用 profiler 跑一遍模型看看时间分布是算子占大头还是 kernel launch 占大头再决定优化的主攻方向。如果启动开销已经占了大头那你首先应该做算子融合而不是去手工调某个算子的 tile size。5.3 团队协作与论文库维护经验最后分享一点维护“ML Compiler 系列论文”这个资料库的心得。很多人会把论文下载到本地看完就堆在文件夹里过几个月想找某个 idea 却无从下手。我采用的框架是每篇论文一个 Markdown 文件包含“problem/idea/method/experiments/notes”五个字段用标签管理比如graph-opt、auto-tuning、ir-design、gpu-kernel。用 Git 仓库维护既方便版本管理也能在上面做一些交叉链接。我还会在每篇笔记的开头加上“关联工作”把相关的论文串成一条线。每三个月做一次整体复盘把看过的论文按主题重新归类淘汰掉过时的想法再把新发现的经典工作补充进去。这套方法让我在写综述、做方案选型和面试准备时都能快速调用到对应的资料。如果你也在整理这个系列强烈建议从一开始就建立一套自己的索引体系而不是等到攒了几百篇再回头整理那会非常痛苦。6. 下一步还能怎么扩展如果你已经读完了上面提到的经典论文也动手跑过一些实验接下来可以考虑这几个扩展方向关注编译器与 LLM 的结合尤其是用大模型辅助生成调度代码或 kernel 代码的工作。我预测这个方向未来一两年会持续升温突破口可能出现在 cost model 的构造而不是直接让 LLM 写 kernel。关注异构计算和多设备编译问题比如 CPU/GPU/NPU 混合场景下的任务切分与数据搬移优化。现在更迫切的问题是 IO 和调度的协同设计单纯要求峰值算力提升意义已经不大了。关注训练场景下的编译器优化。很多经典的 ML 编译工作都聚焦推理场景因为推理没有反向传播、梯度更新这些复杂依赖。但训练场景的收益也非常大比如 PyTorch 2.x 的 TorchInductor 就在训练性能上下了很多工夫。关注编译器自动差分能力的实现。自动微分和编译器的结合能让我们在自定义算子时同时得到高效的前向和反向 kernel这在框架开发中极为有用。我觉得 ML Compiler 这个方向最吸引人的地方在于它既有理论深度又有非常强的工程反馈。你写出的每一行优化代码都能直接体现在延迟或吞吐的数字上这种即时的、量化的正反馈在整个软件工程里都比较少见。如果你已经把这条路走到了这里那剩下的就是持续动手保持对最新论文和工程实践的好奇心我希望这份持续更新的系列论文解读能成为你在这条路上的一位靠谱同行者。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价