Метод wait_stream

Метод wait_stream класса torch.cuda.Stream используется для синхронизации работы двух потоков CUDA. Он заставляет текущий поток ожидать завершения всех операций, запущенных в другом потоке, переданном в качестве аргумента. Это необходимо, когда результат вычислений в одном потоке требуется для дальнейшей работы в другом потоке. Метод принимает один параметр - объект потока, которого нужно дождаться.

Синтаксис

stream.wait_stream(other_stream)

Пример

Давайте создадим два потока и выполним операции с ожиданием:

import torch

torch.manual_seed(0)

s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

t = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device='cuda')

with torch.cuda.stream(s1):
	t1 = t * 2
	t2 = t + 10

s2.wait_stream(s1)

with torch.cuda.stream(s2):
	res = t1 + t2

s2.synchronize()
print(res)

Результат выполнения кода:

tensor([13., 16., 19., 22., 25.], device='cuda:0')

Пример

Покажем, как метод влияет на порядок выполнения операций в разных потоках:

import torch

torch.manual_seed(0)

stream_a = torch.cuda.Stream()
stream_b = torch.cuda.Stream()

t = torch.tensor([1, 2, 3, 4, 5], device='cuda')

with torch.cuda.stream(stream_a):
	t_a = t * 3

stream_b.wait_stream(stream_a)

with torch.cuda.stream(stream_b):
	t_b = t_a + 5

stream_b.synchronize()
print(t_b)

Результат выполнения кода:

tensor([8, 11, 14, 17, 20], device='cuda:0')

Пример

Покажем работу метода с несколькими потоками:

import torch

torch.manual_seed(0)

main_stream = torch.cuda.Stream()
worker_stream = torch.cuda.Stream()

t = torch.tensor([[1, 2, 3], [4, 5, 6]], device='cuda')

with torch.cuda.stream(main_stream):
	t_main = t * 2

worker_stream.wait_stream(main_stream)

with torch.cuda.stream(worker_stream):
	res = t_main + 1

worker_stream.synchronize()
print(res)

Результат выполнения кода:

tensor([
	[3, 5, 7],
	[9, 11, 13],
], device='cuda:0')

Смотрите также

класс Stream,
который представляет поток выполнения CUDA
метод wait_event,
который ожидает завершения события в потоке
метод record_event,
который записывает событие в поток
метод synchronize,
который синхронизирует текущий поток