Triton编译器及Pass管理器：你写的GPU代码，是怎么变成机器指令的？

本文整理自视频《【AI实操 · 优化篇】03 Triton编译器及Pass管理器》是先进编译实验室系列课程优化篇第三讲。前两集我们分别聊了 Triton 在 PyTorch 编译加速里的角色、以及怎么用性能工具找瓶颈。这一集终于要揭开最神秘的部分编译器到底对你的代码做了什么Pass又是个什么东西面向小白所有难懂的概念我都会掰开揉碎讲。痛点开场你写代码凭什么指望机器听懂很多同学被编译器这三个字劝退过。一听到编译“IR”“LLVM”第一反应就是这是搞底层的大佬才配聊的东西跟我这种调 API 的有什么关系但真相是只要你在 PyTorch 里碰过 Triton写过一个triton.jit的 Kernel你的每一行代码都已经悄悄走完了一场编译器之旅。而这场旅程里最核心的调度者就是这一集的主角——Pass 管理器。搞懂了它你会发现一件很爽的事以前觉得是黑魔法的性能优化其实都藏在编译器的某个Pass里。今天我们就把它一层层扒开。一、先打地基编译器、中间表示、LLVM 到底是啥这一节全是小白最怕的概念我一个个用大白话讲清楚。1.1 编译器Compiler人类的翻译官编译器就是一个翻译器。它的工作简单粗暴把人类写的高级语言Python、C、Triton翻译成机器能直接执行的低级指令机器码。类比你想让一个只会说英语的人帮你做事你得自己先学英语吗不用。找个翻译官你说中文他翻成英语对面照做。编译器就是这个翻译官——你说算个矩阵乘法它翻译成 GPU 听得懂的底层指令。1.2 中间表示IR / Intermediate Representation翻译中间的草稿你可能觉得翻译就翻译呗一步到位不就行了翻译官也得先在心里把话过一遍才能翻得准。编译器也一样它不会直接把高级语言翻译成机器码而是先翻译成一种**“中间状态”**这个状态就叫IRIntermediate Representation中间表示。IR 长啥样它不完全是高级语言也不完全是机器码是介于两者之间的东西。特点就是既保留足够的信息类型、运算逻辑又远离了具体的硬件。类比你要从中国去美国不是一步飞到底而是先到中转站、再转机。IR 就是这个中转站。1.3 LLVM编译器界的乐高积木LLVM 不是某个具体产品而是一套开源的编译器基础设施/工具链。它本身自带一套 IRLLVM IR和一堆现成的优化模块。为什么要提它因为Triton 编译器最后也要借助 LLVM——写完 Triton 代码后先翻译成 Triton 自己的 IR最后再落到底层的 LLVM IR剩下的活交给 LLVM 去优化、生成机器码。相当于你的翻译官接了个外包。1.4 Triton 编译器在模型和硬件之间的位置先用一张图看清大局Triton 编译器就站在模型级别PyTorch、LLM和硬件级别GPU、CPU、DSA 芯片中间负责把上层运算翻译成下层指令。这就是 Triton 的核心野心你写一套代码它能帮你适配各种硬件不用像写 CUDA 那样死磕 NVIDIA。二、为什么要自己折腾编译器聊聊 CUDA 的痛你可能想问既然 CUDA 早就有了为啥还要 Triton 编译器视频里讲得很直白CUDA 生态有三大限制CUDA 生态的限制大白话解释硬件编程接口限制只能用 NVIDIA 那套接口换芯片就要重写迁移任务繁重现有 CUDA 代码迁到新硬件工作量巨大底层生态相互隔离 / 以 NVIDIA GPU 为基础生态都是围绕 NVIDIA 造的其它芯片很难接入说白了CUDA 很牛但它是 NVIDIA 的亲儿子别人家的芯片想用很难。而 Triton 想做的是让一套代码通吃多种硬件。同时 Triton 也搞了个FlagGems 算子库联合一大票国产芯片厂商寒武纪、天数智芯、沐曦这些和自己的合作伙伴用算子库的形式衔接 AI 框架和编译器把生态做起来。三、核心来了Triton 编译器到底怎么干活这是全片的精华。我们跟着一段 Triton 代码走完它从人看的语言到机器跑的指令的全过程。3.1 一次编译的完整路线视频里给了这样一张编译流程图图上这条主线请你记牢从左到右你写的 Triton 语言代码 ↓ Triton IRTriton 自己的中间表示 ↓ 经过一堆优化 Pass LLVM IRLLVM 的中间表示 ↓ PTX / 其它后端指令 ↓ 机器码GPU 真正执行的东西V 关键点Triton IR 是和上层的 Triton DSL 语言一一对应的——你写的每一条 Triton 语句都会先变成一个对应的 IR 节点。3.2 Dialect方言IR 也分派系视频里反复出现一个词Dialect直译方言。这是什么意思在 MLIR一个更底层的编译器框架Triton 编译器就是基于它重构的体系里IR 不是只有一种而是分了很多种方言各管一摊有的方言管 CPU 相关的运算比如 Linalg有的方言专门管 GPU比如 Triton GPU Dialect有的方言管更底层的 LLVM可以理解成一个公司里分了很多部门前端的部门处理你写的代码中台部门做分析优化后端部门对特定硬件发货。编译器就是一个部门一个部门地把代码交接过去每交接一次代码就更接近机器码一步。注意上面这张图里标出的主要讲解路线虚线框起来的 GPU 相关部分Triton Language → Triton Dialect → Triton GPU Dialect → LLVM最后落到 GPU。这是本片的主角路线。3.3 关键优化访存合并Coalescing编译器能做的优化之一就是访存合并Coalescing。GPU 里有成千上万个线程一起干活。如果这些线程各读各的地址数据就得一块一块慢慢搬但如果让它们合并着读一段连续的内存一次就能搬一大块速度突飞猛进。类比一个班 50 个人都要去食堂如果大家各走各的路乱作一团如果排好队一次走过去连续访问又快又整齐。Coalescing 就是让内存访问排队走的优化。除此之外编译器还会做Layout布局分配这一系列优化——比如你的张量到底怎么映射到线程CTA Layout、放进共享内存时怎么摆放Shared Layout、怎么避开银行冲突Bank、Swizzle这些都是决定性能的细节这些细节非常硬核小白只需记得一句话编译器里有一堆这样的优化工序每一道工序就是一个 Pass。四、重头戏Pass 到底是什么4.1 Pass 的定义大白话版Pass优化趟/优化通道就是编译器里一道独立的处理工序。编译过程不是一步到位的而是拆成很多很多步每一步只做一件小事改完之后把结果交给下一步。这样一步就叫一个 Pass。类比一条汽车生产线。一个工人只装轮胎一个 Pass另一个工人只装车门另一个 Pass第三个工人只做喷漆又一个 Pass。每一个 Pass 只管自己的活改完交给下一个。把无数个 Pass 串起来一辆车就造好了。4.2 Pass 的正式定义视频原话视频里给了一张专门讲 Pass 定义的图简化后的核心意思Pass 是对程序IR进行的一次变换输入一个中间表示经过处理输出一个变换后的同类结构。它做的事情要么是优化让代码更快、更省内存要么是转换把代码从一种表示换成另一种表示。在 MLIR 体系里Pass 可以作用在不同层级的东西上Operation操作、Function函数、Module模块、Loop循环等还可以做Analysis分析——不改变代码只是诊断一下代码状况为后面的 Pass 提供情报。4.3 Pass 管理器PassManager总调度者好的Pass 有这么多、有先后依赖关系、还要传递数据前一个的分析结果要给后一个用谁来安排——Pass ManagerPass 管理器。它就是编译器的流水线班长决定注册哪些 Pass决定这些 Pass 的顺序依赖关系不能乱一个一个地把 Pass调度执行起来视频把它的原理概括成了清晰的步骤核心就是注册 → 调度 → 执行一环扣一环保证上一道工序的产出正好是下一道工序的输入。五、手把手实操在 Triton 里加一个自己的 Pass光讲概念不过瘾视频还带着我们真的写了一个自定义 Pass——用来检查代码里有没有加法运算AddIOp的 PrintAddIOp Pass。我们捋一下它的完整加装流程。5.1 写 Pass 的 C 代码先写一个 C 类继承 MLIR 的机制比如mlir::PassWrapper实现runOnOperation()方法遍历 IR 里的运算看到 AddIOp 就打印出来5.2 声明 Pass 注册在Passes.td文件里用 TableGen 声明这个 Pass比如def TritonGPURintAddIOp : Pass...声明它依赖哪些 DialectdependentDialects在triton.cc老版本或passes.cc新版本里把 C Pass 注册、并暴露给 Python5.3 挂到编译流程上最后在compiler.py里把新 Pass 插进 PassManager 的流水线讲完编译原理Python 侧再调用它就被执行了。编译好的 Triton 装上之后写个测试脚本一跑终端里就输出了“Matched AddIOp”——说明你的自定义 Pass 真的在编译期干上活了5.4 注意Triton 3.2 的流程变化视频专门点了最新的Triton 3.2 版本添加 Pass 的三大变化别照着旧教程踩坑改文件从老的triton.cc改到passes.cc等新入口文件Pass 注册方式变化一部分变成自动定义 Pass 接口compiler.py 路径变化要从新的backend/compiler.py里挂载六、一张表消化全部核心概念概念大白话解释编译器把人类高级语言翻译成机器码的翻译官IR中间表示编译过程中的中转站既保留语义又贴近底层不绑定具体硬件中间表示就是 IR 的中文叫法同一概念LLVM一套开源编译器基础设施自带 IR 和优化模块Triton 最终借它落地Dialect方言IR 里的不同派系各管一类运算CPU/GPU/底层编译器按部门交接代码Pass优化趟编译过程中的一道独立工序只做一件小优化/转换输入 IR 输出 IRPass 管理器总调度者负责 Pass 的注册、排序、调度执行访存合并Coalescing让线程排队访问连续内存一次搬一大块提升访存速度优化趟即 Pass 的另一种翻译同一概念七、总结这一集我们把 Triton 编译器的内部结构摸了个底朝天编译器 翻译官把 Triton 代码一步步翻译成机器码中间过程靠IR / Dialect逐层交接Triton Language → Triton IR → LLVM IR → 机器码每道优化工序就是一个PassPass 管理器负责统筹调度想真正懂性能优化就去编译器里看那些 Pass 在干什么——访存合并、Layout 分配这些都是性能的秘密武器你甚至可以自己写一个 Pass插进编译流程观察、修改、优化你的 Kernel。一句话记住全片你写的 Triton 代码不是被一步翻译的而是被一大串 Pass接力调优后才最终成为 GPU 上飞驰的机器指令的。参考【AI实操 · 优化篇】04 Triton算子关键参数优化