sim_lib_compute_cuda/
storage.rs

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

sim_lib_compute_cuda/
storage.rs
