1 triton 3.0.0 源码结构
- triton
-
docs/:项目文档
-
cmake/:构建配置相关
-
bin/:工具、脚本
-
CmakeLists.txt:cmake 配置文件
-
LSCENSE
-
README.md
-
Pyproject.toml:python 项目配置文件
-
utils/:项目配置文件目录
-
unittest/:单元测试代码
-
third_party/:第三方资源
- amd/
- f2reduce/
-
test/:测试代码
-
python/:python 接口代码
-
lib/:核心逻辑实现,
.cc/.cpp
,核心功能的具体实现- Analysis:相关分析
- Alias.cpp:内存别名分析
- Allocation.cpp:共享内存分配相关分析
- Axisinfo.cpp:轴分析相关
- Membar.cpp:线程同步、内存屏障相关
- Conversion:dialect 之间的转换
- TritonGPUToLLVM:tritonGPU dialect 降级到 LLVM dialect
- TritonToTritonGPU:triton dialect 降级到 tritonGPU dialect
- Dialect:各级中间表示 dialect 的定义,以及在对应 dialect 上进行的优化 pass
- triton
- IR:dialect/算子/属性/类型的定义
- Transforms:相应中间表示上的优化 pass
- Combine.cpp:优化 select 和 load 操作的组合
- ReorderBroadcast.cpp:将通过 broad 和乘法生成的规约操作优化为点积操作
- tritonGPU
- TritonNvidiaGPU
- triton
- Target:将 llvm dialect 降级到 llvm ir,为 llvm ir 添加元数据,链接外部数学库
- LLVMIR
- Tools:辅助工具头文件,分析、调试、优化生成的代码
CMakeLists.txt
- Analysis:相关分析
-
include/:核心逻辑定义,核心功能的
.h
头文件,提供约定和规范- triton
- Analysis
- Alias.h
- …
- Conversion
- TritonGPUToLLVM
- TritonToTritonGPU
- Dialect
- triton
- tritonGPU
- TritonNvidiaGPU
- Target
- LLVMIR
- Tools
- Analysis
- CMakeLists.txt
CMakeLists.txt
- triton
-
注:transforms 用于各级 dialect 之上,conversion 用于各级 dialect 之间
编译流程对应源码位置:
2 常用源码位置
Triton dialect:
- OP 定义:include/triton/Dialect/IR/TritonOps.td
tt.call, tt.func, tt.return, tt.load, tt.store, tt.dot 等 OP - 优化 pass:lib/Dialect/Triton/Transforms
CombineOpsPass, ReorderBroadcastPass, RewriterTensorPointerPass, LoopUnrollPass
triton dialect --> tritonGPU Dialect:
转换 pass:lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
TritonGPU Dialect:
- Op 定义:include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
async_wait, alloc_tensor, insert_slice_async, convert_Layout 等 OP - Layout 属性定义:include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
Blocked Layout, MMA Layout, DotOperand Layout, Slice Layout, Shared Layout - 优化 pass:lib/Dialect/TritonGPU/Transforms
AccelerateMatmul, Coalesce, CombineTensorSelectAndIf