spatialrust_gpu/
runtime.rs

1#[cfg(not(target_arch = "wasm32"))]
2use std::sync::Arc;
3use std::sync::OnceLock;
4
5#[cfg(feature = "gpu-image")]
6use std::collections::HashMap;
7#[cfg(feature = "gpu-image")]
8use std::sync::Mutex;
9
10use spatialrust_core::{DeviceKind, SpatialError, SpatialResult, SpatialRuntime};
11
12use crate::pipeline_cache::ComputePipelineCache;
13use crate::upload_cache::GpuBufferPool;
14
15/// Headless wgpu runtime for compute-only workloads.
16#[cfg(feature = "gpu-wgpu")]
17pub struct WgpuRuntime {
18    _instance: wgpu::Instance,
19    device: wgpu::Device,
20    queue: wgpu::Queue,
21    pipelines: OnceLock<ComputePipelineCache>,
22    max_gather_channels: u32,
23    upload_pool: GpuBufferPool,
24    adapter_info: WgpuAdapterInfo,
25    #[cfg(feature = "gpu-image")]
26    image_texture_pool: Mutex<HashMap<(u32, u32), Vec<wgpu::Texture>>>,
27    #[cfg(feature = "gpu-image")]
28    pub(crate) image_gray_pipeline: OnceLock<crate::image::kernels::gray::GrayPipeline>,
29    #[cfg(feature = "gpu-image")]
30    pub(crate) image_blur_pipeline: OnceLock<crate::image::kernels::box_blur::BlurPipeline>,
31    #[cfg(feature = "gpu-image")]
32    pub(crate) image_spatial_pipelines: OnceLock<crate::image::kernels::spatial::SpatialPipelines>,
33    #[cfg(feature = "gpu-image")]
34    pub(crate) image_ai_pack_pipeline: OnceLock<crate::image::ai_tensor::AiPackPipeline>,
35}
36
37/// Stable, serializable-friendly identity for the selected wgpu adapter.
38#[derive(Clone, Debug, PartialEq, Eq)]
39pub struct WgpuAdapterInfo {
40    /// Human-readable adapter name.
41    pub name: String,
42    /// Graphics API backend such as Vulkan, Metal, or DirectX 12.
43    pub backend: String,
44    /// Adapter class such as integrated, discrete, CPU, or virtual GPU.
45    pub device_type: String,
46    /// Driver name reported by wgpu.
47    pub driver: String,
48    /// Additional driver version/details reported by wgpu.
49    pub driver_info: String,
50}
51
52/// Adapter power preference used when creating a headless wgpu runtime.
53#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
54pub enum WgpuPowerPreference {
55    /// Prefer a discrete or otherwise high-performance adapter.
56    #[default]
57    HighPerformance,
58    /// Prefer an integrated or otherwise power-efficient adapter.
59    LowPower,
60}
61
62/// Minimum storage buffers required for the 4-channel gather kernel.
63#[cfg(feature = "gpu-wgpu")]
64pub const MULTI_GATHER4_STORAGE_BUFFERS: u32 = 10;
65
66/// Minimum storage buffers required for the 2-channel gather kernel.
67#[cfg(feature = "gpu-wgpu")]
68pub const MULTI_GATHER2_STORAGE_BUFFERS: u32 = 6;
69
70#[cfg(feature = "gpu-wgpu")]
71#[cfg(not(target_arch = "wasm32"))]
72static SHARED_RUNTIME: OnceLock<Result<Arc<WgpuRuntime>, String>> = OnceLock::new();
73
74#[cfg(feature = "gpu-wgpu")]
75impl WgpuRuntime {
76    /// Creates a headless wgpu runtime preferring a high-performance adapter.
77    ///
78    /// Prefer [`Self::shared`] when running multiple GPU filters in one process.
79    pub fn new_headless() -> SpatialResult<Self> {
80        Self::new_headless_with_preference(WgpuPowerPreference::HighPerformance)
81    }
82
83    /// Creates a headless wgpu runtime with an explicit adapter power preference.
84    pub fn new_headless_with_preference(preference: WgpuPowerPreference) -> SpatialResult<Self> {
85        pollster::block_on(Self::new_headless_async(preference))
86    }
87
88    /// Returns a process-wide shared headless runtime, initializing it on first use.
89    #[cfg(not(target_arch = "wasm32"))]
90    pub fn shared() -> SpatialResult<Arc<Self>> {
91        match SHARED_RUNTIME.get_or_init(init_shared_runtime) {
92            Ok(runtime) => Ok(Arc::clone(runtime)),
93            Err(message) => Err(SpatialError::InvalidArgument(message.clone())),
94        }
95    }
96
97    /// Returns the underlying wgpu device.
98    #[must_use]
99    pub fn device(&self) -> &wgpu::Device {
100        &self.device
101    }
102
103    /// Returns the underlying wgpu queue.
104    #[must_use]
105    pub fn queue(&self) -> &wgpu::Queue {
106        &self.queue
107    }
108
109    /// Returns the selected adapter/backend receipt.
110    #[must_use]
111    pub const fn adapter_info(&self) -> &WgpuAdapterInfo {
112        &self.adapter_info
113    }
114
115    /// Blocks until all previously submitted work on this device is complete.
116    ///
117    /// Normal device-resident chains do not synchronize implicitly. This is an
118    /// explicit profiling/testing boundary or host-coordination primitive.
119    pub fn wait_idle(&self) {
120        self.device.poll(wgpu::Maintain::Wait);
121    }
122
123    #[cfg(feature = "gpu-image")]
124    pub(crate) fn acquire_image_texture(
125        &self,
126        width: u32,
127        height: u32,
128        label: &'static str,
129    ) -> wgpu::Texture {
130        if let Some(texture) = self
131            .image_texture_pool
132            .lock()
133            .expect("image texture pool poisoned")
134            .get_mut(&(width, height))
135            .and_then(Vec::pop)
136        {
137            return texture;
138        }
139        self.device.create_texture(&wgpu::TextureDescriptor {
140            label: Some(label),
141            size: wgpu::Extent3d { width, height, depth_or_array_layers: 1 },
142            mip_level_count: 1,
143            sample_count: 1,
144            dimension: wgpu::TextureDimension::D2,
145            format: wgpu::TextureFormat::Rgba8Uint,
146            usage: wgpu::TextureUsages::COPY_SRC
147                | wgpu::TextureUsages::COPY_DST
148                | wgpu::TextureUsages::TEXTURE_BINDING
149                | wgpu::TextureUsages::STORAGE_BINDING,
150            view_formats: &[],
151        })
152    }
153
154    #[cfg(feature = "gpu-image")]
155    pub(crate) fn recycle_image_texture(&self, width: u32, height: u32, texture: wgpu::Texture) {
156        let mut pool = self.image_texture_pool.lock().expect("image texture pool poisoned");
157        let textures = pool.entry((width, height)).or_default();
158        if textures.len() < 8 {
159            textures.push(texture);
160        } else {
161            texture.destroy();
162        }
163    }
164
165    /// Returns the number of image textures retained for steady-state reuse.
166    #[cfg(feature = "gpu-image")]
167    #[must_use]
168    pub fn cached_image_texture_count(&self) -> usize {
169        self.image_texture_pool.lock().map(|pool| pool.values().map(Vec::len).sum()).unwrap_or(0)
170    }
171
172    /// Returns the number of initialized GPU image pipeline families.
173    #[cfg(feature = "gpu-image")]
174    #[must_use]
175    pub fn initialized_image_pipeline_count(&self) -> usize {
176        usize::from(self.image_gray_pipeline.get().is_some())
177            + usize::from(self.image_blur_pipeline.get().is_some())
178            + usize::from(self.image_spatial_pipelines.get().is_some())
179            + usize::from(self.image_ai_pack_pipeline.get().is_some())
180    }
181
182    /// Returns cached compute pipelines for this runtime's device.
183    #[must_use]
184    pub fn pipelines(&self) -> &ComputePipelineCache {
185        self.pipelines.get_or_init(|| ComputePipelineCache::new(&self.device))
186    }
187
188    /// Returns the maximum attribute channels gatherable in one multi dispatch.
189    #[must_use]
190    pub fn max_gather_channels(&self) -> u32 {
191        self.max_gather_channels
192    }
193
194    /// Returns the reusable storage-buffer pool owned by this runtime.
195    #[must_use]
196    pub fn buffer_pool(&self) -> &GpuBufferPool {
197        &self.upload_pool
198    }
199
200    /// Uploads a POD slice into a reusable pooled storage buffer.
201    pub fn upload_pod_storage<T: bytemuck::Pod>(
202        &self,
203        label: &'static str,
204        data: &[T],
205    ) -> SpatialResult<wgpu::Buffer> {
206        self.upload_pool.upload_pod_storage(self, label, data)
207    }
208
209    /// Uploads `f32` values into a reusable pooled storage buffer.
210    pub fn upload_f32_storage(
211        &self,
212        label: &'static str,
213        data: &[f32],
214    ) -> SpatialResult<wgpu::Buffer> {
215        self.upload_pool.upload_f32_storage(self, label, data)
216    }
217
218    /// Uploads `u32` values into a reusable pooled storage buffer.
219    pub fn upload_u32_storage(
220        &self,
221        label: &'static str,
222        data: &[u32],
223    ) -> SpatialResult<wgpu::Buffer> {
224        self.upload_pool.upload_u32_storage(self, label, data)
225    }
226
227    /// Returns a storage buffer to the upload pool for reuse.
228    pub fn recycle_storage(&self, byte_len: u64, buffer: wgpu::Buffer) {
229        self.upload_pool.recycle(byte_len, buffer);
230    }
231
232    /// Clears all buffers cached in the upload pool.
233    pub fn clear_buffer_pool(&self) {
234        self.upload_pool.clear();
235    }
236
237    /// Creates a headless runtime asynchronously.
238    ///
239    /// This is the browser/WebGPU construction path because blocking a WASM
240    /// main thread is not supported. Native callers may also use it from their
241    /// async executor.
242    pub async fn new_headless_async(preference: WgpuPowerPreference) -> SpatialResult<Self> {
243        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
244            backends: wgpu::Backends::PRIMARY,
245            ..Default::default()
246        });
247
248        let adapter = instance
249            .request_adapter(&wgpu::RequestAdapterOptions {
250                power_preference: match preference {
251                    WgpuPowerPreference::HighPerformance => wgpu::PowerPreference::HighPerformance,
252                    WgpuPowerPreference::LowPower => wgpu::PowerPreference::LowPower,
253                },
254                compatible_surface: None,
255                force_fallback_adapter: false,
256            })
257            .await
258            .ok_or_else(|| {
259                SpatialError::InvalidArgument(
260                    "no compatible wgpu adapter found for headless compute".to_owned(),
261                )
262            })?;
263
264        let raw_adapter_info = adapter.get_info();
265        let adapter_info = WgpuAdapterInfo {
266            name: raw_adapter_info.name,
267            backend: format!("{:?}", raw_adapter_info.backend),
268            device_type: format!("{:?}", raw_adapter_info.device_type),
269            driver: raw_adapter_info.driver,
270            driver_info: raw_adapter_info.driver_info,
271        };
272        let (device, queue) = adapter
273            .request_device(
274                &wgpu::DeviceDescriptor {
275                    label: Some("spatialrust-wgpu"),
276                    required_features: wgpu::Features::empty(),
277                    required_limits: adapter.limits(),
278                    memory_hints: wgpu::MemoryHints::Performance,
279                },
280                None,
281            )
282            .await
283            .map_err(|error| {
284                SpatialError::InvalidArgument(format!("failed to create wgpu device: {error}"))
285            })?;
286
287        let max_gather_channels =
288            max_gather_channels_for_limit(device.limits().max_storage_buffers_per_shader_stage);
289
290        Ok(Self {
291            _instance: instance,
292            device,
293            queue,
294            pipelines: OnceLock::new(),
295            max_gather_channels,
296            upload_pool: GpuBufferPool::default(),
297            adapter_info,
298            #[cfg(feature = "gpu-image")]
299            image_texture_pool: Mutex::new(HashMap::new()),
300            #[cfg(feature = "gpu-image")]
301            image_gray_pipeline: OnceLock::new(),
302            #[cfg(feature = "gpu-image")]
303            image_blur_pipeline: OnceLock::new(),
304            #[cfg(feature = "gpu-image")]
305            image_spatial_pipelines: OnceLock::new(),
306            #[cfg(feature = "gpu-image")]
307            image_ai_pack_pipeline: OnceLock::new(),
308        })
309    }
310}
311
312#[cfg(feature = "gpu-wgpu")]
313impl SpatialRuntime for WgpuRuntime {
314    fn device_kind(&self) -> DeviceKind {
315        DeviceKind::Wgpu
316    }
317}
318
319#[cfg(feature = "gpu-wgpu")]
320fn max_gather_channels_for_limit(storage_buffers_per_stage: u32) -> u32 {
321    if storage_buffers_per_stage >= MULTI_GATHER4_STORAGE_BUFFERS {
322        4
323    } else if storage_buffers_per_stage >= MULTI_GATHER2_STORAGE_BUFFERS {
324        2
325    } else {
326        1
327    }
328}
329
330#[cfg(feature = "gpu-wgpu")]
331#[cfg(not(target_arch = "wasm32"))]
332fn init_shared_runtime() -> Result<Arc<WgpuRuntime>, String> {
333    WgpuRuntime::new_headless().map(Arc::new).map_err(|error| error.to_string())
334}
335
336#[cfg(all(feature = "gpu-wgpu", test, not(target_arch = "wasm32")))]
337mod tests {
338    use super::WgpuRuntime;
339    use crate::pipeline_cache::ComputePipelineCache;
340    use std::sync::Arc;
341
342    #[test]
343    fn shared_runtime_is_singleton() {
344        let first = WgpuRuntime::shared().expect("shared runtime");
345        let second = WgpuRuntime::shared().expect("shared runtime");
346        assert!(Arc::ptr_eq(&first, &second));
347    }
348
349    #[test]
350    fn shared_and_headless_use_same_device_type() {
351        let shared = WgpuRuntime::shared().expect("shared runtime");
352        let local = WgpuRuntime::new_headless().expect("local runtime");
353        assert_eq!(
354            shared.device().limits().max_storage_buffers_per_shader_stage,
355            local.device().limits().max_storage_buffers_per_shader_stage
356        );
357    }
358
359    #[test]
360    fn pipeline_cache_is_initialized_once_per_runtime() {
361        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
362        let first = runtime.pipelines() as *const ComputePipelineCache;
363        let second = runtime.pipelines() as *const ComputePipelineCache;
364        assert_eq!(first, second);
365    }
366
367    #[test]
368    fn adapter_limits_enable_multi_channel_gather() {
369        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
370        let limit = runtime.device().limits().max_storage_buffers_per_shader_stage;
371        assert!(
372            limit >= super::MULTI_GATHER2_STORAGE_BUFFERS,
373            "expected at least {} storage buffers per stage, got {limit}",
374            super::MULTI_GATHER2_STORAGE_BUFFERS
375        );
376        assert!(runtime.max_gather_channels() >= 2);
377        assert_eq!(
378            runtime.max_gather_channels(),
379            runtime.pipelines().voxel_gather.multi_max_channels
380        );
381    }
382}
spatialrust_gpu/runtime.rs

spatialrust_gpu/
runtime.rs
