我们今天写一个最简单的GPU程序:
Ci = Ai + Bi
c++
for (int i = 0; i < N; i++) {
C[i] = A[i] + B[i];
}
CUDA一个thread负责一个元素
first CUDA kernel
c++
__global__ void vector_add(
float* A,
float* B,
float* C,
int N
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < N) {
C[i] = A[i] + B[i];
}
}
//__global__ 表示这是一个CUDA kernel,可以有CPU发起,然后再GPU上执行.还需要告诉GPU启动多少个block,每个block多少个thread,CUDA有一种非常特殊的调用语句vector_add<<<blocks, threads>>>(A, B, C, N);,代表blocks个blocks,每个block有threads个threads。
//参数AB为输入数组,C为输出数组 N为数组数量
//i负责计算当前thread到底应该处理哪个元素
完整的CUDA程序
c++
#include <iostream>
#include <cuda_runtime.h>
__global__ void vector_add(
float* A,
float* B,
float* C,
int N
) {
//GPU程序
int i = blockIdx.x * blockDim.x + threadIdx.x;//Thread Block block中thread数
if (i < N) {
C[i] = A[i] + B[i];
}
}
int main() {
int N = 1000;
size_t size = N * sizeof(float);
float* A;
float* B;
float* C;
cudaMalloc(&A, size);
//在pytorch中直接torch。rnadn(N,device="cuda"),pytorch帮我们做了大量底层工作,在cuda c++中,我们自己申请GPU memory,cudaMalloc(&A,size),意思是在GPU global memory中分配一块内存,
cudaMalloc(&B, size);
cudaMalloc(&C, size);
vector_add<<<
(N + 255) / 256,
256
>>>(A, B, C, N);//grid
cudaDeviceSynchronize();
//CPU等待GPU完成,和torch.cuda.synchronize()同源
cudaFree(A);
cudaFree(B);
cudaFree(C);
return 0;
}
实现流程:
CPU
│
│ vector_add<<<blocks, threads>>>
▼
CUDA Runtime
│
▼
GPU
│
├── Block 0
│ ├── Thread 0 → C[0]
│ ├── Thread 1 → C[1]
│ ├── Thread 2 → C[2]
│ └── ...
│
├── Block 1
│ ├── Thread 0 → C[256]
│ ├── Thread 1 → C[257]
│ └── ...
│
└── Block 2
...