sim_lib_compute_wgpu/
storage.rs

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

sim_lib_compute_wgpu/
storage.rs
