1use sim_kernel::{Cx, Error, Expr, NumberLiteral, Result, Symbol, Value};
4use sim_lib_compute_model::ModeledComputeProfile;
5
6use crate::{
7 ComputeDeviceIdentity, ComputeProfileLimits, ComputeProfileProvenance, ComputeProfileSamples,
8 ComputeThermalPowerContext, MeasuredComputeProfile,
9};
10
11const PROFILE_SCHEMA: &str = "sim.compute.auto.profile.v1";
12const DEFAULT_MAX_KEY_BYTES: usize = 96;
13const DEFAULT_MAX_PROFILE_BYTES: usize = 8192;
14
15#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct ProfileStorePolicy {
18 pub max_key_bytes: usize,
20 pub max_profile_bytes: usize,
22}
23
24impl Default for ProfileStorePolicy {
25 fn default() -> Self {
26 Self {
27 max_key_bytes: DEFAULT_MAX_KEY_BYTES,
28 max_profile_bytes: DEFAULT_MAX_PROFILE_BYTES,
29 }
30 }
31}
32
33pub struct ProfileStore {
35 table: Value,
36 policy: ProfileStorePolicy,
37}
38
39impl ProfileStore {
40 pub fn new(table: Value, policy: ProfileStorePolicy) -> Result<Self> {
42 if table.object().as_table_impl().is_none() {
43 return Err(Error::TypeMismatch {
44 expected: "Table or Dir",
45 found: "non-table",
46 });
47 }
48 Ok(Self { table, policy })
49 }
50
51 pub fn save(&self, cx: &mut Cx, key: Symbol, profile: &MeasuredComputeProfile) -> Result<()> {
53 self.validate_key(&key)?;
54 if profile.estimated_profile_bytes() > self.policy.max_profile_bytes {
55 return Err(Error::Eval(
56 "compute profile exceeds byte policy".to_owned(),
57 ));
58 }
59 let value = profile_to_value(cx, profile)?;
60 self.table
61 .object()
62 .as_table_impl()
63 .expect("profile store table")
64 .set(cx, key, value)
65 }
66
67 pub fn load(&self, cx: &mut Cx, key: Symbol) -> Result<Option<MeasuredComputeProfile>> {
69 self.validate_key(&key)?;
70 let table = self
71 .table
72 .object()
73 .as_table_impl()
74 .expect("profile store table");
75 if !table.has(cx, key.clone())? {
76 return Ok(None);
77 }
78 let value = table.get(cx, key)?;
79 let profile = profile_from_value(cx, &value)?;
80 if profile.estimated_profile_bytes() > self.policy.max_profile_bytes {
81 return Err(Error::Eval(
82 "stored compute profile exceeds byte policy".to_owned(),
83 ));
84 }
85 Ok(Some(profile))
86 }
87
88 pub fn keys(&self, cx: &mut Cx) -> Result<Vec<Symbol>> {
90 let keys = self
91 .table
92 .object()
93 .as_table_impl()
94 .expect("profile store table")
95 .keys(cx)?;
96 for key in &keys {
97 self.validate_key(key)?;
98 }
99 Ok(keys)
100 }
101
102 fn validate_key(&self, key: &Symbol) -> Result<()> {
103 if key.to_string().len() > self.policy.max_key_bytes {
104 return Err(Error::Eval(
105 "compute profile key exceeds byte policy".to_owned(),
106 ));
107 }
108 Ok(())
109 }
110}
111
112fn profile_to_value(cx: &mut Cx, profile: &MeasuredComputeProfile) -> Result<Value> {
113 let entries = vec![
114 (sym("schema"), string(cx, PROFILE_SCHEMA)?),
115 (sym("adapter"), string(cx, &profile.identity.adapter)?),
116 (sym("driver"), string(cx, &profile.identity.driver)?),
117 (sym("backend"), string(cx, &profile.identity.backend)?),
118 (
119 sym("max-resident-bytes"),
120 number(cx, profile.limits.max_resident_bytes)?,
121 ),
122 (
123 sym("max-storage-binding-bytes"),
124 number(cx, profile.limits.max_storage_binding_bytes)?,
125 ),
126 (
127 sym("max-queue-depth"),
128 number(cx, profile.limits.max_queue_depth as u64)?,
129 ),
130 (
131 sym("max-queue-bytes"),
132 number(cx, profile.limits.max_queue_bytes)?,
133 ),
134 (
135 sym("submission-deadline-ticks"),
136 number(cx, profile.limits.submission_deadline_ticks)?,
137 ),
138 (
139 sym("upload-bytes-per-tick"),
140 number_list(cx, &profile.samples.upload_bytes_per_tick)?,
141 ),
142 (
143 sym("download-bytes-per-tick"),
144 number_list(cx, &profile.samples.download_bytes_per_tick)?,
145 ),
146 (
147 sym("launch-ticks"),
148 number_list(cx, &profile.samples.launch_ticks)?,
149 ),
150 (
151 sym("element-elements-per-tick"),
152 number_list(cx, &profile.samples.element_elements_per_tick)?,
153 ),
154 (
155 sym("reduction-elements-per-tick"),
156 number_list(cx, &profile.samples.reduction_elements_per_tick)?,
157 ),
158 (
159 sym("matmul-ops-per-tick"),
160 number_list(cx, &profile.samples.matmul_ops_per_tick)?,
161 ),
162 (sym("thermal"), string(cx, &profile.context.thermal)?),
163 (sym("power"), string(cx, &profile.context.power)?),
164 (sym("tile-bytes"), number(cx, profile.tile_bytes)?),
165 (
166 sym("allocation-bytes"),
167 number_list(cx, &profile.allocation_bytes)?,
168 ),
169 (sym("producer"), string(cx, &profile.provenance.producer)?),
170 (
171 sym("measured-at-tick"),
172 number(cx, profile.provenance.measured_at_tick)?,
173 ),
174 (
175 sym("stale-after-ticks"),
176 number(cx, profile.provenance.stale_after_ticks)?,
177 ),
178 (
179 sym("modeled-provider"),
180 string(cx, &profile.modeled.provider)?,
181 ),
182 ];
183 cx.factory().table(entries)
184}
185
186fn profile_from_value(cx: &mut Cx, value: &Value) -> Result<MeasuredComputeProfile> {
187 let entries = value
188 .object()
189 .as_table_impl()
190 .ok_or(Error::TypeMismatch {
191 expected: "profile table",
192 found: "non-table",
193 })?
194 .entries(cx)?;
195 let get = |name: &str| -> Result<Value> {
196 entries
197 .iter()
198 .find(|(key, _)| *key == sym(name))
199 .map(|(_, value)| value.clone())
200 .ok_or_else(|| Error::Eval(format!("compute profile missing {name}")))
201 };
202 if string_value(cx, &get("schema")?)? != PROFILE_SCHEMA {
203 return Err(Error::Eval("unsupported compute profile schema".to_owned()));
204 }
205 let max_queue_depth = number_value(cx, &get("max-queue-depth")?)?;
206 Ok(MeasuredComputeProfile {
207 identity: ComputeDeviceIdentity {
208 adapter: string_value(cx, &get("adapter")?)?,
209 driver: string_value(cx, &get("driver")?)?,
210 backend: string_value(cx, &get("backend")?)?,
211 },
212 limits: ComputeProfileLimits {
213 max_resident_bytes: number_value(cx, &get("max-resident-bytes")?)?,
214 max_storage_binding_bytes: number_value(cx, &get("max-storage-binding-bytes")?)?,
215 max_queue_depth: usize::try_from(max_queue_depth)
216 .map_err(|_| Error::Eval("max-queue-depth exceeds usize".to_owned()))?,
217 max_queue_bytes: number_value(cx, &get("max-queue-bytes")?)?,
218 submission_deadline_ticks: number_value(cx, &get("submission-deadline-ticks")?)?,
219 },
220 samples: ComputeProfileSamples {
221 upload_bytes_per_tick: number_vec(cx, &get("upload-bytes-per-tick")?)?,
222 download_bytes_per_tick: number_vec(cx, &get("download-bytes-per-tick")?)?,
223 launch_ticks: number_vec(cx, &get("launch-ticks")?)?,
224 element_elements_per_tick: number_vec(cx, &get("element-elements-per-tick")?)?,
225 reduction_elements_per_tick: number_vec(cx, &get("reduction-elements-per-tick")?)?,
226 matmul_ops_per_tick: number_vec(cx, &get("matmul-ops-per-tick")?)?,
227 },
228 context: ComputeThermalPowerContext {
229 thermal: string_value(cx, &get("thermal")?)?,
230 power: string_value(cx, &get("power")?)?,
231 },
232 tile_bytes: number_value(cx, &get("tile-bytes")?)?,
233 allocation_bytes: number_vec(cx, &get("allocation-bytes")?)?,
234 provenance: ComputeProfileProvenance {
235 producer: string_value(cx, &get("producer")?)?,
236 measured_at_tick: number_value(cx, &get("measured-at-tick")?)?,
237 stale_after_ticks: number_value(cx, &get("stale-after-ticks")?)?,
238 },
239 modeled: ModeledComputeProfile {
240 provider: string_value(cx, &get("modeled-provider")?)?,
241 max_queue_depth: usize::try_from(max_queue_depth)
242 .map_err(|_| Error::Eval("max-queue-depth exceeds usize".to_owned()))?,
243 max_queue_bytes: number_value(cx, &get("max-queue-bytes")?)?,
244 max_resident_bytes: number_value(cx, &get("max-resident-bytes")?)?,
245 segment_tile_bytes: number_value(cx, &get("tile-bytes")?)?,
246 max_storage_binding_bytes: number_value(cx, &get("max-storage-binding-bytes")?)?,
247 submission_deadline_ticks: number_value(cx, &get("submission-deadline-ticks")?)?,
248 fault: None,
249 auto_flush_batches: false,
250 },
251 })
252}
253
254fn number(cx: &mut Cx, value: u64) -> Result<Value> {
255 cx.factory()
256 .number_literal(Symbol::qualified("numbers", "u64"), value.to_string())
257}
258
259fn string(cx: &mut Cx, value: &str) -> Result<Value> {
260 cx.factory().string(value.to_owned())
261}
262
263fn number_list(cx: &mut Cx, values: &[u64]) -> Result<Value> {
264 let values = values
265 .iter()
266 .map(|value| number(cx, *value))
267 .collect::<Result<Vec<_>>>()?;
268 cx.factory().list(values)
269}
270
271fn string_value(cx: &mut Cx, value: &Value) -> Result<String> {
272 match value.object().as_expr(cx)? {
273 Expr::String(value) => Ok(value),
274 _ => Err(Error::TypeMismatch {
275 expected: "string",
276 found: "non-string",
277 }),
278 }
279}
280
281fn number_value(cx: &mut Cx, value: &Value) -> Result<u64> {
282 match value.object().as_expr(cx)? {
283 Expr::Number(NumberLiteral { canonical, .. }) => canonical
284 .parse::<u64>()
285 .map_err(|_| Error::Eval("invalid u64 profile value".to_owned())),
286 _ => Err(Error::TypeMismatch {
287 expected: "number",
288 found: "non-number",
289 }),
290 }
291}
292
293fn number_vec(cx: &mut Cx, value: &Value) -> Result<Vec<u64>> {
294 match value.object().as_expr(cx)? {
295 Expr::List(values) => values
296 .iter()
297 .map(|expr| match expr {
298 Expr::Number(NumberLiteral { canonical, .. }) => canonical
299 .parse::<u64>()
300 .map_err(|_| Error::Eval("invalid u64 profile sample".to_owned())),
301 _ => Err(Error::TypeMismatch {
302 expected: "number list",
303 found: "non-number list",
304 }),
305 })
306 .collect(),
307 _ => Err(Error::TypeMismatch {
308 expected: "list",
309 found: "non-list",
310 }),
311 }
312}
313
314fn sym(name: &str) -> Symbol {
315 Symbol::qualified("compute-profile", name)
316}