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                    .as_chunks::<4>()
137                    .0
138                    .iter()
139                    .take(self.len)
140                    .map(|bytes| {
141                        let value = f32::from_ne_bytes(*bytes);
142                        DefaultFactory.number_literal(self.dtype.clone(), value.to_string())
143                    })
144                    .collect::<Result<Vec<_>>>()?;
145                Ok(Arc::new(BoxedWgpuStorage::new(
146                    self.dtype.clone(),
147                    cells.into(),
148                )))
149            })
150            .clone()
151    }
152
153    fn as_any(&self) -> &dyn Any {
154        self
155    }
156}
157
158struct BoxedWgpuStorage {
159    dtype: Symbol,
160    cells: Arc<[Value]>,
161}
162
163impl BoxedWgpuStorage {
164    fn new(dtype: Symbol, cells: Arc<[Value]>) -> Self {
165        Self { dtype, cells }
166    }
167}
168
169impl TensorStorage for BoxedWgpuStorage {
170    fn dtype(&self) -> &Symbol {
171        &self.dtype
172    }
173
174    fn len(&self) -> usize {
175        self.cells.len()
176    }
177
178    fn location(&self) -> TensorLocation {
179        TensorLocation::Host
180    }
181
182    fn cell(&self, index: usize) -> Result<Value> {
183        self.cells
184            .get(index)
185            .cloned()
186            .ok_or_else(|| Error::Eval("wgpu tensor cell index was out of bounds".to_owned()))
187    }
188
189    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
190        Ok(Arc::new(Self {
191            dtype: self.dtype.clone(),
192            cells: self.cells.clone(),
193        }))
194    }
195
196    fn as_any(&self) -> &dyn Any {
197        self
198    }
199}
200
201impl WgpuArenaAllocation {
202    fn id_symbol(&self) -> Symbol {
203        Symbol::qualified("compute.alloc.wgpu", format!("{:?}", self.id))
204    }
205}
206
207impl sim_kernel::Object for WgpuArenaAllocation {
208    fn display(&self, _cx: &mut sim_kernel::Cx) -> Result<String> {
209        Ok(format!("#<wgpu-allocation {:?}>", self.id))
210    }
211
212    fn as_any(&self) -> &dyn Any {
213        self
214    }
215}
216
217impl sim_kernel::ObjectCompat for WgpuArenaAllocation {
218    fn class(&self, _cx: &mut sim_kernel::Cx) -> Result<sim_kernel::ClassRef> {
219        DefaultFactory.class_stub(
220            sim_kernel::CORE_FUNCTION_CLASS_ID,
221            Symbol::qualified("compute", "WgpuAllocation"),
222        )
223    }
224}
sim_lib_compute_wgpu/storage.rs

sim_lib_compute_wgpu/
storage.rs
