CUDA-MODE 课程笔记第 29 课 Triton 内部机制

我的课程笔记，欢迎关注：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/cuda-mode 。
这节课详细介绍了 Triton 编译器的内部工作原理。文章首先介绍了 CUDA 编译器（NVCC）的工作流程，然后深入探讨了 Triton 编译器的架构设计，包括其如何将 Python DSL 代码通过多个中间表示（IR）最终编译成 CUDA 可执行文件。课程重点讲解了 Triton 与 MLIR（Multi-Level Intermediate Representation）的关系，展示了各种优化 Pass 的实现和作用，如 TritonCombineOps、TritonReorderBroadcast 等。通过一个 vector_add 的具体示例，详细展示了代码从 Python 层面到 GPU IR 的转换过程，以及在这个过程中各种中间产物（如 .cubin、.ptx 等）的生成。最后还介绍了如何在 Triton 中实现新的编译器 Pass，为读者提供了深入理解和扩展 Triton 编译器的实践指导。

第 29 课，Triton Internals

课程笔记

这是一场由 Meta 软件工程师 Kapil Sharma 主讲的技术分享，主题是 Triton 编译器的内部工作原理。演讲者目前在 Meta 的 RecSys/Ranking 基础设施团队工作，他在 Slide 中分享了自己的社交媒体和代码仓库链接，包括 LinkedIn、Twitter 和 GitHub。

这张 Slide 介绍了关于 Triton 的演讲概要。Triton 是一个复杂的编译器/代码生成机制，之前已经有过一些相关演讲（包括 Triton 101、Kernel 融合和 Liger kernel 等）。研究人员普遍喜欢使用 Triton，它既可用于研究也可用于生产环境。Triton 的工作流程是从 PyTorch 开始，通过 torch.compile 编译成 Triton 内核，最后部署到目标硬件上。相关内容可以在三个系列博客文章中找到更详细的信息。这是一个将 PyTorch 代码优化并编译到 GPU 上运行的重要工具。

三个系列博客的链接：

这是本次演讲的目录。主要内容包括 CUDA 编译、Triton 编译器、示例代码、JIT 编译、Triton 和 MLIR(机器学习中间表示)的关系、IR(中间表示)的详细介绍、MLIR Pass 示例等主题。如果时间允许，还会介绍一个新的编译器 pass。

这张 Slide 展示了 NVIDIA CUDA 编译器(NVCC)的完整工作流程。NVCC 的主要作用是将 CUDA 代码分离成主机代码（C/C++）和设备代码（CUDA 内核）。主机代码通过标准的 C/C++编译器（如 g++或 clang）进行编译，而设备代码则被编译成 PTX 或 cubin 格式。整个流程通过一系列的中间步骤，包括预处理、编译和链接，最终生成可执行文件。图中清晰地展示了从.cu 源文件到最终可执行文件的完整转换过程。

这张 Slide 来自 NVIDIA 的演讲，详细说明了 NVCC 编译器的内部结构和组件。NVCC 是 CUDA 的主编译器，而 CICC 是基于 LLVM 的高级优化器和 PTX 生成器。PTX（Parallel Thread Execution）是一种虚拟指令集。图示展示了从.cu 文件开始，经过 CUDA C++前端处理，然后通过 CICC 生成 PTX 汇编代码，最后通过 PTXAS 和主机编译器生成最终的 CUDA 可执行文件的完整流程。

这张 Slide 重点介绍了 NVCC 编译器处理设备代码的两个关键阶段。第一阶段使用 C++预处理器和 CICC 程序生成中间形式，第二阶段 CICC 程序优化代码并生成 PTX。生成的 PTX 代码随后传递给 ptxas，用于生成 SASS（实际的 GPU 机器代码）。整个流程图清晰地展示了从源代码到最终 GPU 可执行代码的转换过程，并提到了 Godbolt 示例作为参考，我们也可以参考 Compiler Explorer，这两个网站都是很好的在线 CUDA 编译器可视化工具。

这张 Slide 是 OpenAI Triton 介绍博客里面的图，主要介绍了 Triton 编译器的工作流程。它展示了代码从 Python DSL(领域特定语言)到最终 CUDA 执行文件的转换过程。具体来说，Triton 编译器会将 DSL 代码经过多个阶段的编译，最终生成 CUBIN/fatbinary 格式的可执行文件。这个生成的 CUBIN 文件是一个 Slide 格式的文件，包含了可以在 CUDA kernel 中内联加载的可执行代码。Slide 通过展示了三个不同层次的代码(Python 层、Triton-IR 层和 PTX 层)来说明这个编译过程。

这张 Slide 展示了 Triton 编译器的架构设计图。从顶层看,它始于 Triton 语言,然后通过 Triton IR(中间表示)向下流转。架构分为几个主要分支:

左侧分支通过 SPIRV 处理,最终编译到 x86 和 Intel GPU 后端
中间分支经过 LLVM 和 AMD GPU 后端处理
右侧分支通过 LLVM 和 PTX 处理,最终到达 SASS(NVIDIA GPU 汇编)
最右侧还有一个单独的分支用于处理各类加速器(Accelerators)

这种架构设计允许 Triton 代码能够被编译到不同的硬件平台上运行,包括 CPU(x86)、各种 GPU(Intel/AMD/NVIDIA)以及其他加速器,体现了良好的跨平台兼容性。

这是 Triton 的官方教程的第一个例子 vector_add，然后我们可以使用 Triton 提供的编译工具来看一下生成可执行文件的过程中 dump 了哪些东西。

这里展示了一下 vector_add 在生成可执行文件过程中产生了哪些中间参数，包括编译和构建过程中生成的文件,如.cubin、.json、.llir、.ptx等，还有包含 cuda kernel 的最终源代码文件 add_kernel.9969bdda_0123.c (.c 文件)和对应的头文件 (.h 文件)。

我们可以看一下add_kernel.9969bdda_0123.c的内容，由于篇幅原因 unsigned char CUBIN_NAME[10960] 的内容只展示了一小部分：

/_ clang-format off _/  
//  包含必要的头文件  
#include <stdio.h>  
#include <stdint.h>  
#include <inttypes.h>  
#include <string.h>  
#include <cuda.h>

//  定义 CUDA 错误检查宏  
#define CUDA_CHECK(ans) {\  
    gpuAssert((ans), **FILE**, **LINE**);\  
  }\

// CUDA 错误检查辅助函数  
static inline void gpuAssert(CUresult code, const char *file, int line) {  
  if (code != CUDA_SUCCESS) {  
    const char *prefix = "Triton Error [CUDA]: ";  
    const char *str;  
    cuGetErrorString(code, &str);  //  获取错误字符串  
    char err[1024] = {0};  
    strcat(err, prefix);  //  拼接错误前缀  
    strcat(err, str);     //  拼接错误信息  
    printf("%s\\n", err); //  打印错误信息  
    exit(code);          //  退出程序  
  }  
}

//  全局变量定义  
#define CUBIN_NAME add_kernel_9969bdda_0123_cubin  
CUmodule add_kernel_9969bdda_0123_mod = NULL;    // CUDA 模块句柄  
CUfunction add_kernel_9969bdda_0123_func = NULL;  // CUDA 函数句柄  
// CUBIN 二进制数据,包含编译后的 CUDA 内核代码  
unsigned char CUBIN_NAME[10960] = { 0x7f, 0x45, 0x4c, 0x46, 0x02, 0x01, 0x01, 0x33, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0xbe, 0x00, 0x7c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xc0, 0x14.... };

//  卸载 CUDA 内核模块  
void unload_add_kernel_9969bdda_0123(void) {  
    CUDA_CHECK(cuModuleUnload(add_kernel_9969bdda_0123_mod));  
}

//  加载 CUDA 内核模块  
void load_add_kernel_9969bdda_0123() {  
    int dev = 0;  
    void *bin = (void *)&CUBIN_NAME;  
    int shared = 0;  
    //  加载 CUBIN 数据到 CUDA 模块  
    CUDA_CHECK(cuModuleLoadData(&add_kernel_9969bdda_0123_mod, bin));  
    //  获取 add_kernel 函数句柄  
    CUDA_CHECK(cuModuleGetFunction(&add_kernel_9969bdda_0123_func, add_kernel_9969bdda_0123_mod, "add_kernel"));

  
    //  配置共享内存  
    int shared_optin;  
    //  获取设备支持的最大共享内存大小  
    CUDA_CHECK(cuDeviceGetAttribute(&shared_optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, dev));  
    //  如果需要的共享内存大于默认值且设备支持,则设置更大的共享内存  
    if (shared > 49152 && shared_optin > 49152) {  
      CUDA_CHECK(cuFuncSetCacheConfig(add_kernel_9969bdda_0123_func, CU_FUNC_CACHE_PREFER_SHARED));  
      CUDA_CHECK(cuFuncSetAttribute(add_kernel_9969bdda_0123_func, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shared_optin))  
    }  
}

/_  
内核配置参数:  
['BLOCK_SIZE=64', 'num_warps=1', 'num_stages=3']  
_/  
// CUDA 内核启动函数  
CUresult add_kernel_9969bdda_0123(CUstream stream, CUdeviceptr x_ptr, CUdeviceptr y_ptr, CUdeviceptr output_ptr, int32_t n_elements) {  
    //  如果函数未加载则先加载  
    if (add_kernel_9969bdda_0123_func == NULL)  
       load_add_kernel_9969bdda_0123();  
    //  设置网格维度  
    unsigned int gX = 1024;  
    unsigned int gY = 1024;  
    unsigned int gZ = 1024;  
    //  准备内核参数  
    void *args[4] = { &x_ptr, &y_ptr, &output_ptr, &n_elements };  
    //  启动 CUDA 内核  
    if(gX * gY * gZ > 0)  
      return cuLaunchKernel(add_kernel_9969bdda_0123_func, gX, gY, gZ, 1 * 32, 1, 1, 0, stream, args, NULL);  
}

这里展示了add_kernel相关编译产物(Artifacts)。包括 Triton IR、Triton GPU IR、LLVM IR、PTX 和 CUBIN。同时作者还发现了一些有用的工具命令，比如使用readelf查看 elf 格式文件、用cuobjdump导出 sass 和 ptx 代码，以及使用nvidisasm工具让cubin文件变得可读。最后提到更多关于 Python 绑定的详细信息可以在博客系列的第 2 部分找到。这些内容主要是为了帮助开发者理解和分析 CUDA kernel 的编译过程和中间产物。

这张 Slides 展示了 Triton 另外一种即时编译的方式。代码示例展示了 vector_add 的 PyTorch CUDA kernel 的配置和执行过程：首先设定 size 参数，创建输入张量 x 和 y（都在 CUDA 设备上），定义计算网格（grid），创建输出张量，然后编译并执行 kernel。最后通过compiled_kernel.asm.keys()方法可以获取所有的代码生成键值，这些键值包括了各种中间表示形式（如'ttir'、'ttgir'、'ptx'、'cubin'等）。

这张 Slide 介绍了 JIT Compiled Kernel (即时编译内核)的工作原理，它展示了从 Python DSL 开始，经过 IR、PTX 到 CUBIN/fatbinary 和 launcher.so 的多层代码生成过程。系统会通过 TRITON_CACHE_DIR 和缓存管理器(https://github.com/triton-lang/triton/blob/4348109b0a8e1aac748aa9b1bbbcd858e9488940/python/triton/runtime/cache.py#L50-L71)将编译结果保存到磁盘，并以 fatbinary (cubin) 格式内联加载。目标硬件都有各自的驱动程序，这些驱动提供了被 Python 模块封装的原生代码，并负责将 cubin 加载到 CUDA 驱动中(https://github.com/triton-lang/triton/blob/main/third_party/nvidia/backend/driver.c#L389-L406)。此外，还导出了 cuda_utils.so(https://github.com/triton-lang/triton/blob/main/third_party/nvidia/backend/driver.py#L72-L86) 和 triton_launcher.so(https://github.com/triton-lang/triton/blob/main/third_party/nvidia/backend/driver.py#L413-L426) 两个共享库，并提供了 compile_module_from_src(https://github.com/triton-lang/triton/blob/main/third_party/nvidia/backend/driver.py#L48-L64) 和 triton.runtime.build(https://github.com/triton-lang/triton/blob/main/python/triton/runtime/build.py#L21-L80) 两个重要的代码关键点。

这张 Slide 主要介绍了 Triton 和 MLIR(Multi-Level Intermediate Representation)的关系和基本概念。MLIR 是在 2022 年完全重写的现代优化编译器基础设施，是 LLVM 生态系统的一部分。它包含中间表示(IR)规范和转换工具包，通过方言(dialects)机制来扩展 MLIR 框架(https://mlir.llvm.org/docs/Passes/)。TensorFlow是第一个使用MLIR的主要机器学习框架。在Triton中，所有方言都使用table-gen（一种DSL/代码生成工具）来处理MLIR的样板代码，并且可以通过设置MLIR_ENABLE_DUMP=1来dump每次编译过程中的IR信息。

tablegen 的优质资源推荐：https://www.jeremykun.com/2023/08/10/mlir-using-tablegen-for-passes/
MLIR 应用到深度学习框架的资源：https://www.youtube.com/watch?v=R5LLIj8EMxw

这里展示了 MLIR(Multi-Level Intermediate Representation)的基本使用示例。第一张展示了 MLIR 的"Hello World"示例代码，其中创建了一个 ModuleOp 并使用 PassManager 添加了两个优化 pass(CSEPass 和 DeadCodeEliminationPass)来处理模块。第二张则通过一个简单的 Python 函数优化示例，展示了如何使用 MLIR 进行代码优化 - 将一个包含多余计算的函数(a = b + c; e = b + c; d = e; return d)优化简化为更简洁的形式(a = b + c; d = a; return d)，体现了 MLIR 在代码优化方面的应用。

这张 Slide 展示了 MLIR 的常见 Pass(https://mlir.llvm.org/doxygen/namespacemlir.html)，我们可以将其应用到Triton中。

这张 slide 介绍了 Triton 编译器中的几个重要优化 Pass（编译优化阶段）：

TritonCombineOps：用于合并基础运算操作，比如点乘和地址计算
TritonReorderBroadcast：重排广播和分片操作的顺序，使其更高效
TritonRewriteTensorPointer：移除 tensor 指针相关的操作
TritonLoopUnroll：根据指定的因子展开循环结构

这些 Pass 都是 Triton 编译器用来优化 GPU 代码性能的重要步骤，通过这些优化可以生成更高效的 GPU 代码。可以参考这里的代码：https://github.com/triton-lan...

这里展示了 Triton GPU 的编译优化流程（Passes）。代码显示了如何使用 Triton 的编译器优化管道来处理 GPU 代码。具体包括了几个重要的优化 Pass，如线程本地化优化（optimize_thread_locality）、布局转换（layout_conversions）、矩阵乘法加速（accelerate_matmul）以及张量操作优化（optimize_dot_operands）等。这些优化 passes 的目的是为了提高 GPU 上代码的执行效率，是 Triton 编译器优化框架的核心组成部分。我们可以通过 https://github.com/triton-lang/triton/blob/main/include/triton/Dialect/TritonGPU/Transforms/Passes.td 和 https://github.com/triton-lang/triton/blob/main/include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td 这两个 Tablegen 文件查到这 Triton GPU 优化 Pass 的定义。

其实从这两个文件的命名我们可以发现 Triton 在 GPU Passes 应该还在做通用和专用的抽象，比如 TritonGPU 和 TritonNvidiaGPU 的命名。

这里列了一些 Triton GPU 优化 Pass：

内存合并访问（Coalescing）
F32 点积运算优化
CTA（合作线程阵列）规划
线程局部性优化
矩阵乘法加速
点运算操作数优化
数据去重
指令重排序
TMA（Tensor Memory Access）lowering 等优化

接下来我们可以浏览一下 vector_add 的 GPU IR：

这里展示了一下 triton vector_add lowering 到 gpu IR 的结果，我们可以看到其中一些很容易对应到 Python code 的内容，例如：

Module Attributes:

triton_gpu.num-ctas = 1：指定了一个 CTA (Cooperative Thread Array)
triton_gpu.num-warps = 4：每个 CTA 包含 4 个 warp
triton_gpu.threads-per-warp = 32：每个 warp 包含 32 个线程
目标平台是 "cuda:89"（计算能力 8.9 的 CUDA 设备）

Program ID and Range Creation:

%c1024_i32 创建常量 1024
%0 获取程序 ID
%1 将程序 ID 乘以 1024
%2 创建一个从 0 到 1024 的范围，生成一个 1024xi32 的张量

Ops (Splat and add operations):

%3 使用 splat 操作将标量广播到 1024xi32 张量
%4 执行加法操作
%13 显示了浮点数加法操作 (addf)

Load and store operations:

%7-%9 展示了数据加载操作：
splat 操作将指针广播
addptr 计算偏移后的地址
load 从计算出的地址加载数据
最后使用 tt.store 将结果存回内存
tt.return 标志着 kernel 的结束

可以看到这个 IR 展示了向量加法操作的底层实现步骤。

如果你想进一步了解优化 Pass 的实现，例如 TritonGPUAccelerateMatmul 这个 Pass，你在 Triton 代码仓库搜索的时候应该能看到 3 个相关的地方(https://github.com/search?q=repo%3Atriton-lang%2Ftriton%20TritonGPUAccelerateMatmul&type=code)：

分别是这个 Pass 的 Tablegen 定义，具体的 MLIR 实现和 Python Binding。

接下来作者简单展示了一下在 Triton 的 codebase 基础上实现了 2 个新的不需要传参数的简单 Pass，打印 OpGraph 和记录 Op 数量的 Pass，大家感兴趣也可以尝试下。

总结

这节课详细介绍了 Triton 编译器的内部工作原理。文章首先介绍了 CUDA 编译器（NVCC）的工作流程，然后深入探讨了 Triton 编译器的架构设计，包括其如何将 Python DSL 代码通过多个中间表示（IR）最终编译成 CUDA 可执行文件。课程重点讲解了 Triton 与 MLIR（Multi-Level Intermediate Representation）的关系，展示了各种优化 Pass 的实现和作用，如 TritonCombineOps、TritonReorderBroadcast 等。通过一个 vector_add 的具体示例，详细展示了代码从 Python 层面到 GPU IR 的转换过程，以及在这个过程中各种中间产物（如 .cubin、.ptx 等）的生成。最后还介绍了如何在 Triton 中实现新的编译器 Pass，为读者提供了深入理解和扩展 Triton 编译器的实践指导。

END

作者：BBuf
来源：GiantPandaCV

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。

CUDA-MODE 课程笔记 第 29 课 Triton 内部机制

第 29 课，Triton Internals

课程笔记

总结

推荐阅读

目录