henad-compute 0.3.0

//! GPU engine machinery, the sibling of [`crate::cpu`], for models whose state lives in GPU
//! buffers and never round-trips to the CPU.
//!
//! Nothing here ever *creates* a `wgpu::Device`. A host passes one in through a [`GpuContext`].
//! A model contributes shaders, seed data and metadata, and every wgpu object is built here.

pub mod agent_engine;
pub mod capacity;
mod contracts;
pub mod fault;
pub mod grid_engine;
pub mod limits;
pub mod primitives;
pub mod sim_thread;
#[cfg(not(target_arch = "wasm32"))]
pub mod stepping;
pub mod timing;
pub mod view;

#[cfg(test)]
mod tests;

pub use agent_engine::GpuAgentState;
pub use capacity::Demand;
pub use grid_engine::GpuGridState;
pub use limits::GpuNeeds;
pub use primitives::readback::StatsPoll;
pub use primitives::spatial_hash::{GpuSpatialHash, HashGrid};
pub use sim_thread::{GpuSimState, GpuStats};
pub use view::agents::GpuAgents;
pub use view::display::{DisplayTarget, GpuDisplay};
/// The wgpu release Henad builds on. Its types sit in [`GpuContext`]'s fields and in device requests, and a caller
/// refers to them through this path instead of its own `wgpu` dependency.
pub use wgpu;

#[cfg(test)]
use tests::support::headless_context;

pub use sim_thread::GpuSimThread;

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use crate::fault::{Fault, FaultSink};
use crate::runtime_info::RuntimeInfo;

/// Largest number of steps one command buffer holds.
///
/// Enough passes in one submission trip the OS GPU watchdog, with no error, and every later readback reads zero.
pub const MAX_STEPS_PER_SUBMISSION: u32 = 64;

/// GPU handles that a host passes to the engines, with the sink for errors no scope catches.
///
/// A clone is cheap, and shares the fault sink and the record of a lost device.
#[derive(Debug, Clone)]
pub struct GpuContext {
    /// Device every model builds on.
    pub device: wgpu::Device,
    /// Queue the engine submits its work to.
    pub queue: wgpu::Queue,
    /// Colour format of the target a display pipeline renders into.
    ///
    /// A model builds its display render pipeline once, at construction, and a pipeline is tied to its colour target
    /// format.
    pub target_format: wgpu::TextureFormat,
    /// Sink for errors that nothing else caught. See [`GpuContext::new`].
    pub faults: FaultSink,
    /// Set once the device is lost.
    lost: Arc<AtomicBool>,
    /// Facts about the host and the adapter, when whoever acquired the device attached them.
    runtime_info: Option<Arc<RuntimeInfo>>,
}

impl GpuContext {
    /// Creates a context over `device` and `queue`, and takes over the device's error handling. Left to wgpu, every
    /// error is fatal.
    ///
    /// Errors raised inside a [`fault::catching_on`] go to that scope. Everything else, including
    /// egui's own rendering on the same device, is stored in `faults` for the host to pick up. On the
    /// web this is the only route. [`fault::catching_on`] pushes no scopes there.
    ///
    /// A `GPUInternalError` still ends the web build. wgpu converts an error with
    /// `Error::from_js`. Anything other than a `GPUValidationError` or a `GPUOutOfMemoryError`
    /// panics there. A model provokes those two, and the handler reports them normally.
    ///
    /// The context also records the loss of the device. [`Self::is_lost`] reports it.
    ///
    /// Note that a second `new` on the same device takes over the error handler from the first context, and on native
    /// targets the lost callback too. From then on only the newest context receives the device's unscoped errors, and
    /// on native targets its loss. In a browser every context on the device records the loss. A clone shares the fault
    /// sink and the loss record with the context it came from.
    ///
    /// The context carries no [`RuntimeInfo`]. A host that passes its context to a sweep attaches one with
    /// `.with_runtime_info(RuntimeInfo::collect(&adapter, &device))`. Without it the sweep's manifest records no
    /// adapter.
    pub fn new(
        device: wgpu::Device,
        queue: wgpu::Queue,
        target_format: wgpu::TextureFormat,
        faults: FaultSink,
    ) -> Self {
        let sink = faults.clone();
        // wgpu requires `Send + Sync` here even where nothing can be sent. Under atomics a
        // `wgpu::Error` is neither, and the sink holds one. `SendWrapper` panics the moment a
        // second thread touches it. On the web that check is the whole guarantee.
        #[cfg(all(target_arch = "wasm32", target_feature = "atomics"))]
        let sink = send_wrapper::SendWrapper::new(sink);
        device.on_uncaptured_error(std::sync::Arc::new(move |error: wgpu::Error| {
            log::error!("unhandled GPU error: {error}");
            sink.set_once(Fault::device("running on the GPU", error));
        }));
        let lost = Arc::new(AtomicBool::new(false));
        let callback_lost = Arc::clone(&lost);
        device.set_device_lost_callback(move |reason, message| {
            log::error!("GPU device lost ({reason:?}): {message}");
            callback_lost.store(true, Ordering::Release);
        });
        Self {
            device,
            queue,
            target_format,
            faults,
            lost,
            runtime_info: None,
        }
    }

    /// Returns the context with `info` attached, as [`Self::runtime_info`] reads it back.
    pub fn with_runtime_info(self, info: RuntimeInfo) -> Self {
        Self {
            runtime_info: Some(Arc::new(info)),
            ..self
        }
    }

    /// Facts about the host and the adapter, `None` when nothing attached them.
    pub fn runtime_info(&self) -> Option<&RuntimeInfo> {
        self.runtime_info.as_deref()
    }

    /// Returns whether the device is lost. A lost device runs no more work.
    ///
    /// Note that a destroyed device counts as lost once a poll finds its queue empty.
    pub fn is_lost(&self) -> bool {
        self.lost.load(Ordering::Acquire)
    }
}