第三章 数据复用与局部性原理

各位同学,今天我们来聊聊缓存策略里最核心的两个概念——时间局部性和空间局部性。说实话,这两个词听起来有点学术,但说白了,它们就是芯片设计里最朴素的“偷懒”哲学。

我刚开始做AI芯片那会儿,总觉得DRAM带宽够大就行。直到第一次流片回来,跑矩阵乘法时发现性能只有理论值的30%。嗯,从那以后,我再也不敢小看局部性原理了。

3.1 时间局部性:同一个数据,反复用

时间局部性是什么意思?简单说就是:如果一个数据被访问了一次,那么它很可能在不久的将来再次被访问。

你想想看,在矩阵乘法里,一个输入元素会被多个输出元素共享。比如C = A × B,A矩阵的某一行会被B矩阵的所有列反复使用。这就是典型的时间局部性。

核心观点:时间局部性越强,缓存命中率越高。我们只需要把数据留在缓存里,就能避免反复去DRAM里取。

我在项目中遇到过这样一个案例:一个卷积层,输入特征图是224×224×64,卷积核是3×3×64×128。如果不做任何优化,每个输入像素会被卷积核滑动窗口重复访问9次(3×3)。但如果我们把输入特征图的一整行缓存下来,时间局部性就得到了充分利用。

3.2 空间局部性:相邻数据,一起取

空间局部性更直白:如果一个地址被访问了,它附近的地址很快也会被访问。

为什么?因为内存访问模式往往是连续的。比如卷积操作,卷积核在输入特征图上滑动,每次移动一个像素。相邻的卷积窗口,它们访问的输入数据高度重叠。

我记得有一次调试一个3×3卷积,发现缓存命中率只有40%。后来一查,问题出在数据布局上——输入特征图是按行存储的,但卷积核是按列滑动的。空间局部性被破坏了。改成按行滑动后,命中率直接飙到85%。

我的习惯:设计缓存预取策略时,先看数据访问模式。如果是连续访问,就预取一个cache line(通常64字节);如果是随机访问,就别预取,否则会污染缓存。

3.3 矩阵乘法中的局部性体现

我们来看一个具体的矩阵乘法例子。假设C = A × B,A是M×K,B是K×N。

最朴素的实现是三重循环:

for i in range(M):
    for j in range(N):
        for k in range(K):
            C[i][j] += A[i][k] * B[k][j]

这个实现有什么问题?B矩阵是按列访问的。如果B是按行存储的,那么每次访问B[k][j]时,k变化而j不变,这会导致严重的缓存缺失。因为B[k][j]和B[k+1][j]在内存中相隔N个元素,空间局部性极差。

我建议改成循环交换:

for i in range(M):
    for k in range(K):
        for j in range(N):
            C[i][j] += A[i][k] * B[k][j]

这样,内层循环j是连续变化的,B[k][j]就是按行访问的。空间局部性一下子就变好了。同时,A[i][k]在内层循环中保持不变,时间局部性也得到了利用。

我曾经踩过的坑:以为循环交换只是编译器优化的事。结果发现,编译器不一定能识别出所有可交换的循环。尤其是当矩阵维度不是2的幂次时,编译器可能会放弃优化。所以,手动调整循环顺序是值得的。

3.4 卷积中的局部性体现

卷积比矩阵乘法更复杂。因为卷积涉及多个维度:输入通道、输出通道、高度、宽度、卷积核大小。

一个典型的卷积操作:

for oc in range(O):
    for ic in range(I):
        for h in range(H):
            for w in range(W):
                for kh in range(KH):
                    for kw in range(KW):
                        output[oc][h][w] += 
                            input[ic][h+kh][w+kw] * weight[oc][ic][kh][kw]

这个六重循环,局部性很差。为什么?因为内层循环kh和kw变化时,input的访问是随机的(取决于h和w的偏移)。

我常用的优化方法是im2col + GEMM。把卷积转换成矩阵乘法:

  • 把输入特征图展开成一个大矩阵,每一列对应一个卷积窗口
  • 把卷积核也展开成矩阵
  • 然后做矩阵乘法

这样做的好处是:矩阵乘法有成熟的优化库(比如BLAS),局部性已经被优化好了。但代价是内存占用变大(因为im2col会复制数据)。

数据复用率计算:对于一个3×3卷积,输入特征图的每个像素会被复用9次(卷积核大小)。如果输入通道是64,输出通道是128,那么每个输入像素会被复用9×128=1152次。这就是为什么卷积的局部性优化空间巨大。

3.5 如何利用缓存减少DRAM访问

好了,理论讲完了。我们来看看实际怎么做。

减少DRAM访问的核心思路就两条:

  1. 提高缓存命中率——让数据尽量留在缓存里
  2. 减少数据搬运次数——一次搬运,多次使用

具体策略:

策略 适用场景 效果
循环分块(Tiling) 矩阵乘法、卷积 将大矩阵切分成小块,让每个块能完全放入缓存
数据重排(Data Layout) 卷积、池化 把不连续的数据重新排列成连续存储
预取(Prefetching) 连续访问模式 提前把即将用到的数据加载到缓存
写合并(Write Combining) 输出累加 多个写操作合并成一次DRAM写

我个人最常用的是循环分块。以矩阵乘法为例:

// 分块大小,根据L1缓存大小调整
int BI = 64, BJ = 64, BK = 64;

for (int i = 0; i < M; i += BI) {
    for (int j = 0; j < N; j += BJ) {
        for (int k = 0; k < K; k += BK) {
            // 处理子块 C[i:i+BI][j:j+BJ] += A[i:i+BI][k:k+BK] * B[k:k+BK][j:j+BJ]
            for (int ii = i; ii < i+BI; ii++) {
                for (int jj = j; jj < j+BJ; jj++) {
                    for (int kk = k; kk < k+BK; kk++) {
                        C[ii][jj] += A[ii][kk] * B[kk][jj];
                    }
                }
            }
        }
    }
}

分块后,A的子块和B的子块都能放进L1缓存。这样,内层循环的所有数据访问都在缓存内完成,只有分块边界处才会访问DRAM。

避坑指南:分块大小不是越大越好。我曾经试过把分块设成L1缓存大小的一半,结果因为缓存行冲突,命中率反而下降了。建议分块大小设为缓存大小的1/4到1/3,同时保证分块的行列数是2的幂次。

3.6 知识体系图

下面这张图总结了本章的核心逻辑:

数据复用与局部性原理 时间局部性 空间局部性 矩阵乘法 A矩阵行复用、B矩阵列复用 卷积 输入特征图复用、卷积核复用 优化策略:循环分块 + 数据重排 + 预取 + 写合并 最终目标:减少DRAM访问,提升能效比 数据复用率越高,缓存命中率越高,DRAM访问次数越少

这张图把时间局部性和空间局部性分别映射到矩阵乘法和卷积两个典型场景,然后引出四种优化策略,最终指向减少DRAM访问这个核心目标。

好了,这一章的内容就到这里。记住一句话:缓存优化的本质,就是让数据在缓存里多待一会儿,少去DRAM串门。

总结:

  • 时间局部性:同一个数据反复用 → 循环分块、数据复用
  • 空间局部性:相邻数据一起取 → 数据重排、预取
  • 矩阵乘法:循环交换 + 分块,命中率提升2-3倍
  • 卷积:im2col + GEMM,或直接优化六重循环
  • 减少DRAM访问:一次搬运,多次使用

想要更多资料,联系可可