|
| 1 | +import random |
| 2 | +from tinygrad import Device |
| 3 | +from tinygrad.helpers import getenv, DEBUG |
| 4 | + |
| 5 | +def main(): |
| 6 | + seed = getenv("SEED", 1337) |
| 7 | + n_gpus = getenv("GPUS", 3) |
| 8 | + iters = getenv("ITERS", 10000000) |
| 9 | + only_compute = bool(getenv("ONLY_COMPUTE", 0)) |
| 10 | + |
| 11 | + print(f"{n_gpus} GPUs for {iters} iterations, {only_compute=}, seed {seed}") |
| 12 | + devs = tuple([Device[f"{Device.DEFAULT}:{x}"] for x in range(n_gpus)]) |
| 13 | + |
| 14 | + for i in range(iters): |
| 15 | + dev = random.choice(devs) |
| 16 | + q_t = random.choice([dev.hw_copy_queue_t, dev.hw_compute_queue_t] if not only_compute else [dev.hw_compute_queue_t]) |
| 17 | + |
| 18 | + deps_sigs = random.randint(0, len(devs)) |
| 19 | + wait_devs = random.sample(devs, deps_sigs) |
| 20 | + |
| 21 | + q = q_t() |
| 22 | + for d in wait_devs: q.wait(d.timeline_signal, d.timeline_value - 1) |
| 23 | + q.wait(dev.timeline_signal, dev.timeline_value - 1).signal(dev.timeline_signal, dev.timeline_value).submit(dev) |
| 24 | + dev.timeline_value += 1 |
| 25 | + |
| 26 | + if sync:=random.randint(0, 10) < 3: dev.synchronize() |
| 27 | + if DEBUG >= 2: print(f"{i}: {q_t.__name__} {dev.device_id} timeline {dev.timeline_value}, wait for {[d.device_id for d in wait_devs]}, {sync=}") |
| 28 | + elif i % 100 == 0: print(f"\rCompleted {i} iterations", end='') |
| 29 | + |
| 30 | +if __name__ == "__main__": |
| 31 | + main() |
0 commit comments