metrale_circuit/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Architecture circuits. A model is a graph of ops (`kernels/circuits/<arch>.toml`)
4//! whose edges the fuser marks fused or materialised by applying the hardware's fusion rules
5//! (`kernels/<hw>/common/FUSIONS.toml`) for one mode and row count.
6//!
7//! Owner: metrale-circuit.
8//! Invariants:
9//! - Pure: no GPU, no file I/O, no environment, no clock and no randomness. Callers read the
10//!   TOMLs and pass their text in ([`load`]).
11//! - A plan is deterministic, and its digest is recorded beside the closure hash.
12
13pub mod checkpoint;
14pub mod circuit_toml;
15pub mod config_map;
16pub mod declared_precision;
17pub mod digest;
18pub mod dims;
19pub mod display;
20pub mod format;
21pub mod fuser;
22pub mod hardware;
23pub mod instances;
24pub mod instantiate;
25pub mod ir;
26pub mod memory;
27pub mod model_buffer;
28pub mod pipeline;
29pub mod planner;
30pub mod precision;
31pub mod precision_policy;
32pub mod render;
33pub mod rules;
34pub mod runs;
35pub mod runtime;
36pub mod state;
37pub mod state_ops;
38pub mod venn;
39
40#[cfg(test)]
41mod test_toy;
42
43pub use checkpoint::{
44    CheckpointError, QuantMetadata, ResolvedCheckpoint, ServePrecision,
45    instantiate_from_checkpoint, map_checkpoint, recurrent_states, resolve_checkpoint,
46};
47pub use circuit_toml::{CircuitError, includes_of};
48pub use format::{Format, Scale};
49pub use fuser::{
50    AvailableKernels, EdgeState, FuseError, FusionPlan, Group, Policy, fuse, fuse_table,
51};
52pub use instances::{Instance, InstanceError, PrecisionSpec, parse_instances};
53pub use instantiate::instantiate;
54pub use ir::{ArchShape, Circuit, LayerKind, LinearRole, OpKind, Section};
55pub use precision::{EdgePrecision, LinearFormats, PrecisionError, PrecisionTable};
56pub use rules::{KernelId, Mode, Numerics, Rule, RuleError, parse_rules};
57pub use runs::{RowTable, VerifyRun};
58pub use runtime::{RuleSet, RuntimeRoute, parse_rule_set};
59
60/// 2026-09-28: Any failure between the TOML texts and a rendered plan.
61#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
62pub enum LoadError {
63    /// 2026-09-28: The circuit.
64    #[error(transparent)]
65    Circuit(#[from] CircuitError),
66    /// 2026-09-28: FUSIONS.toml.
67    #[error(transparent)]
68    Rules(#[from] RuleError),
69    /// 2026-09-28: The precision table.
70    #[error(transparent)]
71    Precision(#[from] PrecisionError),
72    /// 2026-09-28: Fusion.
73    #[error(transparent)]
74    Fuse(#[from] FuseError),
75    /// 2026-09-28: The display renderer.
76    #[error(transparent)]
77    Display(#[from] display::DisplayError),
78    /// 2026-09-28: The buffer planner.
79    #[error(transparent)]
80    Plan(#[from] planner::PlanError),
81    /// 2026-10-02: A node's kernel declares another pipeline than the plan requires, or a
82    /// requirement or declaration is missing ([`pipeline`]).
83    #[error(transparent)]
84    Pipeline(#[from] pipeline::PipelineError),
85    /// 2026-09-28: The precision table describes another checkpoint.
86    #[error("precision table is for `{table}`, the instance serves `{instance}`")]
87    CheckpointMismatch {
88        /// 2026-09-28: The table's checkpoint.
89        table: String,
90        /// 2026-09-28: The instance's checkpoint.
91        instance: String,
92    },
93}
94
95/// 2026-09-28: The texts one instance is built from.
96#[derive(Debug, Clone, Copy)]
97pub struct Sources<'a> {
98    /// 2026-09-28: `kernels/circuits/<arch>.toml`.
99    pub circuit: &'a str,
100    /// 2026-09-28: The file the instance's [`PrecisionSpec`] names: a precision table or a
101    /// checkpoint plan fixture (`kernels/circuits/checkpoints/<name>.toml`).
102    pub precision: &'a str,
103    /// 2026-09-28: `kernels/<hw>/common/FUSIONS.toml`.
104    pub rules: &'a str,
105    /// 2026-09-28: The block libraries the circuit includes, by name:
106    /// `kernels/circuits/blocks/<name>.toml`.
107    pub blocks: &'a [(&'a str, &'a str)],
108}
109
110/// 2026-09-28: An instance's circuit and rules, ready to fuse.
111#[derive(Debug, Clone)]
112pub struct Loaded {
113    /// 2026-09-28: The instantiated circuit.
114    pub circuit: Circuit,
115    /// 2026-09-28: The rules, in file order.
116    pub rules: Vec<Rule>,
117    /// 2026-09-30: The runtime routes, in file order.
118    pub runtime: Vec<RuntimeRoute>,
119    /// 2026-09-28: SHA-256 of the FUSIONS.toml text ([`digest::rules_digest`]).
120    pub rules_digest: String,
121}
122
123/// 2026-09-28: Parse and instantiate `instance` from `src`, with the precision its spec names.
124pub fn load(instance: &Instance, src: Sources<'_>) -> Result<Loaded, LoadError> {
125    let mismatch = |table: String| LoadError::CheckpointMismatch {
126        table,
127        instance: instance.checkpoint.clone(),
128    };
129    match &instance.precision {
130        PrecisionSpec::Table(_) => {
131            let table = PrecisionTable::parse(src.precision)?;
132            if table.checkpoint != instance.checkpoint {
133                return Err(mismatch(table.checkpoint));
134            }
135            load_with(instance, src, &table)
136        }
137        PrecisionSpec::Policy {
138            tier, caps, engine, ..
139        } => {
140            let plan = precision_policy::CheckpointPlan::parse(src.precision)?;
141            if plan.checkpoint != instance.checkpoint {
142                return Err(mismatch(plan.checkpoint));
143            }
144            let policy = metrale_config::WeightQuantPolicy::new(
145                precision_policy::tier_named(tier)?,
146                &plan.plan,
147                precision_policy::caps_named(caps)?,
148            );
149            load_with(
150                instance,
151                src,
152                &precision_policy::PolicyPrecision::new(policy, engine),
153            )
154        }
155    }
156}
157
158/// 2026-09-28: Parse and instantiate `instance` from `src` with `precision`: the executor
159/// passes the served model's own policy here, the checked-in fixture being only the offline
160/// copy of it.
161pub fn load_with(
162    instance: &Instance,
163    src: Sources<'_>,
164    precision: &dyn EdgePrecision,
165) -> Result<Loaded, LoadError> {
166    let circuit = instantiate(src.circuit, src.blocks, &instance.shape, precision)?;
167    let RuleSet { rules, runtime } = parse_rule_set(src.rules)?;
168    Ok(Loaded {
169        circuit,
170        rules,
171        runtime,
172        rules_digest: digest::rules_digest(src.rules),
173    })
174}
175
176/// 2026-09-28: The header lines a rendering of `instance` carries.
177pub fn header(instance: &Instance) -> render::Header {
178    let mut h = vec![
179        ("recipe".to_string(), instance.recipe.clone()),
180        ("checkpoint".to_string(), instance.checkpoint.clone()),
181        ("target".to_string(), instance.target.clone()),
182    ];
183    let settings: Vec<String> = instance
184        .policy
185        .settings
186        .iter()
187        .map(|(k, v)| format!("{k}={v}"))
188        .collect();
189    h.push(("settings".to_string(), settings.join(" ")));
190    let levers: Vec<&str> = instance
191        .policy
192        .opt_in_levers
193        .iter()
194        .map(String::as_str)
195        .collect();
196    let levers = if levers.is_empty() {
197        "none".to_string()
198    } else {
199        levers.join(" ")
200    };
201    h.push(("opt-in levers".to_string(), levers));
202    h
203}
204
205/// 2026-09-28: Fuse and render one plan of `instance`. 2026-10-02: each group line names the
206/// compute unit its kernels run on ([`venn::compute`]), from the target class's kernel families,
207/// and each member node its pipeline; a node whose kernel declares another pipeline than the
208/// plan requires refuses the plan ([`pipeline::check_plan`]).
209pub fn render_plan(
210    instance: &Instance,
211    loaded: &Loaded,
212    available: &AvailableKernels,
213    mode: Mode,
214    rows: u64,
215    families: &venn::Families,
216) -> Result<String, LoadError> {
217    let plan = fuse(
218        &loaded.circuit,
219        &loaded.rules,
220        available,
221        &instance.policy,
222        mode,
223        rows,
224    )?;
225    let head = header(instance);
226    let c = &loaded.circuit;
227    let group = |g: &fuser::Group| hardware::tc_policy::unit_tag(families, &g.kernels);
228    let pipes = checked(loaded, &plan, &instance.policy, families)?;
229    let node = |n: ir::NodeIdx| pipes.line(n);
230    let notes = render::Notes {
231        group: &group,
232        node: &node,
233    };
234    let mut text = render::render_noted(c, &plan, &head, &notes);
235    let set = (loaded.rules.as_slice(), loaded.runtime.as_slice());
236    for (route, arm) in runtime::route_arms(c, set, available, &instance.policy, &plan)? {
237        let policy = route.policy(&instance.policy);
238        let h = render::with_settings(&head, &policy);
239        let arm_pipes = checked(loaded, &arm, &policy, families)?;
240        let node = |n: ir::NodeIdx| arm_pipes.line(n);
241        let notes = render::Notes {
242            group: &group,
243            node: &node,
244        };
245        text.push_str(&render::route_section(c, &route, &arm, &h, &notes));
246    }
247    Ok(text)
248}
249
250/// 2026-10-02: The pipelines of `plan`, once its kernels' declarations are checked against
251/// what it requires under `policy`.
252fn checked(
253    loaded: &Loaded,
254    plan: &FusionPlan,
255    policy: &Policy,
256    families: &venn::Families,
257) -> Result<pipeline::PlanPipelines, LoadError> {
258    Ok(pipeline::check_plan(
259        &loaded.circuit,
260        plan,
261        &loaded.rules,
262        families,
263        &policy.settings,
264    )?)
265}
266
267/// 2026-09-30: Fuse and render one batched-verify plan of `instance`, for `table`.
268/// 2026-10-02: Checked and annotated as [`render_plan`].
269pub fn render_table_plan(
270    instance: &Instance,
271    loaded: &Loaded,
272    available: &AvailableKernels,
273    table: &RowTable,
274    families: &venn::Families,
275) -> Result<String, LoadError> {
276    let plan = fuse_table(
277        &loaded.circuit,
278        &loaded.rules,
279        available,
280        &instance.policy,
281        table,
282    )?;
283    let pipes = checked(loaded, &plan, &instance.policy, families)?;
284    let group = |g: &fuser::Group| hardware::tc_policy::unit_tag(families, &g.kernels);
285    let node = |n: ir::NodeIdx| pipes.line(n);
286    let notes = render::Notes {
287        group: &group,
288        node: &node,
289    };
290    Ok(render::render_noted(
291        &loaded.circuit,
292        &plan,
293        &header(instance),
294        &notes,
295    ))
296}
297
298/// 2026-09-28: Fuse, lay out and draw one plan of `instance`: the view `met circuit display`
299/// prints, over the same plan [`render_plan`] renders. 2026-10-02: Checked and drawn with each
300/// node's pipeline, as [`render_plan`].
301pub fn display_plan(
302    instance: &Instance,
303    loaded: &Loaded,
304    available: &AvailableKernels,
305    (mode, rows): (Mode, u64),
306    opts: &display::DisplayOpts,
307    families: &venn::Families,
308) -> Result<display::Document, LoadError> {
309    let plan = fuse(
310        &loaded.circuit,
311        &loaded.rules,
312        available,
313        &instance.policy,
314        mode,
315        rows,
316    )?;
317    let buffers = planner::plan_buffers(&loaded.circuit, &plan, rows)?;
318    let pipes = checked(loaded, &plan, &instance.policy, families)?;
319    let info = display::DisplayInfo {
320        checkpoint: instance.checkpoint.clone(),
321        recipe: instance.recipe.clone(),
322        bytes: Some((buffers.materialized_bytes, buffers.arena_bytes)),
323        pipelines: (0..loaded.circuit.nodes.len())
324            .map(|n| pipes.facts(n))
325            .collect(),
326    };
327    Ok(display::display(&loaded.circuit, &plan, &info, opts)?)
328}