添加 gpu_benchmark.py
This commit is contained in:
@@ -0,0 +1,985 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
NVIDIA GPU Compute Benchmark
|
||||
|
||||
Evaluates matrix multiplication (GEMM) throughput in TFLOPS across all NVIDIA GPUs
|
||||
for float32, float16, bfloat16, and int8 data types using PyTorch.
|
||||
|
||||
Matrix multiplication is the core compute operation in deep learning and the
|
||||
standard metric for measuring GPU compute performance.
|
||||
|
||||
Usage:
|
||||
python gpu_benchmark.py # All GPUs, all dtypes, medium size
|
||||
python gpu_benchmark.py -g 0 # GPU 0 only
|
||||
python gpu_benchmark.py -d float16 int8 # float16 and int8 only
|
||||
python gpu_benchmark.py -s large # 16384^2 large matrices
|
||||
python gpu_benchmark.py --check # Self-check mode
|
||||
python gpu_benchmark.py -o results.json # Export JSON
|
||||
python gpu_benchmark.py -v # Verbose output
|
||||
|
||||
Dependencies: PyTorch (with CUDA), rich
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
import sys
|
||||
import textwrap
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.rule import Rule
|
||||
from rich.table import Table
|
||||
|
||||
console = Console(highlight=False)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Constants
|
||||
# ============================================================================
|
||||
|
||||
VERSION = "0.0.1"
|
||||
|
||||
# Matrix size presets (M, N, K)
|
||||
SIZE_PRESETS: Dict[str, Tuple[int, int, int]] = {
|
||||
"small": (4096, 4096, 4096),
|
||||
"medium": (8192, 8192, 8192),
|
||||
"large": (16384, 16384, 16384),
|
||||
}
|
||||
|
||||
# int8 tensor core alignment requirements: K aligned to 16, M/N aligned to 4
|
||||
INT8_K_ALIGN = 16
|
||||
INT8_MN_ALIGN = 4
|
||||
|
||||
# dtype name -> (torch.dtype, description, bytes per element)
|
||||
DTYPE_MAP: Dict[str, Tuple[torch.dtype, str, int]] = {
|
||||
"float32": (torch.float32, "32-bit float", 4),
|
||||
"float16": (torch.float16, "16-bit float", 2),
|
||||
"bfloat16": (torch.bfloat16, "16-bit brain float", 2),
|
||||
"int8": (torch.int8, "8-bit integer", 1),
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# GPU Theoretical Peak TFLOPS Database
|
||||
# Source: NVIDIA official whitepapers / TechPowerUp GPU Database
|
||||
# Keys: GPU name substring (case-sensitive match); Values: peak TFLOPS per dtype
|
||||
# ============================================================================
|
||||
|
||||
# fmt: off
|
||||
GPU_PEAK_DB: Dict[str, Dict[str, float]] = {
|
||||
# Data Center (Hopper)
|
||||
"H200": {"float32": 494, "float16": 989, "bfloat16": 989, "int8": 1979},
|
||||
"H100": {"float32": 494, "float16": 989, "bfloat16": 989, "int8": 1979},
|
||||
"H800": {"float32": 494, "float16": 989, "bfloat16": 989, "int8": 1979},
|
||||
# Data Center (Ada Lovelace)
|
||||
"L40S": {"float32": 183, "float16": 366, "bfloat16": 366, "int8": 733},
|
||||
"L40": {"float32": 90.5, "float16": 181, "bfloat16": 181, "int8": 362},
|
||||
"L20": {"float32": 59.8, "float16": 119.5, "bfloat16": 119.5, "int8": 239},
|
||||
"L4": {"float32": 30.3, "float16": 60.6, "bfloat16": 60.6, "int8": 121},
|
||||
# Data Center (Ampere)
|
||||
"A100": {"float32": 156, "float16": 312, "bfloat16": 312, "int8": 624},
|
||||
"A800": {"float32": 156, "float16": 312, "bfloat16": 312, "int8": 624},
|
||||
"A40": {"float32": 149.7, "float16": 149.7, "bfloat16": 149.7, "int8": 299},
|
||||
"A30": {"float32": 82.6, "float16": 165.2, "bfloat16": 165.2, "int8": 330},
|
||||
"A10": {"float32": 62.5, "float16": 125, "bfloat16": 125, "int8": 250},
|
||||
"A2": {"float32": 9, "float16": 18, "bfloat16": 18, "int8": 36},
|
||||
# Data Center (Volta)
|
||||
"V100": {"float32": 15.7, "float16": 112, "bfloat16": 0, "int8": 0},
|
||||
# Data Center (Turing)
|
||||
"T4": {"float32": 8.1, "float16": 65, "bfloat16": 0, "int8": 130},
|
||||
# Data Center (Pascal)
|
||||
"P100": {"float32": 10.6, "float16": 21.2, "bfloat16": 0, "int8": 0},
|
||||
"P40": {"float32": 12, "float16": 0.18, "bfloat16": 0, "int8": 0},
|
||||
"P4": {"float32": 5.5, "float16": 0.09, "bfloat16": 0, "int8": 0},
|
||||
# GeForce RTX 50 (Blackwell)
|
||||
"RTX 5090": {"float32": 104.8, "float16": 209.7, "bfloat16": 209.7, "int8": 419.4},
|
||||
"RTX 5080": {"float32": 56.3, "float16": 112.5, "bfloat16": 112.5, "int8": 225},
|
||||
"RTX 5070 Ti": {"float32": 44.5, "float16": 89, "bfloat16": 89, "int8": 178},
|
||||
"RTX 5070": {"float32": 31.2, "float16": 62.4, "bfloat16": 62.4, "int8": 124.8},
|
||||
# GeForce RTX 40 (Ada Lovelace)
|
||||
"RTX 4090": {"float32": 82.6, "float16": 330.3, "bfloat16": 330.3, "int8": 660.6},
|
||||
"RTX 4080 SUPER": {"float32": 52.2, "float16": 208.9, "bfloat16": 208.9, "int8": 417.8},
|
||||
"RTX 4080": {"float32": 48.7, "float16": 194.9, "bfloat16": 194.9, "int8": 389.8},
|
||||
"RTX 4070 Ti SUPER": {"float32": 44.1, "float16": 176.4, "bfloat16": 176.4, "int8": 352.8},
|
||||
"RTX 4070 Ti": {"float32": 40.1, "float16": 160.4, "bfloat16": 160.4, "int8": 320.8},
|
||||
"RTX 4070 SUPER": {"float32": 35.5, "float16": 141.8, "bfloat16": 141.8, "int8": 283.6},
|
||||
"RTX 4070": {"float32": 29.1, "float16": 116.8, "bfloat16": 116.8, "int8": 233.6},
|
||||
"RTX 4060 Ti": {"float32": 22.1, "float16": 88.3, "bfloat16": 88.3, "int8": 176.6},
|
||||
"RTX 4060": {"float32": 15.1, "float16": 60.4, "bfloat16": 60.4, "int8": 120.8},
|
||||
# GeForce RTX 30 (Ampere)
|
||||
"RTX 3090 Ti": {"float32": 40, "float16": 160, "bfloat16": 160, "int8": 320},
|
||||
"RTX 3090": {"float32": 35.6, "float16": 142.3, "bfloat16": 142.3, "int8": 284.6},
|
||||
"RTX 3080 Ti": {"float32": 34.1, "float16": 136.4, "bfloat16": 136.4, "int8": 272.8},
|
||||
"RTX 3080": {"float32": 29.8, "float16": 119, "bfloat16": 119, "int8": 238},
|
||||
"RTX 3070 Ti": {"float32": 21.8, "float16": 87, "bfloat16": 87, "int8": 174},
|
||||
"RTX 3070": {"float32": 20.3, "float16": 81.2, "bfloat16": 81.2, "int8": 162.4},
|
||||
"RTX 3060 Ti": {"float32": 16.2, "float16": 64.8, "bfloat16": 64.8, "int8": 129.6},
|
||||
"RTX 3060": {"float32": 12.7, "float16": 51, "bfloat16": 51, "int8": 102},
|
||||
"RTX 3050": {"float32": 9.1, "float16": 36.4, "bfloat16": 36.4, "int8": 72.8},
|
||||
# GeForce RTX 20 (Turing)
|
||||
"RTX 2080 Ti": {"float32": 13.4, "float16": 26.9, "bfloat16": 0, "int8": 53.8},
|
||||
"RTX 2080 SUPER": {"float32": 11.2, "float16": 22.3, "bfloat16": 0, "int8": 44.6},
|
||||
"RTX 2080": {"float32": 10.1, "float16": 20.1, "bfloat16": 0, "int8": 40.2},
|
||||
"RTX 2070 SUPER": {"float32": 9.1, "float16": 18.1, "bfloat16": 0, "int8": 36.2},
|
||||
"RTX 2070": {"float32": 7.5, "float16": 14.9, "bfloat16": 0, "int8": 29.8},
|
||||
"RTX 2060 SUPER": {"float32": 7.2, "float16": 14.4, "bfloat16": 0, "int8": 28.8},
|
||||
"RTX 2060": {"float32": 6.5, "float16": 12.9, "bfloat16": 0, "int8": 25.8},
|
||||
# GeForce GTX 16 (Turing, no Tensor Core)
|
||||
"GTX 1660 Ti": {"float32": 5.4, "float16": 10.8, "bfloat16": 0, "int8": 0},
|
||||
"GTX 1660": {"float32": 5.0, "float16": 10.0, "bfloat16": 0, "int8": 0},
|
||||
"GTX 1650": {"float32": 3.0, "float16": 6.0, "bfloat16": 0, "int8": 0},
|
||||
# Quadro / RTX Professional (Ada)
|
||||
"RTX 6000 Ada": {"float32": 91.1, "float16": 364.2, "bfloat16": 364.2, "int8": 728.4},
|
||||
"RTX 5000 Ada": {"float32": 52.2, "float16": 208.9, "bfloat16": 208.9, "int8": 417.8},
|
||||
"RTX 4000 Ada": {"float32": 26.7, "float16": 106.8, "bfloat16": 106.8, "int8": 213.6},
|
||||
# Quadro / RTX Professional (Ampere)
|
||||
"RTX A6000": {"float32": 38.7, "float16": 154.8, "bfloat16": 154.8, "int8": 309.6},
|
||||
"RTX A5000": {"float32": 27.8, "float16": 111, "bfloat16": 111, "int8": 222},
|
||||
"RTX A4000": {"float32": 19.2, "float16": 76.8, "bfloat16": 76.8, "int8": 153.6},
|
||||
# Quadro RTX (Turing)
|
||||
"Quadro RTX 8000": {"float32": 16.3, "float16": 32.6, "bfloat16": 0, "int8": 65.2},
|
||||
"Quadro RTX 6000": {"float32": 16.3, "float16": 32.6, "bfloat16": 0, "int8": 65.2},
|
||||
"Quadro RTX 5000": {"float32": 11.2, "float16": 22.3, "bfloat16": 0, "int8": 44.6},
|
||||
"Quadro RTX 4000": {"float32": 7.1, "float16": 14.2, "bfloat16": 0, "int8": 28.4},
|
||||
# GRID / vGPU
|
||||
"GRID A100": {"float32": 156, "float16": 312, "bfloat16": 312, "int8": 624},
|
||||
}
|
||||
# fmt: on
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Utility Functions
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def align_dim(val: int, alignment: int) -> int:
|
||||
"""Round val up to the nearest multiple of alignment."""
|
||||
return ((val + alignment - 1) // alignment) * alignment
|
||||
|
||||
|
||||
def flops_gemm(m: int, n: int, k: int) -> int:
|
||||
"""Floating-point operations for C(MxN) = A(MxK) x B(KxN): 2*M*N*K."""
|
||||
return 2 * m * n * k
|
||||
|
||||
|
||||
def tflops_from_ms(flops: int, time_ms: float) -> float:
|
||||
"""Compute TFLOPS from total FLOPs and elapsed time in milliseconds."""
|
||||
if time_ms <= 0:
|
||||
return 0.0
|
||||
return flops / (time_ms * 1e9)
|
||||
|
||||
|
||||
def calc_compute_intensity(m: int, n: int, k: int, dtype_bytes: int) -> float:
|
||||
"""
|
||||
Compute arithmetic intensity (FLOPs/Byte) for matrix multiplication.
|
||||
|
||||
Data transfer ~ (M*K + K*N + M*N) * bytes_per_element.
|
||||
Arithmetic intensity = total_flops / total_bytes.
|
||||
Returns FLOPs/Byte.
|
||||
"""
|
||||
total_bytes = (m * k + k * n + m * n) * dtype_bytes
|
||||
total_flops = flops_gemm(m, n, k)
|
||||
return total_flops / total_bytes if total_bytes > 0 else 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# GPU Info Detection
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def get_gpu_info(device_idx: int) -> Dict[str, Any]:
|
||||
"""Get detailed information for a specific GPU."""
|
||||
props = torch.cuda.get_device_properties(device_idx)
|
||||
name = props.name
|
||||
mem_gb = props.total_memory / (1024**3)
|
||||
sm_count = props.multi_processor_count
|
||||
cc_major = props.major
|
||||
cc_minor = props.minor
|
||||
|
||||
return {
|
||||
"index": device_idx,
|
||||
"name": name,
|
||||
"memory_gb": mem_gb,
|
||||
"sm_count": sm_count,
|
||||
"cc_major": cc_major,
|
||||
"cc_minor": cc_minor,
|
||||
"cc": f"{cc_major}.{cc_minor}",
|
||||
}
|
||||
|
||||
|
||||
def estimate_peak_tflops(gpu_name: str, sm_count: int, dtype_name: str) -> Optional[float]:
|
||||
"""
|
||||
Estimate theoretical peak TFLOPS for a given dtype on a specific GPU.
|
||||
Looks up from the database first; falls back to architecture-based estimation.
|
||||
Returns None if the dtype is not supported by the hardware.
|
||||
"""
|
||||
# Match by substring in database
|
||||
for key, peaks in GPU_PEAK_DB.items():
|
||||
if key in gpu_name:
|
||||
val = peaks.get(dtype_name)
|
||||
if val is not None:
|
||||
return val if val > 0 else None # 0 means unsupported
|
||||
# Fallback: estimate by architecture and SM count
|
||||
return _estimate_by_arch(sm_count, dtype_name, gpu_name)
|
||||
|
||||
|
||||
def _estimate_by_arch(sm_count: int, dtype_name: str, gpu_name: str) -> Optional[float]:
|
||||
"""
|
||||
Rough peak estimation based on SM count and GPU architecture.
|
||||
Last resort; accuracy is limited.
|
||||
"""
|
||||
name_upper = gpu_name.upper()
|
||||
|
||||
# Infer architecture generation and approximate clock
|
||||
# fmt: off
|
||||
if any(t in name_upper for t in ["H200", "H100", "H800"]):
|
||||
# Hopper, ~1.8 GHz
|
||||
per_sm = {"float32": 4.57, "float16": 9.14, "bfloat16": 9.14, "int8": 18.29}
|
||||
elif any(t in name_upper for t in ["L40", "L20", "L4", "RTX 40", "RTX 6000 ADA", "RTX 5000 ADA", "RTX 4000 ADA"]):
|
||||
# Ada Lovelace, ~2.5 GHz consumer / ~2.0 GHz pro
|
||||
per_sm = {"float32": 0.52, "float16": 2.08, "bfloat16": 2.08, "int8": 4.16}
|
||||
elif any(t in name_upper for t in ["RTX 50"]):
|
||||
# Blackwell consumer, ~2.2 GHz
|
||||
per_sm = {"float32": 0.66, "float16": 2.63, "bfloat16": 2.63, "int8": 5.27}
|
||||
elif any(t in name_upper for t in ["A100", "A800", "A40", "A30", "A10", "A2", "RTX 30", "RTX A"]):
|
||||
# Ampere, ~1.4 GHz DC / ~1.8 GHz consumer
|
||||
per_sm = {"float32": 0.74, "float16": 1.48, "bfloat16": 1.48, "int8": 2.95}
|
||||
elif any(t in name_upper for t in ["V100"]):
|
||||
per_sm = {"float32": 0.20, "float16": 1.40, "bfloat16": 0, "int8": 0}
|
||||
elif any(t in name_upper for t in ["T4", "RTX 20", "QUADRO RTX"]):
|
||||
per_sm = {"float32": 0.18, "float16": 0.35, "bfloat16": 0, "int8": 0.70}
|
||||
elif any(t in name_upper for t in ["GTX 16"]):
|
||||
per_sm = {"float32": 0.35, "float16": 0.70, "bfloat16": 0, "int8": 0}
|
||||
elif any(t in name_upper for t in ["P100", "P40", "P4"]):
|
||||
per_sm = {"float32": 0.19, "float16": 0.38, "bfloat16": 0, "int8": 0}
|
||||
else:
|
||||
# Generic conservative estimate
|
||||
per_sm = {"float32": 0.25, "float16": 0.50, "bfloat16": 0, "int8": 0}
|
||||
# fmt: on
|
||||
|
||||
if dtype_name not in per_sm:
|
||||
return None
|
||||
val = per_sm[dtype_name] * sm_count
|
||||
return val if val > 0 else None
|
||||
|
||||
|
||||
def check_bf16_hw_support(cc_major: int, cc_minor: int) -> bool:
|
||||
"""Check if GPU has native bfloat16 hardware support (Ampere SM 8.0+)."""
|
||||
return cc_major >= 8
|
||||
|
||||
|
||||
def check_int8_tc_support(cc_major: int, cc_minor: int) -> bool:
|
||||
"""Check if GPU has int8 Tensor Core support (Turing SM 7.5+)."""
|
||||
return (cc_major == 7 and cc_minor >= 5) or cc_major >= 8
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# int8 GEMM Support
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _find_int_mm() -> Optional[Callable]:
|
||||
"""Locate torch._int_mm function, compatible across PyTorch versions."""
|
||||
# Prefer torch._int_mm (public but undocumented API)
|
||||
if hasattr(torch, "_int_mm") and callable(torch._int_mm):
|
||||
return torch._int_mm
|
||||
# Fallback to torch.ops.aten._int_mm (torch.ops is a runtime dynamic namespace,
|
||||
# not visible to static analyzers, so use getattr)
|
||||
try:
|
||||
torch_ops = getattr(torch, "ops", None)
|
||||
if torch_ops is not None:
|
||||
aten = getattr(torch_ops, "aten", None)
|
||||
if aten is not None and hasattr(aten, "_int_mm"):
|
||||
return aten._int_mm
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
_INT_MM_FN = _find_int_mm()
|
||||
|
||||
|
||||
def has_int8_support() -> bool:
|
||||
"""Check if current PyTorch supports int8 GEMM."""
|
||||
return _INT_MM_FN is not None
|
||||
|
||||
|
||||
def run_int8_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""Execute int8 matrix multiplication. Input int8, output int32."""
|
||||
if _INT_MM_FN is None:
|
||||
raise RuntimeError("torch._int_mm is not available in this PyTorch version")
|
||||
return _INT_MM_FN(a.contiguous(), b.contiguous())
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Core Benchmarking
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_gemm_kernel(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
gemm_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
|
||||
warmup: int,
|
||||
iters: int,
|
||||
verbose: bool = False,
|
||||
) -> Dict[str, float]:
|
||||
"""
|
||||
Precision-timed GEMM operation using CUDA Events.
|
||||
Returns min/median/mean/std (milliseconds) and p95.
|
||||
|
||||
Design notes:
|
||||
- Uses torch.cuda.Event for GPU-side timing, excluding kernel launch overhead
|
||||
- Warmup eliminates JIT compilation, cuBLAS heuristics, and GPU boost ramp-up
|
||||
- Multiple iterations capture min (interference-free peak) and median (typical)
|
||||
"""
|
||||
stream = torch.cuda.current_stream()
|
||||
|
||||
# Warmup
|
||||
for _ in range(warmup):
|
||||
_ = gemm_fn(a, b)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Timed iterations
|
||||
start_ev = torch.cuda.Event(enable_timing=True)
|
||||
end_ev = torch.cuda.Event(enable_timing=True)
|
||||
timings_ms: List[float] = []
|
||||
|
||||
for _ in range(iters):
|
||||
start_ev.record(stream)
|
||||
_ = gemm_fn(a, b)
|
||||
end_ev.record(stream)
|
||||
torch.cuda.synchronize()
|
||||
timings_ms.append(start_ev.elapsed_time(end_ev))
|
||||
|
||||
sorted_times = sorted(timings_ms)
|
||||
p95_idx = int(len(sorted_times) * 0.95)
|
||||
|
||||
result = {
|
||||
"min": sorted_times[0],
|
||||
"median": statistics.median(sorted_times),
|
||||
"mean": statistics.mean(sorted_times),
|
||||
"std": statistics.stdev(sorted_times) if len(sorted_times) > 1 else 0.0,
|
||||
"p95": sorted_times[p95_idx] if p95_idx < len(sorted_times) else sorted_times[-1],
|
||||
}
|
||||
|
||||
if verbose:
|
||||
console.print(
|
||||
f"[dim]\\[verbose] Timing samples ({iters} iters): "
|
||||
f"min={result['min']:.4f}ms median={result['median']:.4f}ms "
|
||||
f"mean={result['mean']:.4f}ms std={result['std']:.4f}ms[/dim]"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _create_fp_tensors(dtype: torch.dtype, m: int, n: int, k: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Create random float32/float16/bfloat16 matrices."""
|
||||
a = torch.randn(m, k, dtype=dtype, device=device)
|
||||
b = torch.randn(k, n, dtype=dtype, device=device)
|
||||
return a, b
|
||||
|
||||
|
||||
def _create_int8_tensors(m: int, n: int, k: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Create random int8 matrices (aligned dimensions)."""
|
||||
m_aligned = align_dim(m, INT8_MN_ALIGN)
|
||||
n_aligned = align_dim(n, INT8_MN_ALIGN)
|
||||
k_aligned = align_dim(k, INT8_K_ALIGN)
|
||||
a = torch.randint(-128, 127, (m_aligned, k_aligned), dtype=torch.int8, device=device)
|
||||
b = torch.randint(-128, 127, (k_aligned, n_aligned), dtype=torch.int8, device=device)
|
||||
return a, b
|
||||
|
||||
|
||||
def _validate_input_size(m: int, n: int, k: int, element_bytes: int, device: torch.device) -> Tuple[int, int, int]:
|
||||
"""
|
||||
Validate whether matrices fit in GPU memory. Auto-downsize if needed.
|
||||
Returns (m, n, k).
|
||||
"""
|
||||
total_bytes = (m * k + k * n) * element_bytes
|
||||
# Reserve 50% margin for intermediate results
|
||||
needed_bytes = int(total_bytes * 1.5)
|
||||
|
||||
free_bytes = 0
|
||||
try:
|
||||
mem_info = torch.cuda.mem_get_info(device)
|
||||
free_bytes = mem_info[0] if isinstance(mem_info, tuple) else mem_info
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if free_bytes > 0 and needed_bytes > free_bytes * 0.75:
|
||||
# Downsize: scale to 70% of available memory
|
||||
scale = math.sqrt((free_bytes * 0.7) / needed_bytes)
|
||||
new_dim = max(1024, int(m * scale))
|
||||
new_dim = align_dim(new_dim, 16)
|
||||
console.print(f"[yellow]\\[WARN][/] Insufficient memory, downscaling from {m} to {new_dim}")
|
||||
return new_dim, new_dim, new_dim
|
||||
|
||||
return m, n, k
|
||||
|
||||
|
||||
def benchmark_dtype(
|
||||
device: torch.device,
|
||||
dtype_name: str,
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
verbose: bool,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Benchmark GEMM performance for a single dtype on one GPU.
|
||||
Returns None if the dtype is unavailable.
|
||||
"""
|
||||
# Handle int8 separately
|
||||
if dtype_name == "int8":
|
||||
if not has_int8_support():
|
||||
console.print("[yellow]\\[SKIP][/] int8: torch._int_mm unavailable")
|
||||
return None
|
||||
|
||||
m_v, n_v, k_v = _validate_input_size(m, n, k, 1, device)
|
||||
try:
|
||||
a, b = _create_int8_tensors(m_v, n_v, k_v, device)
|
||||
except torch.cuda.OutOfMemoryError:
|
||||
console.print(f"[yellow]\\[SKIP][/] int8: out of memory ({m_v}x{n_v}x{k_v})")
|
||||
return None
|
||||
|
||||
actual_m, actual_k = a.shape
|
||||
actual_n = b.shape[1]
|
||||
|
||||
timing = benchmark_gemm_kernel(a, b, run_int8_gemm, warmup, iters, verbose)
|
||||
|
||||
total_flops = flops_gemm(actual_m, actual_n, actual_k)
|
||||
return {
|
||||
"dtype": dtype_name,
|
||||
"m": actual_m,
|
||||
"n": actual_n,
|
||||
"k": actual_k,
|
||||
"flops": total_flops,
|
||||
"timing_ms": timing,
|
||||
"tflops_min": tflops_from_ms(total_flops, timing["min"]),
|
||||
"tflops_median": tflops_from_ms(total_flops, timing["median"]),
|
||||
"tflops_mean": tflops_from_ms(total_flops, timing["mean"]),
|
||||
"tflops_std": tflops_from_ms(total_flops, timing["std"]),
|
||||
}
|
||||
|
||||
# Floating-point types
|
||||
torch_dtype, _, elem_bytes = DTYPE_MAP[dtype_name]
|
||||
|
||||
m_v, n_v, k_v = _validate_input_size(m, n, k, elem_bytes, device)
|
||||
|
||||
try:
|
||||
a, b = _create_fp_tensors(torch_dtype, m_v, n_v, k_v, device)
|
||||
except torch.cuda.OutOfMemoryError:
|
||||
console.print(f"[yellow]\\[SKIP][/] {dtype_name}: out of memory ({m_v}x{n_v}x{k_v})")
|
||||
return None
|
||||
except RuntimeError as e:
|
||||
if "not supported" in str(e).lower():
|
||||
console.print(f"[yellow]\\[SKIP][/] {dtype_name}: {e}")
|
||||
return None
|
||||
raise
|
||||
|
||||
# GEMM: use @ operator (cuBLAS backend)
|
||||
def _fp_gemm(x, y):
|
||||
return x @ y
|
||||
|
||||
timing = benchmark_gemm_kernel(a, b, _fp_gemm, warmup, iters, verbose)
|
||||
|
||||
total_flops = flops_gemm(m_v, n_v, k_v)
|
||||
ci = calc_compute_intensity(m_v, n_v, k_v, a.element_size())
|
||||
|
||||
result = {
|
||||
"dtype": dtype_name,
|
||||
"m": m_v,
|
||||
"n": n_v,
|
||||
"k": k_v,
|
||||
"flops": total_flops,
|
||||
"timing_ms": timing,
|
||||
"tflops_min": tflops_from_ms(total_flops, timing["min"]),
|
||||
"tflops_median": tflops_from_ms(total_flops, timing["median"]),
|
||||
"tflops_mean": tflops_from_ms(total_flops, timing["mean"]),
|
||||
"tflops_std": tflops_from_ms(total_flops, timing["std"]),
|
||||
}
|
||||
|
||||
if verbose:
|
||||
console.print(
|
||||
f"[dim]\\[verbose] Arithmetic intensity: {ci:.1f} FLOPs/Byte "
|
||||
f"(elem size {a.element_size()} bytes)[/dim]"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Single GPU Full Benchmark
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_single_gpu(
|
||||
device_idx: int,
|
||||
dtypes: List[str],
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
verbose: bool,
|
||||
) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
|
||||
"""Run full benchmark suite on a single GPU. Returns (gpu_info, [result_dict, ...])."""
|
||||
gpu_info = get_gpu_info(device_idx)
|
||||
device = torch.device(f"cuda:{device_idx}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
console.rule(f"GPU {device_idx}: {gpu_info['name']}")
|
||||
console.print(
|
||||
f"Memory: {gpu_info['memory_gb']:.2f} GB | "
|
||||
f"SMs: {gpu_info['sm_count']} | "
|
||||
f"Compute Capability: {gpu_info['cc']}"
|
||||
)
|
||||
console.print(f"Matrix size: {m}x{n}x{k} (MxNxK)")
|
||||
|
||||
results: List[Dict[str, Any]] = []
|
||||
|
||||
for dtype_name in dtypes:
|
||||
if dtype_name == "bfloat16" and not check_bf16_hw_support(
|
||||
gpu_info["cc_major"], gpu_info["cc_minor"]
|
||||
):
|
||||
console.print(
|
||||
f"[yellow]\\[WARN][/] bfloat16: no native hardware support "
|
||||
f"(SM {gpu_info['cc_major']}.{gpu_info['cc_minor']} < 8.0), "
|
||||
f"software emulation will be slow"
|
||||
)
|
||||
|
||||
console.print(f"[bold cyan]\\[{dtype_name}][/] benchmarking...", end="")
|
||||
|
||||
result = benchmark_dtype(device, dtype_name, m, n, k, warmup, iters, verbose)
|
||||
|
||||
if result is None:
|
||||
continue
|
||||
|
||||
# Look up peak
|
||||
peak = estimate_peak_tflops(gpu_info["name"], gpu_info["sm_count"], dtype_name)
|
||||
result["peak_tflops"] = peak
|
||||
if peak and peak > 0:
|
||||
result["efficiency"] = (result["tflops_median"] / peak) * 100
|
||||
else:
|
||||
result["efficiency"] = None
|
||||
|
||||
results.append(result)
|
||||
console.print(f" done: median={result['tflops_median']:.2f} TFLOPS", end="")
|
||||
if result["efficiency"] is not None:
|
||||
console.print(f", efficiency={result['efficiency']:.1f}%")
|
||||
else:
|
||||
console.print()
|
||||
|
||||
return gpu_info, results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Self-Check Mode
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def run_check_mode(gpu_indices: List[int]) -> None:
|
||||
"""Self-check mode: verify dtype availability on each GPU without full benchmarking."""
|
||||
console.rule("GPU Benchmark — Self-Check Mode (--check)")
|
||||
|
||||
console.print(f"PyTorch: {torch.__version__}")
|
||||
console.print(f"CUDA available: {torch.cuda.is_available()}")
|
||||
if torch.cuda.is_available():
|
||||
console.print(f"CUDA version: {torch.version.cuda}")
|
||||
console.print(f"cuDNN version: {torch.backends.cudnn.version()}")
|
||||
console.print(f"GPU count: {torch.cuda.device_count()}")
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
console.print("[red]\\[ERROR][/] CUDA not available, cannot continue.")
|
||||
return
|
||||
|
||||
all_dtypes = ["float32", "float16", "bfloat16", "int8"]
|
||||
|
||||
for idx in gpu_indices:
|
||||
gpu_info = get_gpu_info(idx)
|
||||
device = torch.device(f"cuda:{idx}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
console.rule(f"GPU {idx}: {gpu_info['name']}", style="dim")
|
||||
console.print(
|
||||
f"Memory: {gpu_info['memory_gb']:.2f} GB | "
|
||||
f"SMs: {gpu_info['sm_count']} | "
|
||||
f"CC: {gpu_info['cc']}"
|
||||
)
|
||||
|
||||
for dtype_name in all_dtypes:
|
||||
status = "OK"
|
||||
note = ""
|
||||
|
||||
if dtype_name == "bfloat16":
|
||||
if not check_bf16_hw_support(gpu_info["cc_major"], gpu_info["cc_minor"]):
|
||||
note = "(no native hardware, software emulation)"
|
||||
|
||||
if dtype_name == "int8":
|
||||
if not has_int8_support():
|
||||
status = "N/A"
|
||||
note = "(torch._int_mm unavailable)"
|
||||
elif not check_int8_tc_support(gpu_info["cc_major"], gpu_info["cc_minor"]):
|
||||
note = "(no Tensor Core acceleration)"
|
||||
|
||||
if status == "OK":
|
||||
# Try a small matrix
|
||||
try:
|
||||
_ = benchmark_dtype(device, dtype_name, 256, 256, 256, 2, 3, False)
|
||||
except Exception as e:
|
||||
status = "FAIL"
|
||||
note = str(e)[:60]
|
||||
|
||||
peak = estimate_peak_tflops(gpu_info["name"], gpu_info["sm_count"], dtype_name)
|
||||
peak_str = f"{peak:.1f} TFLOPS" if peak else "--"
|
||||
|
||||
status_style = "green" if status == "OK" else ("red" if status == "FAIL" else "dim")
|
||||
console.print(
|
||||
f" {dtype_name:<10s} [{status_style}]{status:<4s}[/{status_style}] "
|
||||
f"Peak ~ {peak_str} {note}"
|
||||
)
|
||||
|
||||
console.rule("Self-check complete.")
|
||||
console.print("[dim]If no FAIL items, you may proceed with full benchmark.[/dim]\n")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Formatted Output
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def print_single_gpu_table(gpu_info: Dict[str, Any], results: List[Dict[str, Any]]) -> None:
|
||||
"""Print benchmark result table for a single GPU."""
|
||||
if not results:
|
||||
console.print("[dim]No results available.[/dim]")
|
||||
return
|
||||
|
||||
table = Table(
|
||||
title=f"GPU {gpu_info['index']}: {gpu_info['name']} "
|
||||
f"[dim]\\[{gpu_info['memory_gb'] * 1024:.0f}MB] "
|
||||
f"\\[SM {gpu_info['cc']}][/]",
|
||||
show_header=True,
|
||||
header_style="bold",
|
||||
)
|
||||
table.add_column("Data Type", style="cyan", width=12)
|
||||
table.add_column("MxNxK", width=18)
|
||||
table.add_column("Min TFLOPS", justify="right", width=11)
|
||||
table.add_column("Median", justify="right", width=10)
|
||||
table.add_column("Mean", justify="right", width=10)
|
||||
table.add_column("Peak", justify="right", width=10)
|
||||
table.add_column("Eff%", justify="right", width=7)
|
||||
|
||||
for r in results:
|
||||
table.add_row(
|
||||
r["dtype"],
|
||||
f"{r['m']}x{r['n']}x{r['k']}",
|
||||
f"{r['tflops_min']:.2f}" if r["tflops_min"] is not None else "N/A",
|
||||
f"{r['tflops_median']:.2f}" if r["tflops_median"] is not None else "N/A",
|
||||
f"{r['tflops_mean']:.2f}" if r["tflops_mean"] is not None else "N/A",
|
||||
f"{r['peak_tflops']:.2f}" if r.get("peak_tflops") else "N/A",
|
||||
f"{r['efficiency']:.1f}%" if r.get("efficiency") is not None else "N/A",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
console.print("[dim]* Peak = theoretical peak | Eff% = Median / Peak x 100%[/dim]")
|
||||
console.print("[dim]* Min reflects interference-free peak; Median reflects typical performance[/dim]")
|
||||
|
||||
|
||||
def print_summary_table(
|
||||
all_results: List[Tuple[Dict[str, Any], List[Dict[str, Any]]]],
|
||||
) -> None:
|
||||
"""Print multi-GPU summary comparison table (by median TFLOPS)."""
|
||||
if len(all_results) <= 1:
|
||||
return
|
||||
|
||||
# Collect all dtypes seen
|
||||
all_dtypes: List[str] = []
|
||||
for _, res_list in all_results:
|
||||
for r in res_list:
|
||||
if r["dtype"] not in all_dtypes:
|
||||
all_dtypes.append(r["dtype"])
|
||||
|
||||
if not all_dtypes:
|
||||
return
|
||||
|
||||
console.rule("SUMMARY — Multi-GPU Comparison (Median TFLOPS)")
|
||||
|
||||
table = Table(show_header=True, header_style="bold")
|
||||
table.add_column("GPU", style="cyan", no_wrap=True)
|
||||
for dt in all_dtypes:
|
||||
table.add_column(dt, justify="right")
|
||||
|
||||
for gpu_info, res_list in all_results:
|
||||
res_by_dtype = {r["dtype"]: r for r in res_list}
|
||||
row = [gpu_info["name"]]
|
||||
for dt in all_dtypes:
|
||||
r = res_by_dtype.get(dt)
|
||||
if r and r["tflops_median"] is not None:
|
||||
row.append(f"{r['tflops_median']:.1f}")
|
||||
else:
|
||||
row.append("N/A")
|
||||
table.add_row(*row)
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Command-Line Parsing
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="NVIDIA GPU Compute Benchmark -- measure GPU GEMM throughput (TFLOPS) with PyTorch",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=textwrap.dedent("""\
|
||||
Usage examples:
|
||||
%(prog)s # All GPUs, all dtypes, medium size
|
||||
%(prog)s -g 0 # GPU 0 only
|
||||
%(prog)s -g 0 1 -d float16 int8 # GPU 0,1; fp16 and int8 only
|
||||
%(prog)s -s large # 16384^2 large matrices
|
||||
%(prog)s -s 8192 8192 8192 # Custom dimensions
|
||||
%(prog)s -w 20 -i 100 # 20 warmup, 100 timing iterations
|
||||
%(prog)s --check # Self-check mode
|
||||
%(prog)s -v # Verbose output
|
||||
%(prog)s -o results.json # Export JSON
|
||||
"""),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-g",
|
||||
"--gpus",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=None,
|
||||
help="GPU indices to benchmark (default: all), e.g. -g 0 1 2",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-d",
|
||||
"--dtypes",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=["float32", "float16", "bfloat16", "int8"],
|
||||
choices=["float32", "float16", "bfloat16", "int8"],
|
||||
help="Data types to benchmark (default: all four)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-s",
|
||||
"--sizes",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=["medium"],
|
||||
help="Matrix size: small(4096) medium(8192) large(16384) or custom 'M N K' (e.g. -s 4096 4096 4096)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-w",
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Warmup iterations (default: 10)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-i",
|
||||
"--iters",
|
||||
type=int,
|
||||
default=50,
|
||||
help="Timing iterations (default: 50)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Export JSON results to file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
help="Verbose output (show per-iteration timing, arithmetic intensity, etc.)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--check",
|
||||
action="store_true",
|
||||
help="Self-check mode: verify dtype availability without full benchmarking",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tf32",
|
||||
action="store_true",
|
||||
help="Enable TF32 Tensor Core acceleration for float32 (default: disabled, native FP32)",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def parse_sizes(size_args: List[str]) -> Tuple[int, int, int]:
|
||||
"""Parse --sizes argument, returning (M, N, K)."""
|
||||
# Try preset name
|
||||
if len(size_args) == 1 and size_args[0] in SIZE_PRESETS:
|
||||
return SIZE_PRESETS[size_args[0]]
|
||||
# Try three integers
|
||||
if len(size_args) == 3:
|
||||
try:
|
||||
return int(size_args[0]), int(size_args[1]), int(size_args[2])
|
||||
except ValueError:
|
||||
pass
|
||||
# Parse failure
|
||||
console.print(f"[red]\\[ERROR][/] Invalid size argument: {size_args}", stderr=True)
|
||||
console.print(
|
||||
"[dim]Available presets: small medium large or custom: M N K[/dim]",
|
||||
stderr=True,
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Main Entry Point
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
# Check CUDA availability
|
||||
if not torch.cuda.is_available():
|
||||
console.rule("[red]CUDA Error[/]")
|
||||
console.print("[red]\\[ERROR][/] CUDA is not available.")
|
||||
console.print(
|
||||
Panel(
|
||||
"1. NVIDIA GPU drivers are installed\n"
|
||||
"2. PyTorch was installed with CUDA support (not CPU-only)\n"
|
||||
"3. CUDA version is compatible with the driver",
|
||||
title="Please verify",
|
||||
border_style="red",
|
||||
)
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
num_gpus = torch.cuda.device_count()
|
||||
if num_gpus == 0:
|
||||
console.print("[red]\\[ERROR][/] No CUDA GPUs detected.")
|
||||
sys.exit(1)
|
||||
|
||||
# Determine GPU list
|
||||
if args.gpus is None:
|
||||
gpu_indices = list(range(num_gpus))
|
||||
else:
|
||||
gpu_indices = [g for g in args.gpus if 0 <= g < num_gpus]
|
||||
invalid = [g for g in args.gpus if g < 0 or g >= num_gpus]
|
||||
if invalid:
|
||||
console.print(f"[yellow]\\[WARN][/] Invalid GPU indices {invalid} (total {num_gpus} GPUs), ignored")
|
||||
if not gpu_indices:
|
||||
console.print(f"[red]\\[ERROR][/] No valid GPUs. Available: 0..{num_gpus - 1}")
|
||||
sys.exit(1)
|
||||
|
||||
# TF32: disabled by default (native FP32), enable only with --tf32
|
||||
# New API (PyTorch >= 2.9): fp32_precision = 'tf32'/'ieee'
|
||||
# Legacy API (deprecated): allow_tf32 = True/False
|
||||
precision = "tf32" if args.tf32 else "ieee"
|
||||
tf32_enabled = False
|
||||
try:
|
||||
torch.backends.cuda.matmul.fp32_precision = precision
|
||||
torch.backends.cudnn.conv.fp32_precision = precision
|
||||
tf32_enabled = args.tf32
|
||||
except (AttributeError, TypeError):
|
||||
# Fallback to legacy API (PyTorch < 2.9)
|
||||
try:
|
||||
torch.backends.cuda.matmul.allow_tf32 = args.tf32
|
||||
torch.backends.cudnn.allow_tf32 = args.tf32
|
||||
tf32_enabled = args.tf32
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
# Self-check mode
|
||||
if args.check:
|
||||
run_check_mode(gpu_indices)
|
||||
return
|
||||
|
||||
# Parse matrix size
|
||||
m, n, k = parse_sizes(args.sizes)
|
||||
|
||||
# Header
|
||||
console.rule(f"GPU Benchmark v{VERSION} — NVIDIA GPU Compute Benchmark")
|
||||
|
||||
tf32_note = (
|
||||
"enabled (Ampere+ float32 will use TF32 acceleration)"
|
||||
if tf32_enabled
|
||||
else "disabled (native FP32; use --tf32 to enable)"
|
||||
)
|
||||
config_text = (
|
||||
f"Detected {num_gpus} GPU(s) | Benchmarking: {gpu_indices}\n"
|
||||
f"Data types: {args.dtypes}\n"
|
||||
f"Matrix size: {m}x{n}x{k} | Warmup: {args.warmup} iters | Timing: {args.iters} iters\n"
|
||||
f"TF32: {tf32_note}\n"
|
||||
f"Note: int8 GEMM uses torch._int_mm low-level API"
|
||||
)
|
||||
console.print(Panel(config_text, title="Configuration", border_style="blue"))
|
||||
|
||||
# Benchmark each GPU
|
||||
all_results: List[Tuple[Dict[str, Any], List[Dict[str, Any]]]] = []
|
||||
|
||||
for idx in gpu_indices:
|
||||
gpu_info, results = benchmark_single_gpu(idx, args.dtypes, m, n, k, args.warmup, args.iters, args.verbose)
|
||||
print_single_gpu_table(gpu_info, results)
|
||||
all_results.append((gpu_info, results))
|
||||
|
||||
# Summary
|
||||
print_summary_table(all_results)
|
||||
|
||||
# JSON export
|
||||
if args.output:
|
||||
export = {
|
||||
"version": VERSION,
|
||||
"pytorch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"config": {
|
||||
"m": m,
|
||||
"n": n,
|
||||
"k": k,
|
||||
"warmup": args.warmup,
|
||||
"iters": args.iters,
|
||||
"dtypes": args.dtypes,
|
||||
},
|
||||
"gpus": [],
|
||||
}
|
||||
for gpu_info, res_list in all_results:
|
||||
export["gpus"].append(
|
||||
{
|
||||
"index": gpu_info["index"],
|
||||
"name": gpu_info["name"],
|
||||
"memory_gb": round(gpu_info["memory_gb"], 2),
|
||||
"sm_count": gpu_info["sm_count"],
|
||||
"compute_capability": gpu_info["cc"],
|
||||
"results": [
|
||||
{
|
||||
"dtype": r["dtype"],
|
||||
"m": r["m"],
|
||||
"n": r["n"],
|
||||
"k": r["k"],
|
||||
"flops": r["flops"],
|
||||
"tflops_min": round(r["tflops_min"], 2),
|
||||
"tflops_median": round(r["tflops_median"], 2),
|
||||
"tflops_mean": round(r["tflops_mean"], 2),
|
||||
"tflops_std": round(r["tflops_std"], 2),
|
||||
"peak_tflops": round(r["peak_tflops"], 2) if r["peak_tflops"] else None,
|
||||
"efficiency_pct": round(r["efficiency"], 1) if r["efficiency"] else None,
|
||||
"timing_ms": {k: round(v, 4) for k, v in r["timing_ms"].items()},
|
||||
}
|
||||
for r in res_list
|
||||
],
|
||||
}
|
||||
)
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(export, f, indent=2, ensure_ascii=False)
|
||||
console.print(f"\n[green]\\[JSON][/] Results exported to: {args.output}")
|
||||
|
||||
console.rule("Benchmark complete")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user