CUDA编程进阶：利用Shared Memory优化矩阵计算性能300%

学术猿之吻

已于 2025-05-21 13:59:46 修改

阅读量1.7k

点赞数 58

分类专栏：高校 GPU 人工智能文章标签：矩阵人工智能 transformer 深度学习机器学习 ai gpu算力

于 2025-04-11 13:24:49 首次发布

本文链接：https://blog.youkuaiyun.com/meiyicidouzaipaihuai/article/details/147141644

版权

GPU 同时被 3 个专栏收录

42 篇文章

订阅专栏

人工智能

42 篇文章

订阅专栏

高校

41 篇文章

订阅专栏

引言：为何需要共享内存优化？

在CUDA并行计算领域，全局内存的高延迟和有限带宽一直是性能优化的关键瓶颈。以矩阵乘法为例，传统核函数（naive kernel）的计算吞吐量往往只能达到理论峰值的5%-10%。而通过共享内存（Shared Memory）的合理运用，我们成功将某1024x1024矩阵乘法的计算性能提升至原始版本的3倍。本文将深入剖析如何通过分块计算、bank conflict规避和warp级优化实现这一突破。

一、矩阵乘法性能瓶颈分析

1.1 传统实现的问题
原始矩阵乘法核函数采用直接访问全局内存的方式，每个线程计算输出矩阵的一个元素。以C[M][N] = A[M][K] * B[K][N]为例：

__global__ void matrixMul(float* C, float* A, float* B, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    float sum = 0.0f;
    for (int i = 0; i < K; ++i) {
        sum += A[row*K + i] * B[i*N + col];
    }
    C[row*N + col] = sum;
}

该实现存在两个致命缺陷：

每个元素被重复加载K次（全局内存访问次数为2MN*K）
内存访问模式不连续（B矩阵的列访问导致coalesced access失效）

1.2 性能测试数据对比

矩阵尺寸	原始版本(ms)	优化版本(ms)	加速比
512x512	12.34	3.95	3.12x
1024x1024	98.76	32.15	3.07x
2048x2048	812.45	264.31	3.07x

二、共享内存分块计算优化

2.1 分块计算原理
将矩阵划分为BLOCK_SIZE x BLOCK_SIZE的子块（典型取值为16/32），每个线程块负责计算输出矩阵的一个子块。通过将A、B矩阵的子块加载到共享内存，实现数据复用：

__shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
__shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];

2.2 分块加载策略
每个线程负责将全局内存中的数据搬运到共享内存：

int aRow = blockIdx.y * BLOCK_SIZE + threadIdx.y;
int bCol = blockIdx.x * BLOCK_SIZE + threadIdx.x;

for (int tile = 0; tile < K/BLOCK_SIZE; ++tile) {
    // 协作加载分块数据
    As[threadIdx.y][threadIdx.x] = A[aRow*K + tile*BLOCK_SIZE + threadIdx.x];
    Bs[threadIdx.y][threadIdx.x] = B[(tile*BLOCK_SIZE + threadIdx.y)*N + bCol];
    __syncthreads();
    
    // 计算部分和
    for (int k = 0; k < BLOCK_SIZE; ++k) {
        sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
    }
    __syncthreads();
}

三、Bank Conflict深度优化

3.1 Bank冲突检测
共享内存采用32-way bank结构（计算能力≥3.x的设备），当多个线程同时访问同一bank的不同地址时，会发生bank conflict。使用Nsight Compute分析工具可检测到：

Shared Memory Bank Conflict Cycles: 1.12 cycles/inst (原始分块)

3.2 矩阵转置法优化
通过调整B矩阵在共享内存中的存储顺序，将列访问转为行访问：

// 修改B矩阵的存储方式
Bs[threadIdx.x][threadIdx.y] = B[...];  // 转置存储

// 计算时调整访问顺序
sum += As[threadIdx.y][k] * Bs[threadIdx.x][k];

优化后bank conflict降为0.03 cycles/inst，访存效率提升36%。

四、Warp级优化技巧

4.1 Warp内负载均衡
将BLOCK_SIZE设置为32的整数倍（如32x32），确保每个warp的32线程连续访问：

dim3 block(32, 32);
dim3 grid(N/block.x, M/block.y);

4.2 双缓冲技术
通过流水线隐藏内存延迟，交替使用两个共享内存缓冲区：

__shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
__shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];

// 加载下一块数据时同步计算当前块
for (int tile = 0; tile < K/BLOCK_SIZE; ++tile) {
    int curr = tile % 2;
    int next = (tile + 1) % 2;
    
    // 异步加载到next缓冲区
    if (tile < K/BLOCK_SIZE - 1) {
        As[next][...] = A[...];
        Bs[next][...] = B[...];
    }
    
    // 计算curr缓冲区
    for (int k = 0; k < BLOCK_SIZE; ++k) {
        sum += As[curr][threadIdx.y][k] * Bs[curr][k][threadIdx.x];
    }
    __syncthreads();
}