CUDA编程进阶：从GEMM优化到逼近GPU理论算力极限

原创已于 2025-05-21 13:43:37 修改

· 862 阅读

13 ·

版权

文章标签：

#人工智能 #transformer #深度学习 #机器学习 #python

于 2025-04-06 15:29:03 首次发布

GPU 同时被 3 个专栏收录

42 篇文章

订阅专栏

人工智能

42 篇文章

订阅专栏

高校

41 篇文章

订阅专栏

引言：GPU算力压榨的艺术

在深度学习和科学计算领域，GEMM（通用矩阵乘法）占据超过70%的计算量。英伟达A100 GPU的理论FP16算力达312 TFLOPS，但未经优化的GEMM实现往往只能达到理论值的10-30%。本文通过SAXPY基准分析→GEMM优化层次→逼近硬件极限的递进式路径，揭示CUDA性能调优的核心方法论。

一、性能评估基础：SAXPY案例与带宽测试

1.1 SAXPY：内存带宽的试金石
SAXPY（单精度αX+Y）是典型的内存带宽受限型操作：

__global__ void saxpy(int n, float a, float* x, float* y) {  
    int i = blockIdx.x * blockDim.x + threadIdx.x;  
    if (i < n) y[i] = a * x[i] + y[i];  
}

性能公式：
有效带宽（GB/s）=（数据量×2）/(时间×1e9)
（每个元素需读取x[i]和y[i]，写入y[i]，共3次操作，但现代GPU通过L2缓存合并访问）

1.2 A100 GPU带宽瓶颈分析

理论带宽：1555 GB/s（HBM2e）
SAXPY实测值：约1300 GB/s（达到理论值83.6%）
优化关键：确保全局内存访问合并（Coalesced Access）

二、GEMM优化层次化拆解

2.1 优化层次金字塔

2.2 Level 1：基础核函数优化

// 基础矩阵乘法核函数  
__global__ void gemm_naive(float *A, float *B, float *C, int M, int N, int K) {  
    int row = blockIdx.y * blockDim.y + threadIdx.y;  
    int col = blockIdx.x * blockDim.x + threadIdx.x;  
    if (row < M && col < N) {  
        float sum = 0.0f;  
        for (int k = 0; k < K; ++k)  
            sum += A[row*K + k] * B[k*N + col];  
        C[row*N + col] = sum;  
    }  
}

性能缺陷：

全局内存非合并访问（B矩阵列遍历）
未利用共享内存，重复加载数据

2.3 Level 2：共享内存分块优化
将数据分块加载到共享内存，减少全局内存访问：

__global__ void gemm_tiled(float *A, float *B, float *C, int M, int N, int K) {  
    __shared__ float As[TILE_SIZE][TILE_SIZE];  
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];  

    int bx = blockIdx.x, by = blockIdx.y;  
    int tx = threadIdx.x, ty = threadIdx.y;  

    int row = by * TILE_SIZE + ty;  
    int col = bx * TILE_SIZE + tx;  

    float sum = 0.0f;  
    for (int t = 0; t < K/TILE_SIZE; ++t) {  
        // 加载分块到共享内存  
        As[ty][tx] = A[row*K + t*TILE_SIZE + tx];  
        Bs[ty][tx] = B[(t*TILE_SIZE + ty)*N + col];  
        __syncthreads();  

        for (int k = 0; k < TILE_SIZE; ++k)  
            sum += As[ty][k] * Bs[k][tx];  
        __syncthreads();  
    }  
    if (row < M && col < N)  
        C[row*N + col] = sum;  
}

优化效果：性能提升5-8倍，但仍有寄存器瓶颈

2.4 Level 3：双缓冲与寄存器优化
使用双缓冲技术隐藏内存延迟，最大化寄存器利用率：

float a[2][THREAD_PER_TILE];  
float b[2][THREAD_PER_TILE];  
#pragma unroll  
for (int t = 0; t < K; t += TILE_SIZE) {  
    // 异步加载下一块到缓冲区  
    load_tile_to_registers(a[(t/TILE_SIZE)%2], ...);  
    load_tile_to_registers(b[(t/TILE_SIZE)%2], ...);  
    // 计算当前块  
    compute_tile(a[(t/TILE_SIZE-1)%2], b[(t/TILE_SIZE-1)%2]);  
}

性能提升：进一步获得2-3倍加速

三、逼近极限：高级调优技巧

3.1 指令级并行（ILP）
通过循环展开和寄存器重用提升指令吞吐：

#pragma unroll 4  
for (int k = 0; k < TILE_SIZE; k += 4) {  
    sum0 += a0 * b0;  
    sum1 += a1 * b1;  
    sum2 += a2 * b2;  
    sum3 += a3 * b3;  
}

3.2 避免Bank冲突
调整共享内存访问模式，确保同一Warp内线程访问不同Bank：

// 列主序存储 + 添加Padding  
__shared__ float As[TILE_SIZE][TILE_SIZE + 1];

3.3 使用Tensor Core
调用WMMA API利用Tensor Core加速：

#include <mma.h>  
using namespace nvcuda;  

wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;  
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;  
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;  

wmma::load_matrix_sync(a_frag, A + ...);  
wmma::load_matrix_sync(b_frag, B + ...);  
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);  
wmma::store_matrix_sync(C + ..., c_frag);