bench.py

#!/usr/bin/env python3
"""
bench.py -- Single-kernel (flash_attention) benchmark harness.
FIXED. The agent NEVER modifies this file.

Provides:
  1. GPU detection + roofline
  2. 5-stage correctness (smoke, shape sweep, numerical stability, determinism, edge cases)
  3. Performance benchmarking (triton.testing.do_bench)
  4. Greppable output for the agent loop

Usage:
  python bench.py                 # full run
  python bench.py --quick         # skip correctness stages 3-5, bench only large
  python bench.py --sizes large   # bench only one size
  python bench.py --profile       # emit torch profiler trace
"""

from __future__ import annotations

import argparse
import importlib
import os
import signal
import sys
import time
import traceback
from dataclasses import dataclass
from typing import Any, Callable, Dict, Tuple

import torch
import torch.nn.functional as F

# ---------------------------------------------------------------------------
# Timeout helper (Unix SIGALRM)
# ---------------------------------------------------------------------------

class BenchTimeoutError(Exception):
    pass


class _Timeout:
    def __init__(self, seconds: int):
        self.seconds = seconds

    def _handler(self, signum, frame):
        raise BenchTimeoutError(f"Timed out after {self.seconds}s")

    def __enter__(self):
        if hasattr(signal, "SIGALRM"):
            self._old = signal.signal(signal.SIGALRM, self._handler)
            signal.alarm(self.seconds)
        return self

    def __exit__(self, *exc):
        if hasattr(signal, "SIGALRM"):
            signal.alarm(0)
            signal.signal(signal.SIGALRM, self._old)
        return False


# =========================================================================
# 1. GPU DETECTION
# =========================================================================

@dataclass
class GPUSpec:
    name: str = "Unknown"
    sm_count: int = 0
    memory_gb: float = 0.0
    peak_tflops_fp16: float = 0.0
    peak_bandwidth_gb_s: float = 0.0
    l2_cache_mb: float = 0.0
    compute_capability: Tuple[int, int] = (0, 0)


# (peak_fp16_tflops, peak_bandwidth_gb_s, l2_cache_mb)
_KNOWN_GPUS: Dict[str, Tuple[float, float, float]] = {
    "H100 SXM":  (989.5,  3352.0, 50.0),
    "H100 PCIe": (756.0,  2039.0, 50.0),
    "H100":      (756.0,  2039.0, 50.0),
    "A100-SXM":  (312.0,  2039.0, 40.0),
    "A100-PCIE": (312.0,  1935.0, 40.0),
    "A100":      (312.0,  2039.0, 40.0),
    "L40S":      (362.05, 864.0,  48.0),
    "L4":        (121.0,  300.0,  48.0),
    "A10":       (125.0,  600.0,  6.0),
    "4090":      (330.0,  1008.0, 72.0),
    "4080":      (305.0,  716.8,  64.0),
    "3090":      (142.0,  936.2,  6.0),
}


def detect_gpu() -> GPUSpec:
    if not torch.cuda.is_available():
        print("WARNING: No CUDA GPU detected.")
        return GPUSpec()

    props = torch.cuda.get_device_properties(0)
    name = props.name
    matched = next((v for k, v in _KNOWN_GPUS.items() if k in name), None)

    if matched is not None:
        peak_fp16, peak_bw, l2 = matched
    else:
        if hasattr(props, "clock_rate") and props.clock_rate > 0:
            ops_per_clock = 256 if props.major >= 8 else 128
            clock_ghz = props.clock_rate / 1e6
            peak_fp16 = props.multi_processor_count * ops_per_clock * clock_ghz * 2 / 1e3
            peak_bw = max(clock_ghz * 256 / 8 * 2, 500.0)
        else:
            peak_fp16 = 500.0
            peak_bw = 2000.0
        l2 = props.L2_cache_size / (1024 * 1024) if hasattr(props, "L2_cache_size") else 0.0

    return GPUSpec(
        name=name,
        sm_count=props.multi_processor_count,
        memory_gb=round(props.total_memory / (1024 ** 3), 1),
        peak_tflops_fp16=peak_fp16,
        peak_bandwidth_gb_s=peak_bw,
        l2_cache_mb=l2,
        compute_capability=(props.major, props.minor),
    )


# =========================================================================
# 2. KERNEL CONFIG (flash_attention only)
# =========================================================================

def _gen_inputs(size: dict, dtype: torch.dtype, device: str, seed: int = 42) -> dict:
    torch.manual_seed(seed)
    B, H, S, D = size["batch"], size["heads"], size["seq_len"], size["head_dim"]
    return {
        "Q": torch.randn(B, H, S, D, device=device, dtype=dtype),
        "K": torch.randn(B, H, S, D, device=device, dtype=dtype),
        "V": torch.randn(B, H, S, D, device=device, dtype=dtype),
    }


def _ref_fn(inputs: dict) -> torch.Tensor:
    import reference
    return reference.flash_attention_ref(inputs["Q"], inputs["K"], inputs["V"])


def _dtype_bytes(dtype: torch.dtype) -> int:
    return torch.tensor([], dtype=dtype).element_size()


CONFIG: Dict[str, Any] = {
    "kernel_type": "flash_attention",
    "test_sizes": [
        ("tiny",   {"batch": 1, "heads": 4,  "seq_len": 64,   "head_dim": 64}),
        ("small",  {"batch": 2, "heads": 8,  "seq_len": 256,  "head_dim": 64}),
        ("medium", {"batch": 2, "heads": 16, "seq_len": 512,  "head_dim": 64}),
        ("large",  {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}),
        ("xlarge", {"batch": 2, "heads": 32, "seq_len": 2048, "head_dim": 64}),
        ("long",   {"batch": 1, "heads": 32, "seq_len": 4096, "head_dim": 64}),
        ("gqa",    {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}),
        ("llm_7b", {"batch": 1, "heads": 32, "seq_len": 2048, "head_dim": 128}),
    ],
    "test_dtypes": [torch.float16, torch.bfloat16],
    "tolerances": {
        torch.float16:  {"atol": 1e-2, "rtol": 1e-2},
        torch.bfloat16: {"atol": 2e-2, "rtol": 2e-2},
    },
    # 4 * B * H * S^2 * D
    "flops_fn":  lambda s: 4 * s["batch"] * s["heads"] * (s["seq_len"] ** 2) * s["head_dim"],
    # 3*Q + O bytes (rough estimate)
    "bytes_fn":  lambda s, dt: 4 * s["batch"] * s["heads"] * s["seq_len"] * s["head_dim"] * _dtype_bytes(dt),
    "input_generator": _gen_inputs,
    "reference_fn":    _ref_fn,
    "edge_sizes": [
        ("edge_127",  {"batch": 1, "heads": 8, "seq_len": 127,  "head_dim": 64}),
        ("edge_1023", {"batch": 1, "heads": 8, "seq_len": 1023, "head_dim": 64}),
    ],
}


# =========================================================================
# 3. CORRECTNESS (5 stages)
# =========================================================================

def _compare(output: torch.Tensor, expected: torch.Tensor, atol: float, rtol: float) -> dict:
    if output.shape != expected.shape:
        return {"match": False, "reason": f"shape mismatch: {output.shape} vs {expected.shape}",
                "max_abs_error": float("inf"), "pct_within_tol": 0.0}

    out_f = output.float()
    exp_f = expected.float()
    abs_diff = (out_f - exp_f).abs()
    max_abs = abs_diff.max().item()
    within = (abs_diff <= atol + rtol * exp_f.abs()).float().mean().item() * 100.0
    match = torch.allclose(out_f, exp_f, atol=atol, rtol=rtol)
    return {
        "match": match,
        "reason": "" if match else f"max_abs_error={max_abs:.6e} exceeds tol(atol={atol}, rtol={rtol})",
        "max_abs_error": max_abs,
        "pct_within_tol": within,
    }


def _has_nan_inf(t: torch.Tensor) -> bool:
    return bool(torch.isnan(t).any().item() or torch.isinf(t).any().item())


def run_correctness(kernel_fn: Callable, config: dict, quick: bool = False) -> dict:
    device = "cuda"
    results = {k: "SKIP" for k in
               ("smoke_test", "shape_sweep", "numerical_stability", "determinism", "edge_cases")}
    results["correctness"] = "FAIL"
    details = []
    all_pass = True

    gen_fn = config["input_generator"]
    ref_fn = config["reference_fn"]
    sizes = config["test_sizes"]
    dtypes = config["test_dtypes"]
    tols = config["tolerances"]

    # ----- Stage 1: smoke -----
    print("\n--- Stage 1: Smoke Test ---")
    try:
        _, tiny_sz = sizes[0]
        dt0 = dtypes[0]
        inputs = gen_fn(tiny_sz, dt0, device, seed=42)
        expected = ref_fn(inputs)
        with _Timeout(30):
            output = kernel_fn(**inputs)

        if _has_nan_inf(output):
            results["smoke_test"] = "FAIL"
            details.append("  smoke: NaN/Inf in output")
            all_pass = False
            print("  FAIL: NaN/Inf in output")
        else:
            cmp = _compare(output, expected, **tols.get(dt0, {"atol": 1e-2, "rtol": 1e-2}))
            if cmp["match"]:
                results["smoke_test"] = "PASS"
                print(f"  PASS (max_abs_error={cmp['max_abs_error']:.6e})")
            else:
                results["smoke_test"] = "FAIL"
                details.append(f"  smoke: {cmp['reason']}")
                all_pass = False
                print(f"  FAIL: {cmp['reason']}")
    except BenchTimeoutError:
        results["smoke_test"] = "FAIL"
        details.append("  smoke: TIMEOUT")
        all_pass = False
        print("  FAIL: TIMEOUT")
    except Exception as e:
        results["smoke_test"] = "FAIL"
        details.append(f"  smoke: CRASH ({type(e).__name__}: {e})")
        all_pass = False
        print(f"  FAIL: CRASH ({type(e).__name__}: {e})")

    if results["smoke_test"] == "FAIL":
        results["correctness"] = "FAIL"
        results["details"] = details
        print("\ncorrectness: FAIL (smoke failed; aborting)")
        return results

    # ----- Stage 2: shape sweep -----
    print("\n--- Stage 2: Shape Sweep ---")
    sweep_pass = True
    sweep_count = 0
    sweep_fail = 0
    worst_err = 0.0
    worst_case = ""

    for label, sz in sizes:
        for dtype in dtypes:
            sweep_count += 1
            try:
                inputs = gen_fn(sz, dtype, device, seed=42)
                expected = ref_fn(inputs)
                with _Timeout(30):
                    output = kernel_fn(**inputs)

                if _has_nan_inf(output):
                    sweep_pass = False
                    sweep_fail += 1
                    details.append(f"  sweep {label}/{dtype}: NaN/Inf")
                    print(f"  FAIL: {label} {dtype} -> NaN/Inf")
                    continue

                cmp = _compare(output, expected, **tols.get(dtype, {"atol": 1e-2, "rtol": 1e-2}))
                if cmp["max_abs_error"] > worst_err:
                    worst_err = cmp["max_abs_error"]
                    worst_case = f"{label}/{dtype}"

                if cmp["match"]:
                    print(f"  PASS: {label} {dtype} (max_err={cmp['max_abs_error']:.2e})")
                else:
                    sweep_pass = False
                    sweep_fail += 1
                    details.append(f"  sweep {label}/{dtype}: {cmp['reason']}")
                    print(f"  FAIL: {label} {dtype} -> {cmp['reason']}")
            except torch.cuda.OutOfMemoryError:
                print(f"  SKIP: {label} {dtype} -> OOM")
                torch.cuda.empty_cache()
            except BenchTimeoutError:
                sweep_pass = False
                sweep_fail += 1
                details.append(f"  sweep {label}/{dtype}: TIMEOUT")
                print(f"  FAIL: {label} {dtype} -> TIMEOUT")
            except Exception as e:
                sweep_pass = False
                sweep_fail += 1
                details.append(f"  sweep {label}/{dtype}: {type(e).__name__}: {e}")
                print(f"  FAIL: {label} {dtype} -> {type(e).__name__}: {e}")
            finally:
                torch.cuda.empty_cache()

    if sweep_pass:
        results["shape_sweep"] = f"PASS ({sweep_count} configs, worst_err={worst_err:.2e} at {worst_case})"
    else:
        results["shape_sweep"] = f"FAIL ({sweep_fail}/{sweep_count} failed)"
        all_pass = False

    if quick:
        results["numerical_stability"] = "SKIP (quick mode)"
        results["determinism"] = "SKIP (quick mode)"
        results["edge_cases"] = "SKIP (quick mode)"
        results["correctness"] = "PASS" if all_pass else "FAIL"
        results["details"] = details
        print(f"\ncorrectness: {results['correctness']} (quick mode: stages 3-5 skipped)")
        return results

    # ----- Stage 3: numerical stability -----
    print("\n--- Stage 3: Numerical Stability ---")
    stab_pass = True
    stab_size = next((sz for label, sz in sizes if label == "small"), sizes[min(1, len(sizes) - 1)][1])
    stab_dtype = dtypes[0]

    cases = [
        ("near_max", lambda t: t * 60000.0 if t.dtype == torch.float16 else t * 1e30),
        ("near_zero", lambda t: t * 1e-6),
        ("all_zeros", lambda t: torch.zeros_like(t)),
        ("all_same", lambda t: torch.ones_like(t) * 0.5),
    ]
    for case_name, transform in cases:
        try:
            inputs = gen_fn(stab_size, stab_dtype, device, seed=42)
            transformed = {k: (transform(v) if isinstance(v, torch.Tensor) and v.is_floating_point() else v)
                           for k, v in inputs.items()}
            expected = ref_fn(transformed)
            with _Timeout(30):
                output = kernel_fn(**transformed)

            if _has_nan_inf(output) and not _has_nan_inf(expected):
                stab_pass = False
                details.append(f"  stability {case_name}: NaN/Inf (reference clean)")
                print(f"  FAIL: {case_name} -> NaN/Inf (reference clean)")
            elif _has_nan_inf(output) and _has_nan_inf(expected):
                print(f"  PASS: {case_name} -> both NaN/Inf (expected overflow)")
            else:
                tol = tols.get(stab_dtype, {"atol": 1e-2, "rtol": 1e-2})
                cmp = _compare(output, expected, atol=tol["atol"] * 10, rtol=tol["rtol"] * 10)
                if cmp["match"]:
                    print(f"  PASS: {case_name} (max_err={cmp['max_abs_error']:.2e})")
                else:
                    stab_pass = False
                    details.append(f"  stability {case_name}: {cmp['reason']}")
                    print(f"  FAIL: {case_name} -> {cmp['reason']}")
        except torch.cuda.OutOfMemoryError:
            print(f"  SKIP: {case_name} -> OOM")
            torch.cuda.empty_cache()
        except BenchTimeoutError:
            stab_pass = False
            details.append(f"  stability {case_name}: TIMEOUT")
            print(f"  FAIL: {case_name} -> TIMEOUT")
        except Exception as e:
            stab_pass = False
            details.append(f"  stability {case_name}: {type(e).__name__}: {e}")
            print(f"  FAIL: {case_name} -> {type(e).__name__}: {e}")
        finally:
            torch.cuda.empty_cache()

    results["numerical_stability"] = "PASS" if stab_pass else "FAIL"
    if not stab_pass:
        all_pass = False

    # ----- Stage 4: determinism -----
    print("\n--- Stage 4: Determinism ---")
    det_pass = True
    try:
        outputs = []
        for _ in range(3):
            inputs_i = gen_fn(stab_size, dtypes[0], device, seed=42)
            with _Timeout(30):
                outputs.append(kernel_fn(**inputs_i))
        for i in range(1, 3):
            if not torch.equal(outputs[0], outputs[i]):
                det_pass = False
                diff = (outputs[0].float() - outputs[i].float()).abs()
                details.append(f"  determinism: run 0 vs {i} differ (max_diff={diff.max().item():.6e})")
                print(f"  FAIL: run 0 vs {i} differ (max_diff={diff.max().item():.6e})")
        if det_pass:
            print("  PASS: 3 runs are bitwise identical")
        results["determinism"] = "PASS" if det_pass else "FAIL"
    except Exception as e:
        results["determinism"] = f"FAIL ({type(e).__name__})"
        details.append(f"  determinism: {type(e).__name__}: {e}")
        det_pass = False
        print(f"  FAIL: {type(e).__name__}: {e}")
    finally:
        torch.cuda.empty_cache()

    if not det_pass:
        all_pass = False

    # ----- Stage 5: edge cases -----
    print("\n--- Stage 5: Edge Cases ---")
    edge_pass = True
    for label, sz in config.get("edge_sizes", []):
        try:
            inputs = gen_fn(sz, dtypes[0], device, seed=42)
            expected = ref_fn(inputs)
            with _Timeout(30):
                output = kernel_fn(**inputs)

            if _has_nan_inf(output) and not _has_nan_inf(expected):
                edge_pass = False
                details.append(f"  edge {label}: NaN/Inf")
                print(f"  FAIL: {label} -> NaN/Inf")
            else:
                cmp = _compare(output, expected, **tols.get(dtypes[0], {"atol": 1e-2, "rtol": 1e-2}))
                if cmp["match"]:
                    print(f"  PASS: {label} (max_err={cmp['max_abs_error']:.2e})")
                else:
                    edge_pass = False
                    details.append(f"  edge {label}: {cmp['reason']}")
                    print(f"  FAIL: {label} -> {cmp['reason']}")
        except torch.cuda.OutOfMemoryError:
            print(f"  SKIP: {label} -> OOM")
            torch.cuda.empty_cache()
        except BenchTimeoutError:
            edge_pass = False
            details.append(f"  edge {label}: TIMEOUT")
            print(f"  FAIL: {label} -> TIMEOUT")
        except Exception as e:
            edge_pass = False
            details.append(f"  edge {label}: {type(e).__name__}: {e}")
            print(f"  FAIL: {label} -> {type(e).__name__}: {e}")
        finally:
            torch.cuda.empty_cache()
    results["edge_cases"] = "PASS" if edge_pass else "FAIL"
    if not edge_pass:
        all_pass = False

    results["correctness"] = "PASS" if all_pass else "FAIL"
    results["details"] = details
    print(f"\ncorrectness: {results['correctness']}")
    return results


# =========================================================================
# 4. PERFORMANCE
# =========================================================================

def _do_bench(fn: Callable, warmup: int = 25, rep: int = 100) -> float:
    """Median time in ms."""
    try:
        from triton.testing import do_bench
        return do_bench(fn, warmup=warmup, rep=rep)
    except ImportError:
        pass
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    times = []
    for _ in range(rep):
        s = torch.cuda.Event(enable_timing=True)
        e = torch.cuda.Event(enable_timing=True)
        s.record()
        fn()
        e.record()
        torch.cuda.synchronize()
        times.append(s.elapsed_time(e))
    times.sort()
    return times[len(times) // 2]


def run_performance(kernel_fn: Callable, config: dict, gpu: GPUSpec, sizes_filter: str = "all") -> dict:
    device = "cuda"
    gen_fn = config["input_generator"]
    ref_fn = config["reference_fn"]
    flops_fn = config["flops_fn"]
    bytes_fn = config["bytes_fn"]
    dtypes = config["test_dtypes"]
    sizes = config["test_sizes"]

    if sizes_filter == "all":
        bench_sizes = sizes
    else:
        bench_sizes = [(l, s) for l, s in sizes if l == sizes_filter]
        if not bench_sizes:
            bench_sizes = [next(((l, s) for l, s in sizes if l == "large"), sizes[-1])]

    primary_label, primary_size = next(((l, s) for l, s in sizes if l == "large"), sizes[-1])
    dtype = dtypes[0]

    all_results = []
    primary_result = None

    for label, sz in bench_sizes:
        print(f"\n  Benchmarking: {label} ...")
        try:
            flops = flops_fn(sz)
            nbytes = bytes_fn(sz, dtype)
            inputs = gen_fn(sz, dtype, device, seed=42)

            with _Timeout(30):
                kernel_ms = _do_bench(lambda: kernel_fn(**inputs), warmup=25, rep=100)
            with _Timeout(30):
                ref_ms = _do_bench(lambda: ref_fn(inputs), warmup=25, rep=100)

            kernel_us = kernel_ms * 1000.0
            ref_us = ref_ms * 1000.0
            tflops = flops / (kernel_ms / 1000.0) / 1e12 if kernel_ms > 0 else 0.0
            bw = nbytes / (kernel_ms / 1000.0) / 1e9 if kernel_ms > 0 else 0.0
            ref_tflops = flops / (ref_ms / 1000.0) / 1e12 if ref_ms > 0 else 0.0

            ai = flops / nbytes if nbytes > 0 else 0.0
            ridge = (gpu.peak_tflops_fp16 * 1e12) / (gpu.peak_bandwidth_gb_s * 1e9) if gpu.peak_bandwidth_gb_s > 0 else 0.0
            bottleneck = "memory_bound" if ai < ridge else "compute_bound"
            pct_compute = (tflops / gpu.peak_tflops_fp16 * 100.0) if gpu.peak_tflops_fp16 > 0 else 0.0
            pct_bw = (bw / gpu.peak_bandwidth_gb_s * 100.0) if gpu.peak_bandwidth_gb_s > 0 else 0.0
            speedup = ref_ms / kernel_ms if kernel_ms > 0 else 0.0

            entry = {
                "label": label, "size": sz, "dtype": str(dtype),
                "flops": flops, "bytes": nbytes,
                "kernel_latency_us": kernel_us, "pytorch_latency_us": ref_us,
                "throughput_tflops": tflops, "bandwidth_gb_s": bw,
                "ref_throughput_tflops": ref_tflops,
                "pct_peak_compute": pct_compute, "pct_peak_bandwidth": pct_bw,
                "arithmetic_intensity": ai, "ridge_point": ridge,
                "bottleneck": bottleneck, "speedup_vs_pytorch": speedup,
            }
            all_results.append(entry)
            if label == primary_label:
                primary_result = entry

            print(f"    kernel: {kernel_us:.2f} us | pytorch: {ref_us:.2f} us | "
                  f"speedup: {speedup:.3f}x | {tflops:.3f} TFLOPS | {pct_compute:.1f}% peak")

        except torch.cuda.OutOfMemoryError:
            print(f"    SKIP: {label} -> OOM")
            torch.cuda.empty_cache()
        except BenchTimeoutError:
            print(f"    SKIP: {label} -> TIMEOUT")
        except Exception as e:
            print(f"    ERROR: {label} -> {type(e).__name__}: {e}")
            traceback.print_exc()
        finally:
            torch.cuda.empty_cache()

    if primary_result is None and all_results:
        primary_result = all_results[-1]
    return {"primary": primary_result, "all": all_results}


# =========================================================================
# 5. PROFILER (optional)
# =========================================================================

def run_profile(kernel_fn: Callable, config: dict):
    device = "cuda"
    sizes = config["test_sizes"]
    prof_size = next((sz for label, sz in sizes if label == "medium"), sizes[0][1])
    dtype = config["test_dtypes"][0]
    inputs = config["input_generator"](prof_size, dtype, device, seed=42)

    os.makedirs("./traces", exist_ok=True)
    print("\n=== PROFILING ===")
    print(f"Profiling size: {prof_size}, dtype: {dtype}")

    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
        record_shapes=True, with_stack=True,
    ) as prof:
        for _ in range(5):
            kernel_fn(**inputs)
        torch.cuda.synchronize()
        for _ in range(10):
            kernel_fn(**inputs)
        torch.cuda.synchronize()

    trace_path = "./traces/kernel_trace.json"
    prof.export_chrome_trace(trace_path)
    print(f"profile_trace: {trace_path}")
    try:
        print(prof.key_averages().table(sort_by="self_device_time_total", row_limit=20))
    except Exception:
        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))


# =========================================================================
# 6. MAIN
# =========================================================================

def main():
    t_start = time.time()

    parser = argparse.ArgumentParser(description="autokernel-mini benchmark harness")
    parser.add_argument("--sizes", type=str, default="all",
                        help="Which sizes to benchmark (default: all)")
    parser.add_argument("--quick", action="store_true",
                        help="Skip correctness stages 3-5; bench only large")
    parser.add_argument("--profile", action="store_true",
                        help="Emit torch profiler trace")
    args = parser.parse_args()

    print("=" * 60)
    print("autokernel-mini Benchmark Harness")
    print("=" * 60)

    # ---- Import kernel ----
    try:
        if os.getcwd() not in sys.path:
            sys.path.insert(0, os.getcwd())
        script_dir = os.path.dirname(os.path.abspath(__file__))
        if script_dir not in sys.path:
            sys.path.insert(0, script_dir)

        kernel_module = importlib.import_module("kernel")
        kernel_fn = kernel_module.kernel_fn
        kernel_type = getattr(kernel_module, "KERNEL_TYPE", None)

        if kernel_type != "flash_attention":
            print(f"ERROR: kernel.py KERNEL_TYPE is {kernel_type!r}, expected 'flash_attention'")
            sys.exit(1)

        print(f"kernel_type: {kernel_type}")
        print("kernel_module: kernel.py loaded successfully")
    except SyntaxError as e:
        print(f"\nERROR: kernel.py has a syntax error: {e}")
        traceback.print_exc()
        print("\ncorrectness: FAIL")
        print("throughput_tflops: 0.000")
        sys.exit(1)
    except Exception as e:
        print(f"\nERROR: failed to import kernel.py: {type(e).__name__}: {e}")
        traceback.print_exc()
        print("\ncorrectness: FAIL")
        print("throughput_tflops: 0.000")
        sys.exit(1)

    # ---- GPU info ----
    gpu = detect_gpu()
    print("\n=== GPU INFO ===")
    print(f"gpu_name: {gpu.name}")
    print(f"gpu_sm_count: {gpu.sm_count}")
    print(f"gpu_memory_gb: {gpu.memory_gb}")
    print(f"gpu_peak_tflops_fp16: {gpu.peak_tflops_fp16}")
    print(f"gpu_peak_bandwidth_gb_s: {gpu.peak_bandwidth_gb_s}")
    print(f"gpu_l2_cache_mb: {gpu.l2_cache_mb}")
    print(f"gpu_compute_capability: {gpu.compute_capability[0]}.{gpu.compute_capability[1]}")

    # ---- Correctness ----
    print("\n=== CORRECTNESS ===")
    try:
        correctness = run_correctness(kernel_fn, CONFIG, quick=args.quick)
    except Exception as e:
        print(f"\nFATAL: correctness crashed: {type(e).__name__}: {e}")
        traceback.print_exc()
        correctness = {"correctness": "FAIL", "smoke_test": "CRASH", "shape_sweep": "CRASH",
                       "numerical_stability": "CRASH", "determinism": "CRASH", "edge_cases": "CRASH"}

    print("\n--- Correctness Summary ---")
    for k in ("smoke_test", "shape_sweep", "numerical_stability", "determinism", "edge_cases"):
        print(f"{k}: {correctness.get(k, 'N/A')}")
    print(f"correctness: {correctness['correctness']}")

    # ---- Performance ----
    sizes = CONFIG["test_sizes"]
    primary_label, primary_size = next(((l, s) for l, s in sizes if l == "large"), sizes[-1])
    perf_dtype = CONFIG["test_dtypes"][0]
    size_params = ", ".join(f"{k}={v}" for k, v in primary_size.items())
    print(f"\n=== PERFORMANCE ({primary_label}: {size_params}, dtype={perf_dtype}) ===")

    perf = {"primary": None, "all": []}
    peak_vram_mb = 0.0
    try:
        sizes_filter = "large" if args.quick else args.sizes
        torch.cuda.reset_peak_memory_stats()
        perf = run_performance(kernel_fn, CONFIG, gpu, sizes_filter=sizes_filter)
        peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
    except Exception as e:
        print(f"\nFATAL: performance crashed: {type(e).__name__}: {e}")
        traceback.print_exc()

    primary = perf.get("primary")
    if primary is not None:
        print(f"\n--- Performance Summary (primary: {primary['label']}) ---")
        print(f"latency_us: {primary['kernel_latency_us']:.2f}")
        print(f"latency_ms: {primary['kernel_latency_us'] / 1000.0:.4f}")
        print(f"throughput_tflops: {primary['throughput_tflops']:.3f}")
        print(f"bandwidth_gb_s: {primary['bandwidth_gb_s']:.1f}")
        print(f"pct_peak_compute: {primary['pct_peak_compute']:.1f}%")
        print(f"pct_peak_bandwidth: {primary['pct_peak_bandwidth']:.1f}%")
        print(f"arithmetic_intensity: {primary['arithmetic_intensity']:.2f}")
        print(f"ridge_point: {primary['ridge_point']:.2f}")
        print(f"bottleneck: {primary['bottleneck']}")
        print(f"flops: {primary['flops']}")
        print(f"bytes: {primary['bytes']}")
        print(f"peak_vram_mb: {peak_vram_mb:.1f}")

        print("\n=== COMPARISON VS PYTORCH ===")
        print(f"pytorch_latency_us: {primary['pytorch_latency_us']:.2f}")
        print(f"kernel_latency_us: {primary['kernel_latency_us']:.2f}")
        print(f"speedup_vs_pytorch: {primary['speedup_vs_pytorch']:.3f}x")
        print(f"pytorch_tflops: {primary['ref_throughput_tflops']:.3f}")
        print(f"kernel_tflops: {primary['throughput_tflops']:.3f}")
    else:
        print("\nlatency_us: 0.00")
        print("throughput_tflops: 0.000")
        print(f"peak_vram_mb: {peak_vram_mb:.1f}")
        print("\n=== COMPARISON VS PYTORCH ===")
        print("speedup_vs_pytorch: 0.000x")

    all_perf = perf.get("all", [])
    if len(all_perf) > 1:
        print("\n=== SIZE SWEEP ===")
        print(f"{'size':<12} {'kernel_us':>12} {'pytorch_us':>12} {'speedup':>10} {'tflops':>10} {'%peak':>8}")
        print("-" * 66)
        for entry in all_perf:
            print(f"{entry['label']:<12} {entry['kernel_latency_us']:>12.2f} "
                  f"{entry['pytorch_latency_us']:>12.2f} {entry['speedup_vs_pytorch']:>9.3f}x "
                  f"{entry['throughput_tflops']:>10.3f} {entry['pct_peak_compute']:>7.1f}%")

    if args.profile:
        try:
            run_profile(kernel_fn, CONFIG)
        except Exception as e:
            print(f"\nWARNING: Profiling failed: {type(e).__name__}: {e}")

    elapsed = time.time() - t_start
    throughput = primary["throughput_tflops"] if primary else 0.0
    print("\n=== FINAL ===")
    print(f"kernel_type: {CONFIG['kernel_type']}")
    print(f"correctness: {correctness['correctness']}")
    print(f"throughput_tflops: {throughput:.3f}")
    if primary:
        print(f"speedup_vs_pytorch: {primary['speedup_vs_pytorch']:.3f}x")
        print(f"pct_peak_compute: {primary['pct_peak_compute']:.1f}%")
    else:
        print("speedup_vs_pytorch: 0.000x")
        print("pct_peak_compute: 0.0%")
    print(f"bench_time_seconds: {elapsed:.1f}")
    if elapsed > 90:
        print(f"WARNING: bench.py took {elapsed:.1f}s (budget: 90s)")


if __name__ == "__main__":
    main()