metrale_circuit/hardware/
avail.rs1use std::collections::{BTreeMap, BTreeSet};
17
18use super::HwError;
19use super::class::ClassInfo;
20use super::device::{Device, Guard, Instr, Polarity};
21use super::sources::ClassSources;
22use crate::fuser::{AvailableKernels, pattern_fits};
23use crate::ir::{Circuit, NodeIdx};
24use crate::rules::{KernelId, PatternOp, Rule};
25use crate::venn::Families;
26use crate::venn::classify::implements;
27use crate::venn::roofline::nvfp4_mma;
28
29#[derive(Debug, Clone, PartialEq, Eq)]
31pub enum Absence {
32 NotCompiled,
34 ExpectedAbsent(String),
36 CompiledOut {
38 macro_name: String,
40 requires: Instr,
42 },
43 NotNative(Instr),
45}
46
47impl Absence {
48 pub fn describe(&self) -> String {
50 match self {
51 Absence::NotCompiled => "not compiled for this class".into(),
52 Absence::ExpectedAbsent(why) => format!("expected absent: {why}"),
53 Absence::CompiledOut {
54 macro_name,
55 requires,
56 } => format!("compiled out by -D{macro_name} (needs {})", requires.name()),
57 Absence::NotNative(i) => format!("needs {}, which the device lacks", i.name()),
58 }
59 }
60}
61
62#[derive(Debug, Clone, Default, PartialEq, Eq)]
64pub struct Availability {
65 pub kernels: AvailableKernels,
67 pub absent: BTreeMap<KernelId, Absence>,
69 pub requires: BTreeMap<KernelId, Instr>,
71}
72
73pub fn check_build(device: &Device, class: &ClassInfo, guards: &[Guard]) -> Result<(), HwError> {
75 for g in guards.iter().filter(|g| g.polarity == Polarity::Ifndef) {
76 let defined = class.defines.contains(&g.macro_name);
77 let native = device.native.contains(&g.requires);
78 if defined == native {
79 return Err(HwError::BuildContradictsDevice {
80 device: device.id.clone(),
81 class: class.name.clone(),
82 macro_name: g.macro_name.clone(),
83 requires: g.requires.name(),
84 defined,
85 });
86 }
87 }
88 Ok(())
89}
90
91pub fn names_word(text: &str, word: &str) -> bool {
93 text.match_indices(word).any(|(at, _)| {
94 let ok = |c: Option<char>| c.is_none_or(|c| !(c.is_ascii_alphanumeric() || c == '_'));
95 ok(text[..at].chars().next_back()) && ok(text[at + word.len()..].chars().next())
96 })
97}
98
99pub fn guard_of<'g>(text: &str, func: &str, guards: &'g [Guard]) -> Option<&'g Guard> {
102 let lines: Vec<&str> = text.lines().collect();
103 let regions = regions(&lines, guards);
104 let macros = entry_macros(&lines);
105 let mut found: Option<&Guard> = None;
106 let mut any = false;
107 for (i, line) in lines.iter().enumerate() {
108 let direct = defines_here(line, func)
109 && lines[i.saturating_sub(3)..=i]
110 .join(" ")
111 .contains("__global__");
112 if !direct && !instantiates(line, func, ¯os) {
113 continue;
114 }
115 any = true;
116 match regions[i] {
117 Some(g) => found = found.or(Some(&guards[g])),
118 None => return None,
119 }
120 }
121 if any { found } else { None }
122}
123
124fn defines_here(line: &str, func: &str) -> bool {
125 line.match_indices(func).any(|(at, _)| {
126 let before = line[..at].chars().next_back();
127 let after = line[at + func.len()..].trim_start();
128 before.is_none_or(|c| !(c.is_ascii_alphanumeric() || c == '_')) && after.starts_with('(')
129 })
130}
131
132fn entry_macros(lines: &[&str]) -> Vec<(String, usize)> {
137 let ident = |c: &char| c.is_ascii_alphanumeric() || *c == '_';
138 let mut out = Vec::new();
139 let mut i = 0;
140 while i < lines.len() {
141 let define = lines[i]
142 .trim_start()
143 .strip_prefix('#')
144 .map(str::trim_start)
145 .and_then(|r| r.strip_prefix("define"));
146 let mut body = define.unwrap_or_default().to_string();
147 while define.is_some() && body.trim_end().ends_with('\\') && i + 1 < lines.len() {
148 body.truncate(body.trim_end().len() - 1);
149 i += 1;
150 body.push(' ');
151 body.push_str(lines[i]);
152 }
153 i += 1;
154 let head = body.trim_start();
155 let name: String = head.chars().take_while(ident).collect();
156 let Some((params, rest)) = head[name.len()..]
157 .strip_prefix('(')
158 .and_then(|r| r.split_once(')'))
159 else {
160 continue;
161 };
162 if !rest.contains("__global__") {
163 continue;
164 }
165 if let Some(at) = params
166 .split(',')
167 .map(str::trim)
168 .position(|p| !p.is_empty() && defines_here(rest, p))
169 {
170 out.push((name, at));
171 }
172 }
173 out
174}
175
176fn instantiates(line: &str, func: &str, macros: &[(String, usize)]) -> bool {
178 let line = line.trim_start();
179 macros.iter().any(|(m, at)| {
180 line.strip_prefix(m.as_str())
181 .and_then(|r| r.trim_start().strip_prefix('('))
182 .is_some_and(|args| args.split([',', ')']).nth(*at).map(str::trim) == Some(func))
183 })
184}
185
186fn regions(lines: &[&str], guards: &[Guard]) -> Vec<Option<usize>> {
188 let mut stack: Vec<Option<(usize, bool)>> = Vec::new();
191 let mut out = Vec::with_capacity(lines.len());
192 for line in lines {
193 let d = line.trim_start();
194 let directive = d.strip_prefix('#').map(str::trim_start);
195 match directive {
196 Some(rest) if rest.starts_with("if") => stack.push(opening(rest, guards)),
197 Some(rest) if rest.starts_with("elif") || rest.starts_with("else") => {
198 if let Some(Some((_, needing))) = stack.last_mut() {
199 *needing = !*needing && rest.starts_with("else");
200 }
201 }
202 Some(rest) if rest.starts_with("endif") => {
203 stack.pop();
204 }
205 _ => {}
206 }
207 out.push(
208 stack
209 .iter()
210 .rev()
211 .find_map(|f| f.and_then(|(g, needing)| needing.then_some(g))),
212 );
213 }
214 out
215}
216
217fn opening(rest: &str, guards: &[Guard]) -> Option<(usize, bool)> {
218 let (neg, name) = if let Some(n) = rest.strip_prefix("ifndef") {
219 (true, n.trim())
220 } else if let Some(n) = rest.strip_prefix("ifdef") {
221 (false, n.trim())
222 } else {
223 let e = rest.strip_prefix("if")?.trim();
224 let (neg, e) = match e.strip_prefix('!') {
225 Some(x) => (true, x.trim_start()),
226 None => (false, e),
227 };
228 let e = e.strip_prefix("defined")?.trim();
229 let e = e.trim_start_matches('(').trim_end_matches(')').trim();
230 (neg, e)
231 };
232 let name: String = name
233 .chars()
234 .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
235 .collect();
236 let g = guards.iter().position(|g| g.macro_name == name)?;
237 let needing = match guards[g].polarity {
238 Polarity::Ifndef => neg,
239 Polarity::Ifdef => !neg,
240 };
241 Some((g, needing))
242}
243
244pub fn kernel_status(
247 device: &Device,
248 class: &ClassInfo,
249 sources: &ClassSources,
250 guards: &[Guard],
251 k: &KernelId,
252) -> (Result<(), Absence>, Option<Instr>) {
253 let text = sources.modules.get(&k.module).map(|m| m.text.as_str());
254 let Some(text) = text.filter(|t| names_word(t, &k.func)) else {
255 return (Err(Absence::NotCompiled), None);
256 };
257 if let Some(why) = sources
258 .expected_absent
259 .get(&k.module)
260 .and_then(|m| m.get(&k.func))
261 {
262 return (Err(Absence::ExpectedAbsent(why.clone())), None);
263 }
264 let Some(g) = guard_of(text, &k.func, guards) else {
265 return (Ok(()), None);
266 };
267 let status = if g.polarity == Polarity::Ifndef && class.defines.contains(&g.macro_name) {
268 Err(Absence::CompiledOut {
269 macro_name: g.macro_name.clone(),
270 requires: g.requires,
271 })
272 } else if !device.native.contains(&g.requires) {
273 Err(Absence::NotNative(g.requires))
274 } else {
275 Ok(())
276 };
277 (status, Some(g.requires))
278}
279
280pub fn availability(
282 device: &Device,
283 class: &ClassInfo,
284 sources: &ClassSources,
285 rules: &[Rule],
286 guards: &[Guard],
287) -> Availability {
288 let mut out = Availability::default();
289 for r in rules {
290 out.kernels.caps.extend(r.requires.iter().cloned());
291 for k in &r.kernels {
292 if out.kernels.kernels.contains(k) || out.absent.contains_key(k) {
293 continue;
294 }
295 let (status, req) = kernel_status(device, class, sources, guards, k);
296 if let Some(i) = req {
297 out.requires.insert(k.clone(), i);
298 }
299 match status {
300 Ok(()) => {
301 out.kernels.kernels.insert(k.clone());
302 }
303 Err(a) => {
304 out.absent.insert(k.clone(), a);
305 }
306 }
307 }
308 }
309 out
310}
311
312pub fn without_fp4_kernel(
320 c: &Circuit,
321 rules: &[Rule],
322 families: &Families,
323 kernel: &dyn Fn(&KernelId) -> Option<bool>,
324) -> BTreeSet<NodeIdx> {
325 let fp4 = |k: &KernelId| kernel(k) == Some(true);
326 let patterns: Vec<&PatternOp> = rules
327 .iter()
328 .filter(|r| r.kernels.iter().all(|k| kernel(k).is_some()) && r.kernels.iter().any(fp4))
329 .flat_map(|r| &r.pattern)
330 .collect();
331 let fams: Vec<_> = families
332 .families
333 .iter()
334 .filter(|f| f.kernels.iter().any(fp4))
335 .collect();
336 c.nodes
337 .iter()
338 .enumerate()
339 .filter(|(_, n)| nvfp4_mma(c, n))
340 .filter(|(_, n)| {
341 !patterns.iter().any(|p| pattern_fits(c, p, n))
342 && !fams.iter().any(|f| implements(f, c, n))
343 })
344 .map(|(i, _)| i)
345 .collect()
346}