sim_lib_numbers_tensor/implementation/
storage.rs

1//! Open tensor storage and the boxed host-storage implementation.
2
3use std::any::Any;
4use std::sync::Arc;
5
6use half::{bf16, f16};
7use sim_kernel::{DefaultFactory, Error, Factory, Result, Symbol, Value};
8use sim_lib_numbers_core::domains;
9
10/// The observable placement of tensor storage.
11///
12/// Host storage can be read directly. Resident storage belongs to a loadable
13/// execution site and names an opaque allocation supplied by that site.
14#[derive(Clone, Debug, Eq, PartialEq)]
15pub enum TensorLocation {
16    /// Storage is directly observable by the host runtime.
17    Host,
18    /// Storage resides at a loadable execution site.
19    Resident {
20        /// The site that owns the allocation.
21        site: Symbol,
22        /// An opaque allocation name meaningful to that site.
23        allocation: Symbol,
24    },
25}
26
27/// Storage behind the canonical [`Tensor`](super::value::Tensor) value.
28///
29/// Implementations may keep cells in host-native or resident layouts.
30/// Observation is fallible: a resident implementation performs and caches its
31/// checked readback in [`materialize`](TensorStorage::materialize). A successful
32/// materialization must return host storage with the same dtype and length.
33pub trait TensorStorage: Send + Sync + 'static {
34    /// The logical scalar domain of every cell.
35    fn dtype(&self) -> &Symbol;
36
37    /// The logical row-major cell count.
38    fn len(&self) -> usize;
39
40    /// Whether the storage has no logical cells.
41    fn is_empty(&self) -> bool {
42        self.len() == 0
43    }
44
45    /// The current storage placement.
46    fn location(&self) -> TensorLocation;
47
48    /// Observes one row-major cell.
49    fn cell(&self, index: usize) -> Result<Value>;
50
51    /// Returns a host-observable form of this storage.
52    ///
53    /// Resident implementations cache both success and failure so repeated or
54    /// concurrent observations perform at most one readback.
55    fn materialize(&self) -> Result<Arc<dyn TensorStorage>>;
56
57    /// Exposes the concrete storage for safe downcasting by typed adapters.
58    fn as_any(&self) -> &dyn Any;
59}
60
61/// Boxed host storage for scalar runtime values.
62///
63/// This is the default storage used by tensor constructors. Its cells are
64/// reference counted so tensor clones and repeated observations preserve
65/// value identity without copying the cell vector.
66pub struct BoxedTensorStorage {
67    dtype: Symbol,
68    cells: Arc<[Value]>,
69}
70
71impl BoxedTensorStorage {
72    pub(crate) fn new(dtype: Symbol, cells: Vec<Value>) -> Self {
73        Self {
74            dtype,
75            cells: cells.into(),
76        }
77    }
78
79    pub(crate) fn cells(&self) -> Arc<[Value]> {
80        self.cells.clone()
81    }
82}
83
84impl TensorStorage for BoxedTensorStorage {
85    fn dtype(&self) -> &Symbol {
86        &self.dtype
87    }
88
89    fn len(&self) -> usize {
90        self.cells.len()
91    }
92
93    fn location(&self) -> TensorLocation {
94        TensorLocation::Host
95    }
96
97    fn cell(&self, index: usize) -> Result<Value> {
98        self.cells
99            .get(index)
100            .cloned()
101            .ok_or_else(|| Error::Eval("tensor cell index was out of bounds".to_owned()))
102    }
103
104    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
105        Ok(Arc::new(Self {
106            dtype: self.dtype.clone(),
107            cells: self.cells.clone(),
108        }))
109    }
110
111    fn as_any(&self) -> &dyn Any {
112        self
113    }
114}
115
116/// A scalar cell type that can be held in typed host tensor storage.
117pub trait TensorCell: Clone + Send + Sync + 'static {
118    /// The scalar number domain this cell encodes.
119    fn dtype() -> Symbol;
120
121    /// Encodes the typed cell as a scalar runtime value.
122    fn to_value(&self) -> Result<Value>;
123}
124
125/// Host storage for typed scalar cells behind the canonical
126/// [`Tensor`](super::value::Tensor).
127///
128/// Typed tensor adapters use this storage so their public wrapper and the
129/// uniform `Tensor` share one storage lifetime. The adapter still gets a native
130/// cell slice for fast operations, while the runtime observes cells through the
131/// ordinary [`TensorStorage`] contract.
132pub struct TypedTensorStorage<T: TensorCell> {
133    dtype: Symbol,
134    cells: Arc<[T]>,
135}
136
137impl<T: TensorCell> TypedTensorStorage<T> {
138    /// Builds typed host storage from a flat row-major cell buffer.
139    pub fn new(cells: Vec<T>) -> Self {
140        Self::from_shared(cells.into())
141    }
142
143    /// Builds typed host storage from an already shared cell buffer.
144    pub fn from_shared(cells: Arc<[T]>) -> Self {
145        Self {
146            dtype: T::dtype(),
147            cells,
148        }
149    }
150
151    /// Borrows the typed row-major cells.
152    pub fn cell_slice(&self) -> &[T] {
153        &self.cells
154    }
155
156    /// Clones the shared typed row-major cells.
157    pub fn cells(&self) -> Arc<[T]> {
158        self.cells.clone()
159    }
160}
161
162impl<T: TensorCell> TensorStorage for TypedTensorStorage<T> {
163    fn dtype(&self) -> &Symbol {
164        &self.dtype
165    }
166
167    fn len(&self) -> usize {
168        self.cells.len()
169    }
170
171    fn location(&self) -> TensorLocation {
172        TensorLocation::Host
173    }
174
175    fn cell(&self, index: usize) -> Result<Value> {
176        self.cells
177            .get(index)
178            .ok_or_else(|| Error::Eval("tensor cell index was out of bounds".to_owned()))?
179            .to_value()
180    }
181
182    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
183        Ok(Arc::new(Self {
184            dtype: self.dtype.clone(),
185            cells: self.cells.clone(),
186        }))
187    }
188
189    fn as_any(&self) -> &dyn Any {
190        self
191    }
192}
193
194impl TensorCell for f64 {
195    fn dtype() -> Symbol {
196        domains::f64()
197    }
198
199    fn to_value(&self) -> Result<Value> {
200        DefaultFactory.number_literal(domains::f64(), self.to_string())
201    }
202}
203
204impl TensorCell for f32 {
205    fn dtype() -> Symbol {
206        domains::f32()
207    }
208
209    fn to_value(&self) -> Result<Value> {
210        DefaultFactory.number_literal(domains::f32(), self.to_string())
211    }
212}
213
214impl TensorCell for f16 {
215    fn dtype() -> Symbol {
216        domains::f16()
217    }
218
219    fn to_value(&self) -> Result<Value> {
220        DefaultFactory.number_literal(domains::f16(), self.to_f32().to_string())
221    }
222}
223
224impl TensorCell for bf16 {
225    fn dtype() -> Symbol {
226        domains::bf16()
227    }
228
229    fn to_value(&self) -> Result<Value> {
230        DefaultFactory.number_literal(domains::bf16(), self.to_f32().to_string())
231    }
232}
233
234impl TensorCell for i64 {
235    fn dtype() -> Symbol {
236        domains::i64()
237    }
238
239    fn to_value(&self) -> Result<Value> {
240        DefaultFactory.number_literal(domains::i64(), self.to_string())
241    }
242}
243
244impl TensorCell for bool {
245    fn dtype() -> Symbol {
246        domains::bool()
247    }
248
249    fn to_value(&self) -> Result<Value> {
250        DefaultFactory.number_literal(domains::bool(), self.to_string())
251    }
252}
253
254impl TensorCell for (f64, f64) {
255    fn dtype() -> Symbol {
256        domains::complex()
257    }
258
259    fn to_value(&self) -> Result<Value> {
260        DefaultFactory.number_literal(domains::complex(), format!("{}{:+}i", self.0, self.1))
261    }
262}
263
264impl TensorCell for (i64, i64) {
265    fn dtype() -> Symbol {
266        domains::rational()
267    }
268
269    fn to_value(&self) -> Result<Value> {
270        DefaultFactory.number_literal(domains::rational(), format!("{}/{}", self.0, self.1))
271    }
272}
sim_lib_numbers_tensor/implementation/storage.rs

sim_lib_numbers_tensor/implementation/
storage.rs
