CUDA算子优化(4)：矩阵乘法

矩阵乘法的数学定义

矩阵乘法是线性代数中的核心运算。给定两个矩阵：

• 矩阵A：维度为$M \times K$

• 矩阵B：维度为$ K \times N $

它们的乘积矩阵$ C = A \times B $的维度为$ M \times N $。其中，C中的第i行第j列的元素$ C[i][j] $由以下公式计算：

$$C[i][j] = \sum_{k = 0}^{K - 1}{A[i][k] * B[k][j]}$$

即，C的每一个元素是A的一行与B的一列对应元素乘积之和

计算强度：

计算强度是衡量计算任务中浮点运算量与内存访问量比率的核心指标，单位为FLOPs/byte（每字节浮点运算次数），其计算公式为：

$$I = \frac{总浮点运算次数(FLOPs)}{从主存(DRAM)读取/写入的总字节数}$$

对于上述矩阵乘法来说，每个C[i][j]需要K次乘法+(K- 1)次加法，约等于2K FLOPs，总的FLOPs=M x N x 2K = 2MNK；对于主存访问量，读取A，需要M x K x 4字节；读取B，需要K x N x 4字节；写入C，需要M x N x 4字节，总字节数=4(MK + KN + MN)

所以计算强度为$I = \frac{2MNK}{4(MK + KN + MN)} = \frac{MNK}{2(MK + KN + MN)}$

当M=K=N时，I = N / 6（FLOPs/byte），矩阵乘法的计算强度与矩阵大小N成正比。

版本V0矩阵乘法cpu实现

void multi_cpu(float *C, float *A, float *B, int M, int K, int N) {
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += A[k + i * K] * B[j + k * N];
            }
            C[j + i * N] = sum;
        }
    }
}

当M=N=K时，时间复杂度为$O(MNK)=O(N^3)$，空间复杂度为$O(MK+KN+MN)=O(N^2)$。

版本V1朴素矩阵乘法

实现原理：全局内存访问

其核心思想

1. 任务并行分解：一个线程负责计算结果矩阵C的一个元素C[i][j]
2. 全局内存访问：每个线程直接从全局内存中读取矩阵A的一整行和矩阵B的一整列数据
3. 计算与累加：线程执行内积运算，将结果写回全局内存中的C矩阵。

核函数设计思路

• 线程网络与块结构：使用二维网格和二维线程块，自然地映射到二维结果矩阵C

• 索引计算：每个线程通过blockIdx、blockDim和threadIdx计算出自己负责的C矩阵元素坐标(i, j)

• 循环计算：线程通过一个循环，累加A[i][k] * B[k][j]的结果

• 边界检查：确保线程索引不会超出矩阵维度

核函数启动配置

// 线程块大小：16x16 = 256 线程/block（GPU warp 友好的选择）
dim3 blockDim(32, 32);
// Grid 大小：覆盖整个 C 矩阵
dim3 gridDim((N + 32 - 1) / 32, (M + 32 - 1) / 32);

matmul_v1_naive<<<gridDim, blockDim>>>(d_C, d_A, d_B, M, K, N);

每个线程块包含 $32 \times 32 = 1024$ 个线程，正好是 32 个 warp（$1024/32=32$），对 GPU 调度友好。Grid 的 x 方向覆盖 C 矩阵的列数 N，y 方向覆盖行数 M，每个线程独立计算一个 $C[i][j]$。

CUDA Kernel 代码实现

#include "utils.h"
#include <stdio.h>
#include <stdlib.h>

#define WS 32

// 每个线程处理一个输出元素
__global__ void sgemm_kernel1(float *C, const float *A, const float *B, int M, int N, int K) {
  int col = threadIdx.x + blockDim.x * blockIdx.x;
  int row = threadIdx.y + blockDim.y * blockIdx.y;
  if (row >= M || col >= N)
    return;

  float val = 0.0f;
  for (int k = 0; k < K; ++k) {
    // C[row][col] += A[row][k] * B[k][col]
    val += A[k + row * K] * B[col + k * N];
  }
  C[col + row * N] = val;
}

int main(int argc, char **argv) {
  int M = 2048, K = 4096, N = 3072;

  size_t Abytes = (size_t)M * K * sizeof(float);
  size_t Bbytes = (size_t)K * N * sizeof(float);
  size_t Cbytes = (size_t)M * N * sizeof(float);

  float *h_A = (float *)malloc(Abytes);
  float *h_B = (float *)malloc(Bbytes);
  float *h_C = (float *)malloc(Cbytes);
  float *h_ref = (float *)malloc(Cbytes);
  if (!h_A || !h_B || !h_C || !h_ref) {
    printf("A/B/C malloc failed!\n");
    return 1;
  }
  init_float_normal(h_A, M * K, 0.0f, 1.0f);
  init_float_normal(h_B, K * N, 0.0f, 1.0f);

  float *d_A, *d_B, *d_C;
  CHECK_CUDA(cudaMalloc(&d_A, Abytes));
  CHECK_CUDA(cudaMalloc(&d_B, Bbytes));
  CHECK_CUDA(cudaMalloc(&d_C, Cbytes));
  CHECK_CUDA(cudaMemcpy(d_A, h_A, Abytes, cudaMemcpyHostToDevice));
  CHECK_CUDA(cudaMemcpy(d_B, h_B, Bbytes, cudaMemcpyHostToDevice));

  // CPU 参考计算
  sgemm_cpu(h_ref, h_A, h_B, M, N, K);

  dim3 blockDim(WS, WS);
  int gx = (N + WS - 1) / WS;
  int gy = (M + WS - 1) / WS;
  dim3 gridDim(gx, gy);

  // warmup
  sgemm_kernel1<<<gridDim, blockDim>>>(d_C, d_A, d_B, M, N, K);
  CHECK_CUDA(cudaGetLastError());
  CHECK_CUDA(cudaMemcpy(h_C, d_C, Cbytes, cudaMemcpyDeviceToHost));
  check_results(h_C, h_ref, M * N, 5e-4f);

  int repeat = argc > 1 ? atoi(argv[1]) : 20;
  CudaTimer timer;
  cuda_timer_start(&timer);
  for (int r = 0; r < repeat; r++) {
    sgemm_kernel1<<<gridDim, blockDim>>>(d_C, d_A, d_B, M, N, K);
    CHECK_CUDA(cudaGetLastError());
  }
  float ms = cuda_timer_stop(&timer);
  float avg_ms = ms / repeat;
  double gflops = 2.0 * M * N * K / (avg_ms * 1e6);
  printf("avg %.3f ms | %.1f GFLOPS | ", avg_ms, gflops);

  free(h_ref);
  free(h_A);
  free(h_B);
  free(h_C);
  CHECK_CUDA(cudaFree(d_A));
  CHECK_CUDA(cudaFree(d_B));
  CHECK_CUDA(cudaFree(d_C));
  return 0;
}

内存访问模式分析

朴素实现虽然逻辑清晰，但内存访问效率极低。以 $M=N=K=1024$ 为例逐一分析：

单线程视角

每个线程负责一个 $C[i][j]$，需要：

• 读取 A 的第 i 行：$K$ 个 float = $4K$ 字节

• 读取 B 的第 j 列：$K$ 个 float = $4K$ 字节

• 写入 C 的 1 个元素：$4$ 字节

单线程总访存量 = $8K + 4$ 字节，计算量为 $2K$ FLOPs。

全局视角：重复读取问题

问题在于不同线程之间存在大量重复的全局内存访问：

• A 矩阵的行复用：A 的第 i 行被负责 $C[i][0], C[i][1], \dots, C[i][N-1]$ 的 $N$ 个线程各自完整读取一遍，同一行数据被重复读取 $N$ 次。

• B 矩阵的列复用：B 的第 j 列被负责 $C[0][j], C[1][j], \dots, C[M-1][j]$ 的 $M$ 个线程各自完整读取一遍，同一列数据被重复读取 $M$ 次。

总全局内存读取量：

$$Read_{global} = M \times N \times (K + K) \times 4 = 8MNK \text{ 字节}$$

但实际只需 $4(MK + KN)$ 字节。当 $M=N=K$ 时：

$$Read_{global} = 8N^3,\quad Read_{ideal}=8N^2,\quad 冗余比 = N$$

以 $N=1024$ 为例，全局内存读取量是理论最小值的 1024 倍，绝大部分带宽被浪费在重复读取上。

有效计算强度

理论计算强度 $I_{peak}=N/6$，而朴素实现的实际计算强度：

$$I_{effective} = \frac{2MNK}{8MNK} = \frac{1}{4} \text{ FLOPs/byte}$$

与矩阵规模 N 无关！无论矩阵多大，有效计算强度恒定为 0.25 FLOPs/byte，完全受限于内存带宽。

性能瓶颈与 Roofline 模型

Roofline 模型是分析计算瓶颈的经典工具。横轴为算术强度（FLOPs/byte），纵轴为可达到的性能（FLOPs/s）。模型由两部分组成：

• 内存带宽边界（斜线）：$P = bs \times I$，当计算强度较低时，性能受限于内存带宽。

• 计算峰值边界（水平线）：$P = P_{peak}$，当计算强度足够高时，性能受限于算力天花板。

以 NVIDIA A100 为例（显存带宽约 1555 GB/s，FP32 峰值约 19.5 TFLOPS），两条线的交点（ridge point）约为：

$$I_{ridge} = \frac{19.5 \times 10^{12}}{1555 \times 10^9} \approx 12.5 \text{ FLOPs/byte}$$

即需要 $I > 12.5$ 才能突破带宽瓶颈。而朴素 V1 的有效计算强度仅为 $0.25$，远低于 ridge point。这意味着：

• 朴素实现处于 Roofline 模型的陡峭斜线区域，性能完全由内存带宽决定。

• 即使 GPU 有强大的计算能力，绝大部分时间也在等待数据从全局内存加载。

• 要提升性能，核心思路是提高有效计算强度——减少对全局内存的访问，通过数据复用将 $I_{effective}$ 推向 ridge point。

实测性能参考（$M=N=K=4096$，FP32，A100）：

实现版本	耗时	GFLOPS	带宽利用率
V1 Naive	~120 ms	~1150	~8%
cuBLAS	~0.8 ms	~172000	—

朴素实现仅达到理论峰值的约 5.9%，与高度优化的 cuBLAS 相比差距超过 100 倍，优化空间巨大。

优化方向

从上述分析可以看出，朴素实现的根本瓶颈在于全局内存的重复访问。后续优化版本将围绕"数据复用"这一核心思想展开：

• V2：共享内存分块（Tiling）——将矩阵分块加载到共享内存中，块内线程复用同一份数据，大幅减少全局内存访问。

版本V1共享内存分块

分块计算的基本原理

分块（Tiling）是高性能计算中一种关键的优化技术，其核心思想是将大规模的数据集（如大型矩阵）分解为更小的、可管理的“块”，使得每个数据块能够完全放入高速但容量有限的片上内存（如CUDA的共享内存）中。

其核心目标

• 提高数据局部性：让计算单元反复使用已加载到高速内存中的数据

• 减少全局内存访问：全局内存访问延迟高、带宽有限，是性能的主要瓶颈。

• 隐藏内存延迟：通过让更多的线程在等待数据时执行计算，提高硬件利用率。

具体做法

针对矩阵乘法 $C = A \times B$，每个线程块负责计算 $C$ 的一个分块（tile）。将 $A$ 和 $B$ 的对应分块加载到共享内存中，线程块内的各线程协力从共享内存读取数据，完成该分块的累加计算。

具体做法：

• 将 $C$ 分解为大小为 $TILE_M \times TILE_N$ 的子块，每个线程块计算一个子块。

• 计算该子块需要 $A$ 的 $TILE_M$ 行与 $B$ 的 $TILE_N$ 列。沿公共维度 $K$ 再划分步长为 $TILE_K$ 的滑动窗口，每次将 $A$ 的 $TILE_M \times TILE_K$ 分块和 $B$ 的 $TILE_K \times TILE_N$ 分块从全局内存加载到共享内存中。

• 线程块内所有线程协作，将这两个共享内存分块中的值相乘并累加到局部寄存器中，一个分块计算完毕后，滑动窗口向前移动，直到覆盖整个 $K$ 维。

• 计算过程中，每个 $A$ 元素被 $TILE_N$ 次复用，每个 $B$ 元素被 $TILE_M$ 次复用，显著提高了数据局部性。

通过这种方式，分块计算将原本 $O(MNK)$ 次全局内存访问大幅降低，有效计算强度 $I_{effective}$ 提升至 $\frac{2 \times TILE_M \times TILE_N \times TILE_K}{(TILE_M \times TILE_K + TILE_N \times TILE_K) \times 4}$（FP32），只要分块尺寸选择得当，就能使 $I_{effective}$ 接近 ridge point 甚至进入计算密集区域，从而突破带宽瓶颈。

CUDA Kernel 代码实现

#include "utils.h"
#include <stdio.h>
#include <stdlib.h>

#define WS 32

__global__ void sgemm_kernel2(float *C, const float *A, const float *B, int M,
                              int N, int K) {
  int tx = threadIdx.x;
  int ty = threadIdx.y;
  int col = tx + blockDim.x * blockIdx.x;
  int row = ty + blockDim.y * blockIdx.y;

  __shared__ float As[WS][WS];
  __shared__ float Bs[WS][WS];

  int ntiles = (K + WS - 1) / WS;
  float val = 0.0f;
  for (int nt = 0; nt < ntiles; ++nt) {
    // 1. 全局的A/B -> 共享 As/Bs
    if (row < M && (tx + nt * WS) < K) {
      As[ty][tx] = A[row * K + tx + nt * WS];
    } else {
      As[ty][tx] = 0.0f;
    }

    if (col < N && (ty + nt * WS) < K) {
      Bs[ty][tx] = B[(ty + nt * WS) * N + col];
    } else {
      Bs[ty][tx] = 0.0f;
    }

    __syncthreads();

    // 2. 局部累积
    for (int k = 0; k < WS; ++k) {
      val += As[ty][k] * Bs[k][tx];
    }
    __syncthreads();
  }

  if (row < M || col < N)
    C[col + row * N] = val;
}

int main(int argc, char **argv) {
  int M = 2048, K = 4096, N = 3072;

  size_t Abytes = (size_t)M * K * sizeof(float);
  size_t Bbytes = (size_t)K * N * sizeof(float);
  size_t Cbytes = (size_t)M * N * sizeof(float);

  float *h_A = (float *)malloc(Abytes);
  float *h_B = (float *)malloc(Bbytes);
  float *h_C = (float *)malloc(Cbytes);
  float *h_ref = (float *)malloc(Cbytes);
  if (!h_A || !h_B || !h_C || !h_ref) {
    printf("A/B/C malloc failed!\n");
    return 1;
  }
  init_float_normal(h_A, M * K, 0.0f, 1.0f);
  init_float_normal(h_B, K * N, 0.0f, 1.0f);

  float *d_A, *d_B, *d_C;
  CHECK_CUDA(cudaMalloc(&d_A, Abytes));
  CHECK_CUDA(cudaMalloc(&d_B, Bbytes));
  CHECK_CUDA(cudaMalloc(&d_C, Cbytes));
  CHECK_CUDA(cudaMemcpy(d_A, h_A, Abytes, cudaMemcpyHostToDevice));
  CHECK_CUDA(cudaMemcpy(d_B, h_B, Bbytes, cudaMemcpyHostToDevice));

  // CPU 参考计算
  sgemm_cpu(h_ref, h_A, h_B, M, N, K);

  dim3 blockDim(WS, WS);
  int gx = (N + WS - 1) / WS;
  int gy = (M + WS - 1) / WS;
  dim3 gridDim(gx, gy);

  // warmup
  sgemm_kernel2<<<gridDim, blockDim>>>(d_C, d_A, d_B, M, N, K);
  CHECK_CUDA(cudaGetLastError());
  CHECK_CUDA(cudaMemcpy(h_C, d_C, Cbytes, cudaMemcpyDeviceToHost));
  check_results(h_C, h_ref, M * N, 5e-4f);

  int repeat = argc > 1 ? atoi(argv[1]) : 20;
  CudaTimer timer;
  cuda_timer_start(&timer);
  for (int r = 0; r < repeat; r++) {
    sgemm_kernel2<<<gridDim, blockDim>>>(d_C, d_A, d_B, M, N, K);
    CHECK_CUDA(cudaGetLastError());
  }
  float ms = cuda_timer_stop(&timer);
  float avg_ms = ms / repeat;
  double gflops = 2.0 * M * N * K / (avg_ms * 1e6);
  printf("avg %.3f ms | %.1f GFLOPS | ", avg_ms, gflops);

  free(h_ref);
  free(h_A);
  free(h_B);
  free(h_C);
  CHECK_CUDA(cudaFree(d_A));
  CHECK_CUDA(cudaFree(d_B));
  CHECK_CUDA(cudaFree(d_C));
  return 0;
}