Triton(2)——Triton源码接结构

1 triton 3.0.0 源码结构

• triton
• triton
• CMakeLists.txt
• LLVMIR
• triton
• tritonGPU
• TritonNvidiaGPU
• TritonGPUToLLVM
• TritonToTritonGPU
• Alias.h
• …
• Analysis
• Conversion
• Dialect
• Target
• Tools
• Analysis：相关分析
• Conversion：dialect 之间的转换
• Dialect：各级中间表示 dialect 的定义，以及在对应 dialect 上进行的优化 pass
• Target：将 llvm dialect 降级到 llvm ir，为 llvm ir 添加元数据，链接外部数学库
• Tools：辅助工具头文件，分析、调试、优化生成的代码 CMakeLists.txt
• Alias.cpp：内存别名分析
• Allocation.cpp：共享内存分配相关分析
• Axisinfo.cpp：轴分析相关
• Membar.cpp：线程同步、内存屏障相关
• TritonGPUToLLVM：tritonGPU dialect 降级到 LLVM dialect
• TritonToTritonGPU：triton dialect 降级到 tritonGPU dialect
• IR：dialect/算子/属性/类型的定义
• Transforms：相应中间表示上的优化 pass
• Combine.cpp：优化 select 和 load 操作的组合
• ReorderBroadcast.cpp：将通过 broad 和乘法生成的规约操作优化为点积操作
• triton
• tritonGPU
• TritonNvidiaGPU
• LLVMIR
• amd/
• f2reduce/

• docs/：项目文档

• cmake/：构建配置相关

• bin/：工具、脚本

• CmakeLists.txt：cmake 配置文件

• LSCENSE

• README.md

• Pyproject.toml：python 项目配置文件

• utils/：项目配置文件目录

• unittest/：单元测试代码

• third_party/：第三方资源

• test/：测试代码

• python/：python 接口代码

• lib/：核心逻辑实现，.cc/.cpp，核心功能的具体实现

• include/：核心逻辑定义，核心功能的.h头文件，提供约定和规范

  CMakeLists.txt

注：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