4. 矩阵分块与内存层次

各位同学，今天我们来聊聊矩阵分块。说实话，这是张量加速核性能调优里最核心的一环。我个人习惯把这个问题拆成两个层面来看：数据怎么放，以及数据怎么搬。

先问大家一个问题：为什么矩阵乘法要分块？

答案其实很简单——因为内存有层次。GPU 里有三级存储：全局内存（Global Memory）、共享内存（Shared Memory）、寄存器文件（Register File）。每一级的速度和容量都不一样。全局内存很大，但慢；共享内存快，但小；寄存器最快，但更小。

你想想看，如果直接把两个大矩阵扔进全局内存，然后让每个线程去读，那性能会怎样？

嗯，我告诉你——带宽瓶颈直接卡死。全局内存的延迟是几百个周期，而共享内存只有几个周期。所以，分块的目的，就是把数据从慢的全局内存，搬到快的共享内存里，让计算单元能快速访问。

三级存储的层次结构

我画了一张图，帮你理解这三层的关系：

这张图你看懂了吗？数据从全局内存搬到共享内存，再从共享内存搬到寄存器，最后喂给计算单元。每一层搬运都有代价，所以我们要尽量减少搬运次数。

Bank Conflict 与 Shared Memory Padding

共享内存虽然快，但它有个坑——Bank Conflict。

共享内存内部被分成了 32 个 Bank（在大多数现代 GPU 上）。每个 Bank 每个周期只能响应一次访问。如果多个线程同时访问同一个 Bank 的不同地址，就会发生冲突，导致访问串行化。

我在项目中遇到过这样一个场景：写一个矩阵转置的 kernel，性能死活上不去。后来一查，发现共享内存的访问模式导致了严重的 Bank Conflict。每个 warp 里的线程都在访问同一个 Bank 的不同行，结果延迟直接翻了好几倍。

举个例子，假设共享内存有 32 个 Bank，每行 32 个 float：

// 不填充：每行 32 个元素
__shared__ float smem[32][32];
// 访问 smem[0][0] 和 smem[1][0] 会冲突（都在 Bank 0）

// 填充：每行 33 个元素（多一个 padding）
__shared__ float smem[32][33];
// 访问 smem[0][0] 和 smem[1][0] 分别在 Bank 0 和 Bank 1，不冲突

你看，就多了一个元素，性能可能提升 30% 以上。嗯，这里要注意：padding 会浪费一点共享内存空间，但通常值得。

数据预取（Prefetching）技巧

数据预取，说白了就是提前把数据搬到共享内存里，让计算和搬运重叠。

我习惯用双缓冲（Double Buffering）来实现预取。思路很简单：

1. 准备两块共享内存缓冲区：A 和 B
2. 当前计算使用缓冲区 A 的数据
3. 同时，异步地把下一块数据从全局内存搬到缓冲区 B
4. 下一轮计算时，交换角色

这样，计算和搬运就重叠起来了。你想想看，如果计算需要 100 个周期，搬运需要 80 个周期，那总时间就是 max(100, 80) = 100 个周期，而不是 180 个周期。

我曾经在一个项目中，把预取和分块结合起来，让一个矩阵乘法的 kernel 性能提升了 2.3 倍。具体做法是：

• 把矩阵分成 16x16 的块
• 每个块使用双缓冲预取
• 共享内存里用 padding 避免 Bank Conflict
• 寄存器里缓存最常用的数据

结果呢？全局内存的访问次数减少了 80%，共享内存的利用率提高了 50%。

实战建议

最后，给你几个实战建议：

场景	建议
矩阵乘法	分块大小 16x16 或 32x32，配合双缓冲预取
卷积	分块大小根据卷积核大小调整，注意 Bank Conflict
Reduce 操作	尽量用寄存器，减少共享内存的使用
Stencil 计算	使用 halo 区域预取，避免边界条件

记住一句话：数据在哪里，性能就在哪里。把数据放在离计算最近的地方，你的 kernel 就成功了一半。