scripts/_bench_common.py

"""Vendored PerfSkills-style harness helpers for GEAK Triton tasks."""
from __future__ import annotations

import argparse
import math
import sys
import time
from typing import Any, Callable, Dict, List

import torch


def make_argparser(title: str) -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=title)
    p.add_argument("--correctness", action="store_true")
    p.add_argument("--benchmark", action="store_true")
    p.add_argument("--full-benchmark", dest="full_benchmark", action="store_true")
    return p


def _bench_once(fn: Callable[[], Any], *, warmup: int, repeat: int) -> float:
    for _ in range(warmup):
        fn()
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(repeat):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / repeat * 1000.0


def _geomean(xs: List[float]) -> float:
    xs = [x for x in xs if x > 0 and not math.isnan(x)]
    if not xs:
        return float("nan")
    return math.exp(sum(math.log(x) for x in xs) / len(xs))


def run_modes(args: argparse.Namespace, cases: List[Dict[str, Any]]) -> None:
    if args.correctness:
        all_ok = True
        for c in cases:
            name = c["name"]
            print(f"CASE={name} checking...")
            try:
                ok = bool(c["check"]())
            except Exception as e:
                print(f"CASE={name} exception: {e}")
                ok = False
            st = "PASS" if ok else "FAIL"
            print(f"CASE={name} {st}")
            all_ok = all_ok and ok
        if all_ok:
            print("CORRECTNESS_OVERALL: PASS")
            sys.exit(0)
        print("CORRECTNESS_OVERALL: FAIL")
        sys.exit(1)

    if args.benchmark or args.full_benchmark:
        warmup, repeat = (5, 10) if args.benchmark else (10, 100)
        lats: List[float] = []
        for c in cases:
            name = c["name"]
            ms = _bench_once(c["run"], warmup=warmup, repeat=repeat)
            lats.append(ms)
            print(f"CASE={name} GEAK_RESULT_LATENCY_MS={ms:.6f}")
        gm = _geomean(lats)
        print(f"GEAK_SHAPES_USED={list(range(len(cases)))}")
        print(f"GEAK_RESULT_LATENCY_MS={gm:.6f}")
        if args.full_benchmark:
            print("GEAK_RESULT_SPEEDUP=1.000000")
        sys.exit(0)

    print("No mode selected (--correctness | --benchmark | --full-benchmark).", file=sys.stderr)
    sys.exit(2)