open-standards/wai/wai-ops-gpu/src/splat.rs

//! The compositing of the splat renderer `wai-splat-render:i1` (extensions/splat-exact.md
//! §5.6–§5.7) on a GPU: `wgsl/splat.wgsl`, one invocation a pixel and one workgroup a
//! 16 × 16 tile, after `wgsl/arith.wgsl`, whose 64-bit products the device has passed the
//! arithmetic probe on ([`crate::probe`]).
//!
//! The host projects, culls and orders the Gaussians (§5.3–§5.5) and gives the device the
//! kept ones in order, [`SCREEN_WORDS`] words each; the device composes every pixel over
//! them. Every pixel is independent, so the frame is the host's whatever order the device
//! runs its workgroups in.

use std::borrow::Cow;
use std::num::NonZeroU64;

use crate::device::{Gpu, GpuError, Scopes};

/// The kernel's text, without the arithmetic it follows.
pub const SPLAT: &str = include_str!("wgsl/splat.wgsl");

/// Words of one Gaussian: `m_x`, `m_y`, `x0`, `x1`, `y0`, `y1`, `k1`, `k2`, `k3`, the opacity
/// and the three colour values (§5.5).
pub const SCREEN_WORDS: usize = 13;

/// The largest frame side the kernel composes (§2).
pub const MAX_SIDE: u32 = 1024;

/// The shipped module: the arithmetic, then the kernel.
pub fn source() -> String {
    format!("{}\n{}", crate::kernels::ARITH, SPLAT)
}

impl Gpu {
    /// The frame of `screens` (the kept Gaussians in order), `width × height` pixels of RGB8,
    /// rows top to bottom.
    pub async fn splat_composite(&self, screens: &[[i32; SCREEN_WORDS]], width: u32, height: u32) -> Result<Vec<u8>, GpuError> {
        self.splat_composite_with(&source(), screens, width, height).await
    }

    /// [`Gpu::splat_composite`] with the module compiled from `text` (a test's mutant).
    pub async fn splat_composite_with(&self, text: &str, screens: &[[i32; SCREEN_WORDS]], width: u32, height: u32) -> Result<Vec<u8>, GpuError> {
        if width == 0 || height == 0 || width > MAX_SIDE || height > MAX_SIDE || screens.len() > u32::MAX as usize / SCREEN_WORDS {
            return Err(GpuError::Device(format!("a frame of {width} × {height} is outside the kernel's sizes")));
        }
        let scopes = Scopes::push(&self.device);
        let entry = |binding: u32, ty: wgpu::BufferBindingType| wgpu::BindGroupLayoutEntry {
            binding,
            visibility: wgpu::ShaderStages::COMPUTE,
            ty: wgpu::BindingType::Buffer { ty, has_dynamic_offset: false, min_binding_size: None },
            count: None,
        };
        let layout = self.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("wai-ops-gpu splat"),
            entries: &[
                entry(0, wgpu::BufferBindingType::Uniform),
                entry(1, wgpu::BufferBindingType::Storage { read_only: true }),
                entry(2, wgpu::BufferBindingType::Storage { read_only: false }),
            ],
        });
        let pl = self.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("wai-ops-gpu splat"), bind_group_layouts: &[Some(&layout)], immediate_size: 0 });
        let module = self.device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("wai-ops-gpu splat"), source: wgpu::ShaderSource::Wgsl(Cow::Owned(text.to_string())) });
        let pipe = self.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("splat"),
            layout: Some(&pl),
            module: &module,
            entry_point: Some("splat"),
            compilation_options: Default::default(),
            cache: None,
        });
        let params = self.buffer("wai-ops-gpu splat params", 16, wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST);
        let words: Vec<u8> = [screens.len() as u32, width, height, 0].iter().flat_map(|w| w.to_le_bytes()).collect();
        self.queue.write_buffer(&params, 0, &words);
        let list: Vec<u8> = screens.iter().flatten().flat_map(|w| w.to_le_bytes()).collect();
        let list_size = (list.len() as u64).max(16);
        let list_buf = self.buffer("wai-ops-gpu splat list", list_size.div_ceil(4) * 4, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST);
        if !list.is_empty() {
            self.queue.write_buffer(&list_buf, 0, &list);
        }
        let pixels = u64::from(width) * u64::from(height);
        let out = self.buffer("wai-ops-gpu splat frame", pixels * 4, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC);
        let staging = self.buffer("wai-ops-gpu splat staging", pixels * 4, wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST);
        let bind = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("wai-ops-gpu splat"),
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { buffer: &params, offset: 0, size: NonZeroU64::new(16) }) },
                wgpu::BindGroupEntry { binding: 1, resource: list_buf.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 2, resource: out.as_entire_binding() },
            ],
        });
        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("wai-ops-gpu splat") });
        {
            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("wai-ops-gpu splat"), timestamp_writes: None });
            pass.set_pipeline(&pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(width.div_ceil(16), height.div_ceil(16), 1);
        }
        enc.copy_buffer_to_buffer(&out, 0, &staging, 0, pixels * 4);
        let index = self.queue.submit([enc.finish()]);
        scopes.close().await?;
        let bytes = self.read(&staging, pixels * 4, index).await?;
        self.settled()?;
        let mut rgb = Vec::with_capacity(pixels as usize * 3);
        for w in bytes.as_chunks::<4>().0 {
            rgb.extend_from_slice(&w[..3]);
        }
        Ok(rgb)
    }
}

/// The calls for a native caller without an async runtime.
#[cfg(not(target_arch = "wasm32"))]
pub mod blocking {
    use super::SCREEN_WORDS;
    use crate::device::GpuError;
    use crate::device::blocking::resolve;

    /// [`crate::Gpu::splat_composite`].
    pub fn composite(g: &crate::Gpu, screens: &[[i32; SCREEN_WORDS]], width: u32, height: u32) -> Result<Vec<u8>, GpuError> {
        resolve(g.splat_composite(screens, width, height))?
    }

    /// [`crate::Gpu::splat_composite_with`].
    pub fn composite_with(g: &crate::Gpu, text: &str, screens: &[[i32; SCREEN_WORDS]], width: u32, height: u32) -> Result<Vec<u8>, GpuError> {
        resolve(g.splat_composite_with(text, screens, width, height))?
    }
}