250 lines
9.5 KiB
Bash
Executable File
250 lines
9.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# run_track.sh — Outer loop: invoke AI agent to optimize a CUDA kernel track.
|
|
#
|
|
# Usage: run_track.sh <track-name> --experiments <N> [--quick-until-improvement] [--max-hours <H>]
|
|
|
|
set -uo pipefail
|
|
# Note: NOT using set -e — we handle errors explicitly to keep the loop running
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
|
REPO_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
|
|
|
# ---- Parse args ----
|
|
TRACK="${1:?Usage: run_track.sh <track-name> --experiments <N>}"
|
|
shift
|
|
EXPERIMENTS=10
|
|
QUICK_UNTIL_IMPROVEMENT=false
|
|
MAX_HOURS=8
|
|
|
|
while [[ $# -gt 0 ]]; do
|
|
case "$1" in
|
|
--experiments) EXPERIMENTS="$2"; shift 2 ;;
|
|
--quick-until-improvement) QUICK_UNTIL_IMPROVEMENT=true; shift ;;
|
|
--max-hours) MAX_HOURS="$2"; shift 2 ;;
|
|
*) echo "Unknown arg: $1" >&2; exit 1 ;;
|
|
esac
|
|
done
|
|
|
|
TRACK_DIR="$SCRIPT_DIR/$TRACK"
|
|
if [[ ! -f "$TRACK_DIR/program.md" ]]; then
|
|
echo "Error: $TRACK_DIR/program.md not found" >&2
|
|
exit 1
|
|
fi
|
|
|
|
TARGET_FILE=$(jq -r '.target_file' "$TRACK_DIR/baseline.json")
|
|
HISTORY_FILE="$TRACK_DIR/history.jsonl"
|
|
BASELINE_FILE="$TRACK_DIR/baseline.json"
|
|
START_TIME=$(date +%s)
|
|
MAX_SECONDS=$((MAX_HOURS * 3600))
|
|
|
|
# ---- Create experiment branch ----
|
|
BRANCH_NAME="autoresearch/$TRACK/$(date +%Y%m%d-%H%M%S)"
|
|
cd "$REPO_DIR"
|
|
git checkout -b "$BRANCH_NAME" 2>/dev/null || true
|
|
BASELINE_SHA=$(git rev-parse HEAD)
|
|
echo ">>> Branch: $BRANCH_NAME (base: ${BASELINE_SHA:0:8})" >&2
|
|
|
|
# ---- Establish baseline ----
|
|
echo ">>> Establishing baseline..." >&2
|
|
BASELINE_RESULT=$("$SCRIPT_DIR/run_experiment.sh" "$TRACK")
|
|
echo "$BASELINE_RESULT" | jq . >&2
|
|
|
|
BASELINE_STATUS=$(echo "$BASELINE_RESULT" | jq -r '.status')
|
|
if [[ "$BASELINE_STATUS" != "success" ]]; then
|
|
echo "Error: baseline benchmark failed: $BASELINE_STATUS" >&2
|
|
exit 1
|
|
fi
|
|
|
|
# Update baseline with current measurements
|
|
BASELINE_TG=$(echo "$BASELINE_RESULT" | jq '.tg128')
|
|
BASELINE_PPL=$(echo "$BASELINE_RESULT" | jq '.ppl // empty')
|
|
jq --argjson tg "$BASELINE_TG" '.tg128 = $tg' "$BASELINE_FILE" > "$BASELINE_FILE.tmp" && mv "$BASELINE_FILE.tmp" "$BASELINE_FILE"
|
|
if [[ -n "${BASELINE_PPL:-}" ]]; then
|
|
jq --argjson ppl "$BASELINE_PPL" '.ppl = $ppl' "$BASELINE_FILE" > "$BASELINE_FILE.tmp" && mv "$BASELINE_FILE.tmp" "$BASELINE_FILE"
|
|
fi
|
|
|
|
echo ">>> Baseline: ${BASELINE_TG} t/s" >&2
|
|
|
|
# ---- Counters ----
|
|
CONSECUTIVE_FAILURES=0
|
|
CONSECUTIVE_NO_IMPROVEMENT=0
|
|
TOTAL_KEPT=0
|
|
TOTAL_REVERTED=0
|
|
BEST_TG=$BASELINE_TG
|
|
|
|
# ---- Experiment loop ----
|
|
for i in $(seq 1 "$EXPERIMENTS"); do
|
|
ELAPSED=$(( $(date +%s) - START_TIME ))
|
|
if (( ELAPSED > MAX_SECONDS )); then
|
|
echo ">>> Time limit reached (${MAX_HOURS}h). Stopping." >&2
|
|
break
|
|
fi
|
|
|
|
echo "" >&2
|
|
echo "================================================================" >&2
|
|
echo ">>> Experiment $i / $EXPERIMENTS" >&2
|
|
echo "================================================================" >&2
|
|
|
|
# Determine if we should skip PPL this round
|
|
QUICK_FLAG=""
|
|
if [[ "$QUICK_UNTIL_IMPROVEMENT" == "true" ]]; then
|
|
QUICK_FLAG="--quick"
|
|
fi
|
|
# Every 10th experiment always runs PPL
|
|
if (( i % 10 == 0 )); then
|
|
QUICK_FLAG=""
|
|
fi
|
|
|
|
# Build the agent prompt
|
|
LAST_RESULT=""
|
|
if [[ -f "$HISTORY_FILE" ]]; then
|
|
LAST_RESULT=$(tail -1 "$HISTORY_FILE" 2>/dev/null || echo "")
|
|
fi
|
|
|
|
STALL_HINT=""
|
|
if (( CONSECUTIVE_NO_IMPROVEMENT >= 5 )); then
|
|
STALL_HINT="IMPORTANT: The last $CONSECUTIVE_NO_IMPROVEMENT experiments showed no improvement. Try a fundamentally different approach — different algorithm, different memory access pattern, different thread mapping."
|
|
CONSECUTIVE_NO_IMPROVEMENT=0
|
|
fi
|
|
|
|
AGENT_PROMPT="You are optimizing a CUDA kernel. Read the program file and make ONE modification to improve performance.
|
|
|
|
Read: $TRACK_DIR/program.md
|
|
Read: $REPO_DIR/$TARGET_FILE
|
|
|
|
$(if [[ -n "$LAST_RESULT" ]]; then echo "Last experiment result: $LAST_RESULT"; fi)
|
|
$(if [[ -n "$STALL_HINT" ]]; then echo "$STALL_HINT"; fi)
|
|
|
|
Current best: ${BEST_TG} t/s (baseline: ${BASELINE_TG} t/s)
|
|
|
|
Rules:
|
|
- Modify ONLY: $TARGET_FILE
|
|
- Make exactly ONE conceptual change per experiment
|
|
- Do NOT run any builds, benchmarks, or git commands
|
|
- After editing, briefly describe what you changed and why (one line to stderr)"
|
|
|
|
# Invoke the agent
|
|
echo ">>> Invoking agent..." >&2
|
|
claude -p "$AGENT_PROMPT" \
|
|
--allowedTools Read,Edit \
|
|
--max-turns 20 \
|
|
2>&1 | tee /dev/stderr | tail -1 > /dev/null || true
|
|
|
|
# Check if the file was actually modified
|
|
if ! git diff --quiet -- "$TARGET_FILE" 2>/dev/null; then
|
|
# Check no OTHER source files were modified (ignore config/untracked)
|
|
OTHER_CHANGES=$(git diff --name-only -- '*.cu' '*.cuh' '*.cpp' '*.h' '*.c' | grep -vF "$TARGET_FILE" | head -5)
|
|
if [[ -n "$OTHER_CHANGES" ]]; then
|
|
echo ">>> SAFETY: Agent modified non-target files: $OTHER_CHANGES — reverting all" >&2
|
|
git checkout -- .
|
|
RESULT='{"status": "safety_revert", "error": "modified non-target files"}'
|
|
else
|
|
# Run the experiment
|
|
echo ">>> Running experiment..." >&2
|
|
RESULT=$("$SCRIPT_DIR/run_experiment.sh" "$TRACK" $QUICK_FLAG)
|
|
echo "$RESULT" | jq . >&2
|
|
fi
|
|
else
|
|
echo ">>> Agent made no changes. Skipping." >&2
|
|
RESULT='{"status": "no_change"}'
|
|
fi
|
|
|
|
STATUS=$(echo "$RESULT" | jq -r '.status')
|
|
TG=$(echo "$RESULT" | jq -r '.tg128 // 0')
|
|
|
|
# Decide: keep or revert
|
|
KEPT=false
|
|
if [[ "$STATUS" == "success" ]]; then
|
|
# Check if it's actually faster
|
|
IMPROVEMENT=$(echo "$TG > $BEST_TG" | bc -l 2>/dev/null || echo 0)
|
|
if [[ "$IMPROVEMENT" == "1" ]]; then
|
|
# Speed improvement found. If we were in quick mode, validate PPL now.
|
|
if [[ -n "$QUICK_FLAG" && "$QUICK_UNTIL_IMPROVEMENT" == "true" ]]; then
|
|
echo ">>> Speed improvement found (+$(echo "$RESULT" | jq -r '.delta_pct')). Validating PPL..." >&2
|
|
PPL_RESULT=$("$SCRIPT_DIR/run_experiment.sh" "$TRACK")
|
|
PPL_STATUS=$(echo "$PPL_RESULT" | jq -r '.status')
|
|
if [[ "$PPL_STATUS" == "ppl_regression" ]]; then
|
|
echo ">>> PPL regression — reverting despite speed gain" >&2
|
|
git checkout -- "$TARGET_FILE"
|
|
KEPT=false
|
|
STATUS="ppl_regression"
|
|
RESULT="$PPL_RESULT"
|
|
else
|
|
KEPT=true
|
|
fi
|
|
else
|
|
KEPT=true
|
|
fi
|
|
|
|
if [[ "$KEPT" == "true" ]]; then
|
|
BEST_TG="$TG"
|
|
git add "$TARGET_FILE"
|
|
SUMMARY=$(git diff --cached --stat | head -1)
|
|
git commit -m "autoresearch($TRACK): +$(echo "$RESULT" | jq -r '.delta_pct') tg128 (experiment $i)" --no-verify
|
|
# Update baseline
|
|
jq --argjson tg "$TG" '.tg128 = $tg' "$BASELINE_FILE" > "$BASELINE_FILE.tmp" && mv "$BASELINE_FILE.tmp" "$BASELINE_FILE"
|
|
PPL_VAL=$(echo "$RESULT" | jq '.ppl // empty')
|
|
if [[ -n "${PPL_VAL:-}" ]]; then
|
|
jq --argjson ppl "$PPL_VAL" '.ppl = $ppl' "$BASELINE_FILE" > "$BASELINE_FILE.tmp" && mv "$BASELINE_FILE.tmp" "$BASELINE_FILE"
|
|
fi
|
|
TOTAL_KEPT=$((TOTAL_KEPT + 1))
|
|
CONSECUTIVE_NO_IMPROVEMENT=0
|
|
CONSECUTIVE_FAILURES=0
|
|
echo ">>> KEPT: ${TG} t/s (+$(echo "$RESULT" | jq -r '.delta_pct'))" >&2
|
|
fi
|
|
else
|
|
echo ">>> No improvement (${TG} vs best ${BEST_TG}). Reverting." >&2
|
|
git checkout -- "$TARGET_FILE"
|
|
CONSECUTIVE_NO_IMPROVEMENT=$((CONSECUTIVE_NO_IMPROVEMENT + 1))
|
|
fi
|
|
elif [[ "$STATUS" == "build_failed" ]]; then
|
|
echo ">>> Build failed. Reverting." >&2
|
|
git checkout -- "$TARGET_FILE"
|
|
CONSECUTIVE_FAILURES=$((CONSECUTIVE_FAILURES + 1))
|
|
elif [[ "$STATUS" == "runtime_crash" ]]; then
|
|
echo ">>> Runtime crash. Reverting." >&2
|
|
git checkout -- "$TARGET_FILE"
|
|
CONSECUTIVE_FAILURES=$((CONSECUTIVE_FAILURES + 1))
|
|
else
|
|
echo ">>> Status: $STATUS. Reverting." >&2
|
|
git checkout -- "$TARGET_FILE"
|
|
fi
|
|
|
|
if [[ "$KEPT" == "false" ]]; then
|
|
TOTAL_REVERTED=$((TOTAL_REVERTED + 1))
|
|
fi
|
|
|
|
# Log to history
|
|
TIMESTAMP=$(date -u +%Y-%m-%dT%H:%M:%SZ)
|
|
echo "$RESULT" | jq -c --arg exp "$i" --arg ts "$TIMESTAMP" --argjson kept "$KEPT" \
|
|
'. + {experiment: ($exp|tonumber), timestamp: $ts, kept: $kept}' >> "$HISTORY_FILE"
|
|
|
|
# Safety: too many consecutive failures
|
|
if (( CONSECUTIVE_FAILURES >= 3 )); then
|
|
echo ">>> WARNING: 3 consecutive failures. Pausing for review." >&2
|
|
echo ">>> Last error: $(echo "$RESULT" | jq -r '.error // .status')" >&2
|
|
CONSECUTIVE_FAILURES=0
|
|
# Don't exit — just reset counter and let the stall hint kick in
|
|
fi
|
|
|
|
# GPU cooldown
|
|
sleep 5
|
|
done
|
|
|
|
# ---- Summary ----
|
|
echo "" >&2
|
|
echo "================================================================" >&2
|
|
echo ">>> AUTORESEARCH COMPLETE" >&2
|
|
echo ">>> Track: $TRACK" >&2
|
|
echo ">>> Experiments: $((TOTAL_KEPT + TOTAL_REVERTED))" >&2
|
|
echo ">>> Kept: $TOTAL_KEPT" >&2
|
|
echo ">>> Reverted: $TOTAL_REVERTED" >&2
|
|
echo ">>> Baseline: ${BASELINE_TG} t/s → Best: ${BEST_TG} t/s" >&2
|
|
if [[ "$BEST_TG" != "$BASELINE_TG" ]]; then
|
|
TOTAL_GAIN=$(echo "scale=1; ($BEST_TG - $BASELINE_TG) / $BASELINE_TG * 100" | bc -l 2>/dev/null || echo "?")
|
|
echo ">>> Total improvement: +${TOTAL_GAIN}%" >&2
|
|
fi
|
|
echo ">>> Branch: $BRANCH_NAME" >&2
|
|
echo ">>> History: $HISTORY_FILE" >&2
|
|
echo "================================================================" >&2
|