launchbound_space/
constraint.rs

1//! A deliberately small constraint language over integer dimensions.
2//!
3//! Grammar (no parentheses, two precedence levels):
4//!   constraint := arith CMP arith
5//!   arith      := term (('+' | '-') term)*
6//!   term       := atom (('*' | '/' | '%') atom)*
7//!   atom       := integer | dimension-name
8//!   CMP        := '==' | '!=' | '<=' | '>=' | '<' | '>'
9//!
10//! Evaluation is checked u64 arithmetic (division truncates); overflow or
11//! division/modulo by zero makes the constraint an error, never silently
12//! true or false.
13
14use crate::spec::Value;
15use crate::{Config, SpaceError};
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18enum Token {
19    Num(u64),
20    Ident(String),
21    Op(char),
22    Cmp(&'static str),
23}
24
25/// One `[constraints]` expression, parsed once and evaluated per candidate.
26///
27/// A constraint prunes the cartesian product before anything is compiled or
28/// measured: `tile % block_x == 0` removes the combinations a kernel cannot
29/// use. It is the author's statement about their own kernel, not an
30/// analysis — launchbound evaluates it and attaches no meaning to the
31/// dimension names.
32#[derive(Debug, Clone)]
33pub struct Constraint {
34    text: String,
35    lhs: Vec<Token>,
36    cmp: &'static str,
37    rhs: Vec<Token>,
38}
39
40impl Constraint {
41    /// The expression as written in `kernel.toml`, for error messages.
42    pub fn text(&self) -> &str {
43        &self.text
44    }
45
46    /// Parse one expression, checking every identifier against `dims`.
47    ///
48    /// Unknown identifiers are refused here rather than at evaluation, so a
49    /// typo in `kernel.toml` is reported once at load instead of once per
50    /// candidate.
51    pub fn parse(expr: &str, dims: &[&str]) -> Result<Self, SpaceError> {
52        let err = |reason: &str| SpaceError::Constraint {
53            expr: expr.to_string(),
54            reason: reason.to_string(),
55        };
56        let tokens = tokenize(expr).map_err(|r| err(&r))?;
57        let cmp_pos = tokens
58            .iter()
59            .position(|t| matches!(t, Token::Cmp(_)))
60            .ok_or_else(|| err("no comparison operator"))?;
61        let Token::Cmp(cmp) = tokens[cmp_pos] else {
62            unreachable!()
63        };
64        if tokens.iter().filter(|t| matches!(t, Token::Cmp(_))).count() != 1 {
65            return Err(err("exactly one comparison operator required"));
66        }
67        let lhs = tokens[..cmp_pos].to_vec();
68        let rhs = tokens[cmp_pos + 1..].to_vec();
69        for side in [&lhs, &rhs] {
70            if side.is_empty() {
71                return Err(err("empty side of comparison"));
72            }
73            for t in side {
74                if let Token::Ident(name) = t
75                    && !dims.contains(&name.as_str())
76                {
77                    return Err(err(&format!("unknown dimension `{name}`")));
78                }
79            }
80        }
81        Ok(Constraint {
82            text: expr.to_string(),
83            lhs,
84            cmp,
85            rhs,
86        })
87    }
88
89    /// Does this configuration satisfy the constraint?
90    ///
91    /// Arithmetic that would overflow or divide by zero is an
92    /// [`SpaceError::Constraint`], never a silent `false`: a candidate
93    /// dropped because the constraint could not be computed is
94    /// indistinguishable from one the author meant to exclude, and the two
95    /// need different fixes.
96    pub fn eval(&self, config: &Config) -> Result<bool, SpaceError> {
97        let resolve = config_resolver(config);
98        let l = eval_arith(&self.lhs, &resolve, &self.text)?;
99        let r = eval_arith(&self.rhs, &resolve, &self.text)?;
100        Ok(match self.cmp {
101            "==" => l == r,
102            "!=" => l != r,
103            "<=" => l <= r,
104            ">=" => l >= r,
105            "<" => l < r,
106            ">" => l > r,
107            _ => unreachable!(),
108        })
109    }
110}
111
112fn tokenize(expr: &str) -> Result<Vec<Token>, String> {
113    let mut tokens = Vec::new();
114    let bytes = expr.as_bytes();
115    let mut i = 0;
116    while i < bytes.len() {
117        let c = bytes[i] as char;
118        match c {
119            ' ' | '\t' => i += 1,
120            '0'..='9' => {
121                let start = i;
122                while i < bytes.len() && bytes[i].is_ascii_digit() {
123                    i += 1;
124                }
125                let n: u64 = expr[start..i]
126                    .parse()
127                    .map_err(|_| "integer literal too large".to_string())?;
128                tokens.push(Token::Num(n));
129            }
130            'a'..='z' | '_' => {
131                let start = i;
132                while i < bytes.len()
133                    && (bytes[i].is_ascii_lowercase()
134                        || bytes[i].is_ascii_digit()
135                        || bytes[i] == b'_')
136                {
137                    i += 1;
138                }
139                tokens.push(Token::Ident(expr[start..i].to_string()));
140            }
141            '*' | '/' | '%' | '+' | '-' => {
142                tokens.push(Token::Op(c));
143                i += 1;
144            }
145            '=' | '!' | '<' | '>' => {
146                let two = &expr[i..(i + 2).min(expr.len())];
147                let cmp = match two {
148                    "==" => Some("=="),
149                    "!=" => Some("!="),
150                    "<=" => Some("<="),
151                    ">=" => Some(">="),
152                    _ => None,
153                };
154                if let Some(cmp) = cmp {
155                    tokens.push(Token::Cmp(cmp));
156                    i += 2;
157                } else if c == '<' {
158                    tokens.push(Token::Cmp("<"));
159                    i += 1;
160                } else if c == '>' {
161                    tokens.push(Token::Cmp(">"));
162                    i += 1;
163                } else {
164                    return Err(format!("unexpected character `{c}`"));
165                }
166            }
167            other => return Err(format!("unexpected character `{other}`")),
168        }
169    }
170    Ok(tokens)
171}
172
173fn eval_arith(
174    tokens: &[Token],
175    resolve: &dyn Fn(&str) -> Result<u64, String>,
176    text: &str,
177) -> Result<u64, SpaceError> {
178    let err = |reason: String| SpaceError::Constraint {
179        expr: text.to_string(),
180        reason,
181    };
182    let atom = |t: &Token| -> Result<u64, SpaceError> {
183        match t {
184            Token::Num(n) => Ok(*n),
185            Token::Ident(name) => resolve(name).map_err(err),
186            Token::Op(_) | Token::Cmp(_) => Err(err("misplaced operator".into())),
187        }
188    };
189
190    // First pass: fold * and % into a term list separated by +/-.
191    let mut terms: Vec<(char, u64)> = Vec::new(); // (sign-op, value)
192    let mut pending_op: Option<char> = None; // within-term * or %
193    let mut sign: char = '+';
194    let mut current: Option<u64> = None;
195    for t in tokens {
196        match t {
197            Token::Op(op @ ('*' | '/' | '%')) => {
198                if current.is_none() {
199                    return Err(err(format!("`{op}` with no left operand")));
200                }
201                pending_op = Some(*op);
202            }
203            Token::Op(op @ ('+' | '-')) => {
204                let value = current
205                    .take()
206                    .ok_or_else(|| err(format!("`{op}` with no left operand")))?;
207                terms.push((sign, value));
208                sign = *op;
209                pending_op = None;
210            }
211            atom_token => {
212                let v = atom(atom_token)?;
213                current = Some(match (current, pending_op.take()) {
214                    (None, None) => v,
215                    (Some(acc), Some('*')) => acc
216                        .checked_mul(v)
217                        .ok_or_else(|| err("multiplication overflow".into()))?,
218                    (Some(acc), Some('%')) => {
219                        if v == 0 {
220                            return Err(err("modulo by zero".into()));
221                        }
222                        acc % v
223                    }
224                    (Some(acc), Some('/')) => {
225                        if v == 0 {
226                            return Err(err("division by zero".into()));
227                        }
228                        acc / v
229                    }
230                    (Some(_), None) => {
231                        return Err(err("two operands with no operator".into()));
232                    }
233                    (None, Some(_)) => unreachable!(),
234                    (Some(_), Some(_)) => unreachable!(),
235                });
236            }
237        }
238    }
239    let value = current.ok_or_else(|| err("trailing operator".into()))?;
240    terms.push((sign, value));
241
242    let mut acc: u64 = 0;
243    for (op, v) in terms {
244        acc = match op {
245            '+' => acc
246                .checked_add(v)
247                .ok_or_else(|| err("addition overflow".into()))?,
248            '-' => acc
249                .checked_sub(v)
250                .ok_or_else(|| err("subtraction underflow".into()))?,
251            _ => unreachable!(),
252        };
253    }
254    Ok(acc)
255}
256
257fn config_resolver(config: &Config) -> impl Fn(&str) -> Result<u64, String> + '_ {
258    move |name: &str| match config.get(name) {
259        Some(Value::Int(n)) => Ok(*n),
260        Some(Value::Str(_)) => Err(format!(
261            "dimension `{name}` is a string and cannot be used in arithmetic"
262        )),
263        None => Err(format!("dimension `{name}` missing from config")),
264    }
265}
266
267/// Evaluate a comparison-free arithmetic expression against a candidate's
268/// dimensions plus extra named variables (bench plans use this for grid
269/// shapes and buffer sizes, e.g. `elements / block_x`).
270pub fn eval_arith_expr(
271    expr: &str,
272    config: &Config,
273    extra: &std::collections::BTreeMap<String, u64>,
274) -> Result<u64, SpaceError> {
275    let err = |reason: &str| SpaceError::Constraint {
276        expr: expr.to_string(),
277        reason: reason.to_string(),
278    };
279    let tokens = tokenize(expr).map_err(|r| err(&r))?;
280    if tokens.iter().any(|t| matches!(t, Token::Cmp(_))) {
281        return Err(err("comparison operators are not allowed here"));
282    }
283    let base = config_resolver(config);
284    let resolve = move |name: &str| -> Result<u64, String> {
285        if let Some(v) = extra.get(name) {
286            return Ok(*v);
287        }
288        base(name)
289    };
290    eval_arith(&tokens, &resolve, expr)
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296    use crate::KernelSpec;
297
298    fn config_with(block: u64, tile: u64) -> Config {
299        let spec = KernelSpec::from_toml_str(
300            "t",
301            &format!(
302                r#"
303                [kernel]
304                name = "t"
305                entry = "t"
306                domain = 1
307                [dims.block_x]
308                values = [{block}]
309                [dims.tile]
310                values = [{tile}]
311                "#
312            ),
313        )
314        .unwrap();
315        crate::enumerate(&spec).unwrap().into_iter().next().unwrap()
316    }
317
318    #[test]
319    fn arithmetic_and_comparisons() {
320        let dims = ["block_x", "tile"];
321        let c = config_with(64, 256);
322        for (expr, expected) in [
323            ("tile % block_x == 0", true),
324            ("tile % block_x != 0", false),
325            ("block_x * tile <= 16384", true),
326            ("block_x * tile < 16384", false),
327            ("tile - block_x == 192", true),
328            ("tile + block_x >= 320", true),
329            ("block_x > 32", true),
330        ] {
331            let parsed = Constraint::parse(expr, &dims).unwrap();
332            assert_eq!(parsed.eval(&c).unwrap(), expected, "{expr}");
333        }
334    }
335
336    #[test]
337    fn rejects_unknown_dimension_and_junk() {
338        let dims = ["block_x"];
339        assert!(Constraint::parse("bogus == 1", &dims).is_err());
340        assert!(Constraint::parse("block_x == ", &dims).is_err());
341        assert!(Constraint::parse("block_x", &dims).is_err());
342        assert!(Constraint::parse("block_x == 1 == 2", &dims).is_err());
343        assert!(Constraint::parse("block_x @ 2", &dims).is_err());
344    }
345
346    #[test]
347    fn division_by_zero_is_an_error_not_a_verdict() {
348        let dims = ["block_x", "tile"];
349        let c = config_with(64, 0);
350        let parsed = Constraint::parse("block_x % tile == 0", &dims).unwrap();
351        assert!(parsed.eval(&c).is_err());
352    }
353}
launchbound_space/constraint.rs

launchbound_space/
constraint.rs
