第四讲：TensorCore矩阵乘法实现

各位同学，今天我们来聊聊TensorCore矩阵乘法的实战。说实话，这个主题我讲了不下二十次，但每次备课还是会发现新的细节。TensorCore这东西，说白了就是NVIDIA给深度学习量身定做的加速器。你想想看，传统的CUDA Core做矩阵乘法，一次只能处理一个标量，而TensorCore一次能处理一个4x4的小矩阵——这差距可不是一星半点。

4.1 WMMA API 基础

WMMA，全称是Warp Matrix Multiply-Accumulate。这是CUDA 9.0之后引入的一套API，专门用来操作TensorCore。我个人习惯把它理解成「给warp级别的矩阵运算准备的工具包」。

先看一个最简单的例子：

#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

// 声明wmma片段
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;

// 初始化累加器
wmma::fill_fragment(c_frag, 0.0f);

// 加载数据到片段
wmma::load_matrix_sync(a_frag, d_a, 16);
wmma::load_matrix_sync(b_frag, d_b, 16);

// 执行矩阵乘法
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

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

这里有几个关键点要注意。fragment的模板参数，第一个是矩阵类型，第二个是M，第三个是N，第四个是K。M、N、K必须是16的倍数——这是TensorCore的硬件限制。我在项目中遇到过有人把M设成8，结果编译通过了但运行结果全错，排查了半天才发现是这个问题。

注意：fragment的尺寸不是随便选的。16x16x16是最常用的配置，但如果你用Turing架构的卡，还可以用8x8x4的配置。Ampere之后又加了16x8x8和8x16x8。选哪个？看你的矩阵规模。

4.2 分块矩阵乘法策略

实际应用中，矩阵规模往往远大于16。比如一个1024x1024的矩阵乘法，你不可能一次性塞进TensorCore。这时候就需要分块——把大矩阵切成16x16的小块，逐个送入TensorCore计算。

我常用的分块策略是这样的：

const int BLOCK_SIZE = 16;
const int TILE_SIZE = 128;  // 每个block处理128x128的tile

__global__ void tensorcore_gemm(half* A, half* B, float* C, 
                                int M, int N, int K) {
    // 每个block负责一个tile
    int baseRow = blockIdx.y * TILE_SIZE;
    int baseCol = blockIdx.x * TILE_SIZE;
    
    // 每个warp处理tile内的16x16子块
    int warpRow = (threadIdx.x / 32) * 16;
    int warpCol = (threadIdx.x % 32) * 16;
    
    // 累加器
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc;
    wmma::fill_fragment(acc, 0.0f);
    
    // 沿着K维度循环
    for (int k = 0; k < K; k += 16) {
        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::load_matrix_sync(a_frag, A + (baseRow + warpRow) * K + k, K);
        wmma::load_matrix_sync(b_frag, B + k * N + (baseCol + warpCol), N);
        
        // 矩阵乘累加
        wmma::mma_sync(acc, a_frag, b_frag, acc);
    }
    
    // 写回结果
    wmma::store_matrix_sync(C + (baseRow + warpRow) * N + (baseCol + warpCol), 
                            acc, N, wmma::mem_row_major);
}

嗯，这里要注意一个细节。每个warp处理一个16x16的子块，但一个block里通常有多个warp。我习惯让一个block处理128x128的tile，这样就有8x8=64个warp同时工作，刚好占满一个SM。

4.3 Bank Conflict 的坑

说到Shared Memory，就不得不提Bank Conflict。Shared Memory被分成32个bank，每个bank宽度4字节。如果同一个warp里的线程访问了同一个bank的不同地址，就会发生冲突——说白了就是串行化了，性能直接打折扣。

我曾经在一个项目里，矩阵乘法跑出来只有理论峰值的30%。排查了三天，最后发现是Shared Memory的访问模式有问题。看这个例子：

// 有bank conflict的写法
__shared__ float s_data[32][32];
int tid = threadIdx.x;
float val = s_data[tid][tid];  // 所有线程访问同一列，bank冲突！

// 无bank conflict的写法
__shared__ float s_data[32][32 + 4];  // padding
float val = s_data[tid][tid];  // 加上padding后，每行错开，无冲突

为什么会这样？因为s_data[tid][tid]这种访问模式，所有线程都在访问同一列。32个线程访问32个不同的行，但列相同——这就意味着它们访问的是同一个bank的不同地址。加上4个元素的padding后，每行的起始地址错开了，自然就没有冲突了。

小技巧：给Shared Memory数组的第二维加上几个元素的padding，是最简单有效的bank conflict避免方法。我一般加4个元素，既不影响逻辑，又能保证错开。

4.4 Shared Memory 优化实战

Shared Memory的优化，说白了就是两件事：减少bank conflict，提高数据复用。在TensorCore的GEMM里，数据复用尤其重要。因为每个TensorCore操作需要从Shared Memory里读取16x16的数据块，如果每次都要从Global Memory加载，那带宽根本不够用。

我常用的优化策略是这样的：

数据预取：在计算当前tile的同时，预取下一个tile的数据到Shared Memory
双缓冲：用两块Shared Memory交替使用，计算和加载可以流水线进行
向量化加载：用float4或half2一次性加载多个元素，减少指令数

看一个双缓冲的示例：

__shared__ half s_A[2][TILE_SIZE][TILE_SIZE + 4];  // 双缓冲 + padding
__shared__ half s_B[2][TILE_SIZE][TILE_SIZE + 4];

int ping = 0;
// 预取第一个tile
load_tile_to_shared(s_A[ping], s_B[ping], ...);
__syncthreads();

for (int k = 0; k < K; k += TILE_SIZE) {
    // 预取下一个tile（异步）
    if (k + TILE_SIZE < K) {
        load_tile_to_shared(s_A[ping ^ 1], s_B[ping ^ 1], ...);
    }
    
    // 计算当前tile
    compute_tensorcore(s_A[ping], s_B[ping], acc);
    __syncthreads();
    
    ping ^= 1;  // 切换缓冲区
}

这个模式我用了好几年，效果一直很好。关键点在于load_tile_to_shared和compute_tensorcore要能重叠执行。如果计算时间比加载时间长，那预取就白做了——你得调整tile大小让两者匹配。

核心要点总结：

WMMA fragment尺寸必须符合硬件限制（16的倍数）
分块策略要匹配SM的warp调度能力
Shared Memory加padding是避免bank conflict的银弹
双缓冲+预取是隐藏内存延迟的有效手段

最后说一句，TensorCore的优化没有银弹。不同的矩阵规模、不同的GPU架构，最优参数都不一样。我建议你拿到一块新卡后，先跑一遍profiling，看看实际瓶颈在哪——是计算单元没跑满，还是内存带宽不够，还是warp调度有气泡。对症下药，才能事半功倍。

这张图展示了我个人习惯的优化流程。从输入矩阵开始，经过分块、Shared Memory加载、WMMA计算，最后输出结果。右侧的优化要点，每一个都值得单独开一节课来讲。不过今天我们先到这里，后面的章节会逐一深入。

想要更多资料，联系可可

📚 课程目录

第四讲：TensorCore矩阵乘法实现

4.1 WMMA API 基础

4.2 分块矩阵乘法策略

4.3 Bank Conflict 的坑

4.4 Shared Memory 优化实战