metakernel / dev /report_local.py
phanerozoic's picture
metakernel v1: device dossier, throttle-rejected bench, ULP compare, fuzz, sweep, stamps
0c6aadc verified
Raw
History Blame
937 Bytes
"""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))