Quartz 4

协程组

7 min read

协程组

在 CUDA 中，Cooperative Groups（协作组） 是一种用于灵活组织线程协作的编程模型，而 Tensor Core 是 GPU 上的专用硬件单元，用于加速矩阵乘法 - 累加（MMA）等张量运算。结合两者编程时，Cooperative Groups 可用于精细化管理线程协作（如 warp 级、块级同步），配合 WMMA（Warp Matrix Multiply-Accumulate）API 高效调用 Tensor Core。

核心思路

Tensor Core 的运算以 warp（32 线程） 为基本单位（例如，一个 warp 可处理 16x16x16 的半精度矩阵乘法）。Cooperative Groups 可用于：

  1. 显式定义 warp 级或块级线程组，确保 Tensor Core 操作的线程同步；
  2. 协调多个 warp 协作处理更大规模的矩阵（超出单个 warp 的处理能力）；
  3. 简化线程与 Tensor Core 硬件的映射逻辑。

具体步骤与示例代码

以下是使用 Cooperative Groups 配合 WMMA API 编程 Tensor Core 的典型流程（以半精度矩阵乘法为例）：

1. 必要头文件

需包含协作组和 WMMA 的头文件：

 
#include <cooperative_groups.h>
 
#include <mma.h>
 
using namespace cooperative_groups;
 
using namespace nvcuda;

2. 定义 Tensor Core 支持的矩阵参数

Tensor Core 对矩阵尺寸有固定要求（如 16x16x16、32x8x16 等，取决于数据类型）。以最常用的 16x16x16 半精度（half）矩阵乘法 为例：

  • 输入矩阵 A：形状为 M×K（此处 M=16，K=16）
  • 输入矩阵 B：形状为 K×N（此处 K=16，N=16）
  • 输出矩阵 C：形状为 M×N（此处 16×16）

3. 用 Cooperative Groups 组织 Warp 级协作

Tensor Core 的操作由 warp 内线程共同完成（每个线程负责矩阵的一部分）。通过 cooperative_groups::warp 获取当前 warp 组，确保线程同步：

 
__global__ void tensor_core_mmul(half *C, const half *A, const half *B, int M, int N, int K) {
 
// 1. 获取当前 warp 协作组（32 线程）
 
auto warp = cooperative_groups::warp::current_thread();
 
  
 
// 2. 定义 WMMA 矩阵片段（fragment），用于存储 Tensor Core 操作的矩阵块
 
wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag; // A 矩阵片段（列主序）
 
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag; // B 矩阵片段（行主序）
 
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag; // 累加器（用 float 避免精度损失）
 
  
 
// 3. 初始化累加器为 0
 
wmma::fill_fragment(c_frag, 0.0f);
 
  
 
// 4. 加载 A 和 B 的子矩阵到片段（每个 warp 处理 16x16x16 的块）
 
// 计算当前 warp 处理的矩阵块坐标（简化示例，假设 M、N、K 均为 16 的倍数）
 
int warp_row = (blockIdx.y * blockDim.y + threadIdx.y) / 16; // 每个 warp 处理 16 行
 
int warp_col = (blockIdx.x * blockDim.x + threadIdx.x) / 16; // 每个 warp 处理 16 列
 
  
 
// 加载 A 的子块（M×K）和 B 的子块（K×N）
 
wmma::load_matrix_sync(a_frag, A + warp_row * 16 * K, K); // A 的步长为 K（列主序）
 
wmma::load_matrix_sync(b_frag, B + warp_col * 16, N); // B 的步长为 N（行主序）
 
  
 
// 5. 执行 Tensor Core 矩阵乘法-累加（warp 内线程协作完成）
 
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); // c_frag = a_frag * b_frag + c_frag
 
  
 
// 6. 存储结果到输出矩阵 C
 
wmma::store_matrix_sync(C + warp_row * 16 * N + warp_col * 16, c_frag, N, wmma::mem_row_major);
 
  
 
// 7. （可选）warp 级同步（确保所有线程完成存储）
 
warp.sync();
 
}

4. 核函数配置与启动

需确保线程块大小适配 warp 数量（每个 warp 32 线程），且 GPU 架构支持 Tensor Core（Volta 及以上，-arch=sm_70 及更高）：

 
int main() {
 
// 假设 A、B、C 为设备端半精度矩阵（M=16, N=16, K=16）
 
half *d_A, *d_B, *d_C;
 
cudaMalloc(&d_A, 16*16*sizeof(half));
 
cudaMalloc(&d_B, 16*16*sizeof(half));
 
cudaMalloc(&d_C, 16*16*sizeof(half));
 
  
 
// 启动核函数：1 个块，32 线程（刚好 1 个 warp）
 
dim3 block(32); // 32 线程 = 1 个 warp
 
dim3 grid(1);
 
tensor_core_mmul<<<grid, block>>>(d_C, d_A, d_B, 16, 16, 16);
 
  
 
cudaFree(d_A);
 
cudaFree(d_B);
 
cudaFree(d_C);
 
return 0;
 
}

关键细节解析

  1. 协作组的作用：
  • warp::current_thread() 获取当前 warp 组，确保 load_matrix_sync、mma_sync 等操作在 warp 内线程间同步（这些函数依赖 warp 级协作，必须所有 32 线程参与）。
  • 对于更大的矩阵（如 1024x1024），可通过 块级协作组（thread_block）协调多个 warp 分工（每个 warp 处理 16x16 子块，块级同步确保整体正确性）。
  1. WMMA 与 Tensor Core 的绑定：

wmma::mma_sync 是直接调用 Tensor Core 的接口，其实现依赖 warp 内线程的协作（每个线程处理子矩阵的特定元素）。Cooperative Groups 在此处的价值是显式管理这种协作，避免线程不同步导致的错误。

  1. 硬件限制：
  • 仅 Volta（sm_70）及以上架构支持 Tensor Core。
  • 矩阵尺寸必须符合 WMMA 要求（如 16x16x16、32x8x16 等），否则无法利用 Tensor Core。

总结

结合 Cooperative Groups 编程 Tensor Core 的核心是：

  • 用 warp 组管理 Tensor Core 操作的基本线程单元（32 线程），确保同步；
  • 用 thread_block 等更高层级的协作组协调多个 warp 处理大规模矩阵；
  • 配合 WMMA API 完成矩阵加载、计算、存储的全流程，充分利用 Tensor Core 的硬件加速能力。

编译时需指定支持 Tensor Core 的架构（如 nvcc -arch=sm_80…）。

编译与运行说明

上述代码需通过支持 Tensor Core 的 CUDA 工具链编译，命令示例：

 
nvcc -arch=sm_80 tensor_core_mmul.cu -o tensor_core_mmul # sm_80对应Ampere架构（如A100）

代码中，Cooperative Groups 的 warp 组确保了 Tensor Core 操作（load_matrix_sync、mma_sync）的线程同步，而 WMMA API 直接调用硬件加速单元完成高效矩阵乘法。对于更大规模的矩阵，可扩展网格和块尺寸，并通过块级协作组（thread_block）协调多个 warp 的工作。

Graph View

Backlinks