Skip to main content

launchbound_space/
spec.rs

1//! Kernel spec: the `kernel.toml` sitting next to a corpus kernel's source.
2
3use crate::SpaceError;
4use crate::constraint::Constraint;
5use serde::Deserialize;
6use std::fmt;
7use std::path::{Path, PathBuf};
8
9/// CUDA's per-axis limits on a thread block, from the Programming Guide's
10/// compute-capability table. They have been these numbers since 2.x and are
11/// not capability-dependent, unlike the shared-memory and occupancy figures
12/// that `launchbound-model` keeps per device.
13const BLOCK_AXIS_LIMITS: [(&str, u64); 3] = [("block_x", 1024), ("block_y", 1024), ("block_z", 64)];
14
15/// One value a dimension can take.
16#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
17pub enum Value {
18    /// A non-negative integer, e.g. a `block_x` or `tile` size.
19    Int(u64),
20    /// A string, for dimensions whose values are names rather than sizes.
21    Str(String),
22}
23
24impl fmt::Display for Value {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        match self {
27            Value::Int(n) => write!(f, "{n}"),
28            Value::Str(s) => f.write_str(s),
29        }
30    }
31}
32
33/// What a dimension controls.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum DimRole {
36    /// Launch geometry only: no source change, no recompile. `block_x`,
37    /// `block_y`, `block_z` default to this.
38    Launch,
39    /// Compile-time specialization: a distinct generated source, therefore a
40    /// distinct reconverge verdict and a distinct compiled artifact.
41    Spec,
42}
43
44/// One tunable dimension.
45#[derive(Debug, Clone)]
46pub struct Dim {
47    /// Dimension name, matching the `params.rs` constant it rewrites and
48    /// the identifier constraints refer to. `[a-z0-9_]` only.
49    pub name: String,
50    /// Whether this changes the launch shape or the compiled source.
51    pub role: DimRole,
52    /// The values to enumerate, in declaration order, deduplicated at load.
53    pub values: Vec<Value>,
54}
55
56/// What the corpus documents about a kernel's expected gate behaviour, so
57/// the gate is tested in both directions (corpus/README.md).
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum SafetyExpectation {
60    /// Known to flip safety with block size (the §2.2 family).
61    Flip,
62    /// Known not to flip at any block size in its space.
63    Stable,
64    /// No documented expectation.
65    None,
66}
67
68/// A kernel's declared tuning space, loaded from `kernel.toml`.
69#[derive(Debug, Clone)]
70pub struct KernelSpec {
71    /// Kernel name, matching the corpus directory.
72    pub name: String,
73    /// The `#[kernel]` entry function name.
74    pub entry: String,
75    /// Launch-contract domain (1, 2 or 3).
76    pub domain: u8,
77    /// Minimum compute capability the kernel needs, e.g. "7.0".
78    pub needs_cc: Option<String>,
79    /// What the corpus asserts this kernel's gate result should be, so a
80    /// regression in the analyzer shows up as a corpus failure.
81    pub known: SafetyExpectation,
82    /// Directory holding the kernel crate (where kernel.toml lives).
83    pub dir: PathBuf,
84    /// Every declared dimension, sorted by name — the order that makes
85    /// [`crate::Config::id`] canonical.
86    pub dims: Vec<Dim>,
87    /// Expressions every enumerated configuration must satisfy.
88    pub constraints: Vec<Constraint>,
89}
90
91#[derive(Deserialize)]
92struct RawSpec {
93    kernel: RawKernel,
94    #[serde(default)]
95    dims: toml::value::Table,
96    #[serde(default)]
97    constraints: RawConstraints,
98}
99
100#[derive(Deserialize)]
101struct RawKernel {
102    name: String,
103    entry: String,
104    domain: u8,
105    #[serde(default)]
106    needs_cc: Option<String>,
107    #[serde(default)]
108    known: Option<String>,
109}
110
111#[derive(Deserialize, Default)]
112struct RawConstraints {
113    #[serde(default)]
114    exprs: Vec<String>,
115}
116
117#[derive(Deserialize)]
118struct RawDim {
119    #[serde(default)]
120    role: Option<String>,
121    values: Vec<toml::Value>,
122}
123
124impl KernelSpec {
125    /// Load `<dir>/kernel.toml`.
126    pub fn load(dir: &Path) -> Result<Self, SpaceError> {
127        let path = dir.join("kernel.toml");
128        let text = std::fs::read_to_string(&path).map_err(|source| SpaceError::Io {
129            path: path.display().to_string(),
130            source,
131        })?;
132        Self::parse(&path.display().to_string(), &text, dir.to_path_buf())
133    }
134
135    /// Parse from a string (tests).
136    pub fn from_toml_str(origin: &str, text: &str) -> Result<Self, SpaceError> {
137        Self::parse(origin, text, PathBuf::from("."))
138    }
139
140    fn parse(origin: &str, text: &str, dir: PathBuf) -> Result<Self, SpaceError> {
141        let raw: RawSpec = toml::from_str(text).map_err(|source| SpaceError::Parse {
142            path: origin.to_string(),
143            source: Box::new(source),
144        })?;
145
146        if !(1..=3).contains(&raw.kernel.domain) {
147            return Err(SpaceError::Invalid(format!(
148                "domain must be 1..=3, got {}",
149                raw.kernel.domain
150            )));
151        }
152        let known = match raw.kernel.known.as_deref() {
153            Some("flip") => SafetyExpectation::Flip,
154            Some("stable") => SafetyExpectation::Stable,
155            None => SafetyExpectation::None,
156            Some(other) => {
157                return Err(SpaceError::Invalid(format!(
158                    "known must be \"flip\" or \"stable\", got {other:?}"
159                )));
160            }
161        };
162
163        let mut dims = Vec::new();
164        for (name, value) in raw.dims {
165            if !name
166                .chars()
167                .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
168            {
169                return Err(SpaceError::Invalid(format!(
170                    "dimension name {name:?} must be [a-z0-9_]"
171                )));
172            }
173            let raw_dim: RawDim = value
174                .try_into()
175                .map_err(|e| SpaceError::Invalid(format!("dimension {name}: {e}")))?;
176            let role = match raw_dim.role.as_deref() {
177                Some("launch") => DimRole::Launch,
178                Some("spec") => DimRole::Spec,
179                None if name.starts_with("block_") => DimRole::Launch,
180                None => DimRole::Spec,
181                Some(other) => {
182                    return Err(SpaceError::Invalid(format!(
183                        "dimension {name}: role must be \"launch\" or \"spec\", got {other:?}"
184                    )));
185                }
186            };
187            let mut values = Vec::new();
188            for v in raw_dim.values {
189                match v {
190                    toml::Value::Integer(n) if n >= 0 => values.push(Value::Int(n as u64)),
191                    toml::Value::String(s) => values.push(Value::Str(s)),
192                    other => {
193                        return Err(SpaceError::Invalid(format!(
194                            "dimension {name}: values must be non-negative integers or strings, got {other}"
195                        )));
196                    }
197                }
198            }
199            if values.is_empty() {
200                return Err(SpaceError::Invalid(format!(
201                    "dimension {name} has no values"
202                )));
203            }
204            // CUDA's per-axis block limits (Programming Guide, "Compute
205            // Capability" table: max x/y 1024, max z 64, and at most 1024
206            // threads per block overall). Refusing here means `block_threads`
207            // is never asked to multiply anything that could overflow, and it
208            // means the operator hears about a typo when the spec loads
209            // rather than as a launch failure ten minutes into a sweep. The
210            // product limit is not checkable here -- it depends on which
211            // values are drawn together, which is enumeration's job.
212            if let Some(axis) = BLOCK_AXIS_LIMITS.iter().find(|(n, _)| *n == name) {
213                let (_, limit) = axis;
214                for v in &values {
215                    if let Value::Int(n) = v
216                        && *n > *limit
217                    {
218                        return Err(SpaceError::Invalid(format!(
219                            "dimension {name}: {n} exceeds the CUDA limit of {limit} \
220                             threads on this axis (max 1024 per block overall)"
221                        )));
222                    }
223                }
224            }
225            let mut seen = values.clone();
226            seen.sort();
227            seen.dedup();
228            if seen.len() != values.len() {
229                return Err(SpaceError::Invalid(format!(
230                    "dimension {name} has duplicate values"
231                )));
232            }
233            dims.push(Dim { name, role, values });
234        }
235        if dims.is_empty() {
236            return Err(SpaceError::Invalid("spec declares no dimensions".into()));
237        }
238        dims.sort_by(|a, b| a.name.cmp(&b.name));
239
240        let dim_names: Vec<&str> = dims.iter().map(|d| d.name.as_str()).collect();
241        let constraints = raw
242            .constraints
243            .exprs
244            .iter()
245            .map(|e| Constraint::parse(e, &dim_names))
246            .collect::<Result<Vec<_>, _>>()?;
247
248        Ok(KernelSpec {
249            name: raw.kernel.name,
250            entry: raw.kernel.entry,
251            domain: raw.kernel.domain,
252            needs_cc: raw.kernel.needs_cc,
253            known,
254            dir,
255            dims,
256            constraints,
257        })
258    }
259
260    /// One dimension by name, or `None` if the spec does not declare it.
261    pub fn dim(&self, name: &str) -> Option<&Dim> {
262        self.dims.iter().find(|d| d.name == name)
263    }
264
265    /// Dimensions in canonical (name-sorted) order.
266    pub fn dims_sorted(&self) -> Vec<&Dim> {
267        self.dims.iter().collect()
268    }
269}