A script to detect GPU TFLOPS using torch

benchmark.py

import torch
import time
import platform

def get_system_info():
    props = torch.cuda.get_device_properties(0)
    return {
        "OS": f"{platform.system()} {platform.release()}",
        "CPU": platform.processor(),
        "GPU": props.name,
        "VRAM": f"{round(props.total_memory / (1024**3), 2)} GB",
        "Arch": f"Compute Cap {props.major}.{props.minor}",
        "Backend": f"Torch {torch.__version__} / {torch.version.cuda}"
    }

def benchmark_compute_cell(target_vram_gb, dtype):
    device = "cuda"
    matrix_size = 4096 
    
    chunks, A, B = None, None, None
    # 1. Physical VRAM Allocation (Fix: differentiate between float and int)
    elem_size = torch.tensor([], dtype=dtype).element_size()
    chunk_elements = (256 * 1024 * 1024) // elem_size
    num_chunks = int((target_vram_gb * 1024**3) / (256 * 1024**2))
    
    chunks = []
    try:
        for _ in range(num_chunks):
            # FIX: Use randint for integer types, randn for floating points
            if dtype.is_floating_point:
                chunks.append(torch.randn(chunk_elements, device=device, dtype=dtype))
            else:
                # Fill with random integers in a safe range
                chunks.append(torch.randint(-10, 10, (chunk_elements,), device=device, dtype=dtype))
    except Exception:
        pass

    if not chunks: return 0.0

    # 2. Compute Performance Measurement
    try:
        # Prepare operands (Fix: randint for integers)
        if dtype.is_floating_point:
            A = torch.randn((matrix_size, matrix_size), device=device, dtype=dtype)
            B = torch.randn((matrix_size, matrix_size), device=device, dtype=dtype)
        else:
            A = torch.randint(-10, 10, (matrix_size, matrix_size), device=device, dtype=dtype)
            B = torch.randint(-10, 10, (matrix_size, matrix_size), device=device, dtype=dtype)
        
        res = None
        torch.cuda.synchronize()
        iters = 15
        start = time.perf_counter()
        
        for i in range(iters):
            if dtype.is_floating_point:
                if res is not None:
                    B[0, 0] += res.mean() 
                res = torch.matmul(A, B)
                ops_per_iter = 2 * matrix_size**3
            else:
                if res is not None:
                    B[0, 0] += res[0, 0]
                res = (A * B) + A
                ops_per_iter = 2 * matrix_size**3 / 16 # Scaled for normalization

            # VRAM Dependency Injection
            A += chunks[i % len(chunks)][:matrix_size*matrix_size].view(matrix_size, matrix_size)
            A += chunks[i % len(chunks)][:matrix_size*matrix_size].view(matrix_size, matrix_size)
            torch.cuda.synchronize()
            
        end = time.perf_counter()
        duration_sec = (end - start) / iters
        perf_score = (ops_per_iter / duration_sec) / 1e12
        return perf_score

    except Exception:
        return 0.0
    finally:
        if chunks is not None: del chunks
        if A is not None: del A
        if B is not None: del B
        torch.cuda.empty_cache()


# --- Execution ---
info = get_system_info()
print("="*80)
print("SYSTEM DIAGNOSTICS")
for k, v in info.items(): print(f"{k:<12}: {v}")
print("="*80)

gpu_capacity = torch.cuda.get_device_properties(0).total_memory
available_vram_gb = gpu_capacity // (1024 ** 3)
test_steps = list(range(1, available_vram_gb+1))
dtypes = {
    "INT8": torch.int8,
    "INT16": torch.int16,
    "INT32": torch.int32,
    "FP16": torch.float16,
    "BF16": torch.bfloat16,
    "FP32": torch.float32,
    "FP64": torch.float64
}

# Run Benchmarks
results_matrix = {}
for name, dt in dtypes.items():
    print(f"Benchmarking {name}...", end="\r")
    col_results = {}
    for gb in test_steps:
        col_results[gb] = benchmark_compute_cell(gb, dt)
    results_matrix[name] = col_results

print(f"All benchmarking done", end="\r")
# --- Final Table Output ---
print("\n" + "="*80)
print(f"COMPUTE PERFORMANCE REPORT (TFLOPS) - Higher is Better")
print("-" * 80)
header = f"{'VRAM GB':<10}"
for name in dtypes.keys(): header += f" | {name:<12}"
print(header)
print("-" * 80)

for gb in test_steps:
    row = f"{gb:<10.1f}"
    for name in dtypes.keys():
        val = results_matrix[name].get(gb, 0.0)
        row += f" | {val:<12.4f}"
    print(row)
print("="*80)

================================================================================
SYSTEM DIAGNOSTICS
OS          : Linux 7.0.0-15-generic
CPU         : 
GPU         : NVIDIA GeForce RTX 5090 D v2
VRAM        : 23.41 GB
Arch        : Compute Cap 12.0
Backend     : Torch 2.11.0+cu130 / 13.0
================================================================================
All benchmarking done
================================================================================
COMPUTE PERFORMANCE REPORT (TFLOPS) - Higher is Better
--------------------------------------------------------------------------------
VRAM GB    | INT8         | INT16        | INT32        | FP16         | BF16         | FP32         | FP64        
--------------------------------------------------------------------------------
1.0        | 10.2301      | 44.7920      | 15.0907      | 21.4661      | 29.6092      | 27.6435      | 1.5918      
2.0        | 144.8313     | 49.1349      | 15.4851      | 186.2803     | 191.6524     | 58.7142      | 1.6383      
3.0        | 128.3673     | 44.6270      | 15.3233      | 185.9901     | 190.2927     | 59.5430      | 1.6383      
4.0        | 140.2818     | 50.2781      | 15.5195      | 189.5303     | 189.9085     | 59.5667      | 1.6340      
5.0        | 137.4829     | 45.4310      | 15.2576      | 184.2244     | 190.0584     | 59.5024      | 1.6395      
6.0        | 136.4229     | 49.3035      | 14.7178      | 189.7964     | 189.2351     | 59.3024      | 1.6439      
7.0        | 137.5033     | 47.8476      | 15.4067      | 188.7432     | 189.8978     | 59.4527      | 1.6436      
8.0        | 137.7749     | 49.0841      | 15.3194      | 184.3960     | 189.6070     | 59.3933      | 1.6441      
9.0        | 137.6978     | 48.2174      | 15.3767      | 189.4217     | 190.0001     | 59.3818      | 1.6391      
10.0       | 137.7415     | 49.7333      | 15.4279      | 189.6808     | 189.2064     | 59.5064      | 1.6390      
11.0       | 135.1158     | 47.8973      | 14.8607      | 188.8723     | 190.1686     | 59.4420      | 1.6440      
12.0       | 137.0954     | 49.8163      | 15.0624      | 189.8682     | 189.8156     | 59.3243      | 1.6391      
13.0       | 136.2664     | 47.4085      | 14.9469      | 189.2087     | 189.8753     | 59.4051      | 1.5912      
14.0       | 132.3044     | 50.0050      | 15.2279      | 188.6305     | 189.1356     | 59.4203      | 1.6296      
15.0       | 122.6620     | 48.0158      | 15.3992      | 188.0039     | 190.0929     | 59.5259      | 1.6441      
16.0       | 134.4475     | 49.5539      | 15.4654      | 189.2408     | 189.9341     | 59.5567      | 1.6441      
17.0       | 131.7314     | 47.8250      | 14.7391      | 189.4077     | 189.7696     | 59.4809      | 1.6437      
18.0       | 134.9713     | 49.8745      | 15.0967      | 188.5701     | 190.0609     | 59.3144      | 1.6388      
19.0       | 135.3311     | 47.4569      | 15.3348      | 188.6584     | 189.4480     | 59.4533      | 1.6392      
20.0       | 134.9556     | 49.0999      | 15.0326      | 189.2268     | 189.4351     | 58.4601      | 1.6442      
21.0       | 0.0000       | 0.0000       | 0.0000       | 189.0807     | 189.2822     | 0.0000       | 0.0000      
22.0       | 0.0000       | 0.0000       | 0.0000       | 189.0104     | 189.8917     | 0.0000       | 0.0000      
23.0       | 0.0000       | 0.0000       | 0.0000       | 189.3388     | 189.7761     | 0.0000       | 0.0000      
================================================================================