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
10use crate::{ComputeEvidenceKind, ComputePhysicalEvidence, verify_physical};
11
12const DEFAULT_STALE_AFTER_TICKS: u64 = 10_000;
13const CELL_BYTES: u64 = 8;
14
15/// Adapter, driver, and backend identity for measured compute evidence.
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct ComputeDeviceIdentity {
18    /// Stable adapter label.
19    pub adapter: String,
20    /// Driver label or version string.
21    pub driver: String,
22    /// Backend label, such as `wgpu`, `cuda`, `rocm`, or `modeled`.
23    pub backend: String,
24}
25
26impl ComputeDeviceIdentity {
27    /// Builds a device identity.
28    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/// Hardware and scheduling limits captured with a measured profile.
46#[derive(Clone, Debug, PartialEq, Eq)]
47pub struct ComputeProfileLimits {
48    /// Maximum resident bytes accepted by the provider.
49    pub max_resident_bytes: u64,
50    /// Maximum storage binding bytes accepted by the provider.
51    pub max_storage_binding_bytes: u64,
52    /// Maximum queued submissions.
53    pub max_queue_depth: usize,
54    /// Maximum queued bytes.
55    pub max_queue_bytes: u64,
56    /// Maximum submission deadline in logical ticks.
57    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/// Transfer, launch, arithmetic, reduction, and matmul sample distributions.
73#[derive(Clone, Debug, PartialEq, Eq)]
74pub struct ComputeProfileSamples {
75    /// Upload bandwidth samples, in bytes per logical tick.
76    pub upload_bytes_per_tick: Vec<u64>,
77    /// Download bandwidth samples, in bytes per logical tick.
78    pub download_bytes_per_tick: Vec<u64>,
79    /// Launch latency samples, in logical ticks.
80    pub launch_ticks: Vec<u64>,
81    /// Element-wise throughput samples, in elements per logical tick.
82    pub element_elements_per_tick: Vec<u64>,
83    /// Reduction throughput samples, in elements per logical tick.
84    pub reduction_elements_per_tick: Vec<u64>,
85    /// Matmul throughput samples, in multiply-adds per logical tick.
86    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/// Thermal and power context captured beside benchmark samples.
105#[derive(Clone, Debug, PartialEq, Eq)]
106pub struct ComputeThermalPowerContext {
107    /// Thermal context label.
108    pub thermal: String,
109    /// Power context label.
110    pub power: String,
111}
112
113/// Provenance for a measured profile.
114#[derive(Clone, Debug, PartialEq, Eq)]
115pub struct ComputeProfileProvenance {
116    /// Evidence kind for this profile.
117    pub evidence_kind: ComputeEvidenceKind,
118    /// Tool or library that produced the profile.
119    pub producer: String,
120    /// Logical measurement tick.
121    pub measured_at_tick: u64,
122    /// Profile validity horizon in logical ticks.
123    pub stale_after_ticks: u64,
124    /// Device identity observed by the producer before caller naming.
125    pub observed_identity: Option<ComputeDeviceIdentity>,
126}
127
128/// Checked measured profile used by automatic routing.
129#[derive(Clone, Debug, PartialEq, Eq)]
130pub struct MeasuredComputeProfile {
131    /// Device identity.
132    pub identity: ComputeDeviceIdentity,
133    /// Provider limits.
134    pub limits: ComputeProfileLimits,
135    /// Transfer and operation sample distributions.
136    pub samples: ComputeProfileSamples,
137    /// Thermal and power measurement context.
138    pub context: ComputeThermalPowerContext,
139    /// Selected tile byte size.
140    pub tile_bytes: u64,
141    /// Allocation probe byte sizes that succeeded.
142    pub allocation_bytes: Vec<u64>,
143    /// Provenance for this record.
144    pub provenance: ComputeProfileProvenance,
145    /// Modeled profile to use when the evidence is accepted.
146    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    /// Returns true when the profile has usable, bounded measurement evidence.
179    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    /// Returns true when this profile still applies to `identity` at `now_tick`.
187    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    /// Returns the modeled executor profile selected by this measured evidence.
196    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
233/// Citizen class symbol for measured compute profile read-construct records.
234pub fn measured_compute_profile_citizen_symbol() -> Symbol {
235    Symbol::qualified("compute-profile", "MeasuredProfile")
236}
237
238/// Shape symbol for checked measured compute profile records.
239pub fn measured_compute_profile_shape_symbol() -> Symbol {
240    Symbol::qualified("compute-profile", "MeasuredProfileShape")
241}
242
243/// Bounded synthetic benchmark inputs.
244#[derive(Clone, Debug, PartialEq, Eq)]
245pub struct BenchmarkBounds {
246    /// Transfer byte sizes to sample.
247    pub transfer_bytes: Vec<u64>,
248    /// Element counts to sample.
249    pub element_counts: Vec<u64>,
250    /// Square matrix edges to sample.
251    pub matrix_edges: Vec<u64>,
252    /// Maximum accepted byte size in any sample.
253    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
267/// Runs a deterministic bounded profile harness from supplied provider facts.
268pub 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/// Reason an automatic route selected CPU or device placement.
324#[derive(Clone, Debug, PartialEq, Eq)]
325pub enum AutoRouteDecision {
326    /// No measured profile was supplied or loaded.
327    Absent,
328    /// The measured profile is stale for the current logical tick.
329    Stale,
330    /// Adapter, driver, or backend identity does not match.
331    Incompatible,
332    /// The profile lacks required bounded samples or limits.
333    Inconclusive,
334    /// The profile is modeled, host-emulated, or caller-renamed.
335    NonPhysical,
336    /// Device evidence was accepted.
337    Device,
338}
339
340/// Auto router that accepts device placement only with fresh compatible evidence.
341#[derive(Clone, Debug)]
342pub struct AutoComputeRouter {
343    expected: ComputeDeviceIdentity,
344    now_tick: u64,
345}
346
347impl AutoComputeRouter {
348    /// Builds a router for `expected` device identity at `now_tick`.
349    pub fn new(expected: ComputeDeviceIdentity, now_tick: u64) -> Self {
350        Self { expected, now_tick }
351    }
352
353    /// Returns the decision and modeled profile when device placement is proven.
354    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/// One provider-selection ledger row.
382#[derive(Clone, Debug, PartialEq, Eq)]
383pub struct AutoRoutingEvent {
384    /// Provider label chosen for the request or flush.
385    pub provider: String,
386    /// Routing decision.
387    pub decision: AutoRouteDecision,
388    /// Estimated materialization bytes.
389    pub materialization_bytes: u64,
390    /// Synchronization count represented by the event.
391    pub synchronizations: usize,
392}
393
394/// In-memory routing ledger for explainable placement decisions.
395#[derive(Clone, Debug, Default)]
396pub struct AutoRoutingLedger {
397    events: Arc<Mutex<Vec<AutoRoutingEvent>>>,
398}
399
400impl AutoRoutingLedger {
401    /// Records one routing event.
402    pub fn record(&self, event: AutoRoutingEvent) {
403        self.events
404            .lock()
405            .expect("auto routing ledger poisoned")
406            .push(event);
407    }
408
409    /// Returns recorded routing events.
410    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}