内存层次结构：全局内存、共享内存与寄存器文件，数据搬移的代价
做AI芯片这么久，我越来越觉得一件事很关键——内存层次结构。说白了，你算力再强，数据喂不进去，一切都是白搭。很多刚入行的朋友喜欢盯着峰值算力看，觉得TOPS越高越好。嗯，我当年也这么想，直到被一个项目狠狠教育了一回。
今天咱们就聊聊AI芯片里的三级存储：全局内存、共享内存、寄存器文件。这三者的关系，就像你办公桌上的水杯、手边的笔记本，还有你脑子里的短期记忆。哪个快，哪个慢，哪个容量大，你得心里有数。
一、三级存储的物理真相
先看一张图，这是我画的内存层次结构。你一看就明白。
这张图我每次讲课都会放。你看，从寄存器到全局内存，延迟差了上百倍。我见过不少团队，算法调得飞起，结果一上芯片，性能惨不忍睹。为什么？数据搬移的代价没算清楚。
二、寄存器文件：最快但最贵
寄存器文件是离计算单元最近的地方。它有多快？一个时钟周期，数据就能到位。但代价是什么？面积大、功耗高。所以芯片里寄存器文件通常只有几十KB。
核心要点：寄存器文件是AI芯片的"黄金地段"。能放进去的数据，千万别往外挪。
我在项目中遇到过一个问题：某个卷积层的权重反复从共享内存读取，每次读都要等十几个周期。后来我把最常用的权重直接塞进寄存器文件，性能直接提升了30%。你想想看，这多划算。
但要注意，寄存器文件不是无限大的。每个线程或处理单元能用的寄存器数量是有限的。比如NVIDIA的GPU，每个线程最多255个寄存器。一旦超出，就会发生"寄存器溢出"——数据被迫存到慢速内存里，性能断崖式下跌。
避坑指南：我曾经调试过一个模型，发现某个kernel的性能比预期差5倍。查了半天，原来是寄存器使用超标，编译器自动做了spill。从那以后，我写kernel都会先检查寄存器用量。
三、共享内存：手动管理的缓存
共享内存，也叫L1缓存或片上SRAM。它的延迟在10~30个时钟周期，比寄存器慢，但比全局内存快得多。关键是——你可以手动控制。
为什么手动控制很重要？因为AI芯片的访存模式往往有规律可循。比如矩阵乘法，数据可以被分块（tiling）加载到共享内存，然后反复使用。这比每次都去全局内存读，效率高太多了。
我举个例子。假设你要做两个256x256矩阵的乘法。如果不做分块，每次计算都要从全局内存读数据，延迟200周期起步。但如果用共享内存做分块，把16x16的小块先搬上来，计算完再换下一块，延迟就降到了20周期左右。
个人习惯：我一般把共享内存分成两个bank，一个用于读，一个用于写。这样能避免bank conflict，提高带宽利用率。
共享内存的容量通常有限。以NVIDIA A100为例，每个SM有192KB的共享内存。听起来不少，但如果你要跑大模型，这点空间根本不够。所以你得精打细算——哪些数据该放共享内存，哪些该放寄存器，哪些只能放全局内存。
四、全局内存：大但慢
全局内存就是HBM或DDR。容量大，几个GB到几十个GB，但延迟高得吓人——200到800个时钟周期。说白了，CPU去内存读一次数据，GPU能算几百次乘加。
所以，全局内存的优化核心就一句话：减少访问次数，提高访问效率。
怎么减少访问次数？
- 数据复用：同一个数据尽量多用几次。比如卷积的权重，被多个输出通道共享，那就一次加载，反复使用。
- 合并访问：连续的内存地址一起读，利用总线带宽。我见过有人写代码，每个线程读一个不连续的位置，结果带宽利用率不到10%。
- 预取：在计算当前数据的同时，提前把下一批数据加载到缓存里。
我记得有一次优化一个Transformer模型，发现瓶颈在全局内存的访问。后来用了预取和合并访问，推理速度提升了2倍。嗯，这就是数据搬移优化的威力。
五、数据搬移的代价：一个量化分析
咱们来算一笔账。假设你要处理一个1MB的数据块，从全局内存搬到寄存器里做计算。
| 搬移路径 | 延迟（周期） | 带宽（GB/s） | 耗时（微秒） | 能耗（pJ/bit） |
|---|---|---|---|---|
| 全局内存 → 寄存器 | 200~800 | ~900 (HBM2e) | ~1.1 | ~10 |
| 全局内存 → 共享内存 | 200~800 | ~900 | ~1.1 | ~10 |
| 共享内存 → 寄存器 | 10~30 | ~2000 | ~0.5 | ~2 |
| 寄存器 → 计算单元 | 1 | N/A | ~0.001 | ~0.1 |
你看，从全局内存搬数据到寄存器，耗时1.1微秒，能耗10pJ/bit。而从共享内存搬，能耗只有2pJ/bit。如果你能减少一次全局内存访问，换成共享内存访问，能效就能提升5倍。
关键结论：数据搬移的代价，不仅仅是延迟，还有能耗。在AI芯片里，数据搬移消耗的能量，往往比计算本身还高。这就是为什么我们要优化内存层次结构。
六、实战中的优化策略
说了这么多理论，咱们来点实际的。我在做AI芯片部署时，常用的优化策略有这几条：
- 数据分块（Tiling）：把大矩阵切成小块，逐块加载到共享内存。这样全局内存的访问次数从N次降到N/块大小次。
- 寄存器缓存：把最常访问的变量声明为register类型，强制编译器用寄存器存储。但别滥用，小心寄存器溢出。
- 双缓冲（Double Buffering）：用两个共享内存缓冲区，一个用于计算，一个用于加载下一批数据。这样计算和访存可以重叠。
- 向量化加载：用SIMD指令一次加载多个数据，提高带宽利用率。
我举个例子，双缓冲的伪代码大概长这样：
// 双缓冲示例
__shared__ float buffer[2][BLOCK_SIZE];
// 初始化：加载第一批数据到buffer[0]
load_data(buffer[0]);
for (int i = 0; i < num_tiles; i++) {
int current = i % 2;
int next = (i + 1) % 2;
// 异步加载下一批数据到buffer[next]
async_load(buffer[next], i + 1);
// 计算当前buffer[current]中的数据
compute(buffer[current]);
// 等待异步加载完成
sync();
}这段代码我用了很多次。它能让计算和访存完全重叠，理论上性能翻倍。当然，实际中还要考虑同步开销和bank conflict，但思路是对的。
一个小技巧：在写双缓冲代码时，记得把同步操作放在计算之后、下一次加载之前。我曾经搞反了顺序，结果数据还没加载完就开始计算，出了bug。调试了一整天才发现。
七、总结一下
内存层次结构，说白了就是用容量换速度，用手动管理换效率。寄存器最快但最小，全局内存最大但最慢，共享内存是中间的平衡点。
做AI芯片优化，你得时刻问自己三个问题：
- 这个数据应该放在哪一层？
- 能不能减少一次全局内存访问？
- 数据搬移的能耗，值不值得？
我见过太多人，算法模型调得完美，但部署到芯片上就卡成PPT。原因往往不是算力不够，而是数据搬移的代价没算清楚。嗯，希望今天的分享能帮你少走一些弯路。
想要更多资料，联系可可