1pub mod avail;
23pub mod class;
24pub mod device;
25pub mod estimate;
26pub mod exec;
27pub mod gaps;
28pub mod model;
29pub mod model_checkpoint;
30pub mod plan;
31mod render;
32mod render_matrix;
33mod render_routes;
34pub mod sources;
35pub mod tc_policy;
36mod tc_policy_render;
37
38#[cfg(test)]
39mod fp4_costing_tests;
40#[cfg(test)]
41mod hardware_tests;
42#[cfg(test)]
43#[path = "runtime_tests.rs"]
44mod runtime_tests;
45#[cfg(test)]
46mod test_fixture;
47
48use std::collections::BTreeMap;
49
50use crate::fuser::Policy;
51use crate::render::Header;
52use crate::rules::KernelId;
53use crate::venn::Class;
54use crate::venn::roofline::CostError;
55
56pub use device::{Device, Registry, parse_devices};
57pub use model::{CircuitSource, InstancesSource, ModelSpec, ModelUnderPlan, PrecisionChoice};
58pub use model_checkpoint::CheckpointSource;
59pub use render::render_report;
60pub use render_matrix::{PortList, port_lists, summary_row};
61pub use render_routes::plan_text;
62pub use sources::{ClassSources, KernelTree, Module};
63
64#[derive(Debug, Clone, PartialEq, thiserror::Error)]
66pub enum HwError {
67 #[error("DEVICES.toml: {0}")]
69 Registry(String),
70 #[error("unknown device `{id}` (kernels/DEVICES.toml lists: {})", .known.join(", "))]
72 UnknownDevice {
73 id: String,
75 known: Vec<String>,
77 },
78 #[error("{0}")]
80 Class(String),
81 #[error(
83 "kernels/{class}/HARDWARE.toml {} -D{macro_name}, but device `{device}` {} {requires}: \
84 the build and the registry contradict each other",
85 if *.defined { "defines" } else { "does not define" },
86 if *.defined { "claims" } else { "lacks" }
87 )]
88 BuildContradictsDevice {
89 device: String,
91 class: String,
93 macro_name: String,
95 requires: String,
97 defined: bool,
99 },
100 #[error(
102 "kernels/{class}/HARDWARE.toml [tensor_core_policy]: these ops must run on tensor cores \
103 and no exemption lists them:\n {violations}"
104 )]
105 TensorCore {
106 class: String,
108 violations: String,
110 },
111 #[error("{0}")]
113 Model(String),
114 #[error("{0}")]
116 Plan(String),
117 #[error(transparent)]
119 Cost(#[from] CostError),
120}
121
122#[derive(Debug, Clone)]
124pub struct HwReport {
125 pub model: ModelUnderPlan,
127 pub resolved: plan::Resolved,
129 pub policy: Policy,
131 pub class_settings: BTreeMap<String, String>,
133 pub header: Header,
135 pub tables: Vec<gaps::GapTable>,
137 pub tensor_core: Vec<tc_policy::TcAudit>,
140 pub routes: Vec<render_routes::RouteRow>,
142 pub prefill: Vec<(u64, f64)>,
144 pub footprint: estimate::Footprint,
146 pub weight_floor: f64,
148 pub exec: BTreeMap<(String, String, String, exec::Exec), usize>,
150 pub absent: Vec<(KernelId, avail::Absence)>,
153 pub others: Vec<(String, f64)>,
155 pub command: String,
157}
158
159impl HwReport {
160 pub fn covered(&self, table: usize, measured_only: bool) -> f64 {
162 let classes: &[Class] = if measured_only {
163 &[Class::Shared]
164 } else {
165 &[Class::Shared, Class::SharedUnmeasured]
166 };
167 self.tables.get(table).map_or(0.0, |t| t.share_of(classes))
168 }
169}
170
171#[derive(Debug, Clone)]
175pub struct OnePlan {
176 pub resolved: plan::Resolved,
178 pub policy: Policy,
180 pub header: Header,
182 pub planned: plan::Planned,
184}
185
186pub fn plan_one(
188 registry: &Registry,
189 device_id: &str,
190 tree: &dyn KernelTree,
191 model: &ModelUnderPlan,
192 run: crate::venn::Run,
193) -> Result<OnePlan, HwError> {
194 let resolved = plan::resolve(registry, device_id, tree, model)?;
195 let (policy, _) = model::policy_on_class(
196 &model.policy,
197 model.settings_class.as_deref(),
198 &resolved.chain[0],
199 )?;
200 let header =
201 crate::render::with_settings(&model::header_on(model, &resolved.device.class), &policy);
202 let planned = plan::fuse_on(&resolved, &model.circuit, &policy, run)?;
203 tc_policy::enforce(&resolved, &model.circuit, &planned)?;
204 Ok(OnePlan {
205 resolved,
206 policy,
207 header,
208 planned,
209 })
210}
211
212pub fn display_on(
215 model: &ModelUnderPlan,
216 one: &OnePlan,
217 opts: &crate::display::DisplayOpts,
218) -> Result<crate::display::Document, crate::LoadError> {
219 let plan = &one.planned.plan;
220 let buffers = crate::planner::plan_buffers(&model.circuit, plan, plan.rows)?;
221 let info = crate::display::DisplayInfo {
222 checkpoint: model.checkpoint.clone(),
223 recipe: format!("{} on {}", model.label, one.resolved.device.id),
224 bytes: Some((buffers.materialized_bytes, buffers.arena_bytes)),
225 pipelines: (0..model.circuit.nodes.len())
226 .map(|n| one.planned.pipelines.facts(n))
227 .collect(),
228 };
229 Ok(crate::display::display(&model.circuit, plan, &info, opts)?)
230}
231
232fn heads_a_node(r: &crate::rules::Rule, c: &crate::ir::Circuit) -> bool {
235 r.pattern.first().is_some_and(|p| {
236 c.nodes.iter().any(|n| match n.op {
237 crate::ir::OpKind::Linear(role) if !p.roles.is_empty() => p.roles.contains(&role),
238 op => op == p.op,
239 })
240 })
241}
242
243pub const PREFILL_TOKENS: [u64; 2] = [4096, 32768];
245
246pub fn build_report(
248 registry: &Registry,
249 device_id: &str,
250 tree: &dyn KernelTree,
251 model: ModelUnderPlan,
252 command: String,
253) -> Result<HwReport, HwError> {
254 let resolved = plan::resolve(registry, device_id, tree, &model)?;
255 let (policy, class_settings) = model::policy_on_class(
256 &model.policy,
257 model.settings_class.as_deref(),
258 &resolved.chain[0],
259 )?;
260 let header =
261 crate::render::with_settings(&model::header_on(&model, &resolved.device.class), &policy);
262 let c = &model.circuit;
263 let mut tables = Vec::new();
264 let mut tensor_core = Vec::new();
265 for run in plan::report_runs() {
266 let planned = plan::fuse_on(&resolved, c, &policy, run)?;
267 tensor_core.extend(tc_policy::enforce(&resolved, c, &planned)?);
268 tables.push(gaps::gap_table(
269 &resolved,
270 c,
271 &policy.settings,
272 &model.label,
273 planned,
274 )?);
275 }
276 let routes = render_routes::route_rows(&resolved, c, &policy, &model.label, &tables)?;
277 let rf = |n| resolved.roofline_of(n);
278 let prefill = PREFILL_TOKENS
279 .iter()
280 .map(|&t| Ok((t, estimate::prefill_us(c, &policy.settings, &rf, t)?)))
281 .collect::<Result<Vec<_>, HwError>>()?;
282 let footprint = estimate::footprint(c, &policy.settings).map_err(HwError::Model)?;
283 let weight_floor = estimate::weight_floor_bytes(c).map_err(HwError::Model)?;
284 let mut exec = BTreeMap::new();
285 for (i, n) in c.nodes.iter().enumerate() {
286 let (Some(w), Some(a), Some(e)) =
287 (n.weight, estimate::activation_of(c, n), resolved.exec[i])
288 else {
289 continue;
290 };
291 *exec
292 .entry((exec::pair_name(w, a), w.name(), a.name(), e))
293 .or_insert(0) += 1;
294 }
295 let absent = resolved
296 .rule_list()
297 .iter()
298 .filter(|r| {
299 r.when
300 .iter()
301 .all(|(k, v)| policy.settings.get(k) == Some(v))
302 && heads_a_node(r, c)
303 })
304 .flat_map(|r| r.kernels.iter())
305 .filter_map(|k| {
306 resolved
307 .availability
308 .absent
309 .get(k)
310 .map(|a| (k.clone(), a.clone()))
311 })
312 .collect::<BTreeMap<_, _>>()
313 .into_iter()
314 .collect();
315 Ok(HwReport {
316 model,
317 resolved,
318 policy,
319 class_settings,
320 header,
321 tables,
322 tensor_core,
323 routes,
324 prefill,
325 footprint,
326 weight_floor,
327 exec,
328 absent,
329 others: registry
330 .devices
331 .iter()
332 .map(|d| (d.id.clone(), d.memory_bytes * d.usable_fraction))
333 .collect(),
334 command,
335 })
336}