GPU 线程、warp、线程块与网格

GPU 内核暴露了大量并行工作。执行模型将这些工作组织成线程、warp、线程块和网格的层级结构。每一层回答不同的问题:单个逻辑工作者做什么、哪些工作者一起执行、哪些工作者可以协作,以及整个内核启动如何覆盖输入。

线程

线程(thread)(在 AMD GPU 上也被称为工作单元(work unit))是内核函数内最小的逻辑执行单元。每个线程都运行相同的内核代码,但内置的坐标让每个线程可以选择不同的数据。这种模型称为 SIMT:单指令多线程(Single Instruction, Multiple Threads)。

你不需要手动为每个线程指定要做什么。相反,每个线程通常会把在块内的位置与块的位置结合起来,计算出要处理输入的哪一部分。更多信息请参见下面的网格与索引。

同一个内核函数体可以处理包含数千或数百万个元素的向量。线程还可以使用二维或三维坐标,这些坐标可以自然地映射到矩阵、图像和分块计算。

Warp

warp(在 AMD GPU 上也被称为波前(wavefront))是线程块中一起执行的一小组线程。在 NVIDIA GPU 上,一个 warp 包含 32 个线程。AMD 波前传统上是 64 个,在 RDNA 上则为 32 或 64 个。

warp 是实际的调度单元。它被呈现给 warp 调度器,一个 warp 一次只能被发出一条指令。当 warp 收到一条指令时,其中活动的线程各自在它们自己的寄存器和数据上执行该指令。不同的 warp 相互独立,可以在同一时刻执行不同的指令。

warp 中的线程不一定要产生相同的结果,但当它们走相同的控制流路径时,执行效率最高。如果一条分支把一些线程引向一个方向、另一些引向另一个方向,warp 就会分别执行这些路径。每条路径都有不同的活动线程掩码,只有当前路径上的线程是活动的。这被称为 warp 发散(warp divergence),当发散路径包含大量工作时,它会降低效率,因为不在当前路径上的线程处于空闲状态。

线程块

线程块(thread block,或简称 block)(在 AMD GPU 上也被称为工作组(workgroup))是网格(grid)内的一小组线程,而网格是执行内核函数的线程的顶层组织结构。作为工作负载分配的主要构建块,线程块有多个关键用途:

1. 它们把内核函数的整体工作负载(由网格管理)分解成更小、更易于管理的部分,这些部分可以被独立处理。这种划分使得 GPU 上多个流式多处理器(SM)之间能获得更好的资源利用和调度灵活性。
2. 线程块为线程提供了通过共享内存和同步原语进行协作的作用域,从而支持高效的并行算法和数据共享模式。
3. 线程块有助于可扩展性,它允许同一个程序在不同 GPU 架构上高效运行,因为硬件可以根据可用资源自动分配线程块。

你可以指定网格中线程块的数量,以及它们在一维、二维或三维维度上的排布方式。网格中的每个线程块都会被分配一个唯一的块索引,用于确定它在网格中的位置。同样,你也可以指定每个线程块中的线程数量以及它们在一维、二维或三维维度上的排布方式。当前 NVIDIA 和 AMD GPU 通常允许每个线程块最多 1024 个线程,具体受设备相关的维度、寄存器和共享内存限制。

GPU 会把网格中的每个线程块分配给一个流式多处理器(SM),线程块通常在该 SM 上保持驻留直至完成;线程块不会在 SM 之间迁移。当寄存器、共享内存和调度器槽位允许时,多个线程块可以同时驻留在同一个 SM 上。

线程块内的线程可以通过共享内存在线程之间共享数据,并使用内建机制同步,但它们无法直接与其它线程块中的线程通信。

网格与索引

**网格(grid)**是定义内核启动的顶层组织结构——它包含所有线程块,而线程块包含所有执行线程。

要编写 GPU 内核,你必须通过创建网格来指定要执行的工作。在 NVIDIA GPU 上,CUDA 编程模型允许你使用以下值定义内核的启动维度和坐标:

• gridDim:网格维度,即每个维度上的线程块数量。
• blockDim:线程块维度,即线程块每个维度上的线程数量。
• blockIdx:当前线程块在网格中的位置。
• threadIdx:当前线程在线程块中的位置。

内核会组合这些值(通常通过计算全局线程 ID),来确定每个线程应处理哪个元素或哪部分数据。

编写内核函数

现在你已经理解了 GPU 工作是如何组织为线程、warp、线程块和网格的,就能更好地理解内核代码了。

下面是一个执行向量加法的简单 CUDA 内核,每个线程处理一个元素:

__global__ void vecAdd(float* A, float* B, float* C, int vectorLength)
{
    int workIndex = threadIdx.x + blockDim.x * blockIdx.x;
    if (workIndex < vectorLength)
    {
        C[workIndex] = A[workIndex] + B[workIndex];
    }
}

这里的边界检查很重要,因为网格很少能整除输入。启动本身使用位于三尖括号之间的执行配置(execution configuration),其中第一个值设置网格中的线程块数量,第二个值设置每个线程块中的线程数量:

int threads = 256;
int blocks = (vectorLength + threads - 1) / threads;

vecAdd<<<blocks, threads>>>(devA, devB, devC, vectorLength);

网格通常包含的线程块数量超过 GPU 一次能运行的数量。随着线程块完成,GPU 会把新线程块分配给空闲的 SM 容量。这种调度模型让同一次启动可以在拥有不同 SM 数量的 GPU 上扩展,而无需修改内核。

然而,上面的 CUDA 代码只在 NVIDIA GPU 上有效。如果你想为 AMD GPU 编程,可以使用 AMD 的原生 ROCm 软件。这两个框架都以 C/C++ 作为核心编程语言,并带有一些额外的关键字和语法扩展。

另外,MAX 使用 Mojo 编程语言为 GPU 和其他加速器提供了与硬件无关的编程模型。该框架提供了与上述相同的网格与线程块编程模型,只不过代码可以面向 NVIDIA GPU、AMD GPU、Apple silicon 等更多平台。例如,下面是用 Mojo 编写的同一个向量加法函数:

from std.gpu import block_dim, block_idx, thread_idx
from max.gpu.host import DeviceContext

def kernel():
    var i = block_idx.x * block_dim.x + thread_idx.x
    # Process element i.

def main() raises:
    var ctx = DeviceContext()
    # Same geometry as the CUDA launch above.
    ctx.enqueue_function[kernel](grid_dim=4096, block_dim=256)
    ctx.synchronize()

grid_dim 和 block_dim 对应 CUDA 启动时的两个值,索引运算完全相同。Mojo 还提供了 global_idx,一个替你计算 block_idx.x * block_dim.x + thread_idx.x 的简写,所以同一行可以写成 var i = global_idx.x。

要学习如何为 NVIDIA GPU 编程,请参阅 CUDA 编程指南。

要学习如何为任何 GPU 编程,请参阅 Mojo GPU 编程指南。

层级结构如何影响内核性能

启动几何结构会改变内核使用硬件的方式:

• 线程块通常需要足够的线程来提供多个 warp,但过大的线程块会消耗过多的寄存器或共享内存。
• 相邻线程应尽可能访问相邻的全局内存地址,这让 GPU 能把内存请求合并成高效的事务。
• 在算法允许的情况下,分支应让 warp 内的线程保持在同一路径上。
• 网格应暴露足够多的独立线程块,以保持所有 SM 忙碌。

最佳几何结构取决于每个线程的工作量以及每个线程块所需的资源。资源使用还会影响占用率。像"总是用 256 个线程"这样的固定规则只是起点,而非普适答案。

扩展阅读

• CUDA C++ Programming Guide
• AMD ROCm HIP programming model
• Mojo GPU programming fundamentals