存储层次优化：从全局内存到寄存器文件的性能博弈

存储层次，说白了就是芯片上的数据搬运系统。我做了这么多年AI芯片优化，最深的体会就是：计算单元再快，数据喂不上去也是白搭。很多同学把精力都花在优化计算上，结果发现性能瓶颈其实在内存访问——这坑我踩过不止一次。

全局内存、共享内存、寄存器文件：三者的带宽与延迟特性

先看一张我整理的对比表，这玩意儿你最好记在心里：

存储层级 典型容量 延迟（周期） 带宽（相对值） 作用域
寄存器文件 ~256KB/核 1~2 极高（~100x） 单个线程/PE
共享内存 ~48~96KB/块 ~20~30 高（~10x） 线程块内
全局内存 ~8~80GB ~400~800 低（1x基准） 所有线程

看到这个延迟差距了吗？全局内存访问一次，够寄存器做几百次运算了。我个人习惯是：能放寄存器的绝不放共享内存，能放共享内存的绝不碰全局内存。这话听起来简单，实际做起来全是细节。

寄存器文件：最贵也最快

寄存器文件是离计算单元最近的存储。每个线程私有，延迟只有1~2个周期。但它的容量极其有限——一个SM里通常也就几万个寄存器，分到每个线程头上可能就几十个。

我在项目中遇到过一个问题：某个卷积核的代码，寄存器压力太大导致编译器把变量溢出（spill）到局部内存。结果呢？性能直接掉了30%。后来我手动拆分了计算逻辑，把中间结果复用起来，才把寄存器占用降下来。

核心原则：寄存器是黄金，能复用就复用，别浪费。编译器有时候并不聪明，你得帮它一把。

共享内存：可编程的缓存

共享内存是芯片上的一块SRAM，由程序员显式管理。它的延迟比全局内存低一个数量级，而且同一个线程块内的所有线程都能访问。

但共享内存有个大坑——Bank Conflict。我刚开始做优化时，被这玩意儿折磨得不轻。

Bank Conflict 的根源

共享内存被划分成32个Bank（以NVIDIA为例），每个Bank宽度4字节。如果多个线程同时访问同一个Bank的不同地址，就会发生冲突，导致访问串行化。

举个例子：

// 假设共享内存数组 shared[32][32]
// 线程 i 访问 shared[i][j]
// 如果按行访问：线程0访问[0][0]，线程1访问[1][0]...
// 这些地址映射到不同的Bank，无冲突

// 但如果按列访问：线程0访问[0][j]，线程1访问[1][j]...
// 这些地址可能映射到同一个Bank，产生冲突
我曾经踩过的坑：写了一个矩阵转置的kernel，按列读取共享内存，结果性能比预期差了3倍。查了半天才发现是Bank Conflict在作祟。后来改成按行读取+转置存储，冲突就没了。

如何避免Bank Conflict

几个实用技巧：

  • 填充（Padding）：在数组末尾加几个无用元素，改变地址映射关系。比如 shared[32][33] 而不是 shared[32][32]
  • 改变访问模式：尽量让连续线程访问连续地址，这样自然映射到不同Bank
  • 使用向量化访问：一次读取128位（4个float），减少访问次数

你想想看，Bank Conflict的本质是什么？是地址映射的哈希函数出了问题。理解了这一点，你就能自己推导出解决方案。

数据预取：让数据在路上

数据预取的核心思想很简单：别让计算单元闲着等数据。在计算当前数据的同时，把下一批数据提前加载到缓存或共享内存里。

我习惯用软件预取（software prefetch），因为硬件预取对不规则访问模式基本没用。看个例子：

// 朴素版本：计算和加载串行
for (int i = 0; i < N; i++) {
    load_data(&input[i]);      // 等待内存
    compute(input[i]);         // 计算
}

// 预取版本：计算和加载重叠
load_data(&input[0]);          // 先加载第一个
for (int i = 0; i < N-1; i++) {
    prefetch(&input[i+1]);     // 预取下一个
    compute(input[i]);         // 计算当前
}
compute(input[N-1]);          // 处理最后一个
小技巧：预取的距离要合适。太近了，数据还没到；太远了，可能把有用的数据踢出缓存。我一般设预取距离为计算延迟的1.5~2倍，具体得调。

双缓冲：流水线化的极致

双缓冲是预取的升级版。它用两块缓冲区交替工作：一块用于计算，一块用于数据加载。说白了就是让内存访问和计算完全流水线化。

我画了一张流程图，帮你理解这个机制：

双缓冲流水线示意图 时间→ 加载到Buffer A Buffer B空闲 计算Buffer A 加载到Buffer B Buffer A空闲 计算Buffer B 加载到Buffer A Buffer B空闲 加载和计算完全重叠，隐藏了内存访问延迟 加载 计算 空闲

双缓冲的代码实现大概是这样的：

// 初始化
load(buffer[0], data_block_0);
current = 0;

for (int i = 0; i < num_blocks - 1; i++) {
    int next = 1 - current;  // 切换缓冲区
    
    // 异步加载下一块数据
    async_load(buffer[next], data_block[i+1]);
    
    // 计算当前数据
    compute(buffer[current]);
    
    // 等待加载完成
    sync_wait();
    
    current = next;
}

// 处理最后一块
compute(buffer[current]);
关键点：双缓冲的效果取决于加载时间和计算时间的匹配程度。如果加载时间远大于计算时间，双缓冲也只能帮你隐藏一部分延迟。这时候你得考虑更激进的数据复用策略。

综合优化策略：从全局视角看存储层次

说了这么多，其实就一句话：让数据尽量靠近计算单元。具体怎么做？

  1. 数据分块（Tiling）：把大矩阵切成小块，让每个块能完全塞进共享内存
  2. 寄存器重用：同一个数据尽量在寄存器里多算几次，别急着写回
  3. 合并访问（Coalescing）：全局内存访问要连续，这样硬件能一次加载一大块
  4. 异步传输：用DMA或者异步拷贝，让数据传输和计算并行

我记得有一次优化一个Transformer的推理，就是靠这三板斧——分块+双缓冲+寄存器重用，把端到端延迟降了40%。说白了，存储优化没有银弹，你得根据具体的数据流和计算模式来组合使用这些技术。

我的建议：拿到一个kernel，先看它的访存模式。如果是连续访问，优先优化全局内存的合并访问；如果是随机访问，考虑用共享内存做缓存；如果数据复用率高，寄存器重用是王道。

嗯，存储层次优化就聊到这儿。这些技术你多练几次，慢慢就能形成直觉——看到一段代码，大概就能猜到它的访存瓶颈在哪儿。到那时候，你离性能优化高手就不远了。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →