launchbound_space/
lib.rs

1//! Configuration space model for launchbound.
2//!
3//! A kernel declares its tunable dimensions in a `kernel.toml` next to its
4//! source. This crate loads that spec, enumerates the (constraint-filtered)
5//! configuration space deterministically, and gives every configuration a
6//! canonical, stable, hashable ID. Enumeration is a pure function of the
7//! spec: same spec, same order, byte for byte.
8
9mod constraint;
10mod spec;
11
12pub use constraint::{Constraint, eval_arith_expr};
13pub use spec::{Dim, DimRole, KernelSpec, SafetyExpectation, Value};
14
15use sha2::{Digest, Sha256};
16use std::collections::BTreeMap;
17use std::fmt;
18
19#[derive(Debug, thiserror::Error)]
20pub enum SpaceError {
21    #[error("failed to read {path}: {source}")]
22    Io {
23        path: String,
24        source: std::io::Error,
25    },
26    #[error("failed to parse {path}: {source}")]
27    Parse {
28        path: String,
29        source: Box<toml::de::Error>,
30    },
31    #[error("invalid kernel spec: {0}")]
32    Invalid(String),
33    #[error("invalid constraint `{expr}`: {reason}")]
34    Constraint { expr: String, reason: String },
35}
36
37/// One point in a kernel's configuration space: a total assignment of every
38/// declared dimension. Dimensions are kept sorted by name, which is what
39/// makes the canonical ID canonical.
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct Config {
42    kernel: String,
43    values: BTreeMap<String, Value>,
44}
45
46impl Config {
47    pub fn kernel(&self) -> &str {
48        &self.kernel
49    }
50
51    pub fn get(&self, dim: &str) -> Option<&Value> {
52        self.values.get(dim)
53    }
54
55    pub fn values(&self) -> impl Iterator<Item = (&str, &Value)> {
56        self.values.iter().map(|(k, v)| (k.as_str(), v))
57    }
58
59    /// Total threads per block implied by this configuration. Absent block
60    /// dimensions default to 1, matching CUDA launch semantics.
61    pub fn block_threads(&self) -> u64 {
62        ["block_x", "block_y", "block_z"]
63            .iter()
64            .map(|d| match self.values.get(*d) {
65                Some(Value::Int(n)) => *n,
66                _ => 1,
67            })
68            .product()
69    }
70
71    /// The canonical, stable ID: a versioned SHA-256 over the kernel name
72    /// and the sorted dimension assignments. Changing the encoding is a
73    /// breaking change and must bump the `config.v1` tag.
74    pub fn id(&self) -> ConfigId {
75        let mut hasher = Sha256::new();
76        hasher.update(b"launchbound.config.v1\0");
77        hasher.update(self.kernel.as_bytes());
78        hasher.update(b"\0");
79        for (name, value) in &self.values {
80            hasher.update(name.as_bytes());
81            hasher.update(b"=");
82            match value {
83                Value::Int(n) => hasher.update(n.to_string().as_bytes()),
84                Value::Str(s) => hasher.update(s.as_bytes()),
85            }
86            hasher.update(b"\n");
87        }
88        let digest = hasher.finalize();
89        let mut hex = String::with_capacity(16);
90        for byte in &digest[..8] {
91            hex.push_str(&format!("{byte:02x}"));
92        }
93        ConfigId(format!("c1-{hex}"))
94    }
95
96    /// Only the compile-time specialization dimensions, sorted. Candidates
97    /// sharing this key share generated source, and therefore share one
98    /// reconverge verdict and one compiled artifact.
99    pub fn spec_key(&self, spec: &KernelSpec) -> String {
100        let mut parts = Vec::new();
101        for (name, value) in &self.values {
102            if spec.dim(name).is_some_and(|d| d.role == DimRole::Spec) {
103                parts.push(format!("{name}={value}"));
104            }
105        }
106        parts.join(",")
107    }
108}
109
110impl fmt::Display for Config {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        let mut first = true;
113        for (name, value) in &self.values {
114            if !first {
115                write!(f, " ")?;
116            }
117            write!(f, "{name}={value}")?;
118            first = false;
119        }
120        Ok(())
121    }
122}
123
124/// Canonical configuration identifier, e.g. `c1-9f2a4c1e77b0d3a5`.
125#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
126pub struct ConfigId(String);
127
128impl ConfigId {
129    pub fn as_str(&self) -> &str {
130        &self.0
131    }
132}
133
134impl fmt::Display for ConfigId {
135    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
136        f.write_str(&self.0)
137    }
138}
139
140/// Enumerate the full constraint-filtered space, in canonical order.
141///
142/// Order: dimensions sorted by name; values in declared order; the last
143/// dimension varies fastest (odometer). Constraints filter, never reorder.
144pub fn enumerate(spec: &KernelSpec) -> Result<Vec<Config>, SpaceError> {
145    let dims: Vec<&Dim> = spec.dims_sorted();
146    let mut out = Vec::new();
147    if dims.is_empty() {
148        return Ok(out);
149    }
150    let mut indices = vec![0usize; dims.len()];
151    'outer: loop {
152        let mut values = BTreeMap::new();
153        for (dim, &idx) in dims.iter().zip(&indices) {
154            values.insert(dim.name.clone(), dim.values[idx].clone());
155        }
156        let config = Config {
157            kernel: spec.name.clone(),
158            values,
159        };
160        if spec
161            .constraints
162            .iter()
163            .try_fold(true, |ok, c| c.eval(&config).map(|v| ok && v))?
164        {
165            out.push(config);
166        }
167        // Odometer increment, last dimension fastest.
168        for pos in (0..dims.len()).rev() {
169            indices[pos] += 1;
170            if indices[pos] < dims[pos].values.len() {
171                continue 'outer;
172            }
173            indices[pos] = 0;
174        }
175        break;
176    }
177    Ok(out)
178}
179
180/// The size of the unfiltered space (product of value counts).
181pub fn raw_size(spec: &KernelSpec) -> u64 {
182    spec.dims_sorted()
183        .iter()
184        .map(|d| d.values.len() as u64)
185        .product()
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    fn toy_spec() -> KernelSpec {
193        KernelSpec::from_toml_str(
194            "toy",
195            r#"
196            [kernel]
197            name = "toy"
198            entry = "toy"
199            domain = 1
200            [dims.block_x]
201            values = [32, 64, 128]
202            [dims.tile]
203            role = "spec"
204            values = [128, 256]
205            [constraints]
206            exprs = ["tile % block_x == 0"]
207            "#,
208        )
209        .unwrap()
210    }
211
212    #[test]
213    fn enumeration_is_deterministic_and_filtered() {
214        let spec = toy_spec();
215        let a = enumerate(&spec).unwrap();
216        let b = enumerate(&spec).unwrap();
217        assert_eq!(a, b);
218        // 3*2 = 6 raw; tile % block_x == 0 removes (128, tile=128)? no:
219        // 128 % 128 == 0 keeps it; removed are none for 32/64; block_x=128
220        // with tile=128 ok, tile=256 ok. Everything passes here except none.
221        assert_eq!(raw_size(&spec), 6);
222        assert_eq!(a.len(), 6);
223    }
224
225    #[test]
226    fn constraint_actually_filters() {
227        let spec = KernelSpec::from_toml_str(
228            "toy",
229            r#"
230            [kernel]
231            name = "toy"
232            entry = "toy"
233            domain = 1
234            [dims.block_x]
235            values = [32, 48]
236            [dims.tile]
237            values = [64]
238            [constraints]
239            exprs = ["tile % block_x == 0"]
240            "#,
241        )
242        .unwrap();
243        let configs = enumerate(&spec).unwrap();
244        assert_eq!(configs.len(), 1);
245        assert_eq!(configs[0].get("block_x"), Some(&Value::Int(32)));
246    }
247
248    #[test]
249    fn ids_are_stable_and_distinct() {
250        let spec = toy_spec();
251        let configs = enumerate(&spec).unwrap();
252        let ids: Vec<_> = configs.iter().map(|c| c.id()).collect();
253        let mut unique = ids.clone();
254        unique.sort();
255        unique.dedup();
256        assert_eq!(unique.len(), ids.len(), "duplicate config IDs");
257        // Golden: the first canonical config of this exact spec. If this
258        // changes, the ID encoding changed and config.v1 must be bumped.
259        let first = &configs[0];
260        assert_eq!(first.get("block_x"), Some(&Value::Int(32)));
261        assert_eq!(first.get("tile"), Some(&Value::Int(128)));
262        assert_eq!(first.id().as_str(), configs[0].id().as_str());
263        assert!(first.id().as_str().starts_with("c1-"));
264        assert_eq!(first.id().as_str().len(), 3 + 16);
265    }
266
267    #[test]
268    fn block_threads_multiplies_and_defaults() {
269        let spec = toy_spec();
270        let configs = enumerate(&spec).unwrap();
271        assert_eq!(configs[0].block_threads(), 32);
272    }
273
274    #[test]
275    fn spec_key_covers_only_spec_dims() {
276        let spec = toy_spec();
277        let configs = enumerate(&spec).unwrap();
278        assert_eq!(configs[0].spec_key(&spec), "tile=128");
279    }
280}
launchbound_space/lib.rs

launchbound_space/
lib.rs
