#!/usr/bin/env python3
import json
import statistics
import torch


def measure(n, dtype, repeats=12):
    torch.manual_seed(1101)
    a = torch.randn((n, n), device="cuda", dtype=dtype)
    b = torch.randn((n, n), device="cuda", dtype=dtype)
    for _ in range(4):
        c = a @ b
    torch.cuda.synchronize()
    samples = []
    for _ in range(repeats):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        c = a @ b
        end.record()
        end.synchronize()
        samples.append(start.elapsed_time(end))
    median_ms = statistics.median(samples)
    return {
        "n": n,
        "dtype": str(dtype),
        "samples_ms": samples,
        "median_ms": median_ms,
        "tflops": (2 * n**3) / (median_ms * 1e-3) / 1e12,
        "checksum": float(c.float().mean()),
    }


print(
    json.dumps(
        {
            "gpu": torch.cuda.get_device_name(),
            "capability": torch.cuda.get_device_capability(),
            "torch": torch.__version__,
            "results": [
                measure(16384, torch.bfloat16),
                measure(16384, torch.float16),
            ],
        },
        indent=2,
    )
)