"""Measured numbers for the card, run on the local GPU from source."""
import json
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import torch

import load_local

mk = load_local.load()

print("GPU:", torch.cuda.get_device_name(0))
d = mk.probe_device()
print(json.dumps(d, indent=1))

n = 8192
x = torch.randn(n, n, device="cuda", dtype=torch.float16)
y = torch.randn(n, n, device="cuda", dtype=torch.float16)
rep = mk.bench(lambda: x @ y, iters=64, warmup=8, bytes=3 * n * n * 2,
               flops=2 * n ** 3, dossier=d, dtype="fp16")
print("\nbench torch.mm fp16 8192^3:", json.dumps(rep, indent=1))

big = torch.randn(128 << 20, device="cuda")
dst = torch.empty_like(big)
rep2 = mk.bench(lambda: dst.copy_(big), iters=48, warmup=6,
                bytes=2 * big.numel() * 4, flops=big.numel(), dossier=d)
print("\nbench 1 GiB copy:", json.dumps(rep2, indent=1))