写在前面
学一下 CUDA. 关键在于并行思维. 当然还要了解一些矩阵/线性代数相关的内容, 还有算法/深度学习的一些公式的推导. 这里先做一下 leetgpu 里面的第二题. 矩阵乘法.
https://leetgpu.com/challenges/matrix-multiplication
$M\times N$的矩阵 A 还有$N\times K$的矩阵 B 相乘, 得到$M\times K$ 的矩阵 C. 用 CUDA 实现. 传入的就是每一个矩阵线性化之后的一维数组.
基本实现
#include <cuda_runtime.h>
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N, int K) {
int c = threadIdx.x + blockDim.x * blockIdx.x;
int r = threadIdx.y + blockDim.y * blockIdx.y;
if (r < M and c < K) {
float tmp = 0.0f;
for (int i = 0; i < N; ++i) {
tmp += A[r * N + i] * B[c + i * K];
}
C[r * K + c] = tmp;
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
dim3 threadsPerBlock(16, 16);
dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
(M + threadsPerBlock.y - 1) / threadsPerBlock.y);
matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
cudaDeviceSynchronize();
}