最近在系统性过一遍Triton源码第一篇就挑了Combine这个Pass。起因很实际我写的Triton kernel里按逻辑拆了两个triton.jit函数编译完dump出MLIR发现IR里是两个tt.func互相call。如果按这个IR直接发射到GPU就是两次kernel launch中间结果得写一遍全局内存再读回来。Combine这个Pass干的事就是把这些跨函数的调用重新拼回一个kernel省掉中间那次显存往返。这篇文章把它的设计思路、源码主流程、一些边界问题讲清楚适合正在读Triton源码的人也适合被kernel性能折腾得想了解编译期优化的人。1. 从碎片化kernel说起Combine为什么值得专门读1.1 我看到碎kernel的方式先描述一下场景。我调试一个带elementwise和累加逻辑的kernel开了IR dump之后看到一个主函数里面稀稀拉拉挂着几个小函数返回值传来传去。第一反应是编译器在搞什么第二反应是这要直接在GPU上跑中间张量全写显存再从显存读回来带宽白花花地流走。具体形态大概是这样主函数里连续调用两个子函数子函数的返回张量直接作为下一个调用的参数。从计算逻辑上看这就是一条挺自然的数据流链但在TTIRTriton Dialect层面它被表示成了tt.call的嵌套。如果不做任何处理后端拿到这个IR会倾向于把它们发射成多个独立的kernel每个kernel对应一次GPU端的完整执行流程。这种碎不是我手工制造的而是Triton前端编译的自然产物。当Python侧把一个大kernel拆成多个triton.jit函数时前端就是各自生成独立的tt.func调用点生成tt.call。至于这些tt.func最后是合成一个GPU kernel还是各发各的取决于优化Pass其中最关键的一个就是Combine。1.2 碎片化的来源不只是用户写了多个函数很多人会有一个误解只要Python侧只写一个triton.jit函数IR里就只有一个kernel。其实不一定。Triton的优化过程里循环转换、条件分支处理、类型扩展都可能在中间层制造新的函数调用。比如带scf.for的循环体里如果编译器认为某个子计算值得复用就可能把它单独提取成一个函数。这种情况在老的Triton版本里更常见新版本逐渐改成以模式匹配和region内联为主但tt.call依然会出现在很多合法场景中。Combine要处理的核心问题就是两个函数之间如果只通过返回值传递数据且这个中间值只被一处使用那么编译器可以把被调用函数的body内联到调用者里让两个逻辑上分离的函数在IR层面先拼起来。这种合并不是传统意义上的函数内联优化——它更多是冲着kernel融合去的。因为一旦拼起来中间结果就不再需要落到全局内存。还要注意Combine和TritonGPU层的layout转换Pass是两码事。Combine工作在TTIR层这个时候还没有layout信息也不关心线程块布局。它只关心数据流关系这个call的结果有没有被其他call消费有没有跨过函数边界。1.3 开销账为什么这值得一个专门的Pass有人可能觉得多个kernel启动也就是几次CPU调用的事能有多大开销我简单算过一笔账。一次kernel launch的CPU侧开销包括参数打包、命令缓冲区提交、驱动调度大约在5到10微秒。这个数字看起来不大但中间结果的传输开销往往比它大一个数量级。假设kernel A产出了一个1M元素的float32中间张量kernel B要消费它。那么A要把这4MB数据写进全局内存B要从全局内存读4MB回来总共8MB的流量。按GPU显存有效带宽50GB/s算这一来一回大约160微秒是启动开销的二十倍以上。如果把这个中间张量合并掉A和B变成一个kernel中间数据完全可以留在寄存器或片上内存里全局内存流量直接归零。所以Combine的收益大头在带宽不在启动次数。这也是我读这个Pass时最先想明白的一点它不是一个少调几次函数的优化而是一个省一次显存往返的优化。想通这一点后面读源码就不会被那些琐碎的IR操作带偏。中间结果大小全局内存往返耗时按50GB/s估算kernel启动开销合并后可节省64KB约2.6us约5-10us启动少量带宽1MB约40us约5-10us主要是带宽16MB约640us约5-10us带宽占绝对主导2. 搭框架从Pass注册到主循环2.1 文件位置与Pass基类以我读的主线版本为例Combine的实现文件在lib/Dialect/Triton/Transforms/Combine.cpp。注意不同版本路径可能不一样老版本可能在lib/Transforms/下面但文件名叫Combine.cpp这一点基本没变过。这个Pass注册名在不同的tablegen定义里可能略有差异常见的是triton-combine。它继承的是TritonCombineBaseCombine这是MLIR tablegen自动生成的一个CRTP基类。重点是它的runOnOperation()里拿到的operation类型是FuncOp也就是说这是一个Function级别的Pass不是Module级别的。这意味着什么它每次只处理一个tt.func对这个函数内部的tt.call做变换。跨函数的符号解析比如根据call的callee符号名字找到对应的FuncOp需要另外通过symbolTable来做。理解这个粒度很重要它决定了你能合并的边界你只能把某个函数调用的body搬进当前函数不能直接跨Module重排两个独立函数。2.2 主循环的三步走runOnOperation的逻辑大致可以概括成三步。第一步遍历当前函数里所有的tt.call把它们收集到一个std::vectorCallOp里。第二步遍历当前函数的ReturnOp从return的操作数开始往回找看这些返回值是不是某个call的结果这个call能不能被合并。第三步一旦合并成功标记changed并触发后续的清理逻辑。这里有个细节值得停下来想一想为什么从ReturnOp往回找而不是直接在收集到的所有call里逐个判断因为跨kernel的数据流本质上是返回值在流动。一个call的结果如果只是被存储到全局内存或者被计算使用那它并没有真正穿过函数边界只有它成为外层函数的return操作数或者成为另一个call的操作数时才发生了跨函数的数据传递。以return为起点刚好能对准这个语义。伪代码看起来是这样void runOnOperation() override { FuncOp funcOp getOperation(); std::vectorCallOp callOps; collectCall(funcOp, callOps); bool changed false; for (auto op : llvm::reverse(funcOp.getBody().getOps())) { auto ret dyn_castReturnOp(op); if (!ret) continue; for (auto operand : ret.getOperands()) { changed | tryCombineFrom(operand, callOps); } } if (changed) { // 清理无用的 call 和被孤立出来的 FuncOp } }逆序遍历getBody().getOps()是有讲究的。这样能保证先处理靠近return的调用把更下游的计算先合并进来然后随着IR的改写原本更靠前的call也会变成新的返回值来源可以在同一轮里继续合并。如果正序遍历容易出现迭代器失效的问题而且合并顺序会变成从producer推到consumer处理多个返回值时会很麻烦。2.3 collectCall在收集什么collectCall这个函数本身不复杂它就是遍历当前函数body里的所有tt.call把Operation*塞进一个vector。但它有一个隐含的作用把后续要修改的操作先快照下来避免在walk过程中修改IR导致迭代器悬空。实际处理的时候并不是每个收集到的call都值得尝试。我看到的代码里会逐个检查call的返回值use情况只有满足条件的才会进入combineImpl。这个筛选是核心一个call的结果必须恰好只被一个地方使用而且这个地方必须是另一个call的操作数或者是当前函数的return操作数。为什么要求恰好一个使用点因为合并的本质是把producer的body搬到consumer的调用位置。如果producer的结果被两个地方使用搬过去之后另一个使用点就拿不到值了。编译器不会自动帮你复制一份body因为那可能带来成倍的重复计算。所以这里的保守判断非常关键。3. combineImpl的合并策略三类场景的分支处理3.1 最基础的场景call结果喂给另一个call在我读到的源码里核心递归逻辑集中在combineImpl。它接收一个待处理的callOp和当前是否发生改变的引用返回是否成功。最典型的场景是这样一个链式调用%a是call producer的结果%b是call consumer(%a, ...)的结果而%b被最终return。这种情况下%a只有一个使用者就是那个consumer调用。于是combineImpl会把producer的body内联到consumer调用之前让%a从一个函数的返回值变成当前函数内部的一个中间值。内联之后原来的call producer就不需要了可以直接删除。而call consumer里对%a的引用也会自动指向内联进来的计算结果的最终值。如果consumer原本是一个tt.func它的body也要被继续内联进来——这就涉及到多层合并。combineImpl会递归处理直到某条数据流链上不再有call为止。这里还有一个判断维度callee是否只有一个调用点。如果producer被多个地方调用严格来说也可以做内联把body复制多份但Triton的Combine选择了一种更省事的策略只有当callee只有一个调用者或者内联后不会破坏其他调用点时才做合并。这样避免了对同一个函数做多份内存拷贝也避免了符号表引用混乱。3.2 分支普通call、launch_kernel、scf.forcombineImpl内部并不是一刀切而是会根据callOp所在的上下文走不同分支。我梳理下来主要有三类这也是读源码时最容易绕晕的地方。第一类是普通call也就是callOp直接位于FuncOp的body里结果被另一个call或return使用。这种最简单直接复制callee body做一下block argument的映射替换然后把call删掉。第二类涉及launch_kernel这种包装调用。Triton里有些kernel启动不是一个普通的call而是一个更外层的包装操作它除了被调用的函数名还携带grid维数、stream等启动参数。合并这种调用时不光要把callee body搬进来还要处理启动参数和kernel属性的重写。如果只是一个普通内联启动参数会丢失后面发射阶段就找不着grid配置了。第三类是scf.for循环体里的调用。如果callOp出现在循环体的region中它对应的返回值可能被scf.yield当作循环结果带出循环。此时合并的复杂度明显上升要把callee的body插入循环体末尾还要处理迭代变量和循环携带的iter_args的映射关系。我之前一度觉得这跟普通内联没区别后来发现block argument的处理很容易出错因为循环体用的值未必来自函数形参可能来自上一轮迭代的yield值。3.3 合并顺序的讲究从下游往上游推读这段代码的时候我建议特别关注合并顺序。源码里的一个特点是它倾向于从下游往上游合并也就是先处理更靠近return的调用。举个例子一个调用链是call A的结果喂给call Bcall B的结果再喂给call C最终return C的结果。正确做法是先尝试合并B和C把B的body并入C得到一个合并后的BC函数接着再把A并入BC。这相当于沿着数据流反向回溯。如果反过来从A开始合并把A并入B得到的AB再并入C也不是不行但问题在于A的结果可能被多个下游消费或者A的返回值个数和B的参数列表对不上回溯时需要维护的状态就很多。从下游往上游推每步只需要盯住当前返回值对应哪个call状态机更简单。这也是一个通用经验读这类数据流驱动的Pass先搞清楚它的遍历方向再去看具体匹配逻辑思路会顺很多。4. 合并落地时的细节参数映射、属性与清理4.1 参数映射要用ValueMap合并不是把callee的body文字复制一遍就完事。callee里有block argument也就是函数形参这些形参在callee内部是被replaceAllUsesWith替换成调用点的实际参数的。也就是说把callee body搬过来之前要先建一个ValueToValueMap建立callee形参 - caller实参的映射然后对callee body里的所有operation做remap。有一个容易忽略的点如果callee的某个参数在调用点传入的是一个常量比如arith.constant那它就不需要被提升为合并后函数的外部参数。这个常量可以直接保留在body里。所以在做参数映射的时候源码里往往会有一步匹配常量参数的逻辑把它们从形参变成内联常量。这个细节在普通函数内联里也有但Triton Combine这里特别重要因为kernel的参数列表直接对应host端传入的buffer指针和标量值。多一个参数就多一份host侧的配置能少则少。4.2 新函数参数顺序为什么重要合并两个函数之后按道理会生成一个新的tt.func它的参数列表是producer的参数和consumer的参数的并集。但参数顺序不能随便排因为对这个新函数的外部调用者来说参数顺序就是ABI的一部分排错了host端传参就全乱了。我看到的实现逻辑里通常会把producer侧需要的参数排列在前consumer侧原有的参数排列在后并且在合并过程中使用一个统一的参数收集列表来保证不遗漏、不重复。如果producer和consumer都依赖同一个指针参数需要去重否则同一个buffer会变成两个形参后面lowering会很别扭。这个排序逻辑不算复杂但很容易写错。我自己曾经在改一个类似Pass时因为参数去重没做好导致合并后在另一个模块里出现两个同名buffer被当成不同参数的问题调试了快一下午。后来才意识到合并类Pass里的参数整理一定要单独拎出来做一轮显式的处理不能随手塞在某个分支里。4.3 死代码与符号表清理合并完成的瞬间IR里会留下不少垃圾被替换掉的callOp、失去所有调用者的callee函数。所以Combine在合并后会触发一个清理阶段把无用的call删掉并检查符号表里某个tt.func是否还有调用者。没有调用者的函数且它不是入口函数就一并删掉。这里有个陷阱删除函数时要小心入口标记。Triton里某些函数会被标记为kernel entry point它们即使没有调用者也不能删因为编译产物要从它开始发射。通常通过检查函数的可见性或某种entry属性来判断。源码里如果不做这一步很容易把入口函数当dead code清掉导致编译产物为空。清理完callOp和孤立函数之后还得留意函数属性。比如callee body里可能带有与kernel发射相关的属性这些属性在合并后需要被复制到新的函数上否则后面某个Pass会找不到它们而报错。这一块通常比较琐碎但决定了合并后的IR能否顺利进入下一轮优化。5. 哪些情况Combine不会合并保守也是优化策略5.1 多使用者场景直接放弃前面提到过中间结果如果被两个地方使用Combine不会合并。这是第一个硬边界。比如%a call producer(...)之后%a既被call consumer(%a)使用又被一个tt.store写进显存。这种情况下producer的body没办法只在一个位置展开所以放弃合并。有人可能会问为什么不把consumer也复制一份让两个使用点各自内联一份producer从编译器角度如果producer很廉价重复计算可能划算但Triton的Combine不赌这个。它的策略是只做最安全的合并把收益不确定的场景留给后面的Pass处理。这种宁可少做不做错的思路在大型编译器里其实是值得学习的盲目的激进优化往往带来的是可维护性灾难。5.2 调用图里的递归问题递归调用是每个函数内联类Pass都要处理的。假设函数A里调用了BB的body里又调用了A这种循环调用如果不加限制合并过程会无限递归下去直到栈溢出。源码里规避这个问题的常用做法是维护一个已经访问过的函数集合或者在合并递归时设置深度上限。一旦发现当前正在处理的函数已经在合并栈里就直接放弃合并。这种循环调用从语义上说要实现真正的kernel融合也很麻烦而且GPU上的device函数递归本来就不是一个被鼓励的写法所以放弃并不可惜。5.3 别名与全局内存依赖还有一个更隐蔽的边界涉及别名分析。假设producer从一个指针参数%pload数据计算结果通过return传出consumer对另一个指针参数%q做了store之后再使用producer的结果。如果%p和%q实际上指向同一块显存合并前后语义可能就变了。本来producer和consumer是两个独立kernelkernel A执行完所有读写kernel B才开始执行天然有一个全局同步。合并之后变成同一个kernelproducer的load和consumer的store之间的顺序可能被重新安排甚至被一起向量化这就可能在别名存在时踩到数据竞争。Triton的Combine没有引入专门的别名分析模块所以遇到这种潜在冲突时它的选择是保持原样。这个边界在源码里不一定有显式的检查更多是通过只合并return数据流这种保守结构天然避开的。5.4 与后续layout优化的交互最后还有一个大局层面的原因让Combine保持克制它工作在TTIR层此时没有线程块内layout信息。如果过早把两个函数暴力合并成一个后面TritonGPU层的layout选择、共享内存分配就失去了自由发挥的空间。比如producer可能更适合一个按行划分的layoutconsumer可能更适合按列划分的layout。如果它们分开编译器可以分别给它们最优的layout中间通过全局内存衔接。而一旦在TTIR层强行合并编译器只能给整体选一个折中的layout可能反而更慢。所以Combine只处理那些在数据流上紧密耦合、明显应该合在一起的调用。真正复杂的融合优化在Triton里是交给后续更专门的Pass去做的。想明白这一点就不会抱怨Combine太弱了。6. 自己动手验证Combine复现步骤与调试心得6.1 用test用例学习期望行为读源码最快的方式其实是先看测试。Triton源码树里通常有test/Dialect/Triton/combine.mlir之类的文件里面直接写了输入IR和期望输出IR的对照。拿到这个文件先跑一遍看Pass到底把哪些call合并了、哪些没合并比看代码快得多。本地验证方式也不复杂。把测试文件里的输入部分提取出来保存成input.mlir然后运行mlir-opt --triton-combine input.mlir不同版本的Pass注册名可能不同如果--triton-combine提示找不到就先跑mlir-opt --help搜一下combine相关的flag。输出里如果tt.call消失、函数数量减少说明Pass生效了。6.2 一个最小复现用例我自己调试时用的一个最小IR大概是下面这个样子。它构造了一个producer函数返回一个从指针load出来的值一个consumer函数把producer的返回值和一个常量相加后返回。正常情况下Combine应该能把producer的body内联进consumertt.func public producer(%p: !tt.ptrf32) - f32 { %0 tt.load %p : !tt.ptrf32 tt.return %0 : f32 } tt.func public consumer(%p: !tt.ptrf32) - f32 { %a tt.call producer(%p) : (!tt.ptrf32) - f32 %c arith.addf %a, %a : f32 tt.return %c : f32 }注意这是一个示意性的IR真实Triton的tt.load返回类型可能带张量shape但不影响理解。跑完mlir-opt之后consumerbody里应该能看到tt.load直接出现tt.call producer消失producer函数如果不再有调用者也会被删除。这个before/after的差异足够判断行为是否符合预期。6.3 三个实用的调试手段第一用--mlir-print-ir-after-all跑Pass管线把每一步的IR都打印出来。这是最直观的方式。你能看到Combine之后IR变成什么样也能看到它后面紧接着的Pass是否有异常。第二在源码里临时加日志。combineImpl的入口处加一行llvm::dbgs() combine candidate: op-getName() \n;重新编译后跑测试能精确看到哪些操作被当成了候选。这种临时的插桩调试比在IR层面猜行为高效得多。第三结合Python端验证。如果把mlir-opt单独跑通还不够可以用triton.compile编译一个拆成两个triton.jit函数的例子通过设置Triton的环境变量输出中间MLIR再观察最终的编译产物。我自己的体会是调试这种Pass最忌讳一上来就追大IR。先把最小case跑通再用mlir-opt一行命令断开前后端的依赖效率会高很多。源码里那些一眼看不懂的函数配合before/after的IR对比基本都能猜出七八分意图再去代码里找实现细节就很快了。读Combine这个Pass最大的收获不是记住了它处理了哪几种调用场景而是理解了一个边界清晰的Pass如何在值得合并和不能合并之间做取舍。这种逻辑比单纯堆优化技巧更值得花时间琢磨。后续我还会沿着这条线继续往下看下一篇准备看看TTIR到TritonGPU的layout dialect转换前后发生了什么。