做 AI Infra 这行聊到算子优化和内核编写几乎绕不开“寄存器 tiling”这个概念。它是把参与计算的小数据块锁在寄存器里反复使用、压低访存指令占比的关键手段但真正把它落到不同硬件上你会发现它长得完全不一样。这篇就借“AI Infra 每日一问 · Day 6”的契机把寄存器 tiling 在 NVIDIA、AMD、昇腾、CPU 这几类常见架构上的不同形态、设计逻辑和实战取舍一次说透适合正在做 GPU 算子开发、性能优化、加速器内核移植的同学当一份对照笔记。1. 寄存器 tiling 到底在解决什么问题1.1 一句话定义把“热数据”锁死在最快存储层寄存器 tiling简单说就是把计算要反复用到的数据切成一个小块tile一次性塞进寄存器文件让计算指令直接操作寄存器里的数据而不是每次运算都去 shared memory 或 global memory 里搬。这个思路和“切菜时把油盐酱醋摆到手边而不是每次用都回柜子里拿”是一个道理。寄存器的访问速度和带宽是存储层级里最顶级的能在这里面完成的数据复用绝不会浪费在更高延迟的存储上。所以在各种高性能计算内核里寄存器 tiling 是最靠近计算单元的一层数据复用手段用来保证计算单元的利用率不被访存拖垮。很多人会把它和 shared memory tiling、global memory tiling 混在一起聊实际上这三者是数据分级缓存的不同层面global memory 的 tile 决定了一个计算块从哪里搬进芯片shared memory 的 tile 解决同一 block 内线程间的数据共享和复用寄存器 tiling 则是每个线程把自己的那部分数据进一步锁死在最靠近计算单元的位置。一套完整的高性能内核通常三层都会用上。1.2 为什么是寄存器而不是只在共享内存里搬这个问题经常有人问共享内存带宽已经很高了为什么还要费劲做寄存器 tiling答案藏在硬件规模和指令开销里。以 NVIDIA 的 A100 为例一个 SM 的寄存器文件总容量大约是 256KB每周期能提供的带宽远高于 shared memory。更关键的是从寄存器取操作数几乎没有额外指令开销而访问 shared memory 每条指令都要带上地址计算和可能的 bank conflict 风险。当一段计算要在同一个数据上执行多次乘加时把它留在寄存器里省掉的不仅是访存时间还有地址生成的指令数。但寄存器不是无限的所以寄存器 tiling 的目标从来不是“tile 越大越好”而是“让访存次数降到计算单元刚好吃满的水平”。一旦 tile 尺寸超过寄存器容量编译器就会把数据溢出到 local memory性能会出现断崖式下跌。这也是为什么寄存器 tiling 是内核优化里最需要精算的一环。2. 寄存器 tiling 在几种主流架构上的长相2.1 NVIDIAwarp 内每线程私有一份 fragmentNVIDIA GPU 的执行模型是 SIMT一个 warp 有 32 个线程每个线程拥有自己独立的寄存器。所以这里的寄存器 tiling 形态非常直接每个线程各自负责一个子块sub-tile寄存器里同时保存 A 的小切片、B 的小切片和 C 的累加器。比如一个最经典的 SIMT GEMM 内核对齐方式每个线程计算 C tile 的 8x8 区域那它就需要 64 个寄存器放 C 累加结果再准备若干寄存器放 A 和 B 的切片。整个 warp 合起来就是 32 份这样的结构并行推进。你可以在反汇编或者 PTX 里清楚看到每个线程的寄存器使用列表这种“一坨扁平寄存器数组”的感觉就是 NVIDIA 上寄存器 tiling 的第一印象。到了 Tensor Core 时代寄存器 tiling 变得更抽象。以mma.sync.aligned.m16n8k16指令为例虽然它在语义上是整个 warp 一起执行矩阵乘加但实际硬件是把一个 16x8 的输出矩阵按固定规则切分到 32 个线程的寄存器里。每个线程持有 4 个 32 位寄存器作为 C fragment另外还要按硬件规定的位置持有 A/B fragment。这时候你不能像写普通 SIMT 内核那样自由规划寄存器布局必须按照指令 ABI 的要求来摆放数据。到了 Hopper 上的 wgmmawarpgroup mma粒度又上升了一个级别整个 warpgroup 共 128 个线程一起参与一个线程持有的数据量更大kernel 也需要在 shared memory 里准备更规整的 fragment 布局然后一次指令把大量数据喂进张量核心。寄存器 tiling 的规划单元从单个 warp 扩展到了 warpgroup对数据摆放的规整度要求也更高。2.2 AMD CDNA64 lane 的 wavefront 加上 MFMAAMD 的 GPU 在不同代际差异不小CDNA 系列里面向 AI 计算的 Matrix Core 和 MFMA 指令是和 NVIDIA Tensor Core 对标的东西。AMD 的基本执行粒度是 wavefront一个 wavefront 包含 64 个 lane比 NVIDIA 的 warp 宽一倍。这个差异直接决定了寄存器 tiling 切分出来的 fragment 长什么样。比如v_mfma_f32_16x16x16_f16这条指令计算的是 16x16 的输出矩阵由 64 个 lane 共同分担平均每个 lane 只要承担 4 个元素的累加。相比 NVIDIA 32 线程分担相同大小的 tile、每个线程要承担 8 个元素AMD 的每线程寄存器压力会明显小一些。但 lane 数多也意味着线程间的同步和数据交换成本更高fragment 分布的索引计算也更复杂。AMD 的寄存器体系还有一个特点除了每个 lane 私有的 VGPR还有一套供整个 wavefront 共享读写的 SGPR标量寄存器。在设计寄存器 tiling 时有些循环不变量、地址偏移可以放进 SGPR从而减少 VGPR 的压力。这个设计让 AMD 上的 tiling 比 NVIDIA 多了“私有/共享”两个维度写内核时可以利用这一点做更细的分配。另一个影响是 AMD 的 VGPR 数量往往比 NVIDIA 更宽松比如 RDNA/CDNA 架构可以给每个线程分配 256 个甚至更多 VGPR。这让大 tile 在寄存器层面有了更多发挥空间但宽松的寄存器预算也容易让人忽略 occupancy——寄存器占得越多能同时驻留的 wavefront 就越少延迟掩盖能力可能下降。2.3 昇腾 Ascend片上 buffer 驱动的分层 tiling昇腾的 AI Core 架构和 GPU 很不一样。它的核心是 Cube 单元做矩阵计算Vector 单元做向量计算Scalar 单元做标量控制数据流更多是围绕各级片上缓冲区展开比如 Unified BufferUB、L1 Buffer 等。在这种架构里你几乎不会像写 CUDA 那样显式地定义“每个线程私有的寄存器数组”更多是把 tiling 当作一个“如何在多级 buffer 间分配数据切块”的问题。具体到算子开发昇腾的 tiling 策略更接近数据流调度的概念你要决定一个 AI Core 负责输出 feature map 或矩阵的哪一块这一块的数据能否塞进 UB输入切片要分多少次从 L1 搬进 UBCube 和 Vector 之间如何用双 buffer 交替隐藏延迟。可以说昇腾把传统意义上“寄存器 tiling shared memory tiling”的职责统一到了片上 buffer 的分层管理里。当然这并不代表昇腾完全不做近计算单元的数据复用。它的 L0 buffer 离 Cube 足够近角色上类似 NVIDIA 的寄存器加共享内存的混合体只是软件编程模型没有把这一层暴露成线程私有寄存器而是以 buffer 的地址空间形式出现。所以在移植 GPU kernel 到昇腾时不能想当然地把float c[8][8]这种寄存器数组搬过去而要先梳理清楚数据需要在哪个 buffer 里保持存活、何时被 Cube 消费、何时被 Vector 复用。2.4 CPU SIMD寄存器阻塞和向量长度的故事CPU 上不会有 GPU 那种成百上千个线程抢寄存器文件的场景但在数值计算内核里寄存器 tiling 同样存在只不过它常被叫做“寄存器阻塞register blocking”或“micro-kernel 展开”。以 x86 为例AVX-512 提供 32 个 512 位的 zmm 寄存器每个能放 16 个 float。一个典型的 GEMM micro-kernel 会让多个 zmm 寄存器同时保存 C 矩阵的多行多列累加器比如 6x16 的累加块就要用掉 6 个 zmm 管 6 行再用若干 zmm 保存 A 和 B 的切片。这样的布局让最内层循环的乘加指令能直接从寄存器取操作数而不是反复访问 L1 cache。ARM 平台也有类似逻辑NEON 的向量寄存器是 128 位SVE 更是把向量长度做成可变的短到 128 位、长到 2048 位都能跑。这意味着同一份寄存器 tiling 代码在不同 CPU 上展开后的寄存器占有率和循环次数会不一样做性能优化时必须跟着实际硬件能力走。CPU 的寄存器总量比 GPU 小得多所以寄存器 tiling 的余地也小。x86 上如果 tile 选得太大编译器很快会把寄存器变量溢出到栈上性能掉得比 GPU 还夸张。这也是为什么 CPU 上的高性能内核常常把更多精力放在 L1/L2 cache 阻塞上寄存器层只是一个“尽量用满但不贪多”的辅助手段。3. 实操一次寄存器 tiling 内核的切分与参数推演3.1 从最朴素的 GEMM 开始算寄存器账与其空谈概念不如实打实算一笔账。假设要写一个最简单的 CUDA GEMM kernel每个线程计算 C 的一个 TM x TN 子块数据类型是 fp32。如果选 TM8、TN8那么C 的累加寄存器8x8 64 个。每一轮 k 循环需要从 shared memory 读 8 个 A 元素和 8 个 B 元素至少需要 16 个暂存寄存器。合计大约 80 个 32 位寄存器。看起来离 NVIDIA 255 个寄存器的上限还很远但如果把 TM 改成 16、TN 改成 16C 累加器就变成 256 个直接超过了硬件上限编译器只能把一部分累加器溢出到 local memory性能就会断崖下跌。所以设计寄存器 tiling 的第一原则是C 累加器的寄存器数量是第一约束A/B 的暂存寄存器是第二约束流水线双缓冲需要的额外寄存器是第三约束。很多人上来就选一个大 tile结果连 C 累加器都放不下更别提缓存 A/B 数据。正确做法是先估算累加器数量再叠加暂存需求最后看是否还有余量做双缓冲或更深度的流水线。3.2 用代码把 tiling 结构写出来下面是一个用 CUDA C 表达的寄存器 tiling 内核骨架重点看数据如何从 shared memory 进入寄存器并复用// 每个线程负责计算 C tile 的 TM x TN 子块 // sA、sB 是已经搬到 shared memory 的矩阵分块 // BK 是 shared memory tile 的 K 维尺寸 __global__ void gemm_reg_tiling(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float sA[BM * BK]; __shared__ float sB[BK * BN]; float c[TM][TN] {0.0f}; float a_reg[TM]; float b_reg[TN]; int row blockIdx.y * BM threadIdx.y * TM; int col blockIdx.x * BN threadIdx.x * TN; for (int k0 0; k0 K; k0 BK) { // 从 global memory 加载到 shared memory此处省略边界处理 __syncthreads(); #pragma unroll for (int kk 0; kk BK; kk) { #pragma unroll for (int i 0; i TM; i) { a_reg[i] sA[(threadIdx.y * TM i) * BK kk]; } #pragma unroll for (int j 0; j TN; j) { b_reg[j] sB[kk * BN threadIdx.x * TN j]; } #pragma unroll for (int i 0; i TM; i) { #pragma unroll for (int j 0; j TN; j) { c[i][j] a_reg[i] * b_reg[j]; } } } __syncthreads(); } // 写回 global memory此处省略 }这段代码里c[TM][TN]就是寄存器 tiling 的核心载体a_reg和b_reg是从 shared memory 到寄存器之间的中转。最内层的三重循环展开后编译器会把所有c[i][j]映射到寄存器只要 TM 和 TN 选得合理就不会产生 local memory 访问。需要注意的是虽然代码里写了浮点数组编译器未必一定把它们放进寄存器。如果数组索引在编译期无法判断或者循环没有被完全展开编译器可能把数组降级到 local memory。工程上常用#pragma unroll配合固定编译期常量来强制展开。3.3 寄存器数量、Occupancy、Tile 大小的三角博弈接下来要面对的是寄存器 tiling 最经典的博弈tile 越大访存削减越明显tile 越大寄存器占用越高能同时驻留的 warp 数量越少最终可能反过来掩盖不了访存延迟。我用一张表来直观表示这种关系选择方向寄存器数Occupancy访存削减适用场景小 tile 4x4约 40高有限访存敏感、数据复用少的算子中 tile 8x8约 80中良好通用 GEMM、卷积大 tile 16x8约 130低很好计算密集型、可容忍低占用实际操作里A100 这样的 Ampere 架构一个 SM 约有 65536 个 32 位寄存器。如果每线程用 128 个寄存器一个 warp 要消耗 4096 个寄存器一个 SM 最多驻留 16 个 warp。如果每线程压到 64 个寄存器同样寄存器预算下可以驻留 32 个 warp。所以当你用__launch_bounds__(256, 4)告诉编译器“线程数 256、每个 SM 至少 4 个 block”时编译器会自动把每线程寄存器数压在 65536 / (256 * 4) 64 以内。这里有个经验不是所有 kernel 都要追求 100% occupancy。像 GEMM 这种计算密集内核只要 tile 足够大、循环足够长低占用率也能通过长周期计算掩盖访存延迟反而比高占用率更稳。但像 elementwise 或小矩阵算子访存延迟占比高这时候高占用率比大 tile 重要得多。选 tile 大小之前先想清楚瓶颈在计算还是访存。4. 不同架构下寄存器 tiling 的差异根源4.1 线程模型不同fragment 归属方式就不同同样是一块 16x16 的矩阵数据NVIDIA 用 32 线程的 warp 去分AMD 用 64 lane 的 wavefront 去分昇腾更彻底直接用 Cube 单元和 buffer 地址空间来切到了 CPU 则根本没有“线程共同持有 fragment”的概念。线程模型决定了“这份数据归谁管、每人管多少、怎么告诉别人自己管了哪一份”。这个差异带来的实际影响是你在 CUDA 上写的寄存器 tiling 内核靠threadIdx.x和threadIdx.y来索引自己负责的 fragment换到 AMD 上索引公式就会因为 lane 数翻倍而改变换到昇腾上可能连 thread 的概念都不存在。所以做跨架构移植绝对不是把threadIdx改成workgroup_id就行而是要重新按照目标架构的执行模型设计 fragment 归属。4.2 存储层次不同寄存器能承接的工作就不一样NVIDIA 的寄存器文件巨大且读取带宽极高可以容纳大量累加器所以它能支持很大的寄存器 tile。CPU 的向量寄存器总数也就几十个寄存器 tiling 只能做得很克制更多依赖 cache 的多级阻塞。昇腾则把最近端的存储做成了 L0 buffer 的形式编程模型里不暴露细粒度寄存器而是靠显式的 buffer 分配和搬运指令管理数据。这些差异本质上是硬件设计者的取舍GPU 愿意为大量并发线程铺设海量寄存器文件CPU 更倾向把晶体管预算花在 cache 和控制逻辑上ASIC 加速器则直接把数据流和控制流固化到 buffer 流水线里。软件层的寄存器 tiling 形态只是这些硬件取舍的外在反映。4.3 指令集抽象高度不同寄存器 tiling 的“颗粒度”不同在 NVIDIA 上PTX 的 mma 指令把 fragment 布局写进了指令语义你只需要声明“我要一个 m16n8k16 的 fragment”编译器会按规范帮你分配寄存器。在 AMD 上MFMA 指令同样有严格的 lane 映射规则但你有很多手动调整 SGPR/VGPR 分配的空间。昇腾则把 tiling 拆成了 host 侧的 tiling 参数计算和 device 侧的 buffer 搬运内核里基本是参数驱动的“流水线状态机”。CPU 上编译器自动完成大多数寄存器分配你主要通过循环展开和微架构调优来间接影响寄存器使用。所以“寄存器 tiling 长什么样”这个问题在不同架构上的答案差异很大核心不只是硬件寄存器数量不同还有指令集和编译器到底把这层抽象暴露到了哪一层。5. 常见问题与排查技巧实录5.1 寄存器溢出spill的定位与缓解寄存器 tiling 最常见的坑就是改大 tile 后性能反而暴跌。本质原因是寄存器不够用编译器把变量放到了 local memory。定位方法很直接编译时加--ptxas-options-v看输出里的 spill 统计如果出现大量 local load/store 字节数恭喜你编译器已经在拖地板了。缓解手段有四个减小 TM 或 TN直接降低 C 累加寄存器压力。使用__launch_bounds__限制每线程寄存器数逼迫编译器精简分配。去掉或减少 A/B 双缓冲代价是可能增加等待周期。确保所有循环都展开避免动态索引导致编译器放弃寄存器分配。我自己的经验是在 Ampere 上写 fp32 GEMM每线程 112 个寄存器是一个比较安全的区间超过 160 就很容易 spill得尽早查汇编。5.2 共享内存 bank conflict 的次生问题有时候寄存器 tiling 调好了性能还是没有预期高这时候要看看 shared memory 的读取有没有发生 bank conflict。比如sA[(threadIdx.y * TM i) * BK kk]这个索引如果TM取值刚好让同一 warp 的不同线程命中同一个 bank延迟就会翻几倍。排查技巧是NVIDIA Nsight Compute 会直接标出 bank conflict 次数AMD 的 ROCm 工具链也有类似分析。修复方式通常是给 shared memory 数组加一列 padding让 stride 不再对齐到 bank 数量。这个坑特别容易在大 tile 内核里出现因为大家都盯着寄存器优化反而忽略了底层 shared memory 的访存模式。5.3 编译器“不听话”怎么办还有一类问题代码明明写了寄存器 tiling但生成的 SASS 里寄存器用量和你预期差很多。这通常有三个原因数组索引里有运行时变量导致编译器放弃寄存器分配循环没有被完全展开快速数学选项或精度设置改变了指令选择。处理方法是直接看汇编。用cuobjdump --dump-sass导出内核的 SASS重点看有没有LDL、STL这类 local memory 指令。只要有说明某个数组没被放进寄存器。然后再回代码里检查索引是不是编译期常量循环是不是被#pragma unroll全展开了。这一步看着麻烦但比瞎猜有效得多。5.4 一个小技巧别只看理论寄存器数最后分享一个我自己的实战习惯选 tile 参数时不要只盯着“理论寄存器上限”还要考虑目标架构的寄存器分配粒度和编译器行为。同一条mma指令不同编译器版本生成的寄存器分配可能差异很大同一个 SM不同核心频率下最佳 occupancy 也可能不一样。我的做法是搭建一个小型的参数扫描脚本把 TM、TN、BK 和__launch_bounds__组合跑一遍记录每个组合的吞吐和寄存器数再对照 Nsight Compute 的报告做选择。这种方法虽然简单但往往能找到比“照抄论文参数”好 10% 到 20% 的配置尤其是面对新架构时特别管用。实际跑过一轮之后你会发现寄存器 tiling 的最终形态不是由某个单一公式决定的而是由线程模型、存储层级、指令集抽象和编译器行为共同塑造的。摸清这些约束再根据瓶颈选参数比死记某个固定的 tile 大小可靠得多。
