metrale_circuit/hardware/
model.rs1use std::collections::BTreeMap;
19
20use super::HwError;
21use super::class::ClassInfo;
22use crate::fuser::Policy;
23use crate::instances::{Instance, PrecisionSpec, parse_instances};
24use crate::ir::Circuit;
25use crate::render::Header;
26use crate::venn::repo::{Repo, load_instance, resolve};
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum PrecisionChoice {
31 Recipe,
33 Declared,
35}
36
37impl PrecisionChoice {
38 pub fn name(self) -> &'static str {
40 match self {
41 PrecisionChoice::Recipe => "recipe",
42 PrecisionChoice::Declared => "declared",
43 }
44 }
45}
46
47#[derive(Debug, Clone, Copy)]
49pub struct ModelSpec<'a> {
50 pub checkpoint: &'a str,
52 pub config_json: Option<&'a str>,
54 pub hf_quant: Option<&'a str>,
56 pub precision: PrecisionChoice,
58}
59
60#[derive(Debug, Clone)]
62pub struct ModelUnderPlan {
63 pub label: String,
65 pub checkpoint: String,
67 pub kernel_model: String,
69 pub kernel_quant: String,
71 pub circuit: Circuit,
73 pub policy: Policy,
75 pub header: Header,
77 pub precision: String,
79 pub precision_choice: PrecisionChoice,
81 pub policy_sources: Vec<(String, String)>,
84 pub settings_class: Option<String>,
88}
89
90pub trait CircuitSource {
92 fn model(&self, spec: &ModelSpec<'_>) -> Result<ModelUnderPlan, HwError>;
94}
95
96pub struct InstancesSource<'a> {
98 pub repo: &'a dyn Repo,
100}
101
102impl CircuitSource for InstancesSource<'_> {
103 fn model(&self, spec: &ModelSpec<'_>) -> Result<ModelUnderPlan, HwError> {
104 let text = self
105 .repo
106 .read("kernels/circuits/INSTANCES.toml")
107 .map_err(HwError::Model)?;
108 let all = parse_instances(&text).map_err(|e| HwError::Model(e.to_string()))?;
109 let inst = resolve(&all, spec.checkpoint).map_err(|e| HwError::Model(e.to_string()))?;
110 let (inst, precision) = match (spec.precision, &inst.precision) {
111 (PrecisionChoice::Recipe, p) => {
112 let line = format!("recipe {} ({})", inst.recipe, describe(p));
113 (inst, line)
114 }
115 (PrecisionChoice::Declared, PrecisionSpec::Policy { .. }) => {
116 let mut i = inst.clone();
117 if let PrecisionSpec::Policy { tier, .. } = &mut i.precision {
118 *tier = "declared".into();
119 }
120 let line = format!("declared ({})", describe(&i.precision));
121 (i, line)
122 }
123 (PrecisionChoice::Declared, PrecisionSpec::Table(t)) => {
124 let line = format!("declared, as stated by the precision table `{t}`");
125 (inst, line)
126 }
127 };
128 model_of(self.repo, &inst, precision, spec.precision)
129 }
130}
131
132fn describe(p: &PrecisionSpec) -> String {
133 match p {
134 PrecisionSpec::Table(t) => format!("precision table `{t}`"),
135 PrecisionSpec::Policy {
136 checkpoint_plan,
137 tier,
138 ..
139 } => format!("checkpoint plan `{checkpoint_plan}` at tier `{tier}`"),
140 }
141}
142
143pub fn model_of(
145 repo: &dyn Repo,
146 inst: &Instance,
147 precision: String,
148 precision_choice: PrecisionChoice,
149) -> Result<ModelUnderPlan, HwError> {
150 let loaded = load_instance(repo, inst).map_err(|e| HwError::Model(e.to_string()))?;
151 let mut parts = inst.target.split('/');
152 let (hw, model, quant) = (parts.next(), parts.next(), parts.next());
153 let (Some(hw), Some(model), Some(quant)) = (hw, model, quant) else {
154 return Err(HwError::Model(format!(
155 "{}: target `{}` is not hw/model/quant",
156 inst.recipe, inst.target
157 )));
158 };
159 Ok(ModelUnderPlan {
160 label: inst.recipe.clone(),
161 checkpoint: inst.checkpoint.clone(),
162 kernel_model: model.to_string(),
163 kernel_quant: quant.to_string(),
164 circuit: loaded.circuit,
165 policy: inst.policy.clone(),
166 header: crate::header(inst),
167 precision,
168 precision_choice,
169 policy_sources: Vec::new(),
170 settings_class: Some(hw.to_string()),
171 })
172}
173
174pub const CLASS_DEFAULT_SETTINGS: [(&str, &str); 3] = [
178 ("ssm_batched_recurrent", "ssm_batched_recurrent"),
179 ("ssm_ba_gates_hopper", "ssm_ba_gates_hopper"),
180 ("decode_split_silu", "decode_split_silu"),
181];
182
183pub fn policy_on_class(
187 policy: &Policy,
188 settings_class: Option<&str>,
189 class: &ClassInfo,
190) -> Result<(Policy, BTreeMap<String, String>), HwError> {
191 let own = settings_class == Some(class.name.as_str());
192 let mut out = policy.clone();
193 let mut changed = BTreeMap::new();
194 for (setting, key) in CLASS_DEFAULT_SETTINGS {
195 if !policy.settings.contains_key(setting) {
196 continue;
197 }
198 let raw = class.defaults.get(key).ok_or_else(|| {
199 HwError::Class(format!(
200 "kernels/{}/HARDWARE.toml [defaults] does not state `{key}`",
201 class.name
202 ))
203 })?;
204 let v = match raw.as_str() {
205 "true" => "on",
206 "false" => "off",
207 other => other,
208 }
209 .to_string();
210 if own {
211 continue;
212 }
213 if policy.settings.get(setting) != Some(&v) {
214 changed.insert(setting.to_string(), v.clone());
215 }
216 out.settings.insert(setting.to_string(), v);
217 }
218 Ok((out, changed))
219}
220
221pub fn header_on(model: &ModelUnderPlan, class: &str) -> Header {
223 model
224 .header
225 .iter()
226 .map(|(k, v)| {
227 if k == "target" {
228 (
229 k.clone(),
230 format!("{class}/{}/{}", model.kernel_model, model.kernel_quant),
231 )
232 } else {
233 (k.clone(), v.clone())
234 }
235 })
236 .collect()
237}