3. TensorCore编程基础

各位同学，今天我们来聊聊TensorCore编程最核心的部分——WMMA API。说实话，我刚接触TensorCore那会儿，也被这套API搞得有点懵。但用顺手之后你会发现，它其实就是一套帮你管理TensorCore硬件单元的C++接口。

3.1 WMMA API概览

WMMA全称是Warp Matrix Multiply-Accumulate。说白了，它让CUDA程序员能在warp级别直接调用TensorCore。我习惯把它理解成「TensorCore的遥控器」——你按按钮，它算矩阵。

核心API都在 nvcuda::wmma 命名空间下。常用的几个：

• load_matrix_sync — 从显存加载数据到TensorCore的寄存器
• mma_sync — 执行矩阵乘加操作：D = A * B + C
• store_matrix_sync — 把结果写回显存
• fill_fragment — 初始化矩阵片段（比如清零）

3.2 矩阵分片策略

TensorCore一次处理一个「tile」（瓦片/分片）。你不能把整个大矩阵直接扔给它。你得把矩阵切成小块，一块一块喂。

为什么会这样？因为TensorCore的寄存器资源有限。一个tile的大小决定了你需要多少寄存器。我刚开始做优化时，总想用大tile来提升吞吐，结果寄存器溢出，性能反而下降。

常见的分片策略有三种：

1. 均匀分片 — 把矩阵切成大小相同的tile，适合规整的矩阵乘法
2. 自适应分片 — 根据矩阵形状动态调整tile大小，适合变长场景
3. 混合精度分片 — 不同精度用不同tile大小，比如FP16用16x16x16，INT8用8x32x16

3.3 Tile大小选择

TensorCore支持三种tile尺寸：

Tile尺寸	M	N	K	适用场景
16x16x16	16	16	16	通用矩阵乘，FP16/FP32混合
8x32x16	8	32	16	宽度较大的矩阵，比如全连接层
32x8x16	32	8	16	高度较大的矩阵，比如某些卷积

怎么选？我个人的经验法则：

• 如果M和N差不多，用16x16x16。这是最通用的选择。
• 如果N远大于M（比如batch size小但特征维度大），用8x32x16。
• 如果M远大于N，用32x8x16。

你想想看，选错tile大小会怎样？寄存器浪费、计算单元闲置、带宽利用率低。我在一个推荐系统项目里，就因为用了16x16x16处理一个8x1024的矩阵，性能比8x32x16差了30%。

3.4 数据类型转换

TensorCore原生支持的数据类型有限。你喂进去的数据往往需要先转换。常见的转换路径：

• FP32 → FP16：直接截断，精度损失可控
• FP32 → TF32：TensorFloat-32，精度介于FP16和FP32之间
• INT8 → FP16：需要量化/反量化，通常带scale因子
• BF16 → FP16：位宽相同但指数位不同，需要特殊处理

代码示例：

// 加载FP16数据到TensorCore片段
nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 16, 
                        __half, nvcuda::wmma::row_major> a_frag;

// 从显存加载，自动转换
nvcuda::wmma::load_matrix_sync(a_frag, d_a, 16);

// 执行矩阵乘，结果累加到float片段
nvcuda::wmma::fragment<nvcuda::wmma::accumulator, 16, 16, 16, float> c_frag;
nvcuda::wmma::fill_fragment(c_frag, 0.0f);

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

// 写回显存，float转回FP16
nvcuda::wmma::store_matrix_sync(d_c, c_frag, 16, 
                                 nvcuda::wmma::mem_row_major);

这里有个细节：mma_sync的累加器是float，但输入可以是FP16或TF32。为什么？因为TensorCore内部用高精度做累加，防止精度溢出。我见过有人直接用FP16做累加器，结果训练loss震荡——精度不够。

3.5 核心知识体系

下面这张图总结了本章的核心逻辑：

从这张图你能看到，整个流程是串行的。WMMA API是入口，分片策略决定怎么切数据，Tile大小影响硬件利用率，数据类型转换保证精度和性能的平衡。最后，所有操作都落到TensorCore硬件单元上执行。

好了，这一章的内容就到这里。记住，TensorCore编程的核心就是「分而治之」——把大矩阵切成小tile，用WMMA API喂给硬件。下一章我们会深入实战，写一个完整的矩阵乘法kernel。