"""Proof: the CPU runs ahead of the GPU.

Queues 50 large matmuls on the MPS device and measures two times:
how long Python took to *ask* for the work, and how long the work
actually took. The difference is the gap the chapter draws.
"""
import time
import torch

assert torch.backends.mps.is_available(), "needs an Apple-silicon GPU"
dev = torch.device("mps")

a = torch.randn(2048, 2048, device=dev)
b = torch.randn(2048, 2048, device=dev)
for _ in range(5):          # warmup
    (a @ b)
torch.mps.synchronize()

t0 = time.perf_counter()
for _ in range(50):
    c = a @ b
t_queue = time.perf_counter() - t0
torch.mps.synchronize()
t_done = time.perf_counter() - t0

print(f"torch {torch.__version__}, mps")
print(f"time to queue 50 matmuls : {t_queue*1e3:8.2f} ms")
print(f"time until work finished : {t_done*1e3:8.2f} ms")
print(f"python was free for      : {(t_done-t_queue)*1e3:8.2f} ms ({(t_done-t_queue)/t_done:.0%} of the wall time)")

# The three ways to read the loss, measured, for the two-clocks
# widget: never, once at the end, after every step.
import json
def run_mode(mode, iters=50):
    for _ in range(5):
        (a @ b)
    torch.mps.synchronize()
    t0 = time.perf_counter()
    t_free = 0.0
    for i in range(iters):
        c = a @ b
        if mode == "every":
            c[0, 0].item()
    t_q = time.perf_counter() - t0
    if mode == "once":
        c[0, 0].item()
    torch.mps.synchronize()
    total = time.perf_counter() - t0
    return {"mode": mode, "queue_ms": round(t_q * 1e3, 2),
            "total_ms": round(total * 1e3, 2),
            "free_ms": round((total - t_q) * 1e3, 2)}

modes = [run_mode(m) for m in ("never", "once", "every")]
for m in modes:
    print(f"read {m['mode']:>5}: total {m['total_ms']:8.2f} ms, "
          f"python busy {m['queue_ms']:8.2f} ms")
print("JSON_MODES=" + json.dumps(modes))