effectcraft/crates/gpu/src/lib.rs

//! The EffectCraft GPU compositor (Mercury GPU Acceleration's counterpart), on wgpu compute
//! shaders: Metal, Vulkan, Direct3D 12 and WebGPU.
//!
//! The CPU [`Renderer`] stays the reference and keeps rendering layer *content* (sources, masks,
//! CPU effects, layer styles) into its layer cache. [`Gpu`] composites: it uploads the cached
//! layer buffers (once per buffer), applies the layer transforms with the CPU's sampling
//! (nearest / bilinear / bicubic, minification pre-filter), motion-blur sub-samples, track
//! mattes, Preserve Transparency, layer styles' passes, all 38 blend modes and the 8/16 bpc
//! clamping and quantisation, and converts colour spaces. Classic 3D runs composite here too
//! (`classic3d`: per-pixel fragment sort, lights, ray-cast shadows), and adjustment layers run
//! their effect stacks on the GPU-resident comp. Wireframe-quality outlines are drawn here too.
//!
//! Advanced 3D comps rasterise on a render pipeline (`advanced3d.wgsl`: depth buffer, PBR,
//! image-based light, shadow maps) and finish in compute kernels (`adv3d.wgsl`: supersampling
//! resolve, motion-blur sub-samples, iris depth of field, encoding, compositing), see
//! [`Accelerator::render_3d`] and `Renderer::prepare_adv_run`. Runs with layers on the 2D
//! compositing path (blend modes, track mattes, Preserve Transparency:
//! `Renderer::split_adv_run`) render each such layer alone, hide it behind the nearer main
//! scene and composite it through the 2D path here too; environment backgrounds (Classic and
//! Advanced 3D) draw their sky in a kernel (`Renderer::sky_draw`).
//!
//! GPU effects ([`effectcraft_effects::GPU_EFFECTS`]) run as compute kernels with the CPU
//! effect's exact steps (padding, box-blur radii, parameter conversions); chains of them are
//! uploaded and read back once. Each family lives in its own module with its own WGSL file
//! (`fx_color`, `fx_depth`, `fx_distort`, `fx_extra`, `fx_gen2`, `fx_generate`, `fx_key`,
//! `fx_light`, `fx_lut`, `fx_noise`, `fx_particles`, `fx_pixel2`, `fx_sim`, `fx_stylize`,
//! `fx_text`, `fx_time`, `fx_tone`, `fx_transition`, `fx_vr`, `fx_warp`); settings a kernel
//! cannot match fall back to the CPU (`catalog::gpu_supported`, or `None` from the family's
//! `apply`).
//! The 3D Channel effects read the layer's aux channels as an extra texture (`fx_depth`), the
//! colour-management effects interpret colour programs built next to the CPU effects
//! (`fx_lut`), simulations keep their state on the CPU and rasterise its per-frame plan here
//! (`fx_sim`, `fx_particles`), and so do the geometry generators (`fx_gen2`: bolts, strokes,
//! glyphs, audio marks, waves).
//!
//! Working textures come from a pool (`context::Pool`): a frame allocates several full-frame
//! RGBA f32 images per layer, and creating and zeroing those cost more than compositing a small
//! layer. A released texture is reused once every encoder that could still read it has been
//! submitted. 8/16 bpc quantisation after each layer is fused into that layer's composite.
//! [`Backend::Auto`](effectcraft_render::Backend) renders each comp on whichever compositor
//! measured faster for it ([`effectcraft_render::AutoPick`], kept by [`Gpu`]): a light comp
//! that the CPU composites in a millisecond stays there rather than paying for a full-frame
//! readback.
//!
//! GPU particles (`particles`): the stepped particle effects hand their simulation to
//! [`effectcraft_effects::psim::ParticleSim`], implemented here with one invocation per particle
//! and GPU-resident checkpoints.
//!
//! Plug a [`Gpu`] into [`Renderer::accel`] (it implements [`Accelerator`]); renders then use it
//! when [`RenderOpts::backend`](effectcraft_render::RenderOpts) asks for it. The viewer can
//! skip readback entirely with [`Gpu::render_display`], which leaves an RGBA8 texture for
//! egui-wgpu to draw.

#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)]

mod adv3d;
mod bokeh;
mod classic3d;
mod context;
pub mod deferred;
mod effects;
mod fx_color;
mod fx_depth;
mod fx_distort;
mod fx_extra;
mod fx_gen2;
mod fx_generate;
mod fx_key;
mod fx_light;
mod fx_lut;
mod fx_noise;
mod fx_particles;
mod fx_pixel2;
mod fx_sim;
mod fx_stylize;
mod fx_text;
mod fx_time;
mod fx_tone;
mod fx_transition;
mod fx_vr;
mod fx_warp;
mod inpaint;
mod memcap;
mod ops;
mod particles;
mod readback;
mod walk;
mod yuv;

use std::sync::Arc;

pub use context::{GpuContext, GpuImage, TransferStats};
use effectcraft_effects::Buf;
use effectcraft_project::ItemId;
use effectcraft_raster::Image;
use effectcraft_render::{Accelerator, FxStep, Renderer};
use effectcraft_time::Tick;
pub use wgpu;

use crate::context::Enc;

/// Why [`Gpu::within_memory`] discarded a frame.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FrameError {
    /// The GPU context is retired (device lost, a readback timed out, an uncaptured error).
    Retired(String),
    /// Out of video memory (or a driver reporting a failed allocation as an internal error),
    /// with any validation errors it went on to cause.
    Memory(String),
    /// The device rejected a command with no allocation failure before it (an invalid texture
    /// or parameters, #1319): a compositor bug, not memory pressure.
    Validation(String),
}

impl std::fmt::Display for FrameError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            FrameError::Retired(e) | FrameError::Memory(e) | FrameError::Validation(e) => f.write_str(e),
        }
    }
}

impl std::error::Error for FrameError {}

/// The GPU compositor (cheap to clone; one device shared by all clones).
#[derive(Clone)]
pub struct Gpu {
    ctx: Arc<GpuContext>,
    /// Backend::Auto's per-comp CPU / GPU timings.
    auto: Arc<effectcraft_render::AutoPick>,
    /// (comp, reason) pairs already warned about when the GPU declined a frame.
    declined: Arc<std::sync::Mutex<std::collections::HashSet<(u64, String)>>>,
}

/// A viewer frame left on the GPU: premultiplied RGBA8 (`wgpu::TextureFormat::Rgba8Unorm`) with
/// a full mip chain (for minified views), ready for
/// `egui_wgpu::Renderer::register_native_texture_with_sampler_options`.
#[derive(Clone, Debug)]
pub struct DisplayFrame {
    pub texture: wgpu::Texture,
    pub width: u32,
    pub height: u32,
}

impl Gpu {
    /// On a device of its own (CLI, tests, benchmarks); `None` without a usable adapter.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn headless() -> Option<Gpu> {
        Gpu::try_headless().map_err(|e| log::info!("gpu: {e}")).ok()
    }

    /// [`Gpu::headless`], with why there is none.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn try_headless() -> Result<Gpu, String> {
        pollster::block_on(GpuContext::request()).map(Gpu::from_context)
    }

    /// On an existing device (the desktop app shares egui-wgpu's).
    pub fn new(adapter: &wgpu::Adapter, device: wgpu::Device, queue: wgpu::Queue) -> Result<Gpu, String> {
        GpuContext::new(adapter, device, queue).map(Gpu::from_context)
    }

    pub fn from_context(ctx: GpuContext) -> Gpu {
        Gpu { ctx: Arc::new(ctx), auto: Default::default(), declined: Default::default() }
    }

    /// A device of its own without blocking, with deferred readbacks (a browser worker, see
    /// [`deferred`]); natively the same device works too (tests). `Err` says why there is none.
    pub async fn request_deferred() -> Result<Gpu, String> {
        let mut ctx = GpuContext::request().await?;
        ctx.set_deferred(true);
        Ok(Gpu::from_context(ctx))
    }

    /// [`Gpu::request_deferred`] blocking (tests of the browser worker's path, natively).
    #[cfg(not(target_arch = "wasm32"))]
    pub fn headless_deferred() -> Option<Gpu> {
        pollster::block_on(Gpu::request_deferred()).map_err(|e| log::info!("gpu: {e}")).ok()
    }

    /// The deferred readback state, when readbacks are deferred.
    pub fn deferred(&self) -> Option<&Arc<deferred::Deferred>> {
        self.ctx.deferred()
    }

    /// Resolves once no deferred readback is in flight (immediately without deferred readbacks).
    pub async fn settled(&self) {
        if let Some(d) = self.ctx.deferred() {
            d.settled().await;
        }
    }

    /// Deliver finished readbacks without blocking (native; the browser delivers them itself).
    pub fn poll(&self) {
        #[cfg(not(target_arch = "wasm32"))]
        let _ = self.ctx.poll_readbacks(wgpu::PollType::Poll);
    }

    /// `{passes, readbacks, inFlight}` of deferred rendering (diagnostics).
    pub fn deferred_stats(&self) -> Option<(u64, u64, usize)> {
        let d = self.ctx.deferred()?;
        Some((d.passes.load(std::sync::atomic::Ordering::Relaxed), d.readbacks.load(std::sync::atomic::Ordering::Relaxed), d.in_flight()))
    }

    pub fn context(&self) -> &GpuContext {
        &self.ctx
    }

    pub fn device(&self) -> &wgpu::Device {
        &self.ctx.device
    }

    /// Render a top-level comp frame and leave it on the GPU as a display texture (no
    /// readback unless a CPU fallback step needs one). `None` = render on the CPU.
    pub fn render_display(&self, r: &Renderer, comp: ItemId, t: Tick) -> Option<DisplayFrame> {
        self.ctx.check_health().map_err(|e| log::error!("gpu display: {e}")).ok()?;
        let mut e = Enc::new(&self.ctx);
        let img = walk::render(&mut e, r, comp, t)?;
        let texture = e.display(&img);
        e.submit();
        // Reject failures observed during rendering/submission before publishing the frame.
        // A later device fault still requires the presentation host's lifecycle handling.
        self.ctx.check_health().map_err(|e| log::error!("gpu display: {e}")).ok()?;
        Some(DisplayFrame { texture, width: img.width, height: img.height })
    }

    /// Read a display frame back as premultiplied RGBA8 (Info panel sampling, eyedroppers).
    pub fn read_display(&self, f: &DisplayFrame) -> Option<Vec<u8>> {
        Enc::new(&self.ctx).read_texture(&f.texture, f.width, f.height, 4)
    }

    /// [`Gpu::read_display`] without blocking (the browser's main thread): `done` gets the bytes
    /// once the GPU has them.
    pub fn read_display_async(&self, f: &DisplayFrame, done: impl FnOnce(Option<Vec<u8>>) + wgpu::WasmNotSend + 'static) {
        if let Err(e) = self.ctx.check_health() {
            log::error!("gpu display readback: {e}");
            crate::readback::reject(done, e);
            return;
        }
        Enc::new(&self.ctx).read_texture_async(&f.texture, f.width, f.height, 4, done);
    }

    /// Render a frame on the GPU and read it back (`None` = not handled).
    pub fn render(&self, r: &Renderer, comp: ItemId, t: Tick) -> Option<Image> {
        self.ctx.check_health().ok()?;
        if !self.ctx.can_readback() {
            return None;
        }
        if let Some(d) = self.ctx.deferred.clone() {
            return self.render_deferred(&d, r, comp, t);
        }
        let mut e = Enc::new(&self.ctx);
        let img = walk::render(&mut e, r, comp, t)?;
        e.download(&img)
    }

    /// [`Gpu::render`] with deferred readbacks: the frame's readback is keyed by comp, time and
    /// render options; until it arrives (or while layer steps still miss) the pass gets a
    /// placeholder.
    fn render_deferred(&self, d: &Arc<deferred::Deferred>, r: &Renderer, comp: ItemId, t: Tick) -> Option<Image> {
        let mut k = deferred::Key::new(2);
        k.u64(comp.0);
        k.u64(t.0 as u64);
        k.debug(&r.opts);
        let key = k.finish();
        let image = |rb: deferred::Readback| -> Option<Image> {
            let mut img = Image::new(rb.width, rb.height);
            let dst: &mut [u8] = bytemuck::cast_slice_mut(&mut img.data);
            (dst.len() == rb.bytes.len()).then(|| dst.copy_from_slice(&rb.bytes))?;
            Some(img)
        };
        match d.lookup(key) {
            deferred::Lookup::Ready(rb) => return rb.and_then(image),
            deferred::Lookup::InFlight => {
                d.miss();
                return Some(Image::new(1, 1));
            }
            deferred::Lookup::Absent => {}
        }
        let misses = d.misses();
        let mut e = Enc::new(&self.ctx);
        let img = walk::render(&mut e, r, comp, t)?;
        if d.misses() != misses {
            // Layer steps are still in flight: this pass's canvas is not the frame.
            return Some(Image::new(1, 1));
        }
        match e.read_keyed(&img, key, [0.0; 2], 1.0) {
            Some(rb) => rb.and_then(image),
            None => Some(Image::new(1, 1)),
        }
    }

    /// Run work in native thread-local out-of-memory, internal-error and validation scopes. A
    /// failed result is discarded and the textures and buffers kept for reuse are released;
    /// callers can retry with a fresh CPU renderer. Some drivers report a failed allocation as an
    /// internal error rather than out-of-memory (Adreno's `ERROR_INITIALIZATION_FAILED`, #591),
    /// so both count. A failed allocation leaves an invalid texture whose use is then a
    /// validation error (`Texture::create_view`, #1319): caught here, it fails this frame only
    /// instead of retiring the GPU compositor for the session as an uncaptured error.
    /// Browser callers use the host/owned-device failure bridge and bounded readbacks instead
    /// of blocking their event loop on an error-scope future.
    pub fn within_memory<T>(&self, f: impl FnOnce() -> T) -> Result<T, FrameError> {
        self.ctx.check_health().map_err(FrameError::Retired)?;
        #[cfg(not(target_arch = "wasm32"))]
        {
            let oom = self.ctx.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
            let internal = self.ctx.device.push_error_scope(wgpu::ErrorFilter::Internal);
            let validation = self.ctx.device.push_error_scope(wgpu::ErrorFilter::Validation);
            let out = f();
            let validation = pollster::block_on(validation.pop());
            let internal = pollster::block_on(internal.pop());
            let oom = pollster::block_on(oom.pop());
            let memory = oom.is_some() || internal.is_some();
            // The allocation failure first: a validation error is often its consequence.
            let errors: Vec<String> = [oom, internal, validation].into_iter().flatten().map(|e| e.to_string()).collect();
            if errors.is_empty() {
                return self.ctx.check_health().map(|()| out).map_err(FrameError::Retired);
            }
            // Its textures go back to the pool when dropped: drop them before emptying it.
            drop(out);
            self.ctx.discard_kept();
            let errors = errors.join("; then: ");
            Err(if memory { FrameError::Memory(errors) } else { FrameError::Validation(errors) })
        }
        #[cfg(target_arch = "wasm32")]
        {
            let out = f();
            self.ctx.check_health().map(|()| out).map_err(FrameError::Retired)
        }
    }

    /// Why a frame was not rendered (best effort: device health, else size / support).
    fn decline_reason(&self) -> String {
        match self.ctx.check_health() {
            Err(e) => format!("device failure / out of memory: {e}"),
            Ok(()) if !self.ctx.can_readback() => "no GPU readback is available here".into(),
            Ok(()) => format!(
                "the GPU compositor did not handle the frame (a region of interest, unsupported content, a frame over the {} px texture limit, or an allocation failed)",
                self.ctx.max_dim
            ),
        }
    }

    /// Wait for submitted GPU work with the readback deadline (benchmarks).
    pub fn wait(&self) {
        let _ = self.ctx.poll_readbacks(wgpu::PollType::Wait { submission_index: None, timeout: Some(crate::readback::TIMEOUT) });
    }
}

impl Accelerator for Gpu {
    fn name(&self) -> String {
        self.ctx.name.clone()
    }

    fn available(&self) -> bool {
        self.ctx.check_health().is_ok()
    }

    fn comp_frame(&self, r: &Renderer, comp: ItemId, t: Tick) -> Option<Image> {
        let img = self.render(r, comp, t);
        if img.is_none() {
            // The caller renders the frame on the CPU; say so and why, or an export that falls
            // back frame after frame is only mysteriously slow. Once per comp and reason as a
            // warning (the log keeps only the last few hundred, for the System Report), then at
            // debug level.
            let reason = self.decline_reason();
            let first = self.declined.lock().map(|mut seen| seen.len() < 256 && seen.insert((comp.0, reason.clone()))).unwrap_or(false);
            if first {
                log::warn!(target: "effectcraft_gpu", "GPU declined the frame at {:.3} s, the CPU renders it: {reason}", t.seconds());
            } else {
                log::debug!(target: "effectcraft_gpu", "GPU declined the frame at {:.3} s: {reason}", t.seconds());
            }
        }
        img
    }

    fn comp_frame_yuv420(&self, r: &Renderer, comp: ItemId, t: Tick, bg: [f32; 3]) -> Option<Vec<u8>> {
        self.ctx.check_health().ok()?;
        if !self.ctx.can_readback() || self.ctx.deferred.is_some() {
            return None;
        }
        // Sizes the kernel can't take fall back before any GPU work, not after a render.
        let (w, h) = walk::frame_size(r, comp, t)?;
        if !yuv::supported(w, h) || !yuv::fits_limits(w, h, &self.ctx.device.limits()) {
            return None;
        }
        let yp = self.ctx.yuv_pipe()?;
        let mut e = Enc::new(&self.ctx);
        let img = walk::render(&mut e, r, comp, t)?;
        let planes = yuv::convert(&mut e, yp, &img, bg)?;
        // A device fault seen while the frame recorded makes the planes untrustworthy.
        self.ctx.check_health().ok()?;
        Some(planes)
    }

    fn supports_effect(&self, id: &str) -> bool {
        self.ctx.check_health().is_ok() && self.ctx.can_readback() && effects::supports(id)
    }

    fn effects(&self, chain: &[FxStep], buf: &Buf, levels: Option<f32>) -> Option<Buf> {
        self.ctx.check_health().ok()?;
        effects::run_chain(&mut Enc::lent(&self.ctx), chain, buf, levels)
    }

    fn raster_3d(&self, scene: &effectcraft_render::three_d::adv::Scene) -> Option<effectcraft_render::three_d::adv::Target> {
        self.ctx.check_health().ok()?;
        adv3d::render(&self.ctx, scene)
    }

    fn render_3d(&self, run: &effectcraft_render::three_d::adv::Prepared) -> Option<effectcraft_render::three_d::adv::Rendered> {
        self.ctx.check_health().ok()?;
        adv3d::render_prepared(&self.ctx, run)
    }

    fn particles(&self) -> Option<&dyn effectcraft_effects::psim::ParticleSim> {
        self.ctx.check_health().ok()?;
        self.ctx.can_readback().then_some(self as &dyn effectcraft_effects::psim::ParticleSim)
    }

    fn inpaint(&self) -> Option<&dyn effectcraft_raster::inpaint::FillBackend> {
        self.ctx.check_health().ok()?;
        // Readback is synchronous-only, and a kernel that failed to build would make every stage
        // fall back to the CPU *after* paying the upload and a failed dispatch (and the fill
        // would still report the GPU path). Neither is a backend at all.
        if !self.ctx.can_wait() || self.ctx.unavailable_kernels().iter().any(|(e, _)| crate::inpaint::KERNELS.contains(e)) {
            return None;
        }
        Some(self as &dyn effectcraft_raster::inpaint::FillBackend)
    }

    fn auto_pick(&self) -> Option<&effectcraft_render::AutoPick> {
        Some(&self.auto)
    }

    fn render_end(&self) {
        let (hits, misses) = self.ctx.pool_stats();
        let pct = hits.saturating_mul(100) / hits.saturating_add(misses).max(1);
        log::debug!(
            target: "effectcraft_gpu",
            "gpu pool: {hits} working textures reused, {misses} allocated ({pct}% reuse) so far"
        );
        self.ctx.trim_pool();
    }

    fn frame_begin(&self) {
        if let Some(d) = self.ctx.deferred() {
            d.frame_begin();
        }
    }

    fn frame_end(&self) {
        self.ctx.pool_frame_end();
    }

    fn pass_begin(&self) {
        if let Some(d) = self.ctx.deferred() {
            d.pass_begin();
        }
    }

    fn pass_missed(&self) -> bool {
        self.ctx.deferred().is_some_and(|d| d.missed())
    }

    fn miss_gate(&self) -> Option<Arc<std::sync::atomic::AtomicBool>> {
        self.ctx.deferred().map(|d| d.gate())
    }

    fn settle(&self) -> Option<std::pin::Pin<Box<dyn std::future::Future<Output = ()> + '_>>> {
        let d = self.ctx.deferred()?.clone();
        // Natively the device delivers the readbacks when polled: wait for them here.
        #[cfg(not(target_arch = "wasm32"))]
        if d.in_flight() > 0 {
            let _ = self.ctx.poll_readbacks(wgpu::PollType::Wait { submission_index: None, timeout: Some(crate::readback::TIMEOUT) });
        }
        Some(Box::pin(async move { d.settled().await }))
    }
}

impl effectcraft_effects::psim::ParticleSim for Gpu {
    fn simulate(&self, req: &effectcraft_effects::psim::SimRequest) -> Option<Vec<effectcraft_effects::psim::SimParticle>> {
        self.ctx.check_health().ok()?;
        particles::simulate(&self.ctx, req)
    }
}

#[cfg(test)]
mod tests;