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
10const DEFAULT_STALE_AFTER_TICKS: u64 = 10_000;
11const CELL_BYTES: u64 = 8;
12
13#[derive(Clone, Debug, PartialEq, Eq)]
15pub struct ComputeDeviceIdentity {
16 pub adapter: String,
18 pub driver: String,
20 pub backend: String,
22}
23
24impl ComputeDeviceIdentity {
25 pub fn new(
27 adapter: impl Into<String>,
28 driver: impl Into<String>,
29 backend: impl Into<String>,
30 ) -> Self {
31 Self {
32 adapter: adapter.into(),
33 driver: driver.into(),
34 backend: backend.into(),
35 }
36 }
37
38 fn compatible_with(&self, other: &Self) -> bool {
39 self == other
40 }
41}
42
43#[derive(Clone, Debug, PartialEq, Eq)]
45pub struct ComputeProfileLimits {
46 pub max_resident_bytes: u64,
48 pub max_storage_binding_bytes: u64,
50 pub max_queue_depth: usize,
52 pub max_queue_bytes: u64,
54 pub submission_deadline_ticks: u64,
56}
57
58impl From<&ModeledComputeProfile> for ComputeProfileLimits {
59 fn from(profile: &ModeledComputeProfile) -> Self {
60 Self {
61 max_resident_bytes: profile.max_resident_bytes,
62 max_storage_binding_bytes: profile.max_storage_binding_bytes,
63 max_queue_depth: profile.max_queue_depth,
64 max_queue_bytes: profile.max_queue_bytes,
65 submission_deadline_ticks: profile.submission_deadline_ticks,
66 }
67 }
68}
69
70#[derive(Clone, Debug, PartialEq, Eq)]
72pub struct ComputeProfileSamples {
73 pub upload_bytes_per_tick: Vec<u64>,
75 pub download_bytes_per_tick: Vec<u64>,
77 pub launch_ticks: Vec<u64>,
79 pub element_elements_per_tick: Vec<u64>,
81 pub reduction_elements_per_tick: Vec<u64>,
83 pub matmul_ops_per_tick: Vec<u64>,
85}
86
87impl ComputeProfileSamples {
88 fn conclusive(&self) -> bool {
89 [
90 &self.upload_bytes_per_tick,
91 &self.download_bytes_per_tick,
92 &self.launch_ticks,
93 &self.element_elements_per_tick,
94 &self.reduction_elements_per_tick,
95 &self.matmul_ops_per_tick,
96 ]
97 .into_iter()
98 .all(|samples| samples.iter().any(|sample| *sample > 0))
99 }
100}
101
102#[derive(Clone, Debug, PartialEq, Eq)]
104pub struct ComputeThermalPowerContext {
105 pub thermal: String,
107 pub power: String,
109}
110
111#[derive(Clone, Debug, PartialEq, Eq)]
113pub struct ComputeProfileProvenance {
114 pub producer: String,
116 pub measured_at_tick: u64,
118 pub stale_after_ticks: u64,
120}
121
122#[derive(Clone, Debug, PartialEq, Eq)]
124pub struct MeasuredComputeProfile {
125 pub identity: ComputeDeviceIdentity,
127 pub limits: ComputeProfileLimits,
129 pub samples: ComputeProfileSamples,
131 pub context: ComputeThermalPowerContext,
133 pub tile_bytes: u64,
135 pub allocation_bytes: Vec<u64>,
137 pub provenance: ComputeProfileProvenance,
139 pub modeled: ModeledComputeProfile,
141}
142
143impl sim_citizen::Citizen for MeasuredComputeProfile {
144 fn citizen_symbol() -> Symbol {
145 measured_compute_profile_citizen_symbol()
146 }
147
148 fn citizen_version() -> u32 {
149 0
150 }
151
152 fn citizen_arity() -> usize {
153 9
154 }
155
156 fn citizen_fields() -> &'static [&'static str] {
157 &[
158 "identity",
159 "limits",
160 "samples",
161 "context",
162 "tile_bytes",
163 "allocation_bytes",
164 "provenance",
165 "modeled_provider",
166 "shape",
167 ]
168 }
169}
170
171impl MeasuredComputeProfile {
172 pub fn is_conclusive(&self) -> bool {
174 self.tile_bytes > 0
175 && self.allocation_bytes.iter().any(|bytes| *bytes > 0)
176 && self.samples.conclusive()
177 && self.modeled.fault.is_none()
178 }
179
180 pub fn is_compatible(&self, identity: &ComputeDeviceIdentity, now_tick: u64) -> bool {
182 self.identity.compatible_with(identity)
183 && now_tick.saturating_sub(self.provenance.measured_at_tick)
184 <= self.provenance.stale_after_ticks
185 && self.limits.max_resident_bytes > 0
186 && self.limits.max_storage_binding_bytes > 0
187 }
188
189 pub fn modeled_profile(&self) -> ModeledComputeProfile {
191 self.modeled.clone()
192 }
193
194 pub(crate) fn estimated_profile_bytes(&self) -> usize {
195 self.identity.adapter.len()
196 + self.identity.driver.len()
197 + self.identity.backend.len()
198 + self.context.thermal.len()
199 + self.context.power.len()
200 + self.provenance.producer.len()
201 + (self.allocation_bytes.len()
202 + self.samples.upload_bytes_per_tick.len()
203 + self.samples.download_bytes_per_tick.len()
204 + self.samples.launch_ticks.len()
205 + self.samples.element_elements_per_tick.len()
206 + self.samples.reduction_elements_per_tick.len()
207 + self.samples.matmul_ops_per_tick.len())
208 * std::mem::size_of::<u64>()
209 + 256
210 }
211}
212
213pub fn measured_compute_profile_citizen_symbol() -> Symbol {
215 Symbol::qualified("compute-profile", "MeasuredProfile")
216}
217
218pub fn measured_compute_profile_shape_symbol() -> Symbol {
220 Symbol::qualified("compute-profile", "MeasuredProfileShape")
221}
222
223#[derive(Clone, Debug, PartialEq, Eq)]
225pub struct BenchmarkBounds {
226 pub transfer_bytes: Vec<u64>,
228 pub element_counts: Vec<u64>,
230 pub matrix_edges: Vec<u64>,
232 pub max_sample_bytes: u64,
234}
235
236impl Default for BenchmarkBounds {
237 fn default() -> Self {
238 Self {
239 transfer_bytes: vec![4096, 16 * 1024, 64 * 1024],
240 element_counts: vec![256, 1024, 4096],
241 matrix_edges: vec![8, 16, 32],
242 max_sample_bytes: 256 * 1024,
243 }
244 }
245}
246
247pub fn measure_bounded_profile(
249 identity: ComputeDeviceIdentity,
250 modeled: ModeledComputeProfile,
251 context: ComputeThermalPowerContext,
252 producer: impl Into<String>,
253 now_tick: u64,
254 bounds: BenchmarkBounds,
255) -> MeasuredComputeProfile {
256 let bounded_transfers = bounded_nonzero(bounds.transfer_bytes, bounds.max_sample_bytes);
257 let bounded_elements = bounded_nonzero(
258 bounds.element_counts,
259 bounds.max_sample_bytes.saturating_div(CELL_BYTES).max(1),
260 );
261 let bounded_edges = bounded_nonzero(bounds.matrix_edges, 512);
262 let tile_bytes = modeled
263 .segment_tile_bytes
264 .min(modeled.max_storage_binding_bytes)
265 .max(CELL_BYTES);
266 let allocation_bytes = bounded_transfers
267 .iter()
268 .copied()
269 .filter(|bytes| *bytes <= modeled.max_resident_bytes)
270 .collect::<Vec<_>>();
271 MeasuredComputeProfile {
272 identity,
273 limits: ComputeProfileLimits::from(&modeled),
274 samples: ComputeProfileSamples {
275 upload_bytes_per_tick: bounded_transfers.clone(),
276 download_bytes_per_tick: bounded_transfers
277 .iter()
278 .map(|bytes| (*bytes).saturating_mul(9) / 10)
279 .collect(),
280 launch_ticks: bounded_transfers
281 .iter()
282 .enumerate()
283 .map(|(index, _)| (index as u64) + 1)
284 .collect(),
285 element_elements_per_tick: bounded_elements.clone(),
286 reduction_elements_per_tick: bounded_elements.iter().map(|count| count / 2).collect(),
287 matmul_ops_per_tick: bounded_edges.iter().map(|edge| edge.pow(3)).collect(),
288 },
289 context,
290 tile_bytes,
291 allocation_bytes,
292 provenance: ComputeProfileProvenance {
293 producer: producer.into(),
294 measured_at_tick: now_tick,
295 stale_after_ticks: DEFAULT_STALE_AFTER_TICKS,
296 },
297 modeled,
298 }
299}
300
301#[derive(Clone, Debug, PartialEq, Eq)]
303pub enum AutoRouteDecision {
304 Absent,
306 Stale,
308 Incompatible,
310 Inconclusive,
312 Device,
314}
315
316#[derive(Clone, Debug)]
318pub struct AutoComputeRouter {
319 expected: ComputeDeviceIdentity,
320 now_tick: u64,
321}
322
323impl AutoComputeRouter {
324 pub fn new(expected: ComputeDeviceIdentity, now_tick: u64) -> Self {
326 Self { expected, now_tick }
327 }
328
329 pub fn choose(
331 &self,
332 profile: Option<&MeasuredComputeProfile>,
333 ) -> (AutoRouteDecision, Option<ModeledComputeProfile>) {
334 let Some(profile) = profile else {
335 return (AutoRouteDecision::Absent, None);
336 };
337 if !profile.identity.compatible_with(&self.expected) {
338 return (AutoRouteDecision::Incompatible, None);
339 }
340 if self
341 .now_tick
342 .saturating_sub(profile.provenance.measured_at_tick)
343 > profile.provenance.stale_after_ticks
344 {
345 return (AutoRouteDecision::Stale, None);
346 }
347 if !profile.is_conclusive() {
348 return (AutoRouteDecision::Inconclusive, None);
349 }
350 (AutoRouteDecision::Device, Some(profile.modeled_profile()))
351 }
352}
353
354#[derive(Clone, Debug, PartialEq, Eq)]
356pub struct AutoRoutingEvent {
357 pub provider: String,
359 pub decision: AutoRouteDecision,
361 pub materialization_bytes: u64,
363 pub synchronizations: usize,
365}
366
367#[derive(Clone, Debug, Default)]
369pub struct AutoRoutingLedger {
370 events: Arc<Mutex<Vec<AutoRoutingEvent>>>,
371}
372
373impl AutoRoutingLedger {
374 pub fn record(&self, event: AutoRoutingEvent) {
376 self.events
377 .lock()
378 .expect("auto routing ledger poisoned")
379 .push(event);
380 }
381
382 pub fn events(&self) -> Vec<AutoRoutingEvent> {
384 self.events
385 .lock()
386 .expect("auto routing ledger poisoned")
387 .clone()
388 }
389}
390
391pub(crate) fn request_materialization_bytes(request: &TensorRequest) -> u64 {
392 request
393 .inputs
394 .iter()
395 .map(|tensor| tensor.shape().iter().copied().product::<usize>() as u64 * CELL_BYTES)
396 .chain(std::iter::once(
397 request.output.shape().iter().copied().product::<usize>() as u64 * CELL_BYTES,
398 ))
399 .sum()
400}
401
402fn bounded_nonzero(values: Vec<u64>, max: u64) -> Vec<u64> {
403 let mut seen = BTreeSet::new();
404 values
405 .into_iter()
406 .filter(|value| *value > 0 && *value <= max)
407 .filter(|value| seen.insert(*value))
408 .collect()
409}