Skip to content

Commit

Permalink
remove torch.cuda.event time calculator
Browse files Browse the repository at this point in the history
  • Loading branch information
amirakb89 committed Jan 16, 2025
1 parent 06e75b3 commit 9b1ec89
Showing 1 changed file with 5 additions and 19 deletions.
24 changes: 5 additions & 19 deletions fbgemm_gpu/bench/bench_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,25 +40,11 @@ def warmup(
) -> None:
indices, offsets, weights = request.unpack_3()
if warmup_ms:
if torch.cuda.is_available():
elapsed_time_ms = 0
torch.cuda.synchronize()
start_events = torch.cuda.Event(enable_timing=True)
end_events = torch.cuda.Event(enable_timing=True)
while elapsed_time_ms < warmup_ms:
start_events.record()
out = func(indices, offsets, weights)
if bwd_only:
out.backward(grad)
end_events.record()
torch.cuda.synchronize()
elapsed_time_ms += start_events.elapsed_time(end_events)
else:
start_time_ms = time.time() * 1000
while time.time() * 1000 - start_time_ms < warmup_ms:
out = func(indices, offsets, weights)
if bwd_only:
out.backward(grad)
start_time_ms = time.time() * 1000
while time.time() * 1000 - start_time_ms < warmup_ms:
out = func(indices, offsets, weights)
if bwd_only:
out.backward(grad)
else:
for _ in range(warmup_runs):
out = func(indices, offsets, weights)
Expand Down

0 comments on commit 9b1ec89

Please sign in to comment.