【CUDA】Triton

【CUDA】Triton

1. CUDA 与 Triton 的基本区别

CUDA 编程模型：

在传统的 CUDA 编程中，CUDA 是标量程序，带有阻塞线程（blocked threads）。

• 标量程序（Scalar Program）：表示我们直接为每个线程编写操作代码，每个线程处理一个数据元素。
• 阻塞线程（Blocked Threads）：为了高效地处理大数据集，线程被组织成线程块（thread blocks）。线程块之间的计算是隔离的，线程块内的线程通过共享内存（shared memory）进行协作。

在 CUDA 中，程序员直接处理线程和线程块。你需要自己编写线程级别的细节，如如何加载数据、如何存储结果、如何管理共享内存等。CUDA 程序员需要掌握这些低级细节，以便充分利用 GPU 的硬件特性。

Triton 编程模型：

Triton 是基于 CUDA 的高层抽象。Triton 是带有标量线程（scalar threads）的块化程序（blocked program）。

• 块化程序（Blocked Program）：表示你不再直接操作单个线程，而是关注线程块层次的优化。你编写的程序将自动被编译器转换为低级代码。
• 标量线程（Scalar Threads）：在 Triton 中，程序员不再需要管理线程级别的细节（如内存访问模式），编译器会自动为你处理这些复杂的操作，如数据加载、存储、共享内存使用等。

2. Triton vs CUDA：高层抽象 vs 低层控制

在 CUDA 中，程序员需要掌握线程级操作，包括如何组织线程、如何管理共享内存和同步等。程序员通常需要关注线程块之间如何交互，以及如何利用硬件特性（如共享内存、寄存器）来提高效率。简而言之，CUDA 提供的是较为底层的控制，允许程序员手动优化每个操作的细节。

而在 Triton 中，程序员更多关注高层操作，例如卷积、矩阵乘法等深度学习中的标准操作。Triton 会通过编译器自动处理低级细节，如数据加载、存储、内存调度、线程间同步等。Triton 让深度学习的 GPU 编程更像是高级语言编程，减少了繁琐的底层优化工作。

直观理解：

• CUDA：程序员编写的是标量程序，并通过线程块来管理计算工作。你需要手动优化每个线程的行为以及线程块之间的协调。这意味着你要深入理解 GPU 的硬件架构。
• Triton：程序员编写的是块化程序（关注大块操作），并使用标量线程。Triton 的编译器将自动处理线程级别的操作，使得程序员不再需要处理每个线程的细节，能更专注于高级操作。

3. Triton 的优势：简化深度学习编程

Triton 的一个重要特点是为深度学习程序员提供了一个更高层次的抽象，使得复杂的 GPU 编程变得更简洁。程序员不再需要掌握底层的线程管理或内存优化，可以像写 Python 代码一样编写高效的 GPU 核心操作（如卷积、矩阵乘法等）。这对于深度学习研究者来说是一个很大的优势，因为他们的工作更关注模型设计和算法优化，而不是 GPU 编程的低级细节。

举个例子，传统上使用 CUDA 实现 cuBLAS 或 cuDNN 等高效的深度学习库需要深入理解 GPU 架构和高效的内存管理策略，而 Triton 让 Python 程序员通过一个更直观的 API 编写出与这些库同样高效的代码。

4. 为什么不能直接跳过 CUDA 使用 Triton？

Triton 是建立在 CUDA 之上的，因此，理解 CUDA 的一些基本概念仍然是非常重要的。虽然 Triton 抽象了很多复杂的低级细节，但它的性能仍然依赖于底层 CUDA 的硬件特性。

• CUDA 是 Triton 的基础：Triton 的编译器最终会将代码转换成 CUDA 代码并在 GPU 上执行。因此，了解 CUDA 中的线程组织、内存模型等基础概念仍然是优化性能的重要基础。
• 自定义优化：对于一些特定场景，程序员可能需要手动优化代码，或针对特定硬件架构进行调优，这时可能需要直接编写 CUDA 内核，或者深入理解 CUDA 的底层特性来进一步优化 Triton 生成的代码。

5. 资源学习：

• Triton 文档：可以通过官方文档详细了解 Triton 的编程模型、API 和最佳实践。
• OpenAI 博客：可以深入了解 Triton 的设计思想、底层实现以及与 CUDA 的关系。
• GitHub：可以查看 Triton 的源代码、示例程序和开源项目，帮助你理解如何使用 Triton 编写高效的深度学习程序。

总结：

• CUDA 是一种低级 GPU 编程框架，程序员需要自己处理线程调度、内存访问等底层优化细节。
• Triton 提供了一个更高层次的抽象，简化了深度学习 GPU 编程，让程序员能够专注于算法层次的开发，而不需要担心低级硬件细节。
• Triton 是建立在 CUDA 基础之上的，因此了解 CUDA 的基本概念对深入理解 Triton 及其性能优化非常重要。

---

Code

通过下面简易的代码，学习triton的使用方法，以及了解triton相较于cuda的高层次抽象。

vec_add.py

这个程序使用 Triton 和 PyTorch 实现了向量加法（x + y），并对比了两者的性能。Triton 是一个用于编写高效 GPU 内核的工具，类似于 CUDA 但更简单。程序的核心是一个 Triton 内核 add_kernel，它分块处理数据并支持边界检查。通过性能测试，程序比较了 Triton 和 PyTorch 原生加法操作的吞吐量（GB/s），并生成图表展示结果。最终目的是展示 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,  # Size of the vector. 向量的大小。
               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process. 每个程序应处理的元素数量。
               # NOTE: `constexpr` so it can be used as a shape value. 注意：`constexpr` 因此它可以用作形状值。
               ):
    # There are multiple 'programs' processing different data. We identify which program
    # 有多个“程序”处理不同的数据。需要确定是哪一个程序：
    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0. 使用 1D 启动网格，因此轴为 0。
    # This program will process inputs that are offset from the initial data.
    # 该程序将处理相对初始数据偏移的输入。
    # For instance, if you had a vector of length 256 and block_size of 64, the programs would each access the elements [0:64, 64:128, 128:192, 192:256].
    # 例如，如果有一个长度为 256, 块大小为 64 的向量，程序将各自访问 [0:64, 64:128, 128:192, 192:256] 的元素。
    # Note that offsets is a list of pointers:
    # 注意 offsets 是指针列表：
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # Create a mask to guard memory operations against out-of-bounds accesses.
    # 创建掩码以防止内存操作超出边界访问。
    mask = offsets < n_elements
    # Load x and y from DRAM, masking out any extra elements in case the input is not a multiple of the block size.
    # 从 DRAM 加载 x 和 y，如果输入不是块大小的整数倍，则屏蔽掉任何多余的元素。
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    # Write x + y back to DRAM.
    # 将 x + y 写回 DRAM。
    tl.store(output_ptr + offsets, output, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor):
    # We need to preallocate the output.
    # 需要预分配输出。
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = output.numel()
    # The SPMD launch grid denotes the number of kernel instances that run in parallel.
    # SPMD（单程序多数据） 启动网格表示并行运行的内核实例的数量。
    # It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].
    # 它类似于 CUDA 启动网格。它可以是 Tuple[int]，也可以是 C