IQ.Pilot Release Commit @ b6534c0

artifacts/package_sources/tinygrad/extra/gemm/real_pmatmul.py

@@ -0,0 +1,20 @@
+import time
+from tinygrad import Tensor, Device, TinyJit
+from tinygrad.helpers import getenv
+
+if __name__ == "__main__":
+  DEVS = [f"NV:{i}" for i in range(getenv("GPUS", 2))]
+  N = getenv("N", 8192)
+  A = Tensor.rand(N, N).shard(DEVS, 0).realize()
+  B = Tensor.rand(N, N).shard(DEVS, 1).realize()
+  print("***** MUL *****")
+  jmatmul = TinyJit(Tensor.dot)
+  for i in range(10):
+    Device["NV:0"].synchronize()
+    Device["NV:1"].synchronize()
+    st = time.perf_counter()
+    jmatmul(A, B)
+    Device["NV:0"].synchronize()
+    Device["NV:1"].synchronize()
+    et = time.perf_counter()
+    print(f"{(N*N*N*2*1e-12)/(et-st):.2f} TFLOPS")