metrale_circuit/hardware/
tc_policy_parse.rs1use std::collections::BTreeSet;
9
10use serde::Deserialize;
11
12use super::{Exempt, ExemptKind, OpMatch, Require, TcPolicy, Weights};
13use crate::format::Format;
14use crate::rules::{KernelId, Mode};
15
16#[derive(Deserialize)]
17#[serde(deny_unknown_fields)]
18struct PolicyFile {
19 require: Vec<RequireFile>,
20 #[serde(default)]
21 exempt: Vec<ExemptFile>,
22}
23
24#[derive(Deserialize)]
25#[serde(deny_unknown_fields)]
26struct RequireFile {
27 ops: Vec<String>,
28 modes: Vec<String>,
29 min_rows: u64,
30 weights: Vec<String>,
31}
32
33#[derive(Deserialize)]
34#[serde(deny_unknown_fields)]
35struct ExemptFile {
36 ops: Vec<String>,
37 modes: Vec<String>,
38 rows: [u64; 2],
39 kernels: Vec<String>,
40 kind: String,
41 reason: String,
42 evidence: Option<String>,
43}
44
45fn ops(list: &[String]) -> Result<Vec<OpMatch>, String> {
46 if list.is_empty() {
47 return Err("names no op".into());
48 }
49 list.iter().map(|o| OpMatch::parse(o)).collect()
50}
51
52fn modes(list: &[String]) -> Result<BTreeSet<Mode>, String> {
53 if list.is_empty() {
54 return Err("names no mode".into());
55 }
56 list.iter()
57 .map(|m| Mode::parse(m).ok_or_else(|| format!("unknown mode `{m}`")))
58 .collect()
59}
60
61fn weights(list: &[String]) -> Result<Weights, String> {
62 match list {
63 [one] if one == "any" => Ok(Weights::Any),
64 [] => Err("names no weight format (`[\"any\"]` covers every one)".into()),
65 _ => list
66 .iter()
67 .map(|f| Format::parse(f).map_err(|e| e.to_string()))
68 .collect::<Result<_, _>>()
69 .map(Weights::Only),
70 }
71}
72
73pub fn parse_policy(hardware: &toml::Table) -> Result<Option<TcPolicy>, String> {
76 let Some(v) = hardware.get("tensor_core_policy") else {
77 return Ok(None);
78 };
79 let f: PolicyFile = v
80 .clone()
81 .try_into()
82 .map_err(|e: toml::de::Error| format!("[tensor_core_policy]: {e}"))?;
83 if f.require.is_empty() {
84 return Err("[tensor_core_policy] states no requirement".into());
85 }
86 let require = f
87 .require
88 .iter()
89 .enumerate()
90 .map(|(i, r)| {
91 let at = |e: String| format!("[[tensor_core_policy.require]] #{}: {e}", i + 1);
92 if r.min_rows == 0 {
93 return Err(at("min_rows is at least 1".into()));
94 }
95 Ok(Require {
96 ops: ops(&r.ops).map_err(at)?,
97 modes: modes(&r.modes).map_err(at)?,
98 min_rows: r.min_rows,
99 weights: weights(&r.weights).map_err(at)?,
100 })
101 })
102 .collect::<Result<Vec<_>, String>>()?;
103 let exempt = f
104 .exempt
105 .iter()
106 .enumerate()
107 .map(|(i, e)| {
108 let at = |m: String| format!("[[tensor_core_policy.exempt]] #{}: {m}", i + 1);
109 let kind = match e.kind.as_str() {
110 "measured" => ExemptKind::Measured,
111 "shape" => ExemptKind::Shape,
112 "backlog" => ExemptKind::Backlog,
113 other => {
114 return Err(at(format!("kind `{other}` (measured | shape | backlog)")));
115 }
116 };
117 if kind == ExemptKind::Measured && e.evidence.as_deref().is_none_or(str::is_empty) {
118 return Err(at("a measured exemption names its evidence".into()));
119 }
120 if e.rows[0] == 0 || e.rows[0] > e.rows[1] {
121 return Err(at(format!("rows {:?} is not a range from 1", e.rows)));
122 }
123 if e.reason.trim().is_empty() || e.kernels.is_empty() {
124 return Err(at("states its reason and the kernels it allows".into()));
125 }
126 let kernels = e
127 .kernels
128 .iter()
129 .map(|k| {
130 k.split_once("::")
131 .filter(|(m, n)| !m.is_empty() && !n.is_empty())
132 .map(|(m, n)| KernelId {
133 module: m.to_string(),
134 func: n.to_string(),
135 })
136 .ok_or_else(|| at(format!("kernel `{k}` is not `module::function`")))
137 })
138 .collect::<Result<_, _>>()?;
139 Ok(Exempt {
140 ops: ops(&e.ops).map_err(at)?,
141 modes: modes(&e.modes).map_err(at)?,
142 rows: (e.rows[0], e.rows[1]),
143 kernels,
144 kind,
145 reason: e.reason.clone(),
146 evidence: e.evidence.clone(),
147 })
148 })
149 .collect::<Result<Vec<_>, String>>()?;
150 Ok(Some(TcPolicy { require, exempt }))
151}