ring copy example (#3185)

* ring copy example * use ones for init
2024-01-19 23:34:30 -05:00 · 2024-01-19 23:34:30 -05:00 · 485332935e
parent 254a7372fe
commit 485332935e
1 changed files with 15 additions and 0 deletions
--- a/extra/ring_copy.py
+++ b/extra/ring_copy.py
@ -0,0 +1,15 @@
+from tinygrad import Tensor, Device, GlobalCounters
+from tinygrad.helpers import Timing
+
+N = 512
+GPUS = 5
+ds = tuple([f"{Device.DEFAULT}:{i+1}" for i in range(GPUS)])
+t = [Tensor.ones(N, N, N, device=d).contiguous().realize() for d in ds]
+
+for _ in range(10):
+  GlobalCounters.reset()
+  with Timing():
+    for ti in t:
+      ti.to_(ds[(ds.index(ti.device)+1+len(ds))%len(ds)])
+      # ti.to_(ds[(ds.index(ti.device)-1+len(ds))%len(ds)]) # reversed order
+      ti.realize()