from __future__ import annotations
import argparse
import random
from collections.abc import Sequence
from datetime import datetime
from pathlib import Path
import torch
import torch_lattice
from torch_lattice_bench.catalog import GROUPS, MODES, PRESETS
from torch_lattice_bench.cases import all_cases
from torch_lattice_bench.console import make_console
from torch_lattice_bench.report import write_json, write_summary
_RESULTS_DIR = Path('benchmarks/results')
[docs]
def main() -> None:
args = _parser().parse_args()
if args.smoke:
args.preset = 'smoke'
if not torch.cuda.is_available() or not str(args.device).startswith('cuda'):
raise RuntimeError('torch-lattice benchmark requires a CUDA device.')
_configure_backend(args)
device = torch.device(args.device)
groups = tuple(args.group) if args.group else GROUPS
modes = tuple(args.mode) if args.mode else _default_modes(groups)
n_values = tuple(args.n_values) if args.n_values else None
channels = tuple(args.channels) if args.channels else None
layouts = tuple(args.layout) if args.layout else None
console = make_console(args.color, quiet=args.quiet)
cases = all_cases(
args.preset,
groups=groups,
n_values=n_values,
channels=channels,
layouts=layouts,
dtype=args.dtype,
device=device,
)
if args.list:
for case in cases:
print(f'{case.group}/{case.name}')
return
total = _count_runs(cases, modes, args.case_filter)
console.set_total(total)
json_path, summary_path = _report_paths(args)
from torch_lattice_bench.harness import run_cases
console.heading(f'device {args.device}')
results = run_cases(
cases,
modes=modes,
device=args.device,
warmup=args.warmup,
repeats=args.repeats,
include=args.case_filter,
keep_going=not args.fail_fast,
on_start=console.start,
on_result=console.done,
on_skip=console.skipped,
on_error=console.failed,
)
write_json(json_path, results=results)
write_summary(summary_path, results=results)
console.report(json_path, summary_path)
def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description='Benchmark torch-lattice CUDA sparse operator and module surfaces.',
)
parser.add_argument('--preset', choices=PRESETS, default='smoke')
parser.add_argument('--device', default='cuda')
parser.add_argument('--mode', action='append', choices=MODES)
parser.add_argument('--group', '--groups', dest='group', action='append', choices=GROUPS)
parser.add_argument('--case-filter')
parser.add_argument('--warmup', type=int, default=5)
parser.add_argument('--repeats', '--iters', dest='repeats', type=int, default=20)
parser.add_argument('--size', '--points', dest='n_values', action='append', type=_positive_int)
parser.add_argument('--channels', action='append', type=_positive_int)
parser.add_argument('--layout', '--pattern', '--patterns', dest='layout', action='append', choices=('isolated', 'line', 'plane', 'grid', 'block2', 'block3', 'block4', 'block8'))
parser.add_argument('--dtype', choices=('fp16', 'fp32'), default='fp16')
parser.add_argument('--output')
parser.add_argument('--color', choices=('auto', 'always', 'never'), default='auto')
parser.add_argument('--quiet', action='store_true')
parser.add_argument('--list', action='store_true')
parser.add_argument('--smoke', action='store_true', help='Alias for --preset smoke.')
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--allow-tf32', dest='allow_tf32', action='store_true', default=True)
parser.add_argument('--no-allow-tf32', dest='allow_tf32', action='store_false')
parser.add_argument('--allow-fp16', dest='allow_fp16', action='store_true', default=True)
parser.add_argument('--no-allow-fp16', dest='allow_fp16', action='store_false')
parser.add_argument('--hash-rsv-ratio', type=int, default=64)
parser.add_argument('--fail-fast', action='store_true')
return parser
def _configure_backend(args: argparse.Namespace) -> None:
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
torch.backends.cuda.matmul.allow_tf32 = args.allow_tf32
torch_lattice.backends.allow_tf32 = args.allow_tf32
torch_lattice.backends.allow_fp16 = args.allow_fp16
torch_lattice.backends.hash_rsv_ratio = max(args.hash_rsv_ratio, torch_lattice.backends.hash_rsv_ratio)
torch_lattice.backends.benchmark = True
def _default_modes(groups: Sequence[str]) -> tuple[str, ...]:
if tuple(groups) == ('train',):
return ('backward',)
return ('cold_op', 'hot_op')
def _count_runs(cases, modes: Sequence[str], include: str | None) -> int:
total = 0
for case in cases:
if include is not None and include not in case.name:
continue
for _params in case.params:
total += sum(1 for mode in modes if case.supports(mode))
return total
def _report_paths(args: argparse.Namespace) -> tuple[Path, Path]:
if args.output:
json_path = Path(args.output)
if not json_path.is_absolute() and json_path.parent == Path('.'):
json_path = _RESULTS_DIR / json_path
if json_path.suffix != '.json':
json_path = json_path.with_suffix('.json')
else:
stamp = datetime.now().strftime('%Y%m%d-%H%M%S')
json_path = _RESULTS_DIR / f'torch-lattice-bench-{args.preset}-{stamp}.json'
return json_path, json_path.with_suffix('.summary.txt')
def _positive_int(value: str) -> int:
parsed = int(value)
if parsed <= 0:
raise argparse.ArgumentTypeError('value must be positive')
return parsed
if __name__ == '__main__':
main()