Matrix Multiplication
358 字
2 分钟
Matrix Multiplication
题面
答案
#include <cuda_runtime.h>
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N, int K) { int k = blockDim.x * blockIdx.x + threadIdx.x; // Col int m = blockDim.y * blockIdx.y + threadIdx.y; // Row int idx = m * K + k; // ALL if ((k < K) && (m < M)) { float sum = 0; for (int n = 0; n < N; n++) { sum += A[m * N + n] * B[n * K + k]; } C[idx] += sum; }}
// 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();}核心循环图解
线程 → 目标位置
每个线程负责计算 C 的一个元素：
int k = blockDim.x * blockIdx.x + threadIdx.x; // C 的列号int m = blockDim.y * blockIdx.y + threadIdx.y; // C 的行号int idx = m * K + k; // 在 C 中的线性索引 k=0 k=1 k=2 ┌──────┬──────┬──────┐m=0 │ T0,0 │ T0,1 │ T0,2 │ ← 每个格子一个线程m=1 │ T1,0 │ T1,1 │ T1,2 │ └──────┴──────┴──────┘索引公式对照
| 代码 | 数学含义 | 解释 |
|---|---|---|
m * N + n | A[m][n] | A 有 N 列，行主序存储 |
n * K + k | B[n][k] | B 有 K 列，行主序存储 |
m * K + k | C[m][k] | C 有 K 列，行主序存储 |
注意
代码中 C[idx] += sum 是累加而非赋值。如果 C 未被调用方清零，会叠加到垃圾值上。常规矩阵乘法应使用 =，除非有意做增量累加。
Matrix Multiplication
https://dongyanzhang.com/posts/leetgpu/matrix-multiplication/