IQ.Pilot Release Commit @ 0798119

tinygrad_repo/test/speed/external_test_copy_speed.py

@@ -0,0 +1,91 @@
+import unittest, numpy as np, os
+from tinygrad import Tensor, Device, TinyJit
+from tinygrad.helpers import Timing, getenv
+import multiprocessing.shared_memory as shared_memory
+
+N = getenv("NSZ", 256)
+class TestCopySpeed(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls): Device[Device.DEFAULT].synchronize()
+
+  def testCopySHMtoDefault(self):
+    s = shared_memory.SharedMemory(name="test_X", create=True, size=N*N*4)
+    s.close()
+    if os.path.exists("/dev/shm"):
+      t = Tensor.empty(N, N, device="disk:/dev/shm/test_X").realize()
+    else:
+      t = Tensor.empty(N, N, device="disk:shm:test_X").realize()
+    for _ in range(3):
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        with Timing("queue: "):
+          t.to(Device.DEFAULT).realize()
+        Device[Device.DEFAULT].synchronize()
+    s.unlink()
+
+  def testCopyCPUtoDefault(self):
+    t = Tensor.ones(N, N, device="CPU").contiguous().realize()
+    print(f"buffer: {t.nbytes()*1e-9:.2f} GB")
+    for _ in range(3):
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        with Timing("queue: "):
+          t.to(Device.DEFAULT).realize()
+        Device[Device.DEFAULT].synchronize()
+
+  def testCopyCPUtoDefaultFresh(self):
+    print("fresh copy")
+    for _ in range(3):
+      t = Tensor.ones(N, N, device="CPU").contiguous().realize()
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"): # noqa: F821
+        with Timing("queue: "):
+          t.to(Device.DEFAULT).realize()
+        Device[Device.DEFAULT].synchronize()
+      del t
+
+  def testCopyDefaulttoCPU(self):
+    t = Tensor.ones(N, N).contiguous().realize()
+    print(f"buffer: {t.nbytes()*1e-9:.2f} GB")
+    for _ in range(3):
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        t.to('CPU').realize()
+
+  def testCopyDefaulttoCPUJit(self):
+    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
+
+    @TinyJit
+    def _do_copy(t): return t.to('CPU').realize()
+
+    t = Tensor.randn(N, N).contiguous().realize()
+    Device[Device.DEFAULT].synchronize()
+    for _ in range(5):
+      with Timing(f"copy {Device.DEFAULT} -> CPU {t.nbytes()/(1024**2)}M:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        x = _do_copy(t)
+        Device[Device.DEFAULT].synchronize()
+      np.testing.assert_equal(t.numpy(), x.numpy())
+
+  def testCopyCPUtoDefaultJit(self):
+    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
+
+    @TinyJit
+    def _do_copy(x): return x.to(Device.DEFAULT).realize()
+
+    for _ in range(5):
+      t = Tensor.randn(N, N, device="CPU").contiguous().realize()
+      Device["CPU"].synchronize()
+      with Timing(f"copy CPU -> {Device.DEFAULT} {t.nbytes()/(1024**2)}M:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        x = _do_copy(t)
+        Device[Device.DEFAULT].synchronize()
+      np.testing.assert_equal(t.numpy(), x.numpy())
+
+  @unittest.skipIf(Device.DEFAULT != "CL" or Device[Device.DEFAULT].count() != 6, "only test this on CL, with 6 gpus")
+  def testCopyCPUto6GPUs(self):
+    t = Tensor.ones(N, N, device="CPU").contiguous().realize()
+    print(f"buffer: {t.nbytes()*1e-9:.2f} GB")
+    for _ in range(3):
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s ({t.nbytes()*6/ns:.2f} GB/s total)"):
+        with Timing("queue: "):
+          for g in range(6):
+            t.to(f"CL:{g}").realize()
+        Device["CL"].synchronize()
+
+if __name__ == '__main__':
+  unittest.main()