张量加速核性能调优实战 - 第2节

一、Tensor Core 硬件微架构：矩阵乘累加单元

各位同学，咱们今天聊聊 Tensor Core 的硬件微架构。说白了，就是那个专门做矩阵乘法的“小钢炮”。

我记得第一次接触 Tensor Core 时，心里想的是：“这不就是个加强版的 ALU 吗？”后来被现实狠狠教育了一顿。Tensor Core 的设计哲学，跟传统的 CUDA Core 完全不是一回事。

1.1 从 SIMT 到 Tensor Core 的进化

传统的 CUDA Core 走的是 SIMT 路线——单指令多线程。每个线程处理一个元素，适合做 element-wise 操作。但矩阵乘法呢？它天然就是“块”操作。

Tensor Core 的思路很直接：一次指令，完成一个 4x4 矩阵乘法。你想想看，如果用 CUDA Core 做 4x4 矩阵乘，得发多少条指令？16 个线程各算各的，还得同步。Tensor Core 一条指令搞定。

1.2 微架构细节：那个 4x4 的小方块

Tensor Core 内部到底长什么样？我画个图你就明白了。

看到这个图了吗？Tensor Core 内部有 16 个 MAC（乘累加单元），排成 4x4 的阵列。矩阵 A 的每一行广播到所有列，矩阵 B 的每一列广播到所有行。结果直接累加到 C 矩阵上。

我刚开始做优化时，总想手动控制这些 MAC 单元。后来发现，硬件已经帮你安排好了数据流，你只需要告诉它“我要做矩阵乘”就行。

二、Warp 级矩阵运算：32 个线程的集体舞

Tensor Core 不是给单个线程用的。它服务的对象是 Warp——32 个线程的集体。

2.1 为什么是 Warp 级？

你想想看，一个 4x4 矩阵乘需要 16 个 MAC。一个 Warp 有 32 个线程。每个线程负责一部分计算，大家协同完成一个更大的 Tile。

在 Volta 架构中，一个 Warp 通过 Tensor Core 指令完成 16x16x16 的矩阵乘。怎么做到的？

• 每个线程持有 4 个元素（来自矩阵 A 的一行）
• 每个线程持有 4 个元素（来自矩阵 B 的一列）
• 32 个线程合起来，刚好覆盖 16x16 的矩阵

2.2 指令级细节

CUDA 里调用 Tensor Core 的指令长这样：

// 伪代码：WMMA API (Warp Matrix Multiply Accumulate)
wmma::fragment<matrix_a, 16, 16, 16, half, row_major> a_frag;
wmma::fragment<matrix_b, 16, 16, 16, half, col_major> b_frag;
wmma::fragment<accumulator, 16, 16, 16, float> c_frag;

// 加载数据到 fragment
wmma::load_matrix_sync(a_frag, A_ptr, lda);
wmma::load_matrix_sync(b_frag, B_ptr, ldb);
wmma::load_matrix_sync(c_frag, C_ptr, ldc);

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

// 存储结果
wmma::store_matrix_sync(C_ptr, c_frag, ldc, wmma::mem_row_major);

看到没？没有显式的循环，没有线程索引计算。你只需要定义 fragment（数据片段），然后调用 mma_sync。剩下的交给硬件。

三、数据复用与流水线：别让 Tensor Core 闲着

Tensor Core 算得飞快，但数据搬运跟不上怎么办？这就是数据复用和流水线要解决的问题。

3.1 数据复用的层次

我画个数据流图，帮你理解数据是怎么复用的：

数据复用的核心思想：一次加载，多次使用。

• 全局内存 → 共享内存：加载一个 Tile 的数据到共享内存
• 共享内存 → 寄存器：每个线程从共享内存读取自己需要的元素
• 寄存器 → Tensor Core：数据在寄存器中，直接喂给 Tensor Core

我遇到过最典型的性能问题：共享内存加载次数太多。本来一个 Tile 的数据可以复用 16 次，结果代码写成了每次计算都重新加载。性能直接腰斩。

3.2 流水线：让加载和计算重叠

流水线的思路很简单：计算当前 Tile 的同时，加载下一个 Tile 的数据。

在 CUDA 里，可以用 __pipeline 指令或者手动双缓冲来实现：

// 双缓冲流水线示例
__shared__ half A_buff[2][BLOCK_SIZE][BLOCK_SIZE];
__shared__ half B_buff[2][BLOCK_SIZE][BLOCK_SIZE];

int write_idx = 0;  // 当前写入的缓冲区
int read_idx = 1;   // 当前读取的缓冲区

// 预加载第一个 Tile
load_tile(A_buff[write_idx], B_buff[write_idx]);
__syncthreads();

for (int k = 0; k < K_TILES; k++) {
    // 启动下一个 Tile 的加载（异步）
    write_idx = (write_idx + 1) % 2;
    if (k + 1 < K_TILES) {
        load_tile_async(A_buff[write_idx], B_buff[write_idx]);
    }
    
    // 计算当前 Tile
    compute_tile(A_buff[read_idx], B_buff[read_idx], C_frag);
    
    // 等待异步加载完成
    if (k + 1 < K_TILES) {
        wait_for_load();
        __syncthreads();
    }
    
    read_idx = (read_idx + 1) % 2;
}

四、Tile（分块）策略基础：把大问题切成小方块

矩阵乘法本质上是 O(N³) 的复杂度。对于大矩阵，你不能一次性算完。得切成小块——这就是 Tile 策略。

4.1 为什么需要 Tile？

原因很简单：寄存器不够用。

一个 SM 的寄存器文件大概 64K 到 128K 字节。如果矩阵是 1024x1024，光存一个矩阵就要 2MB（FP16）。根本放不下。

所以，我们把大矩阵切成小 Tile，每个 Tile 刚好能塞进共享内存和寄存器。

4.2 Tile 尺寸的选择

Tile 尺寸不是随便选的。它受限于：

• 共享内存大小：每个 SM 的共享内存有限（通常 48KB 到 164KB）
• 寄存器数量：每个线程能用的寄存器有限（通常 255 个）
• Tensor Core 的硬件约束：fragment 尺寸是固定的（16x16）

常见的 Tile 尺寸组合：

Tile 尺寸 (MxNxK)	共享内存占用	寄存器占用/线程	适用场景
16x16x16	1KB	32 个	小矩阵、延迟敏感
32x32x16	4KB	64 个	中等矩阵、平衡型
64x64x16	16KB	128 个	大矩阵、吞吐优先
128x128x16	64KB	256 个	超大矩阵、极致吞吐

4.3 Tile 策略的黄金法则

我总结了几条经验法则：

1. Tile 越大越好，但不能超过共享内存的一半。留点空间给其他数据。
2. M 和 N 维度要平衡。别让一个维度特别大，另一个特别小。
3. K 维度用 Tensor Core 的倍数。16、32、64 都行，别用奇数。
4. 考虑 Bank Conflict。共享内存的 bank 数量是 32，Tile 的 stride 别是 32 的倍数。

嗯，这些规则不是死的。实际调优时，还是要用 profiler 跑一下，看看瓶颈在哪。

五、总结：从硬件到策略的思维转变

今天的内容，说白了就是一句话：理解硬件，才能用好硬件。

• Tensor Core 是 Warp 级的矩阵乘累加单元，不是线程级的
• 数据复用和流水线是性能的关键，别让 Tensor Core 闲着
• Tile 策略要匹配硬件资源，不是越大越好

我刚开始做 Tensor Core 优化时，总想着“怎么让代码跑得更快”。后来发现，先理解硬件怎么工作，优化自然水到渠成。

下一节我们会深入具体的调优技巧，包括 Bank Conflict 的避免、寄存器压力的控制、以及如何用 profiler 找到瓶颈。今天的内容先到这里，有问题随时交流。