Метод wait_event

Метод wait_event класса Stream добавляет в поток ожидание указанного события. Это означает, что все последующие операции в потоке будут ждать, пока событие не будет зафиксировано. Метод принимает один обязательный параметр - объект события, созданный с помощью Event. Событие должно быть предварительно записано в каком-либо потоке с помощью метода record_event.

Синтаксис

stream.wait_event(event)

Пример

Создадим поток и событие, затем используем wait_event для синхронизации:

import torch

# Создаем поток и событие
stream = torch.cuda.Stream()
event = torch.cuda.Event()

# Запускаем операции в потоке
with torch.cuda.stream(stream):
	t = torch.tensor([1, 2, 3, 4, 5], device='cuda')
	t = t * 2
	# Записываем событие после операций
	event.record(stream)

# В основном потоке ждем завершения события
torch.cuda.current_stream().wait_event(event)

# Теперь операции в основном потоке будут ждать событие
res = torch.tensor([10, 20, 30], device='cuda')
print(res)

Пример

Пример синхронизации двух потоков с помощью wait_event:

import torch

torch.manual_seed(0)

# Создаем два потока
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()
event = torch.cuda.Event()

# Первый поток выполняет вычисления
with torch.cuda.stream(stream1):
	t1 = torch.randn(100, 100, device='cuda')
	t1 = t1 * 2
	event.record(stream1)

# Второй поток ждет завершения первого
with torch.cuda.stream(stream2):
	stream2.wait_event(event)
	t2 = torch.randn(100, 100, device='cuda')
	t2 = t2 + t1

# Синхронизируем все потоки
torch.cuda.synchronize()
print(t2.shape)

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

torch.Size([100, 100])

Пример

Использование wait_event с несколькими событиями для сложной синхронизации:

import torch

torch.manual_seed(0)

# Создаем потоки и события
stream_a = torch.cuda.Stream()
stream_b = torch.cuda.Stream()
stream_main = torch.cuda.current_stream()
event_a = torch.cuda.Event()
event_b = torch.cuda.Event()

# Поток A вычисляет данные
with torch.cuda.stream(stream_a):
	data_a = torch.randn(50, 50, device='cuda')
	data_a = data_a * 3
	event_a.record(stream_a)

# Поток B вычисляет другие данные
with torch.cuda.stream(stream_b):
	data_b = torch.randn(50, 50, device='cuda')
	data_b = data_b + 5
	event_b.record(stream_b)

# Основной поток ждет оба события
stream_main.wait_event(event_a)
stream_main.wait_event(event_b)

# Объединяем результаты
result = data_a + data_b
print(result.shape)

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

torch.Size([50, 50])

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

метод record_event,
который записывает событие в поток
метод query,
который проверяет статус события
метод synchronize,
который блокирует выполнение до завершения всех операций в потоке
метод wait_stream,
который ожидает завершения всех операций в другом потоке