你有没有遇到过这样的情况:跑一个深度学习模型,光是等矩阵运算就花了十几分钟?哪怕是现代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怎么操作浮点矩阵,绝对值得。家里那块打游戏的显卡,也能变成你的计算加速器。