github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[TESTING] change do_bench defaults

Browse Source
This commit is contained in:
Phil Tillet
2023-04-11 22:03:19 -07:00
parent d7d62ddae9
commit 9530d93504
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/triton/testing.py
View File

@@ -19,8 +19,8 @@ def nvsmi(attrs):
def do_bench(fn, warmup=25, rep=100, grad_to_none=None,
quantiles=None,
fast_flush=False,
return_mode="min"):
fast_flush=True,
return_mode="mean"):
assert return_mode in ["min", "max", "mean", "median"]
import torch
"""
@@ -85,7 +85,10 @@ def do_bench(fn, warmup=25, rep=100, grad_to_none=None,
torch.cuda.synchronize()
times = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)])
if quantiles is not None:
return torch.quantile(times, torch.tensor(quantiles)).tolist()
ret = torch.quantile(times, torch.tensor(quantiles)).tolist()
if len(ret) == 1:
ret = ret[0]
return ret
return getattr(torch, return_mode)(times).item()