【infra之路】理解 Register Tiling

先看 Shared Memory 版本里，一个线程在做什么

每个线程只算 C 的 1 个元素 C[row][col]，内部循环长这样：

for (int k = 0; k < TILE_SIZE; k++) {
    sum += sA[threadIdx.y][k] * sB[k][threadIdx.x];
    //       ~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~
    //       每次从 Shared       每次从 Shared
    //       Memory 读 1 次      Memory 读 1 次
}

TILE_SIZE=32 时，这个循环从 Shared Memory 读了 32 + 32 = 64 次。

问题在哪？

Shared Memory 虽然快（~20 cycles），但还是比寄存器慢（~1 cycle）。如果能让多个线程共享同一次 Shared Memory 读取，就能进一步减少访问次数。

看一个具体的场景：假设同一个 Block 里，有 4 个线程分别算 C 的第 0、1、2、3 行、同一列：

线程 A 算 C[0][col]: 读 sA[0][k] * sB[k][col]
线程 B 算 C[1][col]: 读 sA[1][k] * sB[k][col]  ← sB[k][col] 重复读！
线程 C 算 C[2][col]: 读 sA[2][k] * sB[k][col]  ← 又重复读！
线程 D 算 C[3][col]: 读 sA[3][k] * sB[k][col]  ← 又重复读！

4 个线程对 sB[k][col] 读了 4 次，但其实同一个 k 时刻它们的值完全一样。如果其中一个线程把 sB[k][0..3] 缓存到寄存器，其他线程也能复用就好了——但寄存器是线程私有的，不能共享。

所以换个思路：让一个线程同时负责多个 C 的元素。

Register Tiling 的核心思想

每个线程不再算 1 个元素，而是算 4×4 = 16 个元素：

Shared Memory 版（每线程 1 个元素）：    Register Tiling（每线程 4×4 个元素）：

Block 32×32 = 1024 线程                Block 16×16 = 256 线程
每线程算 C 的 1 个元素                  每线程算 C 的 4×4 个小块
总共覆盖 32×32 的 C 区域               总共覆盖 64×64 的 C 区域
                                       (16×4) × (16×4) = 64×64

画出来就是：

C 的一个 Block 负责的区域（64×64）:

    ┌────┬────┬────┬────┬─ ─ ─ ┬────┐
    │ 4×4│ 4×4│ 4×4│ 4×4│      │ 4×4│  ← 16 列小块，每块 4 宽
    ├────┼────┼────┼────┼─ ─ ─ ┼────┤
    │ 4×4│ 4×4│ 4×4│ 4×4│      │ 4×4│
    ├────┼────┼────┼────┼─ ─ ─ ┼────┤
    │ 4×4│ 4×4│ 4×4│ 4×4│      │ 4×4│
    ├────┼────┼────┼────┼─ ─ ─ ┼────┤
    │ ...│    │    │    │      │    │
    └────┴────┴────┴────┴─ ─ ─ ┴────┘
    ↑ 16 行小块，每块 4 高

    每个 4×4 小块 = 1 个线程负责
    总共 16×16 = 256 个线程覆盖 64×64

关键循环：为什么省了 Shared Memory 读取？

for (int k = 0; k < BK; k++) {
    // 从 Shared Memory 读 A 的 TM 个值 → 缓存到寄存器
    float regA[TM];                    // 4 个寄存器
    for (int m = 0; m < TM; m++)
        regA[m] = sA[threadRow + m][k];  // 读 4 次 Shared Memory

    // 从 Shared Memory 读 B 的 TN 个值 → 缓存到寄存器
    float regB[TN];                    // 4 个寄存器
    for (int n = 0; n < TN; n++)
        regB[n] = sB[k][threadCol + n];  // 读 4 次 Shared Memory

    // 纯寄存器运算：4×4 = 16 次乘加，零内存访问！
    for (int m = 0; m < TM; m++)
        for (int n = 0; n < TN; n++)
            threadResults[m][n] += regA[m] * regB[n];
}

算一下账：对每个 k 值——

Shared Memory 读取量减少了 4 倍。而且 16 次乘加全是寄存器操作（~1 cycle），计算密度大幅提升。

用一张图对比三个版本

Naive:
  Global Memory ──直接读──→ 计算（1个元素）
  每元素读 Global: 2N 次

Shared Memory:
  Global Memory ──加载──→ Shared Memory ──读──→ 计算（1个元素）
  每元素读 Shared: 2×TILE = 64 次
  每元素读 Global: 2N/TILE = N/16 次

Register Tiling:
  Global Memory ──加载──→ Shared Memory ──读──→ Register ──算──→ 计算（16个元素）
  每元素读 Register: 2×4 = 8 次（~1 cycle 每次！）
  每元素读 Shared: 8/TM = 0.5 次
  每元素读 Global: 更少

每多一层缓存，速度就快一级，能复用的数据就更多。 这就是整个 GPU 优化范式的精髓。

代码里几个容易困惑的细节

1. #pragma unroll 是什么？

告诉编译器把循环展开（loop unrolling），避免循环开销（条件判断、跳转）。对于 TM=4 这种小常数循环，展开后性能更好：

// 编译器看到 #pragma unroll 后，会把：
for (int m = 0; m < 4; m++) regA[m] = sA[...][k];
// 展开成：
regA[0] = sA[...][k];
regA[1] = sA[...][k];
regA[2] = sA[...][k];
regA[3] = sA[...][k];
// 没有循环判断，全部并行执行

2. 为什么 Block 线程数反而变少了（256 vs 1024）？

因为每个线程要做更多工作（算 16 个元素 vs 1 个），用更少的线程就能覆盖更大的 C 区域。256 线程 × 每线程 16 元素 = 4096 个元素 = 64×64，正好覆盖 Block 负责的 C 子矩阵。线程数少但 occupancy 可能更高（因为每线程用的寄存器更多，但 Block 数也更多了）。

3. threadResults[TM][TN] 存在哪？

在寄存器里。4×4 = 16 个 float = 16 个寄存器。加上 regA（4 个）和 regB（4 个），每个线程约用 24 个寄存器——完全放得下。