sim_lib_compute_model/
storage.rs

1//! Resident tensor storage for the modeled compute provider.
2
3use std::any::Any;
4use std::sync::{Arc, OnceLock};
5
6use sim_kernel::{DefaultFactory, Error, Factory, Result, Symbol, Value};
7use sim_lib_numbers_tensor::{Tensor, TensorLocation, TensorStorage};
8
9use crate::model::{ModeledComputeFault, ModeledResidentSegment, ModeledTensorExecutor};
10
11/// Opaque resident allocation handle owned by the modeled site.
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct ResidentHandle {
14    symbol: Symbol,
15}
16
17impl ResidentHandle {
18    pub(crate) fn new(id: usize) -> Self {
19        Self {
20            symbol: Symbol::qualified("compute.alloc", id.to_string()),
21        }
22    }
23
24    /// Returns the opaque allocation symbol.
25    pub fn symbol(&self) -> &Symbol {
26        &self.symbol
27    }
28}
29
30pub(crate) struct ModeledResidentDescriptor {
31    pub(crate) site: Symbol,
32    pub(crate) allocation: ResidentHandle,
33    pub(crate) segments: Vec<ModeledResidentSegment>,
34    pub(crate) shape: Vec<usize>,
35    pub(crate) dtype: Symbol,
36}
37
38/// Resident storage that models readback into host tensor cells.
39pub struct ModeledResidentStorage {
40    site: Symbol,
41    allocation: ResidentHandle,
42    segments: Arc<[ModeledResidentSegment]>,
43    shape: Arc<[usize]>,
44    dtype: Symbol,
45    cells: Arc<[Value]>,
46    executor: ModeledTensorExecutor,
47    fault: Option<ModeledComputeFault>,
48    materialized: OnceLock<Result<Arc<dyn TensorStorage>>>,
49}
50
51impl ModeledResidentStorage {
52    pub(crate) fn new(
53        descriptor: ModeledResidentDescriptor,
54        cells: Arc<[Value]>,
55        executor: ModeledTensorExecutor,
56        fault: Option<ModeledComputeFault>,
57    ) -> Self {
58        Self {
59            site: descriptor.site,
60            allocation: descriptor.allocation,
61            segments: descriptor.segments.into(),
62            shape: descriptor.shape.into(),
63            dtype: descriptor.dtype,
64            cells,
65            executor,
66            fault,
67            materialized: OnceLock::new(),
68        }
69    }
70
71    /// Returns the resident allocation handle.
72    pub fn allocation(&self) -> &ResidentHandle {
73        &self.allocation
74    }
75
76    /// Returns the segmented resident layout for this allocation.
77    pub fn segments(&self) -> &[ModeledResidentSegment] {
78        &self.segments
79    }
80
81    /// Rebuilds a canonical host tensor without counting a user readback.
82    pub fn resident_tensor(&self) -> Option<Tensor> {
83        Tensor::from_storage(
84            self.shape.to_vec(),
85            self.dtype.clone(),
86            Arc::new(BoxedTensorStorageForResident::new(
87                self.dtype.clone(),
88                self.cells.clone(),
89            )),
90        )
91        .ok()
92    }
93}
94
95impl TensorStorage for ModeledResidentStorage {
96    fn dtype(&self) -> &Symbol {
97        &self.dtype
98    }
99
100    fn len(&self) -> usize {
101        self.cells.len()
102    }
103
104    fn location(&self) -> TensorLocation {
105        TensorLocation::Resident {
106            site: self.site.clone(),
107            allocation: self.allocation.symbol().clone(),
108        }
109    }
110
111    fn cell(&self, index: usize) -> Result<Value> {
112        let storage = self.materialize()?;
113        storage.cell(index)
114    }
115
116    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
117        self.materialized
118            .get_or_init(|| {
119                self.executor.increment_readbacks();
120                if !self.executor.is_resident_active(&self.allocation) {
121                    self.executor.increment_materialization_failures();
122                    return Err(Error::Eval(
123                        "modeled compute resident allocation was evicted".to_owned(),
124                    ));
125                }
126                if self.fault == Some(ModeledComputeFault::ReadbackFailure) {
127                    self.executor.increment_materialization_failures();
128                    return Err(Error::Eval("modeled compute readback failed".to_owned()));
129                }
130                Ok(Arc::new(BoxedTensorStorageForResident::new(
131                    self.dtype.clone(),
132                    self.cells.clone(),
133                )))
134            })
135            .clone()
136    }
137
138    fn as_any(&self) -> &dyn Any {
139        self
140    }
141}
142
143struct BoxedTensorStorageForResident {
144    dtype: Symbol,
145    cells: Arc<[Value]>,
146}
147
148impl BoxedTensorStorageForResident {
149    fn new(dtype: Symbol, cells: Arc<[Value]>) -> Self {
150        Self { dtype, cells }
151    }
152}
153
154impl TensorStorage for BoxedTensorStorageForResident {
155    fn dtype(&self) -> &Symbol {
156        &self.dtype
157    }
158
159    fn len(&self) -> usize {
160        self.cells.len()
161    }
162
163    fn location(&self) -> TensorLocation {
164        TensorLocation::Host
165    }
166
167    fn cell(&self, index: usize) -> Result<Value> {
168        self.cells
169            .get(index)
170            .cloned()
171            .ok_or_else(|| Error::Eval("tensor cell index was out of bounds".to_owned()))
172    }
173
174    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
175        Ok(Arc::new(Self {
176            dtype: self.dtype.clone(),
177            cells: self.cells.clone(),
178        }))
179    }
180
181    fn as_any(&self) -> &dyn Any {
182        self
183    }
184}
185
186impl sim_kernel::Object for ResidentHandle {
187    fn display(&self, _cx: &mut sim_kernel::Cx) -> Result<String> {
188        Ok(format!("#<compute-resident {}>", self.symbol))
189    }
190
191    fn as_any(&self) -> &dyn Any {
192        self
193    }
194}
195
196impl sim_kernel::ObjectCompat for ResidentHandle {
197    fn class(&self, _cx: &mut sim_kernel::Cx) -> Result<sim_kernel::ClassRef> {
198        DefaultFactory.class_stub(
199            sim_kernel::CORE_FUNCTION_CLASS_ID,
200            Symbol::qualified("compute", "ResidentHandle"),
201        )
202    }
203
204    fn as_table(&self, cx: &mut sim_kernel::Cx) -> Result<Value> {
205        cx.factory().table(vec![(
206            Symbol::new("allocation"),
207            cx.factory().symbol(self.symbol.clone())?,
208        )])
209    }
210}
sim_lib_compute_model/storage.rs

sim_lib_compute_model/
storage.rs
