TritonIR剖析：从 Python Kernel 到 GPU 指令，TTIR/TTGIR 全流程拆解

Triton IR 剖析从 Python Kernel 到 GPU 指令TTIR/TTGIR 全流程拆解从头梳理 Triton 编译器从 Python DSL 一路 lowering 到 LLVM IR 的完整链路。上半部分讲 kernel 怎么变成 TTIR、MLIR 语法长什么样下半部分拆四个关键优化 Pass——CSE、Canonicalizer、Inliner、RewriteTensorPointer看完你就能对着triton-opt的输出不再发怵。一、整体框架Triton 编译器到底在干什么Triton 是一个用 Python 写 GPU Kernel 的框架。你写一段类似 NumPy 的代码它帮你编译成能在 NVIDIA/AMD GPU 上跑的高性能机器码。这件事听起来简单中间其实藏了一条很长的 lowering 管线Triton DSL (Python) ↓ AST 解析 代码生成 TTIR (Triton IR) ← 高层张量语义硬件无关 ↓ Pass 优化 lowering TTGIR (Triton GPU IR) ← 引入线程/block/warp绑定GPU执行模型 ↓ lowering LLVM IR ← 通用底层IR ↓ NVPTX backend PTX / GPU Binary ← 最终可执行指令这条管线里TTIR 描述算什么TTGIR 描述怎么在 GPU 上算LLVM IR 负责最终怎么执行。三层 IR 各司其职背后依托的正是 MLIR 这套可扩展编译器基础设施。二、上篇从 Triton Kernel 到 TTIR2.1 Kernel 长什么样 → TTIR 怎么来以一个最简的vector_add为例triton.jitdefadd_kernel(x_ptr,y_ptr,output_ptr,n_elements,BLOCK_SIZE:tl.constexpr):pidtl.program_id(axis0)block_startpid*BLOCK_SIZE offsetsblock_starttl.arange(0,BLOCK_SIZE)maskoffsetsn_elements xtl.load(x_ptroffsets,maskmask)ytl.load(y_ptroffsets,maskmask)outputxy tl.store(output_ptroffsets,output,maskmask)这段 Python 代码经过 Triton 前端处理后会生成对应的 TTIR。TTIR 是一种 MLIR Dialect里面的 op 直接对应你写的tl.load、tl.store、tl.arange等操作但完全不涉及线程模型、shared memory、warp 调度等 GPU 细节。它只关心张量级别的语义从哪里加载、做什么运算、写到哪里去。2.2 Triton 源码里的关键模块Triton 的 Python 前端不是黑盒。几个核心模块的职责模块作用triton/language/core.pyTriton 核心操作实现load/store/dot/reshape 等triton/language/semantic.py语法层负责把 Python AST 翻译成内部 IRtriton/language/standard.py常用算子的标准实现triton/language/math.py数学运算exp/sin/cos 等triton/language/random.py随机数生成triton/language/extra/cuda/libdevice.pyCUDA 硬件相关的底层操作如果你要 debug “为什么我的 kernel 生成的 IR 不对”大概率要从semantic.py和core.py入手追溯。2.3 Kernel 与 Triton IR 的映射关系这一页 PPT 展示了 Python kernel 代码与生成的 TTIR 之间的对应关系。核心映射规律tl.load→tt.loadtl.store→tt.storetl.arange→tt.make_rangetl.program_id→tt.get_program_id算术运算 →arith.addf/arith.mulf等 MLIR 标准 Dialect这里有个关键认知TTIR 不是 Triton 自己从零造的而是基于 MLIR 的 Dialect 机制扩展出来的。算术运算用的是 MLIR 内置的arithDialect内存操作用的是ttTritonDialect。这种混搭正是 MLIR 设计的精髓。三、插播MLIR 是什么为什么 Triton 要用它3.1 MLIR 的核心概念MLIRMulti-Level Intermediate Representation是 LLVM 项目下的可扩展编译器框架。它不像传统编译器只有一种 IR而是允许你定义多种方言Dialect每种方言有自己的一套操作Operation、类型Type和属性Attribute然后通过 Pass 在不同方言之间逐层 lowering。三个核心概念一句话DialectIR 的方言比如arith算术、scf结构化控制流、ttTriton 自定义OperationopIR 的基本计算单元类似汇编里的指令Pass对 IR 做变换的模块可以优化、重写、lowering3.2 为什么 GPU 编译器非 MLIR 不可GPU Kernel 编译涉及张量计算、线程/warp/block 调度、shared memory 分配、向量化、bank conflict 消除……这些概念横跨多个抽象层级。传统做法是为每一层手写一套 IR结果就是 IR 之间信息割裂Pass 复用困难。MLIR 的方案是所有层级共用同一套 IR 基础设施通过 Dialect 表达不同层级的语义通过 Pass 在不同 Dialect 之间渐进式 lowering。Triton 的 Dialect 栈长这样Dialect层级作用tt(TTIR)高层张量语义硬件无关ttg(TTGIR)中层GPU 线程/block/warp/shared memoryarith/math通用算术与数学运算scf/cf通用控制流llvm底层对接 LLVM IRnvvm底层NVIDIA PTX 专项这套分层设计让每个 Pass 只需关注自己层级的优化不用操心跨层问题。四、下篇Triton IR 中的优化 Pass下篇聚焦编译器真正干活的部分——优化 Pass。Triton 在 TTIR 阶段会跑一系列 Pass把 naive 的初始 IR 逐步转化成更适合 GPU 执行的形式。4.1 Pass Pipeline 总览NVIDIA GPU 对应的make_ttir优化 Pass pipeline 定义在triton/backends/nvidia/compiler.py中由python/triton/compiler/compiler.py调用。整个 pipeline 的典型执行顺序TTIR 初始 IR ↓ Canonicalizer ← 先把 IR 规范化 ↓ CSE ← 消除重复计算 ↓ Inliner ← 展开所有函数调用 ↓ RewriteTensorPointer ← 拆解张量指针 ↓ TTGIR (后续还有更多 Pass...)下面逐个拆解。4.2 CSE公共子表达式消除CSECommon Subexpression Elimination做的事情非常朴素如果同一个表达式被计算了多次只算一次后面直接复用结果。比如 Triton kernel 里经常出现大量索引计算%off0 arith.addi %base, %stride %off1 arith.addi %base, %stride // 和上一行完全一样CSE 会把%off1直接替换成%off0后续所有引用%off1的地方都改成引用%off0。在 GPU kernel 场景下CSE 的价值特别大。因为 GPU 的寄存器数量是硬限制重复计算不仅浪费 ALU 指令还会增加寄存器压力间接导致 occupancy 下降。消除一次重复计算可能就多塞一个 warp。Triton 使用的是 MLIR 内置的 CSE Pass核心逻辑是基于 op 类型 操作数 属性的哈希匹配。但 Triton 做了定制化处理tt.load、tt.store、tt.reduce这些有副作用的 op 不会被 CSE防止把读两次内存错误优化成读一次。执行命令triton-opt cse_before.ttir-csecse_after.ttir4.3 Canonicalizer规范化 PassCanonicalizer的目标是把 IR 转成一种规范形式。它不是一个单独的大规则而是大量小规则的集合——常量折叠、代数化简、冗余 op 消除等等x 0→xx * 1→x2 * 3→6空的scf.if块直接移除Canonicalizer 通常放在 CSE 之前执行。原因很简单规范化的 IR 让更多表达式长得一样CSE 就能命中更多匹配。Triton 的自定义 op如tt.expand_dims、tt.broadcast等也都注册了自己的 canonicalization pattern确保 Triton 特有的 IR 结构也能被规范化。执行命令triton-opt canonicalizer_before.ttir-canonicalizecanonicalizer_after.ttir4.4 Inliner内联 PassInliner直接把函数调用替换成函数体代码。这是 GPU 编译里必须做的一步不是什么可选优化——因为 GPU 硬件根本不支持通用 function call。所有逻辑必须在同一个 kernel 里展开否则 NVPTX backend 根本没法生成合法代码。Triton 会生成很多辅助函数数学运算、向量化操作等Inliner 把它们全部展开到 kernel 主函数里。展开之后通常会再跑一轮 Canonicalizer CSE因为内联后经常暴露出新的优化机会。MLIR 的 Inliner 还有一个成本模型如果函数体太大或者调用次数太多它会根据启发式策略决定是否内联。但在 Triton 的 GPU kernel 场景下基本是全量内联——没办法硬件不答应。执行命令triton-opt inliner_before.ttir-inlineinliner_after.ttir4.5 RewriteTensorPointer张量指针重写这是 Triton最核心、最特有的 Pass 之一。Triton 引入了一个叫 “Tensor Pointer” 的概念——tt.ptrtensor128xf32。它不是传统的指针而是一个携带了形状、布局、swizzle 等元信息的智能指针。在早期 IR 中一个tt.load配合 Tensor Pointer 就能表达复杂的向量化加载非常简洁。但问题来了LLVM IR 不认识什么 Tensor Pointer。它只认最原始的i8*加上偏移量。所以 RewriteTensorPointer 的任务就是把高级的 Tensor Pointer 拆解为基地址 偏移计算 向量化 load/store 内存 layout 变换。举个例子Before%ptr tt.make_tensor_ptr %base, [%stride], [%shape] %val tt.load %ptr : tensor128xf32After RewriteTensorPointer%ptr被拆成 base address block offset swizzle offset生成多个llvm.load或 PTX 向量化加载指令如ld.global.v4.f32可能会引入 shared memory promotion 和 coalescing 优化这个 Pass 之所以叫 “Rewrite” 而不是 “Lower”是因为它仍然运行在 TTIR/TTGIR 层面产出的是语义等价但更低级的 Tensor Pointer 表示还没到 LLVM Dialect。执行命令triton-opt rewrite_before.ttir -triton-rewrite-tensor-pointerrewrite_after.ttir4.6 四个 Pass 总结Pass干啥为什么重要Canonicalizer规范化 IR常量折叠 代数化简让后续 Pass 更容易命中模式CSE消除重复计算省寄存器、省 ALU提升 occupancyInliner展开所有函数调用GPU 不支持 call必须内联RewriteTensorPointer拆解张量指针为底层地址计算连接高层语义与 LLVM IR 的关键桥梁五、学完能做什么搞懂这套流程之后你至少能做三件事Debug IR 输出TRITON_PRINT_IR1打印各阶段 IR 后你能看懂每一层在干什么定位问题是出在 TTIR 生成还是某个 Pass 优化手写 Triton Pass如果要加自定义优化比如针对某种 memory layout 做特殊处理你知道应该在哪个 Dialect 层面介入跨框架对比理解了 Triton 的 lowering 管线再看 TVM、XLA、IREE 的 IR 栈会发现思路一脉相承——都是 Dialect Pass 的分层 lowering如果你想把某个环节彻底吃透建议直接去读triton/backends/nvidia/compiler.py里的make_ttir函数它是这条管线的总控台所有 Pass 的注册和顺序一目了然。本文基于董贞汝老师《Triton-IR剖析》分享内容整理融入 MLIR 基础概念与各 Pass 原理补充截图均来自原始 PPT。内容由AI生成仅供参考内容由AI生成仅供参考