1use 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
37pub const NOVEL_EMITTER: &str = "novel";
39
40#[derive(Debug, Clone)]
42pub struct Planned {
43 pub run: Run,
45 pub plan: FusionPlan,
47 pub novel: BTreeSet<NodeIdx>,
49 pub refused: Vec<(String, String)>,
52 pub routes: Vec<RoutePlan>,
55 pub pipelines: crate::pipeline::PlanPipelines,
57}
58
59#[derive(Debug, Clone)]
61pub struct RoutePlan {
62 pub route: RuntimeRoute,
64 pub planned: Planned,
66}
67
68#[derive(Debug, Clone)]
70pub struct Resolved {
71 pub device: Device,
73 pub chain: Vec<ClassInfo>,
75 pub planning: Vec<ClassInfo>,
77 pub rules: ClassRules,
79 pub sources: ClassSources,
81 pub availability: Availability,
83 pub families: Families,
85 pub roofline: DeviceRoofline,
87 pub without_fp4: BTreeSet<NodeIdx>,
90 pub exec: Vec<Option<Exec>>,
93}
94
95impl Resolved {
96 pub fn rule_list(&self) -> &[Rule] {
98 match &self.rules {
99 ClassRules::Rules { rules, .. } => rules,
100 ClassRules::None => &[],
101 }
102 }
103
104 pub fn route_list(&self) -> &[RuntimeRoute] {
106 match &self.rules {
107 ClassRules::Rules { runtime, .. } => runtime,
108 ClassRules::None => &[],
109 }
110 }
111
112 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 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
179pub 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, ®istry.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, ®istry.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, ®istry.guards, k)
212 .0
213 .is_ok()
214 })
215 });
216 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, ®istry.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
250fn 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
268fn 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
278fn 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
316pub 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 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
422pub 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}