最近两年AI编译器这个词在技术圈的热度一路走高。我身边不少做深度学习框架、系统优化甚至后端开发的朋友都在问同一个问题这个东西到底该怎么入门网上的资料要么是论文级别的大部头要么是纯打广告的培训课真正能一条线串起来、照着就能下手的学习资料少之又少。我当时转AI编译器方向的时候也走过不少弯路。拿到一份LLVM代码不知道从哪看起读了三天TVM源码感觉自己什么都懂了一关掉编辑器又什么都想不起来。后来靠着一个一个项目磨、一遍一遍过IR dump才算把这条路的骨架摸清楚。所以这篇东西我就直接把我认为最靠谱的学习路线、项目清单、以及每个环节真正需要注意的坑一次性整理出来。不管你是刚毕业的学生还是从应用层转底层优化这篇文章的思路都可以直接用。1. 先想清楚AI编译器到底解决的是什么问题很多初学者一上来就抱着TVM或MLIR的源码啃结果越看越迷糊。根本原因是你不知道这些代码服务于什么目标自然就看不懂它为什么这么设计。所以在碰任何代码之前先把AI编译器的定位和核心问题搞明白。1.1 为什么传统编译器不适用于深度学习场景传统编译器比如GCC、LLVM解决的核心问题是把高级语言翻译成目标机器指令同时做各种优化。它的输入是C/C这种通用语言输出是CPU或GPU的机器码优化目标是让程序运行得更快、代码更小。但深度学习推理和训练场景有一些完全不同的特点计算模式高度固定。深度学习模型本质上是一张大计算图图中的节点是矩阵乘、卷积、归一化、激活函数这些算子边是张量数据流。这个计算图在执行前就已经完全确定了不需要像通用程序那样面对无穷多的控制流分支。性能瓶颈在数据搬运和访存模式。深度学习算子在GPU上跑得慢通常不是因为ALU利用率不够而是因为数据从全局内存搬到共享内存、再从共享内存搬到寄存器这个过程中出现了大量等待和浪费。这跟传统编译器主要关心指令选择和寄存器分配是很不一样的。目标硬件高度异构。NVIDIA的GPU、AMD的GPU、华为的昇腾、各种NPU每个硬件都有自己的内存层次结构、并行模型和指令集。同一个卷积在这张卡上要用implicit GEMM的写法在另一张卡上可能用Winograd更快。传统编译器针对的是相对统一的指令集架构而AI编译器面对的是碎片化的加速硬件生态。AI编译器做的事情简单说就是给定一个模型描述和一块目标硬件自动生成在这块硬件上最高效的可执行代码。它把模型而不是语言作为编译输入把算子的高效实现而不是通用指令选择作为核心优化目标。1.2 AI编译器的三层典型架构虽然不同框架的AI编译器实现各不相同但整体架构高度相似通常分为三层层次典型组件作用前端PyTorch的TorchScript/TorchDynamo、TensorFlow的Grappler、ONNX把训练好的模型转换成计算图做图级优化算子融合、常量折叠、死代码消除中间层TVM Relay、MLIR的多种Dialect、XLA HLO对计算图做硬件无关的优化把图降低为更底层的中间表示后端LLVM、NVCC、各种Codegen把中间表示转换成目标硬件的机器码或者可调用的算子库初学者脑子里装上这个三层架构再去看任何AI编译器代码心里就有了一张地图你现在看到的这段代码到底是在做前端的图变换还是中端的IR优化还是后端的代码生成。定位清楚阅读效率能提升一倍。2. 打地基阶段三门必修课和一条贯穿主线很多人问我做AI编译器需要把编译原理全部学完吗我的答案是不需要啃完龙书但也有绕不开的底线。下面这三块内容是必修缺一块后面都会卡壳。2.1 硬件体系结构你是在跟机器打交道AI编译器最终生成的代码要跑在GPU/NPU上所以你必须清楚这些硬件的执行模型。这不是嵌入式系统那种了解寄存器的级别而是要理解到看到一个算子就能大致估算出性能瓶颈在哪的程度。以NVIDIA GPU为例最需要掌握的几个核心概念线程层次grid、block、warp以及它们如何映射到SM流式多处理器上内存层次全局内存、共享内存、寄存器、常量内存各自的容量和延迟差异。共享内存通常只有几十KB到上百KB全局内存延迟却有几百个周期。一个算子优化得好不好关键看你有没有把全局内存的复用变成共享内存的复用Tensor Core从Volta架构引入专门做矩阵乘加运算。fp16、bf16、int8这些低精度类型配合Tensor Core吞吐量可以远超普通CUDA CoreSIMT执行模型与warp divergence一个warp里的线程要执行同一指令分支分叉会导致性能骤降。推荐资源NVIDIA的CUDA C Programming Guide里的Programming Model和Memory Hierarchy两章是必读的。另外可以去看陈天石的GTC演讲或者各种GPU体系结构教程对建立硬件直觉很有帮助。判断标准如果你看到一篇优化卷积的文章里说把im2col改成implicit GEMM以减少内存开销你能立刻理解它在说什么这关就算过了。2.2 编译原理核心抓住IR和Pass这两个命门传统编译原理里面对AI编译器最有用的部分不是词法分析也不是自顶向下的语法分析而是中间表示IR的设计和Pass基础设施。AI编译器的优化本质上就是一趟一趟的Pass分析IR发现某种优化机会改写IR得到新的IR。MLIR这个框架把这个思想推向了极致——它定义了一整套可扩展的IR基础设施允许不同抽象层级用不同的Dialect来表达。学这一块我推荐两条路径并行读LLVM的官方文档中关于IR的部分特别是LangRef的前半部分搞清楚什么是basic block、什么是dominance、什么是use-def链去看一下MLIR toy tutorialChapter 1到Chapter 5这是MLIR官方提供的入门教程用Toy语言把MLIR的Dialect、Pass、Lowering全走了一遍做完之后你会对If youre building an IR, youre building a compiler这句话有具身理解。我不建议一上来就啃龙书的词法分析那一章对AI编译器方向性价比太低。2.3 深度学习框架基础要知道模型长什么样AI编译器处理的是模型所以你必须能看懂模型代码理解张量的shape、dtype、layout和算子的语义。不需要你会手写ResNet训练但至少要能做到读懂PyTorch模型定义里forward函数的数据流知道什么是静态shape、什么是动态shape哪些算子是有状态的比如BatchNorm的running_mean了解Tensor的stride和memory formatNCHW/NHWC对算子性能的影响。这些知识在后续做图优化和算子优化时会频繁用到。比如你看到一个FoldScale的pass如果不理解BatchNorm推理时可以把scale和bias融合进卷积的权重你就无法理解这个pass为什么要存在。2.4 贯穿始终的主线计算与访存的平衡如果说有一个思维模式是AI编译器干活的核心主线那就是一句话计算要并行化访存要局部化。为什么这么讲GPU有几千个核心算力很强但数据搬运是瓶颈。一个矩阵乘法理论上需要执行的浮点运算次数是固定的你能优化的是相邻线程访问的地址是否连续合并访存、数据是否能在共享内存中复用减少全局内存访问、指令流水线是否被计算密集型指令占满隐藏访存延迟。学AI编译器的时候每接触一个优化技术都可以问自己两个问题它是让计算更并行了吗还是让访存局部性更好了绝大多数优化都能归到这两类中的一类。有了这条主线你学到的各个零散知识点才能串起来。3. 项目清单从造轮子到读源码六个项目循序渐进一步一个脚印学习编译器这类工程性极强的方向光看书是绝对不行的。项目是检验你理解程度的唯一标准。我把做过的和带人做过的项目整理成下面这个清单从零基础到进阶分成了六个项目每个项目都有明确的目标和验收标准。3.1 项目一手写一个微型C语言子集编译器2-3周这个项目听起来跟AI编译器没关系但它练的是编译器的核心基本功流水线。输入是一个简化版C语言比如只有int类型、加减乘除、赋值、函数调用输出可以是LLVM IR或RISC-V汇编。具体任务拆解写词法分析器把源代码切成token流写递归下降语法分析器识别表达式、语句、函数定义生成AST抽象语法树写一个简单的IR生成器输出LLVM IR调用clang/llc把IR编成可执行文件跑通测试用例。这个项目最大的价值在于你会亲身体会到IR作为前后端分界线的意义前端只需要生成合法的IR剩下的所有优化和代码生成都由后端完成。做完之后你再去看TVM里的Relay IR会有一种这东西我见过的亲切感。参考资源国内很多高校的编译原理课程设计都有类似题目。也可以用llvm的Kaleidoscope教程入门它的语言是一个简单的函数式语言但麻雀虽小五脏俱全从Lexer到JIT都有。3.2 项目二用TVM为自定义算子生成CUDA核函数3-4周这是从编译器理论跨到AI编译器实践的关键一步。目标是不手写CUDA Kernel而是用TVM的TETensor Expression描述一个算子让TVM生成可在GPU上运行的代码。建议从GELU高斯误差线性单元这种简单算子开始然后做LayerNorm最后挑战一个带reduction的算子比如softmax。核心步骤import tvm from tvm import te # 以softmax为例先用TE描述计算逻辑 n, m te.var(n), te.var(m) A te.placeholder((n, m), nameA) A_max te.reduce_max(A, axis1, keepdimsTrue) A_exp te.exp(A - A_max) A_sum te.sum(A_exp, axis1, keepdimsTrue) B te.compute((n, m), lambda i, j: A_exp[i, j] / A_sum[i, j])写完TE描述之后用schedule做优化先加向量化再做线程绑定然后分析生成的CUDA源码对比手动实现版本的性能差异。这个项目练的是算子即程序的思维一个算子的实现不是死的而是可以通过调度策略改变的。你会深刻理解为什么计算定义与调度分离是TVM对深度学习编译器最大的贡献之一。3.3 项目三分析并改掉一个TVM/MLIR的Bug4-6周这个项目要求你主动去GitHub的issue区找一个AI编译器相关的bug理解它、修复它、提交PR。为什么这个项目价值极高因为你在写代码的时候无论读了多少源码都只是被动理解。但修bug要求你主动定位一个问题产生了它是在前端被错误转换了还是中端IR改写时有语义差异还是后端的代码生成没覆盖某个边界条件我记得自己修过的第一个bug是TVM里某个算子schedule在特定shape下的错误绑定。排查过程花了两周但这两周我几乎把schedule绑定的源码读了个遍对TVM的认知深度远超之前三个月的泛读。 操作建议不要挑那种一眼就能看出问题的bug选那种附带了完整issue描述和复现步骤的。在没有把握之前先在issue里回复自己的分析思路维护者有时候会给你非常有价值的提示。3.4 项目四从零搭建一个微型算子编译器6-8周这个项目是进阶路上的硬骨头——用LLVM的框架搭一个只服务于少数几个算子的小型编译器。输入是一个简化的计算图描述JSON就行输出是LLVM IR再通过LLVM生成可执行代码。关键任务定义一个简单的DSL可以直接用JSON或Python不用重新发明语言对每个算子矩阵乘、卷积、ReLU、Add编写Lowering逻辑转成LLVM IR接上优化Pass比如循环不变量外提把生成的代码跑在CPU上和对拍库如Eigen对比推导正确性。这个项目的价值在于它会逼着你把AI编译器的骨架亲手搭一遍。你会发现TVM/MLIR里很多看似抽象的设计其实都是你这个小编译器面临的问题在更大规模下的映射。从项目三到项目四是一道分水岭项目三让你读懂别人的编译器项目四让你能写自己的编译器。能完成项目四的人基本就已经具备做AI编译器研发的能力了。3.5 项目五用MLIR实现一个自定义Dialect和Pass6-8周MLIR是当前AI编译基础设施的事实标准不会MLIR的AI编译器工程师在市场上竞争力会大打折扣。项目四如果已经让你对LLVM有了良好的手感项目五就是专门练MLIR的。任务设计定义自己的Dialect比如叫toy规定它的Operation、Type、Attribute实现从你自己的DSL到toyDialect的Lowering实现若干个Pass比如把一个toy.matmultoy.add的pattern融合成一个toy.fused_matmul_add实现从toyDialect到LLVM Dialect的Lowering用MLIR的单元测试框架给pass写测试。这个项目做完之后你对MLIR的多级IR、Dialect间显式Lowering、基于Pattern的Rewrite这三个核心思想会有很直观的感知。3.6 项目六为一个真实加速器写一个Codegen后端10周以上这是终极项目也是AI编译器工程师日常工作的真实写照你手里的加速器有自己的指令集但没有编译器你的任务是把某种IR翻译成这个指令集。如果手边有NPU开发板当然最好没有的话可以考虑用CPU的SIMD指令集或者RISC-V的向量扩展RVV来模拟。核心是输入是某种中间表示输出是目标平台可执行的指令序列。这个项目练的是全栈能力。你会遇到TVM里CodeGenSourceModule的注册机制、寄存器分配器的实现、指令调度的顺序等等。能从头到尾跑通一个端到端demo模型进来调优代码出去性能还不差就已经具备AI编译器方向的硬实力了。4. 避坑指南与实用技巧AI编译器的资料确实多但正因为多不知道怎么选、怎么读反而成了最大的门槛。下面这几点是我踩过的坑希望你能绕开。4.1 不要一上来就看源码要从顶层设计向下看这个问题我见过太多次了。初学者特别喜欢把TVM或MLIR的仓库clone下来然后从main函数开始一行一行读。结果就是代码里各种Visitor、Pattern、Dialect来回跳一周下来脑子变成一团浆糊。正确方式是从数据结构入手以IR为主轴先弄清楚这个编译器里有哪几种IR比如TVM里的Relay IR和TIRMLIR里不同Dialect再看这些IR之间的转换关系哪个pass把高层IR降低到低层IR最后才需要去关注某个具体的pass实现细节。我自己读TVM源码时的路径是先跑一个端到端的例子用print(relay_ir.astext())把IR打印出来对照着IR结构去源码里找对应的类定义。IR是编译器的用户界面先把用户界面摸熟内部实现就只是时间问题了。4.2 不可盲目追新官方tutorial的版本匹配很重要TVM、MLIR这种项目发展速度极快API经常变动。你网上搜到一篇三年前的博客里面的代码放到现在的版本几乎肯定跑不起来。所以当你决定要用某个框架做项目时第一件事就是去它的官方仓库把版本固定住然后找相同版本的文档和tutorial。不要贪心用master分支也不要Windows下的Visual Studio编译直接在Linux环境里搞定最简单。 注意MLIR的toy tutorial是跟着LLVM主分支走的API变动非常频繁。如果发现代码编译不过优先去查LLVM的release note看这个接口在哪个版本改掉了。4.3 手写算子对比是验证理解最好的方式这个建议主要针对项目二、项目三。很多人用TVM生成了一段CUDA代码跑通之后就觉得自己会写算子编译器了。这是错觉。我建议你在TVM之外再用手写一个CUDA kernel实现同样的算子然后用NVIDIA Nsight或者简单的计时函数去对比性能如果手写版本更快去看TVM生成的代码差距在哪如果TVM版本更快去看它的schedule为什么更优。这个过程会逼着你把生成代码和手写优化两边的知识都补齐。能说清楚某个性能差异是为什么产生的才说明你真正理解了优化策略。4.4 阅读Meetup记录和大会Talk建立产业全局观AI编译器是一个快速的产业。每年LLVM Developer Meeting、CGO、MLSys等会议都有很多前沿的实践分享。这些Talk通常由一线工程师来讲比论文更贴近工程落地。比如去听一下PyTorch团队讲TorchInductor的设计演进你就能知道Meta内部对编译派 vs 算子派路线之争的取舍到底是什么。养成在跑步、通勤时用听Talk的方式积累产业动态的习惯对面试、选方向、甚至判断技术趋势都很有帮助。5. 学习时间规划四个月从入门到可上手下面是我觉得比较现实的四个月周期安排按每周大约15-20小时的学习量来算阶段时间重点内容完成标志基础期第1-4周硬件体系结构、编译原理核心、PyTorch基础动手跑完项目一理解GPU执行模型知道IR是编译器的分水岭框架期第5-8周完成项目二选一个框架TVM或MLIR精读其tutorial和核心源码能独立为自定义算子生成代码能流畅讲出TVM/MLIR的架构分层深潜期第9-12周完成项目三修bug或项目四微型编译器能定位并修复一个真实issue或者能写出一个能跑的小型编译器进阶期第13-16周走MLIR相关项目项目五或项目六能独立开发一个Dialect并实现Pass时间只是一个参考。关键是每一阶段的目标必须可验证比如能跑通能修掉能写出来而不是看了一遍。宁可慢一点不能囫囵吞枣。在这个行业里看上去会了和真的会了之间的差距往往就是这些可验证产物的差距。模块化恐惧症没有用动手打开编辑器才是唯一的解药。最后分享一个最近在带新人的时候常用的小技巧把TVM的Relay IR里认为最晦涩的pass挑出来比如FoldScaleAxis先用一手资料搞清楚它做了什么然后在代码里加上日志亲自跑一个例子观察IR变化最后用一句大白话把自己忽悠明白——如果这句话能让你旁边的人听懂你才是真的懂了。这个方法对所有AI编译器框架都适用。希望这份清单能帮你少走一些弯路四个月后我们可能在同一个IRC频道见。