内存层次结构：从寄存器到全局内存
做AI芯片优化这些年，我最大的感触就是：内存访问的代价，往往比计算本身高出一个数量级。你想想看，一个FMA（融合乘加）操作可能只需要几个时钟周期，但一次全局内存访问，可能要几百甚至上千个周期。这就是为什么我们要花大把精力去搞算子融合——说白了，就是为了减少这些昂贵的内存访问。
一、存储层次全景图
先看一张我手绘的层次结构图，它基本概括了现代AI芯片的内存体系：
这张图我每次培训都会拿出来讲。你看，从寄存器到HBM，延迟差了三个数量级。我在优化一个Transformer模型时遇到过这样的情况：明明计算量不大，但跑起来就是慢。后来一分析，发现80%的时间都花在了等待数据从HBM搬运到计算单元的路上。
二、各层次详解
1. 寄存器（Register）
寄存器是离计算单元最近的地方。我个人习惯把寄存器想象成「工程师桌上的便签纸」——随手就能拿到，但空间有限。在AI芯片中，寄存器通常用来存放当前正在计算的中间结果。
关键数据：
- 容量：几百字节（比如一个线程32个寄存器，每个32位）
- 延迟：1个时钟周期
- 带宽：理论上是芯片内部最高，超过10TB/s
嗯，这里要注意：寄存器是编译器最头疼的资源之一。寄存器溢出（register spilling）一旦发生，性能直接跳水。我曾经在优化一个卷积算子时，就因为寄存器分配不当，导致编译器把中间结果写回了L1 Cache，性能直接掉了30%。
2. L1 Cache
L1 Cache是寄存器的「后备队」。它通常和计算单元在同一个时钟域内，访问速度非常快。在NVIDIA的GPU架构中，每个SM都有自己的L1 Cache，容量一般在16KB到64KB之间。
| 属性 | 典型值 | 说明 |
|---|---|---|
| 容量 | 16-64 KB | 每个SM独立拥有 |
| 延迟 | 2-4 周期 | 和计算单元同频 |
| 带宽 | ~5 TB/s | 远高于HBM |
| 策略 | LRU / 写回 | 硬件自动管理 |
你可能会问：为什么L1 Cache不做得更大？说白了，面积和功耗是硬约束。L1 Cache用的是SRAM，6个晶体管才能存1个bit，面积大、功耗高。做大了，计算单元就得让地方。
3. L2 Cache
L2 Cache是芯片上的「二级仓库」。它通常被多个计算单元共享。我参与过一个项目，芯片的L2 Cache设计成4个bank，每个bank独立访问——这样能提高并行度，但也带来了bank冲突的问题。
避坑指南：我曾经在优化一个矩阵乘法算子时，发现L2 Cache的命中率只有60%。排查后发现，是因为两个线程同时访问同一个bank的不同地址，导致bank冲突。解决办法是给地址加一个偏移量，让访问均匀分布到不同bank上。
4. HBM / Global Memory
HBM（高带宽内存）是AI芯片的主力存储。它通过硅中介层（interposer）和计算芯片堆叠在一起，带宽能做到1-2TB/s。但代价是延迟高——200到800个周期不等。
我记得有一次调试一个大规模推荐系统模型，数据量太大，L2 Cache根本装不下。结果每次迭代都要从HBM重新加载权重，性能瓶颈直接卡在内存带宽上。后来我们用了算子融合，把多个小算子合并成一个，减少了中间结果的写回，才把性能提上来。
三、数据局部性原理
数据局部性，说白了就是让数据尽量待在离计算单元近的地方。它分两种：
- 时间局部性：如果一个数据被访问了，不久后很可能再次被访问。比如循环中的累加变量。
- 空间局部性：如果一个地址被访问了，附近的地址很可能也会被访问。比如数组的连续遍历。
为什么局部性这么重要？因为Cache的替换策略是基于局部性设计的。如果你的代码局部性好，Cache命中率就高，数据就能待在L1或L2里，不用频繁去HBM取数据。
一个经典例子：
// 局部性差的写法（按列访问）
for (int j = 0; j < N; j++) {
for (int i = 0; i < M; i++) {
sum += A[i][j]; // 每次跳一行，Cache行被浪费
}
}
// 局部性好的写法（按行访问）
for (int i = 0; i < M; i++) {
for (int j = 0; j < N; j++) {
sum += A[i][j]; // 连续访问，Cache行被充分利用
}
}我在实际项目中见过太多人踩这个坑。按列访问矩阵，Cache命中率可能只有10%，性能差10倍以上。你想想看，同样的计算量，就因为访问顺序不同，结果天差地别。
四、延迟与带宽的权衡
做算子融合的时候，我经常要权衡一个问题：是把中间结果写回HBM，还是留在寄存器里继续算？
举个例子，一个简单的算子链：Conv → ReLU → Pool。如果不做融合，Conv的输出要写回HBM，ReLU再从HBM读出来，Pool再写回去。这一来一回，内存访问次数翻了一倍。
如果做融合，Conv算完一个输出，ReLU和Pool紧接着算，结果只写一次HBM。代价是寄存器压力变大，可能要多占一些寄存器。
| 方案 | HBM访问次数 | 寄存器使用 | 总延迟估计 |
|---|---|---|---|
| 不融合 | 6次（3读3写） | 低 | 高（受HBM带宽限制） |
| 融合 | 2次（1读1写） | 中 | 低（计算密集） |
注意：融合不是万能的。如果算子链中某个算子的计算量特别小，融合带来的寄存器压力可能反而导致性能下降。我遇到过这种情况：融合后寄存器溢出，编译器把数据写回L1 Cache，结果延迟反而比不融合还高。
五、实战建议
基于我这些年的经验，给你几条实在的建议：
- 优先优化内存访问模式：先看数据是怎么流动的，再看计算怎么优化。很多时候，改一下数据布局，性能就上来了。
- 利用tiling技术：把大矩阵切分成小块，让每个小块能塞进L1 Cache。我习惯把tile大小设为L1 Cache容量的1/4左右，留点余量给其他数据。
- 关注bank冲突：在共享内存或L2 Cache中，尽量让不同线程访问不同的bank。可以用padding技术来避免冲突。
- 不要迷信融合：融合前先估算一下寄存器压力和Cache命中率。如果寄存器不够用，或者融合后Cache命中率下降，那就不如不融合。
好了，这一章的内容就到这里。记住一句话：内存层次结构是AI芯片优化的基石。不理解它，算子融合就是空中楼阁。下一章我们会深入具体的融合策略，到时候再细聊。
公众号:蓝海资料掘金营，微信deep3321