Triton 内核编译阶段 – PyTorch - PyTorch 深度学习库

Triton 开源编程语言和编译器提供了一种高级的、基于 Python 的方法来创建高效的 GPU 代码。在本博客中，我们重点介绍 Triton 程序如何编译以及中间表示的底层细节。有关 Triton 的介绍，我们建议读者参阅此博客。

Triton 语言和编译

Triton 编程语言支持不同类型的现代 GPU，并遵循分块编程方法。作为一个例子，我们将以Triton 向量加法教程为例，并进行少量修改。向量加法内核和辅助函数定义如下

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr,  # *Pointer* to first input vector.
               y_ptr,  # *Pointer* to second input vector.
               output_ptr,  # *Pointer* to output vector.
               n_elements, 
               BLOCK_SIZE: tl.constexpr, 
               ):
  
    pid = tl.program_id(axis=0) 
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
 
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)
 
def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = output.numel()

    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    triton_kernel=add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    torch.cuda.synchronize()

    # Save compilation stages - some of the stages identified here are specific to NVIDIA devices:
    with open('triton_IR.txt', 'w') as f:
        print(triton_kernel.asm['ttir'], file=f)
    with open('triton_TTGIR.txt', 'w') as f:
        print(triton_kernel.asm['ttgir'], file=f)
    with open('triton_LLVMIR.txt', 'w') as f:
        print(triton_kernel.asm['llir'], file=f)
    with open('triton_PTX.ptx', 'w') as f:
        print(triton_kernel.asm['ptx'], file=f)
    with open('triton_cubin.txt', 'w') as f:
        print(triton_kernel.asm['cubin'], file=f)

    return output

torch.manual_seed(0)
size = 98432
x = torch.rand(size, device='cuda')
y = torch.rand(size, device='cuda')
output_torch = x + y
output_triton = add(x, y)
print(output_torch)
print(output_triton)
print(f'The maximum difference between torch and triton is '
      f'{torch.max(torch.abs(output_torch - output_triton))}')

Triton 向量加法内核包含 @triton.jit 装饰器。Triton 编译器将编译由 @triton.jit 标记的函数，它通过多个编译阶段降低函数。辅助函数 add 分配输出张量，计算适当的 GPU 网格大小，并额外保存中间编译阶段。

重点关注编译过程，Triton 内核通过一系列阶段（如下图所示）降低到设备特定的汇编代码。

首先通过遍历装饰后的 Python 函数的抽象语法树 (AST) 来编译内核，以创建 Triton 中间表示 (Triton-IR)。Triton-IR 是一种未优化、与机器无关的中间表示。它引入了瓦片级编程要求，并基于开源 LLVM 编译器项目。接下来，Triton 编译器优化 Triton-IR 并将其转换为 Triton-GPU IR (Triton-TTGIR)，然后是 LLVM-IR。Triton-IR 和 Triton-GPUIR 表示都以 MLIR 语言方言编写，其中 MLIR 是 LLVM 的一个子项目，旨在改进异构硬件的编译。

对于 Triton 向量加法教程内核，Triton IR 示例代码片段如下

module {
  tt.func public @add_kernel(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32} loc("/u/saraks/triton_blog/01-vector-add.py":28:0), %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32} loc("/u/saraks/triton_blog/01-vector-add.py":28:0), %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32} loc("/u/saraks/triton_blog/01-vector-add.py":28:0), %arg3: i32 {tt.divisibility = 16 : i32} loc("/u/saraks/triton_blog/01-vector-add.py":28:0)) attributes {noinline = false} {
    %c1024_i32 = arith.constant 1024 : i32 loc(#loc1)
    %0 = tt.get_program_id x : i32 loc(#loc2)
    %1 = arith.muli %0, %c1024_i32 : i32 loc(#loc3)
    %2 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> loc(#loc4)
    %3 = tt.splat %1 : i32 -> tensor<1024xi32> loc(#loc5)
    %4 = arith.addi %3, %2 : tensor<1024xi32> loc(#loc5)
    %5 = tt.splat %arg3 : i32 -> tensor<1024xi32> loc(#loc6)
    %6 = arith.cmpi slt, %4, %5 : tensor<1024xi32> loc(#loc6)
    %7 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>> loc(#loc7)
    %8 = tt.addptr %7, %4 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32> loc(#loc7)
    %9 = tt.load %8, %6 : tensor<1024x!tt.ptr<f32>> loc(#loc8)
    %10 = tt.splat %arg1 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>> loc(#loc9)
    %11 = tt.addptr %10, %4 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32> loc(#loc9)
    %12 = tt.load %11, %6 : tensor<1024x!tt.ptr<f32>> loc(#loc10)
    %13 = arith.addf %9, %12 : tensor<1024xf32> loc(#loc11)
    %14 = tt.splat %arg2 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>> loc(#loc12)
    %15 = tt.addptr %14, %4 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32> loc(#loc12)
    tt.store %15, %13, %6 : tensor<1024x!tt.ptr<f32>> loc(#loc13)
    tt.return loc(#loc14)
  } loc(#loc)
} loc(#loc)

请注意，Triton 内核中的主要函数现在表示为

Triton 内核	Triton IR
x = tl.load(x_ptr + offsets, mask=mask)	%9 = tt.load %8, %6 : tensor<1024x!tt.ptr<f32>> loc(#loc8)
y = tl.load(y_ptr + offsets, mask=mask)	%12 = tt.load %11, %6 : tensor<1024x!tt.ptr<f32>> loc(#loc10)
output = x + y	%13 = arith.addf %9, %12 : tensor<1024xf32> loc(#loc11)
tl.store(output_ptr + offsets, output, mask=mask)	tt.store %15, %13, %6 : tensor<1024x!tt.ptr<f32>> loc(#loc13)

在 Triton IR 阶段，%arg0: !tt.ptr<f32> 和随后的张量引用表明中间表示已经根据数据类型进行了专门化。

我们在配备 CUDA 12.2 版、Python 3.11.9 版和 PyTorch 2.4.1（使用 PyTorch 随附的默认 Triton 版本）的 Tesla V100-SXM2-32GB GPU 上运行了此示例。在此设备上，简单的向量加法具有以下 Triton GPU IR 代码片段（为清晰起见，省略了一些行）

#blocked = #triton_gpu.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
module attributes {"triton_gpu.num-ctas" = 1 : i32, "triton_gpu.num-warps" = 4 : i32, triton_gpu.target = "cuda:70", "triton_gpu.threads-per-warp" = 32 : i32} {
  tt.func public @add_kernel(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}
    ⋮
    %9 = tt.load %8, %6 : tensor<1024x!tt.ptr<f32>, #blocked> loc(#loc8)
    ⋮
    %12 = tt.load %11, %6 : tensor<1024x!tt.ptr<f32>, #blocked> loc(#loc10)
    %13 = arith.addf %9, %12 : tensor<1024xf32, #blocked> loc(#loc11)
    ⋮
    tt.store %15, %13, %6 : tensor<1024x!tt.ptr<f32>, #blocked> loc(#loc13)
    ⋮
  } loc(#loc)
} loc(#loc)

在此阶段，包含了一些硬件特定信息。例如，计算能力以及张量如何分布到核心和 Warp 或 AMD GPU 上的波前。在此示例中，张量表示为 #blocked 布局。在此编码中，每个 Warp 拥有一部分连续的张量。目前，其他可能的内存优化包括 slice（沿维度重组和分布张量）、dot_op（块矩阵乘法的优化布局）、shared（表示 GPU 共享内存）、nvidia_mma（由 NVIDIA 张量核心生成）、amd_mfma（由 AMD MFMA 矩阵核心生成）和 amd_wmma（由 AMD WMMA 矩阵核心生成）等布局。正如最近 Triton 会议上宣布的那样，这种布局表示将过渡到新的线性布局，以统一后端内部和跨后端的布局。从 Triton-GPUIR 到 LLVM-IR 的阶段将 Triton-GPUIR 转换为 LLVM 的表示。目前，Triton 支持 NVIDIA 和 AMD 设备的第三方后端，但其他设备的支持正在由开源社区积极开发中。

LLVM-IR 向量加法参数的一小部分示例如下

  %19 = extractvalue { i32, i32, i32, i32 } %18, 0, !dbg !16
  %39 = extractvalue { i32, i32, i32, i32 } %38, 0, !dbg !18
  %23 = bitcast i32 %19 to float, !dbg !16
  %43 = bitcast i32 %39 to float, !dbg !18
  %56 = fadd float %23, %43, !dbg !19

在一些指针算术和内联汇编调用从全局内存中检索数据之后，向量元素被提取并转换为正确的类型。最后，它们被加在一起，并通过内联汇编表达式写入全局内存。

Triton 编译过程的最后阶段将 LLVM-IR 降低到特定设备的二进制文件。对于 NVIDIA GPU 上的向量加法示例，下一个中间产物是 PTX（并行线程执行）。低级 PTX 语法指定了 NVIDIA 设备在线程级别的执行，从 CUDA 1.0 版本开始。有关 PTX 的深入指南，请参阅 NVIDIA 的文档。在向量加法中，内核参数从主机传递到内核，分配地址，并且 mov 指令促进线程级数据访问，最终通过 add.f32 表示元素加法调用，例如以下示例

	add.f32 	%f17, %f1, %f9// add type float32, output register, input register for x, input register for y

Triton 编译器与不同的硬件后端协同，协调最终阶段，管理汇编代码如何编译成二进制文件。Triton 内核现在可以使用了。

总结

Triton 提供了一个高级抽象来为不同类型的硬件编程和编译内核。在这篇文章中，我们重点介绍了 Triton 代码表示和 Triton 编译器的不同阶段。有关包含自定义 Triton 内核或使用 Triton 内核加速不同工作负载的详细信息，请查看 PyTorch Triton 教程，关于 Triton GPTQ 内核、使用 Triton 进行 Llama3 FP8 推理和 LLM 的无 CUDA 推理的博客文章，或 PyTorch 2.2 Triton 代码生成部分。

Triton 核编译阶段

Triton 语言和编译

总结

文档

教程

资源

保持联系以获取更新、活动信息和最新新闻