benchmark.py

"""
GPU Benchmark Suite
Measures memory bandwidth, compute throughput, and kernel characteristics.
"""
import argparse
import json
import time
import sys

try:
    import torch
    HAS_CUDA = torch.cuda.is_available()
except ImportError:
    HAS_CUDA = False


def bench_memory_bandwidth(device, iterations=1000):
    """Measure GPU memory bandwidth (H2D, D2H, D2D)"""
    if not HAS_CUDA:
        return {"error": "CUDA not available"}
    
    sizes_mb = [1, 4, 16, 64, 256, 1024]
    results = []
    
    for size_mb in sizes_mb:
        n = size_mb * 1024 * 1024 // 4  # float32
        host_tensor = torch.randn(n)
        device_tensor = torch.randn(n, device=device)
        
        # Host to Device
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(iterations):
            device_tensor.copy_(host_tensor)
        torch.cuda.synchronize()
        h2d_time = (time.perf_counter() - start) / iterations
        
        # Device to Host
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(iterations):
            host_tensor.copy_(device_tensor)
        torch.cuda.synchronize()
        d2h_time = (time.perf_counter() - start) / iterations
        
        # Device to Device
        dst = torch.empty_like(device_tensor)
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(iterations):
            dst.copy_(device_tensor)
        torch.cuda.synchronize()
        d2d_time = (time.perf_counter() - start) / iterations
        
        results.append({
            "size_mb": size_mb,
            "h2d_gbps": size_mb / (h2d_time * 1000),
            "d2h_gbps": size_mb / (d2h_time * 1000),
            "d2d_gbps": size_mb / (d2d_time * 1000),
        })
    
    return results


def bench_compute_throughput(device, iterations=100):
    """Measure FP32 and FP16 compute throughput"""
    if not HAS_CUDA:
        return {"error": "CUDA not available"}
    
    sizes = [1024, 2048, 4096]
    results = []
    
    for n in sizes:
        # FP32 matmul
        a = torch.randn(n, n, device=device)
        b = torch.randn(n, n, device=device)
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(iterations):
            c = torch.mm(a, b)
        torch.cuda.synchronize()
        fp32_time = (time.perf_counter() - start) / iterations
        fp32_gflops = (2 * n**3) / (fp32_time * 1e9)
        
        # FP16 matmul
        a16 = a.half()
        b16 = b.half()
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(iterations):
            c = torch.mm(a16, b16)
        torch.cuda.synchronize()
        fp16_time = (time.perf_counter() - start) / iterations
        fp16_gflops = (2 * n**3) / (fp16_time * 1e9)
        
        results.append({
            "matrix_size": n,
            "fp32_gflops": fp32_gflops,
            "fp16_gflops": fp16_gflops,
            "fp32_time_ms": fp32_time * 1000,
            "fp16_time_ms": fp16_time * 1000,
        })
    
    return results


def main():
    parser = argparse.ArgumentParser(description="GPU Benchmark Suite")
    parser.add_argument("--suite", default="all", choices=["all", "memory", "compute"])
    parser.add_argument("--device", type=int, default=0)
    parser.add_argument("--iterations", type=int, default=100)
    parser.add_argument("--output", default="results.json")
    args = parser.parse_args()
    
    if not HAS_CUDA:
        print("WARNING: CUDA not available. Using CPU measurements.")
        device = "cpu"
    else:
        device = f"cuda:{args.device}"
        print(f"Device: {torch.cuda.get_device_name(args.device)}")
    
    results = {"device": str(device), "benchmarks": {}}
    
    if args.suite in ["all", "memory"]:
        print("\n--- Memory Bandwidth ---")
        results["benchmarks"]["memory"] = bench_memory_bandwidth(device, args.iterations)
        for r in results["benchmarks"]["memory"]:
            print(f"  {r['size_mb']}MB: H2D={r['h2d_gbps']:.1f} D2H={r['d2h_gbps']:.1f} D2D={r['d2d_gbps']:.1f} GB/s")
    
    if args.suite in ["all", "compute"]:
        print("\n--- Compute Throughput ---")
        results["benchmarks"]["compute"] = bench_compute_throughput(device, min(args.iterations, 50))
        for r in results["benchmarks"]["compute"]:
            print(f"  {r['matrix_size']}x{r['matrix_size']}: FP32={r['fp32_gflops']:.1f} FP16={r['fp16_gflops']:.1f} GFLOPS")
    
    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()