launchbound-bench 2.2.1

//! Plan execution: resumable, checkpointed after every candidate (the box
//! dies — idle guard, dead-man switch, spot reclaim — so the harness is
//! resumable or it is broken, docs/ARCHITECTURE.md), with a CPU heartbeat so a
//! GPU-bound sweep never looks idle to the 30-minute CPU alarm.

use crate::cuda::Device;
use crate::plan::{ArgSpec, BenchPlan, Candidate};
use crate::stats::{Summary, summarize};
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;

/// What measuring one candidate produced — including the ways it can fail.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CandidateResult {
    /// Its canonical `config.v1` ID, joining this to the plan and verdicts.
    pub id: String,
    /// Its dimension assignments.
    pub config: String,
    /// `ok`, `error`, or `timeout` — the last meaning a watchdog fired on
    /// a candidate the gate had refused, which is the refusal being right.
    pub status: String,
    /// Why it failed, when it did.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// Untimed launches performed before measuring.
    pub warmup: u32,
    /// Timed launches requested.
    pub repeats: u32,
    /// Every raw sample, in launch order. Kept so a summary can be
    /// recomputed without re-running the GPU.
    #[serde(default)]
    pub times_ms: Vec<f64>,
    /// The reduction of `times_ms`. `None` when nothing was measured.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub summary: Option<Summary>,
    /// Wall-clock seconds this candidate consumed on the GPU host.
    pub gpu_seconds: f64,
}

/// A `results.v1` document: every measurement, and the machine that took it.
///
/// Checkpointed after each candidate, so an interrupted sweep resumes from
/// what it already measured rather than starting over.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Results {
    /// Schema tag; always `results.v1`.
    pub schema: String,
    /// The kernel measured.
    pub kernel: String,
    /// The entry point launched.
    pub entry: String,
    /// The capability the plan was gated at.
    pub plan_cc: String,
    /// Product name of the device, as the driver reports it.
    pub device_name: String,
    /// The device's actual capability, which need not equal `plan_cc`.
    pub device_cc: String,
    /// Driver version. Part of what makes a timing reproducible, and part
    /// of why results do not port (`docs/LIMITATIONS.md`).
    pub driver_version: String,
    /// One entry per candidate visited, in visiting order.
    pub candidates: Vec<CandidateResult>,
    /// GPU seconds the whole sweep consumed.
    pub total_gpu_seconds: f64,
    /// Strategy that produced the visiting order (`exhaustive` | `random:<seed>`).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub strategy: Option<String>,
    /// True when the sweep stopped because the wall budget ran out.
    #[serde(default)]
    pub budget_exhausted: bool,
}

/// How to visit a plan's candidates, and when to stop.
pub struct RunOptions {
    /// Visiting order over plan.candidates indices (a permutation or
    /// prefix); defaults to plan order.
    pub order: Vec<usize>,
    /// Wall-clock budget; the sweep stops (resumably) when it is spent.
    pub budget_secs: Option<f64>,
    /// Recorded in results for provenance.
    pub strategy: Option<String>,
}

impl RunOptions {
    /// Visit every candidate in plan order, with no budget.
    pub fn exhaustive(plan: &BenchPlan) -> Self {
        RunOptions {
            order: (0..plan.candidates.len()).collect(),
            budget_secs: None,
            strategy: Some("exhaustive".into()),
        }
    }
}

/// Execute `plan`, appending to `results_path` (resume: candidates already
/// present are skipped). Writes the results file after every candidate.
pub fn run_plan(
    plan: &BenchPlan,
    plan_dir: &Path,
    results_path: &Path,
    options: &RunOptions,
    progress: &mut dyn FnMut(&str),
) -> Result<Results, String> {
    let device = Device::open()?;
    progress(&format!(
        "device: {} (cc {}, driver {})",
        device.name, device.cc, device.driver_version
    ));
    if device.cc != plan.cc {
        progress(&format!(
            "WARNING: plan was gated at cc {}, device is cc {} — verdicts do not transfer \
             across parts (docs/SAFETY.md); results will be labelled with the device cc",
            plan.cc, device.cc
        ));
    }

    // A checkpoint that cannot be read is not "start over": resuming would
    // silently discard measurements that cost GPU time, and the file is the
    // only record of them. Say what is wrong and stop.
    let existing = Results::load(results_path)?;
    let mut results = match existing {
        Some(existing) => {
            progress(&format!(
                "resuming: {} candidates already measured",
                existing.candidates.len()
            ));
            existing
        }
        None => Results {
            schema: "results.v1".into(),
            kernel: plan.kernel.clone(),
            entry: plan.entry.clone(),
            plan_cc: plan.cc.clone(),
            device_name: device.name.clone(),
            device_cc: device.cc.clone(),
            driver_version: device.driver_version.clone(),
            candidates: Vec::new(),
            total_gpu_seconds: 0.0,
            strategy: options.strategy.clone(),
            budget_exhausted: false,
        },
    };
    results.budget_exhausted = false;

    let heartbeat_stop = start_heartbeat();
    let sweep_started = Instant::now();

    for &index in &options.order {
        let Some(candidate) = plan.candidates.get(index) else {
            return Err(format!("order index {index} out of range"));
        };
        if results.candidates.iter().any(|c| c.id == candidate.id) {
            continue;
        }
        if let Some(budget) = options.budget_secs
            && sweep_started.elapsed().as_secs_f64() >= budget
        {
            results.budget_exhausted = true;
            progress(&format!(
                "budget exhausted after {:.1}s: {} of {} candidates measured (resumable)",
                sweep_started.elapsed().as_secs_f64(),
                results.candidates.len(),
                plan.candidates.len()
            ));
            break;
        }
        // A gate-refused candidate may genuinely hang (that is why it was
        // refused). Checkpoint a `timeout` record BEFORE launching, so a
        // watchdog abort (or a wedged GPU context) leaves a resumable
        // truth on disk; overwrite it with the real outcome if we survive.
        let watchdog = if candidate.unsafe_candidate {
            results.candidates.push(CandidateResult {
                id: candidate.id.clone(),
                config: candidate.config.clone(),
                status: "timeout".into(),
                error: Some(format!(
                    "unsafe candidate did not complete within {UNSAFE_TIMEOUT_SECS}s;                      presumed hung (this is the failure mode the gate predicts)"
                )),
                warmup: candidate.warmup,
                repeats: candidate.repeats,
                times_ms: Vec::new(),
                summary: None,
                gpu_seconds: unsafe_timeout_secs() as f64,
            });
            results.checkpoint(results_path)?;
            progress(&format!(
                "{} UNSAFE candidate: watchdog armed at {}s",
                candidate.id,
                unsafe_timeout_secs()
            ));
            Some(arm_watchdog(unsafe_timeout_secs()))
        } else {
            None
        };
        let started = Instant::now();
        let outcome = run_candidate(&device, plan, plan_dir, candidate);
        if let Some(armed) = watchdog {
            armed.store(true, Ordering::Relaxed); // disarm
            // Replace the pre-checkpointed timeout record with the truth.
            results.candidates.retain(|c| c.id != candidate.id);
        }
        let gpu_seconds = started.elapsed().as_secs_f64();
        let result = match outcome {
            Ok(times_ms) => {
                let summary = summarize(&times_ms);
                progress(&format!(
                    "{} {}: median {} over {} repeats ({:.1}s)",
                    candidate.id,
                    candidate.config,
                    summary
                        .as_ref()
                        .map(|s| format!(
                            "{:.4} ms [{:.4}, {:.4}]",
                            s.median_ms, s.ci95_lo_ms, s.ci95_hi_ms
                        ))
                        .unwrap_or_else(|| "n/a".into()),
                    candidate.repeats,
                    gpu_seconds,
                ));
                CandidateResult {
                    id: candidate.id.clone(),
                    config: candidate.config.clone(),
                    status: "ok".into(),
                    error: None,
                    warmup: candidate.warmup,
                    repeats: candidate.repeats,
                    times_ms,
                    summary,
                    gpu_seconds,
                }
            }
            Err(e) => {
                progress(&format!("{} ERROR: {e}", candidate.id));
                CandidateResult {
                    id: candidate.id.clone(),
                    config: candidate.config.clone(),
                    status: "error".into(),
                    error: Some(e),
                    warmup: candidate.warmup,
                    repeats: candidate.repeats,
                    times_ms: Vec::new(),
                    summary: None,
                    gpu_seconds,
                }
            }
        };
        results.candidates.push(result);
        results.total_gpu_seconds = sweep_started.elapsed().as_secs_f64();
        results.checkpoint(results_path)?;
    }

    heartbeat_stop.store(true, Ordering::Relaxed);
    results.total_gpu_seconds = sweep_started.elapsed().as_secs_f64();
    results.checkpoint(results_path)?;
    Ok(results)
}

fn run_candidate(
    device: &Device,
    plan: &BenchPlan,
    plan_dir: &Path,
    candidate: &Candidate,
) -> Result<Vec<f64>, String> {
    let ptx_path = plan_dir.join(&candidate.ptx);
    let ptx = std::fs::read_to_string(&ptx_path)
        .map_err(|e| format!("reading {}: {e}", ptx_path.display()))?;
    let module = device.load_module(&ptx, &plan.entry)?;

    // Materialize buffers and the param pointer table, in ArgSpec order.
    // Each ArgSpec is exactly one .param slot.
    let mut buffers = Vec::new(); // (arg index, Buffer)
    for (i, arg) in candidate.args.iter().enumerate() {
        match arg {
            ArgSpec::InF32 { len } => {
                let host: Vec<f32> = deterministic_f32(*len);
                let buf = device.alloc(host.len() * 4)?;
                device.copy_in(&buf, cast_bytes(&host))?;
                buffers.push((i, buf));
            }
            ArgSpec::InU32 { len, modulo } => {
                let host: Vec<u32> = deterministic_u32(*len, *modulo);
                let buf = device.alloc(host.len() * 4)?;
                device.copy_in(&buf, cast_bytes(&host))?;
                buffers.push((i, buf));
            }
            ArgSpec::OutF32 { len } | ArgSpec::OutU32 { len } => {
                let zero = vec![0u8; (*len as usize) * 4];
                let buf = device.alloc(zero.len())?;
                device.copy_in(&buf, &zero)?;
                buffers.push((i, buf));
            }
            ArgSpec::LenOf { .. } | ArgSpec::U32 { .. } | ArgSpec::U64 { .. } => {}
        }
    }

    // Scalar storage must outlive the launch; the params table points into
    // these vectors and the buffers' device pointers.
    let mut ptr_slots: Vec<u64> = Vec::new();
    let mut u32_slots: Vec<u32> = Vec::new();
    let mut u64_slots: Vec<u64> = Vec::new();
    #[derive(Clone, Copy)]
    enum Slot {
        Ptr(usize),
        U32(usize),
        U64(usize),
    }
    let mut slots = Vec::with_capacity(candidate.args.len());
    for (i, arg) in candidate.args.iter().enumerate() {
        match arg {
            ArgSpec::InF32 { .. }
            | ArgSpec::InU32 { .. }
            | ArgSpec::OutF32 { .. }
            | ArgSpec::OutU32 { .. } => {
                // The loop above materializes a buffer for exactly these
                // four ArgSpec kinds, so the lookup succeeds — an invariant
                // held one loop away from the code that needs it. This
                // function returns a `Result`; a wrong answer here is a bug
                // report, not a crash mid-benchmark.
                let buf = &buffers
                    .iter()
                    .find(|(idx, _)| *idx == i)
                    .ok_or_else(|| format!("internal: argument {i} has no materialized buffer"))?
                    .1;
                ptr_slots.push(buf.ptr);
                slots.push(Slot::Ptr(ptr_slots.len() - 1));
            }
            ArgSpec::LenOf { of } => {
                let len = match candidate.args.get(*of) {
                    Some(ArgSpec::InF32 { len })
                    | Some(ArgSpec::OutF32 { len })
                    | Some(ArgSpec::OutU32 { len })
                    | Some(ArgSpec::InU32 { len, .. }) => *len,
                    other => return Err(format!("len_of {of} points at {other:?}")),
                };
                u64_slots.push(len);
                slots.push(Slot::U64(u64_slots.len() - 1));
            }
            ArgSpec::U32 { value } => {
                u32_slots.push(*value as u32);
                slots.push(Slot::U32(u32_slots.len() - 1));
            }
            ArgSpec::U64 { value } => {
                u64_slots.push(*value);
                slots.push(Slot::U64(u64_slots.len() - 1));
            }
        }
    }
    let mut params: Vec<*mut std::ffi::c_void> = slots
        .iter()
        .map(|slot| match slot {
            Slot::Ptr(k) => std::ptr::from_mut(&mut ptr_slots[*k]).cast(),
            Slot::U32(k) => std::ptr::from_mut(&mut u32_slots[*k]).cast(),
            Slot::U64(k) => std::ptr::from_mut(&mut u64_slots[*k]).cast(),
        })
        .collect();

    for _ in 0..candidate.warmup {
        device.timed_launch(&module, candidate.grid, candidate.block, &mut params)?;
    }
    device.synchronize()?;

    let mut times = Vec::with_capacity(candidate.repeats as usize);
    for _ in 0..candidate.repeats {
        times.push(device.timed_launch(&module, candidate.grid, candidate.block, &mut params)?);
    }
    device.synchronize()?;
    Ok(times)
}

fn cast_bytes<T>(data: &[T]) -> &[u8] {
    unsafe { std::slice::from_raw_parts(data.as_ptr().cast(), std::mem::size_of_val(data)) }
}

/// Deterministic xorshift-seeded data: reproducible across runs and hosts.
fn deterministic_f32(len: u64) -> Vec<f32> {
    let mut state = 0x9e3779b97f4a7c15u64;
    (0..len)
        .map(|_| {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            ((state >> 40) as f32) / ((1u64 << 24) as f32)
        })
        .collect()
}

fn deterministic_u32(len: u64, modulo: u64) -> Vec<u32> {
    let modulo = modulo.max(1);
    let mut state = 0x2545f4914f6cdd1du64;
    (0..len)
        .map(|_| {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            (state % modulo) as u32
        })
        .collect()
}

impl Results {
    /// Read `results.v1` from a run directory.
    ///
    /// `Ok(None)` means one thing only: **the file is not there**, so the
    /// measurement box has not run yet. Everything else is an error.
    ///
    /// This used to be `read_to_string(path).ok()?` then `from_str().ok()`,
    /// so a truncated file, an empty one, `null`, `[]`, a `results.v2` from
    /// a newer runner and a *directory* named `results.json` all collapsed
    /// into that same `None` — and `report` rendered "nothing measured yet",
    /// exit 0, nothing on stderr, with a JSON report that validated. The run
    /// directory is the hand-off between two machines, and the two
    /// conditions call for opposite actions: wait, or go and look. Nothing
    /// told them apart.
    ///
    /// `verdicts.v1` two lines away in the report builder already checked
    /// its schema tag by name; this is the same check, so a future
    /// `results.v2` is refused by name rather than read as nothing.
    ///
    /// # Errors
    ///
    /// Any I/O error that is not `NotFound`, any parse failure, and any
    /// document whose `schema` is not `results.v1` — each naming the path.
    pub fn load(path: &Path) -> Result<Option<Self>, String> {
        let text = match std::fs::read_to_string(path) {
            Ok(text) => text,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
            Err(e) => return Err(format!("{}: {e}", path.display())),
        };
        // The tag first, so a wrong-schema document is named as one rather
        // than as whatever field happens to be missing from it.
        let value: serde_json::Value = serde_json::from_str(&text)
            .map_err(|e| format!("{}: not JSON: {e}", path.display()))?;
        let declared = value.get("schema").and_then(|s| s.as_str());
        match declared {
            Some("results.v1") => {}
            Some(other) => {
                return Err(format!(
                    "{}: unsupported results schema `{other}` (expected `results.v1`)",
                    path.display()
                ));
            }
            None => {
                return Err(format!(
                    "{}: not a results.v1 document (no `schema` field)",
                    path.display()
                ));
            }
        }
        serde_json::from_str(&text)
            .map(Some)
            .map_err(|e| format!("{}: not a results.v1 document: {e}", path.display()))
    }

    /// Atomic checkpoint: write to a temp file, then rename.
    pub fn checkpoint(&self, path: &Path) -> Result<(), String> {
        let tmp = path.with_extension("json.tmp");
        let json =
            serde_json::to_string_pretty(self).map_err(|e| format!("serializing results: {e}"))?;
        std::fs::write(&tmp, json).map_err(|e| e.to_string())?;
        std::fs::rename(&tmp, path).map_err(|e| e.to_string())
    }
}

const UNSAFE_TIMEOUT_SECS: u64 = 10;

fn unsafe_timeout_secs() -> u64 {
    std::env::var("LAUNCHBOUND_UNSAFE_TIMEOUT_SECS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(UNSAFE_TIMEOUT_SECS)
}

/// Watchdog for unsafe candidates: if not disarmed within the deadline the
/// process exits (a hung kernel cannot be cancelled from user code). The
/// pre-checkpointed `timeout` record makes the rerun skip it.
fn arm_watchdog(deadline_secs: u64) -> std::sync::Arc<AtomicBool> {
    let disarmed = std::sync::Arc::new(AtomicBool::new(false));
    let flag = disarmed.clone();
    std::thread::spawn(move || {
        let start = Instant::now();
        while start.elapsed().as_secs() < deadline_secs {
            if flag.load(Ordering::Relaxed) {
                return;
            }
            std::thread::sleep(std::time::Duration::from_millis(100));
        }
        if !flag.load(Ordering::Relaxed) {
            eprintln!(
                "watchdog: unsafe candidate exceeded {deadline_secs}s; exiting so the                  checkpointed timeout record stands (exit 3, rerun to continue)"
            );
            std::process::exit(3);
        }
    });
    disarmed
}

/// A CPU heartbeat: the box's idle alarm terminates on CPU <5% for 30 min,
/// and a GPU-bound loop can look idle. Burn a configurable duty cycle on
/// one core (LAUNCHBOUND_HEARTBEAT_PCT, default 40) until stopped.
fn start_heartbeat() -> &'static AtomicBool {
    static STOP: AtomicBool = AtomicBool::new(false);
    STOP.store(false, Ordering::Relaxed);
    let duty: u64 = std::env::var("LAUNCHBOUND_HEARTBEAT_PCT")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(40)
        .clamp(1, 100);
    std::thread::spawn(move || {
        let mut sink = 0u64;
        while !STOP.load(Ordering::Relaxed) {
            let spin = Instant::now();
            while spin.elapsed().as_millis() < duty as u128 {
                sink = sink.wrapping_mul(6364136223846793005).wrapping_add(1);
            }
            std::hint::black_box(sink);
            std::thread::sleep(std::time::Duration::from_millis(100 - duty.min(99)));
        }
    });
    &STOP
}