metrale_circuit/hardware/
tc_policy.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The tensor-core policy of a kernel class (`kernels/<class>/HARDWARE.toml`
4//! `[tensor_core_policy]`): which ops, at which modes, row counts and weight formats, a plan MUST
5//! run on a tensor-core kernel, and the stated exemptions. [`audit`] checks one fused plan
6//! against it from the compute units the class's kernel families declare (`venn::compute`).
7//!
8//! Owner: metrale-circuit (hardware).
9//! Invariants:
10//! - Pure: the policy arrives as text, the plan and families as values.
11//! - Nothing defaults: a requirement states its ops, modes, minimum rows and weight formats
12//!   (`["any"]` is stated, not assumed); an exemption states its ops, modes, rows, the kernels it
13//!   allows, its kind and its reason, and a `measured` exemption names its evidence.
14//! - A class's policy is its own; it is not inherited along `HARDWARE.toml` `inherits`, since
15//!   another class's tensor cores and kernels differ. A class without one is reported as such.
16//! - A covered node whose group runs no tensor-core kernel is a violation unless an exemption
17//!   lists the op, the mode, the row count and every kernel of the group. A covered node in a
18//!   placeholder group (no kernel of the class covers the op) is a gap, listed, never a pass.
19
20use std::collections::BTreeSet;
21
22use crate::format::Format;
23use crate::fuser::FusionPlan;
24use crate::ir::{Circuit, LinearRole, OpKind};
25use crate::rules::{KernelId, Mode};
26use crate::venn::compute::{ComputeUnit, group_unit};
27use crate::venn::families::Families;
28
29/// 2026-10-02: An op a policy entry names: a base name (`linear` is every role) or `linear:<role>`.
30#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
31pub struct OpMatch {
32    /// 2026-10-02: Base op name.
33    pub base: String,
34    /// 2026-10-02: For `linear:<role>`, the role.
35    pub role: Option<LinearRole>,
36}
37
38impl OpMatch {
39    fn parse(s: &str) -> Result<Self, String> {
40        let (base, role) = match s.split_once(':') {
41            Some(("linear", r)) => (
42                "linear",
43                Some(LinearRole::parse(r).ok_or_else(|| format!("unknown linear role `{r}`"))?),
44            ),
45            Some(_) => return Err(format!("op `{s}`: only `linear` takes a role")),
46            None => (s, None),
47        };
48        if base != "linear" {
49            OpKind::parse(base, None, None).map_err(|e| e.to_string())?;
50        }
51        Ok(OpMatch {
52            base: base.to_string(),
53            role,
54        })
55    }
56
57    fn matches(&self, op: &OpKind) -> bool {
58        op.base_name() == self.base
59            && match (op, self.role) {
60                (OpKind::Linear(r), Some(want)) => *r == want,
61                _ => true,
62            }
63    }
64}
65
66/// 2026-10-02: The weight formats an entry covers.
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub enum Weights {
69    /// 2026-10-02: Every format, and ops that read no weight (`["any"]`).
70    Any,
71    /// 2026-10-02: Only these.
72    Only(BTreeSet<Format>),
73}
74
75impl Weights {
76    fn covers(&self, w: Option<Format>) -> bool {
77        match self {
78            Weights::Any => true,
79            Weights::Only(set) => w.is_some_and(|w| set.contains(&w)),
80        }
81    }
82}
83
84/// 2026-10-02: Ops that must run on tensor cores at these modes, from `min_rows` rows.
85#[derive(Debug, Clone, PartialEq, Eq)]
86pub struct Require {
87    /// 2026-10-02: The ops.
88    pub ops: Vec<OpMatch>,
89    /// 2026-10-02: The modes.
90    pub modes: BTreeSet<Mode>,
91    /// 2026-10-02: The plan's row count from which the requirement applies.
92    pub min_rows: u64,
93    /// 2026-10-02: The weight formats.
94    pub weights: Weights,
95}
96
97/// 2026-10-02: Why an op may stay off tensor cores.
98#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
99pub enum ExemptKind {
100    /// 2026-10-02: Measured: the CUDA-core kernel beats the tensor-core tile there.
101    Measured,
102    /// 2026-10-02: The shape leaves no MMA tile to fill (e.g. an output of one column).
103    Shape,
104    /// 2026-10-02: A known violation with no tensor-core kernel yet: the backlog.
105    Backlog,
106}
107
108impl ExemptKind {
109    /// 2026-10-02: The spelling in the policy file.
110    pub fn name(self) -> &'static str {
111        match self {
112            ExemptKind::Measured => "measured",
113            ExemptKind::Shape => "shape",
114            ExemptKind::Backlog => "backlog",
115        }
116    }
117}
118
119/// 2026-10-02: A stated exemption.
120#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct Exempt {
122    /// 2026-10-02: The ops.
123    pub ops: Vec<OpMatch>,
124    /// 2026-10-02: The modes.
125    pub modes: BTreeSet<Mode>,
126    /// 2026-10-02: Inclusive rows.
127    pub rows: (u64, u64),
128    /// 2026-10-02: The kernels a non-tensor-core group may launch for those ops.
129    pub kernels: BTreeSet<KernelId>,
130    /// 2026-10-02: Its kind.
131    pub kind: ExemptKind,
132    /// 2026-10-02: Why.
133    pub reason: String,
134    /// 2026-10-02: The measurement, for `measured`.
135    pub evidence: Option<String>,
136}
137
138/// 2026-10-02: A class's policy.
139#[derive(Debug, Clone, PartialEq, Eq)]
140pub struct TcPolicy {
141    /// 2026-10-02: Requirements.
142    pub require: Vec<Require>,
143    /// 2026-10-02: Exemptions.
144    pub exempt: Vec<Exempt>,
145}
146
147#[path = "tc_policy_parse.rs"]
148mod parse;
149pub use parse::parse_policy;
150#[path = "tc_policy_lint.rs"]
151mod lint;
152pub use lint::lint_rules;
153
154/// 2026-10-02: One covered node that runs off tensor cores.
155#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
156pub struct Finding {
157    /// 2026-10-02: `block.node`.
158    pub site: String,
159    /// 2026-10-02: Op name (with its role).
160    pub op: String,
161    /// 2026-10-02: Declared weight format.
162    pub weight: Option<String>,
163    /// 2026-10-02: Mode.
164    pub mode: Mode,
165    /// 2026-10-02: The plan's rows.
166    pub rows: u64,
167    /// 2026-10-02: The group's kernels.
168    pub kernels: Vec<KernelId>,
169    /// 2026-10-02: The unit they run on (`host` for a group without kernels).
170    pub unit: String,
171}
172
173/// 2026-10-02: One plan's audit.
174#[derive(Debug, Clone, Default, PartialEq, Eq)]
175pub struct TcAudit {
176    /// 2026-10-02: Covered nodes.
177    pub covered: usize,
178    /// 2026-10-02: Covered nodes on tensor cores.
179    pub on_tensor_cores: usize,
180    /// 2026-10-02: Covered sites off tensor cores that an exemption allows, with its index.
181    pub exempted: Vec<(Finding, usize)>,
182    /// 2026-10-02: Covered sites off tensor cores that nothing allows.
183    pub violations: Vec<Finding>,
184    /// 2026-10-02: Covered sites no kernel of the class plans (placeholder groups): gaps, which
185    /// the gap table reports; listed here because the kernel they need is tensor-core work.
186    pub gaps: Vec<Finding>,
187}
188
189/// 2026-10-02: The unit a group's kernels run on, from the families; an error names a kernel
190/// no family declares. `None` for a group without kernels.
191pub fn unit_of_group(
192    families: &Families,
193    kernels: &[KernelId],
194) -> Result<Option<ComputeUnit>, String> {
195    let mut units = Vec::with_capacity(kernels.len());
196    for k in kernels {
197        units.push(families.compute_of(k).ok_or_else(|| {
198            format!("kernel `{k}` is in no kernel family, so its compute unit is undeclared")
199        })?);
200    }
201    Ok(group_unit(units))
202}
203
204/// 2026-10-02: A group's unit as a plan line shows it ([`ComputeUnit::tag`]): `None` for a group
205/// without kernels, `undeclared` when a kernel is in no family.
206pub fn unit_tag(families: &Families, kernels: &[KernelId]) -> Option<String> {
207    match unit_of_group(families, kernels) {
208        Ok(u) => u.map(|u| u.tag()),
209        Err(_) => Some("undeclared".to_string()),
210    }
211}
212
213/// 2026-10-02: Audit `plan` of `c` against `policy`, with the units `families` declare.
214/// Covered nodes in groups emitted by `skip_emitter` (the hardware planner's placeholders) are
215/// gaps, not violations.
216pub fn audit(
217    c: &Circuit,
218    plan: &FusionPlan,
219    families: &Families,
220    policy: &TcPolicy,
221    skip_emitter: &str,
222) -> Result<TcAudit, String> {
223    let (mode, rows) = (plan.mode, plan.rows);
224    let mut out = TcAudit::default();
225    let mut seen = BTreeSet::new();
226    for g in &plan.groups {
227        let placeholder = g.emitter == skip_emitter;
228        let mut unit: Option<Option<ComputeUnit>> = None;
229        for &n in &g.nodes {
230            let node = &c.nodes[n];
231            let covered = policy.require.iter().any(|r| {
232                r.modes.contains(&mode)
233                    && rows >= r.min_rows
234                    && r.weights.covers(node.weight)
235                    && r.ops.iter().any(|o| o.matches(&node.op))
236            });
237            if !covered {
238                continue;
239            }
240            if placeholder {
241                let f = Finding {
242                    site: format!("{}.{}", node.block, node.local),
243                    op: node.op.name(),
244                    weight: node.weight.map(|w| w.name()),
245                    mode,
246                    rows,
247                    kernels: Vec::new(),
248                    unit: "no planned kernel".to_string(),
249                };
250                if seen.insert(f.clone()) {
251                    out.gaps.push(f);
252                }
253                continue;
254            }
255            out.covered += 1;
256            let u = match &unit {
257                Some(u) => u.clone(),
258                None => {
259                    let u = unit_of_group(families, &g.kernels)?;
260                    unit = Some(u.clone());
261                    u
262                }
263            };
264            if u.as_ref().is_some_and(ComputeUnit::is_tensor_core) {
265                out.on_tensor_cores += 1;
266                continue;
267            }
268            let f = Finding {
269                site: format!("{}.{}", node.block, node.local),
270                op: node.op.name(),
271                weight: node.weight.map(|w| w.name()),
272                mode,
273                rows,
274                kernels: g.kernels.clone(),
275                unit: u.map_or_else(|| "host".to_string(), |u| u.name()),
276            };
277            if !seen.insert(f.clone()) {
278                continue;
279            }
280            let allowed = policy.exempt.iter().position(|e| {
281                e.modes.contains(&mode)
282                    && (e.rows.0..=e.rows.1).contains(&rows)
283                    && e.ops.iter().any(|o| o.matches(&node.op))
284                    && !g.kernels.is_empty()
285                    && g.kernels.iter().all(|k| e.kernels.contains(k))
286            });
287            match allowed {
288                Some(i) => out.exempted.push((f, i)),
289                None => out.violations.push(f),
290            }
291        }
292    }
293    Ok(out)
294}
295
296/// 2026-10-02: One line per violation, for an error.
297pub fn describe(violations: &[Finding]) -> String {
298    violations
299        .iter()
300        .map(|f| {
301            let k: Vec<String> = f.kernels.iter().map(|k| k.to_string()).collect();
302            format!(
303                "{} {} ({}) at {} n={} runs on {} [{}]",
304                f.site,
305                f.op,
306                f.weight.as_deref().unwrap_or("no weight"),
307                f.mode.name(),
308                f.rows,
309                f.unit,
310                if k.is_empty() {
311                    "no kernel".to_string()
312                } else {
313                    k.join(" + ")
314                }
315            )
316        })
317        .collect::<Vec<_>>()
318        .join("\n  ")
319}
320
321/// 2026-10-02: Audit `planned` (and each runtime route's arm) of `c` on `r`'s class: `None`
322/// when the class states no policy; an error ([`super::HwError::TensorCore`]) on a violation.
323pub fn enforce(
324    r: &super::plan::Resolved,
325    c: &Circuit,
326    planned: &super::plan::Planned,
327) -> Result<Option<TcAudit>, super::HwError> {
328    let Some(policy) = &r.chain[0].tensor_core else {
329        return Ok(None);
330    };
331    let mut total = TcAudit::default();
332    let arms = std::iter::once(planned).chain(planned.routes.iter().map(|rp| &rp.planned));
333    for p in arms {
334        let a = audit(c, &p.plan, &r.families, policy, super::plan::NOVEL_EMITTER)
335            .map_err(super::HwError::Plan)?;
336        total.covered += a.covered;
337        total.on_tensor_cores += a.on_tensor_cores;
338        for e in a.exempted {
339            if !total.exempted.contains(&e) {
340                total.exempted.push(e);
341            }
342        }
343        for v in a.violations {
344            if !total.violations.contains(&v) {
345                total.violations.push(v);
346            }
347        }
348        for g in a.gaps {
349            if !total.gaps.contains(&g) {
350                total.gaps.push(g);
351            }
352        }
353    }
354    if !total.violations.is_empty() {
355        return Err(super::HwError::TensorCore {
356            class: r.chain[0].name.clone(),
357            violations: describe(&total.violations),
358        });
359    }
360    Ok(Some(total))
361}
362
363#[cfg(test)]
364#[path = "tc_policy_tests.rs"]
365mod tc_policy_tests;
metrale_circuit/hardware/tc_policy.rs

metrale_circuit/hardware/
tc_policy.rs
