metrale_circuit/hardware/
avail.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: Which kernels a device can run: a rule's kernel is available when the class
4//! compiles its module and names the function, no MODEL.toml declares it absent, the class's
5//! build defines do not compile its region out, and the device runs natively the instruction
6//! kind its source guard stands for (`kernels/DEVICES.toml` `[[guard]]`).
7//!
8//! Owner: metrale-circuit (hardware).
9//! Invariants:
10//! - A kernel's requirement is read from the source it compiles from: the preprocessor region
11//!   holding its `__global__` definition. No kernel is marked by hand.
12//! - The build and the device must agree: a class that compiles a guarded region out while one
13//!   of its devices claims the instruction, or keeps it while a device lacks it, is an error
14//!   ([`check_build`]), never resolved in either's favour.
15
16use 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/// 2026-09-30: Why a kernel is not available on the device.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub enum Absence {
32    /// 2026-09-30: The class does not compile the module, or the module does not name it.
33    NotCompiled,
34    /// 2026-09-30: A MODEL.toml `[expected_absent]` entry, with its reason.
35    ExpectedAbsent(String),
36    /// 2026-09-30: The class's build flags define the macro that compiles its region out.
37    CompiledOut {
38        /// 2026-09-30: The macro.
39        macro_name: String,
40        /// 2026-09-30: What the region needs.
41        requires: Instr,
42    },
43    /// 2026-09-30: The device does not run the region's instruction kind.
44    NotNative(Instr),
45}
46
47impl Absence {
48    /// 2026-09-30: One line for the report.
49    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/// 2026-09-30: The kernels of a rule set on one device.
63#[derive(Debug, Clone, Default, PartialEq, Eq)]
64pub struct Availability {
65    /// 2026-09-30: What the fuser may select.
66    pub kernels: AvailableKernels,
67    /// 2026-09-30: Every rule kernel that is not available, and why.
68    pub absent: BTreeMap<KernelId, Absence>,
69    /// 2026-09-30: The instruction kind each guarded kernel needs.
70    pub requires: BTreeMap<KernelId, Instr>,
71}
72
73/// 2026-09-30: Fail when the class's build defines contradict the device's instruction kinds.
74pub 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
91/// 2026-09-30: `word` appears in `text` as a whole identifier.
92pub 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
99/// 2026-09-30: The guard whose region holds the `__global__` definition of `func` in `text`,
100/// when every definition sits in one.
101pub 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, &macros) {
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
132/// 2026-10-01: The function-like macros of `lines` whose body defines a `__global__` kernel
133/// named by a parameter (`#define ENTRY(NAME, ...) ... __global__ ... void NAME(...)`), with that
134/// parameter's position. An entry point defined by expanding one sits where the macro is
135/// invoked, so its guard is the invocation's region.
136fn 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
176/// 2026-10-01: `line` invokes one of `macros` with `func` as its kernel-name argument.
177fn 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
186/// 2026-09-30: Per line, the index of the guard whose requiring branch holds it.
187fn regions(lines: &[&str], guards: &[Guard]) -> Vec<Option<usize>> {
188    // 2026-09-30: One frame per open `#if`: the guard it tests and whether the current branch
189    // is the one that needs the guard's instruction.
190    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
244/// 2026-09-30: Whether `device` can run kernel `k` from `sources` of `class`, and the
245/// instruction kind its source guard needs (`None`: unguarded or not compiled).
246pub 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
280/// 2026-09-30: The rule kernels `device` can run from `sources` of `class`.
281pub 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
312/// 2026-10-01: The nodes of `c` that multiply an NVFP4 activation ([`nvfp4_mma`]) and that no
313/// compiled FP4 block-scale kernel of the class runs: no rule with a pattern element that fits
314/// the node has every kernel available and one that needs the FP4 block-scale MMA, and no family
315/// that implements the node has such a kernel available. `kernel` answers for one kernel: `None`
316/// when the device cannot run it, else whether its source guard needs the FP4 block-scale MMA.
317/// Row ranges and policy conditions are not consulted: they choose among the kernels a class
318/// compiles, not which kernels it compiles.
319pub 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}