Benchmark a function and return median time in milliseconds. Uses triton.testing.do_bench if available, otherwise manual implementation.
(fn: Callable, warmup: int = 25, rep: int = 100)
| 971 | # ========================================================================= |
| 972 | |
| 973 | def _do_bench(fn: Callable, warmup: int = 25, rep: int = 100) -> float: |
| 974 | """Benchmark a function and return median time in milliseconds. |
| 975 | Uses triton.testing.do_bench if available, otherwise manual implementation.""" |
| 976 | try: |
| 977 | from triton.testing import do_bench |
| 978 | ms = do_bench(fn, warmup=warmup, rep=rep) |
| 979 | return ms |
| 980 | except ImportError: |
| 981 | pass |
| 982 | |
| 983 | # Fallback: manual benchmark |
| 984 | # Warmup |
| 985 | for _ in range(warmup): |
| 986 | fn() |
| 987 | torch.cuda.synchronize() |
| 988 | |
| 989 | times = [] |
| 990 | for _ in range(rep): |
| 991 | start = torch.cuda.Event(enable_timing=True) |
| 992 | end = torch.cuda.Event(enable_timing=True) |
| 993 | start.record() |
| 994 | fn() |
| 995 | end.record() |
| 996 | torch.cuda.synchronize() |
| 997 | times.append(start.elapsed_time(end)) |
| 998 | |
| 999 | times.sort() |
| 1000 | return times[len(times) // 2] # median |
| 1001 | |
| 1002 | |
| 1003 | def run_performance(kernel_fn: Callable, config: dict, gpu: GPUSpec, |