4. TensorCore编程模型：WMMA API深度解析

各位同学，今天我们来聊聊TensorCore编程模型中最核心的部分——WMMA API。说实话，我第一次接触这个API的时候，也被它那套fragment、load/store、mma_sync的组合拳搞得有点晕。但用久了你会发现，这套设计其实非常巧妙，它把底层复杂的矩阵运算封装得干干净净。

我个人习惯把WMMA API理解成「给TensorCore写的驱动程序」。你不需要关心SM里具体怎么调度，只需要按照API的规则来，TensorCore就能帮你把矩阵乘法跑得飞快。

4.1 WMMA API是什么？

WMMA全称是Warp Matrix Multiply-Accumulate。说白了，它是一套以warp（32个线程）为单位的矩阵运算接口。为什么是warp？因为TensorCore的矩阵运算最小粒度就是warp级别，单个线程是没法直接操作TensorCore的。

我记得刚入行时有个同事问：「能不能让一个线程调用TensorCore？」答案是不行。TensorCore的设计理念就是「协同计算」，32个线程一起上，才能发挥它的威力。

WMMA API主要包含三个核心操作：

• load：从全局内存或共享内存加载矩阵数据到fragment
• mma_sync：执行矩阵乘累加运算（D = A * B + C）
• store：将计算结果从fragment写回内存

4.2 fragment类型详解

fragment是WMMA API中最关键的数据类型。你可以把它想象成一个「分布式矩阵分片」——每个线程只持有矩阵的一部分数据，所有线程合起来才构成完整的矩阵。

fragment有三种类型，对应矩阵乘法的三个操作数：

fragment类型	对应矩阵	数据布局	典型大小（16x16x16）
matrix_a	A矩阵	行主序或列主序	每个线程持有8个元素
matrix_b	B矩阵	行主序或列主序	每个线程持有8个元素
accumulator	C/D矩阵	行主序	每个线程持有16个元素

这里有个细节要注意：accumulator类型的fragment每个线程持有的元素数量是matrix_a或matrix_b的两倍。为什么？因为accumulator既要存输入C，又要存输出D，而且TensorCore内部对累加器的精度要求更高。

4.3 load/store操作

load和store操作负责在内存和fragment之间搬运数据。这里有个容易踩的坑——数据布局必须和fragment声明时一致。

举个例子：

// 声明fragment
nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 16, half, 
                        nvcuda::wmma::row_major> a_frag;

// 加载数据 - 注意ldm参数
nvcuda::wmma::load_matrix_sync(a_frag, d_A, 16);

// 存储结果
nvcuda::wmma::store_matrix_sync(d_C, c_frag, 16, nvcuda::wmma::mem_row_major);

ldm（leading dimension）参数是什么意思？它表示矩阵在内存中的「行跨度」。如果你的矩阵是连续存储的，ldm就等于矩阵宽度。但如果矩阵有padding，ldm就要包含padding的大小。

我曾经在一个项目中，因为ldm参数写错了，导致加载出来的数据全是乱的。排查了半天才发现是ldm设成了矩阵宽度，但实际内存布局有对齐要求，ldm应该是32的倍数。

4.4 mma_sync同步指令

mma_sync是WMMA API的核心——它真正触发了TensorCore的矩阵乘累加运算。这个指令的名字里有个「sync」，意思是所有线程必须同步到达这个指令。

语法很简单：

nvcuda::wmma::mma_sync(d_frag, a_frag, b_frag, c_frag);

这行代码执行的是：d_frag = a_frag * b_frag + c_frag。注意，d_frag和c_frag可以是同一个fragment，这样就是原地累加。

为什么需要sync？因为TensorCore的矩阵运算需要所有32个线程协同完成。如果有些线程快有些线程慢，数据就会错乱。所以mma_sync内部会插入一个warp级别的barrier，确保所有线程都准备好了才开始计算。

我记得有一次优化一个BERT模型，发现mma_sync的调用次数比预期多了一倍。后来仔细看代码，发现有人在循环里重复声明了fragment，导致每次迭代都重新分配内存。优化后把fragment声明提到循环外面，性能提升了30%。

4.5 完整示例：16x16x16矩阵乘法

下面是一个完整的WMMA矩阵乘法示例，我加了详细的注释：

#include <cuda_fp16.h>
#include <cuda_wmma.h>

using namespace nvcuda::wmma;

// 16x16x16的矩阵乘法
__global__ void wmma_matmul(half *A, half *B, float *C, int M, int N, int K) {
    // 声明fragment
    fragment<matrix_a, 16, 16, 16, half, row_major> a_frag;
    fragment<matrix_b, 16, 16, 16, half, col_major> b_frag;
    fragment<accumulator, 16, 16, 16, float> c_frag;
    
    // 初始化累加器为0
    fill_fragment(c_frag, 0.0f);
    
    // 计算当前warp负责的矩阵块位置
    int warp_m = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    int warp_n = blockIdx.y * blockDim.y;
    
    // 循环加载并计算
    for (int k = 0; k < K; k += 16) {
        // 加载A和B的16x16块
        load_matrix_sync(a_frag, A + warp_m * 16 + k, K);
        load_matrix_sync(b_frag, B + k + warp_n * 16, N);
        
        // 执行矩阵乘累加
        mma_sync(c_frag, a_frag, b_frag, c_frag);
    }
    
    // 存储结果
    store_matrix_sync(C + warp_m * 16 + warp_n * 16, c_frag, N, mem_row_major);
}

这段代码虽然简单，但包含了WMMA编程的所有核心要素。你想想看，如果没有TensorCore，实现同样的16x16x16矩阵乘法需要多少指令？至少几百条。而用WMMA API，核心计算就一行mma_sync。

4.6 知识体系总览

为了让大家更直观地理解WMMA API的整体架构，我画了一张图：

从这张图你可以看到，WMMA API处于中间层，向上服务深度学习模型，向下驱动TensorCore硬件。fragment类型决定了数据如何分布到每个线程，数据布局则决定了内存访问模式。

好了，这一章的内容就到这里。WMMA API虽然看起来简单，但用好它需要你对warp协作、数据布局、内存对齐都有深入理解。下一章我们会深入fragment的内部实现，看看每个线程到底是怎么持有数据的。