launchbound-bench 2.2.0

//! Bench plans (`plan.v1`): everything the runner needs, self-contained —
//! candidate configs, PTX paths, launch geometry, and the argument layout
//! matching cuda-oxide's PTX parameter lowering (each slice becomes a
//! `ptr, len` pair of `.param` slots; scalars are single slots).
//!
//! A kernel's `[bench]` section in kernel.toml declares the workload;
//! sizes and grid shapes are arithmetic expressions over the kernel's
//! dimensions plus the `elements` variable.

use launchbound_space::{Config, KernelSpec, eval_arith_expr};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::Path;

/// What can go wrong turning a spec into a benchmark plan.
#[derive(Debug, thiserror::Error)]
pub enum PlanError {
    /// The `[bench]` section is missing, malformed, or names a dimension
    /// or buffer the spec does not declare.
    #[error("kernel.toml [bench]: {0}")]
    Spec(String),
    /// The spec or one of its constraints did not load.
    #[error(transparent)]
    Space(#[from] launchbound_space::SpaceError),
    /// The plan could not be read, written or deserialized.
    #[error("plan io: {0}")]
    Io(String),
}

/// One PTX kernel argument slot group, in PTX parameter order.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ArgSpec {
    /// Device buffer of f32, copied in (seeded deterministic init).
    /// One `.param` slot (the pointer).
    InF32 {
        /// Elements, not bytes.
        len: u64,
    },
    /// Device buffer of u32, values `seed % modulo` (histogram bins etc.).
    InU32 {
        /// Elements, not bytes.
        len: u64,
        /// Values are `seed % modulo`, so a histogram's inputs land in
        /// range without the plan carrying the data.
        modulo: u64,
    },
    /// Device buffer of f32, zero-filled output. One slot.
    OutF32 {
        /// Elements, not bytes.
        len: u64,
    },
    /// Device buffer of u32, zero-filled output. One slot.
    OutU32 {
        /// Elements, not bytes.
        len: u64,
    },
    /// The length of the buffer at `of` (0-based ArgSpec index), as u64.
    LenOf {
        /// 0-based index into the candidate's `args` of the buffer whose
        /// length this passes.
        of: usize,
    },
    /// Scalar u32.
    U32 {
        /// The value, held as `u64` and narrowed at launch.
        value: u64,
    },
    /// Scalar u64.
    U64 {
        /// The value.
        value: u64,
    },
}

/// One configuration, ready to launch: everything the runner needs and
/// nothing it has to recompute.
///
/// A plan is deliberately self-contained. The runner takes a plan and a
/// directory of PTX and needs neither the kernel source, nor cuda-oxide,
/// nor reconverge — which is what lets the measurement leg run on a GPU box
/// that has none of them installed.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Candidate {
    /// Its canonical `config.v1` ID, matching `verdicts.v1`.
    pub id: String,
    /// Its dimension assignments, e.g. `block_x=128 tile=256`.
    pub config: String,
    /// PTX file path, relative to the plan file.
    pub ptx: String,
    /// True for a gate-refused candidate measured under --allow-unsafe:
    /// the runner guards it with a watchdog, because the refusal means it
    /// may genuinely hang.
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub unsafe_candidate: bool,
    /// Grid dimensions `[x, y, z]`, resolved from the `[bench]`
    /// expressions for this configuration.
    pub grid: [u32; 3],
    /// Block dimensions `[x, y, z]` — the launch shape the gate judged.
    pub block: [u32; 3],
    /// Kernel arguments in PTX parameter order.
    pub args: Vec<ArgSpec>,
    /// Untimed launches before measurement, to settle clocks and caches.
    pub warmup: u32,
    /// Timed launches. Each contributes one sample to the summary.
    pub repeats: u32,
}

/// A `plan.v1` document: every candidate to measure, and how.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchPlan {
    /// Schema tag; always `plan.v1`.
    pub schema: String,
    /// The kernel being tuned.
    pub kernel: String,
    /// The `#[kernel]` entry point to launch, as named in the PTX.
    pub entry: String,
    /// The capability the gate ran at, carried so the runner can refuse a
    /// plan built for a different target than the device it holds.
    pub cc: String,
    /// Present iff the plan includes gate-refused candidates: the explicit,
    /// recorded reason the operator gave to --allow-unsafe (see the README). Never a
    /// default.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub allow_unsafe_reason: Option<String>,
    /// The candidates to measure, in plan order. A strategy visits them
    /// in its own order without reordering this list.
    pub candidates: Vec<Candidate>,
}

/// The `[bench]` section of kernel.toml.
#[derive(Debug, Deserialize)]
struct RawBench {
    elements: u64,
    #[serde(default = "default_one")]
    grid_x: toml::Value,
    #[serde(default = "default_one")]
    grid_y: toml::Value,
    #[serde(default = "default_one")]
    grid_z: toml::Value,
    #[serde(default = "default_warmup")]
    warmup: u32,
    #[serde(default = "default_repeats")]
    repeats: u32,
    args: Vec<RawArg>,
}

fn default_one() -> toml::Value {
    toml::Value::Integer(1)
}
fn default_warmup() -> u32 {
    20
}
fn default_repeats() -> u32 {
    100
}

#[derive(Debug, Deserialize)]
struct RawArg {
    kind: String,
    #[serde(default)]
    len: Option<toml::Value>,
    #[serde(default)]
    of: Option<usize>,
    #[serde(default)]
    value: Option<toml::Value>,
    #[serde(default)]
    modulo: Option<toml::Value>,
}

/// The `[bench]` section, parsed once and used to build every candidate.
///
/// It holds the grid expressions and argument declarations, which are the
/// same for the whole space; only the values they resolve against change
/// per configuration.
#[derive(Debug)]
pub struct BenchSpec {
    raw: RawBench,
}

impl BenchSpec {
    /// Load the `[bench]` section from the kernel's kernel.toml.
    pub fn load(spec: &KernelSpec) -> Result<Self, PlanError> {
        let path = spec.dir.join("kernel.toml");
        let text =
            std::fs::read_to_string(&path).map_err(|e| PlanError::Io(format!("{path:?}: {e}")))?;
        let table: toml::Value =
            toml::from_str(&text).map_err(|e| PlanError::Spec(e.to_string()))?;
        let bench = table
            .get("bench")
            .ok_or_else(|| PlanError::Spec("kernel.toml has no [bench] section".into()))?;
        let raw: RawBench = bench
            .clone()
            .try_into()
            .map_err(|e| PlanError::Spec(format!("{e}")))?;
        Ok(BenchSpec { raw })
    }

    /// Render one candidate: evaluate every expression against the config.
    pub fn candidate(
        &self,
        _spec: &KernelSpec,
        config: &Config,
        ptx_relative: &str,
    ) -> Result<Candidate, PlanError> {
        let mut extra = BTreeMap::new();
        extra.insert("elements".to_string(), self.raw.elements);
        let eval = |v: &toml::Value, what: &str| -> Result<u64, PlanError> {
            match v {
                toml::Value::Integer(n) if *n >= 0 => Ok(*n as u64),
                toml::Value::String(expr) => Ok(eval_arith_expr(expr, config, &extra)?),
                other => Err(PlanError::Spec(format!(
                    "{what} must be a non-negative integer or expression string, got {other}"
                ))),
            }
        };

        let grid = [
            eval(&self.raw.grid_x, "grid_x")? as u32,
            eval(&self.raw.grid_y, "grid_y")? as u32,
            eval(&self.raw.grid_z, "grid_z")? as u32,
        ];
        let block_dim = |name: &str| -> u32 {
            match config.get(name) {
                Some(launchbound_space::Value::Int(n)) => *n as u32,
                _ => 1,
            }
        };
        let block = [
            block_dim("block_x"),
            block_dim("block_y"),
            block_dim("block_z"),
        ];

        let mut args = Vec::with_capacity(self.raw.args.len());
        for (i, raw) in self.raw.args.iter().enumerate() {
            let need = |v: &Option<toml::Value>, field: &str| -> Result<u64, PlanError> {
                let v = v.as_ref().ok_or_else(|| {
                    PlanError::Spec(format!("args[{i}] kind {} needs `{field}`", raw.kind))
                })?;
                eval(v, field)
            };
            let arg = match raw.kind.as_str() {
                "in_f32" => ArgSpec::InF32 {
                    len: need(&raw.len, "len")?,
                },
                "in_u32" => ArgSpec::InU32 {
                    len: need(&raw.len, "len")?,
                    modulo: need(&raw.modulo, "modulo")?,
                },
                "out_f32" => ArgSpec::OutF32 {
                    len: need(&raw.len, "len")?,
                },
                "out_u32" => ArgSpec::OutU32 {
                    len: need(&raw.len, "len")?,
                },
                "len_of" => ArgSpec::LenOf {
                    of: raw
                        .of
                        .ok_or_else(|| PlanError::Spec(format!("args[{i}] len_of needs `of`")))?,
                },
                "u32" => ArgSpec::U32 {
                    value: need(&raw.value, "value")?,
                },
                "u64" => ArgSpec::U64 {
                    value: need(&raw.value, "value")?,
                },
                other => {
                    return Err(PlanError::Spec(format!(
                        "args[{i}]: unknown kind {other:?}"
                    )));
                }
            };
            args.push(arg);
        }

        Ok(Candidate {
            id: config.id().as_str().to_string(),
            config: config.to_string(),
            ptx: ptx_relative.to_string(),
            unsafe_candidate: false,
            grid,
            block,
            args,
            warmup: self.raw.warmup,
            repeats: self.raw.repeats,
        })
    }
}

impl BenchPlan {
    /// Write the plan as pretty JSON.
    pub fn write(&self, path: &Path) -> Result<(), PlanError> {
        let json = serde_json::to_string_pretty(self)
            .map_err(|e| PlanError::Io(format!("serializing plan: {e}")))?;
        std::fs::write(path, json).map_err(|e| PlanError::Io(format!("{path:?}: {e}")))
    }

    /// Read a plan, refusing any schema but `plan.v1`.
    ///
    /// The schema check is not a formality: the runner launches whatever a
    /// plan tells it to, on a real device, so a document it half-understands
    /// is worse than one it rejects.
    pub fn load(path: &Path) -> Result<Self, PlanError> {
        let text =
            std::fs::read_to_string(path).map_err(|e| PlanError::Io(format!("{path:?}: {e}")))?;
        let plan: BenchPlan =
            serde_json::from_str(&text).map_err(|e| PlanError::Io(e.to_string()))?;
        if plan.schema != "plan.v1" {
            return Err(PlanError::Spec(format!(
                "unsupported plan schema {:?}",
                plan.schema
            )));
        }
        Ok(plan)
    }

    /// Number of `.param` slots this candidate's args expand to (each
    /// ArgSpec is exactly one slot; slices appear as explicit ptr + len_of
    /// pairs). Used to validate against the PTX entry signature.
    pub fn param_slots(candidate: &Candidate) -> usize {
        candidate.args.len()
    }
}