metrale_circuit/memory/
render.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: `met circuit memory`'s output: a per-node table (one row per block-template node,
4//! summed over the layers that instantiate it, or every node with `per_node`), a per-state and
5//! per-cache table, the totals against the budget and the inverse queries; and the same report
6//! as JSON, every node and every state term listed.
7//!
8//! Owner: metrale-circuit (memory).
9//! Invariants: deterministic text; MiB are 2^20 bytes.
10
11use std::collections::BTreeMap;
12use std::fmt::Write as _;
13
14use serde_json::{Value, json};
15
16use super::MemoryReport;
17use crate::ir::Circuit;
18
19/// 2026-10-02: The answers to the inverse queries, as the caller computed them.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct Inverse {
22    /// 2026-10-02: `(isl, osl, most sequences that fit)`; `None` when one does not fit.
23    pub max_concurrency: Option<(u64, u64, Option<u64>)>,
24    /// 2026-10-02: `(concurrency, osl, longest prompt that fits)`.
25    pub max_isl: Option<(u64, u64, Option<u64>)>,
26    /// 2026-10-02: The KV pool the budget leaves once everything else is placed, as the engine
27    /// sizes it: `(blocks, tokens, blocks this workload needs)`.
28    pub kv_pool: Option<(u64, u64, u64)>,
29}
30
31fn mib(b: u64) -> String {
32    format!("{:.1}", b as f64 / (1u64 << 20) as f64)
33}
34
35/// 2026-10-02: Key of a node's row in the folded table: the block template and the node's id
36/// inside it, with its section.
37fn fold_key(c: &Circuit, i: usize, per_node: bool) -> String {
38    let n = &c.nodes[i];
39    if per_node || n.layer.is_none() {
40        return n.id.clone();
41    }
42    format!("{}.{}", n.block, n.local)
43}
44
45/// 2026-10-02: Key of a state's row in the folded table: its block template and local id,
46/// prefixed `draft.` in the draft head.
47fn state_key(d: &crate::state::StateDecl) -> String {
48    match d.section {
49        crate::ir::Section::Main => format!("{}.{}", d.block, d.local),
50        crate::ir::Section::Draft => format!("draft.{}.{}", d.block, d.local),
51    }
52}
53
54/// 2026-10-02: The report as text. `head` are `(key, value)` header lines.
55pub fn render_text(
56    c: &Circuit,
57    r: &MemoryReport,
58    head: &[(String, String)],
59    inverse: &Inverse,
60    per_node: bool,
61) -> String {
62    let mut s = String::from("# Circuit memory\n\n");
63    for (k, v) in head {
64        let _ = writeln!(s, "{k}: {v}");
65    }
66    let mut rows: BTreeMap<String, (usize, [u64; 5], String)> = BTreeMap::new();
67    let mut order = Vec::new();
68    for m in &r.nodes {
69        let v = [m.stored, m.derived, m.activations, m.workspace, m.state];
70        if v.iter().all(|&x| x == 0) {
71            continue;
72        }
73        let key = fold_key(c, m.node, per_node);
74        let e = rows.entry(key.clone()).or_insert_with(|| {
75            order.push(key.clone());
76            (0, [0; 5], c.nodes[m.node].op.name())
77        });
78        e.0 += 1;
79        for (j, b) in v.into_iter().enumerate() {
80            // 2026-10-02: A family's workspace (column 3) is one buffer: never summed.
81            e.1[j] = if j == 3 { e.1[j].max(b) } else { e.1[j] + b };
82        }
83    }
84    let _ = writeln!(
85        s,
86        "\n## Per node (MiB; x = nodes folded)\n\n| node | op | x | stored | derived | activations | workspace | state |\n|---|---|---:|---:|---:|---:|---:|---:|"
87    );
88    for k in &order {
89        let (n, v, op) = &rows[k];
90        let _ = writeln!(
91            s,
92            "| {k} | {op} | {n} | {} | {} | {} | {} | {} |",
93            mib(v[0]),
94            mib(v[1]),
95            mib(v[2]),
96            mib(v[3]),
97            mib(v[4])
98        );
99    }
100    let _ = writeln!(
101        s,
102        "\n## States and caches (MiB; x = declarations folded)\n\n| state | kind | holding | lifetime | dtype | x | units each | MiB |\n|---|---|---|---|---|---:|---:|---:|"
103    );
104    let decl = |id: &str| c.states.iter().find(|d| d.id == id);
105    let mut srows: BTreeMap<(String, String, String, String, String), (usize, u64, u64)> =
106        BTreeMap::new();
107    for t in &r.states.terms {
108        let Some(d) = decl(&t.state) else { continue };
109        let key = (
110            state_key(d),
111            d.kind.name().to_string(),
112            format!("{:?}", t.holding).to_lowercase(),
113            t.lifetime.name().to_string(),
114            t.dtype.name().to_string(),
115        );
116        let e = srows.entry(key).or_insert((0, t.units, 0));
117        e.0 += 1;
118        e.2 += t.bytes;
119    }
120    for t in &r.caches {
121        let Some(d) = decl(&t.state) else { continue };
122        let key = (
123            state_key(d),
124            d.kind.name().to_string(),
125            if t.host { "host" } else { "cache" }.to_string(),
126            d.lifetime.name().to_string(),
127            t.dtype.name().to_string(),
128        );
129        let e = srows.entry(key).or_insert((0, t.units, 0));
130        e.0 += 1;
131        e.2 += t.bytes;
132    }
133    for ((id, kind, holding, life, dt), (n, units, bytes)) in &srows {
134        if *bytes == 0 {
135            continue;
136        }
137        let _ = writeln!(
138            s,
139            "| {id} | {kind} | {holding} | {life} | {dt} | {n} | {units} | {} |",
140            mib(*bytes)
141        );
142    }
143    if !r.runs.is_empty() {
144        let _ = writeln!(
145            s,
146            "\n## Activation runs (planner)\n\n| run | rows | arena MiB | materialized MiB |\n|---|---:|---:|---:|"
147        );
148        for a in &r.runs {
149            let _ = writeln!(
150                s,
151                "| {} | {} | {} | {} |",
152                a.label,
153                a.rows,
154                mib(a.arena),
155                mib(a.materialized)
156            );
157        }
158    }
159    if !r.workspaces.is_empty() {
160        let _ = writeln!(
161            s,
162            "\n## Workspaces (widest run)\n\n| family | workspace | in legacy arena | nodes | MiB |\n|---|---|---|---:|---:|"
163        );
164        for w in &r.workspaces {
165            let _ = writeln!(
166                s,
167                "| {} | {} | {} | {} | {} |",
168                w.family,
169                w.name,
170                if w.arena { "yes" } else { "no" },
171                w.nodes.len(),
172                mib(w.bytes)
173            );
174        }
175    }
176    let t = &r.totals;
177    let _ = writeln!(s, "\n## Totals (MiB)\n\n| term | MiB |\n|---|---:|");
178    for (k, v) in [
179        ("weights stored (declared formats)", t.weights_stored),
180        ("  of which outside the circuit", t.weights_outside),
181        ("weights derived (load-time copies)", t.weights_derived),
182        ("  of which leaked", t.weights_leaked),
183        ("activations planned (circuit peak)", t.activations_planned),
184        ("activations charged", t.activations),
185        ("workspace charged", t.workspace),
186        ("states (recurrent pool, KV)", t.states),
187        ("caches (device)", t.caches),
188        ("driver", t.driver),
189        ("device total", t.device),
190        ("budget", r.budget_bytes),
191        ("host caches", t.host),
192    ] {
193        let _ = writeln!(s, "| {k} | {} |", mib(v));
194    }
195    let h = r.headroom();
196    let _ = writeln!(
197        s,
198        "\nHeadroom: {:.1} MiB ({}).",
199        h as f64 / (1u64 << 20) as f64,
200        if h >= 0 { "fits" } else { "over budget" }
201    );
202    let ans = |v: &Option<u64>| v.map_or_else(|| "none fits".to_string(), |x| x.to_string());
203    if let Some((isl, osl, c)) = &inverse.max_concurrency {
204        let _ = writeln!(s, "Max concurrency at ISL={isl} OSL={osl}: {}.", ans(c));
205    }
206    if let Some((conc, osl, isl)) = &inverse.max_isl {
207        let _ = writeln!(s, "Max ISL at C={conc} OSL={osl}: {}.", ans(isl));
208    }
209    if let Some((blocks, tokens, need)) = &inverse.kv_pool {
210        let _ = writeln!(
211            s,
212            "KV pool at this budget: {blocks} blocks ({tokens} tokens); this workload needs {need}."
213        );
214    }
215    s
216}
217
218/// 2026-10-02: The report as JSON: every node with memory, every state and cache term.
219pub fn render_json(
220    c: &Circuit,
221    r: &MemoryReport,
222    head: &[(String, String)],
223    inverse: &Inverse,
224) -> Value {
225    let nodes: Vec<Value> = r
226        .nodes
227        .iter()
228        .filter(|m| m.stored + m.derived + m.activations + m.workspace + m.state > 0)
229        .map(|m| {
230            let n = &c.nodes[m.node];
231            let derived: Vec<Value> = r
232                .weights
233                .iter()
234                .find(|w| w.node == m.node)
235                .map(|w| {
236                    w.derived
237                        .iter()
238                        .map(|d| json!({"rule": d.rule, "bytes": d.bytes, "leaked": d.leaked}))
239                        .collect()
240                })
241                .unwrap_or_default();
242            json!({
243                "node": n.id, "op": n.op.name(), "layer": n.layer,
244                "stored": m.stored, "derived": derived, "activations": m.activations,
245                "workspace": m.workspace, "state": m.state,
246            })
247        })
248        .collect();
249    let states: Vec<Value> = r
250        .states
251        .terms
252        .iter()
253        .map(|t| {
254            json!({"state": t.state, "holding": format!("{:?}", t.holding).to_lowercase(),
255                   "lifetime": t.lifetime.name(), "dtype": t.dtype.name(),
256                   "units": t.units, "bytes": t.bytes})
257        })
258        .chain(r.caches.iter().map(|t| {
259            json!({"state": t.state, "kind": t.kind.name(), "host": t.host,
260                   "dtype": t.dtype.name(), "units": t.units, "bytes": t.bytes})
261        }))
262        .collect();
263    let t = &r.totals;
264    json!({
265        "schema": 1,
266        "header": head.iter().map(|(k, v)| (k.clone(), Value::String(v.clone())))
267            .collect::<serde_json::Map<_, _>>(),
268        "nodes": nodes,
269        "states": states,
270        "runs": r.runs.iter().map(|a| json!({"label": a.label, "rows": a.rows,
271            "arena": a.arena, "materialized": a.materialized})).collect::<Vec<_>>(),
272        "workspaces": r.workspaces.iter().map(|w| json!({"family": w.family, "name": w.name,
273            "bytes": w.bytes, "arena": w.arena, "nodes": w.nodes.len()})).collect::<Vec<_>>(),
274        "totals": {
275            "weights_stored": t.weights_stored, "weights_outside": t.weights_outside,
276            "weights_derived": t.weights_derived,
277            "weights_leaked": t.weights_leaked, "activations_planned": t.activations_planned,
278            "activations": t.activations, "workspace": t.workspace, "states": t.states,
279            "caches": t.caches, "host": t.host, "driver": t.driver, "device": t.device,
280            "budget": r.budget_bytes, "headroom": r.headroom().to_string(),
281        },
282        "inverse": {
283            "max_concurrency": inverse.max_concurrency.map(|(i, o, c)|
284                json!({"isl": i, "osl": o, "concurrency": c})),
285            "max_isl": inverse.max_isl.map(|(c, o, i)|
286                json!({"concurrency": c, "osl": o, "isl": i})),
287            "kv_pool": inverse.kv_pool.map(|(b, t, n)|
288                json!({"blocks": b, "tokens": t, "needed_blocks": n})),
289        },
290    })
291}
metrale_circuit/memory/render.rs

metrale_circuit/memory/
render.rs
