metrale_circuit/venn/
report.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: Build a whole Venn report: fuse the compared models at every requested mode and
4//! row count, classify and estimate every target node, group the nodes into sites (one template
5//! node across its layers), rank them by estimated step share, and flag the layer kinds that
6//! cannot batch rows.
7//!
8//! Owner: metrale-circuit (venn).
9//! Invariants:
10//! - The compared models are fused with every kernel their rules name counted as built (the
11//!   offline view `met circuit show` takes). The target is fused only when it is a golden
12//!   instance, whose rules cover it; otherwise its nodes are matched to families by op.
13//! - A flag's cost is the estimate of running the flagged nodes once per row minus running them
14//!   once for all rows, so flags rank by the time the missing multi-row path costs.
15
16use std::collections::BTreeMap;
17
18use super::classify::{candidates, classify_node, usages};
19use super::families::{EvidenceSource, Families};
20use super::measurements::Measurements;
21use super::roofline::node_cost;
22use super::{Finding, Run, Subject, VennError};
23use crate::Loaded;
24use crate::fuser::{AvailableKernels, FusionPlan, fuse, section_of};
25use crate::instances::Instance;
26use crate::ir::{Circuit, LayerKind, NodeIdx};
27use crate::rules::Mode;
28
29/// 2026-09-29: One side of the diagram as loaded.
30#[derive(Debug, Clone, Copy)]
31pub struct Side<'a> {
32    /// 2026-09-29: The instance.
33    pub instance: &'a Instance,
34    /// 2026-09-29: Its circuit and rules.
35    pub loaded: &'a Loaded,
36}
37
38/// 2026-09-29: Everything a report is built from.
39#[derive(Debug, Clone)]
40pub struct VennInputs<'a> {
41    /// 2026-09-29: The model being added.
42    pub target: Side<'a>,
43    /// 2026-09-29: The supported models it is compared with.
44    pub against: Vec<Side<'a>>,
45    /// 2026-09-29: The kernel-family manifest.
46    pub families: &'a Families,
47    /// 2026-09-29: measurements.toml rows.
48    pub measurements: &'a Measurements,
49    /// 2026-09-29: Modes and row counts, in report order.
50    pub runs: Vec<Run>,
51    /// 2026-09-29: The command that regenerates the report, printed in its header.
52    pub command: String,
53}
54
55/// 2026-09-29: One site: a template node across the layers that instantiate it.
56#[derive(Debug, Clone)]
57pub struct Row {
58    /// 2026-09-29: `block.node`.
59    pub site: String,
60    /// 2026-09-29: Op name.
61    pub op: String,
62    /// 2026-09-29: Weight format, for weight-reading ops.
63    pub weight: Option<String>,
64    /// 2026-09-29: First-input format.
65    pub activation: Option<String>,
66    /// 2026-09-29: Nodes (layers) it covers.
67    pub count: usize,
68    /// 2026-09-29: Estimated time over those nodes, microseconds.
69    pub time_us: f64,
70    /// 2026-09-29: `time_us` over the table's total.
71    pub share: f64,
72    /// 2026-09-29: The classification; `None` is novel.
73    pub primary: Option<Finding>,
74    /// 2026-09-29: Other families' opportunities.
75    pub also: Vec<Finding>,
76}
77
78/// 2026-09-29: One mode and row count.
79#[derive(Debug, Clone)]
80pub struct Table {
81    /// 2026-09-29: Mode and rows.
82    pub run: Run,
83    /// 2026-09-29: Estimated step time of the target, microseconds.
84    pub total_us: f64,
85    /// 2026-09-29: Sites, by descending share.
86    pub rows: Vec<Row>,
87}
88
89/// 2026-09-29: A layer kind that cannot batch rows in a mode.
90#[derive(Debug, Clone)]
91pub struct Flag {
92    /// 2026-09-29: Layer kind, or `None` for the prologue, epilogue and draft head.
93    pub layer_kind: Option<LayerKind>,
94    /// 2026-09-29: Mode and rows.
95    pub run: Run,
96    /// 2026-09-29: Where the fact comes from: a legacy citation, or the circuit's sites.
97    pub source: String,
98    /// 2026-09-29: One line.
99    pub note: String,
100    /// 2026-09-29: Estimated time the per-row fallback adds, microseconds.
101    pub added_us: f64,
102    /// 2026-09-29: `added_us` over the multi-row step estimate.
103    pub share: f64,
104}
105
106/// 2026-09-29: A built report.
107#[derive(Debug, Clone)]
108pub struct VennReport {
109    /// 2026-09-29: The target instance.
110    pub target: Instance,
111    /// 2026-09-29: Its circuit's arch and layer kinds.
112    pub arch: String,
113    /// 2026-09-29: Its layer kinds.
114    pub layer_kinds: Vec<LayerKind>,
115    /// 2026-09-29: Weight-reading sites: `site` to (weight, activation).
116    pub precision: BTreeMap<String, (String, String)>,
117    /// 2026-09-29: The compared instances.
118    pub against: Vec<Instance>,
119    /// 2026-09-29: Estimate constants.
120    pub roofline: super::families::Roofline,
121    /// 2026-09-29: Flags, by descending cost.
122    pub flags: Vec<Flag>,
123    /// 2026-09-29: Tables, in run order.
124    pub tables: Vec<Table>,
125    /// 2026-09-29: Measured rows cited by the tables' evidence (`key` to % of floor).
126    pub cited: BTreeMap<String, f64>,
127    /// 2026-09-29: The regenerating command.
128    pub command: String,
129}
130
131fn fused(side: &Side<'_>, run: Run) -> Result<FusionPlan, VennError> {
132    let l = side.loaded;
133    fuse(
134        &l.circuit,
135        &l.rules,
136        &AvailableKernels::all_named_by(&l.rules),
137        &side.instance.policy,
138        run.mode,
139        run.rows,
140    )
141    .map_err(|e| {
142        VennError::Load(format!(
143            "{} {} n={}: {e}",
144            side.instance.recipe,
145            run.mode.name(),
146            run.rows
147        ))
148    })
149}
150
151fn in_section(c: &Circuit, mode: Mode) -> Vec<NodeIdx> {
152    let section = section_of(mode);
153    c.blocks
154        .iter()
155        .filter(|b| b.section == section)
156        .flat_map(|b| b.first..b.end)
157        .collect()
158}
159
160fn site_of(c: &Circuit, n: NodeIdx) -> String {
161    site(&c.nodes[n])
162}
163
164/// 2026-09-30: `block.local`, prefixed `draft.` in the draft head, whose blocks may reuse a
165/// main-stack template (the Nemotron-H draft MoE is the `moe` block at `mtp.layers.1`).
166fn site(node: &crate::ir::Node) -> String {
167    let draft = if node.id.starts_with("draft.") {
168        "draft."
169    } else {
170        ""
171    };
172    format!("{draft}{}.{}", node.block, node.local)
173}
174
175/// 2026-09-29: Build the report.
176pub fn build(inp: &VennInputs<'_>) -> Result<VennReport, VennError> {
177    let mut cited = BTreeMap::new();
178    for f in &inp.families.families {
179        for e in &f.evidence {
180            if let EvidenceSource::Measurement(key) = &e.source {
181                let m = inp
182                    .measurements
183                    .get(key)
184                    .ok_or_else(|| VennError::UnknownMeasurement {
185                        family: f.id.clone(),
186                        key: key.clone(),
187                    })?;
188                cited.insert(key.clone(), m.pct_of_floor);
189            }
190        }
191    }
192    let tc = &inp.target.loaded.circuit;
193    let settings = &inp.target.instance.policy.settings;
194    let mut tables = Vec::with_capacity(inp.runs.len());
195    let mut flags = Vec::new();
196    for &run in &inp.runs {
197        let scope = in_section(tc, run.mode);
198        if scope.is_empty() {
199            return Err(VennError::Run(format!(
200                "the target circuit has no {} section",
201                run.mode.name()
202            )));
203        }
204        let plans: Vec<FusionPlan> = inp
205            .against
206            .iter()
207            .map(|s| fused(s, run))
208            .collect::<Result<_, _>>()?;
209        let target_plan = if inp.target.instance.golden {
210            Some(fused(&inp.target, run)?)
211        } else {
212            None
213        };
214        let against: Vec<Subject<'_>> = inp
215            .against
216            .iter()
217            .zip(&plans)
218            .map(|(s, p)| Subject {
219                recipe: &s.instance.recipe,
220                circuit: &s.loaded.circuit,
221                settings: &s.instance.policy.settings,
222                plan: Some(p),
223            })
224            .collect();
225        let used = usages(&against, inp.families)?;
226        let target = Subject {
227            recipe: &inp.target.instance.recipe,
228            circuit: tc,
229            settings,
230            plan: target_plan.as_ref(),
231        };
232        let mut by_site: BTreeMap<(String, String, Option<String>, Option<String>), Row> =
233            BTreeMap::new();
234        let mut total = 0.0;
235        for &n in &scope {
236            let node = &tc.nodes[n];
237            let cost = node_cost(
238                tc,
239                node,
240                run.mode,
241                run.rows,
242                settings,
243                &inp.families.roofline,
244            )?;
245            total += cost.time_us;
246            let weight = node.weight.map(|w| w.name());
247            let activation = node.inputs.first().map(|&e| tc.edges[e].format.name());
248            let key = (
249                site_of(tc, n),
250                node.op.name(),
251                weight.clone(),
252                activation.clone(),
253            );
254            if let Some(row) = by_site.get_mut(&key) {
255                row.count += 1;
256                row.time_us += cost.time_us;
257                continue;
258            }
259            let (primary, also) = classify_node(&target, n, run.rows, &used, inp.families)?;
260            by_site.insert(
261                key.clone(),
262                Row {
263                    site: key.0,
264                    op: key.1,
265                    weight,
266                    activation,
267                    count: 1,
268                    time_us: cost.time_us,
269                    share: 0.0,
270                    primary,
271                    also,
272                },
273            );
274        }
275        let mut rows: Vec<Row> = by_site.into_values().collect();
276        for r in &mut rows {
277            r.share = if total > 0.0 { r.time_us / total } else { 0.0 };
278        }
279        rows.sort_by(|a, b| {
280            b.time_us
281                .total_cmp(&a.time_us)
282                .then_with(|| a.site.cmp(&b.site))
283        });
284        if run.rows > 1 {
285            flags.extend(flags_for(inp, tc, &scope, run, total)?);
286        }
287        tables.push(Table {
288            run,
289            total_us: total,
290            rows,
291        });
292    }
293    flags.sort_by(|a, b| {
294        b.added_us
295            .total_cmp(&a.added_us)
296            .then_with(|| a.run.cmp(&b.run))
297            .then_with(|| a.source.cmp(&b.source))
298    });
299    let mut precision = BTreeMap::new();
300    for n in &tc.nodes {
301        if let Some(w) = n.weight {
302            let act = n
303                .inputs
304                .first()
305                .map(|&e| tc.edges[e].format.name())
306                .unwrap_or_default();
307            precision.insert(site(n), (w.name(), act));
308        }
309    }
310    Ok(VennReport {
311        target: inp.target.instance.clone(),
312        arch: tc.arch.clone(),
313        layer_kinds: tc.layer_kinds.clone(),
314        precision,
315        against: inp.against.iter().map(|s| s.instance.clone()).collect(),
316        roofline: inp.families.roofline,
317        flags,
318        tables,
319        cited,
320        command: inp.command.clone(),
321    })
322}
323
324fn penalty(
325    inp: &VennInputs<'_>,
326    tc: &Circuit,
327    nodes: &[NodeIdx],
328    run: Run,
329) -> Result<f64, VennError> {
330    let settings = &inp.target.instance.policy.settings;
331    let r = &inp.families.roofline;
332    let mut added = 0.0;
333    for &n in nodes {
334        let one = node_cost(tc, &tc.nodes[n], run.mode, 1, settings, r)?;
335        let all = node_cost(tc, &tc.nodes[n], run.mode, run.rows, settings, r)?;
336        added += run.rows as f64 * one.time_us - all.time_us;
337    }
338    Ok(added)
339}
340
341fn flags_for(
342    inp: &VennInputs<'_>,
343    tc: &Circuit,
344    scope: &[NodeIdx],
345    run: Run,
346    total: f64,
347) -> Result<Vec<Flag>, VennError> {
348    let kind_of = |n: NodeIdx| tc.nodes[n].layer.map(|l| tc.layer_kinds[l]);
349    let mut out = Vec::new();
350    for l in &inp.families.legacy {
351        if l.arch != tc.arch || !l.per_sequence.contains(&run.mode) || run.rows <= l.rows_above {
352            continue;
353        }
354        let nodes: Vec<NodeIdx> = scope
355            .iter()
356            .copied()
357            .filter(|&n| kind_of(n) == Some(l.layer_kind))
358            .filter(|&n| l.sites.is_empty() || l.sites.contains(&site_of(tc, n)))
359            .collect();
360        if nodes.is_empty() {
361            continue;
362        }
363        let added = penalty(inp, tc, &nodes, run)?;
364        out.push(Flag {
365            layer_kind: Some(l.layer_kind),
366            run,
367            source: format!("legacy: {}", l.cite),
368            note: l.note.clone(),
369            added_us: added,
370            share: if total > 0.0 { added / total } else { 0.0 },
371        });
372    }
373    let mut by_kind: BTreeMap<Option<LayerKind>, Vec<NodeIdx>> = BTreeMap::new();
374    for &n in scope {
375        let node = &tc.nodes[n];
376        let fams = inp.families;
377        if candidates(fams, tc, node, run.rows).is_empty()
378            && !candidates(fams, tc, node, 1).is_empty()
379        {
380            by_kind.entry(kind_of(n)).or_default().push(n);
381        }
382    }
383    for (kind, nodes) in by_kind {
384        let added = penalty(inp, tc, &nodes, run)?;
385        let sites: std::collections::BTreeSet<String> =
386            nodes.iter().map(|&n| site_of(tc, n)).collect();
387        out.push(Flag {
388            layer_kind: kind,
389            run,
390            source: format!(
391                "circuit: {}",
392                sites.into_iter().collect::<Vec<_>>().join(", ")
393            ),
394            note: "every family that runs these ops covers one row per launch".into(),
395            added_us: added,
396            share: if total > 0.0 { added / total } else { 0.0 },
397        });
398    }
399    Ok(out)
400}
metrale_circuit/venn/report.rs

metrale_circuit/venn/
report.rs
