I wanna use CUDA stream in Pytorch to parallel some computations, but I don't know how to do it. For instance, if there's 2 tasks, A and B, need to be parallelized, I wanna do the following things:
stream0 = torch.get_stream()
stream1 = torch.get_stream()
with torch.now_stream(stream0):
// task A
with torch.now_stream(stream1):
// task B
torch.synchronize()
// get A and B's answer
How can I achieve the goal in real python code?
A CUDA stream is a linear sequence of execution that belongs to a specific device, independent from other streams. See CUDA semantics for details. device (torch. device or int, optional) – a device on which to allocate the stream. If device is None (default) or a negative integer, this will use the current device.
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()
# Initialise cuda tensors here. E.g.:
A = torch.rand(1000, 1000, device = ‘cuda’)
B = torch.rand(1000, 1000, device = ‘cuda’)
# Wait for the above tensors to initialise.
torch.cuda.synchronize()
with torch.cuda.stream(s1):
C = torch.mm(A, A)
with torch.cuda.stream(s2):
D = torch.mm(B, B)
# Wait for C and D to be computed.
torch.cuda.synchronize()
# Do stuff with C and D.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With