"""Quick peak-throughput benchmark: big square matmuls per precision + a DRAM copy.

Prints one line per test, then a final JSON line (prefix "RESULT ") for machine parsing.
Each test runs ~2 s after warmup; nvidia-smi is sampled in the background for SM clock and power.
"""
import json, statistics, subprocess, threading, time
import torch

torch.manual_seed(0)
dev = torch.device("cuda")
p = torch.cuda.get_device_properties(0)
SMS = p.multi_processor_count
smi = subprocess.run(["nvidia-smi", "--query-gpu=power.limit,clocks.max.sm,driver_version", "--format=csv,noheader,nounits"],
                     capture_output=True, text=True).stdout.strip().split(", ")
print(f"# {p.name} | SMs={SMS} | power limit {smi[0]} W | max SM clock {smi[1]} MHz | driver {smi[2]} | torch {torch.__version__}", flush=True)

# FLOPs per clock per SM if every op runs at the full (pro/data-centre) rate on Ampere/Ada
FULL = {"fp64": 4, "fp32": 256, "tf32": 512, "fp16": 1024, "bf16": 1024, "fp8": 2048, "int8": 2048}

tag = [None]
samples = {}
results = []

def sampler():
    while True:
        t = tag[0]
        try:
            out = subprocess.run(["nvidia-smi", "--query-gpu=clocks.sm,power.draw", "--format=csv,noheader,nounits"],
                                 capture_output=True, text=True, timeout=5).stdout.splitlines()[0]
            clk, pwr = (float(v) for v in out.split(","))
            if t is not None and tag[0] == t:
                samples.setdefault(t, []).append((clk, pwr))
        except Exception:
            pass
        time.sleep(0.15)

threading.Thread(target=sampler, daemon=True).start()

def set_tf32(on):
    try:
        torch.backends.cuda.matmul.fp32_precision = "tf32" if on else "ieee"
    except Exception:
        torch.backends.cuda.matmul.allow_tf32 = on

def set_fp16_acc(on):
    if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"):
        torch.backends.cuda.matmul.allow_fp16_accumulation = on
        return True
    return False

def run(name, key, fn, work, kind=None, unit="TFLOPS", min_time=2.0, warmup=3):
    try:
        for _ in range(warmup):
            fn()
        torch.cuda.synchronize()
        t0 = time.perf_counter(); fn(); torch.cuda.synchronize()
        iters = max(2, int(min_time / max(time.perf_counter() - t0, 1e-5)))
        tag[0] = name
        s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        s.record()
        for _ in range(iters):
            fn()
        e.record(); torch.cuda.synchronize()
        tag[0] = None
        sec = s.elapsed_time(e) / 1e3 / iters
        rate = work / sec / (1e9 if unit == "GB/s" else 1e12)
        sm = samples.get(name, [])
        clk = statistics.median(x[0] for x in sm) if sm else None
        pwr = statistics.median(x[1] for x in sm) if sm else None
        eff = 100 * rate * 1e12 / (SMS * clk * 1e6 * FULL[kind]) if (kind and clk) else None
        results.append({"test": name, "key": key, "value": round(rate, 2), "unit": unit,
                        "sm_mhz": clk, "power_w": pwr, "pct_of_full_rate_at_clock": None if eff is None else round(eff, 1)})
        effs = f"{eff:5.1f}% of full rate @ that clock" if eff is not None else ""
        print(f"{name:36s} {rate:8.2f} {unit:6s} | {clk or 0:5.0f} MHz | {pwr or 0:5.1f} W | {effs}", flush=True)
    except Exception as ex:
        tag[0] = None
        results.append({"test": name, "key": key, "error": f"{type(ex).__name__}: {str(ex)[:160]}"})
        print(f"{name:36s} FAILED: {type(ex).__name__}: {str(ex)[:160]}", flush=True)

def mats(n, dtype):
    return torch.randn(n, n, device=dev).to(dtype), torch.randn(n, n, device=dev).to(dtype)

n = 8192
a, b = mats(n, torch.float32)
set_tf32(False)
run("FP32 matmul (IEEE, no TF32)", "fp32", lambda: a @ b, 2 * n**3, "fp32")
set_tf32(True)
run("TF32 matmul", "tf32", lambda: a @ b, 2 * n**3, "tf32")
set_tf32(False)
del a, b

for n in (8192, 16384):
    a, b = mats(n, torch.float16)
    set_fp16_acc(False)
    run(f"FP16 matmul, FP32 accum, n={n}", "fp16", lambda: a @ b, 2 * n**3, "fp16")
    if set_fp16_acc(True):
        run(f"FP16 matmul, FP16 accum, n={n}", "fp16_acc16", lambda: a @ b, 2 * n**3, "fp16")
        set_fp16_acc(False)
    del a, b
    a, b = mats(n, torch.bfloat16)
    run(f"BF16 matmul (FP32 accum), n={n}", "bf16", lambda: a @ b, 2 * n**3, "bf16")
    del a, b

if torch.cuda.get_device_capability() >= (8, 9):
    one = torch.tensor(1.0, device=dev)
    for n in (8192, 16384):
        a = torch.randn(n, n, device=dev).to(torch.float8_e4m3fn)
        b = torch.randn(n, n, device=dev).to(torch.float8_e4m3fn).t()
        run(f"FP8 e4m3 scaled_mm, n={n}", "fp8", lambda: torch._scaled_mm(a, b, one, one, out_dtype=torch.bfloat16), 2 * n**3, "fp8")
        run(f"FP8 e4m3 scaled_mm fast-acc, n={n}", "fp8_fastacc",
            lambda: torch._scaled_mm(a, b, one, one, out_dtype=torch.bfloat16, use_fast_accum=True), 2 * n**3, "fp8")
        del a, b
else:
    print("FP8: no hardware support (needs sm_89+), skipped", flush=True)

for n in (8192, 16384):
    a = torch.randint(-128, 128, (n, n), device=dev, dtype=torch.int8)
    bt = torch.randint(-128, 128, (n, n), device=dev, dtype=torch.int8).t()
    run(f"INT8 int_mm (B col-major), n={n}", "int8", lambda: torch._int_mm(a, bt), 2 * n**3, "int8", unit="TOPS")
    del a, bt

n = 4096
a, b = mats(n, torch.float64)
run("FP64 matmul", "fp64", lambda: a @ b, 2 * n**3, "fp64", warmup=1)
del a, b

x = torch.empty(2**28, device=dev)
y = torch.empty_like(x)
run("DRAM copy bandwidth (1 GiB)", "dram", lambda: y.copy_(x), 2 * x.numel() * 4, unit="GB/s")

print("RESULT " + json.dumps({"gpu": p.name, "sms": SMS, "power_limit_w": float(smi[0]), "max_sm_mhz": float(smi[1]),
                              "driver": smi[2], "torch": torch.__version__, "results": results}), flush=True)