AI编译器学习路线：从零基础到进阶实战

1. 写在前面AI编译器到底是什么为什么现在这么火这两年大模型爆发之后“AI编译器”这个词从学术圈一路火到了工程圈。很多人一听到“编译器”三个字第一反应还是大学里那门《编译原理》课程觉得不就是把高级语言翻译成机器码吗实际上AI编译器做的事情远不止翻译它更像是“优化器转换器调度器”的结合体专门把深度学习模型从一种描述形式转换并优化成可以在特定硬件上高效运行的代码。举一个最直观的例子你用PyTorch写了一个Transformer模型在GPU上跑得好好的但同样的模型部署到华为昇腾、寒武纪或者嵌入式NPU上如果没有AI编译器的适配和优化性能可能直接掉一个数量级。AI编译器的作用就是在不改变模型语义的前提下通过对算子融合、内存布局重排、循环展开、指令调度等手段让同一套模型在五花八门的硬件上都能跑出接近理论峰值的效果。从行业需求来看几大AI芯片厂商都在疯狂招编译器工程师薪资普遍比普通后端开发高一截。为什么因为芯片本身只是硬件真正让芯片发挥价值的是软件栈而AI编译器是软件栈中最核心、最复杂的一层。国内的华为、阿里、百度、字节国外的Google、Meta、NVIDIA都在AI编译器上投入了巨大的人力。说白了AI编译器这个方向是目前AI基础设施领域里少有的“高门槛、高壁垒、高回报”的赛道。这篇内容我打算给出一套从零基础到进阶的完整学习路线并且附带一份可以直接照着做的项目清单。不管你是计算机专业的学生、想转行的在职开发还是已经在做深度学习推理优化但缺乏系统知识的工程师都可以参考这条路线来规划自己的学习节奏。2. 学习AI编译器之前先想清楚三件事2.1 你的基础决定了你的起点AI编译器是一个交叉领域涉及编译原理、计算机体系结构、深度学习框架、高性能计算等多方面知识。基础不同的人学习路径是完全不一样的。如果你是没有接触过编译原理的纯应用型开发者我的建议是先不要急着去啃TVM源码而是花一个月时间补上编译原理的核心知识点词法分析、语法分析、中间表示、基本块优化、寄存器分配这些概念不需要达到能手写编译器的程度但至少要理解编译器的工作流程和每一层的职责。如果你已经学过编译原理但对深度学习模型的计算图不熟悉那么重点是理解计算图、算子、张量这些概念把深度学习框架和编译器之间的接口关系搞清楚。如果你是做推理优化出身熟悉TensorRT或者OpenVINO的使用那么你的起点其实很高可以直接从中间表示层切入重点研究算子融合的原理和自动调优机制。2.2 学习AI编译器不能只看源码很多新手有个误区觉得学AI编译器就是把TVM的源码从头到尾读一遍。这种想法不能说错了但效率极低。TVM的代码量很大涉及的模块非常多如果没有任何上手实践的机会纯粹看代码很容易陷入细节的泥潭越看越迷茫。更合理的策略是“先会用再深入”。先把TVM或者MLIR这类框架跑起来跑通一个简单的模型优化流程理解各个pass的作用然后再针对自己感兴趣的模块做源码级别的深入。也就是说把源码阅读当作一种按需查询的手段而不是学习的主线。2.3 硬件知识是绕不开的坎AI编译器最终是要生成在特定硬件上运行的代码所以硬件知识绝对绕不开尤其是体系结构层面的知识。你需要了解GPU的线程模型、共享内存、寄存器文件、存储层级需要了解NPU的阵列结构、片上缓存、数据通路否则你根本理解不了为什么一个简单的算子融合能带来那么大的性能提升。这并不是要求你把CPU/GPU/NPU的微架构细节都背下来但至少要建立起“代码性能受硬件资源约束”这一基本观念。我第一次看TVM的调度原语时对bind和fuse这些操作完全无感后来认真学习了GPU的线程组织方式和存储层次之后才真正理解了这些调度操作背后的硬件逻辑。3. 学习路线全景图从入门到进阶的四个阶段3.1 阶段一打好编译原理和体系结构的底子这个阶段的目标是建立必要的知识框架不需要追求精通但核心概念务必扎实。推荐学习的核心内容包括编译器的经典两段式结构前端、优化、后端、抽象语法树与中间表示、数据流分析与循环优化、指令级并行与存储层次结构。学习资料方面国内最经典的教材是《编译原理》龙书和《计算机体系结构量化研究方法》Patterson版。说实话这两本书的内容都偏厚从头啃完需要很大毅力。我的建议是带着问题去读比如学中间表示时就思考“为什么AI编译器需要比传统编译器更高级的IR”学循环优化时就思考“这个优化技术在深度学习模型里对应哪个场景”。如果觉得两本大厚书压力太大可以先看《CSAPP》深入理解计算机系统的处理器和存储层次两章再配合《Engineering a Compiler》中的中间表示与优化章节。这条路更适合工程背景的人内容更紧凑和后续AI编译器学习的衔接也更自然。3.2 阶段二理解计算图和深度学习框架的底层机制这个阶段的重点是搞懂深度学习模型在框架层面的表示形式和执行流程。你需要弄清楚以下问题计算图是如何构建的静态图和动态图的区别在哪里PyTorch的算子是如何注册和调度的TensorFlow的GraphDef和MLIR之间是什么关系。实操建议是阅读PyTorch的源码中关于autograd和torch.fx的部分这两个模块恰好覆盖了动态计算图追踪和静态图捕获两种形态。torch.fx目前是PyTorch做模型变换和编译优化的核心入口理解它的工作机制对后续学习AI编译器非常有帮助。手上可以跑的实验使用torch.fx将一个小型模型转换成GraphModule然后用GraphModule的graph对象打印出所有节点观察算子的连接关系和依赖顺序。这个实验看起来简单但能让你对“计算图”形成直观认知。3.3 阶段三上手主流AI编译器框架这个阶段是学习路线的核心部分目标是掌握一到两个主流AI编译器框架的完整使用流程和核心原理。当前生态中比较值得投入学习的框架有TVM含Apache TVM和其分支Meta TVM、MLIRLLVM社区、XLAGoogle、以及国内厂商的编译器栈如华为的MindSpore AKG、阿里的BladeDISC等。我的建议是首选TVM。虽然TVM的社区活跃度和代码维护情况不如巅峰期但它依然拥有最完整的AI编译器学习素材从计算图级别优化到算子级优化再到代码生成整条链路清晰可学。另外一个重要原因是TVM的Relay IR和TIR的层次划分非常经典理解这个分层思想之后再看MLIR会轻松很多。MLIR是更高阶的内容它的设计理念是“可扩展的编译器基础设施”不只是为AI服务还能用于CPU、GPU、自定义硬件等多种目标。学习MLIR需要一定的LLVM基础建议放在TVM之后再学并且侧重点放在Dialect的转换和通用优化pass的编写上。3.4 阶段四深入算子优化与自动调优到了这个阶段你已经有能力阅读编译器框架的源码并且能自己动手写优化pass了。你的目标是掌握算子级优化的核心手段包括算子融合将多个算子合并成一个kernel减少访存、向量化利用SIMD指令或者GPU的向量单元、循环变换分块、交换、展开、存储层级优化通过共享内存或缓存重用数据。同时要掌握自动调优的机制TVM中的Ansor和MetaSchedule是目前最值得研究的自动调优模块。理解自动调优的搜索空间定义、代价模型、进化搜索策略能让你对“编译器如何自动找到最优实现”这件事有非常深入的理解。在这个阶段可以尝试自己实现一个简单的算子融合pass也可以尝试把手写的高性能卷积kernel与自动调优生成的结果进行对比分析这些实践对建立性能优化直觉非常重要。4. 核心概念精讲IR、Pass和调度4.1 从一份代码到一份高性能kernel的完整旅程一个深度学习模型在AI编译器内部会经历多次形式转换每次转换都发生在不同层级的中间表示之间。以TVM为例一个PyTorch模型先被转换为Relay IRRelay IR经过图级别的一系列pass优化后被lower成TIRTIR再经过一系列的循环变换和存储优化最终通过代码生成模块变成可以在目标硬件上执行的代码。整个过程中最核心的三个概念就是IR中间表示、Pass优化过程、Schedule调度。我把它们对应到普通软件开发的场景里IR是数据的结构定义Pass是对数据结构的变换函数Schedule是变换时的可选参数和策略。这个类比不一定完全准确但能帮助初学者快速建立整体认知。4.2 多层级IR为什么一个编译器需要好几种中间表示传统编译器通常只使用两到三种IR而AI编译器往往需要在多层IR之间转换这是由AI模型的特性决定的。AI模型的计算图级别信息比如卷积和ReLU之间的连接关系依赖图结构来表达适合用图IR而算子内部的循环结构、内存访问模式适合用带语句级的IR来表达。如果只用单一IR要么无法表达高层语义导致丢失优化机会要么无法生成底层高性能代码。所以TVM设计了Relay和TIR两层IRMLIR则通过Dialect机制实现了无限扩展的多层IR。理解了多层级IR的必要性你再看编译器框架时就不会觉得“搞这么多中间表示纯粹是增加复杂度”。4.3 Pass与Schedule的关系Pass在编译器中是一个对IR进行遍历和变换的过程比如“常量折叠pass”就是把可以提前计算的表达式计算掉。每个pass是相对固定的输入一种IR输出同一种或另一种IR。Schedule则不一样它是用户对算子实现策略的显式控制。以TVM的调度原语为例fuse将两个循环合并split将一个循环拆分bind将一个循环绑定到具体的线程轴。这样设计的好处是算子的逻辑表达compute和算子在不同硬件上的具体实现schedule被分离开来同一个算子逻辑可以针对不同硬件编写不同的调度策略。5. 项目清单从验证性项目到研究型项目5.1 入门项目跑通一个端到端的模型优化流程第一个项目不需要太复杂目标是完整跑通“模型-优化-部署-验证”的流程。建议选择TVM作为工具选一个轻量级的模型比如MobileNetV2或者ResNet18完成以下步骤用PyTorch训练或者直接下载预训练权重转换为TVM的Relay格式编译到目标硬件有GPU就用GPU后端没有GPU就用CPU的LLVM后端对比编译前后的推理延迟和精度。这个项目做完你就已经对AI编译器的基本工作流程有了完整的体验。建议在过程中记录每个pass对性能的影响如果性能没有提升也要记录并尝试分析原因这个过程比单纯看文档有价值得多。5.2 进阶项目一手写一个算子融合Pass算子融合是AI编译器中最常见也是收益最明显的优化手段之一。以常见的“卷积批归一化ReLU”融合为例不融合时这三个算子会产生多次内核启动和数据往返融合后只需要一个内核。这个项目的实现路径是在你的IR上识别出满足融合条件的子图模式改写IR将多个节点合并为一个融合节点实现融合节点的代码生成逻辑。如果选择TVM框架可以自注册一个pass插件用TVM的pass infrastructure来完成IR遍历和重写。项目验收的标准是融合后的模型推理速度提升明显精度不下降。5.3 进阶项目二为自定义算子编写Schedule现实世界中很多模型会包含自定义算子比如某些新提出的激活函数或者特殊的attention变体这些算子在编译器框架中往往没有现成的优化实现。这个项目的目标是给一个自定义算子编写高效的schedule。难度适中但需要你同时理解算子的计算逻辑和硬件特性。比如在GPU上需要思考如何分配block和thread如何利用共享内存减少全局内存访问如何在计算和访存之间做到overlap。这个项目做完之后你对“调度为什么重要”这件事会有非常深刻的体会。5.4 挑战项目在MLIR上实现一个方言转换MLIR的魅力在于它的可扩展性你可以定义自己的Dialect也可以在现状Dialect之间做转换。这个项目的建议路径是基于MLIR框架新增一个自定义Dialect写上降级lowering的pass把自定义Dialect的操作降级到已有的标准Dialect上然后通过MLIR的代码生成机制输出LLVM IR并最终生成可执行代码。这个项目的深度足够覆盖硕士论文级别的研究是理解MLIR设计思想的最佳途径。但工作量也不小建议在与TVM相关的项目做完之后再进行。5.5 研究探索型项目自动化调优器的改进如果你对AutoTVM或者MetaSchedule有浓厚的兴趣可以尝试对它们做一些改进实验。方向可以是设计新的搜索空间剪枝策略改进代价模型以提升性能预测的准确率或者接入一个新的硬件后端观察自动调优的迁移效果。这个方向比较适合想要发论文或者做学术研究的人。推荐的做法是先复现当前自动调优器在一个benchmark模型上的优化效果然后针对某个薄弱环节提出改进方案用同一组benchmark对比验证。学术研究的关键在于可复现性和对比实验的设计而不是单纯追求性能数字的提升。6. 实操记录以TVM为例跑通一次完整的优化流程6.1 环境准备建议直接在Linux环境下操作使用GPU机器效果更佳。没有GPU也可以先用CPU后端跑通流程。# 创建独立的Python环境建议Python 3.8-3.10 python -m venv tvm-env source tvm-env/bin/activate # 安装TVM pip install apache-tvm # 安装onnx和PyTorch用于模型转换 pip install onnx torch torchvision如果选择从源码编译TVM需要先安装LLVM依赖。源码编译的好处是可以用到TVM的调试版本方便后续断点调试源码。学习阶段先装release版本跑通流程等需要深入源码时再重新编译也不迟。6.2 模型准备与转换先用PyTorch搭建一个简单的卷积神经网络然后导出为ONNX格式。import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3, padding1) self.bn1 nn.BatchNorm2d(16) self.relu nn.ReLU() self.conv2 nn.Conv2d(16, 32, 3, padding1) self.fc nn.Linear(32 * 8 * 8, 10) def forward(self, x): x self.relu(self.bn1(self.conv1(x))) x torch.flatten(self.relu(self.conv2(x)), 1) return self.fc(x) model SimpleNet() model.eval() dummy_input torch.randn(1, 3, 8, 8) torch.onnx.export(model, dummy_input, simple_net.onnx, input_names[input], output_names[output], opset_version13)6.3 TVM编译与性能对比核心代码逻辑如下import tvm from tvm import relay from tvm.contrib import graph_executor import numpy as np # 加载ONNX模型并转换为relay IR onnx_model simple_net.onnx mod, params relay.frontend.from_onnx(onnx_model) # 定义目标平台 target tvm.target.Target(llvm -mcpuskylake) # 编译 with tvm.transform.PassContext(opt_level3): lib relay.build(mod, targettarget, paramsparams) # 创建推理执行器 dev tvm.cpu() module graph_executor.GraphModule(lib[default](dev)) # 执行推理 input_data np.random.rand(1, 3, 8, 8).astype(float32) module.set_input(input, input_data) module.run() output module.get_output(0).numpy() print(Output shape:, output.shape)用time模块记录执行耗时分别测试opt_level0和opt_level3的性能差异。实测下来一个简单的CNN在低优化等级和高优化等级之间的差距可能在20%-50%不等复杂模型差距更为显著。6.4 分析性能差异从何而来如果优化后性能没有提升不要急着沮丧这正是深入学习的契机。排查思路是先用编译报告分析生成代码的算子数量再用TVM提供的profiler工具定位时间消耗最大的算子最后逐层分析是访存瓶颈、计算瓶颈还是调度问题。我在实战中发现新手最常见的误区是只关心整体延迟不关注算子的执行分布导致优化方向完全错误。6.5 观察Pass输出TVM提供了PrintIR这样的调试工具可以在编译过程中的任意阶段输出IR。你可以在PassContext中选择打印所有pass的IR变化逐段观察模型从Relay IR一路lower到TIR再到最终代码生成的完整过程。我第一次看到这个输出时才真正理解了编译器的“编译”到底做了什么事图结构逐步被拆解算子边界逐渐消失循环结构被重组访存模式被重排最后变成一段段面向底层硬件的指令。这个过程就像看着一张建筑蓝图逐步变成一栋真实的楼非常有启发。7. 常见问题与避坑指南7.1 上手阶段最容易掉的坑不少初学者在刚接触TVM、MLIR这类框架时第一件事就是去clone整个仓库然后尝试从头编译。结果折腾几天环境都没配好热情消耗殆尽最后放弃了事。我的建议是不要一开始就源码编译先把预编译包跑通了解基本流程确认要继续深入了再考虑源码构建的事情。另一个常见问题是并行学习TVM和MLIR两个框架。这两个框架虽然理念上有相通之处但代码风格、设计思想、配套工具差异很大同时上手反而容易思维混乱。建议主线只选一个精通之后再横向对比。7.2 学习深度和广度的平衡AI编译器领域知识面非常广从算子实现到自动调优从图优化到代码生成每个方向都有很深的学问。我的建议是遵循T字形学习策略先横向覆盖全链路跑通端到端的流程再选择一个方向深入到底比如专心研究算子融合或者自动调优算法。深度和广度的权衡取决于你的职业目标。如果是面试求职建议以全链路理解为主核心优化手段能讲清原理即可如果是进组做科研那就要在某个点上形成独特的技术优势。7.3 如何有效阅读源码阅读源码需要带着目标来看。不要像看小说一样从头翻到尾而是先确定自己想知道的问题再根据问题定位到相关模块顺着调用链追踪代码逻辑。比如你想知道Relay中的算子融合是怎么实现的就应该先找到融合pass的入口函数再画出一个函数调用关系图逐层向下追踪。我自己的经验是阅读顺序上先看接口定义再看核心数据结构和关键算法最后才是具体的代码实现细节。另外强烈建议自己动手添加print语句把中间IR打印出来把关键变量的值输出出来。纸上得来终觉浅亲手调试过的代码才会真正内化成自己的能力。7.4 项目实战中容易忽视的细节在做优化实验时一个容易被忽视的问题是硬件频率波动导致性能数据不稳定。GPU在长时间高负载运行时会触发热降频CPU也有类似问题。所以做性能测试时一定要做多轮测试取平均值同时监控功耗和温度数据确保对比数据处于同等硬件状态下。另一个问题是精度验证不充分。很多优化操作在理论上是无损的但实际执行时因为浮点运算顺序的改变会导致结果有细微差异。千万不要只看结果的正确性还要关注数值误差的范围是否在可接受范围内。通常建议使用fp32计算误差在1e-5量级内没有太大问题如果超出这个范围就要检查优化方案是否引入了数值稳定性问题。8. 学习资源推荐资料方面TVM官方文档中的《TVM Deep Dive》系列教程目前依然是最好的入门材料。它从张量表达式讲起逐步深入到调度原语和图优化配合代码示例学习效果非常好。MLIR方面MLIR官网的Toy语言教程是入门的最佳选择。它用实现一门玩具语言的方式完整展示了Dialect定义、Pass编写和LLVM代码生成的全过程认真做完这个教程你对MLIR的基本架构就能有完整的理解。论文方面推荐阅读《TVM: An Automated End-to-End Optimizing Compiler for Deep Learning》、《MLIR: A Compiler Infrastructure for the End of Moores Law》、《Ansor: Generating High-Performance Tensor Programs for Deep Learning》这三篇。它们分别代表了AI编译器的经典框架、基础设施思想、自动调优方向三个里程碑读完可以对领域的发展脉络有深入认识。9. 写在最后的个人体会我在学习AI编译器的过程中最大的感受是这个方向的知识体系非常庞大没有任何一个人能同时精通所有子方向。所以千万不要因为某一块知识不懂就觉得自己不适合这个领域。事实上AI编译器的团队协作模式本身就是高度分工的有人擅长图优化有人擅长算子实现有人专注性能调优各司其职共同协作才能把一条优化链路做深做透。如果你问我有没有一条最快的捷径我的答案是找一个实际存在的性能问题然后想尽一切办法去解决它。解决问题的过程会逼迫你把编译器框架、硬件体系和性能分析工具融会贯通这比按部就班读十篇文章都更有用。我在学习TVM的过程中最快速的成长阶段就是闷头打磨一个自定义算子的调度策略、反复对比性能数据的那几周。最后再分享一个小建议学AI编译器一定要养成写实验记录的习惯。每次修改了什么优化策略、性能数据如何变化、为什么会变化都要认真记录下来。这些记录不仅是后续调优的参考依据也是你对编译器性能优化建立直觉判断能力的最重要材料。