1#![forbid(unsafe_code)]
2#![deny(missing_docs)]
3mod 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
28pub fn auto_executor_symbol() -> Symbol {
30 Symbol::qualified("compute", "executor/auto")
31}
32
33pub fn compute_auto_site_symbol() -> Symbol {
35 Symbol::new("site/compute/auto")
36}
37
38pub fn compute_auto_lib_symbol() -> Symbol {
40 Symbol::qualified("compute", "auto-lib")
41}
42
43#[derive(Clone, Debug, Default, PartialEq, Eq)]
45pub struct AutoComputeProfile {
46 pub modeled: Option<ModeledComputeProfile>,
49 pub measured: Option<MeasuredComputeProfile>,
51 pub expected: Option<ComputeDeviceIdentity>,
53 pub now_tick: u64,
55}
56
57#[derive(Clone)]
59pub struct AutoTensorExecutor {
60 modeled: Option<ModeledTensorExecutor>,
61 decision: AutoRouteDecision,
62 ledger: AutoRoutingLedger,
63}
64
65impl AutoTensorExecutor {
66 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 pub fn uses_cpu_fallback(&self) -> bool {
85 self.modeled.is_none()
86 }
87
88 pub fn route_decision(&self) -> AutoRouteDecision {
90 self.decision.clone()
91 }
92
93 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#[derive(Clone, Debug, Default)]
165pub struct ComputeAutoLib {
166 profile: AutoComputeProfile,
167}
168
169impl ComputeAutoLib {
170 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
203pub static RECIPES: sim_cookbook::EmbeddedDir =
205 include!(concat!(env!("OUT_DIR"), "/cookbook_recipes.rs"));
206
207#[cfg(test)]
208mod tests;