1use std::collections::BTreeSet;
4use std::sync::{Arc, Mutex};
5
6use sim_kernel::Symbol;
7use sim_lib_compute_model::ModeledComputeProfile;
8use sim_lib_numbers_tensor::TensorRequest;
9
10use crate::{ComputeEvidenceKind, ComputePhysicalEvidence, verify_physical};
11
12const DEFAULT_STALE_AFTER_TICKS: u64 = 10_000;
13const CELL_BYTES: u64 = 8;
14
15#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct ComputeDeviceIdentity {
18 pub adapter: String,
20 pub driver: String,
22 pub backend: String,
24}
25
26impl ComputeDeviceIdentity {
27 pub fn new(
29 adapter: impl Into<String>,
30 driver: impl Into<String>,
31 backend: impl Into<String>,
32 ) -> Self {
33 Self {
34 adapter: adapter.into(),
35 driver: driver.into(),
36 backend: backend.into(),
37 }
38 }
39
40 fn compatible_with(&self, other: &Self) -> bool {
41 self == other
42 }
43}
44
45#[derive(Clone, Debug, PartialEq, Eq)]
47pub struct ComputeProfileLimits {
48 pub max_resident_bytes: u64,
50 pub max_storage_binding_bytes: u64,
52 pub max_queue_depth: usize,
54 pub max_queue_bytes: u64,
56 pub submission_deadline_ticks: u64,
58}
59
60impl From<&ModeledComputeProfile> for ComputeProfileLimits {
61 fn from(profile: &ModeledComputeProfile) -> Self {
62 Self {
63 max_resident_bytes: profile.max_resident_bytes,
64 max_storage_binding_bytes: profile.max_storage_binding_bytes,
65 max_queue_depth: profile.max_queue_depth,
66 max_queue_bytes: profile.max_queue_bytes,
67 submission_deadline_ticks: profile.submission_deadline_ticks,
68 }
69 }
70}
71
72#[derive(Clone, Debug, PartialEq, Eq)]
74pub struct ComputeProfileSamples {
75 pub upload_bytes_per_tick: Vec<u64>,
77 pub download_bytes_per_tick: Vec<u64>,
79 pub launch_ticks: Vec<u64>,
81 pub element_elements_per_tick: Vec<u64>,
83 pub reduction_elements_per_tick: Vec<u64>,
85 pub matmul_ops_per_tick: Vec<u64>,
87}
88
89impl ComputeProfileSamples {
90 fn conclusive(&self) -> bool {
91 [
92 &self.upload_bytes_per_tick,
93 &self.download_bytes_per_tick,
94 &self.launch_ticks,
95 &self.element_elements_per_tick,
96 &self.reduction_elements_per_tick,
97 &self.matmul_ops_per_tick,
98 ]
99 .into_iter()
100 .all(|samples| samples.iter().any(|sample| *sample > 0))
101 }
102}
103
104#[derive(Clone, Debug, PartialEq, Eq)]
106pub struct ComputeThermalPowerContext {
107 pub thermal: String,
109 pub power: String,
111}
112
113#[derive(Clone, Debug, PartialEq, Eq)]
115pub struct ComputeProfileProvenance {
116 pub evidence_kind: ComputeEvidenceKind,
118 pub producer: String,
120 pub measured_at_tick: u64,
122 pub stale_after_ticks: u64,
124 pub observed_identity: Option<ComputeDeviceIdentity>,
126}
127
128#[derive(Clone, Debug, PartialEq, Eq)]
130pub struct MeasuredComputeProfile {
131 pub identity: ComputeDeviceIdentity,
133 pub limits: ComputeProfileLimits,
135 pub samples: ComputeProfileSamples,
137 pub context: ComputeThermalPowerContext,
139 pub tile_bytes: u64,
141 pub allocation_bytes: Vec<u64>,
143 pub provenance: ComputeProfileProvenance,
145 pub modeled: ModeledComputeProfile,
147}
148
149impl sim_citizen::Citizen for MeasuredComputeProfile {
150 fn citizen_symbol() -> Symbol {
151 measured_compute_profile_citizen_symbol()
152 }
153
154 fn citizen_version() -> u32 {
155 0
156 }
157
158 fn citizen_arity() -> usize {
159 9
160 }
161
162 fn citizen_fields() -> &'static [&'static str] {
163 &[
164 "identity",
165 "limits",
166 "samples",
167 "context",
168 "tile_bytes",
169 "allocation_bytes",
170 "provenance",
171 "modeled_provider",
172 "shape",
173 ]
174 }
175}
176
177impl MeasuredComputeProfile {
178 pub fn is_conclusive(&self) -> bool {
180 self.tile_bytes > 0
181 && self.allocation_bytes.iter().any(|bytes| *bytes > 0)
182 && self.samples.conclusive()
183 && self.modeled.fault.is_none()
184 }
185
186 pub fn is_compatible(&self, identity: &ComputeDeviceIdentity, now_tick: u64) -> bool {
188 self.identity.compatible_with(identity)
189 && now_tick.saturating_sub(self.provenance.measured_at_tick)
190 <= self.provenance.stale_after_ticks
191 && self.limits.max_resident_bytes > 0
192 && self.limits.max_storage_binding_bytes > 0
193 }
194
195 pub fn modeled_profile(&self) -> ModeledComputeProfile {
197 self.modeled.clone()
198 }
199
200 pub(crate) fn estimated_profile_bytes(&self) -> usize {
201 self.identity.adapter.len()
202 + self.identity.driver.len()
203 + self.identity.backend.len()
204 + self.context.thermal.len()
205 + self.context.power.len()
206 + self.provenance.producer.len()
207 + (self.allocation_bytes.len()
208 + self.samples.upload_bytes_per_tick.len()
209 + self.samples.download_bytes_per_tick.len()
210 + self.samples.launch_ticks.len()
211 + self.samples.element_elements_per_tick.len()
212 + self.samples.reduction_elements_per_tick.len()
213 + self.samples.matmul_ops_per_tick.len())
214 * std::mem::size_of::<u64>()
215 + 256
216 }
217}
218
219impl ComputePhysicalEvidence for MeasuredComputeProfile {
220 fn evidence_kind(&self) -> ComputeEvidenceKind {
221 self.provenance.evidence_kind
222 }
223
224 fn claimed_identity(&self) -> Option<&ComputeDeviceIdentity> {
225 Some(&self.identity)
226 }
227
228 fn observed_identity(&self) -> Option<&ComputeDeviceIdentity> {
229 self.provenance.observed_identity.as_ref()
230 }
231}
232
233pub fn measured_compute_profile_citizen_symbol() -> Symbol {
235 Symbol::qualified("compute-profile", "MeasuredProfile")
236}
237
238pub fn measured_compute_profile_shape_symbol() -> Symbol {
240 Symbol::qualified("compute-profile", "MeasuredProfileShape")
241}
242
243#[derive(Clone, Debug, PartialEq, Eq)]
245pub struct BenchmarkBounds {
246 pub transfer_bytes: Vec<u64>,
248 pub element_counts: Vec<u64>,
250 pub matrix_edges: Vec<u64>,
252 pub max_sample_bytes: u64,
254}
255
256impl Default for BenchmarkBounds {
257 fn default() -> Self {
258 Self {
259 transfer_bytes: vec![4096, 16 * 1024, 64 * 1024],
260 element_counts: vec![256, 1024, 4096],
261 matrix_edges: vec![8, 16, 32],
262 max_sample_bytes: 256 * 1024,
263 }
264 }
265}
266
267pub fn measure_bounded_profile(
269 identity: ComputeDeviceIdentity,
270 modeled: ModeledComputeProfile,
271 context: ComputeThermalPowerContext,
272 producer: impl Into<String>,
273 now_tick: u64,
274 bounds: BenchmarkBounds,
275) -> MeasuredComputeProfile {
276 let bounded_transfers = bounded_nonzero(bounds.transfer_bytes, bounds.max_sample_bytes);
277 let bounded_elements = bounded_nonzero(
278 bounds.element_counts,
279 bounds.max_sample_bytes.saturating_div(CELL_BYTES).max(1),
280 );
281 let bounded_edges = bounded_nonzero(bounds.matrix_edges, 512);
282 let tile_bytes = modeled
283 .segment_tile_bytes
284 .min(modeled.max_storage_binding_bytes)
285 .max(CELL_BYTES);
286 let allocation_bytes = bounded_transfers
287 .iter()
288 .copied()
289 .filter(|bytes| *bytes <= modeled.max_resident_bytes)
290 .collect::<Vec<_>>();
291 MeasuredComputeProfile {
292 identity: identity.clone(),
293 limits: ComputeProfileLimits::from(&modeled),
294 samples: ComputeProfileSamples {
295 upload_bytes_per_tick: bounded_transfers.clone(),
296 download_bytes_per_tick: bounded_transfers
297 .iter()
298 .map(|bytes| (*bytes).saturating_mul(9) / 10)
299 .collect(),
300 launch_ticks: bounded_transfers
301 .iter()
302 .enumerate()
303 .map(|(index, _)| (index as u64) + 1)
304 .collect(),
305 element_elements_per_tick: bounded_elements.clone(),
306 reduction_elements_per_tick: bounded_elements.iter().map(|count| count / 2).collect(),
307 matmul_ops_per_tick: bounded_edges.iter().map(|edge| edge.pow(3)).collect(),
308 },
309 context,
310 tile_bytes,
311 allocation_bytes,
312 provenance: ComputeProfileProvenance {
313 evidence_kind: ComputeEvidenceKind::Modeled,
314 producer: producer.into(),
315 measured_at_tick: now_tick,
316 stale_after_ticks: DEFAULT_STALE_AFTER_TICKS,
317 observed_identity: Some(identity.clone()),
318 },
319 modeled,
320 }
321}
322
323#[derive(Clone, Debug, PartialEq, Eq)]
325pub enum AutoRouteDecision {
326 Absent,
328 Stale,
330 Incompatible,
332 Inconclusive,
334 NonPhysical,
336 Device,
338}
339
340#[derive(Clone, Debug)]
342pub struct AutoComputeRouter {
343 expected: ComputeDeviceIdentity,
344 now_tick: u64,
345}
346
347impl AutoComputeRouter {
348 pub fn new(expected: ComputeDeviceIdentity, now_tick: u64) -> Self {
350 Self { expected, now_tick }
351 }
352
353 pub fn choose(
355 &self,
356 profile: Option<&MeasuredComputeProfile>,
357 ) -> (AutoRouteDecision, Option<ModeledComputeProfile>) {
358 let Some(profile) = profile else {
359 return (AutoRouteDecision::Absent, None);
360 };
361 if !profile.identity.compatible_with(&self.expected) {
362 return (AutoRouteDecision::Incompatible, None);
363 }
364 if self
365 .now_tick
366 .saturating_sub(profile.provenance.measured_at_tick)
367 > profile.provenance.stale_after_ticks
368 {
369 return (AutoRouteDecision::Stale, None);
370 }
371 if !profile.is_conclusive() {
372 return (AutoRouteDecision::Inconclusive, None);
373 }
374 if verify_physical(profile).is_err() {
375 return (AutoRouteDecision::NonPhysical, None);
376 }
377 (AutoRouteDecision::Device, Some(profile.modeled_profile()))
378 }
379}
380
381#[derive(Clone, Debug, PartialEq, Eq)]
383pub struct AutoRoutingEvent {
384 pub provider: String,
386 pub decision: AutoRouteDecision,
388 pub materialization_bytes: u64,
390 pub synchronizations: usize,
392}
393
394#[derive(Clone, Debug, Default)]
396pub struct AutoRoutingLedger {
397 events: Arc<Mutex<Vec<AutoRoutingEvent>>>,
398}
399
400impl AutoRoutingLedger {
401 pub fn record(&self, event: AutoRoutingEvent) {
403 self.events
404 .lock()
405 .expect("auto routing ledger poisoned")
406 .push(event);
407 }
408
409 pub fn events(&self) -> Vec<AutoRoutingEvent> {
411 self.events
412 .lock()
413 .expect("auto routing ledger poisoned")
414 .clone()
415 }
416}
417
418pub(crate) fn request_materialization_bytes(request: &TensorRequest) -> u64 {
419 request
420 .inputs
421 .iter()
422 .map(|tensor| tensor.shape().iter().copied().product::<usize>() as u64 * CELL_BYTES)
423 .chain(std::iter::once(
424 request.output.shape().iter().copied().product::<usize>() as u64 * CELL_BYTES,
425 ))
426 .sum()
427}
428
429fn bounded_nonzero(values: Vec<u64>, max: u64) -> Vec<u64> {
430 let mut seen = BTreeSet::new();
431 values
432 .into_iter()
433 .filter(|value| *value > 0 && *value <= max)
434 .filter(|value| seen.insert(*value))
435 .collect()
436}