metrale_circuit/
instantiate.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Instantiate a circuit TOML against an arch shape and a precision resolver.
4//!
5//! The prologue blocks run first, then each layer's blocks in the order its kind maps to,
6//! then the epilogue, then the draft blocks. A block's `stream_in` binds to the edge the
7//! previous stream-writing block left, so layer `i`'s output edge IS layer `i + 1`'s input
8//! edge and a rule can fuse across the boundary.
9//!
10//! Owner: metrale-circuit.
11//! Invariants: see [`crate::circuit_toml`]; every check there is enforced here, in file
12//! order, and the first failure is returned.
13
14use std::collections::{BTreeMap, BTreeSet};
15
16use crate::circuit_toml::{
17    BlockFile, CircuitError, CircuitFile, LayoutRule, NodeFile, layout_rule, parse_file, when_holds,
18};
19use crate::format::Format;
20use crate::ir::{ArchShape, BlockInstance, Circuit, LayerKind, Node, OpKind, Section};
21use crate::precision::{EdgePrecision, LinearFormats};
22
23/// 2026-09-28: Parse `text`, merge the block libraries it includes (`includes`: name to
24/// text), and instantiate it for `shape`, asking `precision` for every weight-reading node's
25/// formats.
26pub fn instantiate(
27    text: &str,
28    includes: &[(&str, &str)],
29    shape: &ArchShape,
30    precision: &dyn EdgePrecision,
31) -> Result<Circuit, CircuitError> {
32    let file = parse_file(text, includes)?;
33    let rule = layout_rule(&file.layout)?;
34    check_dims(&file, shape)?;
35    let sequence = block_sequence(&file, &rule, &shape.layer_kinds, &shape.dims)?;
36    let mut b = Builder {
37        file: &file,
38        shape,
39        precision,
40        circuit: Circuit {
41            arch: file.arch.clone(),
42            description: file.description.clone(),
43            nodes: Vec::new(),
44            edges: Vec::new(),
45            blocks: Vec::new(),
46            layer_kinds: shape.layer_kinds.clone(),
47            dims: shape.dims.clone(),
48            states: Vec::new(),
49        },
50        stream: None,
51        quantized: BTreeMap::new(),
52    };
53    for (template, layer, section, module) in sequence {
54        b.block(&template, layer, section, module)?;
55    }
56    b.finish()
57}
58
59fn check_dims(file: &CircuitFile, shape: &ArchShape) -> Result<(), CircuitError> {
60    for d in &file.dims {
61        if !shape.dims.contains_key(d) {
62            return Err(CircuitError::ShapeMismatch(format!(
63                "circuit `{}` needs dim `{d}`, which the arch shape does not give",
64                file.arch
65            )));
66        }
67    }
68    Ok(())
69}
70
71/// 2026-09-30: One block to instantiate: template, layer, section, and a module that replaces
72/// the draft module for its `{L}`.
73type Planned = (String, Option<usize>, Section, Option<String>);
74
75fn block_sequence(
76    file: &CircuitFile,
77    rule: &LayoutRule,
78    kinds: &[LayerKind],
79    dims: &BTreeMap<String, u64>,
80) -> Result<Vec<Planned>, CircuitError> {
81    if kinds.is_empty() {
82        return Err(CircuitError::Layout("the arch shape has no layers".into()));
83    }
84    let mut used = BTreeSet::new();
85    let mut out = Vec::new();
86    let known = |name: &str, used: &mut BTreeSet<String>| {
87        if !file.block.contains_key(name) {
88            return Err(CircuitError::Layout(format!("no block template `{name}`")));
89        }
90        used.insert(name.to_string());
91        Ok(())
92    };
93    for name in &file.prologue {
94        known(name, &mut used)?;
95        out.push((name.clone(), None, Section::Main, None));
96    }
97    // 2026-09-30: The layout blocks per kind, with the overrides of every switch that holds.
98    let mut layout = file.layout.blocks.clone();
99    for (w, over) in &file.layout.when {
100        for names in over.values() {
101            for n in names {
102                known(n, &mut used)?;
103            }
104        }
105        if when_holds(w, &file.dims, dims)? {
106            layout.extend(over.clone());
107        }
108    }
109    for (i, kind) in kinds.iter().enumerate() {
110        if let LayoutRule::Interval { period } = rule {
111            let want = if (i + 1) % period == 0 {
112                LayerKind::FullAttention
113            } else {
114                LayerKind::LinearAttention
115            };
116            if *kind != want {
117                return Err(CircuitError::Layout(format!(
118                    "layer {i} is {} but the interval-{period} layout puts {} there",
119                    kind.name(),
120                    want.name()
121                )));
122            }
123        }
124        let blocks = layout.get(kind.name()).ok_or_else(|| {
125            CircuitError::Layout(format!(
126                "layer {i} is {}, which the layout maps to no blocks",
127                kind.name()
128            ))
129        })?;
130        for name in blocks {
131            known(name, &mut used)?;
132            out.push((name.clone(), Some(i), Section::Main, None));
133        }
134    }
135    for name in &file.epilogue {
136        known(name, &mut used)?;
137        out.push((name.clone(), None, Section::Main, None));
138    }
139    let on = match &file.draft_when {
140        Some(w) => when_holds(w, &file.dims, dims)?,
141        None => true,
142    };
143    let mut draft = &file.draft;
144    for (w, list) in &file.draft_variant {
145        for entry in list {
146            known(entry.split('@').next().unwrap_or_default(), &mut used)?;
147        }
148        if when_holds(w, &file.dims, dims)? {
149            draft = list;
150        }
151    }
152    for entry in &file.draft {
153        known(entry.split('@').next().unwrap_or_default(), &mut used)?;
154    }
155    if on {
156        for entry in draft {
157            let (name, module) = match entry.split_once('@') {
158                Some((n, m)) => (n.to_string(), Some(m.to_string())),
159                None => (entry.clone(), None),
160            };
161            out.push((name, None, Section::Draft, module));
162        }
163    }
164    if !file.draft.is_empty() && file.draft_module.is_none() {
165        return Err(CircuitError::Layout(
166            "`draft` blocks need a `draft_module` for their bindings".into(),
167        ));
168    }
169    if let Some(unused) = file.local_blocks.iter().find(|k| !used.contains(*k)) {
170        return Err(CircuitError::Layout(format!(
171            "block template `{unused}` is never used"
172        )));
173    }
174    Ok(out)
175}
176
177struct Builder<'a> {
178    file: &'a CircuitFile,
179    shape: &'a ArchShape,
180    precision: &'a dyn EdgePrecision,
181    circuit: Circuit,
182    stream: Option<usize>,
183    /// 2026-09-29: `(edge, format, owning projection of a static scale)` to the edge holding
184    /// it quantized to that format.
185    quantized: BTreeMap<(usize, Format, Option<String>), usize>,
186}
187
188impl Builder<'_> {
189    fn block(
190        &mut self,
191        template: &str,
192        layer: Option<usize>,
193        section: Section,
194        module_override: Option<String>,
195    ) -> Result<(), CircuitError> {
196        let file = self.file;
197        let tpl: &BlockFile = &file.block[template];
198        let (prefix, module) = match (layer, section) {
199            (Some(i), _) => (
200                format!("l{i}.{template}"),
201                Some(file.layer_module.replace("{i}", &i.to_string())),
202            ),
203            (None, Section::Draft) => (
204                format!("draft.{template}"),
205                module_override.or_else(|| file.draft_module.clone()),
206            ),
207            (None, Section::Main) => (template.to_string(), None),
208        };
209        let first = self.circuit.nodes.len();
210        let mut local: BTreeMap<String, usize> = BTreeMap::new();
211        let stream_in = match &tpl.stream_in {
212            Some(name) => {
213                let edge = self.stream.ok_or_else(|| {
214                    CircuitError::Layout(format!(
215                        "block `{template}` reads the stream before any block writes it"
216                    ))
217                })?;
218                local.insert(name.clone(), edge);
219                Some(edge)
220            }
221            None => None,
222        };
223        let first_state = self.circuit.states.len();
224        for sf in &tpl.state {
225            let decl = state_decl(
226                template,
227                &prefix,
228                layer,
229                section,
230                sf,
231                &tpl.state,
232                &self.shape.dims,
233            )?;
234            if self.circuit.states.iter().any(|s| s.id == decl.id) {
235                return Err(dup(template, &sf.id));
236            }
237            self.circuit.states.push(decl);
238        }
239        let mut node_ids = BTreeSet::new();
240        for nf in &tpl.node {
241            if !node_ids.insert(nf.id.as_str()) {
242                return Err(dup(template, &nf.id));
243            }
244            self.node(template, &prefix, layer, module.as_deref(), nf, &mut local)?;
245        }
246        for out in &tpl.outputs {
247            let name = out.edge();
248            let e = *local.get(name).ok_or_else(|| {
249                CircuitError::Layout(format!(
250                    "block `{template}` declares output `{name}`, which it does not produce"
251                ))
252            })?;
253            self.circuit.edges[e].is_output = true;
254            if let crate::circuit_toml::OutputFile::Bound { buffer, .. } = out {
255                let b = crate::model_buffer::ModelBuffer::parse(buffer).ok_or_else(|| {
256                    CircuitError::Layout(format!(
257                        "block `{template}` output `{name}`: `{buffer}` is no model buffer \
258                         (logits, tokens, draft_embed)"
259                    ))
260                })?;
261                self.circuit.edges[e].binds = Some(b);
262            }
263        }
264        let stream_out = match &tpl.stream_out {
265            Some(name) => {
266                let e = *local.get(name).ok_or_else(|| {
267                    CircuitError::Layout(format!(
268                        "block `{template}` writes stream `{name}`, which it does not produce"
269                    ))
270                })?;
271                self.stream = Some(e);
272                Some(e)
273            }
274            None => None,
275        };
276        check_state_access(template, &self.circuit, first_state, first)?;
277        self.circuit.blocks.push(BlockInstance {
278            template: template.to_string(),
279            layer,
280            section,
281            first,
282            end: self.circuit.nodes.len(),
283            stream_in,
284            stream_out,
285        });
286        Ok(())
287    }
288
289    fn node(
290        &mut self,
291        template: &str,
292        prefix: &str,
293        layer: Option<usize>,
294        module: Option<&str>,
295        nf: &NodeFile,
296        local: &mut BTreeMap<String, usize>,
297    ) -> Result<(), CircuitError> {
298        if let Some(w) = &nf.when
299            && !when_holds(w, &self.file.dims, &self.shape.dims)?
300        {
301            return self.elide(template, nf, local);
302        }
303        let fmt = nf
304            .format
305            .as_deref()
306            .map(Format::parse)
307            .transpose()
308            .map_err(|source| CircuitError::Format {
309                block: template.to_string(),
310                edge: nf.id.clone(),
311                source,
312            })?;
313        let op =
314            OpKind::parse(&nf.op, nf.role.as_deref(), fmt).map_err(|source| CircuitError::Op {
315                block: template.to_string(),
316                node: nf.id.clone(),
317                source,
318            })?;
319        let id = format!("{prefix}.{}", nf.id);
320        let mut inputs = Vec::with_capacity(nf.inputs.len());
321        for name in &nf.inputs {
322            let e = *local.get(name).ok_or_else(|| CircuitError::DanglingInput {
323                block: template.to_string(),
324                node: nf.id.clone(),
325                edge: name.clone(),
326            })?;
327            inputs.push(e);
328        }
329        let mut binding = Vec::with_capacity(nf.binding.len());
330        for b in &nf.binding {
331            binding.push(match (b.contains("{L}"), module) {
332                (true, Some(m)) => b.replace("{L}", m),
333                (true, None) => {
334                    return Err(CircuitError::Binding {
335                        node: id.clone(),
336                        detail: format!("`{b}` names {{L}} outside a layer or draft block"),
337                    });
338                }
339                (false, _) => b.clone(),
340            });
341        }
342        let formats = match nf.unquantized {
343            true => edges::unquantized(
344                &id,
345                &op,
346                inputs.first().map(|&x| self.circuit.edges[x].format),
347            )?,
348            false => self.resolve(&id, &op, &binding, &inputs)?,
349        };
350        if let Some(f) = formats {
351            let x = inputs[0];
352            let have = self.circuit.edges[x].format;
353            // 2026-09-30: A projection whose precision reads F32 where the circuit has a 16-bit
354            // elementwise product (the grouped FP8 expert kernels read the FP32 SiLU product)
355            // makes its producer write F32, while it is the edge's only reader. A 16-bit edge
356            // into a projection declared with quantized activations gets the quantizer node.
357            if have == Format::Bf16 && f.activation == Format::F32 && self.widenable(x) {
358                self.circuit.edges[x].format = Format::F32;
359            } else if have != f.activation {
360                let plain_source = matches!(have, Format::Bf16 | Format::F32);
361                if !plain_source || f.activation.is_plain() || !f.activation.is_edge_format() {
362                    return Err(self.mismatch(&id, &op, x, Some(f.activation)));
363                }
364                inputs[0] =
365                    self.quantized(x, f.activation, template, layer, (&nf.id, binding.first()));
366            }
367        }
368        let idx = self.circuit.nodes.len();
369        for &e in &inputs {
370            self.circuit.edges[e].consumers.push(idx);
371        }
372        let mut outputs = Vec::with_capacity(nf.out.len());
373        for ef in &nf.out {
374            if local.contains_key(&ef.edge) {
375                return Err(dup(template, &ef.edge));
376            }
377            let e = self.edge(template, prefix, idx, ef)?;
378            local.insert(ef.edge.clone(), e);
379            outputs.push(e);
380        }
381        if let OpKind::ActQuant(f) = op {
382            for &o in &outputs {
383                self.expect(&id, &op, o, f)?;
384            }
385        }
386        let weight = formats.map(|f| f.weight);
387        for &e in &inputs {
388            let f = self.circuit.edges[e].format;
389            if weight.is_none() && !f.is_plain() && !matches!(op, OpKind::Copy) {
390                return Err(self.mismatch(&id, &op, e, None));
391            }
392        }
393        let state = self.state_refs(template, prefix, nf)?;
394        self.circuit.nodes.push(Node {
395            id,
396            local: nf.id.clone(),
397            op,
398            inputs,
399            outputs,
400            weight,
401            binding,
402            params: self.params(nf)?,
403            layer,
404            block: template.to_string(),
405            state,
406        });
407        Ok(())
408    }
409
410    fn resolve(
411        &self,
412        id: &str,
413        op: &OpKind,
414        binding: &[String],
415        inputs: &[usize],
416    ) -> Result<Option<LinearFormats>, CircuitError> {
417        if !op.reads_linear_weight() {
418            return Ok(None);
419        }
420        let Some(first) = binding.first() else {
421            return Err(CircuitError::Binding {
422                node: id.to_string(),
423                detail: format!("`{}` reads a weight but binds no module", op.name()),
424            });
425        };
426        // 2026-09-28: Experts share one declared precision; `*` is asked as expert 0.
427        let ask = |m: &str| self.precision.linear(&m.replace('*', "0"));
428        let formats: LinearFormats = ask(first);
429        for m in &binding[1..] {
430            let other = ask(m);
431            if other != formats {
432                return Err(CircuitError::MixedPrecision {
433                    node: id.to_string(),
434                    detail: format!(
435                        "{first} = {}/{}, {m} = {}/{}",
436                        formats.weight, formats.activation, other.weight, other.activation
437                    ),
438                });
439            }
440        }
441        if inputs.is_empty() {
442            return Err(CircuitError::Binding {
443                node: id.to_string(),
444                detail: "a weight-reading node needs an activation input".into(),
445            });
446        }
447        Ok(Some(formats))
448    }
449
450    /// 2026-09-30: A node's params: its own, plus those of each `params_when` switch that
451    /// holds, with every `{dim}` replaced by the dim's value.
452    fn params(&self, nf: &NodeFile) -> Result<BTreeMap<String, String>, CircuitError> {
453        let mut out = nf.params.clone();
454        for (w, extra) in &nf.params_when {
455            if when_holds(w, &self.file.dims, &self.shape.dims)? {
456                out.extend(extra.clone());
457            }
458        }
459        for v in out.values_mut() {
460            if let Some(name) = v.strip_prefix('{').and_then(|r| r.strip_suffix('}')) {
461                let d = self.shape.dims.get(name).ok_or_else(|| {
462                    CircuitError::ShapeMismatch(format!(
463                        "node `{}` param `{v}` names dim `{name}`, which the arch shape lacks",
464                        nf.id
465                    ))
466                })?;
467                *v = d.to_string();
468            }
469        }
470        Ok(out)
471    }
472
473    fn finish(self) -> Result<Circuit, CircuitError> {
474        for e in &self.circuit.edges {
475            if e.consumers.is_empty() && !e.is_output {
476                return Err(CircuitError::DanglingOutput(e.id.clone()));
477            }
478        }
479        Ok(self.circuit)
480    }
481}
482
483fn dup(block: &str, name: &str) -> CircuitError {
484    CircuitError::Duplicate {
485        block: block.to_string(),
486        name: name.to_string(),
487    }
488}
489
490#[path = "instantiate/edges.rs"]
491mod edges;
492mod states;
493use states::check_state_access;
494pub(crate) use states::state_decl;
495
496#[cfg(test)]
497#[path = "instantiate_tests.rs"]
498mod instantiate_tests;
metrale_circuit/instantiate.rs

metrale_circuit/
instantiate.rs
