launchbound_bench/
plan.rs

1//! Bench plans (`plan.v1`): everything the runner needs, self-contained —
2//! candidate configs, PTX paths, launch geometry, and the argument layout
3//! matching cuda-oxide's PTX parameter lowering (each slice becomes a
4//! `ptr, len` pair of `.param` slots; scalars are single slots).
5//!
6//! A kernel's `[bench]` section in kernel.toml declares the workload;
7//! sizes and grid shapes are arithmetic expressions over the kernel's
8//! dimensions plus the `elements` variable.
9
10use launchbound_space::{Config, KernelSpec, eval_arith_expr};
11use serde::{Deserialize, Serialize};
12use std::collections::BTreeMap;
13use std::path::Path;
14
15#[derive(Debug, thiserror::Error)]
16pub enum PlanError {
17    #[error("kernel.toml [bench]: {0}")]
18    Spec(String),
19    #[error(transparent)]
20    Space(#[from] launchbound_space::SpaceError),
21    #[error("plan io: {0}")]
22    Io(String),
23}
24
25/// One PTX kernel argument slot group, in PTX parameter order.
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27#[serde(tag = "kind", rename_all = "snake_case")]
28pub enum ArgSpec {
29    /// Device buffer of f32, copied in (seeded deterministic init).
30    /// One `.param` slot (the pointer).
31    InF32 { len: u64 },
32    /// Device buffer of u32, values `seed % modulo` (histogram bins etc.).
33    InU32 { len: u64, modulo: u64 },
34    /// Device buffer of f32, zero-filled output. One slot.
35    OutF32 { len: u64 },
36    /// Device buffer of u32, zero-filled output. One slot.
37    OutU32 { len: u64 },
38    /// The length of the buffer at `of` (0-based ArgSpec index), as u64.
39    LenOf { of: usize },
40    /// Scalar u32.
41    U32 { value: u64 },
42    /// Scalar u64.
43    U64 { value: u64 },
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct Candidate {
48    pub id: String,
49    pub config: String,
50    /// PTX file path, relative to the plan file.
51    pub ptx: String,
52    /// True for a gate-refused candidate measured under --allow-unsafe:
53    /// the runner guards it with a watchdog, because the refusal means it
54    /// may genuinely hang.
55    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
56    pub unsafe_candidate: bool,
57    pub grid: [u32; 3],
58    pub block: [u32; 3],
59    pub args: Vec<ArgSpec>,
60    pub warmup: u32,
61    pub repeats: u32,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct BenchPlan {
66    pub schema: String,
67    pub kernel: String,
68    pub entry: String,
69    pub cc: String,
70    /// Present iff the plan includes gate-refused candidates: the explicit,
71    /// recorded reason the operator gave to --allow-unsafe (see the README). Never a
72    /// default.
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub allow_unsafe_reason: Option<String>,
75    pub candidates: Vec<Candidate>,
76}
77
78/// The `[bench]` section of kernel.toml.
79#[derive(Debug, Deserialize)]
80struct RawBench {
81    elements: u64,
82    #[serde(default = "default_one")]
83    grid_x: toml::Value,
84    #[serde(default = "default_one")]
85    grid_y: toml::Value,
86    #[serde(default = "default_one")]
87    grid_z: toml::Value,
88    #[serde(default = "default_warmup")]
89    warmup: u32,
90    #[serde(default = "default_repeats")]
91    repeats: u32,
92    args: Vec<RawArg>,
93}
94
95fn default_one() -> toml::Value {
96    toml::Value::Integer(1)
97}
98fn default_warmup() -> u32 {
99    20
100}
101fn default_repeats() -> u32 {
102    100
103}
104
105#[derive(Debug, Deserialize)]
106struct RawArg {
107    kind: String,
108    #[serde(default)]
109    len: Option<toml::Value>,
110    #[serde(default)]
111    of: Option<usize>,
112    #[serde(default)]
113    value: Option<toml::Value>,
114    #[serde(default)]
115    modulo: Option<toml::Value>,
116}
117
118#[derive(Debug)]
119pub struct BenchSpec {
120    raw: RawBench,
121}
122
123impl BenchSpec {
124    /// Load the `[bench]` section from the kernel's kernel.toml.
125    pub fn load(spec: &KernelSpec) -> Result<Self, PlanError> {
126        let path = spec.dir.join("kernel.toml");
127        let text =
128            std::fs::read_to_string(&path).map_err(|e| PlanError::Io(format!("{path:?}: {e}")))?;
129        let table: toml::Value =
130            toml::from_str(&text).map_err(|e| PlanError::Spec(e.to_string()))?;
131        let bench = table
132            .get("bench")
133            .ok_or_else(|| PlanError::Spec("kernel.toml has no [bench] section".into()))?;
134        let raw: RawBench = bench
135            .clone()
136            .try_into()
137            .map_err(|e| PlanError::Spec(format!("{e}")))?;
138        Ok(BenchSpec { raw })
139    }
140
141    /// Render one candidate: evaluate every expression against the config.
142    pub fn candidate(
143        &self,
144        _spec: &KernelSpec,
145        config: &Config,
146        ptx_relative: &str,
147    ) -> Result<Candidate, PlanError> {
148        let mut extra = BTreeMap::new();
149        extra.insert("elements".to_string(), self.raw.elements);
150        let eval = |v: &toml::Value, what: &str| -> Result<u64, PlanError> {
151            match v {
152                toml::Value::Integer(n) if *n >= 0 => Ok(*n as u64),
153                toml::Value::String(expr) => Ok(eval_arith_expr(expr, config, &extra)?),
154                other => Err(PlanError::Spec(format!(
155                    "{what} must be a non-negative integer or expression string, got {other}"
156                ))),
157            }
158        };
159
160        let grid = [
161            eval(&self.raw.grid_x, "grid_x")? as u32,
162            eval(&self.raw.grid_y, "grid_y")? as u32,
163            eval(&self.raw.grid_z, "grid_z")? as u32,
164        ];
165        let block_dim = |name: &str| -> u32 {
166            match config.get(name) {
167                Some(launchbound_space::Value::Int(n)) => *n as u32,
168                _ => 1,
169            }
170        };
171        let block = [
172            block_dim("block_x"),
173            block_dim("block_y"),
174            block_dim("block_z"),
175        ];
176
177        let mut args = Vec::with_capacity(self.raw.args.len());
178        for (i, raw) in self.raw.args.iter().enumerate() {
179            let need = |v: &Option<toml::Value>, field: &str| -> Result<u64, PlanError> {
180                let v = v.as_ref().ok_or_else(|| {
181                    PlanError::Spec(format!("args[{i}] kind {} needs `{field}`", raw.kind))
182                })?;
183                eval(v, field)
184            };
185            let arg = match raw.kind.as_str() {
186                "in_f32" => ArgSpec::InF32 {
187                    len: need(&raw.len, "len")?,
188                },
189                "in_u32" => ArgSpec::InU32 {
190                    len: need(&raw.len, "len")?,
191                    modulo: need(&raw.modulo, "modulo")?,
192                },
193                "out_f32" => ArgSpec::OutF32 {
194                    len: need(&raw.len, "len")?,
195                },
196                "out_u32" => ArgSpec::OutU32 {
197                    len: need(&raw.len, "len")?,
198                },
199                "len_of" => ArgSpec::LenOf {
200                    of: raw
201                        .of
202                        .ok_or_else(|| PlanError::Spec(format!("args[{i}] len_of needs `of`")))?,
203                },
204                "u32" => ArgSpec::U32 {
205                    value: need(&raw.value, "value")?,
206                },
207                "u64" => ArgSpec::U64 {
208                    value: need(&raw.value, "value")?,
209                },
210                other => {
211                    return Err(PlanError::Spec(format!(
212                        "args[{i}]: unknown kind {other:?}"
213                    )));
214                }
215            };
216            args.push(arg);
217        }
218
219        Ok(Candidate {
220            id: config.id().as_str().to_string(),
221            config: config.to_string(),
222            ptx: ptx_relative.to_string(),
223            unsafe_candidate: false,
224            grid,
225            block,
226            args,
227            warmup: self.raw.warmup,
228            repeats: self.raw.repeats,
229        })
230    }
231}
232
233impl BenchPlan {
234    pub fn write(&self, path: &Path) -> Result<(), PlanError> {
235        let json = serde_json::to_string_pretty(self).expect("plan serializes");
236        std::fs::write(path, json).map_err(|e| PlanError::Io(format!("{path:?}: {e}")))
237    }
238
239    pub fn load(path: &Path) -> Result<Self, PlanError> {
240        let text =
241            std::fs::read_to_string(path).map_err(|e| PlanError::Io(format!("{path:?}: {e}")))?;
242        let plan: BenchPlan =
243            serde_json::from_str(&text).map_err(|e| PlanError::Io(e.to_string()))?;
244        if plan.schema != "plan.v1" {
245            return Err(PlanError::Spec(format!(
246                "unsupported plan schema {:?}",
247                plan.schema
248            )));
249        }
250        Ok(plan)
251    }
252
253    /// Number of `.param` slots this candidate's args expand to (each
254    /// ArgSpec is exactly one slot; slices appear as explicit ptr + len_of
255    /// pairs). Used to validate against the PTX entry signature.
256    pub fn param_slots(candidate: &Candidate) -> usize {
257        candidate.args.len()
258    }
259}
launchbound_bench/plan.rs

launchbound_bench/
plan.rs
