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::{TensorLocation, TensorStorage};
10
11use crate::{
12    WgpuArenaAllocation, WgpuMaterializationCache, WgpuPhysicalCounters, WgpuResidentSegment,
13    dispatch::{read_f32s, readback_buffer},
14    site::WgpuExecutionContext,
15};
16
17/// Resident storage that records a validated wgpu pipeline result.
18pub struct WgpuResidentStorage {
19    site: Symbol,
20    allocation: WgpuArenaAllocation,
21    pipeline: Symbol,
22    segments: Arc<[WgpuResidentSegment]>,
23    dtype: Symbol,
24    len: usize,
25    buffer: Arc<wgpu::Buffer>,
26    context: WgpuExecutionContext,
27    counters: WgpuPhysicalCounters,
28    cache: WgpuMaterializationCache,
29    materialized: OnceLock<Result<Arc<dyn TensorStorage>>>,
30}
31
32pub(crate) struct WgpuResidentStorageDescriptor {
33    pub(crate) site: Symbol,
34    pub(crate) allocation: WgpuArenaAllocation,
35    pub(crate) pipeline: Symbol,
36    pub(crate) segments: Vec<WgpuResidentSegment>,
37    pub(crate) dtype: Symbol,
38    pub(crate) len: usize,
39    pub(crate) buffer: Arc<wgpu::Buffer>,
40    pub(crate) context: WgpuExecutionContext,
41    pub(crate) counters: WgpuPhysicalCounters,
42}
43
44impl WgpuResidentStorage {
45    /// Builds resident storage around a real device buffer.
46    pub(crate) fn new(descriptor: WgpuResidentStorageDescriptor) -> Self {
47        Self {
48            site: descriptor.site,
49            allocation: descriptor.allocation,
50            pipeline: descriptor.pipeline,
51            segments: descriptor.segments.into(),
52            dtype: descriptor.dtype,
53            len: descriptor.len,
54            buffer: descriptor.buffer,
55            context: descriptor.context,
56            counters: descriptor.counters,
57            cache: WgpuMaterializationCache::default(),
58            materialized: OnceLock::new(),
59        }
60    }
61
62    /// Returns the resident allocation.
63    pub fn allocation(&self) -> &WgpuArenaAllocation {
64        &self.allocation
65    }
66
67    /// Returns the validated pipeline symbol.
68    pub fn pipeline(&self) -> &Symbol {
69        &self.pipeline
70    }
71
72    /// Returns the segmented resident layout.
73    pub fn segments(&self) -> &[WgpuResidentSegment] {
74        &self.segments
75    }
76
77    /// Returns the bindable resident device buffer.
78    pub(crate) fn buffer(&self) -> Arc<wgpu::Buffer> {
79        self.buffer.clone()
80    }
81
82    /// Returns the retained device context that owns the buffer.
83    pub(crate) fn context(&self) -> &WgpuExecutionContext {
84        &self.context
85    }
86}
87
88impl TensorStorage for WgpuResidentStorage {
89    fn dtype(&self) -> &Symbol {
90        &self.dtype
91    }
92
93    fn len(&self) -> usize {
94        self.len
95    }
96
97    fn location(&self) -> TensorLocation {
98        TensorLocation::Resident {
99            site: self.site.clone(),
100            allocation: self.allocation.id_symbol(),
101        }
102    }
103
104    fn cell(&self, index: usize) -> Result<Value> {
105        self.materialize()?.cell(index)
106    }
107
108    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
109        self.materialized
110            .get_or_init(|| {
111                self.counters.record_full_readback();
112                let size = self.allocation.bytes.max(4);
113                let readback = readback_buffer(&self.context.device, size, "materialize");
114                let mut encoder =
115                    self.context
116                        .device
117                        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
118                            label: Some("sim-compute-wgpu-materialize-encoder"),
119                        });
120                encoder.copy_buffer_to_buffer(&self.buffer, 0, &readback, 0, size);
121                self.context.queue.submit([encoder.finish()]);
122                self.counters.record_submit(&self.segments);
123                let values = self
124                    .cache
125                    .get_or_try_init(|| {
126                        let values = read_f32s(&self.context, &readback, self.len)
127                            .map_err(|err| err.to_string())?;
128                        Ok(values
129                            .iter()
130                            .flat_map(|value| value.to_ne_bytes())
131                            .collect::<Vec<_>>()
132                            .into())
133                    })
134                    .map_err(Error::Eval)?;
135                let cells = values
136                    .chunks_exact(4)
137                    .take(self.len)
138                    .map(|bytes| {
139                        let value = f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
140                        DefaultFactory.number_literal(self.dtype.clone(), value.to_string())
141                    })
142                    .collect::<Result<Vec<_>>>()?;
143                Ok(Arc::new(BoxedWgpuStorage::new(
144                    self.dtype.clone(),
145                    cells.into(),
146                )))
147            })
148            .clone()
149    }
150
151    fn as_any(&self) -> &dyn Any {
152        self
153    }
154}
155
156struct BoxedWgpuStorage {
157    dtype: Symbol,
158    cells: Arc<[Value]>,
159}
160
161impl BoxedWgpuStorage {
162    fn new(dtype: Symbol, cells: Arc<[Value]>) -> Self {
163        Self { dtype, cells }
164    }
165}
166
167impl TensorStorage for BoxedWgpuStorage {
168    fn dtype(&self) -> &Symbol {
169        &self.dtype
170    }
171
172    fn len(&self) -> usize {
173        self.cells.len()
174    }
175
176    fn location(&self) -> TensorLocation {
177        TensorLocation::Host
178    }
179
180    fn cell(&self, index: usize) -> Result<Value> {
181        self.cells
182            .get(index)
183            .cloned()
184            .ok_or_else(|| Error::Eval("wgpu tensor cell index was out of bounds".to_owned()))
185    }
186
187    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
188        Ok(Arc::new(Self {
189            dtype: self.dtype.clone(),
190            cells: self.cells.clone(),
191        }))
192    }
193
194    fn as_any(&self) -> &dyn Any {
195        self
196    }
197}
198
199impl WgpuArenaAllocation {
200    fn id_symbol(&self) -> Symbol {
201        Symbol::qualified("compute.alloc.wgpu", format!("{:?}", self.id))
202    }
203}
204
205impl sim_kernel::Object for WgpuArenaAllocation {
206    fn display(&self, _cx: &mut sim_kernel::Cx) -> Result<String> {
207        Ok(format!("#<wgpu-allocation {:?}>", self.id))
208    }
209
210    fn as_any(&self) -> &dyn Any {
211        self
212    }
213}
214
215impl sim_kernel::ObjectCompat for WgpuArenaAllocation {
216    fn class(&self, _cx: &mut sim_kernel::Cx) -> Result<sim_kernel::ClassRef> {
217        DefaultFactory.class_stub(
218            sim_kernel::CORE_FUNCTION_CLASS_ID,
219            Symbol::qualified("compute", "WgpuAllocation"),
220        )
221    }
222}
sim_lib_compute_wgpu/storage.rs

sim_lib_compute_wgpu/
storage.rs
