Skip to main content

sim_lib_compute_auto/
lib.rs

1#![forbid(unsafe_code)]
2#![deny(missing_docs)]
3//! Automatic tensor compute-site selection.
4
5mod profile;
6mod store;
7
8use std::sync::Arc;
9
10use sim_kernel::{
11    AbiVersion, DefaultFactory, Export, Factory, Lib, LibManifest, LibTarget, Linker, Result,
12    Symbol, Version,
13};
14use sim_lib_compute_model::{ModeledComputeProfile, ModeledTensorExecutor};
15use sim_lib_numbers_tensor::{
16    CpuTensorExecutor, SubmissionEvidence, TensorExecError, TensorExecution, TensorExecutor,
17    TensorExecutorCard, TensorRequest, TensorSite,
18};
19
20pub use profile::{
21    AutoComputeRouter, AutoRouteDecision, AutoRoutingEvent, AutoRoutingLedger, BenchmarkBounds,
22    ComputeDeviceIdentity, ComputeProfileLimits, ComputeProfileProvenance, ComputeProfileSamples,
23    ComputeThermalPowerContext, MeasuredComputeProfile, measure_bounded_profile,
24    measured_compute_profile_citizen_symbol, measured_compute_profile_shape_symbol,
25};
26pub use store::{ProfileStore, ProfileStorePolicy};
27
28/// Stable symbol for the automatic tensor executor.
29pub fn auto_executor_symbol() -> Symbol {
30    Symbol::qualified("compute", "executor/auto")
31}
32
33/// Site symbol exported by the automatic compute provider.
34pub fn compute_auto_site_symbol() -> Symbol {
35    Symbol::new("site/compute/auto")
36}
37
38/// Runtime library symbol for the automatic compute provider.
39pub fn compute_auto_lib_symbol() -> Symbol {
40    Symbol::qualified("compute", "auto-lib")
41}
42
43/// Automatic executor profile.
44#[derive(Clone, Debug, Default, PartialEq, Eq)]
45pub struct AutoComputeProfile {
46    /// Optional legacy modeled profile. When absent, auto uses CPU unless
47    /// measured evidence below proves a compatible device route.
48    pub modeled: Option<ModeledComputeProfile>,
49    /// Optional measured profile used by the conservative router.
50    pub measured: Option<MeasuredComputeProfile>,
51    /// Expected adapter/driver/backend identity for measured routing.
52    pub expected: Option<ComputeDeviceIdentity>,
53    /// Current logical tick used for staleness checks.
54    pub now_tick: u64,
55}
56
57/// Tensor executor that picks the best compatible provider and falls back to CPU.
58#[derive(Clone)]
59pub struct AutoTensorExecutor {
60    modeled: Option<ModeledTensorExecutor>,
61    decision: AutoRouteDecision,
62    ledger: AutoRoutingLedger,
63}
64
65impl AutoTensorExecutor {
66    /// Builds an automatic executor from a profile.
67    pub fn new(profile: AutoComputeProfile) -> Self {
68        let (decision, measured_modeled) = match profile.expected.clone() {
69            Some(expected) => {
70                AutoComputeRouter::new(expected, profile.now_tick).choose(profile.measured.as_ref())
71            }
72            None if profile.measured.is_some() => (AutoRouteDecision::Incompatible, None),
73            None => (AutoRouteDecision::Absent, None),
74        };
75        let modeled = measured_modeled.or(profile.modeled);
76        Self {
77            modeled: modeled.map(ModeledTensorExecutor::new),
78            decision,
79            ledger: AutoRoutingLedger::default(),
80        }
81    }
82
83    /// Returns true when the selector is using the CPU fallback.
84    pub fn uses_cpu_fallback(&self) -> bool {
85        self.modeled.is_none()
86    }
87
88    /// Returns the router decision that produced this executor.
89    pub fn route_decision(&self) -> AutoRouteDecision {
90        self.decision.clone()
91    }
92
93    /// Returns routing ledger events recorded by execute and flush calls.
94    pub fn routing_events(&self) -> Vec<AutoRoutingEvent> {
95        self.ledger.events()
96    }
97
98    fn selected(&self) -> Arc<dyn TensorExecutor> {
99        self.modeled
100            .clone()
101            .map(|executor| Arc::new(executor) as Arc<dyn TensorExecutor>)
102            .unwrap_or_else(|| Arc::new(CpuTensorExecutor::new()))
103    }
104}
105
106impl TensorExecutor for AutoTensorExecutor {
107    fn card(&self) -> TensorExecutorCard {
108        let selected = self.selected().card();
109        TensorExecutorCard::new(
110            auto_executor_symbol(),
111            if self.uses_cpu_fallback() {
112                "auto/cpu"
113            } else {
114                "auto/modeled"
115            },
116            Symbol::qualified("compute", "auto"),
117            selected.operations.to_vec(),
118            selected.device_capability,
119        )
120    }
121
122    fn execute(
123        &self,
124        cx: &mut sim_kernel::Cx,
125        request: TensorRequest,
126    ) -> std::result::Result<TensorExecution, TensorExecError> {
127        let materialization_bytes = profile::request_materialization_bytes(&request);
128        self.ledger.record(AutoRoutingEvent {
129            provider: if self.uses_cpu_fallback() {
130                "auto/cpu".to_owned()
131            } else {
132                "auto/modeled".to_owned()
133            },
134            decision: self.decision.clone(),
135            materialization_bytes,
136            synchronizations: 0,
137        });
138        self.selected().execute(cx, request)
139    }
140
141    fn flush(&self) -> std::result::Result<SubmissionEvidence, TensorExecError> {
142        let evidence = self.selected().flush()?;
143        self.ledger.record(AutoRoutingEvent {
144            provider: if self.uses_cpu_fallback() {
145                "auto/cpu".to_owned()
146            } else {
147                "auto/modeled".to_owned()
148            },
149            decision: self.decision.clone(),
150            materialization_bytes: 0,
151            synchronizations: 1,
152        });
153        Ok(evidence)
154    }
155}
156
157impl Default for AutoTensorExecutor {
158    fn default() -> Self {
159        Self::new(AutoComputeProfile::default())
160    }
161}
162
163/// Loadable library that registers the automatic compute site.
164#[derive(Clone, Debug, Default)]
165pub struct ComputeAutoLib {
166    profile: AutoComputeProfile,
167}
168
169impl ComputeAutoLib {
170    /// Builds an automatic compute library from a profile.
171    pub fn new(profile: AutoComputeProfile) -> Self {
172        Self { profile }
173    }
174}
175
176impl Lib for ComputeAutoLib {
177    fn manifest(&self) -> LibManifest {
178        LibManifest {
179            id: compute_auto_lib_symbol(),
180            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
181            abi: AbiVersion { major: 0, minor: 1 },
182            target: LibTarget::HostRegistered,
183            requires: Vec::new(),
184            capabilities: Vec::new(),
185            exports: vec![Export::Site {
186                symbol: compute_auto_site_symbol(),
187                runtime_id: None,
188            }],
189        }
190    }
191
192    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> Result<()> {
193        let executor = Arc::new(AutoTensorExecutor::new(self.profile.clone()));
194        let site = TensorSite::new(compute_auto_site_symbol(), executor, Vec::new());
195        linker.site_value(
196            compute_auto_site_symbol(),
197            DefaultFactory.opaque(Arc::new(site))?,
198        )?;
199        Ok(())
200    }
201}
202
203/// Cookbook recipes for this lib, embedded at build time.
204pub static RECIPES: sim_cookbook::EmbeddedDir =
205    include!(concat!(env!("OUT_DIR"), "/cookbook_recipes.rs"));
206
207#[cfg(test)]
208mod tests;