【infra之路】03_动手写 CUDA Kernel — 向量加法到矩阵乘法

最新推荐文章于 2026-08-18 23:00:00 发布

原创最新推荐文章于 2026-08-18 23:00:00 发布 · 364 阅读

4 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

收录于

infra

前两课讲了 GPU 线程层次和内存层次，这课全部是代码。我们要写 4 个 kernel，一个比一个快，用 profiler 亲眼看到优化的效果。

环境（WSL2 + RTX 5060 Ti）

1. 安装 CUDA Toolkit

# 在 WSL2 内执行
# 先确认 GPU 驱动已安装（Windows 侧安装 NVIDIA 驱动即可，WSL2 自动共享）
nvidia-smi

# 安装 CUDA Toolkit（如果还没装）
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda-repo-wsl-ubuntu-12-8-local_12.8.0-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-12-8-local_12.8.0-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-12-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-8

# 加入环境变量
echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

# 验证
nvcc --version

2. 确认你的 GPU 信息

nvidia-smi
# 你应该看到 RTX 5060 Ti, 8GB VRAM, Blackwell 架构

RTX 5060 Ti 关键参数（参考）：

参数	值
SM 数量	~36
CUDA Cores	~4608
显存	8 GB GDDR7
显存带宽	~448 GB/s
Shared Memory / SM	100 KB (可配置)

文件结构

本课包含以下文件，都在同一个目录下：

lesson3/
├── 01_vector_add.cu       # 热身：向量加法
├── 02_matmul_naive.cu     # 矩阵乘法 v1：naive
├── 03_matmul_shared.cu    # 矩阵乘法 v2：shared memory tiled
├── 04_matmul_register.cu  # 矩阵乘法 v3：register tiling
└── Makefile               # 一键编译

实验 1：向量加法（热身）

最简单的 kernel，目的是熟悉 CUDA 编程的基本流程：分配内存 → 拷贝数据 → 启动 kernel → 拷回结果。

文件: 01_vector_add.cu

/*
 * 实验 1：向量加法 (Vector Addition)
 *
 * 目的：熟悉 CUDA 编程基本流程
 *   1. cudaMalloc 分配 GPU 内存
 *   2. cudaMemcpy 拷贝数据到 GPU
 *   3. <<<grid, block>>> 启动 kernel
 *   4. cudaMemcpy 拷回结果
 *   5. cudaEvent 计时
 *
 * 编译：nvcc -o vec_add 01_vector_add.cu
 * 运行：./vec_add
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

// ============================================================
// CUDA Kernel：向量加法 C = A + B
// ============================================================
// __global__ 表示这是一个在 GPU 上执行、从 CPU 端调用的函数
// 每个线程计算 C 中的一个元素
__global__ void vectorAdd(const float *A, const float *B, float *C, int N) {
    // 计算当前线程的全局 ID
    // threadIdx.x: 线程在当前 block 内的索引 (0 ~ blockDim.x - 1)
    // blockIdx.x:  当前 block 在 grid 内的索引
    // blockDim.x:  每个 block 的线程数
    int tid = threadIdx.x + blockIdx.x * blockDim.x;

    // 边界检查：线程总数可能 > N（因为 block 大小是固定的）
    if (tid < N) {
        C[tid] = A[tid] + B[tid];
    }
}

// ============================================================
// 辅助函数：用 CUDA Event 精确计时
// ============================================================
// cudaEvent 的精度约 0.5 微秒，是 GPU kernel 计时的标准方式
// 比 CPU 侧的 clock() 准确得多（因为 kernel launch 是异步的）

int main() {
    const int N = 1 << 24;  // 16M 个元素，约 64MB per array
    const size_t bytes = N * sizeof(float);

    // ---- 1. 在 CPU (Host) 上分配和初始化数据 ----
    float *h_A = (float *)malloc(bytes);
    float *h_B = (float *)malloc(bytes);
    float *h_C = (float *)malloc(bytes);

    for (int i = 0; i < N; i++) {
        h_A[i] = sinf(i) * sinf(i);
        h_B[i] = cosf(i) * cosf(i);
    }

    // ---- 2. 在 GPU (Device) 上分配内存 ----
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, bytes);  // cudaMalloc(指针的地址, 字节数)
    cudaMalloc(&d_B, bytes);
    cudaMalloc(&d_C, bytes);

    // ---- 3. 将数据从 CPU 拷贝到 GPU ----
    // cudaMemcpy(目标, 来源, 字节数, 方向)
    cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);

    // ---- 4. 配置 kernel 启动参数并执行 ----
    const int blockSize = 256;  // 每个 block 256 个线程
    // 向上取整：确保线程总数 >= N
    const int gridSize = (N + blockSize - 1) / blockSize;

    // 用 cudaEvent 计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    cudaEventRecord(start);  // 记录开始时间
    vectorAdd<<<gridSize, blockSize>>>(d_A, d_B, d_C, N);  // 启动 kernel
    cudaEventRecord(stop);   // 记录结束时间
    cudaEventSynchronize(stop);  // 等 kernel 跑完

    float milliseconds = 0;
    cudaEventElapsedTime(&milliseconds, start, stop);

    printf("向量加法 kernel 耗时: %.3f ms\n", milliseconds);
    printf("元素数量: %d (%.1f M)\n", N, N / 1e6);
    printf("有效带宽: %.1f GB/s\n",
           (3.0 * bytes) / (milliseconds / 1000.0) / 1e9);  // 读A + 读B + 写C

    // ---- 5. 将结果从 GPU 拷回 CPU ----
    cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

    // ---- 6. 验证结果 ----
    // 注意：sin²(x) + cos²(x) = 1
    float maxError = 0.0f;
    for (int i = 0; i < N; i++) {
        maxError = fmaxf(maxError, fabsf(h_C[i] - 1.0f));
    }
    printf("最大误差: %e (应接近 0)\n", maxError);

    // ---- 7. 释放内存 ----
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    free(h_A);
    free(h_B);
    free(h_C);

    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    return 0;
}

编译运行：

nvcc -o vec_add 01_vector_add.cu
./vec_add

在这里插入图片描述

要点：

cudaMalloc 在 GPU 上分配内存，cudaFree 释放
cudaMemcpy 在 CPU 和 GPU 之间拷贝数据
<<<gridSize, blockSize>>> 是 kernel 启动语法
cudaDeviceSynchronize() 等 GPU 算完（kernel launch 是异步的）
用 cudaEvent 做精确计时（微秒级）

实验 2-4：矩阵乘法三级优化

这是本课的核心。我们用同一个矩阵尺寸（N=2048），实现三个版本：

v1: Naive（每个线程算 C 的一个元素）

每个线程独立从 Global Memory 读 A 的一行和 B 的一列，做点积。

问题：A 的第 i 行被 N 个线程各读一次 → Global Memory 读取次数 = N² × N = N³

/*
 * 实验 2：矩阵乘法 v1 — Naive 版本
 *
 * 每个线程计算 C 的一个元素 C[row][col]
 * 该线程需要读 A 的整行（N 个元素）和 B 的整列（N 个元素）
 *
 * 问题：A 的第 i 行被 N 个线程各读一次 → Global Memory 读取次数 = 2 × N³
 * 这是最慢的实现，但逻辑最清晰，作为优化基线。
 *
 * 编译：nvcc -O3 -o matmul_naive 02_matmul_naive.cu
 * 运行：./matmul_naive
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define N 2048  // 矩阵维度：2048 × 2048

// ============================================================
// Kernel：Naive 矩阵乘法 C = A × B
// ============================================================
// 每个线程计算 C[row][col] 的一个元素
// 线程需要从 Global Memory 读 A[row][*] 整行 + B[*][col] 整列
__global__ void matMulNaive(const float *A, const float *B, float *C, int width) {
    // 用 2D block/grid 组织线程
    int row = threadIdx.y + blockIdx.y * blockDim.y;  // 行索引
    int col = threadIdx.x + blockIdx.x * blockDim.x;  // 列索引

    if (row < width && col < width) {
        float sum = 0.0f;
        // 沿 K 维度（公共维度）做点积
        for (int k = 0; k < width; k++) {
            // A[row][k] = A[row * width + k]  （行优先存储）
            // B[k][col] = B[k * width + col]
            sum += A[row * width + k] * B[k * width + col];
        }
        C[row * width + col] = sum;
    }
}

// ============================================================
// CPU 参考实现（用于验证正确性）
// ============================================================
void matMulCPU(const float *A, const float *B, float *C, int width) {
    for (int i = 0; i < width; i++) {
        for (int j = 0; j < width; j++) {
            float sum = 0.0f;
            for (int k = 0; k < width; k++) {
                sum += A[i * width + k] * B[k * width + j];
            }
            C[i * width + j] = sum;
        }
    }
}

int main() {
    const int width = N;
    const size_t bytes = N * N * sizeof(float);

    printf("=== 矩阵乘法 v1: Naive ===\n");
    printf("矩阵大小: %d × %d (%.1f M 元素)\n", N, N, (float)(N * N) / 1e6);
    printf("FLOPs: %.2f G (2 × N³)\n", 2.0 * N * N * N / 1e9);

    // ---- 分配和初始化 ----
    float *h_A = (float *)malloc(bytes);
    float *h_B = (float *)malloc(bytes);
    float *h_C = (float *)malloc(bytes);

    // 用小数随机初始化，避免数值溢出
    srand(42);
    for (int i = 0; i < N * N; i++) {
        h_A[i] = (float)(rand() % 100) / 100.0f;
        h_B[i] = (float)(rand() % 100) / 100.0f;
    }

    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, bytes);
    cudaMalloc(&d_B, bytes);
    cudaMalloc(&d_C, bytes);

    cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);

    // ---- 启动 kernel ----
    // 2D block: 16×16 = 256 线程/block
    // 2D grid: 覆盖整个 N×N 矩阵
    dim3 blockSize(16, 16);
    dim3 gridSize((N + blockSize.x - 1) / blockSize.x,
                  (N + blockSize.y - 1) / blockSize.y);

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    cudaEventRecord(start);
    matMulNaive<<<gridSize, blockSize>>>(d_A, d_B, d_C, width);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);

    float ms = 0;
    cudaEventElapsedTime(&ms, start, stop);

    // 计算性能指标
    double gflops = 2.0 * N * N * N / (ms / 1000.0) / 1e9;
    double bandwidth = (3.0 * N * N * sizeof(float) * N) / (ms / 1000.0) / 1e9;
    // 带宽计算说明：每个元素需要读 A 的 1 行 + B 的 1 列 ≈ 2N 次 float 读取
    // 简化：总读取 ≈ N³ × sizeof(float)（每个点积读 2N 个 float，共 N² 个点积）

    printf("Kernel 耗时: %.3f ms\n", ms);
    printf("性能: %.1f GFLOPS\n", gflops);

    cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

    // ---- 验证（只检查一小部分，CPU 太慢不能算完整的 2048×2048）----
    // 用小矩阵快速验证
    const int V = 64;
    float vA[V * V], vB[V * V], vC_cpu[V * V], vC_gpu[V * V];
    for (int i = 0; i < V * V; i++) {
        vA[i] = h_A[i];
        vB[i] = h_B[i];
    }
    matMulCPU(vA, vB, vC_cpu, V);

    dim3 vBlock(16, 16);
    dim3 vGrid((V + 15) / 16, (V + 15) / 16);
    float *d_vA, *d_vB, *d_vC;
    cudaMalloc(&d_vA, V * V * sizeof(float));
    cudaMalloc(&d_vB, V * V * sizeof(float));
    cudaMalloc(&d_vC, V * V * sizeof(float));
    cudaMemcpy(d_vA, vA, V * V * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_vB, vB, V * V * sizeof(float), cudaMemcpyHostToDevice);
    matMulNaive<<<vGrid, vBlock>>>(d_vA, d_vB, d_vC, V);
    cudaMemcpy(vC_gpu, d_vC, V * V * sizeof(float), cudaMemcpyDeviceToHost);

    float maxErr = 0;
    for (int i = 0; i < V * V; i++) {
        maxErr = fmaxf(maxErr, fabsf(vC_cpu[i] - vC_gpu[i]));
    }
    printf("验证 (64×64): 最大误差 = %e %s\n", maxErr,
           maxErr < 1e-3 ? "✓ PASS" : "✗ FAIL");

    // ---- 清理 ----
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    cudaFree(d_vA); cudaFree(d_vB); cudaFree(d_vC);
    free(h_A); free(h_B); free(h_C);
    cudaEventDestroy(start); cudaEventDestroy(stop);

    return 0;
}

在这里插入图片描述

v2: Shared Memory Tiled

把矩阵沿 K 维度分块，每个 Block 协作加载 tile 到 Shared Memory，Block 内线程从 Shared Memory 读。

收益：Global Memory 读取次数 = N² × N / TILE_SIZE = N³/32（TILE_SIZE=32 时，减少 32 倍）

/*
 * 实验 3：矩阵乘法 v2 — Shared Memory Tiled
 *
 * 核心优化：把 A 和 B 的 tile 加载到 Shared Memory，
 * Block 内所有线程从 Shared Memory 读取，大幅减少 Global Memory 访问。
 *
 * TILE_SIZE = 32：每个 Block 有 32×32 = 1024 个线程
 * 每个线程仍只算 C 的 1 个元素
 *
 * 编译：nvcc -O3 -o matmul_shared 03_matmul_shared.cu
 * 运行：./matmul_shared
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define N 2048
#define TILE_SIZE 32  // tile 维度，也是 block 的维度

// ============================================================
// Kernel：Shared Memory Tiled 矩阵乘法
// ============================================================
__global__ void matMulShared(const float *A, const float *B, float *C, int width) {
    // ① 声明 Shared Memory：每个 Block 分配两块 TILE×TILE
    // __shared__ 表示这块内存在 SM 芯片上，Block 内所有线程共享
    __shared__ float sA[TILE_SIZE][TILE_SIZE];
    __shared__ float sB[TILE_SIZE][TILE_SIZE];

    // ② 当前线程负责 C 的哪个元素
    int row = threadIdx.y + blockIdx.y * TILE_SIZE;
    int col = threadIdx.x + blockIdx.x * TILE_SIZE;

    float sum = 0.0f;

    // ③ 沿 K 维度分块遍历
    // 每次加载一个 TILE×TILE 的 tile 到 Shared Memory
    // 共需要 width / TILE_SIZE 次迭代
    for (int t = 0; t < (width + TILE_SIZE - 1) / TILE_SIZE; t++) {
        // ④ 协作加载：每个线程负责从 Global Memory 搬一个元素
        //
        // 加载 A 的 tile：当前行 row，第 t 块的列
        // A 的坐标: [row][t * TILE_SIZE + threadIdx.x]
        int aCol = t * TILE_SIZE + threadIdx.x;
        int aRow = row;
        if (aRow < width && aCol < width)
            sA[threadIdx.y][threadIdx.x] = A[aRow * width + aCol];
        else
            sA[threadIdx.y][threadIdx.x] = 0.0f;  // 边界填充 0

        // 加载 B 的 tile：第 t 块的行，当前列 col
        // B 的坐标: [t * TILE_SIZE + threadIdx.y][col]
        int bRow = t * TILE_SIZE + threadIdx.y;
        int bCol = col;
        if (bRow < width && bCol < width)
            sB[threadIdx.y][threadIdx.x] = B[bRow * width + bCol];
        else
            sB[threadIdx.y][threadIdx.x] = 0.0f;

        // ⑤ 屏障同步：等 Block 内所有线程都加载完了
        // 如果没有这个，有些线程可能还没写完 Shared Memory，其他线程就开始读了
        __syncthreads();

        // ⑥ 从 Shared Memory 做计算（~20 cycles vs Global Memory ~300+ cycles）
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += sA[threadIdx.y][k] * sB[k][threadIdx.x];
        }

        // ⑦ 第二次同步：等所有线程算完了，才能加载下一个 tile（会覆盖 sA/sB）
        __syncthreads();
    }

    // ⑧ 写回结果到 Global Memory
    if (row < width && col < width) {
        C[row * width + col] = sum;
    }
}

// ============================================================
// CPU 参考实现
// ============================================================
void matMulCPU(const float *A, const float *B, float *C, int width) {
    for (int i = 0; i < width; i++) {
        for (int j = 0; j < width; j++) {
            float sum = 0.0f;
            for (int k = 0; k < width; k++) {
                sum += A[i * width + k] * B[k * width + j];
            }
            C[i * width + j] = sum;
        }
    }
}

int main() {
    const int width = N;
    const size_t bytes = N * N * sizeof(float);

    printf("=== 矩阵乘法 v2: Shared Memory Tiled ===\n");
    printf("矩阵大小: %d × %d\n", N, N);
    printf("Tile 大小: %d × %d\n", TILE_SIZE, TILE_SIZE);
    printf("每个 Block 线程数: %d\n", TILE_SIZE * TILE_SIZE);

    // ---- 分配和初始化 ----
    float *h_A = (float *)malloc(bytes);
    float *h_B = (float *)malloc(bytes);
    float *h_C = (float *)malloc(bytes);

    srand(42);
    for (int i = 0; i < N * N; i++) {
        h_A[i] = (float)(rand() % 100) / 100.0f;
        h_B[i] = (float)(rand() % 100) / 100.0f;
    }

    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, bytes);
    cudaMalloc(&d_B, bytes);
    cudaMalloc(&d_C, bytes);

    cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);

    // ---- 启动 kernel ----
    // Block: 32×32 = 1024 线程
    // Grid: 覆盖 N×N
    dim3 blockSize(TILE_SIZE, TILE_SIZE);
    dim3 gridSize((N + TILE_SIZE - 1) / TILE_SIZE,
                  (N + TILE_SIZE - 1) / TILE_SIZE);

    printf("Grid: %d × %d blocks\n", gridSize.x, gridSize.y);

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    // warmup：第一次运行可能有驱动初始化开销
    matMulShared<<<gridSize, blockSize>>>(d_A, d_B, d_C, width);
    cudaDeviceSynchronize();

    // 正式计时
    cudaEventRecord(start);
    matMulShared<<<gridSize, blockSize>>>(d_A, d_B, d_C, width);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);

    float ms = 0;
    cudaEventElapsedTime(&ms, start, stop);

    double gflops = 2.0 * N * N * N / (ms / 1000.0) / 1e9;
    printf("Kernel 耗时: %.3f ms\n", ms);
    printf("性能: %.1f GFLOPS\n", gflops);

    cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

    // ---- 验证 ----
    const int V = 64;
    float vA[V * V], vB[V * V], vC_cpu[V * V], vC_gpu[V * V];
    for (int i = 0; i < V * V; i++) {
        vA[i] = h_A[i];
        vB[i] = h_B[i];
    }
    matMulCPU(vA, vB, vC_cpu, V);

    dim3 vBlock(TILE_SIZE, TILE_SIZE);
    dim3 vGrid((V + TILE_SIZE - 1) / TILE_SIZE, (V + TILE_SIZE - 1) / TILE_SIZE);
    float *d_vA, *d_vB, *d_vC;
    cudaMalloc(&d_vA, V * V * sizeof(float));
    cudaMalloc(&d_vB, V * V * sizeof(float));
    cudaMalloc(&d_vC, V * V * sizeof(float));
    cudaMemcpy(d_vA, vA, V * V * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_vB, vB, V * V * sizeof(float), cudaMemcpyHostToDevice);
    matMulShared<<<vGrid, vBlock>>>(d_vA, d_vB, d_vC, V);
    cudaMemcpy(vC_gpu, d_vC, V * V * sizeof(float), cudaMemcpyDeviceToHost);

    float maxErr = 0;
    for (int i = 0; i < V * V; i++) {
        maxErr = fmaxf(maxErr, fabsf(vC_cpu[i] - vC_gpu[i]));
    }
    printf("验证 (64×64): 最大误差 = %e %s\n", maxErr,
           maxErr < 1e-3 ? "✓ PASS" : "✗ FAIL");

    // ---- 清理 ----
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    cudaFree(d_vA); cudaFree(d_vB); cudaFree(d_vC);
    free(h_A); free(h_B); free(h_C);
    cudaEventDestroy(start); cudaEventDestroy(stop);

    return 0;
}

在这里插入图片描述

v3: Register Tiling（进一步优化）

每个线程不只算 C 的一个元素，而是算 TM×TN 个小块。加载到 Shared Memory 后，每个线程把 A 和 B 的一行/列缓存到寄存器，进一步减少 Shared Memory 读取。

收益：Shared Memory 读取次数也大幅减少，计算强度（Arithmetic Intensity）更高。

/*
 * 实验 4：矩阵乘法 v3 — Register Tiling（寄存器分块）
 *
 * 在 Shared Memory Tiled 的基础上进一步优化：
 * - 每个线程不再只算 C 的 1 个元素，而是算 TM × TN 个小块（这里 TM=TN=4）
 * - 从 Shared Memory 加载后，把 A 的一行和 B 的一列缓存到寄存器
 * - 大幅减少 Shared Memory 的读取次数
 *
 * 这是 cuBLAS 等高性能库的基本思路（当然它们还有更多优化）。
 *
 * 编译：nvcc -O3 -o matmul_register 04_matmul_register.cu
 * 运行：./matmul_register
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define N 2048

// ============================================================
// 分块参数
// ============================================================
// BM, BN: Block 负责的 C 的子矩阵大小（一个 Block 算 C 的 BM × BN 区域）
#define BM 64
#define BN 64

// BK: 沿 K 维度的 tile 大小
#define BK 16

// TM, TN: 每个线程负责的子块大小
// 每线程算 TM × TN = 4 × 4 = 16 个 C 的元素
#define TM 4
#define TN 4

// Block 线程数: (BM/TM) × (BN/TN) = 16 × 16 = 256
// 即 Block 是 16×16 的 2D grid，每个线程负责一个 4×4 的小块

// ============================================================
// Kernel：Register Tiled 矩阵乘法
// ============================================================
__global__ void matMulRegister(const float *A, const float *B, float *C, int width) {
    // Block 负责的 C 的子矩阵的起始行列
    const int blockRow = blockIdx.y * BM;
    const int blockCol = blockIdx.x * BN;

    // 线程在 Block 内的坐标
    const int tx = threadIdx.x;  // 0 ~ BN/TN - 1 = 0 ~ 15
    const int ty = threadIdx.y;  // 0 ~ BM/TM - 1 = 0 ~ 15

    // 该线程负责的 C 的 TM×TN 子块的起始位置（相对于 block 起点）
    const int threadRow = ty * TM;
    const int threadCol = tx * TN;

    // ① Shared Memory：加载 A 和 B 的 tile
    __shared__ float sA[BM][BK];  // 64 × 16
    __shared__ float sB[BK][BN];  // 16 × 64

    // ② 寄存器：存储当前线程的累积结果（TM × TN = 16 个 float）
    float threadResults[TM][TN] = {0.0f};

    // ③ 沿 K 维度分块遍历
    for (int t = 0; t < (width + BK - 1) / BK; t++) {
        // ---- 协作加载 A 的 tile 到 Shared Memory ----
        // A 的 tile: [blockRow ~ blockRow+BM][t*BK ~ t*BK+BK] = 64 × 16
        // 256 个线程加载 64×16 = 1024 个元素 → 每线程加载 4 个
        for (int loadOffset = 0; loadOffset < BM * BK; loadOffset += blockDim.x * blockDim.y) {
            int loadIdx = loadOffset + ty * blockDim.x + tx;
            if (loadIdx < BM * BK) {
                int loadRow = loadIdx / BK;
                int loadCol = loadIdx % BK;
                int globalRow = blockRow + loadRow;
                int globalCol = t * BK + loadCol;
                sA[loadRow][loadCol] = (globalRow < width && globalCol < width)
                                       ? A[globalRow * width + globalCol] : 0.0f;
            }
        }

        // ---- 协作加载 B 的 tile 到 Shared Memory ----
        // B 的 tile: [t*BK ~ t*BK+BK][blockCol ~ blockCol+BN] = 16 × 64
        // 16×64 = 1024 个元素 → 每线程加载 4 个
        for (int loadOffset = 0; loadOffset < BK * BN; loadOffset += blockDim.x * blockDim.y) {
            int loadIdx = loadOffset + ty * blockDim.x + tx;
            if (loadIdx < BK * BN) {
                int loadRow = loadIdx / BN;
                int loadCol = loadIdx % BN;
                int globalRow = t * BK + loadRow;
                int globalCol = blockCol + loadCol;
                sB[loadRow][loadCol] = (globalRow < width && globalCol < width)
                                       ? B[globalRow * width + globalCol] : 0.0f;
            }
        }

        __syncthreads();

        // ---- 在寄存器内做计算 ----
        // 对 K 维度内的每个 k：
        for (int k = 0; k < BK; k++) {
            // 把 A 的一行加载到寄存器（TM 个值）
            // A 的行：sA[threadRow + 0..TM-1][k]
            float regA[TM];
            #pragma unroll
            for (int m = 0; m < TM; m++) {
                regA[m] = sA[threadRow + m][k];
            }

            // 把 B 的一列加载到寄存器（TN 个值）
            // B 的列：sB[k][threadCol + 0..TN-1]
            float regB[TN];
            #pragma unroll
            for (int n = 0; n < TN; n++) {
                regB[n] = sB[k][threadCol + n];
            }

            // 外积累加：TM × TN 次乘加
            #pragma unroll
            for (int m = 0; m < TM; m++) {
                #pragma unroll
                for (int n = 0; n < TN; n++) {
                    threadResults[m][n] += regA[m] * regB[n];
                }
            }
        }

        __syncthreads();
    }

    // ④ 写回结果到 Global Memory
    #pragma unroll
    for (int m = 0; m < TM; m++) {
        int globalRow = blockRow + threadRow + m;
        if (globalRow < width) {
            #pragma unroll
            for (int n = 0; n < TN; n++) {
                int globalCol = blockCol + threadCol + n;
                if (globalCol < width) {
                    C[globalRow * width + globalCol] = threadResults[m][n];
                }
            }
        }
    }
}

// ============================================================
// CPU 参考实现
// ============================================================
void matMulCPU(const float *A, const float *B, float *C, int width) {
    for (int i = 0; i < width; i++) {
        for (int j = 0; j < width; j++) {
            float sum = 0.0f;
            for (int k = 0; k < width; k++) {
                sum += A[i * width + k] * B[k * width + j];
            }
            C[i * width + j] = sum;
        }
    }
}

int main() {
    const int width = N;
    const size_t bytes = N * N * sizeof(float);

    printf("=== 矩阵乘法 v3: Register Tiled ===\n");
    printf("矩阵大小: %d × %d\n", N, N);
    printf("Block tile: %d × %d, 每线程 tile: %d × %d\n", BM, BN, TM, TN);
    printf("Block 线程数: %d × %d = %d\n", BN / TN, BM / TM,
           (BN / TN) * (BM / TM));

    // ---- 分配和初始化 ----
    float *h_A = (float *)malloc(bytes);
    float *h_B = (float *)malloc(bytes);
    float *h_C = (float *)malloc(bytes);

    srand(42);
    for (int i = 0; i < N * N; i++) {
        h_A[i] = (float)(rand() % 100) / 100.0f;
        h_B[i] = (float)(rand() % 100) / 100.0f;
    }

    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, bytes);
    cudaMalloc(&d_B, bytes);
    cudaMalloc(&d_C, bytes);

    cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);

    // ---- 启动 kernel ----
    dim3 blockSize(BN / TN, BM / TM);  // 16 × 16 = 256
    dim3 gridSize((N + BN - 1) / BN, (N + BM - 1) / BM);

    printf("Grid: %d × %d blocks\n", gridSize.x, gridSize.y);

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    // warmup
    matMulRegister<<<gridSize, blockSize>>>(d_A, d_B, d_C, width);
    cudaDeviceSynchronize();

    // 正式计时
    cudaEventRecord(start);
    matMulRegister<<<gridSize, blockSize>>>(d_A, d_B, d_C, width);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);

    float ms = 0;
    cudaEventElapsedTime(&ms, start, stop);

    double gflops = 2.0 * N * N * N / (ms / 1000.0) / 1e9;
    printf("Kernel 耗时: %.3f ms\n", ms);
    printf("性能: %.1f GFLOPS\n", gflops);

    cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

    // ---- 验证 ----
    const int V = 64;
    float vA[V * V], vB[V * V], vC_cpu[V * V], vC_gpu[V * V];
    for (int i = 0; i < V * V; i++) {
        vA[i] = h_A[i];
        vB[i] = h_B[i];
    }
    matMulCPU(vA, vB, vC_cpu, V);

    dim3 vBlock(BN / TN, BM / TM);
    dim3 vGrid((V + BN - 1) / BN, (V + BM - 1) / BM);
    float *d_vA, *d_vB, *d_vC;
    cudaMalloc(&d_vA, V * V * sizeof(float));
    cudaMalloc(&d_vB, V * V * sizeof(float));
    cudaMalloc(&d_vC, V * V * sizeof(float));
    cudaMemcpy(d_vA, vA, V * V * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_vB, vB, V * V * sizeof(float), cudaMemcpyHostToDevice);
    matMulRegister<<<vGrid, vBlock>>>(d_vA, d_vB, d_vC, V);
    cudaMemcpy(vC_gpu, d_vC, V * V * sizeof(float), cudaMemcpyDeviceToHost);

    float maxErr = 0;
    for (int i = 0; i < V * V; i++) {
        maxErr = fmaxf(maxErr, fabsf(vC_cpu[i] - vC_gpu[i]));
    }
    printf("验证 (64×64): 最大误差 = %e %s\n", maxErr,
           maxErr < 1e-3 ? "✓ PASS" : "✗ FAIL");

    // ---- 清理 ----
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    cudaFree(d_vA); cudaFree(d_vB); cudaFree(d_vC);
    free(h_A); free(h_B); free(h_C);
    cudaEventDestroy(start); cudaEventDestroy(stop);

    return 0;
}

在这里插入图片描述

编译与运行

# 一键编译所有程序
make

# 或者单独编译
nvcc -O3 -o matmul_naive 02_matmul_naive.cu
nvcc -O3 -o matmul_shared 03_matmul_shared.cu
nvcc -O3 -o matmul_register 04_matmul_register.cu

# 运行
./matmul_naive
./matmul_shared
./matmul_register

性能分析

# 用 nsight compute 分析（如果可用）
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,dram__throughput.avg.pct_of_peak_sustained_elapsed ./matmul_naive
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,dram__throughput.avg.pct_of_peak_sustained_elapsed ./matmul_shared

# 或者用 nvprof（旧版）
nvprof ./matmul_naive

预期结果

版本	大致时间 (N=2048)	相对加速	瓶颈
Naive	~50-80 ms	1x	Global Memory 带宽
Shared Tiled	~10-20 ms	3-5x	Shared Memory bank conflict
Register Tiled	~3-8 ms	8-15x	寄存器数量 / occupancy

注：具体数值取决于你的 GPU 和 CUDA 版本，重点是看相对加速比和瓶颈转移。

本课小结

版本	优化手段	减少的访问
Naive	无	—
Shared Tiled	Shared Memory 缓存	Global Memory → Shared Memory
Register Tiled	寄存器缓存 + 每线程多元素	Shared Memory → Register

这三个版本就是 GPU 优化的核心范式：把数据从慢的内存搬到快的内存，最大化复用。FlashAttention、vLLM 的 PagedAttention 本质上都在做同样的事，只是搬的数据和层级不同。

标签

#矩阵 #线性代数 #infra