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
11/// Opaque CUDA allocation evidence.
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct CudaAllocation {
14    /// Monotonic allocation id within the executor.
15    pub id: usize,
16    /// Logical resident byte count.
17    pub bytes: u64,
18    /// CUDA operation that produced the allocation.
19    pub operation: Symbol,
20}
21
22impl CudaAllocation {
23    fn id_symbol(&self) -> Symbol {
24        Symbol::qualified("compute.alloc.cuda", self.id.to_string())
25    }
26}
27
28/// Resident storage for a CUDA matmul result.
29pub struct CudaResidentStorage {
30    site: Symbol,
31    allocation: CudaAllocation,
32    shape: Arc<[usize]>,
33    dtype: Symbol,
34    cells: Arc<[Value]>,
35    materialized: OnceLock<Result<Arc<dyn TensorStorage>>>,
36}
37
38impl CudaResidentStorage {
39    /// Builds resident CUDA storage around checked host-equivalent cells.
40    pub fn new(
41        site: Symbol,
42        allocation: CudaAllocation,
43        shape: Vec<usize>,
44        dtype: Symbol,
45        cells: Arc<[Value]>,
46    ) -> Self {
47        Self {
48            site,
49            allocation,
50            shape: shape.into(),
51            dtype,
52            cells,
53            materialized: OnceLock::new(),
54        }
55    }
56
57    /// Returns the resident allocation evidence.
58    pub fn allocation(&self) -> &CudaAllocation {
59        &self.allocation
60    }
61
62    /// Rebuilds a canonical tensor without counting a user readback.
63    pub fn resident_tensor(&self) -> Option<Tensor> {
64        Tensor::from_storage(
65            self.shape.to_vec(),
66            self.dtype.clone(),
67            Arc::new(BoxedCudaStorage::new(
68                self.dtype.clone(),
69                self.cells.clone(),
70            )),
71        )
72        .ok()
73    }
74}
75
76impl TensorStorage for CudaResidentStorage {
77    fn dtype(&self) -> &Symbol {
78        &self.dtype
79    }
80
81    fn len(&self) -> usize {
82        self.cells.len()
83    }
84
85    fn location(&self) -> TensorLocation {
86        TensorLocation::Resident {
87            site: self.site.clone(),
88            allocation: self.allocation.id_symbol(),
89        }
90    }
91
92    fn cell(&self, index: usize) -> Result<Value> {
93        self.materialize()?.cell(index)
94    }
95
96    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
97        self.materialized
98            .get_or_init(|| {
99                Ok(Arc::new(BoxedCudaStorage::new(
100                    self.dtype.clone(),
101                    self.cells.clone(),
102                )))
103            })
104            .clone()
105    }
106
107    fn as_any(&self) -> &dyn Any {
108        self
109    }
110}
111
112struct BoxedCudaStorage {
113    dtype: Symbol,
114    cells: Arc<[Value]>,
115}
116
117impl BoxedCudaStorage {
118    fn new(dtype: Symbol, cells: Arc<[Value]>) -> Self {
119        Self { dtype, cells }
120    }
121}
122
123impl TensorStorage for BoxedCudaStorage {
124    fn dtype(&self) -> &Symbol {
125        &self.dtype
126    }
127
128    fn len(&self) -> usize {
129        self.cells.len()
130    }
131
132    fn location(&self) -> TensorLocation {
133        TensorLocation::Host
134    }
135
136    fn cell(&self, index: usize) -> Result<Value> {
137        self.cells
138            .get(index)
139            .cloned()
140            .ok_or_else(|| Error::Eval("cuda tensor cell index was out of bounds".to_owned()))
141    }
142
143    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
144        Ok(Arc::new(Self {
145            dtype: self.dtype.clone(),
146            cells: self.cells.clone(),
147        }))
148    }
149
150    fn as_any(&self) -> &dyn Any {
151        self
152    }
153}
154
155impl sim_kernel::Object for CudaAllocation {
156    fn display(&self, _cx: &mut sim_kernel::Cx) -> Result<String> {
157        Ok(format!("#<cuda-allocation {} {}b>", self.id, self.bytes))
158    }
159
160    fn as_any(&self) -> &dyn Any {
161        self
162    }
163}
164
165impl sim_kernel::ObjectCompat for CudaAllocation {
166    fn class(&self, _cx: &mut sim_kernel::Cx) -> Result<sim_kernel::ClassRef> {
167        DefaultFactory.class_stub(
168            sim_kernel::CORE_FUNCTION_CLASS_ID,
169            Symbol::qualified("compute", "CudaAllocation"),
170        )
171    }
172}
sim_lib_compute_cuda/storage.rs

sim_lib_compute_cuda/
storage.rs
