1pub mod checkpoint;
14pub mod circuit_toml;
15pub mod config_map;
16pub mod declared_precision;
17pub mod digest;
18pub mod dims;
19pub mod display;
20pub mod format;
21pub mod fuser;
22pub mod hardware;
23pub mod instances;
24pub mod instantiate;
25pub mod ir;
26pub mod memory;
27pub mod model_buffer;
28pub mod pipeline;
29pub mod planner;
30pub mod precision;
31pub mod precision_policy;
32pub mod render;
33pub mod rules;
34pub mod runs;
35pub mod runtime;
36pub mod state;
37pub mod state_ops;
38pub mod venn;
39
40#[cfg(test)]
41mod test_toy;
42
43pub use checkpoint::{
44 CheckpointError, QuantMetadata, ResolvedCheckpoint, ServePrecision,
45 instantiate_from_checkpoint, map_checkpoint, recurrent_states, resolve_checkpoint,
46};
47pub use circuit_toml::{CircuitError, includes_of};
48pub use format::{Format, Scale};
49pub use fuser::{
50 AvailableKernels, EdgeState, FuseError, FusionPlan, Group, Policy, fuse, fuse_table,
51};
52pub use instances::{Instance, InstanceError, PrecisionSpec, parse_instances};
53pub use instantiate::instantiate;
54pub use ir::{ArchShape, Circuit, LayerKind, LinearRole, OpKind, Section};
55pub use precision::{EdgePrecision, LinearFormats, PrecisionError, PrecisionTable};
56pub use rules::{KernelId, Mode, Numerics, Rule, RuleError, parse_rules};
57pub use runs::{RowTable, VerifyRun};
58pub use runtime::{RuleSet, RuntimeRoute, parse_rule_set};
59
60#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
62pub enum LoadError {
63 #[error(transparent)]
65 Circuit(#[from] CircuitError),
66 #[error(transparent)]
68 Rules(#[from] RuleError),
69 #[error(transparent)]
71 Precision(#[from] PrecisionError),
72 #[error(transparent)]
74 Fuse(#[from] FuseError),
75 #[error(transparent)]
77 Display(#[from] display::DisplayError),
78 #[error(transparent)]
80 Plan(#[from] planner::PlanError),
81 #[error(transparent)]
84 Pipeline(#[from] pipeline::PipelineError),
85 #[error("precision table is for `{table}`, the instance serves `{instance}`")]
87 CheckpointMismatch {
88 table: String,
90 instance: String,
92 },
93}
94
95#[derive(Debug, Clone, Copy)]
97pub struct Sources<'a> {
98 pub circuit: &'a str,
100 pub precision: &'a str,
103 pub rules: &'a str,
105 pub blocks: &'a [(&'a str, &'a str)],
108}
109
110#[derive(Debug, Clone)]
112pub struct Loaded {
113 pub circuit: Circuit,
115 pub rules: Vec<Rule>,
117 pub runtime: Vec<RuntimeRoute>,
119 pub rules_digest: String,
121}
122
123pub fn load(instance: &Instance, src: Sources<'_>) -> Result<Loaded, LoadError> {
125 let mismatch = |table: String| LoadError::CheckpointMismatch {
126 table,
127 instance: instance.checkpoint.clone(),
128 };
129 match &instance.precision {
130 PrecisionSpec::Table(_) => {
131 let table = PrecisionTable::parse(src.precision)?;
132 if table.checkpoint != instance.checkpoint {
133 return Err(mismatch(table.checkpoint));
134 }
135 load_with(instance, src, &table)
136 }
137 PrecisionSpec::Policy {
138 tier, caps, engine, ..
139 } => {
140 let plan = precision_policy::CheckpointPlan::parse(src.precision)?;
141 if plan.checkpoint != instance.checkpoint {
142 return Err(mismatch(plan.checkpoint));
143 }
144 let policy = metrale_config::WeightQuantPolicy::new(
145 precision_policy::tier_named(tier)?,
146 &plan.plan,
147 precision_policy::caps_named(caps)?,
148 );
149 load_with(
150 instance,
151 src,
152 &precision_policy::PolicyPrecision::new(policy, engine),
153 )
154 }
155 }
156}
157
158pub fn load_with(
162 instance: &Instance,
163 src: Sources<'_>,
164 precision: &dyn EdgePrecision,
165) -> Result<Loaded, LoadError> {
166 let circuit = instantiate(src.circuit, src.blocks, &instance.shape, precision)?;
167 let RuleSet { rules, runtime } = parse_rule_set(src.rules)?;
168 Ok(Loaded {
169 circuit,
170 rules,
171 runtime,
172 rules_digest: digest::rules_digest(src.rules),
173 })
174}
175
176pub fn header(instance: &Instance) -> render::Header {
178 let mut h = vec![
179 ("recipe".to_string(), instance.recipe.clone()),
180 ("checkpoint".to_string(), instance.checkpoint.clone()),
181 ("target".to_string(), instance.target.clone()),
182 ];
183 let settings: Vec<String> = instance
184 .policy
185 .settings
186 .iter()
187 .map(|(k, v)| format!("{k}={v}"))
188 .collect();
189 h.push(("settings".to_string(), settings.join(" ")));
190 let levers: Vec<&str> = instance
191 .policy
192 .opt_in_levers
193 .iter()
194 .map(String::as_str)
195 .collect();
196 let levers = if levers.is_empty() {
197 "none".to_string()
198 } else {
199 levers.join(" ")
200 };
201 h.push(("opt-in levers".to_string(), levers));
202 h
203}
204
205pub fn render_plan(
210 instance: &Instance,
211 loaded: &Loaded,
212 available: &AvailableKernels,
213 mode: Mode,
214 rows: u64,
215 families: &venn::Families,
216) -> Result<String, LoadError> {
217 let plan = fuse(
218 &loaded.circuit,
219 &loaded.rules,
220 available,
221 &instance.policy,
222 mode,
223 rows,
224 )?;
225 let head = header(instance);
226 let c = &loaded.circuit;
227 let group = |g: &fuser::Group| hardware::tc_policy::unit_tag(families, &g.kernels);
228 let pipes = checked(loaded, &plan, &instance.policy, families)?;
229 let node = |n: ir::NodeIdx| pipes.line(n);
230 let notes = render::Notes {
231 group: &group,
232 node: &node,
233 };
234 let mut text = render::render_noted(c, &plan, &head, ¬es);
235 let set = (loaded.rules.as_slice(), loaded.runtime.as_slice());
236 for (route, arm) in runtime::route_arms(c, set, available, &instance.policy, &plan)? {
237 let policy = route.policy(&instance.policy);
238 let h = render::with_settings(&head, &policy);
239 let arm_pipes = checked(loaded, &arm, &policy, families)?;
240 let node = |n: ir::NodeIdx| arm_pipes.line(n);
241 let notes = render::Notes {
242 group: &group,
243 node: &node,
244 };
245 text.push_str(&render::route_section(c, &route, &arm, &h, ¬es));
246 }
247 Ok(text)
248}
249
250fn checked(
253 loaded: &Loaded,
254 plan: &FusionPlan,
255 policy: &Policy,
256 families: &venn::Families,
257) -> Result<pipeline::PlanPipelines, LoadError> {
258 Ok(pipeline::check_plan(
259 &loaded.circuit,
260 plan,
261 &loaded.rules,
262 families,
263 &policy.settings,
264 )?)
265}
266
267pub fn render_table_plan(
270 instance: &Instance,
271 loaded: &Loaded,
272 available: &AvailableKernels,
273 table: &RowTable,
274 families: &venn::Families,
275) -> Result<String, LoadError> {
276 let plan = fuse_table(
277 &loaded.circuit,
278 &loaded.rules,
279 available,
280 &instance.policy,
281 table,
282 )?;
283 let pipes = checked(loaded, &plan, &instance.policy, families)?;
284 let group = |g: &fuser::Group| hardware::tc_policy::unit_tag(families, &g.kernels);
285 let node = |n: ir::NodeIdx| pipes.line(n);
286 let notes = render::Notes {
287 group: &group,
288 node: &node,
289 };
290 Ok(render::render_noted(
291 &loaded.circuit,
292 &plan,
293 &header(instance),
294 ¬es,
295 ))
296}
297
298pub fn display_plan(
302 instance: &Instance,
303 loaded: &Loaded,
304 available: &AvailableKernels,
305 (mode, rows): (Mode, u64),
306 opts: &display::DisplayOpts,
307 families: &venn::Families,
308) -> Result<display::Document, LoadError> {
309 let plan = fuse(
310 &loaded.circuit,
311 &loaded.rules,
312 available,
313 &instance.policy,
314 mode,
315 rows,
316 )?;
317 let buffers = planner::plan_buffers(&loaded.circuit, &plan, rows)?;
318 let pipes = checked(loaded, &plan, &instance.policy, families)?;
319 let info = display::DisplayInfo {
320 checkpoint: instance.checkpoint.clone(),
321 recipe: instance.recipe.clone(),
322 bytes: Some((buffers.materialized_bytes, buffers.arena_bytes)),
323 pipelines: (0..loaded.circuit.nodes.len())
324 .map(|n| pipes.facts(n))
325 .collect(),
326 };
327 Ok(display::display(&loaded.circuit, &plan, &info, opts)?)
328}