From e721be735e4fafe3d5dafa543aa814585d94aef6 Mon Sep 17 00:00:00 2001 From: Whitney Tsang Date: Fri, 5 Jul 2024 21:57:01 -0400 Subject: [PATCH] Fix `do_bench` (#1567) Fixes #1563 Signed-off-by: Whitney Tsang --- python/triton/testing.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/triton/testing.py b/python/triton/testing.py index 6ef6025c48..91d16dbf40 100644 --- a/python/triton/testing.py +++ b/python/triton/testing.py @@ -22,7 +22,7 @@ def __init__(self, **kwargs): self.record() def record(self): - self.timestamp = time.time_ns() + self.timestamp = time.time_ns() // 1_000_000 def elapsed_time(self, end): return end.timestamp - self.timestamp @@ -155,8 +155,11 @@ def do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, fast_flu for _ in range(5): cache.zero_() fn() + if USE_WALL_TIME: + synchronize() end_event.record() - synchronize() + if not USE_WALL_TIME: + synchronize() estimate_ms = start_event.elapsed_time(end_event) / 5 # compute number of warmup and repeat