From ab248b18b1167753b4f992fc083f34ef6c2fd016 Mon Sep 17 00:00:00 2001
From: Francois Fleuret <francois@fleuret.org>
Date: Wed, 2 Sep 2020 22:13:28 +0200
Subject: [PATCH] Updates from Twitter @michaelcarilli

---
 speed.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/speed.py b/speed.py
index f682a16..10a008f 100755
--- a/speed.py
+++ b/speed.py
@@ -4,20 +4,24 @@ import time, torch
 
 if torch.cuda.is_available():
     device = torch.device('cuda')
+    sync = lambda: torch.cuda.synchronize()
 else:
     device = torch.device('cpu')
+    sync = lambda: None
 
 nb_runs = 10000
-d1, d2, d3 = 50000, 256, 512
+d1, d2, d3 = 2048, 2048, 2048
 
 a, b = torch.rand(d1, d2).to(device), torch.rand(d2, d3).to(device)
 
+sync
 start_time = time.perf_counter()
 for k in range(nb_runs):
     c = a @ b
+sync
 duration = time.perf_counter() - start_time
 
-nb_flop = float(nb_runs * d1 * d2 * d3)
+nb_flop = float(nb_runs * d1 * d2 * d3 * 2) # 1 multiply-and-add is 2 ops
 speed = nb_flop / duration
 
 for u in [ '', 'K', 'M', 'G', 'T', 'P' ]:
-- 
2.39.5