#!/usr/bin/env bash
# mi50_power_context_sweep.sh
#
# Two-axis sweep for dual MI50s: power limit x context size.
# For each (power, context) pair it runs llama-bench on gpt-oss and logs
# prefill (pp) and decode (tg) tok/s plus edge temp and actual draw.
#
# The question this answers: the power level that is "free" at small
# context (decode is HBM-bandwidth bound, core clock barely matters) may
# NOT be free at large context, where prefill is compute-heavy and a low
# core clock could bite. This grid shows where, if anywhere, that happens.
#
# Set MODEL below. Assumes llama-bench (llama.cpp ROCm build) in PATH.
# Run in tmux; a full grid takes a while (big model reloads per context).

set -uo pipefail   # not -e: a single failed run should not kill the sweep

# ------------------------------- config -------------------------------
MODEL="${MODEL:-$HOME/.cache/llama.cpp/openai_gpt-oss-120b-Q4_K_M/openai_gpt-oss-120b-Q4_K_M-00001-of-00002.gguf}"
NGL=999
GPUS=(0 1)

# power levels in watts (stock cap 300, sweet spot expected 180-225)
POWER_LEVELS=(225 200 180 160)

# context sizes: (prompt_tokens gen_tokens) pairs.
# small, medium, large, xlarge. gen held at 128 so decode is comparable;
# prompt scales to exercise prefill compute at each context.
CTX_PROMPT=(512 4096 16384 32768)
CTX_GEN=128

COOLDOWN=20                 # seconds between runs to normalize temps
REPS=2                      # runs per cell, averaged (reduces noise)
LOGFILE="mi50_pxc_sweep_$(date +%Y%m%d_%H%M%S).csv"
# ----------------------------------------------------------------------

command -v rocm-smi   >/dev/null || { echo "rocm-smi not found"; exit 1; }
command -v llama-bench >/dev/null || { echo "llama-bench not found in PATH"; exit 1; }
[[ -f "$MODEL" ]] || { echo "Model not found: $MODEL"; exit 1; }

echo "power_w,ctx_prompt,gen_tokens,pp_tok_s,tg_tok_s,edge_temp_c,draw_w" > "$LOGFILE"
echo "Logging to $LOGFILE"
echo "Model: $MODEL"

restore_power() {
  echo ""
  echo "Restoring stock power limit on all GPUs..."
  for g in "${GPUS[@]}"; do
    rocm-smi -d "$g" --resetpoweroverdrive >/dev/null 2>&1 || true
  done
}
trap restore_power EXIT INT TERM

# pull a float from arbitrary rocm-smi output, first match
grab_float() { grep -oE '[0-9]+\.[0-9]+' | head -1; }

# extract the tok/s value (last numeric field) from a llama-bench csv row
# matching a given test tag (e.g. pp512 or tg128)
extract_toks() {
  local out="$1" tag="$2"
  echo "$out" \
    | grep -iE "[\",]${tag}[\",]" \
    | tail -1 \
    | awk -F',' '{gsub(/"/,"",$NF); print $NF}'
}

sample_temp() {
  rocm-smi -d 0 --showtemp 2>/dev/null | grep -i edge | grab_float
}
sample_draw() {
  rocm-smi -d 0 --showpower 2>/dev/null | grab_float
}

total_cells=$(( ${#POWER_LEVELS[@]} * ${#CTX_PROMPT[@]} ))
cell=0

for pw in "${POWER_LEVELS[@]}"; do
  echo ""
  echo "############ POWER = ${pw}W ############"
  for g in "${GPUS[@]}"; do
    rocm-smi -d "$g" --setpoweroverdrive "$pw" >/dev/null 2>&1 \
      || echo "  warn: failed to set ${pw}W on GPU $g"
  done
  sleep 3

  for pp in "${CTX_PROMPT[@]}"; do
    cell=$((cell+1))
    echo ""
    echo "--- [$cell/$total_cells] power=${pw}W  prompt=${pp}  gen=${CTX_GEN} ---"

    pp_sum=0; tg_sum=0; ok=0
    for r in $(seq 1 "$REPS"); do
      OUT=$(llama-bench -m "$MODEL" -ngl "$NGL" \
              -p "$pp" -n "$CTX_GEN" -o csv 2>/dev/null)

      PP=$(extract_toks "$OUT" "pp${pp}")
      TG=$(extract_toks "$OUT" "tg${CTX_GEN}")
      # fallback: positional (rows 2 and 3) if tag match missed
      [[ -z "${PP:-}" ]] && PP=$(echo "$OUT" | awk -F',' 'NR==2{gsub(/"/,"",$NF);print $NF}')
      [[ -z "${TG:-}" ]] && TG=$(echo "$OUT" | awk -F',' 'NR==3{gsub(/"/,"",$NF);print $NF}')

      if [[ -n "${PP:-}" && -n "${TG:-}" ]]; then
        pp_sum=$(awk -v a="$pp_sum" -v b="$PP" 'BEGIN{print a+b}')
        tg_sum=$(awk -v a="$tg_sum" -v b="$TG" 'BEGIN{print a+b}')
        ok=$((ok+1))
      else
        echo "    rep $r: parse miss (pp=${PP:-NA} tg=${TG:-NA})"
      fi
      sleep "$COOLDOWN"
    done

    if [[ "$ok" -gt 0 ]]; then
      PP_AVG=$(awk -v s="$pp_sum" -v n="$ok" 'BEGIN{printf "%.2f", s/n}')
      TG_AVG=$(awk -v s="$tg_sum" -v n="$ok" 'BEGIN{printf "%.2f", s/n}')
    else
      PP_AVG=NA; TG_AVG=NA
    fi

    TEMP=$(sample_temp); DRAW=$(sample_draw)
    echo "    => pp: ${PP_AVG} tok/s | tg: ${TG_AVG} tok/s | temp: ${TEMP:-NA}C | draw: ${DRAW:-NA}W"
    echo "${pw},${pp},${CTX_GEN},${PP_AVG},${TG_AVG},${TEMP:-NA},${DRAW:-NA}" >> "$LOGFILE"
  done
done

echo ""
echo "======================================================================"
echo "Sweep complete. Raw results in $LOGFILE"
echo ""
column -t -s',' "$LOGFILE"

echo ""
echo "Quick read: for each context size, compare tg across power levels."
echo "If tg holds flat as power drops, that context is bandwidth-bound and"
echo "low power is free. If tg falls off at low power for large context,"
echo "that is prefill compute starvation -- keep power higher for that regime."