diff --git a/extra/ring_copy.py b/extra/ring_copy.py new file mode 100644 index 0000000000..1e3863b89a --- /dev/null +++ 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() \ No newline at end of file