forked from ccf-ai-infra/Intro-ops
21 lines
509 B
Python
21 lines
509 B
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
|
|
import torch
|
|
|
|
|
|
def cuda_time_ms(fn: Callable[[], None], *, warmup: int = 10, iterations: int = 100) -> float:
|
|
for _ in range(warmup):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
start.record()
|
|
for _ in range(iterations):
|
|
fn()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
return start.elapsed_time(end) / iterations
|
|
|