之前在实验中我们写过c_gpu = a_gpu + b_gpu,CPU 执行这行代码时,并不是一定要等 GPU 把加法彻底做完,CPU 发出任务后,可以继续干自己的事情,CUDA默认一个非常重要的特性:CPU 发起 GPU 工作,GPU 异步执行,它更像是CPU--CUDA Stream--KernelA,B,C--GPU执行,Stream就是组织这些GPU工作的机制之一,这也是我们写torch.cuda.synchronize()的原因,它实现了CPU 等 GPU,把当前已经提交的工作执行完。
现在我们步入今天的正题:
CUDA Stream 是什么
我们可以把它想象为GPU的一条任务队列,task1--task2--task3,同一个stream中任务默认按照顺序执行Kernel A--Kernel B--Kernel C,事实上通常会有多个stream,假设我们有两个独立的任务,如果只有一个stream,只能完成一个在完成另一个,我们使用两个stream,理论上可以节约一半的时间,这就是overlap。
pytorch中如何创建stream?
python
stream = torch.cuda.Stream()
#use the stream work
with torch.cuda.stream(stream):
x = a + b
#多个stream
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()
with torch.cuda.stream(stream1):
x = a + b
with torch.cuda.stream(stream2):
y = c + d
温馨提示:多个stream提供并发机会,不保证一定并发,因为GPU资源是有限的
之前我们使用time.perf_counter()来测量GPU时间,其实有更适合GPU的方式,torch.cuda.Event,为什么它更好一点呢,因为它是直接在GPU execution timeline上标记,而time.perf_counter()是站在CPU,下面是示例代码
python
import torch
N = 10_000_000
a = torch.randn(N, device="cuda")
b = torch.randn(N, device="cuda")
# warm up
for _ in range(10):
c = a + b
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
c = a + b
end.record()
torch.cuda.synchronize()
print(f"GPU time: {start.elapsed_time(end):.3f} ms")
前几节我们
CPU
↓
CUDA Kernel
↓
Grid
↓
Block
↓
Thread
↓
Warp
↓
Memory
现在我们加入stream:
CPU
│
├───────────────┐
│ │
▼ ▼
Stream 0 Stream 1
│ │
▼ ▼
Kernel A Kernel B
│ │
└───────┬───────┘
▼
GPU
GPU 不只是"有很多核心",它还是一个有任务队列、有调度、有异步执行机制的计算系统。