spatialrust_gpu/
upload_cache.rs

1//! Reusable wgpu storage-buffer pool for host→device uploads.
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5
6use bytemuck::Pod;
7use spatialrust_core::SpatialResult;
8
9use crate::runtime::WgpuRuntime;
10
11/// Reusable GPU storage buffers keyed by byte length.
12///
13/// Buffers are recycled through [`WgpuRuntime::recycle_storage`] or
14/// [`GpuBufferPool::recycle`]. Prefer the runtime helpers for typical use;
15/// access the pool directly when implementing custom kernels.
16#[derive(Default, Debug)]
17pub struct GpuBufferPool {
18    free: Mutex<HashMap<u64, Vec<wgpu::Buffer>>>,
19}
20
21impl GpuBufferPool {
22    /// Acquires an uninitialized reusable storage buffer of the requested size.
23    pub fn acquire_storage(
24        &self,
25        runtime: &WgpuRuntime,
26        label: &'static str,
27        byte_len: u64,
28    ) -> wgpu::Buffer {
29        self.take_storage(runtime, label, byte_len)
30    }
31
32    /// Uploads a POD slice into a pooled storage buffer.
33    pub fn upload_pod_storage<T: Pod>(
34        &self,
35        runtime: &WgpuRuntime,
36        label: &'static str,
37        data: &[T],
38    ) -> SpatialResult<wgpu::Buffer> {
39        let byte_len = std::mem::size_of_val(data) as u64;
40        let buffer = self.take_storage(runtime, label, byte_len);
41        runtime.queue().write_buffer(&buffer, 0, bytemuck::cast_slice(data));
42        Ok(buffer)
43    }
44
45    /// Uploads `f32` values into a pooled storage buffer.
46    pub fn upload_f32_storage(
47        &self,
48        runtime: &WgpuRuntime,
49        label: &'static str,
50        data: &[f32],
51    ) -> SpatialResult<wgpu::Buffer> {
52        self.upload_pod_storage(runtime, label, data)
53    }
54
55    /// Uploads `u32` values into a pooled storage buffer.
56    pub fn upload_u32_storage(
57        &self,
58        runtime: &WgpuRuntime,
59        label: &'static str,
60        data: &[u32],
61    ) -> SpatialResult<wgpu::Buffer> {
62        self.upload_pod_storage(runtime, label, data)
63    }
64
65    /// Returns a storage buffer to the pool for reuse.
66    pub fn recycle(&self, byte_len: u64, buffer: wgpu::Buffer) {
67        if byte_len == 0 {
68            return;
69        }
70        if let Ok(mut free) = self.free.lock() {
71            free.entry(byte_len).or_default().push(buffer);
72        }
73    }
74
75    /// Discards all cached buffers without returning them to the device allocator.
76    pub fn clear(&self) {
77        if let Ok(mut free) = self.free.lock() {
78            free.clear();
79        }
80    }
81
82    /// Returns the number of buffers currently held in the pool.
83    #[must_use]
84    pub fn cached_buffer_count(&self) -> usize {
85        self.free.lock().map(|free| free.values().map(Vec::len).sum()).unwrap_or(0)
86    }
87
88    fn take_storage(
89        &self,
90        runtime: &WgpuRuntime,
91        label: &'static str,
92        byte_len: u64,
93    ) -> wgpu::Buffer {
94        if byte_len == 0 {
95            return runtime.device().create_buffer(&wgpu::BufferDescriptor {
96                label: Some(label),
97                size: 4,
98                usage: wgpu::BufferUsages::STORAGE
99                    | wgpu::BufferUsages::COPY_DST
100                    | wgpu::BufferUsages::COPY_SRC,
101                mapped_at_creation: false,
102            });
103        }
104
105        if let Ok(mut free) = self.free.lock() {
106            if let Some(buffer) = free.get_mut(&byte_len).and_then(|buffers| buffers.pop()) {
107                return buffer;
108            }
109        }
110
111        runtime.device().create_buffer(&wgpu::BufferDescriptor {
112            label: Some(label),
113            size: byte_len,
114            usage: wgpu::BufferUsages::STORAGE
115                | wgpu::BufferUsages::COPY_DST
116                | wgpu::BufferUsages::COPY_SRC,
117            mapped_at_creation: false,
118        })
119    }
120}
121
122#[cfg(all(feature = "gpu-wgpu", test))]
123mod tests {
124    use super::GpuBufferPool;
125    use crate::runtime::WgpuRuntime;
126
127    #[test]
128    fn buffer_pool_reuses_equal_sized_buffers() {
129        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
130        let pool = GpuBufferPool::default();
131        let data = [1.0_f32, 2.0, 3.0];
132
133        let first =
134            pool.upload_f32_storage(&runtime, "upload-pool-test", &data).expect("first upload");
135        assert_eq!(pool.cached_buffer_count(), 0);
136        pool.recycle(first.size(), first);
137        assert_eq!(pool.cached_buffer_count(), 1);
138
139        let second =
140            pool.upload_f32_storage(&runtime, "upload-pool-test", &data).expect("second upload");
141        assert_eq!(second.size(), (data.len() * std::mem::size_of::<f32>()) as u64);
142        assert_eq!(pool.cached_buffer_count(), 0);
143    }
144
145    #[test]
146    fn runtime_buffer_pool_matches_direct_upload() {
147        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
148        let data = [4.0_f32, 5.0];
149        let buffer = runtime.upload_f32_storage("runtime-pool-test", &data).expect("upload");
150        runtime.recycle_storage(buffer.size(), buffer);
151        assert_eq!(runtime.buffer_pool().cached_buffer_count(), 1);
152        runtime.buffer_pool().clear();
153        assert_eq!(runtime.buffer_pool().cached_buffer_count(), 0);
154    }
155}
spatialrust_gpu/upload_cache.rs

spatialrust_gpu/
upload_cache.rs
