寄存器tiling跨架构深度解析：不同GPU/CPU的硬件约束与优化差异

做 AI Infra 或者手写高性能算子这些年，我一直觉得寄存器 tiling（register tiling）是最能体现“架构感知”的一个优化手段。同样一段 GEMM 逻辑，在 A100 上把 block 尺寸调成 128x128 可能带来 15% 的加速，但把这套参数原封不动搬到 AMD MI250 上，性能可能直接倒退回一半。很多人第一反应是“编译器适配没做好”，但真正的原因往往更底层——不同架构的寄存器文件、矩阵指令集和调度方式，对 tiling 的物理约束完全不一样。这一篇就把这个问题掰开揉碎讲清楚：寄存器 tiling 在不同架构上到底长什么样，以及为什么它注定无法“一套代码吃到老”。

1. 寄存器 tiling 到底在解决什么问题

1.1 从“内存墙”说起：为什么算子层一定要做数据分块

先回到最原始的矛盾。GPU 和 CPU 的算力提升速度，远快于内存带宽和访存延迟的改善。以一个常见的 FP32 GEMM 为例，如果直接让每个线程从全局内存读一个 A 元素和一个 B 元素做乘加，把 C 写回，那你实际跑的会是内存带宽瓶颈，ALU 大部分时间在空转。要压满 FP32/FP16/TF32 的峰值算力，唯一出路是提高数据的复用率——每个从全局内存读进来的数据，尽量被计算单元多利用几次。

tiling 做的事情，本质就是把数据切块后搬进离计算单元更近的存储层级，用“片上带宽”抵消“片外带宽”的不足。通常大家会分三级：全局内存 -> 共享内存/一级缓存 -> 寄存器。其中寄存器是离 ALU 最近、带宽最大、延迟最低的一层，也是最“贵”的一层，数量少、容量小、按线程私有。芯片设计者把所有编译期能确定的数据分配都砸在寄存器文件上，就是为了让计算单元可以连续多个周期“自给自足”，不被打断。

这里有个容易混淆的点：很多人把 tiling 等同于共享内存 tiling（shared memory tiling），其实不够全面。共享内存负责的是线程块内部的数据交换与合并访存，而寄存器 tiling 负责的是单个线程在循环体内的数据布局和复用。两者经常配合出现，但寄存器 tiling 更极端——它直接决定了每个线程在一条指令发射前手里握着多少数据和累加器。

1.2 寄存器级 vs 共享内存级 tiling：两级存储体系的分工

我用一个生活化类比解释这两级分工。假设你在流水线上组装玩具，共享内存是一个小型零件柜，就在工位旁边，从零件柜拿零件比去仓库（全局内存）快得多；但更快的办法是直接把接下来几次组装要用的零件全放在手心里（寄存器）。手心的空间有限，但你放对了零件，整个组装过程就可以完全不动脚、不转身。

放到实际 GEMM 里，分工是这样：

• 全局内存装入共享内存：解决“跨线程合并访存”和“数据复用”，一个 tile 的数据被 block 内几百个线程分享；
• 共享内存装入寄存器：解决“单线程内的累加器复用”，A tile 的一行和 B tile 的一列在寄存器里完成多次乘加，不重复回访共享内存；
• 寄存器里的累加器：在整个 inner k 循环期间一直驻留不动，每次只是读入新的 a、b 向量做乘加。

所以共享内存 tiling 做得再好，如果寄存器 tiling 没设计好，累加器频繁溢出到 local memory，性能照样会崩。寄存器 tiling 才是把算力榨干的最后一公里。

1.3 一个最小例子：每线程到底要多少寄存器

套一个具体的数字来感受。假设 block 尺寸是 128x128，每个线程负责 8x8 个 C 元素的累加器，每个元素是 FP32，那么每个线程光累加器就需要 64 个寄存器。再来，inner 循环中每次还需要从共享内存搬运 A 的 8 个元素和 B 的 8 个元素到寄存器（如果做向量化加载可能更多），这就又占了 16 个。一个线程跑到 80 个 + 的状态。

NVIDIA 的编译器通常把单线程寄存器上限控制在 128 或 255 个（不同架构不同代际有差异），AMD 的 CDNA 架构也有类似的寄存器堆限制。一个 8x8 的累加器方案还算轻松，但如果线程负责 16x16 个 C 元素，累加器就要 256 个寄存器，直接撞线。这就是为什么寄存器 tiling 的尺寸上限不是拍脑袋定的，而是被硬件寄存器文件大小、每线程寄存器配额、以及指令集的 operand 承载方式共同锁死的。

2. NVIDIA、AMD、Intel 和 ARM，寄存器 tiling 各长什么样

2.1 NVIDIA CUDA：寄存器片段（fragments）与 mma 指令的一一对应

NVIDIA 是大多数 AI Infra 工程师最熟悉的平台。在 Turing 之后的架构上，真正干活的矩阵指令是mma.sync.aligned.m16n8k8这类指令（Volta 引入，Ampere 和 Hopper 都有演进）。它要求 A、B、C、D 矩阵的数据以特定的“寄存器片段”格式分布在 warp 内的线程里。比如m16n8k16（Tensor Core 使用），每个线程手里持有的不是矩阵里的连续行或列，而是按硬件规定的一个交错布局。

这也是我第一次手写 mma 时最崩溃的地方——寄存器 tiling 在这里不只是“把数据放到寄存器数组中”，而是“把数据按硬件指定的 fragment 布局放到寄存器数组中”。A 矩阵的 4x8 片段（k=16 时）分布在 32 个线程中，每个线程持有 4 个元素；B 矩阵的 8x8 片段每个线程持有 2 个元素；C/D 累加器的 4x8 片段每个线程持有 4 个元素。你在写 CUTLASS 或手写 PTX 时的所有数据搬运和 swizzle 操作，都是在往这个特定布局里塞数据。

以 Hopper/H100 上的 WGMMA（warpgroup MMA）为例，它更进一步——一条指令直接让 warpgroup 从共享内存加载大块 A、B 数据到 Tensor Core 内部的寄存器/数据路径，并对每个线程的寄存器分配做了非常强的假设。编程时就不太能自由定义“我手上哪些 cycle 持有哪些元素”，而是要顺着硬件的意图走。

还有一个 NVIDIA 上特有的 tiling 行为是c[m][n]累加器的“固定驻留”。比如写一个典型的 SGEMM kernel：

// 每线程 4x4 的 C tile float c[4][4] = {0.0f}; for (int k = 0; k < K; k += BK) { float a[4] = load_A_from_shared(k); float b[4] = load_B_from_shared(k); for (int i = 0; i < 4; i++) for (int j = 0; j < 4; j++) c[i][j] += a[i] * b[j]; }

c[4][4]这 16 个 float 在整个循环中永远不离开寄存器。编译器会尽量把这些数据分配到高位寄存器里，防止被内层循环临时计算的a、b冲掉。如果你在 CUDA 里用__launch_bounds__或maxregcount强行限制寄存器数量，最先溢出的往往是这些累加器，性能立刻打折。这是 NVIDIA 架构上寄存器 tiling 最典型的“物理约束”。

2.2 AMD CDNA/RDNA：wavefront 与矩阵核的片段布局差异

AMD 的 GPU 编程模型和 NVIDIA 有个大区别：wavefront 大小是 64 个线程（wave32 仅部分 RDNA 系列支持），而不是 warp 的 32。这一个数字变化，直接影响了所有寄存器 tiling 的 fragment 布局。MFMA（Matrix Fused Multiply Add）指令在 CDNA2/3 上承担了类似 mma 的角色，但它要求的数据排布和 NVIDIA 完全不同。同一个 GEMM 分块，在 NVIDIA 上每个线程持有一份 A 片段，在 AMD 上可能是另一个按 64 个 lane 循环分配的排列。

从寄存器硬件本身看，AMD 的寄存器文件是 256 个 32-bit 寄存器, 按 wave64 来划分 lane。CDNA2 的每个 CU 拥有较大的寄存器文件和更宽的 SIMD 宽度，所以 MFMA 一条指令可以一次把大块的矩阵片段从寄存器中取走。代价是：如果你在 NVIDIA 上把 tiling 调成某个形状，在 AMD 上很可能遇到寄存器 bank conflict——因为两个线程要读的寄存器正好落在同一个硬件 bank 上，一次读取被拆成两次。这种冲突在 NVIDIA 上也有，但频率远低于 AMD，因为 AMD 的寄存器 bank 批次大小和历史设计导致的调度粒度不同。

实操来看，在 ROCm 生态里用 HIP 手写 kernel 时，blockSize的选择往往要配合 wave64 的整数倍，否则寄存器 tiling 的边界会产生大量空转。比如一个 128 线程的 block 映射成两个 wave64，每 wave 处理半个 tile，那每个 wave 的寄存器片段尺寸要重算；如果你直接照搬 32 线程一个 warp 的 NVIDIA 写法，编译器会多生成很多 shuffle 指令来重新排布数据，这些指令本身就是寄存器 tiling 的“隐性杀手”。

2.3 Intel GPU 与 AMX：块级寄存器与 tile 指令的另类结合

Intel 的 AI 加速硬件路线基本以 AMX（Advanced Matrix Extensions）为代表。AMX 里有两组值得注意的结构：一组是 8 个 tile 寄存器，每个 tile 寄存器都是 1KB 的二维寄存器块（一个典型的 tile 是 16 行 64 字节，也就是 16x16 的 FP32）；另一组是 tile 配置寄存器TILECFG，用于提前声明每块 tile 的行列维度。

从编程角度看，AMX 的寄存器 tiling 直接抽象到了“块”级别——不是按单线程的 scalar register 进行 tiling，而是按硬件支持的二维 tile 进行 tiling。TMUL指令做 A 和 B 的矩阵乘，TLOAD从内存加载 tile，TSTORE存回。它省去了 NVIDIA 那种“fragment 元素在 warp 内交错排列”的繁琐映射，因为硬件本身就是按块设计。

但这也带来一个新问题：tile 数量只有 8 个，用满之后就需要插桩TLOAD/TSTORE来替换内容。你没办法像 CUDA 那样在一个 kernel 内保住几十个累加器寄存器。因此 Intel AMX 上的 tiling 策略更接近“软流水 + 块复用”：尽量让 8 个 tile 中大部分作为累加器，少量作为输入 A/B，然后频繁从 L1/L2 刷新输入。这种寄存器资源的稀缺性，决定了 Intel 上 tiling 调优的核心是“减少 tile 切换和重载次数”，而不是“增加每线程累加器数量”。

2.4 ARM CPU / 移动端 GPU：向量寄存器和数据重排的艺术

如果不局限于数据中心 GPU，寄存器 tiling 在 ARM 平台上呈现出完全不同的画风。ARMv8 的 NEON 提供 32 个 128-bit 向量寄存器（v0-v31），SVE 进一步支持可变向量长度（128-bit 到 2048-bit）。无论是 NEON 还是 SVE，tiling 都是围绕“向量寄存器中的 lane”展开的——因为 CPU 没有 GPU 那样大而宽的寄存器数组，所有数据都必须以紧凑的向量形式承载。

在 ARM CPU 上做 GEMM 的经典 tiling 是：把 C 的 8x8 分块拆到 8 个 NEON 寄存器，每个寄存器存一行或一列；A 的 4 个元素、B 的 4 个元素通过复用同一组向量寄存器完成外积。这种 tile 尺寸不像 GPU 那样动辄 128x128，而是深深受限于 32 个 NEON 寄存器的容量。SVE 则允许在“向量长度可扩展”的前提下做更大的聚合——比如 512-bit 的 SVE 可以一次处理 16 个 FP32 lane，配合 gather/scatter 和 predication，能写出与 CPU 核数无关的 tiling 循环。

移动端 GPU（比如 Arm Mali 或 Qualcomm Adreno）更特殊——它们的驱动和编译栈对寄存器分配控制得很死，普通开发者几乎碰不到底层寄存器文件。这时做 tiling 更多是靠 shader 编译器在中间表示层帮我们完成，我们能控制的是工作组尺寸和内存布局。所以移动端 GPU 上的寄存器 tiling 属于“半自动挡”，你只负责把块大小和布局定好，剩下交给驱动。

3. 实操：从 Triton 到 CUDA，定量观察寄存器 tiling 的影响

3.1 用 Triton 快速验证跨架构的 tiling 差异

如果你不想一开始就手写 CUDA PTX，Triton 是目前最快的跨架构验证工具。它的编程模型明确引入了 block 级别的表达，下面这个简单的 GEMM 就是最好的实验场：

import triton import triton.language as tl BLOCK_M = 128 BLOCK_N = 128 BLOCK_K = 32 @triton.jit def gemm_kernel( A_ptr, B_ptr, C_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak b_ptrs = B_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, K, BLOCK_K): a_tile = tl.load(a_ptrs) b_tile = tl.load(b_ptrs) acc = tl.dot(a_tile, b_tile, acc) a_ptrs += BLOCK_K * stride_ak b_ptrs += BLOCK_K * stride_bk c_ptrs = C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn tl.store(c_ptrs, acc)

上面的BLOCK_M/N/K就是 tiling 的最直接参数。Triton 编译器在后端会把tl.dot映射成对应架构的矩阵指令，并把acc分配到寄存器。换到不同 GPU 上运行同一个 kernel 时，你可以明显看到BLOCK_N=128在 NVIDIA 上性能还行，但到某些 AMD 卡上反而比BLOCK_N=64差——一个主要原因就是寄存器片段布局和 wavefront 大小的耦合。

当你在 Triton 里编译这个 kernel 时，可以打印生成的 PTX/SASS 或 AMD GCN 指令来看 threads 使用的寄存器数量：

TRITON_KERNEL_DUMP=1 python gemm_bench.py

你会看到不同架构后端对同一 tiling 参数生成了差异巨大的寄存器使用量。有的后端会为 128x128 tile 的 acc 分配 256 个寄存器并触发大量 local memory spill，有的后端则能把 64x128 的 tile 控制在安全寄存器预算内。这就是为什么我建议把“寄存器溢出”作为跨架构 tiling 移植时的第一排查指标。

3.2 手写 CUDA 版寄存器 tiling：如何在不溢出的前提下压榨算力

如果你要在一线生产环境中做极致优化，只用 Triton 通常不够，得能看懂并手写类 CUTLASS 的 kernel。下面给一个我常用的 4x8 C tile 版本 SGEMM 核心片段，展示寄存器 tiling 在 CUDA 中的物理形态：

// 假设 blockDim = (16, 8)，每线程负责 4x8 C tile // 每线程累加器：c[4][8]，32 个 float，即 32 个寄存器 // 每轮 k 循环需要 a[4] 和 b[8]：再占 12 个寄存器 // 总寄存器预算需要控制在 64 以内，留出地址计算等余量 __global__ void sgemm_reg_tiling( const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int M, int N, int K) { int globalRow = blockIdx.y * blockDim.y * 4 + threadIdx.y * 4; int globalCol = blockIdx.x * blockDim.x * 8 + threadIdx.x * 8; float c[4][8] = {0.0f}; for (int k = 0; k < K; ++k) { float a[4]; float b[8]; #pragma unroll for (int i = 0; i < 4; ++i) a[i] = A[(globalRow + i) * K + k]; #pragma unroll for (int j = 0; j < 8; ++j) b[j] = B[k * N + globalCol + j]; #pragma unroll for (int i = 0; i < 4; ++i) for (int j = 0; j < 8; ++j) c[i][j] += a[i] * b[j]; } #pragma unroll for (int i = 0; i < 4; ++i) for (int j = 0; j < 8; ++j) C[(globalRow + i) * N + globalCol + j] = c[i][j]; }

这个实现没有用共享内存，纯粹演示寄存器 tiling 怎么写。每线程 32 个累加器寄存器，加上 a、b 的 12 个临时寄存器，再算上索引、地址寄存器，大概在 64 个上下，不会溢出。真正的生产级 GEMM 会在此之上再加共享内存 staging、向量化访存（float4）、双缓存等，但内核线程级的数据布局思路完全一致。

这里的核心目标是：每次从全局内存读一个 A 元素和 B 元素，就要被 4x8 个 C 累加器用上 32 次（4*8）。如果你把每线程的 C tile 从 4x8 扩到 8x8，复用率翻倍，但因为累加器达到 64 个，再挤下游 a、b 临时寄存器时可能溢出。这就是“复用率”和“寄存器预算”的权衡，每一代 GPU 架构上的最优解都不同，因为寄存器文件总量和每线程配额在变。

3.3 用 profiling 数据实锤“不同架构的寄存器消耗差异”

理论讲再多，不如看数据。我自己在跨架构迁移一个小型矩阵乘法 kernel 时，记录过这样一组 profiling 结果（同一逻辑、同样的每线程 8x8 C tile，编译器自动分配寄存器数）：

表格里的“溢出”是编译日志中报出的 local memory 使用量，通常用local_size字段就能看到。为什么同一逻辑在 MI210 上会溢出？因为 AMD 的 MFMA 指令要求 A/B/C/D 的寄存器片段布局跟 wave64 强相关，编译器为了生成合理 MFMA 代码，需要多分配一些中间寄存器做数据整形，结果就把 255 的每线程寄存器预算挤爆了。这也印证了前面说的：寄存器 tiling 在不同架构上，物理形态真不一样。

4. 常见问题与排查技巧实录

4.1 寄存器溢出为什么难排查，如何快速定位

寄存器溢出是 AI Infra 算子上最典型的“隐形性能杀手”。它不会报错，只是 kernel 变慢，慢到你以为是调度问题或访存问题。

最常见的定位路径是编译日志。NVCC 加--ptxas-options=-v可以直接看到每核函数用了多少寄存器、多少 stack 字节、多少 local memory。HIP 对应--save-temps或rocprof也能看。Triton 则可以通过TRITON_KERNEL_DUMP=1打印生成的代码与寄存器统计。

一旦发现 local memory 非零，优先尝试这几件事：

• 调低每线程 C tile，从 8x8 降为 4x8；
• 用__launch_bounds__或maxrregcount主动限制寄存器数，强制编译器控制占用率；
• 检查内层循环有没有过度#pragma unroll，过度 unroll 会剧增临时寄存器需求；
• 换用更小的BLOCK_K，减少 staged A/B tile 的寄存器缓冲。

这里有一个我踩过的坑：不要盲目用maxrregcount把所有 kernel 都限制在 128 个寄存器。不同 kernel 的最优预算不同，尤其是指令级并行（ILP）较高的 GEMM，限制寄存器后编译器会牺牲累加器复用，反而更低效。先把寄存器数降低，再观察计算吞吐率，找到那个“不溢出且 ILP 最高”的点。

4.2 同样的 block size 换架构后性能剧烈波动，怎么归因

跨架构移植时性能差异突然拉大，很多人第一时间怀疑缓存不友好或带宽问题。我的排查顺序是：

1. 看编译后的寄存器数和溢出情况，把寄存器问题排掉；
2. 看指令混合比——GEMM kernel 中矩阵乘指令（mma、mfma、tmul）占全部计算指令的比例。如果比例偏低，说明 tiling 调度产生太多数据搬移指令；
3. 看 bank/pipe 冲突。AMD 上特别要关注 wave64 导致的跨 lane 数据交换，这些 shuffle 指令会在性能剖析器里显示为“data movement”占用率；
4. 最后才看全局内存与共享内存的带宽利用率。

如果以上数据都正常，但性能还是差，那就大概率是架构本身的寄存器文件和指令吞吐的匹配问题。这只能靠调整 tiling 形状解决，没有银弹。

4.3 三级存储的联动：寄存器 tiling 不能独立优化

寄存器 tiling 的高低位是需要在硬件资源里抠空间的算法问题，它的上游是源指令看寄存器布局，下游则是缓存、共享内存和指令调度。我习惯把它看成三级存储的联动：L2 cache 负责跨 block 的数据复用，共享内存负责块内复用，寄存器负责线程内复用。如果 L2 命中率低、共享内存 bank conflict 高，即使寄存器 tiling 完美，最终 kernel 也会被卡在更外层。

所以调优时不要只盯寄存器。实际工程中我会先跑一个不tiling的 baseline，再逐步叠加共享内存 tiling，最后再做寄存器 tiling 优化。每一步都记录同一指标的耗时变化，就能定位到到底是哪一层拖了后腿。这个方法可比直接瞎试BLOCK_M和BLOCK_N高效得多。

再分享一个小技巧：新架构出现时，不要急着写复杂的寄存器 tiling 代码。先在目标硬件上跑 Intel oneAPI、AMD rocBLAS、NVIDIA CUTLASS 这些官方验证过的库，看看它们在同等矩阵规模下的寄存器分配特征和指令 mix，再用 profiling 数据反推它们的 tiling 参数。把这个当作“架构测量”，比自己反复试错靠谱多了。

寄存器 tiling 在不同架构上的形态差异，说到底是各家公司对“寄存器文件到底该服务于什么”的路线选择不同。NVIDIA 让开发者显式控制 warp 级片段，AMD 把很多数据布局责任压到编译器身上，Intel 干脆把寄存器抽象成块级 tile，ARM 走的是向量长度动态扩展路线。理解了这些差异，你在做算子跨架构移植时就不会再被“玄学性能波动”牵着走了——每一条差异背后，都有寄存器文件和指令集设计者的工程考量。