| """Small helpers shared by benchmark entry points.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import platform |
| import statistics |
| import time |
| from collections.abc import Callable |
| from typing import Any |
|
|
| import torch |
|
|
|
|
| def common_parser(description: str) -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description=description) |
| parser.add_argument("--iterations", type=int, default=20) |
| parser.add_argument("--warmup", type=int, default=5) |
| parser.add_argument("--device", default="cpu") |
| parser.add_argument("--seed", type=int, default=1234) |
| return parser |
|
|
|
|
| def synchronize(device: torch.device) -> None: |
| if device.type == "cuda": |
| torch.cuda.synchronize(device) |
|
|
|
|
| def measure( |
| operation: Callable[[], Any], *, iterations: int, warmup: int, device: torch.device |
| ) -> dict[str, float]: |
| for _ in range(warmup): |
| operation() |
| synchronize(device) |
| durations = [] |
| for _ in range(iterations): |
| start = time.perf_counter() |
| operation() |
| synchronize(device) |
| durations.append(time.perf_counter() - start) |
| return { |
| "mean_ms": statistics.mean(durations) * 1000, |
| "median_ms": statistics.median(durations) * 1000, |
| "min_ms": min(durations) * 1000, |
| "max_ms": max(durations) * 1000, |
| "stdev_ms": statistics.stdev(durations) * 1000 if len(durations) > 1 else 0.0, |
| } |
|
|
|
|
| def metadata(device: torch.device) -> dict[str, Any]: |
| result: dict[str, Any] = { |
| "python": platform.python_version(), |
| "torch": torch.__version__, |
| "device": str(device), |
| "cuda_available": torch.cuda.is_available(), |
| } |
| if device.type == "cuda" and torch.cuda.is_available(): |
| result["gpu"] = torch.cuda.get_device_name(device) |
| try: |
| import dgl |
|
|
| result["dgl"] = dgl.__version__ |
| except ImportError: |
| result["dgl"] = None |
| return result |
|
|
|
|
| def report(name: str, values: dict[str, Any]) -> None: |
| print(json.dumps({"benchmark": name, **values}, sort_keys=True)) |
|
|
