Tensor Core编程与混合精度加速实战 - 第3节

TensorCore编程模型：从CUDA到WMMA再到cuBLAS

大家好，我是你们的老朋友。今天咱们聊聊TensorCore的编程模型。说实话，我刚接触TensorCore那会儿，也被各种API搞得有点晕。但摸爬滚打几年后，我总结出一条经验：理解编程模型，比背API重要得多。

这一章，我会带大家从CUDA编程模型出发，一步步过渡到WMMA API，最后看看怎么用cuBLAS优雅地调用TensorCore。嗯，咱们不搞虚的，直接上干货。

CUDA编程模型回顾

在深入TensorCore之前，咱们先快速回顾一下CUDA编程模型。你想想看，TensorCore本质上还是GPU上的一个计算单元，它的编程模型必然建立在CUDA的基础之上。

线程层次结构

CUDA把计算任务分解成三个层次：Grid → Block → Thread。我个人习惯把Grid想象成一个工厂，Block是车间，Thread就是工人。每个工人干自己的活，但车间里的工人可以互相帮忙（通过共享内存）。

• Thread：最小的执行单元，每个线程有自己的寄存器和局部内存
• Warp：32个线程组成一个warp，这是硬件调度和执行的基本单位。嗯，这里要注意，warp里的线程执行的是同一指令
• Block：多个warp组成block，block内的线程可以通过共享内存通信
• Grid：多个block组成grid，对应一个完整的kernel launch

内存层次结构

CUDA的内存模型也是分层的。我在项目中遇到过不少性能问题，追根溯源都是内存访问模式没搞对。

内存类型	位置	访问速度	作用域
寄存器	片上	最快	单个线程
共享内存	片上	快	单个block
全局内存	片外	慢（约几百周期）	所有线程
常量/纹理内存	片外（有缓存）	较快（只读场景）	所有线程

说白了，写CUDA程序就是在跟内存延迟做斗争。TensorCore之所以快，很大程度上是因为它把数据直接喂到了寄存器级别，绕过了很多中间环节。

Warp Matrix Multiply-Accumulate (WMMA) API

好，现在进入正题。WMMA API是CUDA 9.0之后引入的一套编程接口，专门用来操作TensorCore。为什么叫"Warp Matrix"？因为一个warp的线程共同完成一个矩阵分块的乘加运算。

核心概念：fragment

WMMA引入了一个新的数据类型叫fragment。你可以把它理解成一个"虚拟的矩阵分块"，它分散存储在各个线程的寄存器中。每个线程只持有这个分块的一部分数据。

// 定义一个16x16的矩阵分块，每个元素是half精度
wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

这里我解释一下模板参数：

• matrix_a / matrix_b / accumulator：标识这个fragment的角色
• 16, 16, 16：分别是M、N、K维度。注意，TensorCore有固定的分块大小
• half / float：数据类型。A和B通常是half，累加器可以是float
• row_major / col_major：矩阵的存储布局

三个核心操作

WMMA API主要就三个操作：load、mma、store。简单到令人发指，对吧？

// 1. 从全局内存加载数据到fragment
wmma::load_matrix_sync(a_frag, d_a, lda);
wmma::load_matrix_sync(b_frag, d_b, ldb);
wmma::load_matrix_sync(c_frag, d_c, ldc);

// 2. 执行矩阵乘加：D = A * B + C
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

// 3. 将结果存回全局内存
wmma::store_matrix_sync(d_d, c_frag, ldc, wmma::mem_row_major);

你看，代码就这么几行。但背后的硬件执行可复杂多了。每个线程只做自己那部分计算，然后通过warp内部的shuffle指令交换数据。整个过程对程序员是透明的。

WMMA的局限性

说实话，WMMA API虽然好用，但也不是万能的。我总结了几点：

• 固定分块大小：只能做16x16x16或8x32x16等固定大小的分块
• 数据类型受限：主要支持half、bfloat16、tf32等
• warp内编程：所有线程必须协同，不能单独控制
• 调试困难：fragment里的数据分布不直观，调试起来很头疼

使用cuBLAS调用TensorCore

如果你觉得WMMA还是太底层，那cuBLAS就是你的菜。cuBLAS是NVIDIA官方提供的BLAS库，它对TensorCore做了很好的封装。你不需要关心fragment，不需要关心warp，只需要调用一个函数。

cublasGemmEx：一键调用TensorCore

cuBLAS从8.0版本开始支持TensorCore。核心函数是cublasGemmEx，它比传统的cublasSgemm多了几个参数，用来指定计算精度和TensorCore的使用策略。

cublasStatus_t cublasGemmEx(
    cublasHandle_t handle,
    cublasOperation_t transa, cublasOperation_t transb,
    int m, int n, int k,
    const void *alpha,
    const void *A, cudaDataType_t Atype, int lda,
    const void *B, cudaDataType_t Btype, int ldb,
    const void *beta,
    void *C, cudaDataType_t Ctype, int ldc,
    cudaDataType_t computeType,
    cublasGemmAlgo_t algo);

关键在最后两个参数：

• computeType：指定计算精度。比如CUBLAS_COMPUTE_32F_FAST_16F表示输入用half，内部计算用float，输出用float
• algo：算法选择。设置成CUBLAS_GEMM_DEFAULT_TENSOR_OP就会自动使用TensorCore

cuBLASLt：更灵活的选择

如果你对性能有极致追求，可以看看cuBLASLt。它提供了更细粒度的控制，比如可以自定义epilogue操作（在矩阵乘之后加一个激活函数或bias）。

// cuBLASLt的典型用法
cublasLtMatmulDesc_t operationDesc;
cublasLtMatmulDescCreate(&operationDesc, computeType, CUDA_R_32F);
// 设置各种属性...
cublasLtMatmul(handle, operationDesc,
               alpha, A, Atype, B, Btype, beta, C, Ctype, D, Dtype,
               algo, workspace, workspaceSize, stream);

我个人建议：能用cuBLAS就别自己写WMMA。除非你有非常特殊的计算模式，或者你想深入理解TensorCore的工作原理。我在一个NLP模型优化项目中，用cuBLAS替换了手写的WMMA代码，性能提升了15%，代码量减少了80%。

性能对比

方法	代码量	性能	灵活性	推荐场景
cuBLAS	少（几行）	高（经过优化）	低	标准矩阵乘
cuBLASLt	中等	高（可调优）	中	需要自定义epilogue
WMMA	多（几十行）	中高（取决于实现）	高	自定义算子、研究
PTX汇编	非常多	最高（理论上）	最高	极致优化、学习

你想想看，从cuBLAS到WMMA再到PTX，抽象层次越来越低，控制力越来越强，但开发成本也越来越高。我的建议是：从cuBLAS开始，性能不够再往下走。

小结

这一章我们走完了从CUDA到WMMA再到cuBLAS的完整路径。说白了，TensorCore编程模型就是一层层抽象：

• 底层：CUDA线程模型 + 共享内存，手动管理数据流动
• 中层：WMMA API，用fragment抽象矩阵分块，warp内自动协同
• 高层：cuBLAS/cuBLASLt，一行代码搞定TensorCore调用

我个人更倾向于用cuBLAS做快速原型，用WMMA做精细调优。至于PTX汇编？嗯，除非你是想写论文或者挑战自我，否则真没必要。

下一章我们会深入TensorCore的硬件架构，看看这些神奇的矩阵运算到底是怎么在芯片上执行的。到时候我会分享一些我在Volta和Turing架构上踩过的坑，保证让你少走弯路。