Benchmark *fn* using CUDA events. Returns median latency in microseconds.
(fn, *args, warmup: int = _WARMUP_ITERS, iters: int = _BENCH_ITERS)
| 58 | |
| 59 | |
| 60 | def _benchmark_fn(fn, *args, warmup: int = _WARMUP_ITERS, iters: int = _BENCH_ITERS): |
| 61 | """ |
| 62 | Benchmark *fn* using CUDA events. Returns median latency in microseconds. |
| 63 | """ |
| 64 | # Warmup |
| 65 | for _ in range(warmup): |
| 66 | fn(*args) |
| 67 | torch.cuda.synchronize() |
| 68 | |
| 69 | start_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] |
| 70 | end_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] |
| 71 | |
| 72 | torch.cuda.synchronize() |
| 73 | for i in range(iters): |
| 74 | start_events[i].record() |
| 75 | fn(*args) |
| 76 | end_events[i].record() |
| 77 | torch.cuda.synchronize() |
| 78 | |
| 79 | times_ms = [s.elapsed_time(e) for s, e in zip(start_events, end_events)] |
| 80 | times_ms.sort() |
| 81 | median_ms = times_ms[len(times_ms) // 2] |
| 82 | return median_ms * 1000.0 # convert to microseconds |
| 83 | |
| 84 | |
| 85 | # --------------------------------------------------------------------------- |
no outgoing calls
no test coverage detected