#!/usr/bin/env bash
# bench-pascal-mmvq.sh -- paired A/B benchmark for the Pascal DP4A mmvq nwarps change.
#
# Builds BASELINE (HEAD~1) and PR (HEAD) into separate binaries, then runs
# llama-bench alternating between them. Consumer Pascal can't lock clocks, so we
# report the PAIRED per-iteration delta, which cancels thermal/boost drift.
# Decode (tg) is the mmvq path this PR changes; prefill (pp) is the untouched MMQ
# path, kept as a no-regression control (expected ~0%) with a -ub batch sweep.
#
#   ./bench-pascal-mmvq.sh                     # ITERS=6 REPS=10 defaults
#   SKIP_BUILD=1 ./bench-pascal-mmvq.sh        # reuse existing builds
#
# reproducing: to reproduce, edit only the two blocks flagged "reproducing:" below
# (MODELS paths, and CUDA_HOST_CC iff your build hits a CUDA/host-GCC mismatch).
set -euo pipefail

REPO="${REPO:-$(cd "$(dirname "$0")" && pwd)}"
BASELINE_REF="${BASELINE_REF:-HEAD~1}"          # "before" commit
PR_REF="${PR_REF:-HEAD}"                         # "after" commit

# reproducing: point these at your own GGUFs (or set MODEL_DIR=/path). "LABEL|path".
MDL="${MODEL_DIR:-$REPO/models/dl}"
MODELS=(
  "Q4_0|$MDL/Llama-3.2-3B-Instruct-Q4_0.gguf"
  "Q4_K_M|$MDL/Llama-3.2-3B-Instruct-Q4_K_M.gguf"
  "Q5_K_M|$MDL/Llama-3.2-3B-Instruct-Q5_K_M.gguf"
  "Q8_0|$MDL/Llama-3.2-3B-Instruct-Q8_0.gguf"
  "Q6_K|$MDL/Llama-3.2-3B-Instruct-Q6_K.gguf"
)

ITERS="${ITERS:-6}"; REPS="${REPS:-10}"; NGEN="${NGEN:-128}"; NGL="${NGL:-99}"; COOLDOWN="${COOLDOWN:-3}"

# "cfgname|llama-bench args" per invocation. tg*=decode (mmvq, the win); pp*=prefill
# (MMQ control, ~0%). reproducing: trim pp rows / shrink -p/-ub if VRAM-limited.
BENCH=(
  "tg128|-p 0 -n $NGEN -ub 512"
  "pp512|-p 512 -n 0 -ub 512"
  "pp512_ub128|-p 512 -n 0 -ub 128"
  "pp2048|-p 2048 -n 0 -ub 512"
)

CUDA_ARCH="${CUDA_ARCH:-61}"                      # sm_61 = GTX 10xx/P40/P4; 60 for P100
# reproducing: leave empty normally. Set to a host GCC <= 11 only if the build fails
# with "parameter packs not expanded with '...'" (CUDA <= 11.5 vs new host GCC).
CUDA_HOST_CC="${CUDA_HOST_CC:-}"
CMAKE_FLAGS="${CMAKE_FLAGS:--DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=$CUDA_ARCH}"
[[ -n "$CUDA_HOST_CC" ]] && CMAKE_FLAGS="$CMAKE_FLAGS -DCMAKE_CUDA_HOST_COMPILER=$CUDA_HOST_CC"

WORK="$REPO/.bench-work"; RESULTS="$WORK/results"
BASE_TREE="$WORK/baseline-src"; BASE_BUILD="$WORK/build-baseline"; PR_BUILD="$REPO/build-pr"
BASE_BIN="$BASE_BUILD/bin/llama-bench"; PR_BIN="$PR_BUILD/bin/llama-bench"

log() { printf '\033[1;36m[bench]\033[0m %s\n' "$*"; }

build_all() {
  [[ "${SKIP_BUILD:-0}" == "1" ]] && { log "SKIP_BUILD=1, reusing binaries"; return; }
  cd "$REPO"
  log "Building PR ($PR_REF)"
  cmake -B "$PR_BUILD" $CMAKE_FLAGS -S "$REPO" >/dev/null
  cmake --build "$PR_BUILD" --target llama-bench -j"$(nproc)" >/dev/null
  log "Building baseline ($BASELINE_REF) in a worktree"
  git worktree remove --force "$BASE_TREE" 2>/dev/null || true
  git worktree add --detach "$BASE_TREE" "$BASELINE_REF" >/dev/null
  cmake -B "$BASE_BUILD" $CMAKE_FLAGS -S "$BASE_TREE" >/dev/null
  cmake --build "$BASE_BUILD" --target llama-bench -j"$(nproc)" >/dev/null
}

# run <bin> <model> <extra_args> <out.json>
run() {
  # shellcheck disable=SC2086  -- $3 is a deliberately word-split arg list
  "$1" -m "$2" $3 -ngl "$NGL" -r "$REPS" -o json > "$4" 2>/dev/null
}

main() {
  local have=0; for e in "${MODELS[@]}"; do [[ -f "${e#*|}" ]] && have=1; done
  [[ "$have" == "1" ]] || { echo "ERROR: no MODELS paths exist; edit MODELS or set MODEL_DIR."; exit 1; }
  rm -rf "$RESULTS"; mkdir -p "$RESULTS"              # always override previous run
  build_all
  [[ -x "$BASE_BIN" && -x "$PR_BIN" ]] || { echo "missing binaries; run without SKIP_BUILD"; exit 1; }

  log "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo unknown)"
  log "ITERS=$ITERS REPS=$REPS NGEN=$NGEN NGL=$NGL"
  log "Warmup..."; run "$PR_BIN" "${MODELS[0]#*|}" "-p 0 -n $NGEN -ub 512" "$RESULTS/warmup.json" || true

  for ((i=1; i<=ITERS; i++)); do
    for e in "${MODELS[@]}"; do
      label="${e%%|*}"; path="${e#*|}"; safe="${label//[^A-Za-z0-9]/_}"
      [[ -f "$path" ]] || { echo "  !! skipping missing model: $path"; continue; }
      (( i % 2 == 1 )) && order=(baseline pr) || order=(pr baseline)   # swap so neither build always runs hotter
      for cfg in "${BENCH[@]}"; do
        cname="${cfg%%|*}"; cargs="${cfg#*|}"; csafe="${cname//[^A-Za-z0-9]/_}"
        for b in "${order[@]}"; do                    # A/B adjacent per (model,cfg) for tight drift cancellation
          bin="$BASE_BIN"; [[ "$b" == "pr" ]] && bin="$PR_BIN"
          log "iter $i/$ITERS  $label  $cname  [$b]"
          run "$bin" "$path" "$cargs" "$RESULTS/${b}__${safe}__${csafe}__${i}.json"
          sleep "$COOLDOWN"
        done
      done
    done
  done

  log "Aggregating..."
  RESULTS="$RESULTS" BENCH_CFGS="$(printf '%s\n' "${BENCH[@]%%|*}")" python3 - "${MODELS[@]}" <<'PY'
import json, glob, os, sys, math, statistics as st

results = os.environ["RESULTS"]
labels  = [a.split("|", 1)[0] for a in sys.argv[1:]]
cfgs    = os.environ["BENCH_CFGS"].split()
san     = lambda s: "".join(c if c.isalnum() else "_" for c in s)

def samples(build, safe, csafe):                        # {iter: avg_ts}, for pairing
    out = {}
    for f in sorted(glob.glob(f"{results}/{build}__{safe}__{csafe}__*.json")):
        try: out[int(f.rsplit("__", 1)[1].split(".")[0])] = float(json.load(open(f))[0]["avg_ts"])
        except Exception: pass
    return out

def row(label, cfg, safe, csafe):
    bd, pd = samples("baseline", safe, csafe), samples("pr", safe, csafe)
    if not bd or not pd: return None
    bm, pm = st.mean(list(bd.values())), st.mean(list(pd.values()))
    diffs = [(pd[i]-bd[i])/bd[i]*100.0 for i in sorted(set(bd) & set(pd))]   # per-iter paired Δ%
    n, up = len(diffs), sum(x > 0 for x in diffs)
    dsd = st.stdev(diffs) if n > 1 else float("nan")
    tval = st.mean(diffs)/(dsd/math.sqrt(n)) if n > 1 and dsd else (float("inf") if n > 1 else float("nan"))
    return (label, cfg, bm, pm, (pm-bm)/bm*100.0, dsd, tval, f"{up}/{n}")

decode, prefill, ratios = [], [], []
for cfg in cfgs:
    dest = decode if cfg.lower().startswith("tg") else prefill
    for label in labels:
        r = row(label, cfg, san(label), san(cfg))
        if r: dest.append(r); (ratios.append(r[3]/r[2]) if dest is decode else None)

def table(title, rows, note):
    W = 100; print("\n" + "="*W); print("  " + title); print("="*W)
    print(f"{'quant':<9}{'config':<13}{'baseline':>11}{'PR t/s':>11}{'mean Δ':>9}{'paired σ%':>11}{'t':>8}{'pairs':>7}")
    print("-"*W)
    for l, c, bm, pm, d, sd, t, p in rows:
        print(f"{l:<9}{c:<13}{bm:>11.2f}{pm:>11.2f}{d:>+8.2f}%{sd:>10.2f}%{t:>8.2f}{p:>7}")
    print("-"*W); print(note)

table("DECODE (tg) : BASELINE vs PR  -- the mmvq path this PR changes", decode,
      "pairs = iters where PR>baseline. t = paired t-statistic of per-iter Δ% (drift-cancelling).")
geo = math.exp(sum(map(math.log, ratios))/len(ratios)) if ratios else None
if geo: print(f"  CUMULATIVE decode (geomean): {(geo-1)*100:+.2f}%")
if prefill:
    table("PREFILL (pp) : BASELINE vs PR  -- untouched MMQ path (control, expect ~0%)", prefill,
          "No-regression / batch-coverage control; this change is scoped to decode.")

def md(title, rows):
    print(f"\n**{title}**\n")
    print("| quant | config | baseline t/s | patched t/s | mean Δ | paired σ% | t | pairs up |")
    print("|-------|--------|--------------|-------------|--------|-----------|---|----------|")
    for l, c, bm, pm, d, sd, t, p in rows:
        print(f"| {l} | {c} | {bm:.2f} | {pm:.2f} | {d:+.2f}% | {sd:.2f}% | {t:.2f} | {p} |")

md("Decode (tg) — mmvq path this PR changes", decode)
if geo: print(f"\nGeomean decode speedup: **{(geo-1)*100:+.2f}%**")
if prefill: md("Prefill (pp) — MMQ control (expected ~0%, no regression)", prefill)
PY
}

main "$@"