metrale_circuit/hardware/
plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: Plan one model on one device: resolve the device's class (chain, rules, kernel
4//! sources, families), check the class's build against the device, fuse every requested mode
5//! and row count with the kernels the device can run, and mark every node no rule covers as a
6//! gap instead of failing.
7//!
8//! Owner: metrale-circuit (hardware).
9//! Invariants:
10//! - A rule that matches a node whose edge format it cannot read is left out of that plan and
11//!   listed ([`Planned::refused`]); the plan never fails on it and never runs it. 2026-10-02: So
12//!   is a rule whose kernels declare another pipeline than the plan requires of a node it
13//!   covers ([`crate::pipeline`]): a kernel never runs a precision the circuit does not ask for.
14//! - A node no available rule covers is planned by a placeholder group (`emitter = "novel"`,
15//!   rule id `novel.<op>`), never by a kernel the device cannot run; the placeholder has the
16//!   lowest priority, so it never displaces a real rule.
17//! - On the class a model's golden plans were written for, with its own kernel target, the
18//!   plan equals the golden plan (crates/server/src/cli/circuit_hw_tests.rs holds it).
19
20use std::collections::{BTreeMap, BTreeSet};
21
22use super::avail::{Availability, availability, check_build, kernel_status, without_fp4_kernel};
23use super::class::{ClassInfo, ClassRules, chain, class_families, class_rules, planning_chain};
24use super::device::{Device, MmaKind, Registry};
25use super::estimate::{DeviceRoofline, device_roofline};
26use super::exec::{Exec, fp4_fallback, node_exec};
27use super::sources::{ClassSources, KernelTree};
28use super::{HwError, ModelUnderPlan};
29use crate::fuser::{FuseError, FusionPlan, fuse};
30use crate::ir::{Circuit, NodeIdx};
31use crate::rules::{KernelId, Mode, Numerics, PatternOp, Repeat, Rule};
32use crate::runtime::RuntimeRoute;
33use crate::venn::Run;
34use crate::venn::families::{Families, Roofline};
35use crate::venn::roofline::nvfp4_mma;
36
37/// 2026-09-30: The emitter of a placeholder group.
38pub const NOVEL_EMITTER: &str = "novel";
39
40/// 2026-09-30: One fused plan and the nodes only a placeholder covers.
41#[derive(Debug, Clone)]
42pub struct Planned {
43    /// 2026-09-30: Mode and rows.
44    pub run: Run,
45    /// 2026-09-30: The plan.
46    pub plan: FusionPlan,
47    /// 2026-09-30: Nodes in placeholder groups.
48    pub novel: BTreeSet<NodeIdx>,
49    /// 2026-09-30: Rules left out of this plan because they matched a node whose edge format
50    /// they cannot read (the pattern matches op and weight, not the input format), and why.
51    pub refused: Vec<(String, String)>,
52    /// 2026-09-30: The runtime routes that apply to this run, each with its own plan
53    /// ([`crate::runtime`]).
54    pub routes: Vec<RoutePlan>,
55    /// 2026-10-02: Every planned node's required pipeline, which its kernel declares too.
56    pub pipelines: crate::pipeline::PlanPipelines,
57}
58
59/// 2026-09-30: A route that applies to a run, with its arm's plan.
60#[derive(Debug, Clone)]
61pub struct RoutePlan {
62    /// 2026-09-30: The route.
63    pub route: RuntimeRoute,
64    /// 2026-09-30: The arm's plan (its own `routes` are empty).
65    pub planned: Planned,
66}
67
68/// 2026-09-30: A device's class, resolved for one model.
69#[derive(Debug, Clone)]
70pub struct Resolved {
71    /// 2026-09-30: The device.
72    pub device: Device,
73    /// 2026-09-30: The class and its ancestors, own first.
74    pub chain: Vec<ClassInfo>,
75    /// 2026-09-30: The classes rules and families come from ([`planning_chain`]).
76    pub planning: Vec<ClassInfo>,
77    /// 2026-09-30: The class's rules.
78    pub rules: ClassRules,
79    /// 2026-09-30: The class's sources for the model.
80    pub sources: ClassSources,
81    /// 2026-09-30: What the device can run.
82    pub availability: Availability,
83    /// 2026-09-30: The families, available on the device, with the class's points.
84    pub families: Families,
85    /// 2026-09-30: Roofline constants.
86    pub roofline: DeviceRoofline,
87    /// 2026-10-01: The model's NVFP4-activation nodes the class compiles no FP4 block-scale
88    /// kernel for ([`without_fp4_kernel`]).
89    pub without_fp4: BTreeSet<NodeIdx>,
90    /// 2026-10-01: How each node of the model runs here ([`node_exec`]; `None`: reads no
91    /// weight). Its execution columns, its "FP4 costing" row and its cost all read this.
92    pub exec: Vec<Option<Exec>>,
93}
94
95impl Resolved {
96    /// 2026-09-30: The rules, empty for a class without any.
97    pub fn rule_list(&self) -> &[Rule] {
98        match &self.rules {
99            ClassRules::Rules { rules, .. } => rules,
100            ClassRules::None => &[],
101        }
102    }
103
104    /// 2026-09-30: The runtime routes, empty for a class without any.
105    pub fn route_list(&self) -> &[RuntimeRoute] {
106        match &self.rules {
107            ClassRules::Rules { runtime, .. } => runtime,
108            ClassRules::None => &[],
109        }
110    }
111
112    /// 2026-10-01: The constants node `n` is costed with: the device's, with the NVFP4 slot
113    /// (which `node_cost` reads only for an NVFP4 activation) at the peak of the node's
114    /// execution.
115    pub fn roofline_of(&self, n: NodeIdx) -> Roofline {
116        let base = self.roofline.roofline;
117        match self.exec[n] {
118            Some(e) => Roofline {
119                nvfp4_tflops: e.peak(&base).0,
120                ..base
121            },
122            None => base,
123        }
124    }
125
126    /// 2026-10-01: The report's "FP4 costing", read from the nodes' executions: the device's note
127    /// where it has no FP4 MMA; otherwise `native` only when every NVFP4-activation op the report
128    /// costs in `c` runs on it, and the ops that do not named with the path they take.
129    pub fn fp4_costing(&self, c: &Circuit) -> String {
130        let base = &self.roofline.roofline;
131        let fallback = fp4_fallback(&self.device);
132        if !self.device.runs(MmaKind::Fp4BlockScale) {
133            let lacks = match fallback {
134                Exec::ExactFp8Emulation => "FP4 MMA",
135                _ => "FP4 and FP8 MMA",
136            };
137            return format!(
138                "{lacks} not native: NVFP4-activation nodes costed at {}",
139                fallback.peak(base).1
140            );
141        }
142        let sections: BTreeSet<_> = report_runs()
143            .iter()
144            .map(|r| crate::fuser::section_of(r.mode))
145            .collect();
146        let mut by_exec: BTreeMap<Exec, BTreeSet<String>> = BTreeMap::new();
147        for b in c.blocks.iter().filter(|b| sections.contains(&b.section)) {
148            for i in (b.first..b.end).filter(|&i| nvfp4_mma(c, &c.nodes[i])) {
149                if let Some(e) = self.exec[i] {
150                    by_exec.entry(e).or_default().insert(c.nodes[i].op.name());
151                }
152            }
153        }
154        let list = |s: &BTreeSet<String>| s.iter().cloned().collect::<Vec<_>>().join(", ");
155        let native = by_exec.remove(&Exec::Native(MmaKind::Fp4BlockScale));
156        let mut parts: Vec<String> = by_exec
157            .iter()
158            .map(|(e, ops)| {
159                format!(
160                    "no path for {} (the class compiles no FP4 block-scale kernel for them): costed at {}",
161                    list(ops),
162                    e.peak(base).1
163                )
164            })
165            .collect();
166        match (native, parts.is_empty()) {
167            (None, true) => "not used: the model has no NVFP4-activation node".into(),
168            (Some(_), true) => "native".into(),
169            (native, false) => {
170                if let Some(ops) = native {
171                    parts.push(format!("native for {}", list(&ops)));
172                }
173                parts.join("; ")
174            }
175        }
176    }
177}
178
179/// 2026-09-30: Resolve `device_id` for `model` from `tree`.
180pub fn resolve(
181    registry: &Registry,
182    device_id: &str,
183    tree: &dyn KernelTree,
184    model: &ModelUnderPlan,
185) -> Result<Resolved, HwError> {
186    let device = registry.device(device_id)?.clone();
187    let repo = tree.as_repo();
188    let chain = chain(repo, &device.class)?;
189    let own = &chain[0];
190    if own.arch != device.arch {
191        return Err(HwError::Class(format!(
192            "device `{}` declares arch {}, kernels/{}/HARDWARE.toml builds {}",
193            device.id, device.arch, own.name, own.arch
194        )));
195    }
196    check_build(&device, own, &registry.guards)?;
197    let rules = class_rules(repo, &device.class)?;
198    let sources = tree
199        .class_sources(&device.class, &model.kernel_model, &model.kernel_quant)
200        .map_err(HwError::Class)?;
201    let rule_list: &[Rule] = match &rules {
202        ClassRules::Rules { rules, .. } => rules,
203        ClassRules::None => &[],
204    };
205    let availability = availability(&device, own, &sources, rule_list, &registry.guards);
206    let planning = planning_chain(repo, &device.class)?;
207    let mut families = class_families(repo, &planning, &sources)?;
208    families.families.retain(|f| {
209        f.kernels.is_empty()
210            || f.kernels.iter().any(|k| {
211                kernel_status(&device, own, &sources, &registry.guards, k)
212                    .0
213                    .is_ok()
214            })
215    });
216    // 2026-09-30: A family whose points the class compiles none of cannot be instantiated here,
217    // whatever its kernel names match.
218    families
219        .families
220        .retain(|f| f.kernels.is_empty() || !f.points.is_empty());
221    let base = base_roofline(repo, &planning)?;
222    let measured = (families_class(repo, &device.class)?).then_some(&base);
223    let roofline = device_roofline(&device, measured, &base);
224    let kernel = |k: &KernelId| match kernel_status(&device, own, &sources, &registry.guards, k) {
225        (Ok(()), req) => Some(req.is_some_and(|i| i.kind == MmaKind::Fp4BlockScale)),
226        (Err(_), _) => None,
227    };
228    let c = &model.circuit;
229    let without_fp4 = without_fp4_kernel(c, rule_list, &families, &kernel);
230    let exec = c
231        .nodes
232        .iter()
233        .enumerate()
234        .map(|(i, n)| node_exec(&device, c, n, without_fp4.contains(&i)))
235        .collect();
236    Ok(Resolved {
237        device,
238        chain,
239        planning,
240        rules,
241        sources,
242        availability,
243        families,
244        roofline,
245        without_fp4,
246        exec,
247    })
248}
249
250/// 2026-09-30: The roofline block of the nearest KERNEL_FAMILIES.toml along the chain (its
251/// `context_tokens` is the estimate's assumption on every device).
252fn base_roofline(
253    repo: &dyn crate::venn::repo::Repo,
254    chain: &[ClassInfo],
255) -> Result<crate::venn::families::Roofline, HwError> {
256    for c in chain {
257        if let Ok(t) = repo.read(&format!("kernels/{}/common/KERNEL_FAMILIES.toml", c.name)) {
258            let f = crate::venn::parse_families(&t).map_err(|e| HwError::Class(e.to_string()))?;
259            return Ok(f.roofline);
260        }
261    }
262    Err(HwError::Class(format!(
263        "no KERNEL_FAMILIES.toml along `{}`'s chain",
264        chain.first().map_or("", |c| c.name.as_str())
265    )))
266}
267
268/// 2026-09-30: The class holds its own KERNEL_FAMILIES.toml, so its `[roofline]` is a
269/// measurement on it.
270fn families_class(repo: &dyn crate::venn::repo::Repo, class: &str) -> Result<bool, HwError> {
271    let Ok(t) = repo.read(&format!("kernels/{class}/common/KERNEL_FAMILIES.toml")) else {
272        return Ok(false);
273    };
274    let f = crate::venn::parse_families(&t).map_err(|e| HwError::Class(e.to_string()))?;
275    Ok(f.hardware == class)
276}
277
278/// 2026-09-30: The placeholder for `op` reading `input`: a quantized read is stated, since the
279/// fuser matches a pattern's quantized input format exactly (`fuser_match.rs`).
280fn placeholder(op: &crate::ir::OpKind, input: Option<crate::format::Format>) -> Rule {
281    let input = input.filter(|f| !f.is_plain());
282    Rule {
283        id: match input {
284            Some(f) => format!("novel.{}.{}", op.name(), f.name()),
285            None => format!("novel.{}", op.name()),
286        },
287        pattern: vec![PatternOp {
288            op: *op,
289            roles: BTreeSet::new(),
290            layer_kind: None,
291            local: None,
292            weight: None,
293            input,
294            writes: None,
295            keep: false,
296            stored: false,
297            sibling: false,
298            holds: None,
299            steps: BTreeMap::new(),
300        }],
301        kernels: Vec::new(),
302        repeat: Repeat::Once,
303        copies: None,
304        emitter: NOVEL_EMITTER.to_string(),
305        rows: (1, u64::MAX),
306        modes: Mode::ALL.into_iter().collect(),
307        requires: BTreeSet::new(),
308        when: Default::default(),
309        numerics: Numerics::Reference,
310        runs: Vec::new(),
311        priority: i64::MIN,
312        cite: "no rule of this class covers the op on this device".into(),
313    }
314}
315
316/// 2026-09-30: Fuse `circuit` at `run` with the rules and kernels of `r`, covering what no rule
317/// covers with placeholders; each runtime route that applies is planned beside it.
318pub fn fuse_on(
319    r: &Resolved,
320    circuit: &Circuit,
321    policy: &crate::fuser::Policy,
322    run: Run,
323) -> Result<Planned, HwError> {
324    let mut primary = fuse_arm(r, circuit, policy, run)?;
325    for route in r
326        .route_list()
327        .iter()
328        .filter(|x| x.applies(policy, run.mode, run.rows))
329    {
330        let planned = fuse_arm(r, circuit, &route.policy(policy), run)?;
331        if planned.plan.digest != primary.plan.digest {
332            primary.routes.push(RoutePlan {
333                route: route.clone(),
334                planned,
335            });
336        }
337    }
338    Ok(primary)
339}
340
341fn fuse_arm(
342    r: &Resolved,
343    circuit: &Circuit,
344    policy: &crate::fuser::Policy,
345    run: Run,
346) -> Result<Planned, HwError> {
347    let mut rules: Vec<Rule> = r.rule_list().to_vec();
348    let mut refused = Vec::new();
349    loop {
350        match fuse(
351            circuit,
352            &rules,
353            &r.availability.kernels,
354            policy,
355            run.mode,
356            run.rows,
357        ) {
358            Ok(plan) => {
359                let pipelines = match crate::pipeline::check_plan(
360                    circuit,
361                    &plan,
362                    &rules,
363                    &r.families,
364                    &policy.settings,
365                ) {
366                    Ok(p) => p,
367                    Err(crate::pipeline::PipelineError::Mismatch(ms)) => {
368                        // 2026-10-02: Refuse each rule whose kernels run another precision than
369                        // the node requires, and plan again without it.
370                        for m in ms {
371                            if rules.iter().any(|x| x.id == m.rule) {
372                                rules.retain(|x| x.id != m.rule);
373                                refused.push((m.rule.clone(), m.describe()));
374                            }
375                        }
376                        continue;
377                    }
378                    Err(e) => return Err(HwError::Plan(e.to_string())),
379                };
380                let novel = plan
381                    .groups
382                    .iter()
383                    .filter(|g| g.emitter == NOVEL_EMITTER)
384                    .flat_map(|g| g.nodes.iter().copied())
385                    .collect();
386                return Ok(Planned {
387                    run,
388                    plan,
389                    novel,
390                    refused,
391                    routes: Vec::new(),
392                    pipelines,
393                });
394            }
395            Err(FuseError::Uncovered { node, .. }) => {
396                let idx = circuit
397                    .node(&node)
398                    .ok_or_else(|| HwError::Plan(format!("uncovered node `{node}` not found")))?;
399                let n = &circuit.nodes[idx];
400                let p = placeholder(&n.op, n.inputs.first().map(|&e| circuit.edges[e].format));
401                if rules.iter().any(|x| x.id == p.id) {
402                    return Err(HwError::Plan(format!(
403                        "node `{node}` stays uncovered with its placeholder `{}`",
404                        p.id
405                    )));
406                }
407                rules.push(p);
408            }
409            Err(e @ FuseError::FormatConflict { .. }) => {
410                let FuseError::FormatConflict { rule, .. } = &e else {
411                    unreachable!("matched above")
412                };
413                let id = rule.clone();
414                rules.retain(|x| x.id != id);
415                refused.push((id, e.to_string()));
416            }
417            Err(e) => return Err(HwError::Plan(e.to_string())),
418        }
419    }
420}
421
422/// 2026-09-30: The runs every hardware report shows: decode at C=1, and multi-sequence decode at
423/// C=16 and C=128.
424pub fn report_runs() -> Vec<Run> {
425    vec![
426        Run {
427            mode: Mode::Decode,
428            rows: 1,
429        },
430        Run {
431            mode: Mode::MultiSeq,
432            rows: 16,
433        },
434        Run {
435            mode: Mode::MultiSeq,
436            rows: 128,
437        },
438    ]
439}