1use crate::SpaceError;
4use crate::constraint::Constraint;
5use serde::Deserialize;
6use std::fmt;
7use std::path::{Path, PathBuf};
8
9const BLOCK_AXIS_LIMITS: [(&str, u64); 3] = [("block_x", 1024), ("block_y", 1024), ("block_z", 64)];
14
15#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
17pub enum Value {
18 Int(u64),
20 Str(String),
22}
23
24impl fmt::Display for Value {
25 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26 match self {
27 Value::Int(n) => write!(f, "{n}"),
28 Value::Str(s) => f.write_str(s),
29 }
30 }
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum DimRole {
36 Launch,
39 Spec,
42}
43
44#[derive(Debug, Clone)]
46pub struct Dim {
47 pub name: String,
50 pub role: DimRole,
52 pub values: Vec<Value>,
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum SafetyExpectation {
60 Flip,
62 Stable,
64 None,
66}
67
68#[derive(Debug, Clone)]
70pub struct KernelSpec {
71 pub name: String,
73 pub entry: String,
75 pub domain: u8,
77 pub needs_cc: Option<String>,
79 pub known: SafetyExpectation,
82 pub dir: PathBuf,
84 pub dims: Vec<Dim>,
87 pub constraints: Vec<Constraint>,
89}
90
91#[derive(Deserialize)]
92struct RawSpec {
93 kernel: RawKernel,
94 #[serde(default)]
95 dims: toml::value::Table,
96 #[serde(default)]
97 constraints: RawConstraints,
98}
99
100#[derive(Deserialize)]
101struct RawKernel {
102 name: String,
103 entry: String,
104 domain: u8,
105 #[serde(default)]
106 needs_cc: Option<String>,
107 #[serde(default)]
108 known: Option<String>,
109}
110
111#[derive(Deserialize, Default)]
112struct RawConstraints {
113 #[serde(default)]
114 exprs: Vec<String>,
115}
116
117#[derive(Deserialize)]
118struct RawDim {
119 #[serde(default)]
120 role: Option<String>,
121 values: Vec<toml::Value>,
122}
123
124impl KernelSpec {
125 pub fn load(dir: &Path) -> Result<Self, SpaceError> {
127 let path = dir.join("kernel.toml");
128 let text = std::fs::read_to_string(&path).map_err(|source| SpaceError::Io {
129 path: path.display().to_string(),
130 source,
131 })?;
132 Self::parse(&path.display().to_string(), &text, dir.to_path_buf())
133 }
134
135 pub fn from_toml_str(origin: &str, text: &str) -> Result<Self, SpaceError> {
137 Self::parse(origin, text, PathBuf::from("."))
138 }
139
140 fn parse(origin: &str, text: &str, dir: PathBuf) -> Result<Self, SpaceError> {
141 let raw: RawSpec = toml::from_str(text).map_err(|source| SpaceError::Parse {
142 path: origin.to_string(),
143 source: Box::new(source),
144 })?;
145
146 if !(1..=3).contains(&raw.kernel.domain) {
147 return Err(SpaceError::Invalid(format!(
148 "domain must be 1..=3, got {}",
149 raw.kernel.domain
150 )));
151 }
152 let known = match raw.kernel.known.as_deref() {
153 Some("flip") => SafetyExpectation::Flip,
154 Some("stable") => SafetyExpectation::Stable,
155 None => SafetyExpectation::None,
156 Some(other) => {
157 return Err(SpaceError::Invalid(format!(
158 "known must be \"flip\" or \"stable\", got {other:?}"
159 )));
160 }
161 };
162
163 let mut dims = Vec::new();
164 for (name, value) in raw.dims {
165 if !name
166 .chars()
167 .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
168 {
169 return Err(SpaceError::Invalid(format!(
170 "dimension name {name:?} must be [a-z0-9_]"
171 )));
172 }
173 let raw_dim: RawDim = value
174 .try_into()
175 .map_err(|e| SpaceError::Invalid(format!("dimension {name}: {e}")))?;
176 let role = match raw_dim.role.as_deref() {
177 Some("launch") => DimRole::Launch,
178 Some("spec") => DimRole::Spec,
179 None if name.starts_with("block_") => DimRole::Launch,
180 None => DimRole::Spec,
181 Some(other) => {
182 return Err(SpaceError::Invalid(format!(
183 "dimension {name}: role must be \"launch\" or \"spec\", got {other:?}"
184 )));
185 }
186 };
187 let mut values = Vec::new();
188 for v in raw_dim.values {
189 match v {
190 toml::Value::Integer(n) if n >= 0 => values.push(Value::Int(n as u64)),
191 toml::Value::String(s) => values.push(Value::Str(s)),
192 other => {
193 return Err(SpaceError::Invalid(format!(
194 "dimension {name}: values must be non-negative integers or strings, got {other}"
195 )));
196 }
197 }
198 }
199 if values.is_empty() {
200 return Err(SpaceError::Invalid(format!(
201 "dimension {name} has no values"
202 )));
203 }
204 if let Some(axis) = BLOCK_AXIS_LIMITS.iter().find(|(n, _)| *n == name) {
213 let (_, limit) = axis;
214 for v in &values {
215 if let Value::Int(n) = v
216 && *n > *limit
217 {
218 return Err(SpaceError::Invalid(format!(
219 "dimension {name}: {n} exceeds the CUDA limit of {limit} \
220 threads on this axis (max 1024 per block overall)"
221 )));
222 }
223 }
224 }
225 let mut seen = values.clone();
226 seen.sort();
227 seen.dedup();
228 if seen.len() != values.len() {
229 return Err(SpaceError::Invalid(format!(
230 "dimension {name} has duplicate values"
231 )));
232 }
233 dims.push(Dim { name, role, values });
234 }
235 if dims.is_empty() {
236 return Err(SpaceError::Invalid("spec declares no dimensions".into()));
237 }
238 dims.sort_by(|a, b| a.name.cmp(&b.name));
239
240 let dim_names: Vec<&str> = dims.iter().map(|d| d.name.as_str()).collect();
241 let constraints = raw
242 .constraints
243 .exprs
244 .iter()
245 .map(|e| Constraint::parse(e, &dim_names))
246 .collect::<Result<Vec<_>, _>>()?;
247
248 Ok(KernelSpec {
249 name: raw.kernel.name,
250 entry: raw.kernel.entry,
251 domain: raw.kernel.domain,
252 needs_cc: raw.kernel.needs_cc,
253 known,
254 dir,
255 dims,
256 constraints,
257 })
258 }
259
260 pub fn dim(&self, name: &str) -> Option<&Dim> {
262 self.dims.iter().find(|d| d.name == name)
263 }
264
265 pub fn dims_sorted(&self) -> Vec<&Dim> {
267 self.dims.iter().collect()
268 }
269}