48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
#!/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,
|
|
)
|
|
)
|