metrale_circuit/venn/
repo.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: A whole `met circuit venn` run over a repository reached through [`Repo`]:
4//! resolve the target and compared instances, load their circuits, check the family manifest
5//! against the kernel sources, build and render the report. The CLI supplies a file-system
6//! [`Repo`]; the tests supply the same one over the checked-out tree, so both run this code.
7//!
8//! Owner: metrale-circuit (venn).
9//! Invariants:
10//! - Business logic only: every byte arrives through [`Repo`] (SBIO), nothing touches the
11//!   environment.
12//! - Manifest drift against the sources is an error before any report is built.
13//! - A target named by a checkpoint directory is checked against that checkpoint's own
14//!   `config.json` (layer kinds, declared precision of every bound module); a disagreement is an
15//!   error, never a note.
16
17use super::discover::{KernelSources, drift, macro_files};
18use super::report::{Side, VennInputs};
19use super::{VennArgs, VennError, build, parse_families, parse_measurements, render};
20use crate::instances::{Instance, PrecisionSpec, parse_instances};
21use crate::ir::LayerKind;
22use crate::{Loaded, Sources, includes_of, load};
23
24/// 2026-09-29: Read access to a repository, by repo-relative path.
25pub trait Repo {
26    /// 2026-09-29: The text of `rel`.
27    fn read(&self, rel: &str) -> Result<String, String>;
28    /// 2026-09-29: Every file under the directory `rel`, recursively, repo-relative with `/`.
29    fn list(&self, rel: &str) -> Result<Vec<String>, String>;
30}
31
32/// 2026-09-29: The instance `spec` names: a recipe id, or an arch that exactly one instance
33/// serves, or a checkpoint id.
34pub fn resolve(all: &[Instance], spec: &str) -> Result<Instance, VennError> {
35    if let Some(i) = all
36        .iter()
37        .find(|i| i.recipe == spec || i.checkpoint == spec)
38    {
39        return Ok(i.clone());
40    }
41    let by_arch: Vec<&Instance> = all.iter().filter(|i| i.arch == spec).collect();
42    match by_arch.as_slice() {
43        [one] => Ok((*one).clone()),
44        [] => Err(VennError::Load(format!(
45            "`{spec}` is no recipe, checkpoint or arch in kernels/circuits/INSTANCES.toml (recipes: {})",
46            all.iter()
47                .map(|i| i.recipe.as_str())
48                .collect::<Vec<_>>()
49                .join(", ")
50        ))),
51        many => Err(VennError::Load(format!(
52            "arch `{spec}` has several instances; name one recipe: {}",
53            many.iter()
54                .map(|i| i.recipe.as_str())
55                .collect::<Vec<_>>()
56                .join(", ")
57        ))),
58    }
59}
60
61/// 2026-09-29: The checkpoint id a directory holds, from its path: the Hugging Face cache
62/// layout (`.../models--<org>--<name>/snapshots/<rev>`) or, failing that, the last two path
63/// components (`.../<org>/<name>`).
64pub fn checkpoint_id_of(dir: &str) -> Option<String> {
65    let parts: Vec<&str> = dir
66        .trim_end_matches('/')
67        .split('/')
68        .filter(|p| !p.is_empty())
69        .collect();
70    if let Some(p) = parts.iter().find_map(|p| p.strip_prefix("models--")) {
71        let (org, name) = p.split_once("--")?;
72        return Some(format!("{org}/{name}"));
73    }
74    match parts.as_slice() {
75        [.., org, name] => Some(format!("{org}/{name}")),
76        _ => None,
77    }
78}
79
80pub(crate) fn load_instance(repo: &dyn Repo, inst: &Instance) -> Result<Loaded, VennError> {
81    let io = |e: String| VennError::Load(format!("{}: {e}", inst.recipe));
82    let circuit = repo
83        .read(&format!("kernels/circuits/{}.toml", inst.arch))
84        .map_err(io)?;
85    let names = includes_of(&circuit).map_err(|e| io(e.to_string()))?;
86    let mut blocks = Vec::with_capacity(names.len());
87    for n in names {
88        let text = repo
89            .read(&format!("kernels/circuits/blocks/{n}.toml"))
90            .map_err(io)?;
91        blocks.push((n, text));
92    }
93    let precision = match &inst.precision {
94        PrecisionSpec::Table(n) => format!("kernels/circuits/precision/{n}.toml"),
95        PrecisionSpec::Policy {
96            checkpoint_plan, ..
97        } => format!("kernels/circuits/checkpoints/{checkpoint_plan}.toml"),
98    };
99    let precision = repo.read(&precision).map_err(io)?;
100    let rules = repo
101        .read(&format!("kernels/{}/common/FUSIONS.toml", hardware(inst)?))
102        .map_err(io)?;
103    let blocks: Vec<(&str, &str)> = blocks
104        .iter()
105        .map(|(n, t)| (n.as_str(), t.as_str()))
106        .collect();
107    load(
108        inst,
109        Sources {
110            circuit: &circuit,
111            precision: &precision,
112            rules: &rules,
113            blocks: &blocks,
114        },
115    )
116    .map_err(|e| io(e.to_string()))
117}
118
119fn hardware(inst: &Instance) -> Result<&str, VennError> {
120    inst.target
121        .split('/')
122        .next()
123        .filter(|h| !h.is_empty())
124        .ok_or_else(|| {
125            VennError::Load(format!(
126                "{}: target `{}` names no hardware",
127                inst.recipe, inst.target
128            ))
129        })
130}
131
132/// 2026-09-29: Build and render the report `args` asks for. `checkpoint`, when the target was
133/// given as a checkpoint directory, is that directory's `config.json` text and, if present, its
134/// `hf_quant_config.json` text.
135pub fn report_text(
136    repo: &dyn Repo,
137    args: &VennArgs,
138    checkpoint: Option<(&str, Option<&str>)>,
139) -> Result<String, VennError> {
140    let io = |e: String| VennError::Load(e);
141    let all = parse_instances(&repo.read("kernels/circuits/INSTANCES.toml").map_err(io)?)
142        .map_err(|e| io(e.to_string()))?;
143    let target = resolve(&all, &args.target)?;
144    let against = args
145        .against
146        .iter()
147        .map(|a| resolve(&all, a))
148        .collect::<Result<Vec<_>, _>>()?;
149    let hw = hardware(&target)?;
150    if let Some(a) = against.iter().find(|a| hardware(a).ok() != Some(hw)) {
151        return Err(io(format!(
152            "`{}` runs on another hardware than the target",
153            a.recipe
154        )));
155    }
156    let lt = load_instance(repo, &target)?;
157    if let Some((config, hfq)) = checkpoint {
158        super::checkpoint::check(&target, &lt, config, hfq)?;
159    }
160    let loaded = against
161        .iter()
162        .map(|a| load_instance(repo, a))
163        .collect::<Result<Vec<_>, _>>()?;
164    let fams = parse_families(
165        &repo
166            .read(&format!("kernels/{hw}/common/KERNEL_FAMILIES.toml"))
167            .map_err(io)?,
168    )
169    .map_err(|e| io(e.to_string()))?;
170    let mut src = KernelSources {
171        paths: repo
172            .list("kernels")
173            .map_err(io)?
174            .into_iter()
175            .chain(repo.list("crates").map_err(io)?)
176            .collect(),
177        ..KernelSources::default()
178    };
179    for f in macro_files(&fams) {
180        let text = repo.read(&f).map_err(io)?;
181        src.texts.insert(f, text);
182    }
183    let problems = drift(&fams, &src);
184    if !problems.is_empty() {
185        return Err(VennError::Drift(problems));
186    }
187    let meas = parse_measurements(
188        &repo
189            .read("docs/kernel-perf/measurements.toml")
190            .map_err(io)?,
191    )
192    .map_err(io)?;
193    let inputs = VennInputs {
194        target: Side {
195            instance: &target,
196            loaded: &lt,
197        },
198        against: against
199            .iter()
200            .zip(&loaded)
201            .map(|(instance, loaded)| Side { instance, loaded })
202            .collect(),
203        families: &fams,
204        measurements: &meas,
205        runs: args.runs()?,
206        command: args.command(),
207    };
208    Ok(render(&build(&inputs)?))
209}
210
211/// 2026-09-29: The layer kinds a `config.json` declares, from `layers_block_type`
212/// (`mamba` / `moe` / `attention`) or `layer_types` (`linear_attention` / `full_attention`).
213pub fn config_layer_kinds(cfg: &serde_json::Value) -> Result<Vec<LayerKind>, VennError> {
214    let bad = |d: String| VennError::Checkpoint(d);
215    let (key, list) = ["layers_block_type", "layer_types"]
216        .iter()
217        .find_map(|k| cfg.get(*k).and_then(|v| v.as_array()).map(|a| (*k, a)))
218        .ok_or_else(|| {
219            bad("config.json has neither `layers_block_type` nor `layer_types`".into())
220        })?;
221    list.iter()
222        .map(|v| {
223            let s = v.as_str().unwrap_or_default();
224            let kind = match s {
225                "attention" => Some(LayerKind::FullAttention),
226                other => LayerKind::parse(other),
227            };
228            kind.ok_or_else(|| {
229                bad(format!(
230                    "`{key}` entry `{s}` is no layer kind a circuit models"
231                ))
232            })
233        })
234        .collect()
235}
metrale_circuit/venn/repo.rs

metrale_circuit/venn/
repo.rs
