Метод wait

Метод wait класса Event предназначен для синхронизации вычислений на GPU. Он блокирует выполнение текущего потока до тех пор, пока событие, для которого он вызван, не будет зафиксировано на устройстве. Это полезно для измерения времени выполнения операций или для синхронизации потоков при работе с несколькими GPU.

Синтаксис

event.wait()

Метод не принимает параметров и не возвращает значения. После вызова текущий поток приостанавливается до тех пор, пока все операции, предшествующие вызову record для данного события, не будут завершены на GPU.

Пример

Давайте создадим событие, запишем его после выполнения операций на GPU и дождёмся его завершения с помощью метода wait:

import torch

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

# Создаем тензор на GPU
t = torch.tensor([1, 2, 3, 4, 5], device='cuda')

# Выполняем операцию
t = t * 2

# Записываем событие после операции
event.record()

# Ожидаем завершения события
event.wait()

print(t)

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

tensor([2, 4, 6, 8, 10], device='cuda:0')

Пример

Рассмотрим использование wait для синхронизации с измерением времени выполнения операций:

import torch

# Создаем два события для измерения времени
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

# Создаем тензоры на GPU
t1 = torch.ones((1000, 1000), device='cuda')
t2 = torch.ones((1000, 1000), device='cuda')

# Записываем начальное событие
start_event.record()

# Выполняем операцию умножения матриц
res = torch.matmul(t1, t2)

# Записываем конечное событие
end_event.record()

# Ожидаем завершения конечного события
end_event.wait()

# Вычисляем время выполнения
time = start_event.elapsed_time(end_event)
print(f"Time: {time:.2f} ms")

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

Time: 12.34 ms

Пример

Метод wait также можно использовать для синхронизации нескольких потоков, работающих с одним GPU:

import torch
import threading

def worker(event, t):
	# Выполняем операцию на GPU
	t = t * 3
	# Записываем событие
	event.record()

# Создаем событие и тензор
event = torch.cuda.Event()
t = torch.tensor([1, 2, 3], device='cuda')

# Запускаем поток
thread = threading.Thread(target=worker, args=(event, t))
thread.start()

# Ожидаем завершения события в основном потоке
event.wait()
print("Event completed, continuing...")

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

Event completed, continuing...

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

класс Event,
который представляет событие CUDA для синхронизации
метод record,
который записывает событие в поток выполнения
метод query,
который проверяет, произошло ли событие
метод elapsed_time,
который вычисляет время между двумя событиями