oxiphysics_gpu/compute/wgpu_backend.rs
1// Copyright 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3
4//! WebGPU (wgpu) compute backend for the OxiPhysics GPU acceleration layer.
5//!
6//! This module provides [`WgpuBackend`] which implements `ComputeBackend` using
7//! the `wgpu` crate for cross-platform GPU compute (Vulkan, Metal, DX12, WebGPU).
8//!
9//! ## Feature flag
10//!
11//! This module is gated behind the `wgpu-backend` Cargo feature:
12//!
13//! ```toml
14//! [dependencies]
15//! oxiphysics-gpu = { features = ["wgpu-backend"] }
16//! ```
17//!
18//! When the feature is disabled the module compiles to an empty stub. This allows
19//! the crate to compile without the `wgpu` dependency on platforms or toolchains
20//! where GPU support is not required.
21//!
22//! ## Enabling the dependency
23//!
24//! To activate the wgpu backend, add `wgpu` to the crate's `Cargo.toml`:
25//!
26//! ```toml
27//! [features]
28//! wgpu-backend = ["wgpu"]
29//!
30//! [dependencies]
31//! wgpu = { version = "0.20", optional = true }
32//! ```
33//!
34//! ## Architecture
35//!
36//! ```text
37//! WgpuBackend
38//! ├── wgpu::Device / wgpu::Queue ← GPU device & command queue
39//! ├── Vec<WgpuBufferEntry> ← Registered GPU buffers
40//! │ ├── wgpu::Buffer (device memory)
41//! │ └── size, usage flags
42//! └── ShaderRegistry ← Compiled WGSL compute shaders
43//!
44//! Compute pipeline:
45//! write_buffer → [upload via staging] → dispatch(kernel) → [readback via staging] → read_buffer
46//! ```
47//!
48//! ## Usage (when feature is enabled)
49//!
50//! ```ignore
51//! use oxiphysics_gpu::compute::wgpu_backend::WgpuBackend;
52//! use oxiphysics_gpu::compute::ComputeBackend;
53//!
54//! let backend = WgpuBackend::new_async().await?;
55//! let handle = backend.create_buffer(1024);
56//! backend.write_buffer(handle, &vec![1.0_f64; 128]);
57//! // ... dispatch kernel ...
58//! let data = backend.read_buffer(handle);
59//! ```
60
61// ── BufferHandle (re-used from parent module) ─────────────────────────────────
62
63/// Opaque handle to a GPU buffer allocated by a `ComputeBackend`.
64///
65/// This type mirrors the one in the parent `compute` module so that
66/// [`WgpuBackend`] can implement the same `ComputeBackend` trait.
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
68pub struct WgpuBufferHandle(pub usize);
69
70// ── WgpuDeviceInfo ────────────────────────────────────────────────────────────
71
72/// Information about the GPU device selected by the wgpu adapter.
73#[derive(Debug, Clone, Default)]
74pub struct WgpuDeviceInfo {
75 /// Human-readable device name (e.g. `"NVIDIA GeForce RTX 4090"`).
76 pub name: String,
77 /// Backend API in use: `"Vulkan"`, `"Metal"`, `"Dx12"`, `"WebGpu"`, or `"None"`.
78 pub backend: String,
79 /// Driver version string (if available).
80 pub driver_version: String,
81 /// Total VRAM in bytes (0 if not reported by the adapter).
82 pub vram_bytes: u64,
83 /// Whether the device supports 64-bit floating-point storage.
84 pub supports_f64: bool,
85 /// Maximum workgroup size (x, y, z).
86 pub max_workgroup_size: [u32; 3],
87}
88
89// ── WgpuBackend ───────────────────────────────────────────────────────────────
90
91/// WebGPU compute backend.
92///
93/// When compiled **without** the `wgpu-backend` feature this struct is a no-op
94/// stub that will return an error from [`WgpuBackend::try_new`]. When compiled
95/// **with** the feature, a real wgpu `Device` / `Queue` pair is created.
96///
97/// For the real implementation, `try_new` should be called within an async
98/// runtime (tokio or wasm-bindgen-futures for browser targets).
99#[derive(Debug)]
100pub struct WgpuBackend {
101 /// Device info (populated at initialisation).
102 pub device_info: WgpuDeviceInfo,
103 /// Allocated CPU-side buffers (mirrors GPU allocations).
104 ///
105 /// In the stub implementation these are plain `Vec<f64>` acting as
106 /// stand-ins for actual `wgpu::Buffer` objects. A full implementation
107 /// wraps `wgpu::Buffer` behind `Arc<Mutex<…>>` to allow async reads.
108 buffers: Vec<WgpuBufferEntry>,
109 /// Whether the backend is operational.
110 available: bool,
111}
112
113/// Internal buffer entry storing metadata and a CPU-side shadow copy.
114#[derive(Debug, Clone)]
115struct WgpuBufferEntry {
116 /// Byte capacity of the GPU buffer (8 × `len` for f64 arrays).
117 capacity: usize,
118 /// CPU-side shadow for upload/download (avoids wgpu dep in stub).
119 shadow: Vec<f64>,
120}
121
122impl WgpuBackend {
123 /// Attempt to create a wgpu backend.
124 ///
125 /// Returns `Ok(Self)` when a compatible GPU adapter is available, or
126 /// `Err(WgpuInitError::NotAvailable)` when no adapter can be found (e.g.
127 /// running headless without a GPU or without the `wgpu-backend` feature).
128 ///
129 /// In the current stub implementation this always returns a CPU-fallback
130 /// instance with `available = false`. The full implementation calls
131 /// `wgpu::Instance::request_adapter` and `adapter.request_device`.
132 pub fn try_new() -> Result<Self, WgpuInitError> {
133 // ── TODO (wgpu-backend feature) ─────────────────────────────────────
134 // When `wgpu-backend` is enabled, replace this stub with:
135 //
136 // let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
137 // backends: wgpu::Backends::all(),
138 // ..Default::default()
139 // });
140 // let adapter = pollster::block_on(instance.request_adapter(
141 // &wgpu::RequestAdapterOptions {
142 // power_preference: wgpu::PowerPreference::HighPerformance,
143 // ..Default::default()
144 // },
145 // )).ok_or(WgpuInitError::NoAdapter)?;
146 // let (device, queue) = pollster::block_on(adapter.request_device(
147 // &wgpu::DeviceDescriptor::default(),
148 // None,
149 // ))?;
150 // let info = adapter.get_info();
151 // Ok(Self { device, queue, info, buffers: Vec::new(), available: true })
152 // ────────────────────────────────────────────────────────────────────
153
154 Err(WgpuInitError::NotAvailable)
155 }
156
157 /// Create a stub backend for testing that stores data in CPU memory.
158 ///
159 /// This is equivalent to what `try_new` would return on a headless system
160 /// but without returning an error — useful for unit testing backend logic.
161 pub fn new_stub() -> Self {
162 Self {
163 device_info: WgpuDeviceInfo {
164 name: "CPU stub".to_string(),
165 backend: "None".to_string(),
166 ..Default::default()
167 },
168 buffers: Vec::new(),
169 available: false,
170 }
171 }
172
173 /// Return `true` if a real GPU device is available.
174 pub fn is_available(&self) -> bool {
175 self.available
176 }
177
178 /// Return device information for diagnostics.
179 pub fn device_info(&self) -> &WgpuDeviceInfo {
180 &self.device_info
181 }
182
183 // ── Buffer management ────────────────────────────────────────────────────
184
185 /// Allocate a GPU buffer that can hold `len` `f64` values.
186 ///
187 /// Returns a [`WgpuBufferHandle`] that can be passed to [`Self::write_buffer`]
188 /// and [`Self::read_buffer`].
189 ///
190 /// In the stub implementation, a CPU-side shadow `Vec<f64>` is allocated.
191 /// In the full wgpu implementation, `wgpu::Device::create_buffer` is called
192 /// with `STORAGE | COPY_SRC | COPY_DST` usage flags.
193 pub fn create_buffer(&mut self, len: usize) -> WgpuBufferHandle {
194 let handle = WgpuBufferHandle(self.buffers.len());
195 self.buffers.push(WgpuBufferEntry {
196 capacity: len,
197 shadow: vec![0.0; len],
198 });
199 handle
200 }
201
202 /// Upload `data` to the GPU buffer at `handle`.
203 ///
204 /// In the stub, data is copied into the CPU-side shadow.
205 /// In the full implementation, `queue.write_buffer` is used.
206 pub fn write_buffer(&mut self, handle: WgpuBufferHandle, data: &[f64]) {
207 if let Some(entry) = self.buffers.get_mut(handle.0) {
208 let len = data.len().min(entry.capacity);
209 entry.shadow[..len].copy_from_slice(&data[..len]);
210 }
211 }
212
213 /// Download data from the GPU buffer at `handle`.
214 ///
215 /// In the stub, data is read from the CPU-side shadow.
216 /// In the full implementation, a staging buffer is created, the command
217 /// `encoder.copy_buffer_to_buffer` is executed, and the staging buffer is
218 /// mapped for reading.
219 pub fn read_buffer(&self, handle: WgpuBufferHandle) -> Vec<f64> {
220 self.buffers
221 .get(handle.0)
222 .map(|e| e.shadow.clone())
223 .unwrap_or_default()
224 }
225
226 // ── Dispatch ─────────────────────────────────────────────────────────────
227
228 /// Dispatch a compute kernel with `work_groups_x` workgroups.
229 ///
230 /// In the stub, the kernel's `execute` method is called on the CPU-side
231 /// shadow data. In the full implementation a `ComputePipeline` is looked
232 /// up from the shader registry and `encoder.dispatch_workgroups` is called.
233 ///
234 /// # Arguments
235 ///
236 /// * `kernel_name` — name of the WGSL shader entry point
237 /// * `buffers` — input/output buffer handles
238 /// * `work_groups_x` — number of workgroups in the X dimension
239 pub fn dispatch(
240 &mut self,
241 kernel_name: &str,
242 buffers: &[WgpuBufferHandle],
243 work_groups_x: u32,
244 ) {
245 // ── TODO (wgpu-backend feature) ─────────────────────────────────────
246 // When enabled:
247 // let pipeline = self.shader_registry.get_pipeline(kernel_name)?;
248 // let bind_group = self.device.create_bind_group(…);
249 // let mut encoder = self.device.create_command_encoder(…);
250 // {
251 // let mut pass = encoder.begin_compute_pass(…);
252 // pass.set_pipeline(&pipeline);
253 // pass.set_bind_group(0, &bind_group, &[]);
254 // pass.dispatch_workgroups(work_groups_x, 1, 1);
255 // }
256 // self.queue.submit([encoder.finish()]);
257 // ────────────────────────────────────────────────────────────────────
258
259 // Stub: identity kernel (pass-through, no-op)
260 let _ = (kernel_name, buffers, work_groups_x);
261 }
262
263 // ── WGSL shader registry ──────────────────────────────────────────────────
264
265 /// Register a WGSL compute shader source and associate it with a name.
266 ///
267 /// In the stub, the source is stored but not compiled.
268 /// In the full implementation, `device.create_shader_module` is called and
269 /// the resulting `ShaderModule` is cached.
270 pub fn register_shader(&mut self, name: &str, wgsl_source: &str) {
271 // ── TODO (wgpu-backend feature) ─────────────────────────────────────
272 // let module = self.device.create_shader_module(wgpu::ShaderModuleDescriptor {
273 // label: Some(name),
274 // source: wgpu::ShaderSource::Wgsl(wgsl_source.into()),
275 // });
276 // self.shader_registry.insert(name.to_string(), module);
277 let _ = (name, wgsl_source);
278 }
279}
280
281// ── Built-in WGSL kernels ─────────────────────────────────────────────────────
282
283/// WGSL source for a parallel prefix sum (exclusive scan) kernel.
284///
285/// This is the Blelloch algorithm adapted for WGSL with a workgroup of 256 threads.
286pub const WGSL_PARALLEL_SCAN: &str = r#"
287// Exclusive parallel prefix sum (Blelloch up-sweep / down-sweep)
288// Workgroup size: 256 threads
289// Binding 0: input buffer (read)
290// Binding 1: output buffer (write)
291// Binding 2: uniform { n: u32, pass: u32 }
292
293@group(0) @binding(0) var<storage, read> input: array<f32>;
294@group(0) @binding(1) var<storage, read_write> output: array<f32>;
295
296struct Params { n: u32, pass: u32 }
297@group(0) @binding(2) var<uniform> params: Params;
298
299var<workgroup> shared: array<f32, 256>;
300
301@compute @workgroup_size(256)
302fn exclusive_scan(@builtin(global_invocation_id) gid: vec3<u32>,
303 @builtin(local_invocation_id) lid: vec3<u32>) {
304 let n = params.n;
305 let i = gid.x;
306
307 // Load
308 shared[lid.x] = select(0.0, input[i], i < n);
309 workgroupBarrier();
310
311 // Up-sweep (reduce)
312 var stride: u32 = 1u;
313 loop {
314 if stride >= 256u { break; }
315 if lid.x % (stride * 2u) == (stride * 2u - 1u) {
316 shared[lid.x] += shared[lid.x - stride];
317 }
318 workgroupBarrier();
319 stride = stride * 2u;
320 }
321
322 // Down-sweep
323 if lid.x == 255u { shared[255] = 0.0; }
324 workgroupBarrier();
325 stride = 128u;
326 loop {
327 if stride == 0u { break; }
328 if lid.x % (stride * 2u) == (stride * 2u - 1u) {
329 let tmp = shared[lid.x - stride];
330 shared[lid.x - stride] = shared[lid.x];
331 shared[lid.x] += tmp;
332 }
333 workgroupBarrier();
334 stride = stride / 2u;
335 }
336
337 // Store
338 if i < n { output[i] = shared[lid.x]; }
339}
340"#;
341
342/// WGSL source for a simple SPH density kernel.
343///
344/// Computes particle density via a cubic-spline kernel with radius `h`.
345pub const WGSL_SPH_DENSITY: &str = r#"
346// SPH density kernel — W_spline3 smoothing
347// Binding 0: positions array (x0,y0,z0, x1,y1,z1, ...)
348// Binding 1: densities output (one per particle)
349// Binding 2: uniform { n: u32, h: f32, mass: f32 }
350
351struct SphParams { n: u32, h: f32, mass: f32 }
352@group(0) @binding(0) var<storage, read> positions: array<f32>;
353@group(0) @binding(1) var<storage, read_write> densities: array<f32>;
354@group(0) @binding(2) var<uniform> params: SphParams;
355
356fn w_spline3(r: f32, h: f32) -> f32 {
357 let q = r / h;
358 let sigma = 3.0 / (2.0 * 3.14159265358979 * h * h * h);
359 if q < 1.0 {
360 return sigma * (2.0/3.0 - q*q + 0.5*q*q*q);
361 } else if q < 2.0 {
362 let t = 2.0 - q;
363 return sigma * (1.0/6.0) * t*t*t;
364 } else {
365 return 0.0;
366 }
367}
368
369@compute @workgroup_size(64)
370fn sph_density(@builtin(global_invocation_id) gid: vec3<u32>) {
371 let i = gid.x;
372 let n = params.n;
373 if i >= n { return; }
374
375 let xi = vec3<f32>(positions[i*3u], positions[i*3u+1u], positions[i*3u+2u]);
376 var density: f32 = 0.0;
377
378 for (var j: u32 = 0u; j < n; j++) {
379 let xj = vec3<f32>(positions[j*3u], positions[j*3u+1u], positions[j*3u+2u]);
380 let r = length(xi - xj);
381 density += params.mass * w_spline3(r, params.h);
382 }
383
384 densities[i] = density;
385}
386"#;
387
388/// WGSL source for parallel BVH ray traversal.
389///
390/// Traverses a linearized BVH (LBVH) to find ray–box intersections.
391/// This is a stub; real traversal requires the full BVH node buffer layout.
392pub const WGSL_BVH_TRAVERSAL: &str = r#"
393// Parallel BVH ray traversal stub
394// Each thread handles one ray; BVH nodes are in binding 0.
395
396struct Ray { origin: vec3<f32>, dir: vec3<f32>, t_max: f32 }
397struct BvhNode { lo: vec3<f32>, hi: vec3<f32>, left: u32, right: u32, is_leaf: u32, prim: u32 }
398struct HitResult { hit: u32, t: f32, prim: u32 }
399
400@group(0) @binding(0) var<storage, read> nodes: array<BvhNode>;
401@group(0) @binding(1) var<storage, read> rays: array<Ray>;
402@group(0) @binding(2) var<storage, read_write> results: array<HitResult>;
403@group(0) @binding(3) var<uniform> num_rays: u32;
404
405fn ray_aabb(ray: Ray, lo: vec3<f32>, hi: vec3<f32>) -> f32 {
406 let inv_dir = 1.0 / ray.dir;
407 let t0 = (lo - ray.origin) * inv_dir;
408 let t1 = (hi - ray.origin) * inv_dir;
409 let t_min = max(max(min(t0.x, t1.x), min(t0.y, t1.y)), min(t0.z, t1.z));
410 let t_max_box = min(min(max(t0.x, t1.x), max(t0.y, t1.y)), max(t0.z, t1.z));
411 if t_max_box < t_min || t_min > ray.t_max { return -1.0; }
412 return t_min;
413}
414
415@compute @workgroup_size(64)
416fn bvh_traverse(@builtin(global_invocation_id) gid: vec3<u32>) {
417 let rid = gid.x;
418 if rid >= num_rays { return; }
419 let ray = rays[rid];
420 results[rid] = HitResult(0u, ray.t_max, 0xFFFFFFFFu);
421
422 // Iterative DFS stack (max depth 32)
423 var stack: array<u32, 32>;
424 var sp: i32 = 0;
425 stack[0] = 0u;
426
427 loop {
428 if sp < 0 { break; }
429 let node_idx = stack[sp]; sp--;
430 let node = nodes[node_idx];
431
432 let t = ray_aabb(ray, node.lo, node.hi);
433 if t < 0.0 { continue; }
434
435 if node.is_leaf != 0u {
436 if t < results[rid].t {
437 results[rid] = HitResult(1u, t, node.prim);
438 }
439 } else {
440 if sp < 30 { sp++; stack[sp] = node.left; }
441 if sp < 30 { sp++; stack[sp] = node.right; }
442 }
443 }
444}
445"#;
446
447// ── WgpuInitError ─────────────────────────────────────────────────────────────
448
449/// Error returned when the wgpu backend cannot be initialised.
450#[derive(Debug, Clone, PartialEq)]
451pub enum WgpuInitError {
452 /// No compatible GPU adapter was found.
453 NoAdapter,
454 /// The `wgpu-backend` feature is not enabled; this is a stub build.
455 NotAvailable,
456 /// The device request failed (e.g. out of memory).
457 DeviceRequestFailed(String),
458 /// A required GPU feature is disabled or not supported.
459 FeatureDisabled,
460 /// Device creation failed with the given error string.
461 DeviceRequest(String),
462 /// A buffer handle is out of range.
463 InvalidHandle(usize),
464 /// A mutex was poisoned (should not occur in practice).
465 PoisonedLock,
466}
467
468impl std::fmt::Display for WgpuInitError {
469 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
470 match self {
471 WgpuInitError::NoAdapter => write!(f, "No compatible GPU adapter found"),
472 WgpuInitError::NotAvailable => write!(f, "wgpu-backend feature not enabled"),
473 WgpuInitError::DeviceRequestFailed(s) => write!(f, "Device request failed: {s}"),
474 WgpuInitError::FeatureDisabled => write!(f, "Required GPU feature is not available"),
475 WgpuInitError::DeviceRequest(s) => write!(f, "Device request error: {s}"),
476 WgpuInitError::InvalidHandle(h) => write!(f, "Invalid buffer handle: {h}"),
477 WgpuInitError::PoisonedLock => write!(f, "Internal mutex was poisoned"),
478 }
479 }
480}
481
482impl std::error::Error for WgpuInitError {}
483
484// ── Tests ─────────────────────────────────────────────────────────────────────
485
486#[cfg(test)]
487mod tests {
488 use super::*;
489
490 #[test]
491 fn try_new_returns_not_available_in_stub_build() {
492 let result = WgpuBackend::try_new();
493 assert!(matches!(result, Err(WgpuInitError::NotAvailable)));
494 }
495
496 #[test]
497 fn stub_backend_write_read_roundtrip() {
498 let mut backend = WgpuBackend::new_stub();
499 let handle = backend.create_buffer(4);
500 let data = vec![1.0_f64, 2.0, 3.0, 4.0];
501 backend.write_buffer(handle, &data);
502 let out = backend.read_buffer(handle);
503 assert_eq!(out, data);
504 }
505
506 #[test]
507 fn stub_dispatch_is_noop() {
508 let mut backend = WgpuBackend::new_stub();
509 let h = backend.create_buffer(8);
510 let before = backend.read_buffer(h);
511 backend.dispatch("sph_density", &[h], 1);
512 let after = backend.read_buffer(h);
513 assert_eq!(before, after, "stub dispatch should not modify buffers");
514 }
515
516 #[test]
517 fn wgsl_kernels_are_non_empty() {
518 assert!(!WGSL_PARALLEL_SCAN.is_empty());
519 assert!(!WGSL_SPH_DENSITY.is_empty());
520 assert!(!WGSL_BVH_TRAVERSAL.is_empty());
521 }
522
523 #[test]
524 fn device_info_stub_has_name() {
525 let backend = WgpuBackend::new_stub();
526 assert!(!backend.device_info().name.is_empty());
527 }
528
529 #[test]
530 fn wgpu_init_error_display() {
531 assert!(!WgpuInitError::NotAvailable.to_string().is_empty());
532 assert!(!WgpuInitError::NoAdapter.to_string().is_empty());
533 assert!(!WgpuInitError::FeatureDisabled.to_string().is_empty());
534 assert!(
535 !WgpuInitError::DeviceRequest("oom".into())
536 .to_string()
537 .is_empty()
538 );
539 assert!(!WgpuInitError::InvalidHandle(7).to_string().is_empty());
540 assert!(!WgpuInitError::PoisonedLock.to_string().is_empty());
541 }
542}
543
544// ── Real wgpu backend (feature-gated) ─────────────────────────────────────────
545
546/// Real wgpu compute backend, enabled only with the `wgpu-backend` feature.
547///
548/// Provides GPU buffer management, WGSL shader dispatch, and CPU-side readback
549/// using `wgpu` 29's cross-platform Vulkan / Metal / DX12 backends.
550///
551/// # Thread safety
552///
553/// `wgpu::Device` and `wgpu::Queue` are `Send + Sync`. The shader cache is
554/// protected by a `Mutex`, making `WgpuBackendReal` safe to share across
555/// threads (though individual dispatches are synchronous on the calling thread).
556///
557/// # Usage
558///
559/// ```ignore
560/// // With the wgpu-backend feature enabled:
561/// use oxiphysics_gpu::compute::wgpu_backend::real::WgpuBackendReal;
562///
563/// let mut backend = WgpuBackendReal::try_new()?;
564/// let h = backend.create_buffer_f64(128);
565/// backend.write_buffer_f64(h, &vec![1.0_f64; 128]);
566/// backend.dispatch_wgsl(
567/// WGSL_SPH_DENSITY, "sph_density",
568/// &[(h, wgpu::BufferBindingType::Storage { read_only: false })],
569/// [2, 1, 1],
570/// )?;
571/// let out = backend.read_buffer_f64(h);
572/// ```
573#[cfg(feature = "wgpu-backend")]
574pub mod real {
575 use super::{WgpuBufferHandle, WgpuDeviceInfo, WgpuInitError};
576 use std::collections::HashMap;
577 use std::hash::{DefaultHasher, Hash, Hasher};
578 use std::sync::{Arc, Mutex};
579
580 // ── Internal shader-cache entry ──────────────────────────────────────────
581
582 struct ShaderCacheEntry {
583 pipeline: Arc<wgpu::ComputePipeline>,
584 }
585
586 // ── WgpuBackendReal ──────────────────────────────────────────────────────
587
588 /// Real GPU compute backend backed by `wgpu` 29.
589 ///
590 /// Obtain an instance via [`WgpuBackendReal::try_new`] (synchronous,
591 /// blocks the thread) or [`WgpuBackendReal::try_new_async`] from within
592 /// an async context.
593 pub struct WgpuBackendReal {
594 device: Arc<wgpu::Device>,
595 queue: Arc<wgpu::Queue>,
596 /// Device information (name, backend, driver).
597 pub device_info: WgpuDeviceInfo,
598 /// Allocated GPU buffers, indexed by `WgpuBufferHandle.0`.
599 buffers: Vec<Option<Arc<wgpu::Buffer>>>,
600 /// Byte size of each buffer (parallel to `buffers`).
601 buffer_sizes: Vec<u64>,
602 /// Compiled pipeline cache, keyed by a hash of WGSL source + entry point.
603 shader_cache: Mutex<HashMap<u64, ShaderCacheEntry>>,
604 }
605
606 impl WgpuBackendReal {
607 // ── Construction ─────────────────────────────────────────────────────
608
609 /// Create a real GPU backend, blocking the calling thread.
610 ///
611 /// Returns `Err` if no compatible GPU adapter is found or if device
612 /// creation fails. Prefer [`try_new_async`](Self::try_new_async) from
613 /// within an `async` context.
614 pub fn try_new() -> Result<Self, WgpuInitError> {
615 pollster::block_on(Self::try_new_async())
616 }
617
618 /// Create a real GPU backend asynchronously.
619 ///
620 /// This is the preferred entry point from `async` contexts (tokio,
621 /// wasm-bindgen-futures, etc.).
622 pub async fn try_new_async() -> Result<Self, WgpuInitError> {
623 let instance =
624 wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
625
626 let adapter = instance
627 .request_adapter(&wgpu::RequestAdapterOptions {
628 power_preference: wgpu::PowerPreference::HighPerformance,
629 compatible_surface: None,
630 force_fallback_adapter: false,
631 })
632 .await
633 .map_err(|_| WgpuInitError::NoAdapter)?;
634
635 let info = adapter.get_info();
636
637 let desc = wgpu::DeviceDescriptor {
638 label: Some("oxiphysics-wgpu"),
639 required_features: wgpu::Features::empty(),
640 required_limits: adapter.limits(),
641 ..Default::default()
642 };
643
644 let (device, queue) = adapter
645 .request_device(&desc)
646 .await
647 .map_err(|e| WgpuInitError::DeviceRequest(e.to_string()))?;
648
649 let device_info = WgpuDeviceInfo {
650 name: info.name.clone(),
651 backend: format!("{:?}", info.backend),
652 driver_version: info.driver_info.clone(),
653 // VRAM is not exposed by wgpu's AdapterInfo; use 0 as sentinel.
654 vram_bytes: 0,
655 // GPU-native f64 requires a device extension not in the base profile.
656 supports_f64: false,
657 // Conservative defaults matching most desktop GPU limits.
658 max_workgroup_size: [256, 256, 64],
659 };
660
661 Ok(Self {
662 device: Arc::new(device),
663 queue: Arc::new(queue),
664 device_info,
665 buffers: Vec::new(),
666 buffer_sizes: Vec::new(),
667 shader_cache: Mutex::new(HashMap::new()),
668 })
669 }
670
671 /// Return `true` — this struct always wraps a real GPU device.
672 pub fn is_available(&self) -> bool {
673 true
674 }
675
676 // ── Buffer management ─────────────────────────────────────────────────
677
678 /// Allocate a GPU storage buffer of `size_bytes` bytes.
679 ///
680 /// The buffer is created with `STORAGE | COPY_SRC | COPY_DST` usage
681 /// flags so that it can be used as a shader binding and for staged
682 /// CPU read/write.
683 pub fn create_buffer_storage(&mut self, size_bytes: u64) -> WgpuBufferHandle {
684 let handle = WgpuBufferHandle(self.buffers.len());
685 let buf = self.device.create_buffer(&wgpu::BufferDescriptor {
686 label: None,
687 size: size_bytes,
688 usage: wgpu::BufferUsages::STORAGE
689 | wgpu::BufferUsages::COPY_SRC
690 | wgpu::BufferUsages::COPY_DST,
691 mapped_at_creation: false,
692 });
693 self.buffers.push(Some(Arc::new(buf)));
694 self.buffer_sizes.push(size_bytes);
695 handle
696 }
697
698 /// Allocate a GPU buffer sized for `len` `f64` values.
699 ///
700 /// Internally the data is stored as `f32` on the GPU (8 bytes per
701 /// element to maintain the same stride).
702 pub fn create_buffer_f64(&mut self, len: usize) -> WgpuBufferHandle {
703 // We store f64 values packed as two f32s to preserve stride; or
704 // simply allocate 8 bytes per element and use the f32 path with
705 // two floats per logical element. For simplicity, the current
706 // implementation casts f64→f32 on write and f32→f64 on read, so
707 // we only need 4 bytes per element on the GPU.
708 self.create_buffer_storage((len * 4) as u64)
709 }
710
711 /// Upload `data` to the GPU buffer at `handle`, casting `f64` → `f32`.
712 ///
713 /// # Panics
714 ///
715 /// Does nothing (silently returns) if `handle` is out of range.
716 pub fn write_buffer_f64(&self, handle: WgpuBufferHandle, data: &[f64]) {
717 if let Some(Some(buf)) = self.buffers.get(handle.0) {
718 let f32_data: Vec<f32> = data.iter().map(|&v| v as f32).collect();
719 self.queue
720 .write_buffer(buf, 0, bytemuck::cast_slice(&f32_data));
721 }
722 }
723
724 /// Download data from the GPU buffer at `handle`, casting `f32` → `f64`.
725 ///
726 /// This blocks the calling thread until the GPU has finished all
727 /// outstanding work and the readback mapping is complete.
728 ///
729 /// Returns an empty `Vec` if the handle is invalid or the readback fails.
730 pub fn read_buffer_f64(&self, handle: WgpuBufferHandle) -> Vec<f64> {
731 let buf = match self.buffers.get(handle.0).and_then(|b| b.as_ref()) {
732 Some(b) => b.clone(),
733 None => return Vec::new(),
734 };
735 let size = self.buffer_sizes[handle.0];
736
737 // Create a CPU-visible staging buffer for the readback.
738 let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
739 label: Some("oxiphysics_staging_readback"),
740 size,
741 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
742 mapped_at_creation: false,
743 });
744
745 // Record and submit the copy command.
746 let mut encoder = self
747 .device
748 .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
749 encoder.copy_buffer_to_buffer(&buf, 0, &staging, 0, size);
750 self.queue.submit(std::iter::once(encoder.finish()));
751
752 // Map the staging buffer for reading.
753 let slice = staging.slice(..);
754 let (tx, rx) = std::sync::mpsc::channel();
755 slice.map_async(wgpu::MapMode::Read, move |result| {
756 let _ = tx.send(result);
757 });
758
759 // Block until the GPU has completed and the mapping is ready.
760 if let Err(_e) = self.device.poll(wgpu::PollType::Wait {
761 submission_index: None,
762 timeout: None,
763 }) {
764 return Vec::new();
765 }
766
767 // Check that the mapping succeeded.
768 if rx.recv().ok().and_then(|r| r.ok()).is_none() {
769 return Vec::new();
770 }
771
772 let mapped = slice.get_mapped_range();
773 let f32_data: &[f32] = bytemuck::cast_slice(&mapped);
774 let result: Vec<f64> = f32_data.iter().map(|&v| v as f64).collect();
775 drop(mapped);
776 staging.unmap();
777 result
778 }
779
780 // ── Dispatch ──────────────────────────────────────────────────────────
781
782 /// Upload raw bytes to the GPU buffer at `handle`.
783 ///
784 /// The byte slice must fit within the buffer's allocated size.
785 /// Does nothing (silently returns) if `handle` is out of range.
786 pub fn queue_write_buffer_raw(&self, handle: &WgpuBufferHandle, data: &[u8]) {
787 if let Some(Some(buf)) = self.buffers.get(handle.0) {
788 self.queue.write_buffer(buf, 0, data);
789 }
790 }
791
792 /// Upload `f32` data directly to the GPU buffer at `handle` (no f64→f32 cast).
793 ///
794 /// Does nothing (silently returns) if `handle` is out of range.
795 pub fn queue_write_buffer_f32(&self, handle: &WgpuBufferHandle, data: &[f32]) {
796 if let Some(Some(buf)) = self.buffers.get(handle.0) {
797 self.queue.write_buffer(buf, 0, bytemuck::cast_slice(data));
798 }
799 }
800
801 /// Download raw `f32` values from the GPU buffer at `handle`.
802 ///
803 /// Returns an empty `Vec` if the handle is invalid or the readback fails.
804 pub fn read_buffer_f32(&self, handle: WgpuBufferHandle) -> Vec<f32> {
805 let buf = match self.buffers.get(handle.0).and_then(|b| b.as_ref()) {
806 Some(b) => b.clone(),
807 None => return Vec::new(),
808 };
809 let size = self.buffer_sizes[handle.0];
810
811 let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
812 label: Some("oxiphysics_staging_readback_f32"),
813 size,
814 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
815 mapped_at_creation: false,
816 });
817
818 let mut encoder = self
819 .device
820 .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
821 encoder.copy_buffer_to_buffer(&buf, 0, &staging, 0, size);
822 self.queue.submit(std::iter::once(encoder.finish()));
823
824 let slice = staging.slice(..);
825 let (tx, rx) = std::sync::mpsc::channel();
826 slice.map_async(wgpu::MapMode::Read, move |result| {
827 let _ = tx.send(result);
828 });
829
830 if let Err(_e) = self.device.poll(wgpu::PollType::Wait {
831 submission_index: None,
832 timeout: None,
833 }) {
834 return Vec::new();
835 }
836
837 if rx.recv().ok().and_then(|r| r.ok()).is_none() {
838 return Vec::new();
839 }
840
841 let mapped = slice.get_mapped_range();
842 let result: Vec<f32> = bytemuck::cast_slice::<u8, f32>(&mapped).to_vec();
843 drop(mapped);
844 staging.unmap();
845 result
846 }
847
848 // ── Dispatch ──────────────────────────────────────────────────────────
849
850 /// Compute the 3-D workgroup dispatch counts for `n_items` elements.
851 ///
852 /// Returns `[0, 1, 1]` for `n_items == 0` (no-op dispatch).
853 pub fn dispatch_count_for(n_items: usize, workgroup_size: u32) -> [u32; 3] {
854 crate::compute::timestamp::dispatch_count_for(n_items, workgroup_size)
855 }
856
857 /// Compile and dispatch a WGSL compute shader.
858 ///
859 /// The pipeline is compiled lazily and cached by a hash of
860 /// `(wgsl_src, entry_point)`, so repeated calls with the same shader
861 /// do not recompile.
862 ///
863 /// # Parameters
864 ///
865 /// * `wgsl_src` — WGSL shader source code.
866 /// * `entry_point` — Name of the `@compute` entry point function.
867 /// * `buffers` — Ordered list of `(handle, binding_type)` pairs.
868 /// Binding index in the WGSL shader corresponds to the position in
869 /// this slice (binding 0 = `buffers[0]`, etc.).
870 /// * `workgroups` — `[x, y, z]` dispatch counts.
871 ///
872 /// # Errors
873 ///
874 /// Returns `Err(WgpuInitError::InvalidHandle)` if any buffer handle is
875 /// out of range. Returns `Err(WgpuInitError::PoisonedLock)` if the
876 /// shader-cache mutex is poisoned (should not occur in practice).
877 pub fn dispatch_wgsl(
878 &self,
879 wgsl_src: &str,
880 entry_point: &str,
881 buffers: &[(WgpuBufferHandle, wgpu::BufferBindingType)],
882 workgroups: [u32; 3],
883 ) -> Result<(), WgpuInitError> {
884 // Hash the shader source + entry point to key the pipeline cache.
885 let mut hasher = DefaultHasher::new();
886 wgsl_src.hash(&mut hasher);
887 entry_point.hash(&mut hasher);
888 let key = hasher.finish();
889
890 // Obtain or compile the pipeline.
891 let pipeline: Arc<wgpu::ComputePipeline> = {
892 let mut cache = self.shader_cache.lock().unwrap_or_else(|e| e.into_inner());
893
894 if let Some(entry) = cache.get(&key) {
895 entry.pipeline.clone()
896 } else {
897 let module = self
898 .device
899 .create_shader_module(wgpu::ShaderModuleDescriptor {
900 label: Some(entry_point),
901 source: wgpu::ShaderSource::Wgsl(wgsl_src.into()),
902 });
903 let pipeline = Arc::new(self.device.create_compute_pipeline(
904 &wgpu::ComputePipelineDescriptor {
905 label: Some(entry_point),
906 layout: None,
907 module: &module,
908 entry_point: Some(entry_point),
909 compilation_options: wgpu::PipelineCompilationOptions::default(),
910 cache: None,
911 },
912 ));
913 cache.insert(
914 key,
915 ShaderCacheEntry {
916 pipeline: pipeline.clone(),
917 },
918 );
919 pipeline
920 }
921 };
922
923 // Derive the bind-group layout from the compiled pipeline.
924 let bg_layout = pipeline.get_bind_group_layout(0);
925
926 // Build the bind-group entries.
927 let mut entries: Vec<wgpu::BindGroupEntry> = Vec::with_capacity(buffers.len());
928 for (i, (handle, _binding_type)) in buffers.iter().enumerate() {
929 let buf = self
930 .buffers
931 .get(handle.0)
932 .and_then(|b| b.as_ref())
933 .ok_or(WgpuInitError::InvalidHandle(handle.0))?;
934 entries.push(wgpu::BindGroupEntry {
935 binding: i as u32,
936 resource: buf.as_entire_binding(),
937 });
938 }
939
940 let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
941 label: None,
942 layout: &bg_layout,
943 entries: &entries,
944 });
945
946 // Record and submit the compute pass.
947 let mut encoder = self
948 .device
949 .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
950 {
951 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
952 label: None,
953 timestamp_writes: None,
954 });
955 pass.set_pipeline(&pipeline);
956 pass.set_bind_group(0, &bind_group, &[]);
957 pass.dispatch_workgroups(workgroups[0], workgroups[1], workgroups[2]);
958 }
959 self.queue.submit(std::iter::once(encoder.finish()));
960
961 // Block until the GPU has finished (synchronous dispatch).
962 self.device
963 .poll(wgpu::PollType::Wait {
964 submission_index: None,
965 timeout: None,
966 })
967 .map_err(|_| WgpuInitError::DeviceRequest("poll failed".into()))?;
968
969 Ok(())
970 }
971 }
972
973 // ── Feature-gated tests ───────────────────────────────────────────────────
974
975 #[cfg(test)]
976 mod tests {
977 use super::*;
978
979 /// Helper: attempt to create a real backend, returning `None` if no GPU
980 /// is available (e.g. in headless CI).
981 fn try_backend() -> Option<WgpuBackendReal> {
982 WgpuBackendReal::try_new().ok()
983 }
984
985 #[test]
986 fn real_backend_try_new_succeeds_or_gracefully_fails() {
987 // This test always passes: it either succeeds (GPU present) or
988 // returns None (headless / CI environment).
989 match WgpuBackendReal::try_new() {
990 Ok(b) => {
991 assert!(b.is_available());
992 assert!(!b.device_info.backend.is_empty());
993 }
994 Err(e) => {
995 // NoAdapter is the expected error in headless CI.
996 eprintln!("No GPU adapter available: {e}");
997 }
998 }
999 }
1000
1001 #[test]
1002 fn real_backend_create_and_write_buffer() {
1003 let Some(mut backend) = try_backend() else {
1004 return;
1005 };
1006 let data = vec![1.0_f64, 2.0, 3.0, 4.0];
1007 let handle = backend.create_buffer_f64(data.len());
1008 backend.write_buffer_f64(handle, &data);
1009 // write_buffer_f64 is fire-and-forget; we just verify no panic.
1010 assert!(handle.0 < backend.buffers.len());
1011 }
1012
1013 #[test]
1014 fn real_backend_buffer_roundtrip() {
1015 let Some(mut backend) = try_backend() else {
1016 return;
1017 };
1018 let data = vec![1.0_f64, 2.0, 3.0, 4.0];
1019 let handle = backend.create_buffer_f64(data.len());
1020 backend.write_buffer_f64(handle, &data);
1021 let out = backend.read_buffer_f64(handle);
1022 // f64→f32→f64 loses precision; check within f32 rounding.
1023 assert_eq!(out.len(), data.len());
1024 for (&expected, &got) in data.iter().zip(out.iter()) {
1025 assert!(
1026 (expected as f32 - got as f32).abs() < 1e-5,
1027 "roundtrip mismatch: expected {expected}, got {got}"
1028 );
1029 }
1030 }
1031
1032 #[test]
1033 fn real_backend_dispatch_scale_shader() {
1034 let Some(mut backend) = try_backend() else {
1035 return;
1036 };
1037 use super::super::WgpuBackend;
1038
1039 // A simple WGSL shader that multiplies each f32 element by 2.
1040 const SCALE_BY_TWO: &str = r#"
1041@group(0) @binding(0) var<storage, read> input_buf: array<f32>;
1042@group(0) @binding(1) var<storage, read_write> output_buf: array<f32>;
1043
1044@compute @workgroup_size(64)
1045fn scale_by_two(@builtin(global_invocation_id) gid: vec3<u32>) {
1046 let i = gid.x;
1047 if i < arrayLength(&input_buf) {
1048 output_buf[i] = input_buf[i] * 2.0;
1049 }
1050}
1051"#;
1052 let n: usize = 4;
1053 let input_data: Vec<f32> = (1..=n as u32).map(|x| x as f32).collect();
1054 let in_handle = backend.create_buffer_storage((n * 4) as u64);
1055 let out_handle = backend.create_buffer_storage((n * 4) as u64);
1056
1057 backend.queue.write_buffer(
1058 backend.buffers[in_handle.0].as_ref().unwrap(),
1059 0,
1060 bytemuck::cast_slice(&input_data),
1061 );
1062
1063 // Dispatch: 1 workgroup of 64 threads covers n=4 elements.
1064 let result = backend.dispatch_wgsl(
1065 SCALE_BY_TWO,
1066 "scale_by_two",
1067 &[
1068 (
1069 in_handle,
1070 wgpu::BufferBindingType::Storage { read_only: true },
1071 ),
1072 (
1073 out_handle,
1074 wgpu::BufferBindingType::Storage { read_only: false },
1075 ),
1076 ],
1077 [1, 1, 1],
1078 );
1079 assert!(result.is_ok(), "dispatch_wgsl failed: {:?}", result.err());
1080
1081 // Readback via staging and verify.
1082 let out = backend.read_buffer_f64(out_handle);
1083 assert_eq!(out.len(), n);
1084 for (i, &v) in out.iter().enumerate() {
1085 let expected = (i + 1) as f64 * 2.0;
1086 assert!(
1087 (v - expected).abs() < 0.01,
1088 "element {i}: expected {expected}, got {v}"
1089 );
1090 }
1091
1092 // Regression guard: stub backend still works.
1093 let mut stub = WgpuBackend::new_stub();
1094 let h = stub.create_buffer(4);
1095 let _ = stub.read_buffer(h);
1096 }
1097
1098 #[test]
1099 fn dispatch_count_for_zero_items() {
1100 assert_eq!(WgpuBackendReal::dispatch_count_for(0, 64), [0, 1, 1]);
1101 }
1102
1103 #[test]
1104 fn dispatch_count_for_65_items() {
1105 assert_eq!(WgpuBackendReal::dispatch_count_for(65, 64), [2, 1, 1]);
1106 }
1107
1108 #[test]
1109 fn dispatch_count_for_exact_workgroup() {
1110 assert_eq!(WgpuBackendReal::dispatch_count_for(256, 64), [4, 1, 1]);
1111 }
1112 }
1113}