第四章 共享内存优化：从入门到精通

共享内存，说白了就是GPU上的一块“手写板”。每个线程块里的线程都能读写它，速度比全局内存快一个数量级。我刚开始做CUDA优化时，总觉得共享内存可有可无——直到有一次，一个简单的矩阵转置，全局内存版本跑了半天，换成共享内存后直接起飞。嗯，从那以后我再也不敢小看它了。

4.1 共享内存的分配与使用

共享内存的分配分两种：静态分配和动态分配。静态分配最简单，直接在核函数里声明就行：

__global__ void kernel() {
    __shared__ float cache[256];  // 静态分配，编译时就定好了
    // 剩下的逻辑...
}

动态分配呢？你可以在调用核函数时指定大小。我个人习惯用动态分配，因为灵活。比如处理不同尺寸的输入数据时，你不需要为每种尺寸都写一个核函数。

__global__ void dynamic_kernel() {
    extern __shared__ float cache[];  // 动态分配，大小由调用时决定
    // 使用 cache[0..size-1]
}

// 调用时指定共享内存大小
dynamic_kernel<<<grid, block, 256 * sizeof(float)>>>();

这里有个坑：共享内存的总大小是有限制的。每个线程块最多能用多少？取决于你的GPU架构。以V100为例，每个SM有96KB共享内存，但每个线程块最多能用48KB（默认配置）。我曾经在项目里贪心了一把，一个块用了64KB，结果核函数直接罢工——嗯，后来查文档才知道上限。

4.2 动态共享内存的实战技巧

动态共享内存最常用的场景是：输入数据大小在运行时才能确定。比如做图像处理时，图片尺寸千变万化，你不可能为每种尺寸都写一个核函数。

来看一个实际例子：向量求和，但每个线程块处理的数据量不固定。

__global__ void flexible_sum(float* input, float* output, int n) {
    extern __shared__ float sdata[];
    
    int tid = threadIdx.x;
    int bid = blockIdx.x;
    int block_size = blockDim.x;
    
    // 每个线程加载一个元素到共享内存
    int idx = bid * block_size + tid;
    sdata[tid] = (idx < n) ? input[idx] : 0.0f;
    __syncthreads();
    
    // 规约求和
    for (int s = 1; s < block_size; s *= 2) {
        if (tid % (2 * s) == 0) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) output[bid] = sdata[0];
}

// 调用时动态分配
int block_size = 256;
int grid_size = (n + block_size - 1) / block_size;
flexible_sum<<<grid_size, block_size, block_size * sizeof(float)>>>(d_in, d_out, n);

为什么用 __syncthreads()？因为共享内存是线程块内共享的，但线程执行顺序不确定。不加同步的话，一个线程读到的数据可能是另一个线程还没写好的——这就是典型的“数据竞争”。我刚开始写CUDA时，就因为这个bug调试了整整两天。

4.3 共享内存与寄存器的权衡

这是个经典问题：用共享内存还是用寄存器？

寄存器更快，但数量有限。每个线程最多能用255个寄存器（不同架构略有差异）。一旦超过这个数，编译器会把多余的数据“溢出”到本地内存——而本地内存其实在全局内存里，速度慢得可怜。

共享内存呢？速度比寄存器慢一点，但容量大得多。而且共享内存是线程块内所有线程共享的，适合做数据交换。

我一般遵循这个原则：

• 线程私有的、频繁访问的数据 → 放寄存器
• 线程间需要共享的数据 → 放共享内存
• 数据量太大，寄存器放不下 → 放共享内存

来看一个矩阵乘法的例子，感受一下权衡：

// 版本1：全部用寄存器（容易溢出）
__global__ void matmul_reg(float* A, float* B, float* C, int N) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    float sum = 0.0f;  // 寄存器
    for (int k = 0; k < N; k++) {
        sum += A[row * N + k] * B[k * N + col];  // 每次都要读全局内存
    }
    C[row * N + col] = sum;
}

// 版本2：使用共享内存做分块
__global__ void matmul_shared(float* A, float* B, float* C, int N) {
    __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];
    
    int row = blockIdx.y * BLOCK_SIZE + threadIdx.y;
    int col = blockIdx.x * BLOCK_SIZE + threadIdx.x;
    
    float sum = 0.0f;
    for (int tile = 0; tile < N / BLOCK_SIZE; tile++) {
        // 加载数据到共享内存
        As[threadIdx.y][threadIdx.x] = A[row * N + tile * BLOCK_SIZE + threadIdx.x];
        Bs[threadIdx.y][threadIdx.x] = B[(tile * BLOCK_SIZE + threadIdx.y) * N + col];
        __syncthreads();
        
        // 从共享内存计算
        for (int k = 0; k < BLOCK_SIZE; k++) {
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        }
        __syncthreads();
    }
    C[row * N + col] = sum;
}

我曾经在一个项目里，把所有的中间变量都塞进共享内存，结果性能反而下降了。为什么？因为共享内存的访问延迟虽然比全局内存低，但比寄存器高。而且共享内存还有bank conflict的问题——这个我们后面会详细讲。

4.4 知识体系总览

下面这张图总结了共享内存优化的核心脉络：

4.5 避坑指南

最后分享几个我踩过的坑：

• 忘记同步：共享内存的数据写完后，一定要加 __syncthreads()。我曾经在一个规约核函数里忘了同步，结果结果时对时错，排查了三天才发现。
• 动态共享内存大小计算错误：调用核函数时，第三个参数是每个块的总共享内存大小（字节）。如果你有多个数组，记得把它们的尺寸加起来。
• 过度使用共享内存：不是所有数据都适合放共享内存。如果每个线程只读一次，直接读全局内存反而更快——因为省去了加载到共享内存的开销。
• 忽略bank conflict：共享内存有32个bank，连续访问同一bank会导致冲突。这个我们下一章会详细讲，但你现在要知道：sdata[threadIdx.x] 这种访问模式是没问题的，但 sdata[threadIdx.x * 32] 就会出问题。

共享内存优化，说白了就是“用空间换时间”。你多花一点心思在数据布局上，GPU就能少等几百个时钟周期。这个投入，绝对值。

想要更多资料，联系可可