Nvidia的”护城河” | Waang’s Blog

type

status

date

slug

summary

本文含少量AI生成内容，请注意鉴别

cuda.ipynb

26.6KB

矩阵乘法

考虑

C=A\cdot B

A\in R^{m\times k}, B\in R^{k\times n}

,利用

C_{ij}=\sum_{r=0}^{K-1}A_{ir}B_{rj}

计算密度为

\frac{2MNK}{MK+KN+MN}

. CPU端串行程序来看,计算访存比为

\frac{1}{2K}

如下内容参考了这个实现：

cuda_code/matrix at main · xgqdut2016/cuda_code

easy cuda code. Contribute to xgqdut2016/cuda_code development by creating an account on GitHub.

https://github.com/xgqdut2016/cuda_code/tree/main/matrix

v1&2.二维线程块实现矩阵乘法使用shared memory

针对当前线程块所有线程可见。仍然存在重复读取的情况，如果重复的不是global memory，而是shared memory。

v3.一个线程处理多个元素

cpp


#define BLOCK_DIM_X 32    // 线程块x方向线程数
#define BLOCK_DIM_Y 32    // 线程块y方向线程数
#define TM 4              // 每个线程处理的行元素数
#define TN 4              // 每个线程处理的列元素数
#define BM (BLOCK_DIM_X * TM)  // A的分块行数: 32*4=128
#define BN (BLOCK_DIM_Y * TN)  // B的分块列数: 32*4=128
#define BK 8                    // 分块共享维度

__global__ void MatrixMulV3(float *A, float *B, float *C, int M, int N, int K) {
    // 共享内存定义（SA: BM×BK, SB: BK×BN）
    __shared__ float SA[BM][BK];
    __shared__ float SB[BK][BN];

    // 线程块索引和线程索引
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;

    // 全局内存起始坐标（每个线程块处理BM×BN的子矩阵）
    int global_m_start = by * BM;  // A的行起始位置
    int global_n_start = bx * BN;  // B的列起始位置

    // 寄存器缓存（每个线程计算TM×TN个元素）
    float c[TM][TN] = {0.0f};

    // 分块循环（总迭代次数为K/BK）
    for (int bk_idx = 0; bk_idx < K; bk_idx += BK) {
        // 协作加载SA和SB到共享内存
        // ------------------------- 加载SA[BM×BK] -------------------------
        for (int load_m = 0; load_m < BM; load_m += BLOCK_DIM_X) {
            int sa_row = load_m + ty;  // 共享内存行索引
            int sa_col = tx;           // 共享内存列索引
            int global_a_row = global_m_start + sa_row;
            int global_a_col = bk_idx + sa_col;
            
            if (global_a_row < M && global_a_col < K) {
                SA[sa_row][sa_col] = A[global_a_row * K + global_a_col];
            } else {
                SA[sa_row][sa_col] = 0.0f;
            }
        }

        // ------------------------- 加载SB[BK×BN] -------------------------
        for (int load_n = 0; load_n < BN; load_n += BLOCK_DIM_Y) {
            int sb_row = ty;           // 共享内存行索引
            int sb_col = load_n + tx;  // 共享内存列索引
            int global_b_row = bk_idx + sb_row;
            int global_b_col = global_n_start + sb_col;
            
            if (global_b_row < K && global_b_col < N) {
                SB[sb_row][sb_col] = B[global_b_row * N + global_b_col];
            } else {
                SB[sb_row][sb_col] = 0.0f;
            }
        }
        __syncthreads();

        // ------------------------- 计算分块乘积 -------------------------
        for (int k = 0; k < BK; ++k) {
            // 加载SA的TM个元素到寄存器
            float a_frag[TM];
            for (int tm = 0; tm < TM; ++tm) {
                int row_in_block = ty * TM + tm;
                a_frag[tm] = SA[row_in_block][k];
            }

            // 加载SB的TN个元素到寄存器
            float b_frag[TN];
            for (int tn = 0; tn < TN; ++tn) {
                int col_in_block = tx * TN + tn;
                b_frag[tn] = SB[k][col_in_block];
            }

            // 外积计算并累加到寄存器
            for (int tm = 0; tm < TM; ++tm) {
                for (int tn = 0; tn < TN; ++tn) {
                    c[tm][tn] += a_frag[tm] * b_frag[tn];
                }
            }
        }
        __syncthreads();
    }

    // ------------------------- 将结果写回全局内存 -------------------------
    for (int tm = 0; tm < TM; ++tm) {
        for (int tn = 0; tn < TN; ++tn) {
            int global_c_row = global_m_start + ty * TM + tm;
            int global_c_col = global_n_start + tx * TN + tn;
            if (global_c_row < M && global_c_col < N) {
                C[global_c_row * N + global_c_col] = c[tm][tn];
            }
        }
    }
}

计算访存比为

\frac{1}{2(\frac{1}{K}+\frac{1}{BM}+\frac{1}{BN})}

v4.优化矩阵load至shared memory过程

线程索引的重新映射

cpp


#define BM 128     // A的分块行数
#define BN 128     // B的分块列数
#define BK 8       // 分块共享维度
#define TM 4       // 每个线程处理的行元素数
#define TN 4       // 每个线程处理的列元素数
#define BLOCK_DIM_X 32  // 线程块x维度
#define BLOCK_DIM_Y 32  // 线程块y维度

__global__ void MatrixMulV4(float *A, float *B, float *C, int M, int N, int K) {
    // 共享内存定义（SA: BM×BK, SB: BK×BN）
    __shared__ float SA[BM][BK];
    __shared__ float SB[BK][BN];

    // 一维线程索引计算
    int tid = threadIdx.x + threadIdx.y * BLOCK_DIM_X;

    // SA的共享内存索引映射
    int smem_a_m = tid % BM;   // SA的行索引（0~127）
    int smem_a_k = tid / BM;   // SA的列索引（0~7）
    
    // SB的共享内存索引映射
    int smem_b_k = tid % BK;   // SB的行索引（0~7）
    int smem_b_n = tid / BK;   // SB的列索引（0~127）

    // 全局内存坐标计算（每个线程块处理BM×BN的子矩阵）
    int block_m = blockIdx.y * BM;  // 当前块在A中的行偏移
    int block_n = blockIdx.x * BN;  // 当前块在B中的列偏移

    // 寄存器缓存（每个线程计算TM×TN个元素）
    float c[TM][TN] = {0.0f};

    // 分块循环（总迭代次数为K/BK）
    for (int bk = 0; bk < K; bk += BK) {
        // 协作加载SA和SB到共享内存
        // 加载SA的BM×BK分块
        int global_a_m = block_m + smem_a_m;
        int global_a_k = bk + smem_a_k;
        if (global_a_m < M && global_a_k < K) {
            SA[smem_a_m][smem_a_k] = A[global_a_m * K + global_a_k];
        } else {
            SA[smem_a_m][smem_a_k] = 0.0f;
        }

        // 加载SB的BK×BN分块（注意B的全局索引需要转置）
        int global_b_k = bk + smem_b_k;
        int global_b_n = block_n + smem_b_n;
        if (global_b_k < K && global_b_n < N) {
            SB[smem_b_k][smem_b_n] = B[global_b_k * N + global_b_n];
        } else {
            SB[smem_b_k][smem_b_n] = 0.0f;
        }
        __syncthreads();

        // 计算分块乘积（外积累加）
        for (int k = 0; k < BK; ++k) {
            // 加载SA的TM元素和SB的TN元素到寄存器
            float a_frag[TM];
            float b_frag[TN];
            
            // 每个线程处理TM×TN元素（TM=4, TN=4）
            #pragma unroll
            for (int tm = 0; tm < TM; ++tm) {
                int row = threadIdx.y * TM + tm;
                a_frag[tm] = SA[row][k];
            }
            #pragma unroll
            for (int tn = 0; tn < TN; ++tn) {
                int col = threadIdx.x * TN + tn;
                b_frag[tn] = SB[k][col];
            }

            // 外积计算并累加到寄存器
            #pragma unroll
            for (int tm = 0; tm < TM; ++tm) {
                #pragma unroll
                for (int tn = 0; tn < TN; ++tn) {
                    c[tm][tn] += a_frag[tm] * b_frag[tn];
                }
            }
        }
        __syncthreads();
    }

    // 将结果写回全局内存C
    #pragma unroll
    for (int tm = 0; tm < TM; ++tm) {
        #pragma unroll
        for (int tn = 0; tn < TN; ++tn) {
            int global_c_m = block_m + threadIdx.y * TM + tm;
            int global_c_n = block_n + threadIdx.x * TN + tn;
            if (global_c_m < M && global_c_n < N) {
                C[global_c_m * N + global_c_n] = c[tm][tn];
            }
        }
    }
int M = 1024, N = 1024, K = 1024;
dim3 blockDim(BLOCK_DIM_X, BLOCK_DIM_Y);  // (32, 32)
dim3 gridDim((N + BN - 1) / BN, (M + BM - 1) / BM);  // (8, 8)
MatrixMulV4<<<gridDim, blockDim>>>(A, B, C, M, N, K);
}

v5.float4类型访存

减少内存事务：单次内存操作可加载/存储 4 个元素，提升带宽利用率。

cpp


float a[4];
(float4 &)a[0] = (float4 &)dA[0];  // 一次性加载4个float到a[0~3]

注意问题：循环展开次数增加，寄存器压力过大（寄存器爆炸），可能引发性能下降，所以需要线程块参数调优，以平衡寄存器使用。

v6.bank conflict

Bank Conflict 定义

共享内存分为 32 个 Bank，每个 Bank 宽度为 4 字节。同一 Warp（32 线程）访问不同 Bank 时性能最佳，否则产生 Bank Conflict。

原始 SA 访问的冲突问题

C 元素索引：[8×threadIdx.x+index.q][8×threadIdx.y+index.v]。

SA 索引公式： $64 \times \text{threadIdx.x} + \text{index.q} \times 8 + \text{index.k}$

SB 索引公式： $\text{index.k} \times 128 + 8 \times \text{threadIdx.y} + \text{index.v}$

固定 index.q 和 index.v 时，SA 的访问模式导致同一 Warp 的线程访问相同 Bank，引发 Bank Conflict。

优化方法

转置 SA 的加载方式：将 SA 的共享内存布局从 [128][BK] 调整为 [8][128]。调整后 SA 的索引模式分散到不同 Bank，消除 Bank Conflict。

代码示意（转置 SA 后的访问逻辑）：

cpp


__shared__ float SA_T[8][128];  // 转置后的 SA 布局

// 加载时按转置方式填充
SA_T[index.k][8*threadIdx.x + index.q] = ...;

// 计算时按转置布局读取
float a_val = SA_T[k][threadIdx.x * 8 + q];

v7.降低shared memory的读取

💡

SA 和 SB 的转置布局使得对 index_k 的访问是非连续的，但对 index_q 和 index_v 是连续的，可以借助寄存器和float4进行优化

核心思想：内积转换成外积

cpp


for (int index_k = 0; index_k < Bk; index_k++) {
    // 从共享内存 SA 中加载 8 个 float 到 com_a（使用 float4 向量化加载）
    (float4 &)com_a[0] = (float4 &)SA[index_k * BM + threadIdx.x * TM];       // 加载前4个元素
    (float4 &)com_a[4] = (float4 &)SA[index_k * BM + threadIdx.x * TM + 4];   // 加载后4个元素

    // 从共享内存 SB 中加载 8 个 float 到 com_b（同样使用 float4）
    (float4 &)com_b[0] = (float4 &)SB[index_k * BM + threadIdx.y * TN];       // 加载前4个元素
    (float4 &)com_b[4] = (float4 &)SB[index_k * BM + threadIdx.y * TN + 4];   // 加载后4个元素

    // 计算外积并累加到 tmp 数组
    for (int index_q = 0; index_q < TN; index_q++) {
        for (int index_v = 0; index_v < TN; index_v++) {
            tmp[index_q * TN + index_v] += com_a[index_q] * com_b[index_v];
        }
    }
}

外积方式将分块乘法分解为多个外积的累加：

A = \begin{bmatrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{bmatrix}

B = \begin{bmatrix} B_{00} & B_{01} \\ B_{10} & B_{11} \end{bmatrix}

可以通过计算

A_{i0} \otimes B_{0j}+A_{i1} \otimes B_{1j}

实现

v8.流水并行

重叠全局内存加载与计算操作，减少内存访问延迟，提升 GPU 计算效率。

原始模式

加载 data[0] → 计算 data[0] → 加载 data[BK-1] → 计算 data[BK-1]

顺序执行加载和计算，存在空闲等待时间.

流水并行模式

加载 data[0] → [加载 data[1] + 计算 data[0]] → ... → [加载 data[BK-1] + 计算 data[BK-2]] → 计算 data[BK-1]

mermaid


graph TD
    subgraph 流水并行模式
        A[加载 data0] --> B[计算 data0]
        A --> C[加载 data1]
        C --> D[计算 data1]
        D --> E[...]
        E --> F[加载 data BK-1]
        F --> G[计算 dataBK-1]
    end

cpp


pu = 0;
{
    加载 data[0] 到共享内存数组 SA[0];  // 初始化第一个分块
}
__syncthreads();  // 确保所有线程完成加载

for (ph = 1; ph < BK / K; ph++) {
    加载 data[ph] 到共享内存数组 SA[ph % 2];  // 双缓冲交替加载（SA[0] 和 SA[1]）
    if (ph - 1) % 2 == 0 {
        __syncthreads();  // 按需同步，确保前一分块计算完成
    }
    计算 SA[(ph - 1) % 2];  // 使用上一个分块的数据
}

v9.cuBLAS库

cuBLAS 是 NVIDIA 提供的 GPU 加速线性代数库，支持矩阵乘法等基础运算。其核心公式为C=α⋅op(A)⋅op(B)+β⋅C，

αβ：标量系数。

op(X)：对矩阵X的转置操作（可选不转置、转置或共轭转置）。

需注意列优先存储与 C 语言行优先的差异！ v10.Tensor Core

warp读取矩阵：tensor core计算过程中，left_frag和right_frag在内存中不连续，这时再结合二位索引重排。

PTX

a low-level Parallel Thread eXecution virtual machine and instruction set architecture (ISA),直面意思是低级并行线程执行虚拟机和指令集架构。PTX是上承GPU编程语言CUDA C++，下启GPU硬件SASS指令，可以借助NVRTC实现运行时优化，某些层面上来说可以称之为GPU设备无关代码，因此PTX可以理解为”CUDA IR“。