做AI Infra这几年来，我面试别人或者被人追问的时候，只要话题绕到算子优化，最后几乎都会落在寄存器 tiling 这个点上。它不像共享内存 tiling 那样在教材里占着完整章节，也没有太多现成模板可以抄，但偏偏是它决定了 GEMM、FlashAttention、MoE 里那些核心 kernel 的性能能到哪一格。
寄存器 tiling，简单说就是把计算要用的数据切块之后，直接塞进寄存器文件里，让同一个数据在寄存器里被反复拉出来乘加，而不是每一轮都从共享内存甚至全局内存里重新捞。它保证计算单元在每一个时钟周期都能以整块硬件能给的最高带宽拿到操作数，也是 Roofline 模型那张图里“ridge point”真正被踩实的地方。
这篇文章是“AI Infra 每日一问”系列第 6 天。我想把寄存器 tiling 在不同硬件架构上的形态摊开讲一遍：NVIDIA 的 SIMT 线程模型里它怎么组织，AMD CDNA 的 wavefront 里它长什么样，昇腾这类 AI ASIC 的 Cube 单元里它变成了什么，再到 TPU 这种脉动阵列为什么干脆绕开了这条路。适合看的人：做算子库或者 kernel 的开发者、推理引擎里做图优化和内存规划的同学、想在 Triton 或者 CUDA 里把 GEMM 跑得更快的人。如果你只是听过 tiling 这个词，想彻底搞明白寄存器这一层为什么重要，这篇文章也能当一份入门地图来看。
1. 先把概念说透：寄存器 tiling 到底在 tile 什么
1.1 三层 tiling：block、shared memory、register 各切一刀
先拿 GEMM 举例。一个 M=4096、N=4096、K=4096 的矩阵乘，数据量在几十 MB 量级，任何片上存储都放不下。所以大家常说的 tiling 实际上分了三层，每一层解决的是不同问题。
第一层是线程块级 tiling，把输出矩阵切成若干 block tile 分给不同的 SM，解决的是并行度和任务分配。第二层是共享内存级 tiling，在 block 内部把 A、B 的切片搬进 shared memory，解决的是全局内存带宽太慢、数据重复从 HBM 读的问题。第三层才是寄存器级 tiling：共享内存里的数据再按每个线程应该负责的输出小块，切到线程私有的寄存器里，然后做乘加。
我把这三层整理成一张表，方便对照看：
| 层级 | 切分单位 | 核心解决 | 典型粒度 |
|---|---|---|---|
| block tiling | 线程块 / SM | 并行度、任务分配 | 128x128 |
| shared memory tiling | block 内所有线程 | 全局内存带宽、跨线程复用 | 128x32 |
| register tiling | 单个线程 | 片上最高带宽、零额外延迟复用 | 8x8 / 16x8 |
寄存器 tiling 关心的就是表里第三行。它在整个优化链路里处在最靠近计算单元的位置，往上一层的数据搬运可以被它掩盖，往下一层的计算能力能不能吃饱也要靠它。
1.2 寄存器凭什么能做 tiling 的主力
用个不太严谨但好懂的生活类比：寄存器就像工位上手边那排工具，共享内存是工位旁边的抽屉，全局内存是整个楼层尽头的仓库。你写 report 的时候，最称手的思路一定是把最常用的资料放在手边，而不是每翻一次就跑到仓库去搬一趟。寄存器 tiling 就是把这个策略用到极致：数据进寄存器一次，留在那里被多条计算指令反复使用，直到这个 tile 算完。
硬件层面，寄存器的地位更特殊。以 A100 为例，每个 SM 的寄存器文件大概 256 KB，分成 4 个处理块，每个处理块内部每个时钟周期都要能同时喂给成千上万个操作数的读取。这种读取带宽是 shared memory 和 L1 完全够不着的，更别提 HBM 了。所以对于计算密度很高的算子，有没有把数据正确放进寄存器，直接决定运算单元每个周期是吃饱还是在干等。
再算一笔账，你会更直观地感受到寄存器复用的威力。GEMM 里每算一个输出元素，朴素做法要从 K 维度读 A 的一行和 B 的一列，总共 2K 个数。如果每个线程持有一个 8x8 的输出 tile，那么 A 的行片段可以被 8 个输出列复用，B 的列片段可以被 8 个输出行复用，需要的访存次数从 2K 降到 K/8 + K/8，也就是 K/4，直接省掉 8 倍。这个复用就发生在寄存器里，不占共享内存的带宽，也不占 LSU 的发射槽位。
1.3 寄存器 tiling 的硬边界
寄存器 tiling 不是无限扩张的。寄存器文件总量有限，分给每个线程多了，能同时驻留的线程数就少，占用率就掉。这里面有三条硬约束，做 kernel 的人必须心里有数。
第一，单线程寄存器数量有上限。NVIDIA 上每个线程最多 255 个 32 位寄存器，超过之后编译器会把多余变量“溢出”到 local memory，其实就是在全局内存里划一块私有区域，性能和直接访问全局内存差不多，属于断崖式下跌。第二，寄存器是线程私有的，线程之间不能直接访问对方的寄存器。跨线程的数据复用只能靠 warp shuffle 或者共享内存中转，所以寄存器 tiling 只解决线程内部的复用。第三，寄存器 tiling 尺寸和占用率永远在打架。A100 一个 SM 最多驻留 2048 个线程，如果每个线程用 128 个寄存器，最多同时放 1024 线程，占用率 50%；要是某个 kernel 把每线程寄存器顶到 256 个，那就只能放 512 线程，占用率掉到 25%。
这就引出一个关键判断：寄存器 tile 不是越大越好，而是要在“单线程计算密度”和“足够多并发的线程”之间找平衡点。后面第 2 章我会用具体例子说明这个平衡怎么找。
2. NVIDIA 系：寄存器 tiling 的“标准答案”
2.1 SIMT 语义下的寄存器 tiling 形态
NVIDIA GPU 是 SIMT 模型，单指令多线程。一个 warp 有 32 个线程，每个线程都有自己独立的寄存器，但同一时刻这 32 个线程执行同一条指令。这种模型下，寄存器 tiling 的天然单位是“每个线程一小块输出”。
经典 GEMM 排布是这样的：warp 里 32 个线程组织成 4x8 的线程阵列，每个线程负责一个 8x8 的 C tile，于是这个 warp 整体覆盖 32x64 的输出块。A 和 B 的片段也按线程各自加载到寄存器里。一个线程的寄存器预算大致是这样：
- 8x8 的 C 累加器：64 个 float 寄存器
- A fragment：8 个寄存器
- B fragment：8 个寄存器
- 基地址、循环计数、临时变量：大约 20 到 40 个
合起来每线程大约 100 到 120 个寄存器，正好落在 128 以内。如果改成每个线程负责 16x16 的 C tile，光累加器就要 256 个寄存器，直接爆掉单线程 255 的上限，必然 spill，性能惨不忍睹。所以 8x8 或者 16x8 这种尺寸在 Ampere 之前几乎是黄金选择，不是拍脑袋定出来的。
伪码层面，这种 tiling 写出来大致是下面这个样子：
// 每个线程负责 8x8 的 C 子块 float c[8][8] = {0}; // 在 K 维上循环，每次取 A 的 8 行片段和 B 的 8 列片段 for (int k0 = 0; k0 < K; k0 += 8) { float a_frag[8]; // 当前线程要用的 A 行片段 float b_frag[8]; // 当前线程要用的 B 列片段 load_a_fragment(a_frag, ...); load_b_fragment(b_frag, ...); // 寄存器里的 8x8x8 乘加 for (int i = 0; i < 8; i++) { for (int j = 0; j < 8; j++) { for (int kk = 0; kk < 8; kk++) { c[i][j] += a_frag[kk] * b_frag[kk]; } } } }注意这只是一个概念示意。真实 CUTLASS 这类库里，fragment 的布局、循环展开、指令调度都是精心安排的，因为编译器默认生成的代码往往达不到最佳状态。但核心思想就是这一段：把数据切成小块，锁进寄存器，拼命乘加。
2.2 从手工 fragment 到 WMMA 指令
从 Turing 架构开始，NVIDIA 引入了硬件矩阵指令，也就是 HMMA/WMMA。这类指令把“寄存器 tiling 的乘加”直接硬件化了，一条指令完成一个 16x8x16 或者 16x16x16 的矩阵乘加，不需要你再写三层循环。
用 WMMA API 写的时候，A、B、C 都以 fragment 对象出现，每个 fragment 内部数据在 warp 内 32 个线程里的分布是硬件规定死的。比如 16x16x16 的 WMMA，每个线程持有的 A fragment 是 8 个寄存器，B fragment 是 8 个，C fragment 是 32 个。这个布局不是程序员能随便改的，你只能确保在调用 load_matrix_sync 之前，数据在共享内存或者全局内存里的排布满足指令要求。
这里有一个新手最容易踩的坑：WMMA 的 fragment 布局不是直观的二维数组字节流。你以为是“线程 0 拿矩阵第一行前 8 个数”，实际上硬件为了减少 bank conflict 和 shuffle 次数，把数据按一种交叉分布放到 32 个线程里。所以千万别手工去推断 fragment 元素的位置，老老实实走 API，要靠内存布局对齐，用不对就等着访问错数据。
到了 Hopper 架构，WGMMA 更进一步。它不再是“一个 warp 一条指令”，而是让 warpgroup（4 个 warp，128 个线程）一起加载更大的寄存器矩阵块，累加器可以到 64x64 甚至更大。这种设计本质上是把寄存器 tiling 的规模从线程级提升到 warpgroup 级，寄存器文件的整体带宽被更高效地利用了。
2.3 算一笔寄存器账
再实算一笔账，看看为什么 16x16x16 的 WMMA 是 Ampere 时代的甜点尺寸。每个线程的数据类寄存器大概 48 个（A 的 8 + B 的 8 + C 的 32），加上地址计算、load 需要的 temporary 寄存器，编译器一般会报到 80 到 120 个寄存器。这个数字既能保证一个 SM 里还有足够多的并发 warp 去隐藏访存延迟，又不会让单线程太少、计算密度上不去。
如果你做的是 double buffering，也就是在 K 维循环里同时维护两套 fragment——一套正在被矩阵指令消费，另一套正在从共享内存加载——那寄存器数会直接翻倍。这时候就要权衡了：流水线掩盖的访存延迟，到底值不值牺牲一半占用率。我在 A100 上调 FlashAttention 的时候就反复试过这个临界点，最终发现 96 到 160 个寄存器之间通常能找到收益拐点，超过 200 寄存器基本就会因为占用率太低反而变慢。
这类经验在工程上非常值钱。当年老一点的优化文章喜欢鼓吹“寄存器越多越能提升单线程 ILP”，但在现代 GPU 上，寄存器文件总量是固定的，单线程分得多，并发线程就少。最好的策略永远是拿着编译器报告做实验，而不是凭空拍一个数字。
3. 别家架构：寄存器 tiling 的更多姿势
3.1 AMD CDNA：SGPR/VGPR 与 MFMA 指令
AMD 的 CDNA 架构，从 MI100 到 MI300，走的也是类似 SIMT 的路子，但有两个明显区别：第一，执行粒度是 wavefront，也就是 64 个线程一组，而不是 NVIDIA 的 32 线程 warp；第二，寄存器文件分成 SGPR（标量寄存器）和 VGPR（向量寄存器）两类，SGPR 在整个 wavefront 里共享，VGPR 才是每个线程私有的。
这种划分带来的直接影响是，寄存器 tiling 的时候，你能把“每个线程都一样”的地址、偏移、循环计数放进 SGPR，把真正随线程变化的矩阵元素放进 VGPR。这比 NVIDIA 把所有东西都堆在通用寄存器里要省不少空间，也让编译器的寄存器分配更灵活。
AMD 的矩阵指令叫 MFMA，比如v_mfma_f32_16x16x16_f32。它的参数布局和 NVIDIA 的 WMMA 完全不同，因为 64 个线程要分四组分别覆盖矩阵的行或列方向。我在把一些 CUDA kernel 往 ROCm 上移植的时候，发现一个非常实际的问题：NVIDIA 上最优的 16x8x16 指令形状，在 AMD 上并不会自动变成同样高效的 MFMA 形状。因为 wavefront 是 64 线程，你往往要把 tile 的边长按 16 的倍数来组织，才能让 MFMA 的 fragment 分布正好落满整个 wavefront。这块只能靠对着汇编和 Matrix Core 的文档一点一点试，没有捷径。
3.2 昇腾达芬奇架构：Cube 单元的类寄存器 Buffer
昇腾的达芬奇架构走的是另一条路。它的 AICore 里有专门做矩阵乘的 Cube 单元，Cube 的输入不是我们熟悉的线程私有寄存器，而是 L0 Buffer，具体分为 L0A、L0B 和 L0C。L0A 存 A 矩阵的 tile，L0B 存 B 矩阵的 tile，L0C 既是累加器又是输出缓冲。
从功能上说，L0 Buffer 扮演的就是“寄存器化存储”的角色：数据在进入 Cube 之前，必须已经按固定形状摆放在 L0 里。但跟 CUDA 有一个本质区别：程序员几乎不手动去控制每个“线程”对应哪几个寄存器。Tiling 策略主要由编译器（比如 AKG）在编译期生成，它会根据算子的 shape 自动推导 Cube 的 tiling 方案、L0 和 L1 之间的搬入搬出节奏。
实际调昇腾算子的时候，你会感觉到“寄存器 tiling”这个词的存在感被编译器包了一层。你要关注的更多是 L0/L1 容量够不够、双 buffer 有没有铺满、Cube 在切 K 维时是不是一次吞一个整数个 L0 tile。这种体验和 CUDA 那种所有布局都攥在自己手里的感觉差异很大，但本质诉求是一致的：让 Cube 单元的每个时钟周期都能从最近的存储拿到它要的操作数。
3.3 脉动阵列与本地 SRAM 的“非典型”寄存器 tiling
TPU 的脉动阵列彻底绕开了“线程 + 寄存器”这套体系。在脉动阵列里，数据一旦从 Unified Buffer 喂进阵列，就在 PE 之间一个节拍一个节拍地流动，权重驻留在每个 PE 边缘的存储里，累加器放在 PE 内部。这里没有传统意义的线程私有寄存器 tiling，真正实现数据复用的是“数据流”，数据流本身就是复用机制。你不需要在寄存器里保留一整块 A 的行片段，因为相邻 PE 会在后续节拍里把这个数据往下传。
Graphcore 的 IPU 又是一种极端。它没有硬件缓存，每个核心的本地 SRAM 是显式管理的，程序员得自己在 SRAM 里规划数据摆放。在这种架构上，tiling 的那点“缓存”空间全都在 SRAM 里，没有隐藏的寄存器层次可以靠编译器兜底，几乎等于把寄存器 tiling 的思想放大成了“SRAM tiling”。
各家架构放到一张表里对比，形态差异会非常明显：
| 架构 | 最小计算粒度 | 寄存器 / buffer 形态 | tiling 主导者 | 典型机制 |
|---|---|---|---|---|
| NVIDIA | 32 线程 warp | 线程私有寄存器 + fragment 布局 | 程序员 + 编译器 | mma.sync / WGMMA |
| AMD CDNA | 64 线程 wavefront | VGPR + SGPR + MFMA fragment | 程序员 + 编译器 | v_mfma |
| 昇腾 | Cube 单元 | L0A / L0B / L0C buffer | 编译器（AKG） | Cube 指令 |
| TPU | 脉动阵列 | PE 内累加器 + 权重驻留 | 编译器 / XLA | 脉动数据流 |
| Graphcore IPU | 单核 | 本地 SRAM 显式管理 | 程序员 / 编译器 | BSP 同步执行 |
看完这张表，你会发现“寄存器 tiling”在不同架构上长得非常不一样，但背后的问题始终只有一个：计算单元下一拍要的操作数，能不能从最近、最快的那层存储直接拿到。
4. 实战：AI Kernel 里的寄存器 tiling 怎么落地
4.1 先画一张寄存器 tiling 草图再写代码
我自己的习惯是，写任何计算密集 kernel 之前，先不碰代码，而是在纸上画一张“数据归属图”：哪个 warp 负责输出矩阵的哪一块，warp 里哪个线程负责哪一小块，每个线程手里同时持有哪些 A、B、C 的 fragment。这个习惯是从 CUTLASS 的文档里学来的，后来发现极其好用，尤其是在跨架构移植的时候。
画图时按这个顺序来：
- 确定输出 tile 的形状，也就是线程块和 warp 的排布。
- 决定单个线程持有多大的寄存器 tile。
- 根据寄存器预算估算每线程大概需要多少寄存器，看看是否和硬件上限匹配。
- 确定 K 维循环里一次处理多大片段，是否要做多级流水线。
- 最后才是写第一版代码。
这个流程里最容易被忽略的是第二步到第三步的衔接。很多人兴致勃勃画了一个 32x32 的大 tile，最后编译一看每线程寄存器冲到 300 多，spill 成灾，性能反而不如十几行朴素写法。记住：草图上看着漂亮没用，寄存器预算这关过不了，所有设计都是空中楼阁。
4.2 Triton、CUTLASS、oneDNN 里怎么体现
不是所有场景都要手写 CUDA。实际工程里，Triton 这类 DSL 已经吸引了很多注意力。你写 Triton 的时候确实不会被要求手动分配寄存器，但这不意味着寄存器 tiling 消失了，它只是被编译器接管了。
你在 Triton 里选的BLOCK_SIZE和num_warps，其实就是在给编译器下达不同的寄存器 tiling 策略。比如一个 64x64 的 block 配上 4 个 warp，编译器倾向于让每个线程持有 16x8 左右的累加器；你把 block 放大到 128x128，每个线程要扛的 fragment 就明显变大，寄存器压力也随之上升。有一次我把一个 Triton GEMM 的num_warps从 4 改成 8，寄存器数从每线程 96 涨到 168，性能不升反降，就是因为并发线程变少后访存延迟盖不住了。
CUTLASS 则是把寄存器 tiling 完全显式化。你通过模板参数指定ThreadblockShape、WarpShape、InstructionShape，每个线程持有的 fragment 就是由这几个参数计算出来的。我建议想做 gemm kernel 深入优化的人，一定先把 CUTLASS 里 fragment 布局的推导逻辑读一遍，那套东西本身就是一部寄存器 tiling 教科书。oneDNN 这边的思路也类似，只是它更多面向推理场景，会把融合算子直接按寄存器布局来规划，尽量避免共享内存里的中转拷贝。
4.3 性能分析里的寄存器关键指标
代码跑起来之后，不能只看端到端时间就完事。我通常这几个指标必看：
- 编译日志里的
registers per thread，这个数字直接反映你的寄存器 tiling 设计得激进不激进。 - 有没有
spill loads/stores，只要有，基本说明寄存器预算爆了。 - 看
local memory的大小，它侧面反映 spill 的严重程度。 - 用 ncu 看 kernel 的 issue 利用率，确认数学指令是不是真的在连续发射。
我举个真实例子。之前在 A100 上调一个 attention 类算子，第一版每线程寄存器用到了 256，编译报告显示有大量 spill。我把寄存器 tiling 从每线程 256 收缩到 168，spill 归零，同时 occupancy 从 25% 回到 37.5%，最终端到端性能反而快了 12%。这个案例很典型：你多占的那些寄存器并没有换来同等的计算密度，反而因为 spill 频繁跑内存，把一切全赔回去了。
提示：用
nvcc --ptxas-options=-v编译时就能看到寄存器数和 spill 情况，不用等运行完；如果在 Triton 里开发，也可以用环境变量打开编译诊断信息。
5. 踩坑清单与排查技巧实录
5.1 RegSpill：寄存器用爆了怎么查
寄存器溢出是寄存器 tiling 最常踩的坑，特征也很明显：kernel 运行时间突然比预期慢一个数量级，ncu或者编译日志里出现大量local memory访问。
排查手法三步走。第一步，确认编译器报告里的每线程寄存器数和 spill 读写次数；第二步，定位是哪个 kernel 爆的，很多时候一个项目里不同 kernel 的寄存器压力完全不一样，你得逐个看编译日志；第三步，考虑用__launch_bounds__(maxThreadsPerBlock, minBlocksPerSM)做上限约束，强制编译器在寄存器分配上收敛，避免它为了激进的指令级并行把寄存器数量推到离谱的地步。
有个小经验：__launch_bounds__不只是限制并发线程数，它其实是在告诉编译器“我要保住多少块 SM 占用率”，编译器会按这个目标调整寄存器分配。你如果目标 block 数是 8，编译器会把每线程寄存器压到刚好能塞 8 个 block 的水平；如果你不设，编译器可能为了 ILP 一路放飞，最后 spill。
5.2 寄存器 tiling 与 Bank Conflict 的真相
很多人会把寄存器 tiling 和共享内存的 bank conflict 混在一起聊，这里澄清一下：寄存器文件本身没有 bank conflict 的概念，至少对程序员不可见。真正会踩 bank conflict 的是你从共享内存往寄存器搬 fragment 的那一步。
比如你有一个 8x8 的 float tile 要从共享内存加载到寄存器，如果共享内存里这 8 个 float 恰好落在同一个 bank 上，那么 32 个线程同时访问时会发生冲突，反而把加载周期拉长。NVIDIA 的 WMMA 指令在设计 fragment 内存布局时已经考虑过这个问题，它会要求你按特定 stride 把数据铺在共享内存里，确保访问不冲突。所以最常见的坑是：你用共享内存存矩阵时完全按行优先铺，然后发现 WMMA 的load_matrix_sync性能很差。这通常不是指令本身的问题，而是矩阵在共享内存里的 tread 排列没遵守硬件偏好。
5.3 不同架构迁移时的“水土不服”
做多架构算子库的同学注意了：把 NVIDIA 上最优的 16x8x16 寄存器 tiling 原封不动搬到 AMD 会水土不服，因为在 64 线程的 wavefront 下，MFMA 的 fragment 分布逻辑变了，16x8 这种形状要么填充不满，要么多出多余的 shuffle 指令。昇腾上更是直接换了一套 L0 Buffer 编程模型，你在 CUDA 里手工安排的 fragment 布局完全没法映射过去。
我的建议是，在库的抽象层定义自己的“tile 描述”，比如用结构体描述每个线程/处理单元应该持有多少 A、B、C 元素，然后在每个后端的代码生成阶段各自翻译成对应架构的寄存器布局。这个抽象层初期写着麻烦，但只要有第二个后端，回报立刻兑现。
5.4 一条快速自查清单
最后给一张我自己常用的 registry tiling 自查表，每次调优半信半疑的时候就拿出来逐项过一遍：
| 检查项 | 怎么看 | 健康标准 |
|---|---|---|
| 每线程寄存器数 | ptxas / 编译器日志 | 40 到 168 之间 |
| 是否发生 spill | local memory 大小、spill 读写次数 | 0 |
| 实际 occupancy | ncu | 和预期设计一致 |
| 数学指令连续发射情况 | issue 利用率 | 越高越好，且无明显 stall |
| fragment 加载指令比例 | 加载指令与矩阵指令比值 | 加载明显少于计算 |
按这张表走一遍，基本能覆盖 80% 的寄存器 tiling 性能问题。
最后再分享一点个人心得。做了这么多架构的矩阵库和算子之后，我最大的体会是：不同硬件对“寄存器”的定义、容量的理解、以及暴露给程序员的控制方式差得非常多，但大家在做的事情本质上是同一件——想办法让计算单元在下一拍开始之前，最快拿到它要的操作数。谁能把这件事做到位，谁写的 kernel 就离这台机器的极限更近。每次写新算子的第一版，我都会先画一张寄存器 tiling 草图，哪怕后来编译器帮我改掉一半，这张草图也是我理解这台架构的起点。
