effectcraft/crates/gpu/src/context.rs

1054 lines
42 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//! Device, pipelines, textures, uploads and readback.
use std::collections::{BTreeSet, HashMap};
use std::sync::{Arc, Mutex, Weak};
use effectcraft_effects::Buf;
use effectcraft_raster::Image;
use wgpu::util::DeviceExt;
/// Compute entry points in `kernels.wgsl` (one pipeline each).
const ENTRIES: &[&str] = &[
"warp_blend",
"blend_full",
"matte",
"preserve",
"knockout",
"channel_mix",
"quantize",
"convert",
"half",
"box_h",
"box_v",
"directional",
"glow_bright",
"glow_combine",
"shadow_make",
"shadow_combine",
"pointwise",
"adjust_mix",
"bokeh_boost",
"bokeh_prefix",
"bokeh_gather",
"fill",
];
/// Entry points that also bind group 1 (four read-only storage buffers; see
/// [`Enc::dispatch_ext`]).
const EXT_ENTRIES: &[&str] = &["classic3d"];
/// Pixel format of every working texture: premultiplied RGBA, 32-bit float (32 bpc headroom;
/// 8/16 bpc are emulated by clamping and quantising, as on the CPU).
pub const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba32Float;
/// Uploaded layer buffers kept on the GPU (bytes).
const UPLOAD_BUDGET: usize = 1536 << 20;
/// Working textures kept for reuse once their last image is dropped (bytes).
const POOL_BUDGET: usize = 768 << 20;
/// A GPU image: premultiplied RGBA f32 in a texture (cheap to clone; immutable once written).
/// Working images come from the context's texture pool and go back to it when the last clone
/// is dropped.
#[derive(Clone, Debug)]
pub struct GpuImage {
pub texture: wgpu::Texture,
pub width: u32,
pub height: u32,
_lease: Option<Arc<Lease>>,
/// Already clamped and quantised to this many levels (8/16 bpc), so quantising again is a
/// no-op (see `ops::quantize`).
pub(crate) levels: Option<f32>,
}
impl GpuImage {
/// An image on a texture of its own (not pooled).
pub fn new(texture: wgpu::Texture, width: u32, height: u32) -> GpuImage {
GpuImage { texture, width, height, _lease: None, levels: None }
}
}
/// Hands a pooled texture back to its [`Pool`] when the last [`GpuImage`] holding it drops.
struct Lease {
texture: wgpu::Texture,
pool: Weak<Mutex<Pool>>,
}
impl std::fmt::Debug for Lease {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("Lease")
}
}
impl Drop for Lease {
fn drop(&mut self) {
if let Some(pool) = self.pool.upgrade()
&& let Ok(mut p) = pool.lock()
{
p.release(self.texture.clone());
}
}
}
/// Recycled working textures: allocating (and zero-initialising) a 1080p RGBA f32 texture
/// costs more than compositing a small layer into it, and a frame makes several per layer.
///
/// A released texture may still be read by commands recorded in an encoder that has not been
/// submitted yet. Reusing it in another encoder that is submitted first would overwrite it
/// before those reads run, so a released texture waits until every encoder that was open when
/// it was released has been submitted (or dropped).
#[derive(Default)]
pub(crate) struct Pool {
free: HashMap<(u32, u32), Vec<wgpu::Texture>>,
/// (texture, newest encoder open at its release).
released: Vec<(wgpu::Texture, u64)>,
/// Encoders with recorded, unsubmitted work.
open: BTreeSet<u64>,
next: u64,
bytes: usize,
/// Working textures reused / newly allocated (diagnostics, tests).
hits: u64,
misses: u64,
}
/// Counts of uploads (CPU → GPU) and readbacks (GPU → CPU).
#[derive(Default)]
pub(crate) struct Transfers {
ups: std::sync::atomic::AtomicU64,
up_bytes: std::sync::atomic::AtomicU64,
downs: std::sync::atomic::AtomicU64,
down_bytes: std::sync::atomic::AtomicU64,
}
/// A snapshot of [`GpuContext::transfer_stats`].
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct TransferStats {
pub uploads: u64,
pub upload_bytes: u64,
pub readbacks: u64,
pub readback_bytes: u64,
}
impl Transfers {
fn count(&self, up: bool, bytes: usize) {
use std::sync::atomic::Ordering::Relaxed;
let (n, b) = if up { (&self.ups, &self.up_bytes) } else { (&self.downs, &self.down_bytes) };
n.fetch_add(1, Relaxed);
b.fetch_add(bytes as u64, Relaxed);
}
}
fn tex_bytes(t: &wgpu::Texture) -> usize {
t.width() as usize * t.height() as usize * 16
}
impl Pool {
fn release(&mut self, t: wgpu::Texture) {
match self.open.last() {
Some(&newest) => self.released.push((t, newest)),
None => self.make_free(t),
}
}
fn make_free(&mut self, t: wgpu::Texture) {
let b = tex_bytes(&t);
if self.bytes + b > POOL_BUDGET {
return;
}
self.bytes += b;
self.free.entry((t.width(), t.height())).or_default().push(t);
}
fn open(&mut self) -> u64 {
self.next += 1;
self.open.insert(self.next);
self.next
}
/// Encoder `id` was submitted or dropped: free what no open encoder can still read.
fn close(&mut self, id: u64) {
self.open.remove(&id);
let oldest = self.open.first().copied().unwrap_or(u64::MAX);
if self.released.iter().all(|(_, w)| *w >= oldest) {
return;
}
let (ready, wait): (Vec<_>, Vec<_>) = std::mem::take(&mut self.released).into_iter().partition(|(_, w)| *w < oldest);
self.released = wait;
for (t, _) in ready {
self.make_free(t);
}
}
fn take(&mut self, w: u32, h: u32) -> Option<wgpu::Texture> {
let t = self.free.get_mut(&(w, h)).and_then(Vec::pop);
match &t {
Some(t) => {
self.bytes -= tex_bytes(t);
self.hits += 1;
}
None => self.misses += 1,
}
t
}
fn clear(&mut self) {
self.free.clear();
self.bytes = 0;
}
}
/// Uniform parameter block shared by every kernel (see `shaders/common.wgsl`).
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct Params {
pub u: [[u32; 4]; 4],
pub f: [[f32; 4]; 12],
}
impl Params {
fn bytes(&self) -> Vec<u8> {
let mut v = Vec::with_capacity(256);
for r in &self.u {
for x in r {
v.extend_from_slice(&x.to_le_bytes());
}
}
for r in &self.f {
for x in r {
v.extend_from_slice(&x.to_le_bytes());
}
}
v
}
}
struct Upload {
buf: Weak<Buf>,
img: GpuImage,
bytes: usize,
last_use: u64,
}
#[derive(Default)]
struct Uploads {
map: HashMap<usize, Upload>,
bytes: usize,
clock: u64,
}
/// The device and everything compiled for it.
pub struct GpuContext {
pub device: wgpu::Device,
pub queue: wgpu::Queue,
pub(crate) name: String,
pub(crate) max_dim: u32,
bgl: wgpu::BindGroupLayout,
bgl_ext: wgpu::BindGroupLayout,
pipelines: HashMap<&'static str, wgpu::ComputePipeline>,
display_bgl: wgpu::BindGroupLayout,
display: wgpu::ComputePipeline,
dummy: wgpu::Texture,
dummy_tex: wgpu::TextureView,
dummy_buf: wgpu::Buffer,
uploads: Mutex<Uploads>,
pool: Arc<Mutex<Pool>>,
/// CPU ↔ GPU traffic (diagnostics: `bench --gpu`).
pub(crate) transfers: Transfers,
/// Readback staging buffers by size.
staging: Mutex<HashMap<u64, Vec<wgpu::Buffer>>>,
/// Transparent read-only images by size (see [`Enc::zeros`]).
zeros: Mutex<HashMap<(u32, u32), GpuImage>>,
/// Advanced 3D render pipelines (built on first use).
pub(crate) adv3d: std::sync::OnceLock<crate::adv3d::Pipes>,
/// The adapter can run the Advanced 3D rasteriser (`adv3d::raster_unsupported`); without
/// it Advanced 3D scenes render on the CPU.
pub(crate) adv3d_raster: bool,
/// GPU particle pipeline (built on first use) and simulation checkpoints.
pub(crate) particles: crate::particles::PipesCell,
pub(crate) particle_states: crate::particles::StatesCell,
/// Deferred readbacks (a browser worker's WebGPU device), see [`crate::deferred`].
pub(crate) deferred: Option<Arc<crate::deferred::Deferred>>,
}
fn tex_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
wgpu::BindGroupLayoutEntry {
binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
}
}
fn storage_tex_entry(binding: u32, format: wgpu::TextureFormat) -> wgpu::BindGroupLayoutEntry {
wgpu::BindGroupLayoutEntry {
binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::StorageTexture { access: wgpu::StorageTextureAccess::WriteOnly, format, view_dimension: wgpu::TextureViewDimension::D2 },
count: None,
}
}
impl GpuContext {
/// Build on an existing device (the desktop app shares egui-wgpu's). `Err` when the
/// adapter cannot run the compositor (no compute shaders, e.g. WebGL2; no float storage
/// textures) or the device was created with limits below what its pipelines bind.
pub fn new(adapter: &wgpu::Adapter, device: wgpu::Device, queue: wgpu::Queue) -> Result<GpuContext, String> {
let info = adapter.get_info();
if !adapter.get_downlevel_capabilities().flags.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS) {
return Err(format!("{} ({:?}): no compute shaders", info.name, info.backend));
}
// egui-wgpu asks for WebGL2's limits on GL (no storage buffers at all): building the
// pipelines there would be a validation error, which release builds only log.
let l = device.limits();
let needs = [
// `adv3d`'s post kernels bind 7 (its rasteriser's fragment stage 6).
("max_storage_buffers_per_shader_stage", l.max_storage_buffers_per_shader_stage, 7),
("max_storage_textures_per_shader_stage", l.max_storage_textures_per_shader_stage, 1),
("max_bind_groups", l.max_bind_groups, 2),
("max_compute_workgroup_size_x", l.max_compute_workgroup_size_x, 64),
("max_compute_workgroup_size_y", l.max_compute_workgroup_size_y, 16),
("max_compute_invocations_per_workgroup", l.max_compute_invocations_per_workgroup, 256),
];
if let Some((limit, have, need)) = needs.iter().find(|(_, have, need)| have < need) {
return Err(format!("{} ({:?}): device limit {limit} is {have}, needs {need}", info.name, info.backend));
}
for f in [FORMAT, wgpu::TextureFormat::Rgba8Unorm] {
if !adapter.get_texture_format_features(f).allowed_usages.contains(wgpu::TextureUsages::STORAGE_BINDING) {
return Err(format!("{} ({:?}): {f:?} storage textures unsupported", info.name, info.backend));
}
}
let name = format!("{} ({:?})", info.name, info.backend);
let adv3d_unsupported = crate::adv3d::raster_unsupported(adapter);
if let Some(why) = &adv3d_unsupported {
log::info!("gpu {name}: Advanced 3D renders on the CPU ({why})");
}
let src = [
include_str!("shaders/common.wgsl"),
include_str!("shaders/kernels.wgsl"),
include_str!("shaders/classic3d.wgsl"),
include_str!("shaders/fx_color.wgsl"),
include_str!("shaders/fx_distort.wgsl"),
include_str!("shaders/fx_generate.wgsl"),
include_str!("shaders/fx_key.wgsl"),
include_str!("shaders/fx_stylize.wgsl"),
include_str!("shaders/fx_noise.wgsl"),
include_str!("shaders/fx_tone.wgsl"),
include_str!("shaders/fx_warp.wgsl"),
include_str!("shaders/fx_extra.wgsl"),
include_str!("shaders/fx_depth.wgsl"),
include_str!("shaders/fx_lut.wgsl"),
include_str!("shaders/fx_sim.wgsl"),
include_str!("shaders/fx_particles.wgsl"),
include_str!("shaders/fx_vr.wgsl"),
include_str!("shaders/fx_light.wgsl"),
include_str!("shaders/fx_transition.wgsl"),
include_str!("shaders/fx_text.wgsl"),
include_str!("shaders/fx_time.wgsl"),
include_str!("shaders/fx_pixel2.wgsl"),
include_str!("shaders/fx_gen2.wgsl"),
include_str!("shaders/sky.wgsl"),
]
.concat();
let module =
device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("effectcraft kernels"), source: wgpu::ShaderSource::Wgsl(src.into()) });
let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("effectcraft kernels"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None },
count: None,
},
tex_entry(1),
tex_entry(2),
storage_tex_entry(3, FORMAT),
wgpu::BindGroupLayoutEntry {
binding: 4,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("effectcraft kernels"),
bind_group_layouts: &[Some(&bgl)],
immediate_size: 0,
});
let storage = |binding: u32| wgpu::BindGroupLayoutEntry {
binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None },
count: None,
};
let bgl_ext = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("effectcraft kernels (ext)"),
entries: &[storage(0), storage(1), storage(2), storage(3)],
});
let layout_ext = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("effectcraft kernels (ext)"),
bind_group_layouts: &[Some(&bgl), Some(&bgl_ext)],
immediate_size: 0,
});
let pipelines = ENTRIES
.iter()
.chain(crate::fx_color::KERNELS)
.chain(crate::fx_distort::KERNELS)
.chain(crate::fx_generate::KERNELS)
.chain(crate::fx_key::KERNELS)
.chain(crate::fx_stylize::KERNELS)
.chain(crate::fx_noise::KERNELS)
.chain(crate::fx_tone::KERNELS)
.chain(crate::fx_warp::KERNELS)
.chain(crate::fx_extra::KERNELS)
.chain(crate::fx_depth::KERNELS)
.chain(crate::fx_lut::KERNELS)
.chain(crate::fx_sim::KERNELS)
.chain(crate::fx_particles::KERNELS)
.chain(crate::fx_vr::KERNELS)
.chain(crate::fx_light::KERNELS)
.chain(crate::fx_transition::KERNELS)
.chain(crate::fx_text::KERNELS)
.chain(crate::fx_time::KERNELS)
.chain(crate::fx_pixel2::KERNELS)
.chain(crate::fx_gen2::KERNELS)
.chain(crate::adv3d::SKY_KERNELS)
.map(|e| (e, &layout))
.chain(EXT_ENTRIES.iter().map(|e| (e, &layout_ext)))
.map(|(e, layout)| {
let p = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(e),
layout: Some(layout),
module: &module,
entry_point: Some(e),
compilation_options: Default::default(),
cache: None,
});
(*e, p)
})
.collect();
let dmodule = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("effectcraft display"),
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/display.wgsl").into()),
});
let display_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("effectcraft display"),
entries: &[tex_entry(1), storage_tex_entry(3, wgpu::TextureFormat::Rgba8Unorm)],
});
let dlayout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("effectcraft display"),
bind_group_layouts: &[Some(&display_bgl)],
immediate_size: 0,
});
let display = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("display"),
layout: Some(&dlayout),
module: &dmodule,
entry_point: Some("display"),
compilation_options: Default::default(),
cache: None,
});
let dummy = device.create_texture(&wgpu::TextureDescriptor {
label: Some("dummy"),
size: wgpu::Extent3d { width: 1, height: 1, depth_or_array_layers: 1 },
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: FORMAT,
usage: wgpu::TextureUsages::TEXTURE_BINDING,
view_formats: &[],
});
let dummy_buf =
device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("dummy"), contents: &[0u8; 16], usage: wgpu::BufferUsages::STORAGE });
let max_dim = device.limits().max_texture_dimension_2d;
Ok(GpuContext {
device,
queue,
name,
max_dim,
bgl,
bgl_ext,
pipelines,
display_bgl,
display,
dummy_tex: dummy.create_view(&Default::default()),
dummy,
dummy_buf,
uploads: Mutex::new(Uploads::default()),
pool: Default::default(),
staging: Default::default(),
transfers: Default::default(),
zeros: Default::default(),
adv3d: std::sync::OnceLock::new(),
adv3d_raster: adv3d_unsupported.is_none(),
particles: Default::default(),
particle_states: Default::default(),
deferred: None,
})
}
/// Switch to deferred readbacks (see [`crate::deferred`]): readbacks are never waited for.
pub fn set_deferred(&mut self, on: bool) {
self.deferred = on.then(Default::default);
}
pub fn deferred(&self) -> Option<&Arc<crate::deferred::Deferred>> {
self.deferred.as_ref()
}
/// A device of its own on the best adapter, without blocking (a browser worker: WebGPU
/// only; natively any backend). `Err` says why there is none.
pub async fn request() -> Result<GpuContext, String> {
#[allow(unused_mut)]
let mut desc = wgpu::InstanceDescriptor::new_without_display_handle();
#[cfg(target_arch = "wasm32")]
{
desc.backends = wgpu::Backends::BROWSER_WEBGPU;
}
let instance = wgpu::Instance::new(desc);
let adapter = instance
.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() })
.await
.map_err(|e| format!("no adapter: {e}"))?;
let limits = adapter.limits();
let required_limits = wgpu::Limits {
max_texture_dimension_2d: limits.max_texture_dimension_2d.min(16384),
max_buffer_size: limits.max_buffer_size,
..wgpu::Limits::default()
};
let (device, queue) = adapter
.request_device(&wgpu::DeviceDescriptor { label: Some("effectcraft gpu"), required_limits, ..Default::default() })
.await
.map_err(|e| format!("no device: {e}"))?;
device.on_uncaptured_error(Arc::new(|e| {
// Under test a validation error (a WGSL typo invalidates every kernel) fails loudly.
#[cfg(test)]
#[allow(clippy::panic)] // test builds only
{
panic!("wgpu: {e}")
}
#[cfg(not(test))]
log::error!("wgpu: {e}")
}));
GpuContext::new(&adapter, device, queue)
}
/// A device of its own on the best adapter (CLI, tests, benchmarks). `None` when no
/// adapter qualifies.
#[cfg(not(target_arch = "wasm32"))]
pub fn headless() -> Option<GpuContext> {
pollster::block_on(GpuContext::request()).map_err(|e| log::info!("gpu: {e}")).ok()
}
/// Readbacks (GPU → CPU) are possible: waited for natively, or deferred (a browser
/// worker, [`crate::deferred`]). Impossible on the browser's main thread.
pub fn can_readback(&self) -> bool {
self.can_wait() || self.deferred.is_some()
}
/// Readbacks can be waited for (natively, unless deferred). Steps that read back without a
/// key ([`crate::deferred`]) need this.
pub fn can_wait(&self) -> bool {
cfg!(not(target_arch = "wasm32")) && self.deferred.is_none()
}
pub(crate) fn fits(&self, w: u32, h: u32) -> bool {
w > 0 && h > 0 && w <= self.max_dim && h <= self.max_dim
}
/// A new (zeroed) texture of its own (not pooled).
pub(crate) fn image(&self, w: u32, h: u32) -> GpuImage {
GpuImage::new(self.texture(w, h), w, h)
}
fn texture(&self, w: u32, h: u32) -> wgpu::Texture {
self.device.create_texture(&wgpu::TextureDescriptor {
label: None,
size: wgpu::Extent3d { width: w.max(1), height: h.max(1), depth_or_array_layers: 1 },
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: FORMAT,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
})
}
/// A pooled working texture: (image, reused). A reused texture holds old pixels; a new
/// one is zeroed.
fn pooled(&self, w: u32, h: u32) -> (GpuImage, bool) {
let (w, h) = (w.max(1), h.max(1));
let reused = self.pool.lock().ok().and_then(|mut p| p.take(w, h));
let was_reused = reused.is_some();
let texture = reused.unwrap_or_else(|| self.texture(w, h));
let lease = Arc::new(Lease { texture: texture.clone(), pool: Arc::downgrade(&self.pool) });
(GpuImage { texture, width: w, height: h, _lease: Some(lease), levels: None }, was_reused)
}
/// Uploads and readbacks so far (images; uploads of cached layers count once).
pub fn transfer_stats(&self) -> TransferStats {
use std::sync::atomic::Ordering::Relaxed;
let t = &self.transfers;
TransferStats {
uploads: t.ups.load(Relaxed),
upload_bytes: t.up_bytes.load(Relaxed),
readbacks: t.downs.load(Relaxed),
readback_bytes: t.down_bytes.load(Relaxed),
}
}
/// Texture pool counters: (working textures reused, newly allocated).
pub fn pool_stats(&self) -> (u64, u64) {
self.pool.lock().map(|p| (p.hits, p.misses)).unwrap_or_default()
}
/// Upload a CPU image (`None` when it does not fit the device).
pub fn upload_image(&self, img: &Image) -> Option<GpuImage> {
if !self.fits(img.width, img.height) {
return None;
}
self.transfers.count(true, img.data.len() * 16);
let g = self.image(img.width, img.height);
self.queue.write_texture(
wgpu::TexelCopyTextureInfo { texture: &g.texture, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All },
bytemuck::cast_slice(&img.data),
wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(img.width * 16), rows_per_image: Some(img.height) },
wgpu::Extent3d { width: img.width, height: img.height, depth_or_array_layers: 1 },
);
Some(g)
}
/// Upload a (layer cache) buffer, reusing the texture while the same `Arc` is alive: static
/// layers upload once, not every frame.
pub(crate) fn upload_buf(&self, buf: &Arc<Buf>) -> Option<GpuImage> {
let key = Arc::as_ptr(buf) as usize;
if let Ok(mut u) = self.uploads.lock() {
u.clock += 1;
let now = u.clock;
if let Some(e) = u.map.get_mut(&key)
&& e.buf.upgrade().is_some_and(|b| Arc::ptr_eq(&b, buf))
{
e.last_use = now;
return Some(e.img.clone());
}
}
let img = self.upload_image(&buf.img)?;
if let Ok(mut u) = self.uploads.lock() {
let bytes = buf.img.data.len() * 16;
let now = u.clock;
if let Some(old) = u.map.insert(key, Upload { buf: Arc::downgrade(buf), img: img.clone(), bytes, last_use: now }) {
u.bytes -= old.bytes;
}
u.bytes += bytes;
// Drop textures of freed buffers, then the least recently used over budget.
if u.bytes > UPLOAD_BUDGET {
let dead: Vec<usize> = u.map.iter().filter(|(_, e)| e.buf.strong_count() == 0).map(|(k, _)| *k).collect();
for k in dead {
if let Some(e) = u.map.remove(&k) {
u.bytes -= e.bytes;
}
}
}
while u.bytes > UPLOAD_BUDGET {
let Some(k) = u.map.iter().min_by_key(|(_, e)| e.last_use).map(|(k, _)| *k) else { break };
if let Some(e) = u.map.remove(&k) {
u.bytes -= e.bytes;
}
}
}
Some(img)
}
/// Forget uploaded layer textures.
pub fn clear_uploads(&self) {
if let Ok(mut u) = self.uploads.lock() {
u.map.clear();
u.bytes = 0;
}
if let Ok(mut p) = self.pool.lock() {
p.clear();
}
if let Ok(mut s) = self.staging.lock() {
s.clear();
}
}
}
/// Records GPU work for one render (a command encoder submitted at readback or the end).
pub(crate) struct Enc<'g> {
pub g: &'g GpuContext,
/// The open command encoder and its id in the texture pool.
enc: Option<(wgpu::CommandEncoder, u64)>,
pending: usize,
}
impl Drop for Enc<'_> {
fn drop(&mut self) {
// Unsubmitted work is discarded, so nothing recorded will read pooled textures.
if let Some((_, id)) = self.enc.take()
&& let Ok(mut p) = self.g.pool.lock()
{
p.close(id);
}
}
}
impl<'g> Enc<'g> {
pub fn new(g: &'g GpuContext) -> Enc<'g> {
Enc { g, enc: None, pending: 0 }
}
pub(crate) fn encoder(&mut self) -> &mut wgpu::CommandEncoder {
let g = self.g;
&mut self
.enc
.get_or_insert_with(|| {
let id = g.pool.lock().map(|mut p| p.open()).unwrap_or(0);
(g.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("effectcraft") }), id)
})
.0
}
/// A new zeroed image (pooled; a reused texture is cleared by a kernel).
pub fn image(&mut self, w: u32, h: u32) -> GpuImage {
let (img, reused) = self.g.pooled(w, h);
if reused {
self.fill(&img, 0.0);
}
img
}
/// A new image for a kernel that writes every pixel of it (pooled; contents undefined).
/// Under test it starts out as NaNs, so a kernel that skips pixels fails the CPU oracle.
pub fn scratch(&mut self, w: u32, h: u32) -> GpuImage {
let (img, _) = self.g.pooled(w, h);
if cfg!(test) {
self.fill(&img, f32::NAN);
}
img
}
/// A transparent image to read from (shared per size; never written).
pub fn zeros(&self, w: u32, h: u32) -> GpuImage {
let (w, h) = (w.max(1), h.max(1));
let Ok(mut z) = self.g.zeros.lock() else { return self.g.image(w, h) };
if z.len() > 8 && !z.contains_key(&(w, h)) {
z.clear();
}
z.entry((w, h)).or_insert_with(|| self.g.image(w, h)).clone()
}
/// Set every pixel of `img` to `v`.
fn fill(&mut self, img: &GpuImage, v: f32) {
let mut p = Params::default();
p.f[0] = [v; 4];
let dummy = GpuImage::new(self.g.dummy.clone(), 1, 1);
self.pixels("fill", &p, &dummy, None, img, None);
}
/// Run kernel `entry` writing `out` (dispatched over `groups` workgroups).
pub fn dispatch(
&mut self,
entry: &str,
p: &Params,
src: &GpuImage,
aux: Option<&GpuImage>,
out: &GpuImage,
data: Option<&wgpu::Buffer>,
groups: (u32, u32),
) {
self.dispatch_ext(entry, p, src, aux, out, data, groups, None);
}
/// [`Enc::dispatch`] for the kernels in `EXT_ENTRIES`, with their group 1 storage buffers.
#[allow(clippy::too_many_arguments)]
pub fn dispatch_ext(
&mut self,
entry: &str,
p: &Params,
src: &GpuImage,
aux: Option<&GpuImage>,
out: &GpuImage,
data: Option<&wgpu::Buffer>,
groups: (u32, u32),
ext: Option<[&wgpu::Buffer; 4]>,
) {
let g = self.g;
let Some(pipe) = g.pipelines.get(entry) else {
log::error!("gpu: no kernel {entry}");
return;
};
let ub = g.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: None, contents: &p.bytes(), usage: wgpu::BufferUsages::UNIFORM });
let sv = src.texture.create_view(&Default::default());
let av = aux.map(|a| a.texture.create_view(&Default::default()));
let ov = out.texture.create_view(&Default::default());
let bg = g.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &g.bgl,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: ub.as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&sv) },
wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::TextureView(av.as_ref().unwrap_or(&g.dummy_tex)) },
wgpu::BindGroupEntry { binding: 3, resource: wgpu::BindingResource::TextureView(&ov) },
wgpu::BindGroupEntry { binding: 4, resource: data.unwrap_or(&g.dummy_buf).as_entire_binding() },
],
});
let bg_ext = ext.map(|b| {
g.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &g.bgl_ext,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: b[0].as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: b[1].as_entire_binding() },
wgpu::BindGroupEntry { binding: 2, resource: b[2].as_entire_binding() },
wgpu::BindGroupEntry { binding: 3, resource: b[3].as_entire_binding() },
],
})
});
let enc = self.encoder();
{
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some(entry), timestamp_writes: None });
pass.set_pipeline(pipe);
pass.set_bind_group(0, &bg, &[]);
if let Some(b) = &bg_ext {
pass.set_bind_group(1, b, &[]);
}
pass.dispatch_workgroups(groups.0.max(1), groups.1.max(1), 1);
}
self.pending += 1;
// Keep command buffers (and the transient textures they hold) bounded.
if self.pending >= 256 {
self.submit();
}
}
/// Per-pixel kernel over `out`'s size (16×16 workgroups).
pub fn pixels(&mut self, entry: &str, p: &Params, src: &GpuImage, aux: Option<&GpuImage>, out: &GpuImage, data: Option<&wgpu::Buffer>) {
let groups = (out.width.div_ceil(16), out.height.div_ceil(16));
self.dispatch(entry, p, src, aux, out, data, groups);
}
/// Copy `src` into a larger zeroed image at (`x`, `y`) (Buf::pad).
pub fn copy_into(&mut self, src: &GpuImage, dst: &GpuImage, x: u32, y: u32) {
let enc = self.encoder();
enc.copy_texture_to_texture(
wgpu::TexelCopyTextureInfo { texture: &src.texture, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All },
wgpu::TexelCopyTextureInfo { texture: &dst.texture, mip_level: 0, origin: wgpu::Origin3d { x, y, z: 0 }, aspect: wgpu::TextureAspect::All },
wgpu::Extent3d { width: src.width, height: src.height, depth_or_array_layers: 1 },
);
}
/// Submit recorded work.
pub fn submit(&mut self) {
if let Some((enc, id)) = self.enc.take() {
self.g.queue.submit([enc.finish()]);
if let Ok(mut p) = self.g.pool.lock() {
p.close(id);
}
}
self.pending = 0;
}
/// Read an image back to the CPU (submits and waits). `None` where readback is impossible.
pub fn download(&mut self, img: &GpuImage) -> Option<Image> {
let mut out = Image::new(img.width, img.height);
let dst: &mut [u8] = bytemuck::cast_slice_mut(&mut out.data);
self.read_texture_into(&img.texture, img.width, img.height, 16, dst)?;
Some(out)
}
/// Tightly packed texel bytes of a texture (`bpp` bytes per pixel).
pub fn read_texture(&mut self, texture: &wgpu::Texture, w: u32, h: u32, bpp: u32) -> Option<Vec<u8>> {
let mut out = vec![0u8; (w * bpp) as usize * h as usize];
self.read_texture_into(texture, w, h, bpp, &mut out)?;
Some(out)
}
/// [`Enc::read_texture`] into `out` (`w × bpp × h` bytes), through a reused staging buffer.
fn read_texture_into(&mut self, texture: &wgpu::Texture, w: u32, h: u32, bpp: u32, out: &mut [u8]) -> Option<()> {
if !self.g.can_wait() {
return None;
}
self.g.transfers.count(false, out.len());
let row = (w * bpp) as usize;
let padded = (w * bpp).div_ceil(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT) * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
let size = padded as u64 * h as u64;
let buf = self.g.staging.lock().ok().and_then(|mut s| s.get_mut(&size).and_then(Vec::pop)).unwrap_or_else(|| {
self.g.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("readback"),
size,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
})
});
self.encoder().copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo { texture, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All },
wgpu::TexelCopyBufferInfo { buffer: &buf, layout: wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(padded), rows_per_image: Some(h) } },
wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 },
);
self.submit();
let (tx, rx) = std::sync::mpsc::channel();
buf.map_async(wgpu::MapMode::Read, .., move |r| {
let _ = tx.send(r);
});
if let Err(e) = self.g.device.poll(wgpu::PollType::wait_indefinitely()) {
log::error!("gpu: poll: {e}");
return None;
}
rx.recv().ok()?.ok()?;
{
let view = buf.get_mapped_range(..).ok()?;
if row == padded as usize {
out.copy_from_slice(&view[..out.len()]);
} else {
for (y, dst) in out.chunks_exact_mut(row).enumerate() {
let s = y * padded as usize;
dst.copy_from_slice(&view[s..s + row]);
}
}
}
buf.unmap();
// Keep a few staging buffers per size (frames read back at the same size every time).
if let Ok(mut s) = self.g.staging.lock() {
let v = s.entry(size).or_default();
if v.len() < 2 && size <= 256 << 20 {
v.push(buf);
}
}
Some(())
}
/// [`Enc::read_texture`] without waiting for the GPU (the browser's main thread can't): `done`
/// gets the bytes once the copy is mapped (from the browser's event loop on the web; from
/// a later device poll natively).
pub fn read_texture_async(&mut self, texture: &wgpu::Texture, w: u32, h: u32, bpp: u32, done: impl FnOnce(Option<Vec<u8>>) + wgpu::WasmNotSend + 'static) {
let row = w * bpp;
let padded = row.div_ceil(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT) * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
let buf = self.g.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("readback-async"),
size: padded as u64 * h as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.encoder().copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo { texture, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All },
wgpu::TexelCopyBufferInfo { buffer: &buf, layout: wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(padded), rows_per_image: Some(h) } },
wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 },
);
self.submit();
let b = buf.clone();
buf.map_async(wgpu::MapMode::Read, .., move |r| {
let out = r.ok().and_then(|_| {
let view = b.get_mapped_range(..).ok()?;
let mut out = Vec::with_capacity((row * h) as usize);
for y in 0..h as usize {
out.extend_from_slice(&view[y * padded as usize..y * padded as usize + row as usize]);
}
Some(out)
});
b.unmap();
done(out);
});
#[cfg(not(target_arch = "wasm32"))]
let _ = self.g.device.poll(wgpu::PollType::Poll);
}
/// A deferred readback of `img` (RGBA f32) under `key` (see [`crate::deferred`]): the bytes
/// once they arrived (`Some(None)`: the readback failed), else `None` after starting the
/// copy if needed and marking the pass missed.
pub(crate) fn read_keyed(&mut self, img: &GpuImage, key: u128, offset: [f64; 2], scale: f64) -> Option<Option<crate::deferred::Readback>> {
use crate::deferred::Lookup;
let d = self.g.deferred.clone()?;
match d.lookup(key) {
Lookup::Ready(r) => return Some(r),
Lookup::InFlight => {}
Lookup::Absent => {
self.g.transfers.count(false, img.width as usize * img.height as usize * 16);
let done = d.start(key, img.width, img.height, offset, scale);
self.read_texture_async(&img.texture, img.width, img.height, 16, done);
}
}
d.miss();
None
}
/// A data buffer for kernels that read tables (curves).
pub fn data(&self, v: &[f32]) -> wgpu::Buffer {
let mut bytes = Vec::with_capacity(v.len().max(4) * 4);
for x in v {
bytes.extend_from_slice(&x.to_le_bytes());
}
self.bytes(bytes)
}
/// Copy an image into a storage buffer (RGBA f32 rows) for kernels that read a third image
/// through `data`: (buffer, row length in pixels).
pub fn image_rows(&mut self, img: &GpuImage) -> (wgpu::Buffer, u32) {
let row = (img.width * 16).div_ceil(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT) * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
let buf = self.g.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("image rows"),
size: row as u64 * img.height as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.encoder().copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo { texture: &img.texture, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All },
wgpu::TexelCopyBufferInfo {
buffer: &buf,
layout: wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(row), rows_per_image: Some(img.height) },
},
wgpu::Extent3d { width: img.width, height: img.height, depth_or_array_layers: 1 },
);
(buf, row / 16)
}
/// A read-only storage buffer holding `bytes` (padded to 16 bytes).
pub fn bytes(&self, mut bytes: Vec<u8>) -> wgpu::Buffer {
while bytes.len() < 16 || !bytes.len().is_multiple_of(4) {
bytes.push(0);
}
self.g.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("data"), contents: &bytes, usage: wgpu::BufferUsages::STORAGE })
}
/// Convert to the viewer's premultiplied RGBA8 texture (egui-wgpu's native texture format).
pub fn display(&mut self, img: &GpuImage) -> wgpu::Texture {
let g = self.g;
let tex = g.device.create_texture(&wgpu::TextureDescriptor {
label: Some("viewer frame"),
size: wgpu::Extent3d { width: img.width, height: img.height, depth_or_array_layers: 1 },
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8Unorm,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
view_formats: &[],
});
let sv = img.texture.create_view(&Default::default());
let ov = tex.create_view(&Default::default());
let bg = g.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &g.display_bgl,
entries: &[
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&sv) },
wgpu::BindGroupEntry { binding: 3, resource: wgpu::BindingResource::TextureView(&ov) },
],
});
let enc = self.encoder();
{
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("display"), timestamp_writes: None });
pass.set_pipeline(&g.display);
pass.set_bind_group(0, &bg, &[]);
pass.dispatch_workgroups(img.width.div_ceil(16), img.height.div_ceil(16), 1);
}
tex
}
}