IQ.Pilot Prebuilt Release @ ab07000

tinygrad_repo/test/backend/test_zero_copy.py

@@ -0,0 +1,27 @@
+import unittest
+from tinygrad import Tensor, Device
+import time
+
+def time_tensor_numpy(out:Tensor):
+  times = []
+  for _ in range(5):
+    st = time.perf_counter()
+    out.uop.base.realized.as_memoryview(allow_zero_copy=True)
+    et = time.perf_counter() - st
+    times.append(et)
+  return min(times)
+
+N = 4096
+class TestZeroCopy(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT not in {"CPU", "METAL"}, "device isn't zero copy")
+  def test_zero_copy_from_default_to_cpu(self):
+    demo = Tensor.rand(1).realize()
+    t1 = time_tensor_numpy(demo)
+    out = Tensor.rand(N, N).realize()
+    t2 = time_tensor_numpy(out)
+    gbps = out.nbytes()*1e-9/max(t2-t1, 1e-10)
+    print(f"time(base): {t1*1e3:.2f} ms, time(copy): {t2*1e3:.2f} ms :  copy speed {gbps:.2f} GB/s")
+    self.assertGreater(gbps, 600)  # more than 600 GB/s = no copy
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)