矩阵乘法优化全集

/**
 * 矩阵乘法 V1: 朴素实现
 * C = A × B
 * 每个线程计算C的一个元素
 * 
 * 编译: nvcc -o matmul_v1 matmul_v1.cu -arch=sm_80
 */

#include <stdio.h>
#include <cuda_runtime.h>

#define CUDA_CHECK(call) do { cudaError_t e = call; if(e != cudaSuccess) { \
    fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(e)); exit(1); } } while(0)

// 朴素矩阵乘法Kernel
// 每个线程计算C的一个元素
__global__ void matmul_naive(const float* A, const float* B, float* C, int N) {
    // 计算全局行列索引
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (row < N && col < N) {
        float sum = 0.0f;
        // 点积计算
        for (int k = 0; k < N; k++) {
            sum += A[row * N + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

// 主函数
int main() {
    const int N = 1024;
    const int size = N * N * sizeof(float);
    
    // Host内存
    float *h_A = (float*)malloc(size);
    float *h_B = (float*)malloc(size);
    float *h_C = (float*)malloc(size);
    
    // 初始化
    for(int i=0; i// Device内存
    float *d_A, *d_B, *d_C;
    CUDA_CHECK(cudaMalloc(&d_A, size));
    CUDA_CHECK(cudaMalloc(&d_B, size));
    CUDA_CHECK(cudaMalloc(&d_C, size));
    
    CUDA_CHECK(cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice));
    
    // 启动配置
    dim3 block(16, 16);  // 256 threads per block
    dim3 grid((N+15)/16, (N+15)/16);
    
    // 计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start); cudaEventCreate(&stop);
    cudaEventRecord(start);
    
    matmul_naive<<<grid, block>>>(d_A, d_B, d_C, N);
    CUDA_CHECK(cudaDeviceSynchronize());
    
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    float ms = 0;
    cudaEventElapsedTime(&ms, start, stop);
    
    // 计算GFLOPS
    double gflops = 2.0 * N * N * N / (ms / 1000.0) / 1e9;
    printf("V1 Naive: %.2f ms, %.1f GFLOPS\n", ms, gflops);
    
    // 清理
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    free(h_A); free(h_B); free(h_C);
    
    return 0;
}

/**
 * 矩阵乘法 V2: Shared Memory Tiling
 * 
 * 核心思想:
 * 1. 将A和B分块加载到Shared Memory
 * 2. Block内线程协作加载
 * 3. 在Shared Memory中计算点积
 * 4. 复用数据，减少Global Memory访问
 */

#define TILE_SIZE 16

__global__ void matmul_tiled(const float* A, const float* B, float* C, int N) {
    // Shared Memory声明 - Block内所有线程共享
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
    
    // Block和Thread索引
    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;
    
    // 遍历所有TILE
    for (int t = 0; t < N / TILE_SIZE; t++) {
        // 协作加载数据到Shared Memory
        // 每个线程加载一个元素
        As[ty][tx] = A[row * N + (t * TILE_SIZE + tx)];
        Bs[ty][tx] = B[(t * TILE_SIZE + ty) * N + col];
        
        // 同步所有线程，确保数据加载完成
        __syncthreads();
        
        // 在Shared Memory中计算点积
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += As[ty][k] * Bs[k][tx];
        }
        
        // 同步，准备加载下一个TILE
        __syncthreads();
    }
    
    // 写入结果
    C[row * N + col] = sum;
}

/**
 * 矩阵乘法 V3: 向量化加载 + Tiling
 * 
 * 优化点:
 * 1. 使用float4一次加载128位数据
 * 2. 提高内存带宽利用率
 * 3. 减少内存事务数量
 */

__global__ void matmul_vectorized(const float* A, const float* B, float* C, int N) {
    __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;
    
    // 向量化加载A - 每次加载4个float
    int a_col_start = (bx * TILE_SIZE + tx * 4) % N;
    float4 a_vec = reinterpret_cast<float4*>(
        &A[row * N + a_col_start]
    )[0];
    
    // 类似地向量化加载B...
    // 计算部分...
}

/**
 * 矩阵乘法 V4: Tensor Core (WMMA)
 * 
 * 利用NVIDIA Tensor Core进行矩阵乘加运算
 * 一条指令完成16×16×16 = 4096次FMA
 * 
 * 需要: GPU架构 >= Volta (sm_70)
 */

#include <mma.h>
using namespace nvcuda;

__global__ void matmul_tensor_core(half *A, half *B, float *C, int N) {
    // 声明WMMA片段
    // fragment: 载入到Tensor Core的矩阵块
    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;
    
    // 初始化累加器为0
    wmma::fill_fragment(c_frag, 0.0f);
    
    // 计算C的16×16块
    int warpM = (blockIdx.y * blockDim.y + threadIdx.y) / 32;
    int warpN = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    
    // 遍历所有16×16×16块
    for (int k = 0; k < N; k += 16) {
        // 加载A的16×16块
        wmma::load_matrix_sync(a_frag, A + warpM * 16 * N + k, N);
        // 加载B的16×16块
        wmma::load_matrix_sync(b_frag, B + k * N + warpN * 16, N);
        // 矩阵乘加: C += A × B
        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    }
    
    // 存储结果
    wmma::store_matrix_sync(C + warpM * 16 * N + warpN * 16, c_frag, N, wmma::mem_row_major);
}

# 编译所有版本
nvcc -o matmul_v1 matmul_v1.cu
nvcc -o matmul_v2 matmul_v2.cu
nvcc -o matmul_v3 matmul_v3.cu -arch=sm_80
nvcc -o matmul_v4 matmul_v4.cu -arch=sm_80

# 运行对比
./matmul_v1
./matmul_v2
./matmul_v3
./matmul_v4

# 使用Nsight Compute分析
ncu --set full ./matmul_v2