metrale_circuit/
dims.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Dimension expressions: a sum of products of named dims and integer literals,
4//! e.g. `q_heads*head_dim*2` or `lin_qk*2+lin_v`. 2026-10-02: A factor may be a dim divided by a
5//! literal, rounded up (`moe_inter/128`: the scale count of a 128-wide group).
6//!
7//! Owner: metrale-circuit.
8//! Invariants:
9//! - An expression keeps its source text, so a plan renders the shape the template wrote.
10//! - Evaluation names every unknown dim and refuses overflow; it never substitutes a value.
11
12use std::collections::BTreeMap;
13
14/// 2026-09-28: A parsed dimension expression.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct DimExpr {
17    text: String,
18    terms: Vec<Vec<Factor>>,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22enum Factor {
23    Name(String),
24    Lit(u64),
25    /// 2026-10-02: `name/lit`, rounded up.
26    CeilDiv(String, u64),
27}
28
29/// 2026-09-28: Why a dimension expression did not parse or evaluate.
30#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
31pub enum DimError {
32    /// 2026-09-28: Empty term, stray character or a literal of zero.
33    #[error("bad dimension expression `{0}`")]
34    Syntax(String),
35    /// 2026-09-28: A name the arch shape does not define.
36    #[error("dimension `{name}` in `{expr}` is not defined by the arch shape")]
37    Unknown {
38        /// 2026-09-28: The missing name.
39        name: String,
40        /// 2026-09-28: The expression that named it.
41        expr: String,
42    },
43    /// 2026-09-28: The value does not fit in u64.
44    #[error("dimension expression `{0}` overflows")]
45    Overflow(String),
46}
47
48impl DimExpr {
49    /// 2026-09-28: Parse `text`. Whitespace is ignored.
50    pub fn parse(text: &str) -> Result<Self, DimError> {
51        let syntax = || DimError::Syntax(text.to_string());
52        let compact: String = text.chars().filter(|c| !c.is_whitespace()).collect();
53        if compact.is_empty() {
54            return Err(syntax());
55        }
56        let mut terms = Vec::new();
57        for term in compact.split('+') {
58            let mut factors = Vec::new();
59            for f in term.split('*') {
60                factors.push(parse_factor(f).ok_or_else(syntax)?);
61            }
62            terms.push(factors);
63        }
64        Ok(DimExpr {
65            text: compact,
66            terms,
67        })
68    }
69
70    /// 2026-09-28: The expression as written, whitespace removed.
71    pub fn text(&self) -> &str {
72        &self.text
73    }
74
75    /// 2026-09-28: Every name the expression reads.
76    pub fn names(&self) -> impl Iterator<Item = &str> {
77        self.terms.iter().flatten().filter_map(|f| match f {
78            Factor::Name(n) | Factor::CeilDiv(n, _) => Some(n.as_str()),
79            Factor::Lit(_) => None,
80        })
81    }
82
83    /// 2026-09-28: The value under `dims`.
84    pub fn eval(&self, dims: &BTreeMap<String, u64>) -> Result<u64, DimError> {
85        let overflow = || DimError::Overflow(self.text.clone());
86        let mut sum: u64 = 0;
87        for term in &self.terms {
88            let mut prod: u64 = 1;
89            for f in term {
90                let dim = |n: &String| {
91                    dims.get(n).copied().ok_or_else(|| DimError::Unknown {
92                        name: n.clone(),
93                        expr: self.text.clone(),
94                    })
95                };
96                let v = match f {
97                    Factor::Lit(v) => *v,
98                    Factor::Name(n) => dim(n)?,
99                    Factor::CeilDiv(n, d) => dim(n)?.div_ceil(*d),
100                };
101                prod = prod.checked_mul(v).ok_or_else(overflow)?;
102            }
103            sum = sum.checked_add(prod).ok_or_else(overflow)?;
104        }
105        Ok(sum)
106    }
107}
108
109fn parse_factor(f: &str) -> Option<Factor> {
110    if let Some((name, lit)) = f.split_once('/') {
111        let d = lit.parse::<u64>().ok().filter(|&v| v > 0)?;
112        return match parse_factor(name)? {
113            Factor::Name(n) => Some(Factor::CeilDiv(n, d)),
114            Factor::Lit(_) | Factor::CeilDiv(..) => None,
115        };
116    }
117    let first = f.chars().next()?;
118    if first.is_ascii_digit() {
119        return f.parse::<u64>().ok().filter(|&v| v > 0).map(Factor::Lit);
120    }
121    let ident = (first.is_ascii_lowercase() || first == '_')
122        && f.chars()
123            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_');
124    ident.then(|| Factor::Name(f.to_string()))
125}
126
127#[cfg(test)]
128#[path = "dims_tests.rs"]
129mod dims_tests;
metrale_circuit/dims.rs

metrale_circuit/
dims.rs
