用CUDA加速浮点矩阵乘法,让计算快到飞起

你有没有遇到过这样的情况:跑一个深度学习模型,光是等矩阵运算就花了十几分钟?哪怕是现代CPU,在面对大规模浮点矩阵乘法时也常常力不从心。这时候,CUDA就能派上大用场了。

为什么选CUDA做矩阵运算?

CUDA是NVIDIA推出的并行计算平台,它能直接调用GPU的成百上千个核心来处理计算任务。而矩阵乘法这种高度可并行的操作,正好是GPU的强项。尤其是浮点型矩阵(比如float32),数据规整、计算密集,特别适合丢给GPU去跑。

举个例子,两个2048×2048的单精度浮点矩阵相乘,CPU可能要忙活好几秒,而一块普通的GeForce RTX 3060,用CUDA加速后,时间能压到几百毫秒以内。

一个简单的CUDA矩阵乘法示例

下面是一个基础的CUDA核函数,用来计算两个N×N浮点矩阵的乘积:

__global__ void matmul_kernel(float* A, 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;
    }
}

这个核函数每个线程负责计算结果矩阵中的一个元素。通过二维线程块和网格的组织方式,把整个矩阵拆成小块并行处理。

主机端代码片段

在调用核函数之前,需要把数据从内存复制到显存:

// 分配显存
float *d_A, *d_B, *d_C;
cudaMalloc(&d_A, N*N*sizeof(float));
cudaMalloc(&d_B, N*N*sizeof(float));
cudaMalloc(&d_C, N*N*sizeof(float));

// 拷贝数据
cudaMemcpy(d_A, h_A, N*N*sizeof(float), cudaMemcpyHostToDevice);
cudaMemcpy(d_B, h_B, N*N*sizeof(float), cudaMemcpyHostToDevice);

// 配置执行配置
dim3 blockSize(16, 16);
dim3 gridSize((N + blockSize.x - 1) / blockSize.x, (N + blockSize.y - 1) / blockSize.y);
matmul_kernel<<<gridSize, blockSize>>>(d_A, d_B, d_C, N);

// 等待完成
cudaDeviceSynchronize();

// 取回结果
cudaMemcpy(h_C, d_C, N*N*sizeof(float), cudaMemcpyDeviceToHost);

虽然写起来比纯C++复杂一点,但一旦跑起来,速度提升是肉眼可见的。

实际应用中的优化思路

上面的例子只是入门级实现。真实项目中,还会用到共享内存来减少全局内存访问次数。比如把矩阵分块加载到shared memory里,避免重复读取显存,效率能再提一截。

另外,cuBLAS库已经提供了高度优化的矩阵乘法函数cublasSgemm,日常开发中直接调用它往往比手写kernel更快更稳定。

如果你经常处理图像处理、神经网络推理或科学计算这类任务,花点时间学学CUDA怎么操作浮点矩阵,绝对值得。家里那块打游戏的显卡,也能变成你的计算加速器。