"""
machine_speed.py: how fast can this computer (1) move bytes from memory,
and (2) do math when the numbers are already close by?
"""
import time
import numpy as np

def best_of(fn, reps=5):
    fn()                                            # warm-up
    ts = []
    for _ in range(reps):
        t = time.perf_counter(); fn(); ts.append(time.perf_counter() - t)
    return min(ts)

# 1) Moving bytes: push ONE token through a 400 MB table (10,000 x 10,000).
#    Each number is used once, for 2 FLOPs, so this is almost pure fetching.
big = np.ones((10_000, 10_000), dtype=np.float32)    # 100 M numbers x 4 bytes
one = np.ones(10_000, dtype=np.float32)
sec = best_of(lambda: big @ one)
bandwidth = big.nbytes / sec
print(f"memory speed : {bandwidth/1e9:6.1f} GB/s   (read {big.nbytes/1e6:.0f} MB in {sec*1e3:.1f} ms)")

# 2) Doing math: multiply two 2048 x 2048 tables (they fit in the chip's cache).
a = np.random.rand(2048, 2048).astype(np.float32)
b = np.random.rand(2048, 2048).astype(np.float32)
sec = best_of(lambda: a @ b)
flops = 2 * 2048**3
print(f"math speed   : {flops/sec/1e9:6.1f} GFLOP/s ({flops/1e9:.1f} GFLOPs in {sec*1e3:.1f} ms)")

balance = (flops / sec) / bandwidth
print(f"balance      : {balance:6.1f} FLOPs per byte")
print(f"tokens a trip needs to keep the math busy: about {2*balance:.0f}")