Skip to main content

sim_lib_compute_auto/
profile.rs

1//! Measured compute profiles, bounded persistence, and conservative routing.
2
3use 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/// Adapter, driver, and backend identity for measured compute evidence.
14#[derive(Clone, Debug, PartialEq, Eq)]
15pub struct ComputeDeviceIdentity {
16    /// Stable adapter label.
17    pub adapter: String,
18    /// Driver label or version string.
19    pub driver: String,
20    /// Backend label, such as `wgpu`, `cuda`, `rocm`, or `modeled`.
21    pub backend: String,
22}
23
24impl ComputeDeviceIdentity {
25    /// Builds a device identity.
26    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/// Hardware and scheduling limits captured with a measured profile.
44#[derive(Clone, Debug, PartialEq, Eq)]
45pub struct ComputeProfileLimits {
46    /// Maximum resident bytes accepted by the provider.
47    pub max_resident_bytes: u64,
48    /// Maximum storage binding bytes accepted by the provider.
49    pub max_storage_binding_bytes: u64,
50    /// Maximum queued submissions.
51    pub max_queue_depth: usize,
52    /// Maximum queued bytes.
53    pub max_queue_bytes: u64,
54    /// Maximum submission deadline in logical ticks.
55    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/// Transfer, launch, arithmetic, reduction, and matmul sample distributions.
71#[derive(Clone, Debug, PartialEq, Eq)]
72pub struct ComputeProfileSamples {
73    /// Upload bandwidth samples, in bytes per logical tick.
74    pub upload_bytes_per_tick: Vec<u64>,
75    /// Download bandwidth samples, in bytes per logical tick.
76    pub download_bytes_per_tick: Vec<u64>,
77    /// Launch latency samples, in logical ticks.
78    pub launch_ticks: Vec<u64>,
79    /// Element-wise throughput samples, in elements per logical tick.
80    pub element_elements_per_tick: Vec<u64>,
81    /// Reduction throughput samples, in elements per logical tick.
82    pub reduction_elements_per_tick: Vec<u64>,
83    /// Matmul throughput samples, in multiply-adds per logical tick.
84    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/// Thermal and power context captured beside benchmark samples.
103#[derive(Clone, Debug, PartialEq, Eq)]
104pub struct ComputeThermalPowerContext {
105    /// Thermal context label.
106    pub thermal: String,
107    /// Power context label.
108    pub power: String,
109}
110
111/// Provenance for a measured profile.
112#[derive(Clone, Debug, PartialEq, Eq)]
113pub struct ComputeProfileProvenance {
114    /// Tool or library that produced the profile.
115    pub producer: String,
116    /// Logical measurement tick.
117    pub measured_at_tick: u64,
118    /// Profile validity horizon in logical ticks.
119    pub stale_after_ticks: u64,
120}
121
122/// Checked measured profile used by automatic routing.
123#[derive(Clone, Debug, PartialEq, Eq)]
124pub struct MeasuredComputeProfile {
125    /// Device identity.
126    pub identity: ComputeDeviceIdentity,
127    /// Provider limits.
128    pub limits: ComputeProfileLimits,
129    /// Transfer and operation sample distributions.
130    pub samples: ComputeProfileSamples,
131    /// Thermal and power measurement context.
132    pub context: ComputeThermalPowerContext,
133    /// Selected tile byte size.
134    pub tile_bytes: u64,
135    /// Allocation probe byte sizes that succeeded.
136    pub allocation_bytes: Vec<u64>,
137    /// Provenance for this record.
138    pub provenance: ComputeProfileProvenance,
139    /// Modeled profile to use when the evidence is accepted.
140    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    /// Returns true when the profile has usable, bounded measurement evidence.
173    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    /// Returns true when this profile still applies to `identity` at `now_tick`.
181    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    /// Returns the modeled executor profile selected by this measured evidence.
190    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
213/// Citizen class symbol for measured compute profile read-construct records.
214pub fn measured_compute_profile_citizen_symbol() -> Symbol {
215    Symbol::qualified("compute-profile", "MeasuredProfile")
216}
217
218/// Shape symbol for checked measured compute profile records.
219pub fn measured_compute_profile_shape_symbol() -> Symbol {
220    Symbol::qualified("compute-profile", "MeasuredProfileShape")
221}
222
223/// Bounded synthetic benchmark inputs.
224#[derive(Clone, Debug, PartialEq, Eq)]
225pub struct BenchmarkBounds {
226    /// Transfer byte sizes to sample.
227    pub transfer_bytes: Vec<u64>,
228    /// Element counts to sample.
229    pub element_counts: Vec<u64>,
230    /// Square matrix edges to sample.
231    pub matrix_edges: Vec<u64>,
232    /// Maximum accepted byte size in any sample.
233    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
247/// Runs a deterministic bounded profile harness from supplied provider facts.
248pub 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/// Reason an automatic route selected CPU or device placement.
302#[derive(Clone, Debug, PartialEq, Eq)]
303pub enum AutoRouteDecision {
304    /// No measured profile was supplied or loaded.
305    Absent,
306    /// The measured profile is stale for the current logical tick.
307    Stale,
308    /// Adapter, driver, or backend identity does not match.
309    Incompatible,
310    /// The profile lacks required bounded samples or limits.
311    Inconclusive,
312    /// Device evidence was accepted.
313    Device,
314}
315
316/// Auto router that accepts device placement only with fresh compatible evidence.
317#[derive(Clone, Debug)]
318pub struct AutoComputeRouter {
319    expected: ComputeDeviceIdentity,
320    now_tick: u64,
321}
322
323impl AutoComputeRouter {
324    /// Builds a router for `expected` device identity at `now_tick`.
325    pub fn new(expected: ComputeDeviceIdentity, now_tick: u64) -> Self {
326        Self { expected, now_tick }
327    }
328
329    /// Returns the decision and modeled profile when device placement is proven.
330    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/// One provider-selection ledger row.
355#[derive(Clone, Debug, PartialEq, Eq)]
356pub struct AutoRoutingEvent {
357    /// Provider label chosen for the request or flush.
358    pub provider: String,
359    /// Routing decision.
360    pub decision: AutoRouteDecision,
361    /// Estimated materialization bytes.
362    pub materialization_bytes: u64,
363    /// Synchronization count represented by the event.
364    pub synchronizations: usize,
365}
366
367/// In-memory routing ledger for explainable placement decisions.
368#[derive(Clone, Debug, Default)]
369pub struct AutoRoutingLedger {
370    events: Arc<Mutex<Vec<AutoRoutingEvent>>>,
371}
372
373impl AutoRoutingLedger {
374    /// Records one routing event.
375    pub fn record(&self, event: AutoRoutingEvent) {
376        self.events
377            .lock()
378            .expect("auto routing ledger poisoned")
379            .push(event);
380    }
381
382    /// Returns recorded routing events.
383    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}