from __future__ import annotations
import statistics
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass, fields, is_dataclass
from typing import Any, Literal
import torch
from torch_lattice import SparseTensor
type Mode = Literal['cold_op', 'hot_op', 'backward']
type Params = Mapping[str, Any]
type WorkloadMetrics = dict[str, int | float]
type MetricFactory = Callable[[Params, Any, Any | None, Any | None], WorkloadMetrics]
[docs]
class SkipCase(RuntimeError):
pass
[docs]
@dataclass(frozen=True, slots=True)
class BenchmarkCase:
name: str
group: str
params: tuple[Params, ...]
setup: Callable[[Params], Any]
prepare: Callable[[Any], Any]
run: Callable[[Any], Any]
backward: Callable[[Any], Any] | None = None
metrics: MetricFactory | None = None
units: tuple[str, ...] = ()
modes: tuple[Mode, ...] | None = None
[docs]
def supports(self, mode: Mode) -> bool:
if self.modes is not None and mode not in self.modes:
return False
return mode != 'backward' or self.backward is not None
[docs]
@dataclass(frozen=True, slots=True)
class BenchmarkResult:
case: str
group: str
mode: Mode
device: str
params: dict[str, Any]
warmup: int
repeats: int
median_ms: float | None
min_ms: float | None
p90_ms: float | None
p95_ms: float | None
samples_ms: tuple[float, ...]
workload: WorkloadMetrics
units: dict[str, float]
skipped: bool = False
notes: str = ''
[docs]
def to_json(self) -> dict[str, Any]:
return {
'case': self.case,
'group': self.group,
'mode': self.mode,
'device': self.device,
'params': _jsonable(self.params),
'warmup': self.warmup,
'repeats': self.repeats,
'median_ms': self.median_ms,
'min_ms': self.min_ms,
'p90_ms': self.p90_ms,
'p95_ms': self.p95_ms,
'samples_ms': list(self.samples_ms),
'workload': self.workload,
'units': self.units,
'skipped': self.skipped,
'notes': self.notes,
}
type ProgressStart = Callable[[BenchmarkCase, Params, Mode, str], None]
type ProgressResult = Callable[[BenchmarkResult, BenchmarkCase, Params, Mode, str], None]
type ProgressSkip = Callable[[BenchmarkCase, Params, Mode, str], None]
type ProgressError = Callable[[BenchmarkCase, Params, Mode, str, BaseException], None]
class CudaTimer:
def __init__(self) -> None:
self.start = torch.cuda.Event(enable_timing=True)
self.end = torch.cuda.Event(enable_timing=True)
def __call__(self, fn: Callable[[], Any], warmup: int, repeats: int) -> tuple[tuple[float, ...], Any | None]:
last_output = None
for _ in range(warmup):
last_output = fn()
force(last_output)
torch.cuda.synchronize()
samples: list[float] = []
for _ in range(repeats):
self.start.record()
last_output = fn()
force(last_output)
self.end.record()
torch.cuda.synchronize()
samples.append(float(self.start.elapsed_time(self.end)))
return tuple(samples), last_output
[docs]
def run_case(
case: BenchmarkCase,
params: Params,
*,
mode: Mode,
device: str,
warmup: int,
repeats: int,
) -> BenchmarkResult | None:
if not case.supports(mode):
return None
fixture = case.setup(params)
force(fixture)
prepared = case.prepare(fixture)
force(prepared)
action = _action(case, fixture, prepared, mode)
timer = CudaTimer()
torch.cuda.reset_peak_memory_stats()
samples, output = timer(action, warmup, repeats)
memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)
workload = _derive_workload_metrics(params, fixture=fixture, prepared=prepared, output=output, extra=case.metrics)
workload['memory_mb'] = memory_mb
return BenchmarkResult(
case=case.name,
group=case.group,
mode=mode,
device=device,
params=dict(params),
warmup=warmup,
repeats=repeats,
median_ms=statistics.median(samples),
min_ms=min(samples),
p90_ms=_percentile(samples, 90),
p95_ms=_percentile(samples, 95),
samples_ms=samples,
workload=workload,
units=_derive_units(samples, params, workload, case.units),
)
[docs]
def run_cases(
cases: Iterable[BenchmarkCase],
*,
modes: Sequence[Mode],
device: str,
warmup: int,
repeats: int,
include: str | None = None,
keep_going: bool = True,
on_start: ProgressStart | None = None,
on_result: ProgressResult | None = None,
on_skip: ProgressSkip | None = None,
on_error: ProgressError | None = None,
) -> list[BenchmarkResult]:
results: list[BenchmarkResult] = []
for case in cases:
if include is not None and include not in case.name:
continue
for params in case.params:
for mode in modes:
if not case.supports(mode):
if on_skip is not None:
on_skip(case, params, mode, device)
continue
if on_start is not None:
on_start(case, params, mode, device)
try:
result = run_case(case, params, mode=mode, device=device, warmup=warmup, repeats=repeats)
except Exception as error:
if on_error is not None:
on_error(case, params, mode, device, error)
if not keep_going:
raise
result = _skip_result(case, params, mode, device, warmup, repeats, error)
if result is not None:
if on_result is not None:
on_result(result, case, params, mode, device)
results.append(result)
return results
def force(value: Any) -> None:
tensors = tuple(_collect_tensors(value))
if tensors:
torch.cuda.synchronize()
def _action(case: BenchmarkCase, fixture: Any, prepared: Any, mode: Mode) -> Callable[[], Any]:
if mode == 'cold_op':
return lambda: case.run(case.prepare(fixture))
if mode == 'hot_op':
return lambda: case.run(prepared)
if case.backward is None:
raise ValueError(f'{case.name} does not support backward.')
return lambda: case.backward(case.prepare(fixture))
def _skip_result(
case: BenchmarkCase,
params: Params,
mode: Mode,
device: str,
warmup: int,
repeats: int,
error: BaseException,
) -> BenchmarkResult:
return BenchmarkResult(
case=case.name,
group=case.group,
mode=mode,
device=device,
params=dict(params),
warmup=warmup,
repeats=repeats,
median_ms=None,
min_ms=None,
p90_ms=None,
p95_ms=None,
samples_ms=(),
workload=_derive_workload_metrics(params, fixture=None, prepared=None, output=None),
units={},
skipped=True,
notes=f'{type(error).__name__}: {error}',
)
def _collect_tensors(value: Any) -> Iterable[torch.Tensor]:
if isinstance(value, torch.Tensor):
yield value
return
if isinstance(value, SparseTensor):
yield value.feats
yield value.coords
return
if isinstance(value, Mapping):
for item in value.values():
yield from _collect_tensors(item)
return
if isinstance(value, tuple | list):
for item in value:
yield from _collect_tensors(item)
return
if is_dataclass(value) and not isinstance(value, type):
for field in fields(value):
yield from _collect_tensors(getattr(value, field.name))
def _derive_workload_metrics(
params: Params,
*,
fixture: Any,
prepared: Any | None,
output: Any | None,
extra: MetricFactory | None = None,
) -> WorkloadMetrics:
metrics: WorkloadMetrics = {}
for source in (params,):
for key, target in (('N', 'N'), ('points', 'points'), ('channels', 'channels_in'), ('kernel', 'kernel_volume')):
value = source.get(key)
if isinstance(value, int | float):
metrics.setdefault(target, value)
for value in (prepared, fixture):
if isinstance(value, SparseTensor):
metrics.setdefault('n_in', int(value.feats.shape[0]))
metrics.setdefault('channels_in', int(value.feats.shape[1]))
break
if isinstance(output, SparseTensor):
metrics['n_out'] = int(output.feats.shape[0])
metrics['channels_out'] = int(output.feats.shape[1])
metrics['elements'] = int(output.feats.numel())
elif isinstance(output, torch.Tensor):
metrics.setdefault('n_out', int(output.shape[0]) if output.ndim else 1)
metrics['elements'] = int(output.numel())
elif isinstance(output, dict):
coords = output.get('coords')
if isinstance(coords, torch.Tensor):
metrics['n_out'] = int(coords.shape[0])
if extra is not None:
metrics.update(extra(params, fixture, prepared, output))
edges = metrics.get('edges')
n_out = metrics.get('n_out')
if isinstance(edges, int | float) and isinstance(n_out, int | float) and n_out > 0:
metrics['avg_neighbors'] = edges / n_out
return metrics
def _derive_units(samples: Sequence[float], params: Params, workload: WorkloadMetrics, units: Sequence[str]) -> dict[str, float]:
median_seconds = statistics.median(samples) / 1000.0
if median_seconds <= 0.0:
return {}
out = {}
for unit in units or _default_units(workload):
raw = workload.get(unit, params.get(unit))
if isinstance(raw, int | float):
out[f'{unit}_per_s'] = float(raw) / median_seconds
return out
def _default_units(workload: WorkloadMetrics) -> tuple[str, ...]:
return tuple(unit for unit in ('edges', 'elements', 'points', 'n_out', 'n_in', 'N') if unit in workload)
def _percentile(samples: Sequence[float], pct: int) -> float:
ordered = sorted(samples)
if len(ordered) == 1:
return ordered[0]
rank = (pct / 100) * (len(ordered) - 1)
lower = int(rank)
upper = min(lower + 1, len(ordered) - 1)
weight = rank - lower
return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
def _jsonable(value: Any) -> Any:
if isinstance(value, Mapping):
return {str(key): _jsonable(item) for key, item in value.items()}
if isinstance(value, tuple | list):
return [_jsonable(item) for item in value]
if isinstance(value, str | int | float | bool) or value is None:
return value
return str(value)
__all__ = ['BenchmarkCase', 'BenchmarkResult', 'Mode', 'SkipCase', 'run_case', 'run_cases']