Skip to main content

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}