Triton源码目录:打开Triton源码的正确姿势:从一头雾水到心里有数

📅 2026/7/5 7:38:28
Triton源码目录:打开Triton源码的正确姿势:从一头雾水到心里有数
你有没有过这种经历——听说Triton写GPU Kernel很香兴冲冲git clone下来然后打开文件夹就懵了include/、lib/、Dialect/、Conversion/……一堆目录扑面而来根本不知道从哪下手。反正我有。第一次看Triton源码的时候光目录结构就研究了半天最后硬着头皮一个个文件点开看效率低得令人发指。后来跟着刘灿老师的分享系统地梳理了一遍才算把Triton的源码骨架理清楚。这篇文章就是我当时的学习笔记希望能帮你少走点弯路。本文基于 triton-lang/triton 仓库https://github.com/triton-lang/triton以 PPT 内容为主线整理。先看全景四大核心目录打开Triton仓库你最先需要关注的是这四个目录别被其他目录分散注意力include/、lib/、python/、runtime/这四个才是主干。其中include/和lib/又是理解编译流程的重中之重——它们承载了Triton基于MLIR构建的整个编译管线。include 和 libMLIR世界的C头文件和实现include/ 里有什么Triton在MLIR框架上定义了两个核心DialectTriton Dialect高层、设备无关的IR。你写的tt.load、tt.store、tt.dot最终都会落到这个Dialect里。TritonGPU DialectGPU相关的抽象引入了Layout概念Blocked、MMA、Slice、Shared等决定了数据在GPU线程和内存中的分布方式。include/下面主要就是这两个Dialect的Op定义TableGen .td文件、类型定义和属性定义。lib/ 里有什么lib/是include/的对应实现包括子目录作用lib/Dialect/Triton/Triton Dialect的IR实现 优化Passlib/Dialect/TritonGPU/TritonGPU Dialect的IR实现 大量优化Passlib/Conversion/Dialect之间的转换降级逻辑lib/Target/最终代码生成LLVM IR → PTX / ROCm简单粗暴地记include定义是什么lib实现怎么做。Dialect 目录Triton编译器的灵魂理解Dialect是理解Triton源码的关键。Triton的编译流程本质上就是Dialect的逐级降级Triton Dialect ──→ TritonGPU Dialect ──→ LLVM Dialect ──→ 机器码 (高层语义) (GPU映射) (通用IR) (PTX/ROCm)为什么要分两个Dialect这是MLIR框架的精髓——分层抽象Triton Dialect完全不关心GPU硬件细节只描述做什么计算load、store、dot、reduce……TritonGPU Dialect负责把计算映射到具体的GPU执行单元上——数据怎么分配到warp、怎么用Tensor Core、shared memory怎么管理这种设计让前端Python DSL和后端各种GPU可以独立演进互不干扰。编译流程对应源码位置一张图搞定这是整份PPT里信息密度最高的一页把Triton的编译流程和源码目录对应起来就长这样编译阶段源码位置做了什么Triton IRlib/Dialect/Triton/定义高层Op语义执行通用优化Triton → TritonGPUlib/Conversion/TritonToTritonGPU/引入GPU Layout分配warpTritonGPU 优化lib/Dialect/TritonGPU/coalesce、pipeline、prefetch等大量GPU优化TritonGPU → LLVMlib/Conversion/TritonGPUToLLVM/MMA lowering、shared memory处理LLVM → 目标代码lib/Target/LLVMIR/生成LLVM IR后面交给LLVM后端生成PTX如果你想跟踪一个triton.dot操作从Python一路变成PTX指令的全过程顺着这个表找对应文件就行。常用文件速查手册Triton Dialect 核心文件Op定义/include/triton/Dialect/Triton/IR/TritonOps.td这里定义了你在Python里调用的所有基础操作tt.call、tt.func、tt.return—— 函数调用机制tt.load、tt.store—— 全局内存读写tt.dot—— 矩阵乘法Triton的招牌操作优化Pass/lib/Dialect/Triton/Transforms/包含CombineOpsPass、ReorderBroadcastPass、RewriteTensorPointerPass、LoopUnrollPass等。这些Pass在GPU无关的层面做图优化——算子融合、死代码消除、循环展开跟传统编译器优化的套路差不多。Triton → TritonGPU 转换lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp这是整个编译流程的第一个硬核跳变。在这里设备无关的Triton IR被注入GPU-specific的Layout信息——数据开始被分配到具体的线程和内存层级。TritonGPU Dialect 核心文件Op定义/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.tdGPU特有的操作async_wait、alloc_tensor、insert_slice_async、convert_layout。Layout属性定义/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td这是TritonGPU最核心的设计——五种LayoutLayout含义Blocked数据按块分配给线程最通用MMA为Tensor Core矩阵乘累加定制DotOperanddot操作的输入数据布局Slice数据切片后的布局Sharedshared memory中的数据布局优化Pass/lib/Dialect/TritonGPU/Transforms/这里有一大堆GPU专用Pass每个都值得单独写一篇文章AccelerateMatmul—— 把普通矩阵乘匹配到Tensor Core指令Pipeline/Prefetch—— 指令流水线和数据预取Triton性能的核心来源Coalesce—— 合并内存访问减少bank conflictRemoveLayoutConversions—— 消除不必要的layout转换开销OptimizeThreadLocality—— 优化线程局部性AllocateSharedMemoryPass—— shared memory分配补充说明Pipeline和Prefetch是Triton自动生成异步拷贝流水线的关键实现。它们会在循环内插入cp_async指令让数据搬运和计算重叠执行——这正是Triton Kernel经常能追平甚至超过手写CUDA的原因。一句话总结Triton的源码结构本质上就是一条MLIR编译管线include/ (定义) → lib/Dialect/ (IR实现) → lib/Conversion/ (Dialect降级) → lib/Target/ (代码生成)如果你也想入手Triton源码我的建议是先看include/triton/Dialect/Triton/IR/TritonOps.td搞清楚Triton有哪些基础操作跟踪triton.dot从 Triton Dialect → TritonGPU → LLVM 的完整降级路径用mlir-opt --print-ir-before-all实际跑一遍 Pass Pipeline直观感受每个Pass的效果源码链接https://github.com/triton-lang/triton本文内容整理自先进编译实验室刘灿的分享感谢。