4. 使用寄存器优化矩阵乘法

任务目标

在GPU矩阵乘法中，全局内存与共享内存的带宽远低于寄存器文件。本案例通过将内层循环中反复使用的矩阵分块数据显式缓存到寄存器，并结合循环展开，消除对共享内存的冗余访问，提升计算密度。完成一个32×32分块的SGEMM内核，并与朴素共享内存版本对比指令数与访存行为。

核心实现

// 编译: nvcc -arch=sm_80 -O3 reg_tiled_sgemm.cu -o reg_tiled_sgemm
// 运行: ./reg_tiled_sgemm  (矩阵尺寸固定为 1024×1024)

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

#define TILE 32
#define BLOCK 16

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

    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;

    float acc[BLOCK][BLOCK] = {{0.0f}};  // 每个线程计算 16×16 输出

    for (int tile = 0; tile < K / TILE; ++tile) {
        // 协作加载 A 与 B 的当前分块到共享内存
        int a_global = (blockIdx.y * TILE + threadIdx.y) * K +
                       (tile * TILE + threadIdx.x);
        int b_global = (tile * TILE + threadIdx.y) * N +
                       (blockIdx.x * TILE + threadIdx.x);
        As[threadIdx.y][threadIdx.x] = A[a_global];
        Bs[threadIdx.y][threadIdx.x] = B[b_global];
        __syncthreads();

        // 寄存器缓存 + 循环展开：每个线程从共享内存读取一次到寄存器
        float a_reg[BLOCK], b_reg[BLOCK];
        #pragma unroll
        for (int k = 0; k < TILE; ++k) {
            #pragma unroll
            for (int i = 0; i < BLOCK; ++i)
                a_reg[i] = As[threadIdx.y * BLOCK + i][k];
            #pragma unroll
            for (int j = 0; j < BLOCK; ++j)
                b_reg[j] = Bs[k][threadIdx.x * BLOCK + j];

            #pragma unroll
            for (int i = 0; i < BLOCK; ++i)
                #pragma unroll
                for (int j = 0; j < BLOCK; ++j)
                    acc[i][j] += a_reg[i] * b_reg[j];
        }
        __syncthreads();
    }

    // 写回全局内存
    #pragma unroll
    for (int i = 0; i < BLOCK; ++i)
        #pragma unroll
        for (int j = 0; j < BLOCK; ++j)
            C[(row + i) * N + col + j] = acc[i][j];
}

int main() {
    const int M = 1024, N = 1024, K = 1024;
    size_t bytes = M * N * sizeof(float);

    float *h_A, *h_B, *h_C;
    cudaMallocHost(&h_A, bytes);
    cudaMallocHost(&h_B, bytes);
    cudaMallocHost(&h_C, bytes);

    for (int i = 0; i < M * K; ++i) h_A[i] = 1.0f;
    for (int i = 0; i < K * N; ++i) h_B[i] = 1.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);

    dim3 block(BLOCK, BLOCK);
    dim3 grid(N / TILE, M / TILE);
    sgemm_reg_tiled<<<grid, block>>>(d_A, d_B, d_C, M, N, K);
    cudaDeviceSynchronize();

    cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

    // 验证：全1矩阵相乘，C[i][j] 应等于 K
    float max_err = 0.0f;
    for (int i = 0; i < M * N; ++i)
        max_err = fmaxf(max_err, fabsf(h_C[i] - K));
    printf("Max error: %f\n", max_err);
    printf("验证%s\n", max_err < 1e-4 ? "通过" : "失败");

    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    cudaFreeHost(h_A); cudaFreeHost(h_B); cudaFreeHost(h_C);
    return 0;
}

关键设计说明

• 寄存器缓存：每个线程负责16×16输出子块。内层循环中，先从共享内存将As的一列（16个元素）和Bs的一行（16个元素）读入寄存器数组a_reg/b_reg，再执行16×16次乘加。这样每个共享内存元素只被读取一次，而非每次乘加都访问共享内存。
• 循环展开：#pragma unroll让编译器将内层循环完全展开，消除循环控制开销，使乘加指令连续发射，隐藏流水线延迟。
• 线程映射：每个block为16×16线程，加载32×32的A/B分块到共享内存，计算32×32输出分块。线程(ty,tx)负责输出行[ty*16, ty*16+16)、列[tx*16, tx*16+16)。

运行与验证

nvcc -arch=sm_80 -O3 reg_tiled_sgemm.cu -o reg_tiled_sgemm
./reg_tiled_sgemm

程序使用全1矩阵作为输入，因此数学上C[i][j] = K = 1024。输出Max error应接近0，并打印“验证通过”。若需对比优化效果，可使用nvcc --ptxas-options=-v查看寄存器用量与局部内存溢出情况，或使用ncu分析共享内存访问次数。

预期输出

Max error: 0.000000
验证通过

常见问题

• 寄存器溢出：若编译时提示寄存器压力过大（spill），可减少BLOCK至8（每个线程计算8×8），或使用__launch_bounds__(256, 2)限制线程数。
• 边界处理：本实现要求M, N, K均为32的整数倍。实际应用中需添加边界判断或补齐数据。
• bank冲突：共享内存访问模式As[ty*16+i][k]在固定k时，同一warp的线程访问不同行，无bank冲突；Bs[k][tx*16+j]同理。