agent.py

#!/usr/bin/env python3
"""
Kernelking — CUDA Kernel Optimization Agent
============================================
An AI agent that automatically optimizes GPU kernels using
Nsight Compute profiling + Claude API analysis.

Run this script on your LOCAL machine. It:
  1. SSH → reads the transpose kernel from the remote GPU server
  2. SSH → generates & runs profile_transpose.sh to get baseline ncu report
  3. Downloads the ncu report details
  4. Calls Claude API to analyze the report and produce an optimized kernel
  5. Uploads the optimized kernel, rebuilds the binary
  6. Re-profiles with the optimized kernel
  7. Compares Duration: improved → accept, degraded → rollback

Usage:
  python agent.py
"""

import os
import sys
import subprocess
import tempfile
import re
import json
import time
from pathlib import Path
from datetime import datetime

# ─────────────────────────────────────────────
# Stage A refactor (2026-04-16): primitives moved to kernelking/
#   - kernelking.config → ROOT, CONFIG_FILE, LOG_FILE, load_config, cfg,
#     class C, log, banner, die, _open_log
#   - kernelking.costs  → _CLAUDE_PRICING_USD_PER_M, _TOKEN_USAGE,
#     _price_for_model, _accumulate_usage, _estimated_cost_usd,
#     _format_cost_summary
#   - kernelking.remote → ssh_run, ssh_run_capture, ssh_capture,
#     scp_download, scp_upload, _close_ssh_master, _check_and_clean_gpu
#   - kernelking.claude → _claude_call, _strip_fences, _NO_ELIDE_RULE,
#     _CLAUDE_CONTINUE_PROMPT, MAX_CLAUDE_RECOVERY,
#     _classify_anthropic_exc, _messages_create_with_retry
# Re-imported here so the rest of agent.py keeps working with bare names.
# ─────────────────────────────────────────────
from kernelking.config import (  # noqa: E402
    ROOT,
    CONFIG_FILE,
    LOG_FILE,
    C,
    load_config,
    cfg,
    log,
    banner,
    die,
    _open_log,
)
from kernelking.costs import (  # noqa: E402
    _CLAUDE_PRICING_USD_PER_M,
    _TOKEN_USAGE,
    _price_for_model,
    _accumulate_usage,
    _estimated_cost_usd,
    _format_cost_summary,
)
from kernelking.remote import (  # noqa: E402
    ssh_run,
    ssh_run_capture,
    ssh_capture,
    scp_download,
    scp_upload,
    _close_ssh_master,
    _check_and_clean_gpu,
    _upload_kernel,
)
from kernelking.claude.client import (  # noqa: E402
    MAX_CLAUDE_RECOVERY,
    _CLAUDE_CONTINUE_PROMPT,
    _NO_ELIDE_RULE,
    _claude_call,
    _classify_anthropic_exc,
    _messages_create_with_retry,
    _strip_fences,
)
from kernelking import rundir  # noqa: E402

# Local plugin chain: pluggable KernelHandler modules that isolate
# kernel-shape concerns (namespace wrapping, Claude prose prepended to
# code, …) out of agent.py itself. See `handlers/__init__.py` for the
# hook contract.
sys.path.insert(0, str(ROOT))
from handlers import (  # noqa: E402
    HandlerContext,
    run_harness_prompt_hints,
    run_harness_fix_prompt_hints,
    run_post_process_code,
)


# ─────────────────────────────────────────────
# Run provenance + Claude token/cost accounting
#
# Inspired by zk-autoresearch's run_provenance.json + experiment cost log.
# Every agent run writes a one-time provenance header (nvcc version, driver,
# target repo SHA, local agent.py SHA, SDK version, model) so that six
# months from now we can still reconstruct exactly what produced a given
# memory/<kernel>.md entry. Claude token counts are also accumulated so we
# can report $ cost at end-of-run + attribute regressions to specific runs.
# ─────────────────────────────────────────────

_PROVENANCE: dict = {}


def _collect_provenance() -> dict:
    """Gather environment fingerprints from the remote GPU host + local repo.

    Called once at startup from preflight(). All calls are best-effort —
    missing values are stored as "?" rather than aborting the run.
    """
    p: dict = {
        "timestamp":      datetime.now().isoformat(timespec="seconds"),
        "claude_model":   cfg("CLAUDE_MODEL", "claude-sonnet-4-6"),
        "anthropic_sdk":  "?",
        "agent_sha":      "?",
        "nvcc_version":   "?",
        "driver_version": "?",
        "gpu_name":       "?",
        "target_repo_sha": "?",
        "target_repo_path": cfg("REMOTE_REPO_DIR", "?"),
    }
    # Local side — never fail.
    try:
        import anthropic
        p["anthropic_sdk"] = getattr(anthropic, "__version__", "?")
    except Exception:
        pass
    try:
        r = subprocess.run(
            ["git", "-C", str(ROOT), "rev-parse", "--short", "HEAD"],
            check=False, capture_output=True, text=True, timeout=3,
        )
        if r.returncode == 0:
            p["agent_sha"] = r.stdout.strip()
    except Exception:
        pass
    # Remote side — one pass, parse separately. Wrap each query with `||:`
    # so a single failure doesn't mask the others.
    remote_repo = cfg("REMOTE_REPO_DIR") or cfg("REMOTE_BASEDIR", "/home/ubuntu")
    query = (
        "echo '--NVCC--'; nvcc --version 2>/dev/null | tail -1 || echo '?'; "
        "echo '--DRIVER--'; nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1 || echo '?'; "
        "echo '--GPU--'; nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || echo '?'; "
        f"echo '--REPOSHA--'; (cd {remote_repo} && git rev-parse --short HEAD 2>/dev/null) || echo '?'"
    )
    try:
        r = ssh_run(query, check=False, capture=True)
        out = (r.stdout or "") if r.returncode == 0 else ""
        def _section(tag: str) -> str:
            m = re.search(rf"--{tag}--\n(.+?)(?:\n--|\Z)", out, re.DOTALL)
            return (m.group(1).strip() if m else "?") or "?"
        p["nvcc_version"]     = _section("NVCC")
        p["driver_version"]   = _section("DRIVER")
        p["gpu_name"]         = _section("GPU")
        p["target_repo_sha"]  = _section("REPOSHA")
    except Exception:
        pass
    return p


def _format_provenance_banner(p: dict) -> str:
    """Compact multi-line string suitable for the startup banner + memory header."""
    return (
        f"model={p['claude_model']}  "
        f"sdk={p['anthropic_sdk']}  "
        f"agent={p['agent_sha']}\n"
        f"gpu={p['gpu_name']}  "
        f"driver={p['driver_version']}  "
        f"nvcc={p['nvcc_version']}\n"
        f"target={p['target_repo_path']}@{p['target_repo_sha']}"
    )


# ─────────────────────────────────────────────
# Local codebase context  (LOCAL_CONTEXT_FILES)
# ─────────────────────────────────────────────

def load_local_context(kernel_code: str) -> str:
    """Read files listed in LOCAL_CONTEXT_FILES and return them as a prompt section.

    config.env format — comma or newline separated paths, relative to
    LOCAL_CODEBASE_PATH (if set) or absolute:

        LOCAL_CODEBASE_PATH=/Users/you/brevis-vm
        LOCAL_CONTEXT_FILES=vm/src/cuda_adaptor/gpuacc_struct/fri_commit.rs,
                            private-pico-gpu/ff/kb31_t.hpp

    Each file is included verbatim under a labeled section so Claude can see
    which extern "C" symbols Rust expects and what data structures the kernel uses.
    """
    base_raw  = cfg("LOCAL_CODEBASE_PATH", "").strip()
    files_raw = cfg("LOCAL_CONTEXT_FILES", "").strip()

    if not files_raw:
        log("LOCAL_CONTEXT_FILES not set — Claude will work without local context", C.DIM)
        return ""

    base = Path(base_raw) if base_raw else None

    parts: list[str] = []
    for entry in re.split(r"[,\n]+", files_raw):
        entry = entry.strip()
        if not entry:
            continue
        path = (base / entry) if (base and not Path(entry).is_absolute()) else Path(entry)
        if not path.exists():
            log(f"LOCAL_CONTEXT_FILES: not found — {path}", C.YELLOW)
            continue
        try:
            content = path.read_text(errors="replace")
            label   = str(path.relative_to(base)) if base and base in path.parents else path.name
            suffix  = path.suffix.lstrip(".")
            lang    = {"rs": "rust", "cu": "cuda", "hpp": "cpp", "h": "cpp"}.get(suffix, suffix)
            parts.append(f"### {label}\n```{lang}\n{content}\n```")
            log(f"Local context: loaded {label} ({len(content)} bytes)", C.GREEN)
        except Exception as e:
            log(f"LOCAL_CONTEXT_FILES: could not read {path}: {e}", C.YELLOW)

    if not parts:
        return ""

    return (
        "## Local codebase files (for context — extern \"C\" function names shown here "
        "MUST be preserved exactly in the kernel)\n\n"
        + "\n\n".join(parts)
    )


# ─────────────────────────────────────────────
# Step 1: Read kernel source from remote
# ─────────────────────────────────────────────
def step_read_kernel() -> tuple[str, str]:
    """Download the kernel .cu from remote.

    Returns (kernel_code, local_context) where local_context is the
    FFI + header snippet built from LOCAL_CODEBASE_PATH (may be empty).
    """
    banner("Step 1 — Read Kernel from Remote")
    remote_cu = cfg("REMOTE_KERNEL_CU")
    kernel_name = cfg("KERNEL_NAME", "transpose_kernel")
    log(f"Kernel : {kernel_name}", C.CYAN)
    log(f"File   : {remote_cu}", C.CYAN)
    code = ssh_capture(f"cat {remote_cu}")
    log(f"Kernel loaded ({len(code)} bytes, {len(code.splitlines())} lines)", C.GREEN)
    (ROOT / "kernel_baseline.cu").write_text(code)

    local_ctx = load_local_context(code)
    if not local_ctx:
        log("LOCAL_CODEBASE_PATH not set — Claude will work without local context", C.DIM)
    return code, local_ctx


# `NCU_KNOWLEDGE` is the ncu CLI flag reference — consumed by the legacy
# profile-script generation and crash-triage paths in
# ``kernelking.profile._legacy``. The live main() path does NOT need it:
# the optimize prompt pulls on-demand per-bottleneck skills from
# ``skills/<rule_id>.md`` via ``_load_skills_for_findings()`` below.
# See ``skills/SKILLS.md`` for the skill vs rule vs prompt vs knowledge
# glossary.


# ─────────────────────────────────────────────
# Stage D refactor (2026-04-16): knowledge subsystem moved to
# kernelking/knowledge/
#   - kernelking.knowledge.skills   → SKILLS_DIR, _SKILL_SIGNAL_RE,
#     _load_skills_for_findings, _build_skill_catalog
#   - kernelking.knowledge.memory   → MEMORY_DIR, CUDA_OPT_FILE, _MEMO_RE,
#     _memory_path, _append_memory, _load_memory_excerpt,
#     _parse_memory_file, _collect_all_memory_records, _extract_memo
#   - kernelking.knowledge.patterns → _group_patterns_by_tag,
#     _format_tag_section_md, _build_cross_kernel_patterns,
#     _load_patterns_for_findings
# Re-imported here so the rest of agent.py keeps working with bare names.
#
# Terminology note (see `skills/SKILLS.md` for the full glossary):
#   - skill  = markdown file in `skills/` describing one bottleneck
#   - rule   = Python `_rule_*` function + entry in `RULES` (fires skills)
#   - prompt = the `f"""…"""` block inside each `_claude_*` call site
# ─────────────────────────────────────────────
from kernelking.knowledge import (  # noqa: E402
    CUDA_OPT_FILE,
    MEMORY_DIR,
    SKILLS_DIR,
    _MEMO_RE,
    _SKILL_SIGNAL_RE,
    _append_memory,
    _build_cross_kernel_patterns,
    _build_skill_catalog,
    _collect_all_memory_records,
    _extract_memo,
    _format_tag_section_md,
    _group_patterns_by_tag,
    _load_memory_excerpt,
    _load_patterns_for_findings,
    _load_skills_for_findings,
    _memory_path,
    _parse_memory_file,
)


# ─────────────────────────────────────────────
# Stage C refactor (2026-04-16): harness subsystem moved to kernelking/harness/
#   - kernelking.harness.nvcc     → STANDALONE_SUBDIR, _standalone_dir,
#     _nvcc_path, _nvcc_arch, _kernel_include_flag, _profile_script_name,
#     _nvcc_build
#   - kernelking.harness.sections → _extract_kernel_section,
#     _extract_harness_section, _assemble_standalone
#   - kernelking.harness.generate → step_generate_harness + preprocess /
#     launch-info / struct-info / extern-C / strip / brace / classify /
#     error-context helpers
#   - kernelking.harness.timing   → KK_TIMING parsers, _run_standalone_timing,
#     _merge_harness_timing_into_metrics, step_profile_standalone
# Re-imported here so the rest of agent.py keeps working with bare names.
# ─────────────────────────────────────────────
from kernelking.harness import (  # noqa: E402
    STANDALONE_SUBDIR,
    _assemble_standalone,
    _check_brace_balance,
    _classify_compile_error,
    _extract_error_context,
    _extract_harness_section,
    _extract_kernel_launch_info,
    _extract_kernel_section,
    _extract_kernel_struct_info,
    _kernel_include_flag,
    _KK_TIMING_RE,
    _merge_harness_timing_into_metrics,
    _nvcc_arch,
    _nvcc_build,
    _nvcc_path,
    _parse_kk_timing_lines,
    _preprocess_for_standalone,
    _profile_script_name,
    _run_standalone_timing,
    _standalone_dir,
    _strip_extern_c_redeclarations,
    _strip_leading_prose,
    _strip_leading_prose_rust,
    step_generate_harness,
    step_profile_standalone,
)
# ─────────────────────────────────────────────
# ncu metric parsing + skill-grouped metric extraction
# ─────────────────────────────────────────────
# Lifted into `kernelking.ncu.metrics` — everything below is re-exported so
# the rest of `agent.py` can keep referencing these names unchanged.
from kernelking.ncu.metrics import (  # noqa: E402
    _FRIENDLY_UNIT_RE,
    _LAUNCH_HEADER_RE,
    _STALL_PREFIX,
    _STALL_SUFFIX,
    _TIMING_UNIT_RE,
    SKILL_METRICS,
    WORKLOAD_METRICS,
    _all_skill_metrics,
    _extract_per_launch_workload_csv,
    _extract_skill_csv,
    _parse_per_launch_data,
    _parse_per_launch_workload_csv,
    _parse_skill_csv,
    _parse_timing_only_metrics,
    _percentile,
    _to_us,
    parse_metrics,
)

# ─────────────────────────────────────────────
# Heuristic rule matrix + skill-grouped summary (moved to kernelking.ncu.rules)
# ─────────────────────────────────────────────
from kernelking.ncu.rules import (  # noqa: E402
    RULES,
    _apply_rules,
    _build_skill_summary,
    _collect_skill_summary,
    _dram_pct,
    _format_skill_summary_markdown,
    _rule_high_waves_per_sm,
    _rule_huge_dram_traffic,
    _rule_low_ipc,
    _rule_low_l1_hit,
    _rule_low_l2_hit,
    _rule_low_occupancy,
    _rule_low_waves_per_sm,
    _rule_near_peak_compute,
    _rule_poor_coalescing,
    _rule_register_limited,
    _rule_severe_memory_bound,
    _rule_smem_bank_conflicts,
    _rule_stall_barrier,
    _rule_stall_long_scoreboard,
    _rule_stall_math_pipe,
    _rule_stall_membar,
    _rule_stall_mio_throttle,
    _rule_stall_short_scoreboard,
    _rule_stall_wait,
    _sm_pct,
    _stall,
)




# ─────────────────────────────────────────────
# Real-env launch bucket analyzer (moved to kernelking.ncu.configs)
# ─────────────────────────────────────────────
from kernelking.ncu.configs import (  # noqa: E402
    _bucket_observed_snapshots,
    _find_representative_configs,
    _pick_representative_snapshot,
)




# ─────────────────────────────────────────────
# Step 4: Apply optimization, build & auto-debug
# ─────────────────────────────────────────────
# IMPORTANT: load config.env into os.environ BEFORE reading the constants below.
# Without this, module-level `os.environ.get(...)` calls run before main() had
# a chance to load config, and config.env overrides would be silently ignored.
load_config()

# Compile-time fix loop (aider-style: run verifier → feed stderr back to the model).
MAX_DEBUG_ATTEMPTS = int(os.environ.get("MAX_COMPILE_DEBUG_ATTEMPTS", "8"))
# Profiling/runtime failure fix loop (SWE-style: observe tool output → act again).
MAX_RUNTIME_DEBUG_ATTEMPTS = int(os.environ.get("MAX_RUNTIME_DEBUG_ATTEMPTS", "6"))
# Optional: after a successful cargo build, run this on the remote (non-interactive).
# Non-zero exit + combined output are fed to the same compile-fix path. Example:
#   REMOTE_VERIFY_CMD=true
#   REMOTE_VERIFY_CMD=cd /path && ./scripts/smoke.sh
MAX_VERIFY_DEBUG_ATTEMPTS = int(os.environ.get("MAX_VERIFY_DEBUG_ATTEMPTS", str(MAX_DEBUG_ATTEMPTS)))
# Remote subdir (under REMOTE_BASEDIR) for all standalone harness files & reports
STANDALONE_SUBDIR = "kernelking_standalone"


# ─────────────────────────────────────────────────────────────
# Stage C/D refactor (2026-04-16): forbidden-diff + comparison + iter
# loop subsystem moved to kernelking/optimize/.  Re-imported below so
# call sites keep working with the bare names.
#
# Note on loop-control constants: ``MAX_OPTIMIZE_ITERATIONS``,
# ``NO_PROGRESS_PATIENCE`` and ``NO_GAIN_THRESHOLD_PCT`` live in
# kernelking.optimize.loop (read once from ``os.environ`` at import
# time). We re-export them here because a couple of top-level logs
# and the final-deploy gate still reference them.
# ─────────────────────────────────────────────────────────────
from kernelking.optimize import (  # noqa: E402
    CUDA_FORBIDDEN_DIFF_PATTERNS,
    MAX_OPTIMIZE_ITERATIONS,
    NO_GAIN_THRESHOLD_PCT,
    NO_PROGRESS_PATIENCE,
    _added_lines,
    _claude_fix_compile_error,
    _claude_fix_compile_error_standalone,
    _claude_fix_harness_runtime,
    _claude_fix_runtime_error,
    _claude_fix_verify_error,
    _claude_optimize_standalone_iter,
    _duration_us,
    _format_per_launch_for_prompt,
    _inspect_cuda_diff,
    _profile_standalone_with_harness_fix,
    _run_verify_cmd,
    _trim_text_block,
    _try_build,
    rollback,
    step_apply_and_rebuild,
    step_compare,
    step_integrate_and_build,
    step_normalize_baseline_binary,
    step_optimize_standalone,
)
# MAX_* retry caps live in both agent.py (above) and kernelking.optimize.apply.
# They read the same env vars and the module-level constants already agree,
# so no re-import is needed — the local copies remain usable.



# ─────────────────────────────────────────────────────────────
# NVIDIA Nsight Compute details-page helpers (moved to kernelking.ncu.recommendations)
# ─────────────────────────────────────────────────────────────
from kernelking.ncu.recommendations import (  # noqa: E402
    _NCU_DROP_SECTION_RE,
    _NCU_RULE_RE,
    _NCU_SECTION_RE,
    _cross_check_findings,
    _dedupe_ncu_recommendations,
    _extract_ncu_recommendations,
    _format_ncu_recommendations_markdown,
    _trim_ncu_details,
)
# ─────────────────────────────────────────────
# Stage E refactor (2026-04-16): observation subsystem moved to
# kernelking/observe/
#   - kernelking.observe.cache      → _observation_cache_path,
#     _observation_load_cache, _observation_save_cache
#   - kernelking.observe.paths      → _resolve_rust_adapter_paths,
#     _upload_text_to_remote
#   - kernelking.observe.instrument → _claude_patch_cpp_for_observation,
#     _claude_patch_rust_for_observation
#   - kernelking.observe.parse      → _KKLOG_LINE_RE, _parse_kklog_line,
#     _parse_observation_log
#   - kernelking.observe.step       → step_observe_real_scalars
# Re-imported here so the rest of agent.py keeps working with bare names.
# ─────────────────────────────────────────────
from kernelking.observe import (  # noqa: E402
    _KKLOG_LINE_RE,
    _claude_patch_cpp_for_observation,
    _claude_patch_rust_for_observation,
    _observation_cache_path,
    _observation_load_cache,
    _observation_save_cache,
    _parse_kklog_line,
    _parse_observation_log,
    _resolve_rust_adapter_paths,
    _upload_text_to_remote,
    step_observe_real_scalars,
)

# ─────────────────────────────────────────────
# Stage E refactor (2026-04-16): real-env profile subsystem moved to
# kernelking/profile/
#   - kernelking.profile.env     → BINARY_ENV (remote Brevis binary env vars)
#   - kernelking.profile.timing  → step_profile_timing_only (the live path)
#   - kernelking.profile._legacy → step_profile, step_profile_with_runtime_fix,
#     step_generate_profile_script, _claude_triage_crash,
#     _read_profile_script, _upload_profile_script,
#     _step_baseline_with_script_fix, step_claude_optimize
#
# Only the live symbols are re-exported here; legacy helpers must be
# imported explicitly via ``from kernelking.profile._legacy import …`` so
# accidental reuse in new code is a visible "do you really mean legacy?"
# signal rather than a silent dependency.
# ─────────────────────────────────────────────
from kernelking.profile import (  # noqa: E402
    BINARY_ENV,
    step_profile_timing_only,
)


# ─────────────────────────────────────────────
# Preflight
# ─────────────────────────────────────────────
def preflight():
    """Sanity checks before starting."""
    host = cfg("SSH_HOST")
    if not host:
        die("SSH_HOST not set in config.env")

    api_key = cfg("ANTHROPIC_API_KEY")
    if not api_key or "YOUR_KEY_HERE" in api_key:
        die(f"ANTHROPIC_API_KEY not set. Edit {CONFIG_FILE}")

    log(f"Remote server          : {host}", C.CYAN)
    log(f"Claude model           : {cfg('CLAUDE_MODEL', 'claude-sonnet-4-6')}", C.CYAN)
    log(f"Kernel path            : {cfg('REMOTE_KERNEL_CU')}", C.CYAN)
    local_path = cfg("LOCAL_CODEBASE_PATH", "").strip()
    if local_path:
        log(f"Local codebase         : {local_path}", C.CYAN)
    else:
        log(f"Local codebase         : (not set — Claude has less context)", C.DIM)
    log(f"Max compile-fix tries  : {MAX_DEBUG_ATTEMPTS}", C.CYAN)
    log(f"Max runtime-fix tries  : {MAX_RUNTIME_DEBUG_ATTEMPTS}", C.CYAN)
    log(f"Max optimize iters     : {MAX_OPTIMIZE_ITERATIONS}  "
        f"(early-stop after {NO_PROGRESS_PATIENCE} consecutive failures)", C.CYAN)
    log(f"Standalone GPU arch    : {_nvcc_arch()}", C.CYAN)
    log(f"nvcc path              : {_nvcc_path()}", C.CYAN)
    log(f"nvcc extra defines     : {cfg('HARNESS_NVCC_DEFINES', '(none)')}", C.CYAN)
    verify_cmd = cfg("REMOTE_VERIFY_CMD", "").strip()
    if verify_cmd:
        log(f"Verify cmd             : {verify_cmd[:80]}", C.CYAN)
        log(f"Max verify-fix tries   : {MAX_VERIFY_DEBUG_ATTEMPTS}", C.CYAN)
    else:
        log(f"Verify cmd             : (not set — skipped)", C.DIM)

    # Test SSH
    log("Testing SSH connection...", C.DIM)
    r = ssh_run("echo ok", check=False, capture=True)
    if r.returncode != 0 or "ok" not in r.stdout:
        die(f"SSH connection to {host} failed! Check your ssh config.")
    log("SSH connection OK.", C.GREEN)

    # Test API key
    log("Testing Claude API key...", C.DIM)
    try:
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)
        # Quick ping — small request
        resp = _messages_create_with_retry(
            client,
            model=cfg("CLAUDE_MODEL", "claude-sonnet-4-6"),
            max_tokens=16,
            messages=[{"role": "user", "content": "Say OK"}],
        )
        log("Claude API key OK.", C.GREEN)
    except Exception as e:
        die(f"Claude API test failed: {e}")

    # GPU must be idle before we start profiling anything
    _check_and_clean_gpu("preflight")

    # Capture build/runtime fingerprints once per run. Used for banner,
    # memory headers, and post-mortem reproducibility.
    global _PROVENANCE
    _PROVENANCE = _collect_provenance()
    log("Run provenance:", C.CYAN)
    for line in _format_provenance_banner(_PROVENANCE).splitlines():
        log(f"  {line}", C.DIM)


# ─────────────────────────────────────────────
# Main — the agent loop
# ─────────────────────────────────────────────
def main():
    load_config()
    banner("🤖 Kernelking — CUDA Kernel Optimization Agent")
    log(f"Target kernel : {cfg('KERNEL_NAME', 'transpose_kernel')}", C.MAG)
    preflight()

    # Rebuild the cross-kernel pattern library from memory/*.md. Pure
    # projection — no network, no GPU. Emits CUDA_OPTIMIZATION.md at the
    # repo root for human inspection; the same data is injected on-demand
    # into step_claude_optimize() prompts for matching tags.
    try:
        _buckets = _build_cross_kernel_patterns()
        _n_tags = len([t for t, b in _buckets.items()
                       if b["deployed"] or b["rolled_back"]])
        if _n_tags:
            log(f"Cross-kernel patterns: {_n_tags} tag(s) with prior "
                f"evidence → CUDA_OPTIMIZATION.md", C.DIM)
    except Exception as _e:
        log(f"[patterns] rebuild skipped: {_e}", C.DIM)

    # ── Step 1: Read kernel from remote + local codebase context ──────────────
    kernel_code, local_ctx = step_read_kernel()

    # Per-run provenance header in the kernel's memory file. Placed here
    # (not in preflight) so the correct resolved KERNEL_NAME is used.
    _append_memory(
        f"\n## Run {_PROVENANCE.get('timestamp', '?')}\n"
        f"- Model: `{_PROVENANCE.get('claude_model', '?')}`  "
        f"SDK: `{_PROVENANCE.get('anthropic_sdk', '?')}`  "
        f"Agent SHA: `{_PROVENANCE.get('agent_sha', '?')}`\n"
        f"- GPU: {_PROVENANCE.get('gpu_name', '?')}  "
        f"Driver: {_PROVENANCE.get('driver_version', '?')}  "
        f"nvcc: {_PROVENANCE.get('nvcc_version', '?')}\n"
        f"- Target: `{_PROVENANCE.get('target_repo_path', '?')}`@"
        f"`{_PROVENANCE.get('target_repo_sha', '?')}`\n"
    )

    # ── Step 1.5: Observe real scalar values by instrumentation (cached) ──────
    # One-off cost (~4-5 min) to patch the kernel wrapper + Rust adapter with
    # diagnostic fprintf/eprintln lines, run the real prover once, then restore
    # + rebuild. Result is cached at refs/<kernel>.observed.json — subsequent
    # runs load the cache and skip this step. Produces ground-truth scalar
    # values per launch, which the harness generator can copy verbatim instead
    # of reverse-engineering scalar args from the ncu workload budget.
    observed = step_observe_real_scalars(kernel_code)

    # ── Step 1.9: Normalize remote binary to current .cu before baseline ──────
    # Without this, a stale `single-node` binary from a prior deploy can leak
    # into the baseline measurement. Observed failure (2026-04-19): source was
    # reverted to unoptimized, but the binary was still from a previous
    # optimized deploy → baseline measured at ~139 ms instead of the true
    # ~650 ms, and a legitimate 78% win got rolled back as "+0.5% degradation".
    #
    # Cost model:
    #   - normal rerun, no source changes → cargo incremental ~1-2 s (free)
    #   - observe just rebuilt (uncached path) → incremental ~1 s (no-op)
    #   - actual state mismatch (the bug case) → ~3 min but now baseline is
    #     trustworthy
    step_normalize_baseline_binary()

    # ── Step 2: Real-env baseline timing (drives harness config choice) ───────
    # We run this BEFORE the standalone harness because the ncu report tells us
    # which (grid, block) configurations actually dominate production time. The
    # standalone harness then tests exactly those configs — so standalone perf
    # becomes a faithful proxy for real-env perf, not a synthetic micro-bench.
    _check_and_clean_gpu("before Step 2 real-env baseline timing")
    baseline_real, _berr = step_profile_timing_only(
        "report_baseline_real", "Original kernel", observed=observed,
    )
    representative_configs: list[dict] = []
    if baseline_real is None:
        log("⚠  Real-env baseline timing failed — harness will fall back to n=1024 only.", C.YELLOW)
    else:
        representative_configs = baseline_real["metrics"].get("representative_configs", []) or []
        if not representative_configs:
            log("⚠  No representative configs derived from real-env — falling back to n=1024.", C.YELLOW)

    # ── Step 3: Generate standalone.cu (kernel + multi-config test harness) ───
    # Harness runs each representative (grid, block) from Step 2 so standalone
    # timings align with what the kernel actually does in the full proof.
    standalone_code = step_generate_harness(kernel_code, local_ctx, representative_configs)

    # ── Step 4: nvcc compile + standalone baseline profiling ──────────────────
    # Fast compile (~5-10s). Fix any harness generation errors before profiling.
    success, compile_log = _nvcc_build()
    for attempt in range(1, MAX_DEBUG_ATTEMPTS + 1):
        if success:
            break
        log(f"Harness compile fix {attempt}/{MAX_DEBUG_ATTEMPTS}...", C.YELLOW)
        standalone_code = _claude_fix_compile_error_standalone(
            standalone_code, compile_log, attempt, local_ctx
        )
        success, compile_log = _nvcc_build(standalone_code)
    if not success:
        die("Could not compile standalone harness — aborting")

    baseline_standalone, standalone_code = _profile_standalone_with_harness_fix(
        standalone_code, "baseline", local_ctx
    )

    # ── Step 5: Fast optimization loop (Claude → nvcc → ncu → repeat) ─────────
    # No cargo build needed; each iteration is seconds not minutes.
    best_code, best_report = step_optimize_standalone(standalone_code, baseline_standalone, local_ctx)

    # ── Short-circuit: if standalone gained nothing, don't waste ~6 min on
    # cargo build + real-env ncu. The Rust-integrated kernel would be
    # byte-identical to baseline, so real-env timing would just re-confirm
    # what we already know.
    baseline_dur = _duration_us(baseline_standalone["metrics"])
    best_dur     = _duration_us(best_report["metrics"])
    gain_pct     = (1 - best_dur / baseline_dur) * 100 if baseline_dur > 0 else 0.0
    # Uses the module-level NO_GAIN_THRESHOLD_PCT constant (default 0.5%) so
    # the short-circuit gate matches the iteration accept/reject threshold.
    if gain_pct < NO_GAIN_THRESHOLD_PCT:
        banner("⏭  Skipping Rust Integration — No Standalone Gain")
        log(f"Standalone best {best_dur:.1f} µs vs baseline {baseline_dur:.1f} µs  "
            f"→ {gain_pct:+.2f}% (below {NO_GAIN_THRESHOLD_PCT}% threshold).", C.YELLOW)
        log("No point spending ~3 min on cargo build + ~2:30 on real-env ncu for "
            "an unchanged kernel.", C.DIM)
        if baseline_real is not None:
            base_us = baseline_real["metrics"]["Total_us"]
            launches = baseline_real["metrics"].get("Count", "?")
            log(f"Real-env baseline was: {base_us:.1f} µs "
                f"({base_us/1000:.2f} ms) over {launches} launches.", C.DIM)
        log("Nothing deployed. Server kernel unchanged.", C.DIM)
        short_circuit_memo = (
            f"\n### Final outcome — SHORT-CIRCUIT (no deploy)\n"
            f"- Standalone gain {gain_pct:+.2f}% below {NO_GAIN_THRESHOLD_PCT}% noise floor.\n"
        )
        if baseline_real is not None:
            short_circuit_memo += (
                f"- Real-env baseline was "
                f"{(baseline_real['metrics']['Total_us']/1000):.2f} ms over "
                f"{baseline_real['metrics'].get('Count', '?')} launches.\n"
            )
        short_circuit_memo += f"- Cost: {_format_cost_summary()}\n"
        _append_memory(short_circuit_memo)
        rundir.refresh_agent_report()
        rundir.end_run()
        log(f"Claude usage this run: {_format_cost_summary()}", C.CYAN)
        sys.exit(2)

    # ── Step 6: Integrate → Rust project + cargo build ────────────────────────
    # Claude strips main(), restores FFI compatibility, then we run cargo build.
    final_code = step_integrate_and_build(best_code, local_ctx)

    # ── Step 7: Optimized real-env timing & final 2×2 comparison ──────────────
    # Same lightweight ncu approach — binary already built above; no extra build.
    kname = cfg("KERNEL_NAME", "kernel")
    optimized_real = None
    if baseline_real is not None:
        _check_and_clean_gpu("before Step 7 optimized timing")
        optimized_real, _oerr = step_profile_timing_only("report_optimized_real", "Optimized kernel")
        if optimized_real is None:
            log("⚠  Real-env optimized timing failed — will decide from standalone only.", C.YELLOW)

    improved = step_compare(baseline_standalone, best_report, baseline_real, optimized_real)

    # ── Memory: write final outcome (includes real-env before/after) ────────
    def _real_ms(report: dict | None) -> str:
        if report is None:
            return "—"
        us = report["metrics"].get("Total_us") or _duration_us(report["metrics"])
        return f"{us/1000:.2f} ms" if us else "—"

    if improved:
        banner("🎉 Optimization Accepted!")
        log("Optimized kernel is live on the server.", C.GREEN)
        log(f"Standalone best  : {ROOT / 'standalone_best.cu'}", C.DIM)
        log(f"Integrated kernel: {ROOT / f'{kname}_integrated.cu'}", C.DIM)
        _append_memory(
            f"\n### Final outcome — DEPLOYED\n"
            f"- Standalone: {baseline_dur:.1f} → {best_dur:.1f} µs ({gain_pct:+.2f}%)\n"
            f"- Real-env : {_real_ms(baseline_real)} → {_real_ms(optimized_real)}\n"
            f"- Integrated kernel file: `{kname}_integrated.cu`\n"
            f"- Cost: {_format_cost_summary()}\n"
        )
        rundir.refresh_agent_report()
        rundir.end_run()
        log(f"Claude usage this run: {_format_cost_summary()}", C.CYAN)
        sys.exit(0)
    else:
        banner("⏪ Rolling Back")
        rollback()
        log("Server restored to original kernel.", C.YELLOW)
        _append_memory(
            f"\n### Final outcome — ROLLED BACK (real-env regression or test failure)\n"
            f"- Standalone: {baseline_dur:.1f} → {best_dur:.1f} µs ({gain_pct:+.2f}%)\n"
            f"- Real-env : {_real_ms(baseline_real)} → {_real_ms(optimized_real)}\n"
            f"- Server restored to original.\n"
            f"- Cost: {_format_cost_summary()}\n"
        )
        rundir.refresh_agent_report()
        rundir.end_run()
        log(f"Claude usage this run: {_format_cost_summary()}", C.CYAN)
        sys.exit(1)


if __name__ == "__main__":
    main()