这是“AI Infra 每日一问”系列的第 6 篇。今天想聊一个特别底层、但只要碰过 CUDA kernel 优化就绕不开的问题：寄存器 tiling 在不同架构上到底长什么样？这里的“寄存器”不是嵌入式里那种 PHY 寄存器地址，而是 CPU、GPU、AI 加速器上离计算单元最近的高速暂存空间；“tiling”也不是图片处理里的瓦片分割，而是把循环切块后，让数据尽量留在寄存器里重复使用。
如果你写过 GEMM、看过 Attention 的 kernel，或者翻过 Cutlass/Triton 生成的代码，一定见过一堆 16×8、64×64、128×128 之类的数字。这些数字往小了看是“一块数据多大”，往大了看是某一种硬件对数据复用方式给出的硬约束。这篇不是教科书，我会从“为什么需要它”开始，分别拆一拆 NVIDIA、AMD、CPU 以及脉动阵列上的寄存器 tiling，每个场景尽量落到“形状怎么来的、代码长什么样、踩过什么坑”上。适合谁读？正在学算子优化的人、想读懂 Cutlass/Triton 代码的人，以及被 occupancy 和 register usage 绕晕的调参选手，应该都能从里面找到自己踩过的坑。
1. 从访存说起：寄存器 tiling 到底在优化什么
1.1 计算与访存的“剪刀差”，逼出了 tiling
先看最朴素的 GEMM 写法：三层循环里，每条语句都从全局内存拿 A[i][k] 和 B[k][j]，对每个输出元素做一次乘加。这个时候数据是“一次性”的——读进来算完就丢，下一个输出继续重新读。
把搬运和计算分开看，问题就来了。假设每个 FP32 计算需要读 4 字节数据，一台 FP32 算力 19.5 TFLOPS 的 A100，需要大约 78 TB/s 的片外带宽才能喂饱。而实际 HBM 带宽只有 2 TB/s 左右，差了一个数量级以上。这还只是每次只读一个操作数的情况，如果考虑读两个操作数，缺口会更大。
这就是大家常说的“内存墙”。tiling 的思路，就是打破“每条数据只用一次”的局面：把循环顺序换一下，让一个数据从 HBM 搬到寄存器之后，被多条指令反复使用。所以 tiling 的本质不是切块，而是通过切块把“访存次数 / 计算次数”这个比例压下来。
寄存器是硬件里最快的存储，但容量极小。一个线程手头上能用的寄存器数量，通常只有几十到几百个。如何在这么小的空间里，让每个数据被尽量多次地复用，其实就是寄存器 tiling 的核心命题。一个很直观的生活类比：把 HBM 想成大仓库，shared memory 是操作台，寄存器就是你手里能捧住的食材。每次从仓库往操作台搬一大块肉，再从操作台往手里抓一小把，手里的量决定你一次能同时处理几道菜。
1.2 寄存器 tiling 和 shared memory tiling 的分工
很多时候大家把“分块”统称为 tiling，但寄存器 tiling 是更靠里的一层。shared memory tiling 解决的是线程间数据共享：一个线程块把一块数据从全局内存搬进 shared memory，块内所有线程都能复用。寄存器 tiling 解决的是线程内数据复用：一个线程的寄存器里同时保有多少个输出累加器、多少个 A/B fragment，以及在 K 方向上展开多长。
两者通常是配对的。一个典型的 GEMM kernel，会把整块 tile 从 HBM 搬进 shared memory，然后再从 shared memory 搬进每个线程自己的寄存器。shared memory 决定“一个 block 一次吃多大”，寄存器 tiling 决定“一个线程一次能攒多少活”。如果把整块 tile 比作一桌菜，shared memory 就是那个可以反复取菜的转盘，寄存器就是你碗里同时堆着的几道菜。转盘上的菜再多，碗里放不下，一次也只能吃那么几口。
所以，寄存器 tiling 的尺寸并不只由“数据有多大”决定，而是由“每个线程手里能捧多少”决定。这也引出了后面最关键的问题：不同架构能“捧”的量，差别很大。
2. NVIDIA、AMD、CPU 与加速器：tile 的“形状”由谁决定
2.1 NVIDIA GPU：从外积寄存器 tile 到 Tensor Core 指令形状
在 Volta 之前的特斯拉架构上，没有 Tensor Core，寄存器 tiling 靠的就是 FMA 指令做外积。一个线程通常在寄存器里维护一个 4×4 甚至 8×4 的 C 累加块，循环内部每次取一个 a 片段和一个 b 片段，用外层积的方式同时更新所有累加器。这样做的好处是：同一个 a[i] 可以被同时用于计算多个不同的输出，b 也一样，数据复用率一下子提上来了。
Volta 之后有了 MMA 指令，事情变得更有意思。以 Ampere 上最常见的mma.m16n8k16为例，一个 warp 的 32 个线程共同完成一个 16×8 的输出块。注意这里的“16×8”不是拍脑袋定的，它跟线程数、寄存器宽度、硬件里的 FMA 阵列布局紧密相关。每个线程的 C fragment 是 4 个 fp32 累加器，A 和 B 的 fragment 按固定模式散布在 32 个线程里。
在实际写 kernel 的时候，你不需要手动去摆每个线程的 A/B fragment，编译器会帮你生成对应的寄存器排布。但这不意味着你可以完全不管形状。如果你把每线程的 C tile 设成 4×4，大概符合 m16n8k16 的粒度；如果设成 8×8，编译器很可能要做更多寄存器重排，或者退化成多次 mma 的拼接。Hopper 上的 wgmma 更进一步，由 4 个 warp 组成的 warpgroup 一起发射，输入形状可以更大，寄存器数据的排布由硬件描述符直接指定，灵活性更小，但效率更高。
这里有个很关键的点：寄存器 tiling 的形状本质上是“线程模型 + 指令集”共同决定的。你不能脱离硬件约束去谈“最优 tile 大小”。同一个 128×128 的 block tile，在 A100 上和 H100 上，最佳线程配置可能完全不同。
2.2 AMD CDNA：wavefront 64 带来的形状变化
AMD 的调度单位是 wavefront，一个 wavefront 是 64 个线程，比 NVIDIA 的 warp 大一倍。这个看似简单的差别，直接影响了寄存器 tiling 的形状。
CDNA 架构上的 MFMA 指令，常见形状是 16×16×16 和 32×8×16，正好由 64 个线程一起协作。以 16×16×16 为例，输出是 16×16 的矩阵，64 个线程每线程拿到 4 个 fp32 累加器，这和 NVIDIA m16n8k16 的每线程 4 个累加器很接近，但因为线程总数不同，fragment 在 lane 上的分布方式完全不一样。如果你只写过 CUDA，第一次看 AMD 的汇编时特别容易懵：明明都是做矩阵乘，为什么每线程持有的 A 片段、B 片段，跟 N 卡对不上？
这就是架构差异最真实的一面。做算子库的人，比如 ROCm 的 rocBLAS 或者 Triton 的 AMD 后端，会专门维护两套 fragment 排布逻辑，原因就在这里。好的一面是，像 Triton 这类编译器已经把这些差异封装掉了，你写一遍计算逻辑，它自动帮你映射到不同硬件。但如果你在做底层性能调优，理解“wavefront 64 导致 tile 更宽”这一点，能帮你少走很多弯路。
2.3 CPU：向量寄存器与 AMX 的整块 tile 寄存器
CPU 这边的情况更“复古”一些，但也更直观。AVX-512 一个向量寄存器能装 512 bit，也就是 16 个 fp32。CPU 上的寄存器 tiling，通常表现为：把内层循环 unroll 8 或 16 次，同时维护 8~16 个 zmm 向量寄存器，每个向量装的是同一行不同列的结果。没有 Tensor Core，全靠 FMA 指令一个一个算，靠的就是向量寄存器之间的数据复用。
Intel AMX 发布之后，x86 CPU 也有了真正的“tile 寄存器”。8 个 1KB 的 tile 寄存器，一条指令可以做 16×16 甚至更大的矩阵乘加。这个思路已经非常接近 GPU 的 wgmma 了：寄存器本身就是按矩阵形状组织的，编译器负责把 tile 映射到物理寄存器文件。不过 AMX 的使用限制更多，比如数据摆放要和 tile 寄存器形状严格对齐，同时 CPU 的访存延迟比 GPU 高很多，寄存器 tiling 的收益往往被主存带宽卡住。
下面这张表能帮你快速建立不同架构的尺寸感：
| 架构 | 典型指令/机制 | 线程/协作单位 | 典型单次累加块 | 每线程 C 累加器数 |
|---|---|---|---|---|
| NVIDIA Ampere | mma.m16n8k16 | 32 线程 warp | 16×8 | 4 个 fp32 |
| NVIDIA Hopper | wgmma | 128 线程 warpgroup | 64×N | 由描述符指定 |
| AMD CDNA | MFMA 16×16×16 | 64 线程 wavefront | 16×16 | 4 个 fp32 |
| x86 AVX-512 | FMA 展开 | 单线程 | 16 个输出 | 取决于 unroll |
| Intel AMX | TMMULT | 单线程 | 16×16 | 由 tile 寄存器布局决定 |
2.4 TPU/脉动阵列：把 tiling 熔进硬件的极端形态
说到极端，就必须提脉动阵列，比如 TPU 的 systolic array。在脉动阵列里，没有我们习惯的“寄存器文件 + 指令”，每个 PE 里的暂存器就是整条数据通路的一部分。矩阵乘法时，A 的数据从左边流进阵列，B 的数据从上面流进，部分和沿对角线一路累积。
在这个模型里，寄存器 tiling 的形状直接由阵列大小决定。硬件是 128×128 的阵列，单次就能累加 128×128 的一块，编译器能做的只是决定数据流进阵列的次序和分块粒度。作为软件开发者，你没法像在 GPU 上那样，自由决定“每个线程手里拿 4×4 还是 8×8”。好处是数据几乎零搬运开销，坏处是灵活性很低，一旦模型尺寸和硬件阵列不匹配，就需要 padding 或者拆分。
对比到这里，你应该能感受到一个光谱：CPU 最自由，寄存器 tiling 靠手动 unroll；GPU 介于中间，指令给了一个硬形状，但线程内还有选择空间；TPU 最受限，tiling 已经变成了物理形状。理解自己在光谱的哪一端，才能合理地设计 kernel。
3. 手把手设计一个 GEMM 的寄存器 tiling：参数怎么定
3.1 从 naive 循环到块状外积：先换循环次序
用一个具体的例子来走一遍。假设我们要算 C[M][N] = A[M][K] * B[K][N]，naive 版本是三层循环：
for i in range(M): for j in range(N): sum = 0 for k in range(K): sum += A[i][k] * B[k][j] C[i][j] = sum这里每个内层迭代只做一个乘加，A[i][k] 读一次、B[k][j] 读一次，数据完全不复用。要把数据复用提上来，第一步是让每个线程同时算多个输出。改成这样：
for i_tile in range(M / MT): for j_tile in range(N / NT): for k in range(K): load a_frag[MT] A[i_tile][k] into registers load b_frag[NT] B[k][j_tile] into registers for i in range(MT): for j in range(NT): C[i][j] += a_frag[i] * b_frag[j]这个过程就是把“每次只算一个输出”变成“每次同时算 MT×NT 个输出”。a_frag 和 b_frag 被复用 MT 次和 NT 次，访存次数从 MT*NT 次降到了 MT+NT 次。这就是寄存器 tiling 的外积形态，也是所有优化版本的基础。
落实到代码上，一个线程内部的累加器往往长这样：
float c[4][4] = {}; // 16 个寄存器，4x4 的 C 累加块 for (int k = 0; k < BK; ++k) { float a_reg[4] = load_a_from_shared(thread_row, k); float b_reg[4] = load_b_from_shared(k, thread_col); #pragma unroll for (int i = 0; i < 4; ++i) { for (int j = 0; j < 4; ++j) { c[i][j] += a_reg[i] * b_reg[j]; } } }这里c[4][4]就是寄存器 tiling 的具体体现。16 个 fp32 累加器，每个外层 k 的迭代里，用 4 个 a 和 4 个 b 同时更新全部 16 个输出。
如果你只是用共享内存分块，却没有在寄存器里同时维护多个累加器，性能会停在很低的水平。因为 shared memory 虽然比 HBM 快很多，但和寄存器还是有数量级的差距。只有把真正频繁更新的数据放到寄存器里，才能把计算流水线压满。
3.2 用“每线程寄存器预算”逆向推导 tile 大小
现在关键问题来了：寄存器 tiling 应该取多大？这里不能凭感觉，必须算账。
以 NVIDIA GPU 为例，每线程最多可用 255 个寄存器，每 SM 有 64K 个 32 位寄存器。假设我们想用一个 block 的 256 个线程（8 warps）处理一个 128×128 的输出块，每个线程平均要负责 128×128 / 256 = 64 个输出。如果排成 8×8，需要 64 个 fp32 累加器，也就是 64 个寄存器。
看起来 255 个寄存器足够放下，但账不能只算 C。A/B 的 fragment、shared memory 的地址指针、循环索引、常驻变量，这些都要占寄存器。一个 8×8 的 C 已经占了 64 个寄存器，再加上 A fragment（比如 8 个半精度或 8 个 fp32）、B fragment、指针和索引，很容易到 100 个以上。如果编译器为了循环展开再预留一些中间变量，突破 128 是分分钟的事。一旦超过 255，硬件只能把多余的数据 spill 到 local memory，local memory 实际是落在 HBM 上的，性能瞬间崩盘。
所以生产级 kernel 往往把 C 控制在 16~32 个寄存器，对应每线程 4×4 或 8×4 的输出块。下面这张表是个参考：
| 每线程 C tile | C 累加器占用 | 单次 k 迭代访问 A/B 次数 | 数据复用比 | 风险 |
|---|---|---|---|---|
| 4×4 | 16 个寄存器 | 8 次 | 4 | 低 |
| 8×4 | 32 个寄存器 | 12 次 | 约 5.3 | 中 |
| 8×8 | 64 个寄存器 | 16 次 | 8 | 高，容易溢出 |
权衡的核心是：C tile 越大，复用比越高，但寄存器压力也越大，同时指令级并行度不一定能跟上。很多时候，你把 C tile 从 8×8 降到 8×4，性能反而更高，因为 spill 消失了，每线程 occupancy 也上来了。
3.3 从伪代码到 Tensor Core：代码不是手摆寄存器
上面的手动外积版本，在 Tensor Core 出来之后已经显得有点“古典”了。如果你在 Ampere 之后的 GPU 上写 GEMM，真正高效的做法是直接发 MMA 指令，让硬件来算 16×8 或更大的一块，而不是自己在寄存器里循环 FMA。
但明白手动外积的过程依然很重要。因为无论是mma还是wgmma，它们本质上都在做同一件事：把若干线程的寄存器片段组成一个大矩阵乘法的输入和输出。只是数据排布的职责从你手里移交给了编译器和硬件描述符。用 CUDA 的 WMMA API 时，你会写类似wmma::fragment<matrix_a, 16, 16, 16, half, row_major>这样的类型，这个 fragment 在语义上就是“分布在 warp 内所有线程寄存器里的一块 A 矩阵”。你不需要手动告诉每个线程做哪几个乘加，硬件结构已经固定了。
对于做算子库的人，这时候真正需要关心的是两个问题：一个是 tile 的几何形状能不能和指令的固定形状对齐，另一个是 fragment 在寄存器里的排布会不会导致额外的 move 指令。例如用mma.m16n8k16时，C tile 的粒度是每个线程 4 个 fp32，A/B 的 fragment 需要在 warp 内做固定排列。如果你上层用了一个让编译器很难生成高效排布的 tile 形状，汇编里就会出现一大堆mov指令，白白浪费吞吐。
所以我的建议是：先用手动外积理解原理，再用 Tensor Core 指令拿性能。不要跳过原理直接抄 WMMA 代码，否则遇到 fragment 排布问题，你根本不知道编译器在对你的 kernel 做什么操作。
4. 寄存器 tiling 的实战排查与常见坑
4.1 寄存器溢出：你以为在读寄存器，其实在读 local memory
我在实际调 kernel 时，遇到最多的坑就是寄存器溢出。典型症状：代码感觉写得很规整，但性能比预期低 20% 到 50%，甚至更多。打开 Nsight Compute 一看，local memory 的访问次数吓人，寄存器使用量已经接近 255，甚至编译日志里直接出现spill stores和spill loads。
寄存器溢出的本质是：每个线程想用的寄存器超过了硬件极限，编译器只能把一部分变量挪到 local memory。而 local memory 在物理上就是 HBM，读写它意味着每个线程都在偷偷访问全局内存。这时候你的寄存器 tiling 再漂亮，也等于白搭，因为数据反复在 HBM 和寄存器之间倒腾。
排查方法，NVIDIA 这边最简单的是编译时加--ptxas-options=-v，直接看寄存器用量和 spill 情况。如果是运行期，用 Nsight Compute 看Register Usage和Local Memory两个指标。一旦确认 spill，优先按这个顺序调：把每线程 C tile 变小（比如从 8×8 降到 8×4）、减少内层循环的 unroll 倍数、用__launch_bounds__或maxrregcount=128强制编译器收敛寄存器用量。AMD 那边对应看 vgpr 占用，用 rocprof 也能看到类似的信息。
有一种情况很容易被忽略：不是 C tile 太大，而是编译器为了生成向量化访存，把 shared memory 里连续的一段数据一次性加载到多个寄存器里，导致峰值寄存器用量暴增。这时候你可以手动把共享内存读取拆成几次独立的 load，让编译器在每次 load 后尽快使用，降低同时存活的寄存器数量。
4.2 tile 形状与指令形状不匹配，Tensor Core 根本没用到
另一个反直觉的坑是：你觉得自己用了 WMMA 或 MMA，实际上代码编译出来跑的却是一堆 FMA，Tensor Core 一个都没参与。原因通常是 tile 形状和指令要求对不上，导致编译器只能退化为标量计算。
例如，mma.m16n8k16需要 C 累加块至少是 16×8 的粒度。如果你把一个 warp 的线程分工设计成每个线程独立算一个 4×4 的小块，而没有在 warp 层面凑出 16×8 的协作关系，编译器就很难把这几条独立的计算合并成一条 MMA，只能老老实实用 FMA 去算。结果就是，Tensor Core 单元空转，SM 的普通 FMA 流水线忙死，性能只比纯手动外积好一点点。
遇到这种情况，建议直接在汇编或者 SASS 层面确认 MMA 指令是否出现。Nsight Compute 的指令统计里可以看到mma类指令的数量。如果比预期少很多，就要回头检查 tile 形状和线程协作方式是否和指令形状对齐。还有一个隐蔽问题：Hopper 的 wgmma 对数据对齐要求更高，A 和 B 的起始地址经常要求 16 字节对齐，你要是只做 4 字节对齐，编译不出错，但性能会差一截。
4.3 不要忘了 shared memory 这条腿
很多初学者把注意力全放在寄存器 tiling 上，结果忘了数据要到寄存器，必须先经过一层 global、再经过 shared。如果 shared memory 的加载模式乱七八糟，每个线程读的地址不连续，或者 bank conflict 频繁，即便寄存器 tiling 做得完美，HBM 带宽照样跑不满。
正确做法是把 tiling 当成一个三级流水来设计：global→shared→register→计算。每一层都有它自己的“颗粒度”。shared memory 层解决线程间的数据复用，寄存器层解决线程内的数据复用，两者缺一不可。我个人的调试顺序是这样的：
- 先保证 global memory 合并访问，用一个最简单的 shared memory tile 建立一个基线；
- 用 profiling 工具看当前是 memory-bound 还是 compute-bound；
- 如果已经是 memory-bound，优化 shared memory 的加载和 bank conflict；
- 如果是 compute-bound，再逐步增大寄存器 tiling，观察寄存器使用量和 spill 变化；
- 每改一版，只动一个变量，对比 profiling 数据，不要一次改一堆。
这个顺序能帮你快速定位瓶颈到底在哪一层。很多人一开始就上大 C tile，发现性能没提升，其实是 shared memory 加载阶段早就把带宽占满了，寄存器层再怎么优化都是白费力。
4.4 常见错误速查表
| 症状 | 可能原因 | 排查手段 |
|---|---|---|
| 性能比预期低，local memory 访问多 | 寄存器溢出，C tile 过大或循环展开过猛 | 看编译日志，查 spill，缩减 C tile |
| compute-bound 但 Tensor Core 利用率低 | tile 形状与 mma/wgmma 形状不匹配，退化成了 FMA | 看 SASS 里 mma 指令数量，调整线程协作模式 |
| 带宽上不去，但寄存器 tiling 看着很合理 | shared memory 层 bank conflict 或非合并访问 | 用 profiler 看 shared memory 吞吐，修加载模式 |
| 改了 tile 大小性能没变化 | 可能根本不是 compute-bound，瓶颈在数据搬运 | 先判断 memory-bound 还是 compute-bound 再动手 |
最后分享一个我自己的习惯：拿到一个新架构，第一件事不是去查理论峰值，而是写一个最简单的 64×64 GEMM，在汇编里数一数每个线程手上同时维护着多少个累加器。因为寄存器 tiling 的“长相”会直接告诉你这个架构的真实脾气：它能塞下多大的 tile，它的 fragment 怎么在 lane 间分布，它的指令对形状有多挑剔。等踩过几次寄存器溢出的坑之后，你对 tile 尺寸的直觉就会变得非常准，再看到别人 kernel 里的128×128 block、8×8 per-thread tile，脑子里马上能浮现出每个线程的寄存器里大概是什么状况。这种直觉，就是调优经验最值钱的部分。
