referee/harness.py

"""OUROBOROS HARNESS — the immutable verify-compile-bench referee. THIS IS THE PRODUCT.

The model is small and replaceable; this harness is the moat. It takes a candidate Triton
kernel (a source string), and returns a grounded verdict that NEITHER the proposer nor any
trainer can fake:

  correctness is a BOOLEAN  (allclose vs a PyTorch reference, across adversarial inputs)
  speed is a NUMBER         (wall-clock on the 4090, CUDA events, vs eager AND torch.compile)

This mirrors `sec_sqli/discovery_specialist/dvwa_oracle.py`: success is a real observed
effect, never a pattern match. There it was a seeded canary reflected in a live HTTP
response; here it is `allclose(out, ref) == True` AND a measured `t_baseline / t_kernel`.
A kernel that merely *looks* fast or *looks* correct gets nothing.

This codebase has been burned before (memory: the matrix-rewrite "win" was a verifier
certifying an ablation artifact). The GPU analog is a benchmark that times launch-async
noise, compilation, or elided work. Every guard below exists to prevent that:

  NON-NEGOTIABLES (the line between this and a toy):
  1. SUBPROCESS ISOLATION + HARD TIMEOUT. Triton kernels segfault and hang. A crash must
     not take down the orchestrator. compile+run+bench all happen in a child process the
     parent kills on timeout.
  2. ADVERSARIAL MULTI-SHAPE CORRECTNESS. Shapes, dtypes AND magnitudes are swept (see
     specs._mk_*). A kernel correct only on benign N(0,1) (e.g. softmax with no
     max-subtraction) FAILS — the negative-control analog.
  3. CUDA EVENTS + WARMUP (both paths) + MEDIAN-of-N. time.time() around a launch is
     meaningless (async). Without warmup you time JIT/inductor compilation. Median because
     the 4090 boost-clocks drift.
  4. HONEST BASELINE = torch.compile, reported even when we LOSE. Beating eager is the
     floor; beating compile is the flex. Losses are printed plainly, never hidden.

Parent API:   evaluate(kernel_src, spec_name, ...) -> Result
Worker mode:  python harness.py --worker      (reads one JSON request on stdin)
Self-test:    python harness.py                (gold kernels pass; wrong kernels REJECTED)
"""
from __future__ import annotations

import argparse
import json
import statistics
import subprocess
import sys
from dataclasses import asdict, dataclass, field
from pathlib import Path

HERE = Path(__file__).resolve().parent


class _BenchVerifyError(Exception):
    """The timed output failed the post-bench allclose — a caching/memoizing kernel."""


# ----------------------------------------------------------------------------------------
@dataclass
class Result:
    status: str                       # ok | compile_fail | runtime_fail | incorrect | timeout | crash
    feedback: str = ""                # structured, teachable signal for the proposer prompt
    correct: bool = False
    n_shapes_passed: int = 0
    max_abs_err: float = 0.0
    latency_ms: float = 0.0           # candidate kernel, median
    eager_ms: float = 0.0
    compile_ms: float = 0.0
    maxauto_ms: float = 0.0           # torch.compile(mode="max-autotune") — the STRONG baseline
    speedup_eager: float = 0.0        # eager_ms / latency_ms   (>1 = faster than eager)
    speedup_compile: float = 0.0      # compile_ms / latency_ms (>1 = faster than torch.compile default)
    speedup_maxauto: float = 0.0      # maxauto_ms / latency_ms (>1 = faster than max-autotune; 0 = not measured)

    def to_dict(self):
        return asdict(self)


# ============================ PARENT: subprocess driver =================================
def evaluate(kernel_src: str, spec_name: str, n_shapes: int = 8, n_iters: int = 100,
             seed: int = 0, strong: bool = False, correctness_only: bool = False,
             rotate: bool = False, bench_override: tuple | None = None,
             timeout: float | None = None) -> Result:
    """Run a candidate kernel through the full referee in an ISOLATED child process.

    The child can segfault or hang freely; we reap it. Only a clean JSON verdict on stdout
    counts as a result — anything else is a crash, reported as such (never silently 'ok').

    strong=True also benchmarks torch.compile(mode="max-autotune") — the strongest honest
    baseline — at the cost of a slow one-time autotune compile (hence the longer timeout).
    correctness_only=True skips ALL benchmarking (returns ok/incorrect from allclose alone) —
    much cheaper, for building/filtering the SFT corpus where the boolean is all that matters.
    rotate=True cycles among 4 distinct input clones inside the timed loop (kernel AND
    baselines alike) so nothing stays L2-resident across iterations — the cache-cold
    cross-check for headline numbers.
    bench_override=(M, N, "fp16"|"bf16"|"fp32") re-benches at an arbitrary shape via
    specs.grid_inputs (the SHAPE-GRID rebench); the override shape also joins the
    correctness sweep, so the anti-special-casing guarantee holds at every grid cell."""
    if timeout is None:
        timeout = 300.0 if strong else (20.0 if correctness_only else 40.0)
    req = json.dumps({"kernel_src": kernel_src, "spec_name": spec_name,
                      "n_shapes": n_shapes, "n_iters": n_iters, "seed": seed, "strong": strong,
                      "correctness_only": correctness_only, "rotate": rotate,
                      "bench_override": list(bench_override) if bench_override else None})
    try:
        proc = subprocess.run([sys.executable, str(HERE / "harness.py"), "--worker"],
                              input=req, capture_output=True, text=True, timeout=timeout)
    except subprocess.TimeoutExpired:
        return Result(status="timeout", feedback=f"killed after {timeout:.0f}s (hang/deadlock in compile or launch)")
    if proc.returncode != 0:
        # segfault / illegal memory access / OOM that took the interpreter down
        tail = (proc.stderr or "").strip().splitlines()
        return Result(status="crash", feedback="child crashed (rc=%d): %s" % (
            proc.returncode, tail[-1] if tail else "no stderr"))
    line = next((l for l in reversed(proc.stdout.splitlines()) if l.startswith("RESULT:")), None)
    if not line:
        return Result(status="crash", feedback="no verdict on stdout (worker produced no RESULT line)")
    return Result(**json.loads(line[len("RESULT:"):]))


def evaluate_inprocess_eager(kernel_src: str, spec_name: str, n_shapes: int = 2,
                             n_iters: int = 30, seed: int = 0,
                             rotate: bool = False,
                             bench_override: tuple | None = None) -> Result:
    """Run a candidate in the current Python process and time it only against eager PyTorch.

    ZeroGPU grants CUDA to the decorated process, not to arbitrary child processes. This keeps
    Triton compilation, correctness, anti-memoization checks, and CUDA-event timing inside the
    `@spaces.GPU` call. It intentionally skips torch.compile/max-autotune so the free-tier
    120s ZeroGPU budget is spent on the candidate and the eager baseline only.
    """
    req = {"kernel_src": kernel_src, "spec_name": spec_name, "n_shapes": n_shapes,
           "n_iters": n_iters, "seed": seed, "strong": False,
           "correctness_only": False, "rotate": rotate, "eager_only": True,
           "bench_override": list(bench_override) if bench_override else None}
    return _worker(req)


def evaluate_inprocess_full(kernel_src: str, spec_name: str, n_shapes: int = 2,
                            n_iters: int = 30, seed: int = 0,
                            rotate: bool = False,
                            bench_override: tuple | None = None) -> Result:
    """Same in-process path as evaluate_inprocess_eager, but ALSO times torch.compile (default)
    and torch.compile max-autotune-no-cudagraphs, so the result carries speedup_compile and
    speedup_maxauto. For these memory-bound norm/activation ops the max-autotune compile is only
    a few seconds (measured: ~1.5-4.5s, inductor-cached after the first op), so it fits the ZeroGPU
    120s budget. This is what lets local mode show the HONEST baseline (vs the compiler) next to the
    eager number, instead of only the inflated vs-eager fusion win.
    """
    req = {"kernel_src": kernel_src, "spec_name": spec_name, "n_shapes": n_shapes,
           "n_iters": n_iters, "seed": seed, "strong": True,
           "correctness_only": False, "rotate": rotate, "eager_only": False,
           "bench_override": list(bench_override) if bench_override else None}
    return _worker(req)


# ============================ CHILD: the actual referee ================================
# Triton's @jit inspects the DEFINING SOURCE FILE, so a kernel must live in a real .py on
# disk — exec'ing a string fails ("@jit functions should be defined in a Python file").
# We write each candidate to a temp module (with a standard import preamble) and import it.
_PREAMBLE = "import torch\nimport triton\nimport triton.language as tl\n\n"


def _worker(req: dict) -> Result:
    import importlib.util
    import os
    import tempfile
    import torch
    sys.path.insert(0, str(HERE))
    from specs import get_spec
    import random

    spec = get_spec(req["spec_name"])

    # ---- 1. COMPILE: materialize the kernel to a temp .py and import it ------------------
    import atexit
    tmp = tempfile.NamedTemporaryFile("w", suffix=".py", dir=str(HERE), delete=False)
    tmp.write(_PREAMBLE + req["kernel_src"])
    tmp.close()
    # keep the file alive: @jit reads it at LAUNCH (during correctness), not at import.
    atexit.register(lambda p=tmp.name: os.path.exists(p) and os.unlink(p))
    try:
        spec_mod = importlib.util.spec_from_file_location("_kern_" + str(os.getpid()), tmp.name)
        mod = importlib.util.module_from_spec(spec_mod)
        spec_mod.loader.exec_module(mod)
    except Exception as e:
        return Result(status="compile_fail", feedback=f"{type(e).__name__}: {str(e)[:200]}")
    run = getattr(mod, "run", None)
    if not callable(run):
        return Result(status="compile_fail", feedback="kernel defines no callable `run(*inputs)`")

    # ---- 2. CORRECTNESS: GUARANTEED stress cases (high-scale fp16/bf16, odd N) FIRST, then
    #         THE BENCH SHAPE ITSELF, then the adversarial random sweep. The stress cases make
    #         the negative controls fail DETERMINISTICALLY — never by seed luck. The bench
    #         inputs are included so a kernel cannot special-case the (public, fixed) timing
    #         shape: anything it returns at the bench shape is allclose-checked here first. --
    # the BENCH inputs (fixed spec shape, or the shape-grid override) — built up front so the
    # correctness sweep ALWAYS covers exactly what gets timed.
    if req.get("bench_override"):
        from specs import grid_inputs
        _M, _N, _dt = req["bench_override"]
        bench = grid_inputs(req["spec_name"], int(_M), int(_N),
                            {"fp16": torch.float16, "bf16": torch.bfloat16,
                             "fp32": torch.float32}[_dt])
    else:
        bench = spec.bench_inputs()

    rng = random.Random(req["seed"])
    cases = (list(spec.stress_inputs()) + [bench]
             + [spec.make_inputs(rng) for _ in range(req["n_shapes"])])
    passed = 0
    worst = 0.0
    for inputs in cases:
        ref = spec.reference(*inputs)
        rtol, atol = spec.tol(inputs[0].dtype)
        clones = [t.clone() for t in inputs]           # kernel gets clones; originals stay pristine
        try:
            out = run(*clones)
        except Exception as e:
            torch.cuda.synchronize()
            return Result(status="runtime_fail", n_shapes_passed=passed,
                          feedback=f"{type(e).__name__} on shape {tuple(inputs[0].shape)}/{inputs[0].dtype}: {str(e)[:160]}")
        # CONTRACT: run() must write a fresh output, never mutate its inputs. Enforced, not
        # assumed — the bench loop reuses identical args across iters, which is only sound
        # because this check rejects in-place kernels.
        for orig, cl in zip(inputs, clones):
            if not torch.equal(orig, cl):
                return Result(status="incorrect", n_shapes_passed=passed,
                              feedback=(f"kernel MUTATES its input ({tuple(orig.shape)}/{orig.dtype}) — "
                                        "run() must write a fresh output tensor"))
        if out is None or out.shape != ref.shape:
            return Result(status="incorrect", n_shapes_passed=passed,
                          feedback=f"wrong shape: got {None if out is None else tuple(out.shape)} want {tuple(ref.shape)}")
        d = (out.float() - ref.float()).abs()
        err = float(d.max())
        worst = max(worst, err)
        if not torch.allclose(out.float(), ref.float(), rtol=rtol, atol=atol):
            bad = int((d > atol + rtol * ref.float().abs()).sum())
            return Result(status="incorrect", n_shapes_passed=passed, max_abs_err=err,
                          feedback=(f"max abs err {err:.3e} on {tuple(inputs[0].shape)}/{inputs[0].dtype} "
                                    f"scale~{inputs[0].float().abs().max():.0f} ({bad} elems over tol {atol:.1e}/{rtol:.1e}) "
                                    f"— likely a reduction/stability bug, not a shape bug"))
        passed += 1
        del clones, out, ref

    # correctness-only mode: the kernel is correct across all adversarial cases; skip the
    # (expensive) benchmark entirely. Used to build/filter the SFT corpus.
    if req.get("correctness_only"):
        return Result(status="ok", correct=True, n_shapes_passed=passed, max_abs_err=worst,
                      feedback=f"PASS {passed} cases (correctness-only)")

    # ---- 3. BENCHMARK (correct kernels only) — CUDA events, warmup both, median ----------
    eager = spec.reference
    eager_only = bool(req.get("eager_only"))
    compiled = None
    if not eager_only:
        try:
            compiled = torch.compile(spec.reference)
        except Exception:
            compiled = spec.reference                   # fall back; still honest vs eager

    rotate = bool(req.get("rotate"))

    def _bench(fn, n_iters, verify_tol=None):
        # Clone the inputs ONCE, OUTSIDE the timed window. Cloning 8192x4096 fp16 tensors is
        # itself a ~134MB GPU memcpy; doing it inside the CUDA-event window would time the
        # copy too and pull every ratio toward 1.0 (the exact contamination this harness
        # exists to prevent). Kernels write a fresh output and never mutate inputs (the
        # correctness phase ENFORCES that now), so identical args across iters is correct.
        #
        # ANTI-MEMOIZATION: before every timed iteration (outside the event window) one
        # element of the first input is overwritten with a LARGE in-distribution value (the
        # stress sweep already uses x64 magnitudes), so a run() that caches its output by
        # input pointer returns a STALE result whose error is far above the tolerance
        # envelope — caught by the verify-after-bench allclose below. (A subtle poke is not
        # enough: perturbing one softmax logit moves outputs ~1e-4, inside fp16 tol.) The
        # poke is identical for the candidate and both baselines (fair), and is a
        # single-element write enqueued before the start event (never inside the window).
        #
        # rotate=True: cycle among 4 distinct clone-sets so inputs cannot stay L2-resident
        # across iterations (the cache-cold cross-check). Applied to candidate AND baselines.
        n_sets = 4 if rotate else 1
        arg_sets = [[t.clone() for t in bench] for _ in range(n_sets)]
        for _ in range(25):
            o = fn(*arg_sets[0])                        # warmup: JIT/inductor compile + clock settle
        torch.cuda.synchronize()
        times = []
        o = None
        flat0 = [a[0].view(-1) for a in arg_sets]
        poke_val = 60.0 if flat0[0].is_floating_point() else 100
        for i in range(n_iters):
            args = arg_sets[i % n_sets]
            f = flat0[i % n_sets]
            f[i % f.numel()] = poke_val                 # poke: stale caches now err >> tol
            a = torch.cuda.Event(enable_timing=True)
            b = torch.cuda.Event(enable_timing=True)
            a.record()
            o = fn(*args)                               # ONLY the kernel is in the timed window
            b.record()
            torch.cuda.synchronize()
            times.append(a.elapsed_time(b))
        if verify_tol is not None:
            # VERIFY-AFTER-BENCH: the final timed iteration's output must match the reference
            # on the final (poked) input state. A memoizing/caching run() fails here.
            rtol, atol = verify_tol
            final_args = arg_sets[(n_iters - 1) % n_sets]
            ref_final = spec.reference(*final_args)
            if o is None or not torch.allclose(o.float(), ref_final.float(), rtol=rtol, atol=atol):
                raise _BenchVerifyError(
                    "bench-output verification FAILED: the timed output does not match the "
                    "reference on the live input state (memoization/caching suspected)")
        return statistics.median(times)

    try:
        t_kernel = _bench(run, req["n_iters"], verify_tol=spec.tol(bench[0].dtype))
        t_eager = _bench(eager, req["n_iters"])
        t_comp = 0.0 if eager_only else _bench(compiled, req["n_iters"])
    except _BenchVerifyError as e:
        return Result(status="incorrect", n_shapes_passed=passed, feedback=str(e))
    except Exception as e:
        return Result(status="runtime_fail", correct=True, n_shapes_passed=passed,
                      feedback=f"correct but bench failed: {type(e).__name__}: {str(e)[:160]}")

    # STRONG baseline (opt-in): torch.compile max-autotune. Slow to compile; only used for
    # the self-test gate and final best-kernel validation, not the inner search loop.
    t_max = 0.0
    if req.get("strong") and not eager_only:
        try:
            cmax = torch.compile(spec.reference, mode="max-autotune-no-cudagraphs")
            t_max = _bench(cmax, req["n_iters"])
        except Exception:
            t_max = 0.0

    return Result(
        status="ok", correct=True, n_shapes_passed=passed, max_abs_err=worst,
        latency_ms=round(t_kernel, 5), eager_ms=round(t_eager, 5), compile_ms=round(t_comp, 5),
        maxauto_ms=round(t_max, 5),
        speedup_eager=round(t_eager / t_kernel, 4),
        speedup_compile=round(t_comp / t_kernel, 4) if t_comp > 0 else 0.0,
        speedup_maxauto=round(t_max / t_kernel, 4) if t_max > 0 else 0.0,
        feedback=(f"PASS {passed} cases | {t_kernel:.4f}ms | {t_eager/t_kernel:.2f}x eager"
                  + ("" if eager_only else f", {t_comp/t_kernel:.2f}x compile")
                  + (f", {t_max/t_kernel:.2f}x max-autotune" if t_max > 0 else "")))


# ============================ self-test (the D1-D2 gate) ===============================
def _selftest():
    """Prove the moat: a hand-written GOLD kernel passes with an honest speedup, and a
    deliberately-WRONG kernel is REJECTED. Mirrors dvwa_oracle.__main__'s positive +
    anti-cheat negatives. Writes the verdict durably to reports/."""
    seeddir = HERE / "seed_kernels"
    cases = [
        ("rmsnorm", "rmsnorm.py", "GOLD", "ok"),
        ("softmax", "softmax.py", "GOLD", "ok"),
        ("swiglu", "swiglu.py", "GOLD", "ok"),
        ("add_rmsnorm", "add_rmsnorm.py", "GOLD", "ok"),
        ("rope", "rope.py", "GOLD", "ok"),
        ("layernorm", "layernorm.py", "GOLD", "ok"),
        ("add_layernorm", "add_layernorm.py", "GOLD", "ok"),
        ("geglu", "geglu.py", "GOLD", "ok"),
        ("qknorm_rope", "qknorm_rope.py", "GOLD (fusion chain)", "ok"),
        # V2 standalone ops
        ("softcap_softmax", "softcap_softmax.py", "GOLD (Gemma2 softcap)", "ok"),
        ("rmsnorm_gemma", "rmsnorm_gemma.py", "GOLD (Gemma 1+w)", "ok"),
        ("glu", "glu.py", "GOLD (original GLU)", "ok"),
        ("rope_interleaved", "rope_interleaved.py", "GOLD (GPT-J RoPE)", "ok"),
        ("cross_entropy", "cross_entropy.py", "GOLD (fused CE)", "ok"),
        # V2.7 invention suite
        ("cumsum", "cumsum.py", "GOLD (prefix scan)", "ok"),
        ("entropy", "entropy.py", "GOLD (entropy-from-logits)", "ok"),
        ("kl_div", "kl_div.py", "GOLD (logit-pair KL)", "ok"),
        ("rmsnorm", "rmsnorm_wrong.py", "NEGATIVE-CONTROL (no rsqrt)", "incorrect"),
        ("softmax", "softmax_wrong.py", "NEGATIVE-CONTROL (no max-subtract)", "incorrect"),
        ("add_rmsnorm", "add_rmsnorm_wrong.py", "NEGATIVE-CONTROL (drops residual in stat)", "incorrect"),
        ("rope", "rope_wrong.py", "NEGATIVE-CONTROL (drops rotate_half negation)", "incorrect"),
        ("layernorm", "layernorm_wrong.py", "NEGATIVE-CONTROL (no mean-subtraction)", "incorrect"),
        ("add_layernorm", "add_layernorm_wrong.py", "NEGATIVE-CONTROL (drops residual in stat)", "incorrect"),
        ("geglu", "geglu_wrong.py", "NEGATIVE-CONTROL (SiLU instead of GELU)", "incorrect"),
        ("qknorm_rope", "qknorm_rope_wrong.py", "NEGATIVE-CONTROL (rms scale dropped on rotated half)", "incorrect"),
        ("softcap_softmax", "softcap_softmax_wrong.py", "NEGATIVE-CONTROL (no softcap)", "incorrect"),
        ("rmsnorm_gemma", "rmsnorm_gemma_wrong.py", "NEGATIVE-CONTROL (drops the +1)", "incorrect"),
        ("glu", "glu_wrong.py", "NEGATIVE-CONTROL (silu not sigmoid)", "incorrect"),
        ("rope_interleaved", "rope_interleaved_wrong.py", "NEGATIVE-CONTROL (drops negation)", "incorrect"),
        ("cross_entropy", "cross_entropy_wrong.py", "NEGATIVE-CONTROL (no max-subtract)", "incorrect"),
        ("cumsum", "cumsum_wrong.py", "NEGATIVE-CONTROL (no carry across blocks)", "incorrect"),
        ("entropy", "entropy_wrong.py", "NEGATIVE-CONTROL (no max-subtract)", "incorrect"),
        ("kl_div", "kl_div_wrong.py", "NEGATIVE-CONTROL (2nd lse not max-subtracted)", "incorrect"),
        # V2 ANTI-GAMING controls — exploits the v1 harness would have ACCEPTED:
        ("softmax", "softmax_cheat_shape.py", "ANTI-GAMING (special-cases bench shape)", "incorrect"),
        ("softmax", "softmax_cheat_memo.py", "ANTI-GAMING (memoizes by input pointer)", "incorrect"),
        ("silu", "silu_mutate.py", "ANTI-GAMING (mutates input in-place)", "incorrect"),
    ]
    report = {"machine": None, "cases": []}
    try:
        import torch
        report["machine"] = torch.cuda.get_device_name(0)
    except Exception:
        pass
    print(f"OUROBOROS harness self-test on {report['machine']}\n" + "=" * 70)
    ok_all = True
    for spec_name, fname, label, expect in cases:
        src = (seeddir / fname).read_text()
        # GOLD kernels get the STRONG (max-autotune) baseline; negative controls are rejected
        # before benchmarking so strong is moot for them (keep them fast).
        res = evaluate(src, spec_name, n_shapes=8, n_iters=100, strong=(expect == "ok"))
        hit = (res.status == expect) or (expect == "incorrect" and res.status in ("incorrect", "runtime_fail"))
        ok_all &= hit
        mark = "OK " if hit else "XX "
        extra = ""
        if res.status == "ok":
            extra = (f"  {res.latency_ms:.4f}ms  {res.speedup_eager:.2f}x eager  "
                     f"{res.speedup_compile:.2f}x compile  "
                     f"{res.speedup_maxauto:.2f}x max-autotune")
        print(f"  {mark}{spec_name:12} {label:34} -> {res.status:12}{extra}")
        if res.status != "ok":
            print(f"        feedback: {res.feedback}")
        report["cases"].append({"spec": spec_name, "kernel": fname, "label": label,
                                "expect": expect, "got": res.status, "pass": hit, **res.to_dict()})
    repath = HERE / "reports" / "harness_selftest.json"
    repath.write_text(json.dumps(report, indent=2))
    verdict = "ALL GREEN — moat proven" if ok_all else "FAILURES — harness not trustworthy yet"
    print("=" * 70 + f"\n{verdict}\nreport -> {repath}")
    return 0 if ok_all else 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--worker", action="store_true", help="internal: run one JSON request from stdin")
    args = ap.parse_args()
    if args.worker:
        req = json.loads(sys.stdin.read())
        res = _worker(req)
        print("RESULT:" + json.dumps(res.to_dict()))
        return
    sys.exit(_selftest())


if __name__ == "__main__":
    main()