1use std::collections::{BTreeMap, BTreeSet};
18
19use serde::Deserialize;
20
21use crate::format::Format;
22use crate::fuser::Policy;
23use crate::ir::{ArchShape, LayerKind};
24use crate::precision::LinearFormats;
25use crate::rules::Mode;
26
27#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct Instance {
30 pub recipe: String,
32 pub checkpoint: String,
34 pub arch: String,
36 pub variant: Option<String>,
40 pub precision: PrecisionSpec,
42 pub target: String,
44 pub golden: bool,
46 pub shape: ArchShape,
48 pub policy: Policy,
50 pub plans: BTreeMap<Mode, Vec<u64>>,
52 pub verify_batch: Vec<crate::runs::RowTable>,
55}
56
57impl Instance {
58 pub fn plan_file(&self, mode: Mode, rows: u64) -> String {
60 format!("{}-{}-n{rows}.txt", self.plan_stem(), mode.name())
61 }
62
63 pub fn table_plan_file(&self, table: &crate::runs::RowTable) -> String {
67 let runs: Vec<String> = table
68 .runs
69 .iter()
70 .map(|r| format!("{}x{}{}", r.k, r.n, if r.contiguous { "" } else { "f" }))
71 .collect();
72 format!(
73 "{}-verify_batch-{}{}.txt",
74 self.plan_stem(),
75 if table.carried { "" } else { "u-" },
76 runs.join("_")
77 )
78 }
79
80 pub fn plan_stem(&self) -> String {
83 match &self.variant {
84 Some(v) => format!("{}.{v}", self.arch),
85 None => self.arch.clone(),
86 }
87 }
88}
89
90#[derive(Debug, Clone, PartialEq, Eq)]
92pub enum PrecisionSpec {
93 Table(String),
95 Policy {
98 checkpoint_plan: String,
100 tier: String,
102 caps: Vec<String>,
104 engine: Vec<(String, LinearFormats)>,
106 },
107}
108
109impl PrecisionSpec {
110 pub fn file_name(&self) -> &str {
112 match self {
113 PrecisionSpec::Table(n) => n,
114 PrecisionSpec::Policy {
115 checkpoint_plan, ..
116 } => checkpoint_plan,
117 }
118 }
119}
120
121#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
123pub enum InstanceError {
124 #[error("INSTANCES.toml: {0}")]
126 Parse(String),
127 #[error("instance `{recipe}`: {detail}")]
129 Field {
130 recipe: String,
132 detail: String,
134 },
135}
136
137#[derive(Deserialize)]
138#[serde(deny_unknown_fields)]
139struct File {
140 schema: u32,
141 instance: Vec<InstanceFile>,
142}
143
144#[derive(Deserialize)]
145#[serde(deny_unknown_fields)]
146struct InstanceFile {
147 recipe: String,
148 checkpoint: String,
149 arch: String,
150 variant: Option<String>,
151 precision: PrecisionFile,
152 target: String,
153 golden: bool,
154 layer_kinds: String,
155 dims: BTreeMap<String, u64>,
156 policy: PolicyFile,
157 plans: BTreeMap<String, Vec<u64>>,
158 #[serde(default)]
159 verify_batch: Vec<String>,
160}
161
162#[derive(Deserialize)]
163#[serde(untagged)]
164enum PrecisionFile {
165 Table(String),
166 Policy(PolicyPrecisionFile),
167}
168
169#[derive(Deserialize)]
170#[serde(deny_unknown_fields)]
171struct PolicyPrecisionFile {
172 checkpoint_plan: String,
173 tier: String,
174 caps: Vec<String>,
175 engine: Vec<EngineFormatFile>,
176}
177
178#[derive(Deserialize)]
179#[serde(deny_unknown_fields)]
180struct EngineFormatFile {
181 #[serde(rename = "match")]
182 pattern: String,
183 weight: String,
184 activation: String,
185 why: String,
187}
188
189#[derive(Deserialize)]
190#[serde(deny_unknown_fields)]
191struct PolicyFile {
192 opt_in_levers: Vec<String>,
193 settings: BTreeMap<String, String>,
194}
195
196pub fn parse_instances(text: &str) -> Result<Vec<Instance>, InstanceError> {
198 let file: File = toml::from_str(text).map_err(|e| InstanceError::Parse(e.to_string()))?;
199 if file.schema != 1 {
200 return Err(InstanceError::Parse(format!(
201 "schema {} (this build reads 1)",
202 file.schema
203 )));
204 }
205 let mut seen = BTreeSet::new();
206 let mut stems = BTreeSet::new();
207 let mut out = Vec::with_capacity(file.instance.len());
208 for f in file.instance {
209 let field = |detail: String| InstanceError::Field {
210 recipe: f.recipe.clone(),
211 detail,
212 };
213 if !seen.insert(f.recipe.clone()) {
214 return Err(field("listed twice".into()));
215 }
216 if f.variant.as_deref().is_some_and(|v| v.is_empty()) {
217 return Err(field("`variant` is empty".into()));
218 }
219 let stem = (f.arch.clone(), f.variant.clone());
220 if f.golden && !stems.insert(stem) {
221 return Err(field(format!(
222 "a second golden `{}` instance without its own `variant`: the golden plan files \
223 would collide",
224 f.arch
225 )));
226 }
227 let mut layer_kinds = Vec::new();
228 for c in f.layer_kinds.chars().filter(|c| !c.is_whitespace()) {
229 layer_kinds.push(match c {
230 'G' => LayerKind::LinearAttention,
231 'A' => LayerKind::FullAttention,
232 'M' => LayerKind::Mamba,
233 'E' => LayerKind::Moe,
234 other => return Err(field(format!("layer kind `{other}` is not G, A, M or E"))),
235 });
236 }
237 let mut plans = BTreeMap::new();
238 for (m, rows) in &f.plans {
239 let mode = Mode::parse(m).ok_or_else(|| field(format!("unknown mode `{m}`")))?;
240 if rows.is_empty() || rows.contains(&0) {
241 return Err(field(format!("mode `{m}` needs row counts of at least 1")));
242 }
243 plans.insert(mode, rows.clone());
244 }
245 let verify_batch = f
246 .verify_batch
247 .iter()
248 .map(|t| {
249 crate::runs::RowTable::parse(t).map_err(|e| field(format!("verify_batch: {e}")))
250 })
251 .collect::<Result<Vec<_>, _>>()?;
252 let precision = match f.precision {
253 PrecisionFile::Table(name) => PrecisionSpec::Table(name),
254 PrecisionFile::Policy(p) => {
255 let mut engine = Vec::with_capacity(p.engine.len());
256 for e in p.engine {
257 if e.why.trim().is_empty() {
258 return Err(field(format!(
259 "engine format `{}` states no `why`",
260 e.pattern
261 )));
262 }
263 let fmt = |s: &str| {
264 Format::parse(s)
265 .map_err(|err| field(format!("engine format `{}`: {err}", e.pattern)))
266 };
267 let formats = LinearFormats {
268 weight: fmt(&e.weight)?,
269 activation: fmt(&e.activation)?,
270 };
271 engine.push((e.pattern, formats));
272 }
273 PrecisionSpec::Policy {
274 checkpoint_plan: p.checkpoint_plan,
275 tier: p.tier,
276 caps: p.caps,
277 engine,
278 }
279 }
280 };
281 out.push(Instance {
282 recipe: f.recipe.clone(),
283 checkpoint: f.checkpoint,
284 arch: f.arch,
285 variant: f.variant,
286 precision,
287 target: f.target,
288 golden: f.golden,
289 shape: ArchShape {
290 layer_kinds,
291 dims: f.dims,
292 },
293 policy: Policy {
294 opt_in_levers: f.policy.opt_in_levers.into_iter().collect(),
295 settings: f.policy.settings,
296 },
297 plans,
298 verify_batch,
299 });
300 }
301 Ok(out)
302}