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

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
    • Target：将 llvm dialect 降级到 llvm ir，为 llvm ir 添加元数据，链接外部数学库
      • LLVMIR
    • Tools：辅助工具头文件，分析、调试、优化生成的代码
      CMakeLists.txt

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

    • triton
      • Analysis
        • Alias.h
        • …
      • Conversion
        • TritonGPUToLLVM
        • TritonToTritonGPU
      • Dialect
        • triton
        • tritonGPU
        • TritonNvidiaGPU
      • Target
        • LLVMIR
      • Tools
    • CMakeLists.txt

    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