GPU: resident effect stacks, in-frame texture reuse, upload-cache budget, stable footprint keys (#687)

crates/gpu/src/context.rs

@@ -37,7 +37,7 @@ const EXT_ENTRIES: &[&str] = &["classic3d"];
 /// 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).
+/// Uploaded layer buffers kept on the GPU by default (bytes; [`GpuContext::set_upload_budget`]).
 const UPLOAD_BUDGET: usize = 1536 << 20;
 
 /// Working textures kept for reuse once their last image is dropped (bytes).
@@ -71,6 +71,8 @@ impl GpuImage {
 struct Lease {
     texture: wgpu::Texture,
     pool: Weak<Mutex<Pool>>,
+    /// The [`Enc`] that took it (see [`Pool::owned`]).
+    owner: u64,
 }
 
 impl std::fmt::Debug for Lease {
@@ -84,7 +86,7 @@ impl Drop for Lease {
         if let Some(pool) = self.pool.upgrade()
             && let Ok(mut p) = pool.lock()
         {
-            p.release(self.texture.clone());
+            p.release_owned(self.texture.clone(), self.owner);
         }
     }
 }
@@ -108,8 +110,18 @@ pub(crate) struct Pool {
     /// Working textures reused / newly allocated (diagnostics, tests).
     hits: u64,
     misses: u64,
+    /// Textures released while the [`Enc`] that took them is alive, kept for that `Enc` only.
+    /// Its later commands run after everything it recorded or submitted before (one queue,
+    /// recording order), so it can overwrite them at once instead of allocating a new texture
+    /// for every intermediate of a frame (a 4K frame holds dozens). Other encoders never see a
+    /// working texture of another `Enc` before it is gone.
+    owned: HashMap<u64, HashMap<(u32, u32), Vec<wgpu::Texture>>>,
+    next_owner: u64,
 }
 
+/// In-`Enc` texture reuse ([`Pool::owned`]); tests turn it off to compare.
+pub(crate) static TEXTURE_REUSE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(true);
+
 /// Counts of uploads (CPU → GPU) and readbacks (GPU → CPU).
 #[derive(Default)]
 pub(crate) struct Transfers {
@@ -178,6 +190,37 @@ impl Pool {
         }
     }
 
+    fn register(&mut self) -> u64 {
+        self.next_owner += 1;
+        if TEXTURE_REUSE.load(std::sync::atomic::Ordering::Relaxed) {
+            self.owned.insert(self.next_owner, HashMap::new());
+        }
+        self.next_owner
+    }
+
+    /// `Enc` `owner` is gone: its kept textures go through the shared pool.
+    fn unregister(&mut self, owner: u64) {
+        if let Some(m) = self.owned.remove(&owner) {
+            for t in m.into_values().flatten() {
+                self.release(t);
+            }
+        }
+    }
+
+    fn release_owned(&mut self, t: wgpu::Texture, owner: u64) {
+        if let Some(m) = self.owned.get_mut(&owner) {
+            m.entry((t.width(), t.height())).or_default().push(t);
+            return;
+        }
+        self.release(t);
+    }
+
+    fn take_owned(&mut self, owner: u64, w: u32, h: u32) -> Option<wgpu::Texture> {
+        let t = self.owned.get_mut(&owner)?.get_mut(&(w, h))?.pop()?;
+        self.hits += 1;
+        Some(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 {
@@ -253,6 +296,8 @@ pub struct GpuContext {
     dummy_tex: wgpu::TextureView,
     dummy_buf: wgpu::Buffer,
     uploads: Mutex<Uploads>,
+    upload_budget: std::sync::atomic::AtomicUsize,
+    upload_full_warned: std::sync::atomic::AtomicBool,
     pool: Arc<Mutex<Pool>>,
     /// CPU ↔ GPU traffic (diagnostics: `bench --gpu`).
     pub(crate) transfers: Transfers,
@@ -598,6 +643,8 @@ impl GpuContext {
             dummy,
             dummy_buf,
             uploads: Mutex::new(Uploads::default()),
+            upload_budget: std::sync::atomic::AtomicUsize::new(UPLOAD_BUDGET),
+            upload_full_warned: Default::default(),
             pool: Default::default(),
             staging: Default::default(),
             transfers: Default::default(),
@@ -817,15 +864,26 @@ impl GpuContext {
 
     /// 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) {
+    fn pooled(&self, w: u32, h: u32, owner: u64) -> (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 reused = self.pool.lock().ok().and_then(|mut p| p.take_owned(owner, w, h).or_else(|| 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) });
+        let lease = Arc::new(Lease { texture: texture.clone(), pool: Arc::downgrade(&self.pool), owner });
         (GpuImage { texture, width: w, height: h, _lease: Some(lease), levels: None }, was_reused)
     }
 
+    /// Bytes of uploaded layer textures the context keeps (default 1536 MB). An application that
+    /// knows the device's memory can raise it (static layers of a big frame upload once) or lower
+    /// it (a small GPU); the least recently used textures go first.
+    pub fn set_upload_budget(&self, bytes: usize) {
+        self.upload_budget.store(bytes.max(1 << 20), std::sync::atomic::Ordering::Relaxed);
+    }
+
+    pub fn upload_budget(&self) -> usize {
+        self.upload_budget.load(std::sync::atomic::Ordering::Relaxed)
+    }
+
     /// Uploads and readbacks so far (images; uploads of cached layers count once).
     pub fn transfer_stats(&self) -> TransferStats {
         use std::sync::atomic::Ordering::Relaxed;
@@ -882,15 +940,23 @@ impl GpuContext {
             }
             u.bytes += bytes;
             // Drop textures of freed buffers, then the least recently used over budget.
-            if u.bytes > UPLOAD_BUDGET {
+            let budget = self.upload_budget();
+            if u.bytes > 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 {
+            if u.bytes > budget && !self.upload_full_warned.swap(true, std::sync::atomic::Ordering::Relaxed) {
+                log::warn!(
+                    target: "effectcraft_gpu",
+                    "GPU upload cache full ({} MB): evicting layer textures, which upload again when next used; a frame whose static layers exceed it re-uploads them every frame",
+                    budget >> 20
+                );
+            }
+            while u.bytes > 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;
@@ -923,22 +989,26 @@ pub(crate) struct Enc<'g> {
     pending: usize,
     /// A dispatch named a kernel the device doesn't have (see [`Enc::take_missing`]).
     missing: bool,
+    /// This `Enc`'s id in the texture pool (textures it may reuse at once).
+    owner: u64,
 }
 
 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);
+        if let Ok(mut p) = self.g.pool.lock() {
+            if let Some((_, id)) = self.enc.take() {
+                p.close(id);
+            }
+            p.unregister(self.owner);
         }
     }
 }
 
 impl<'g> Enc<'g> {
     pub fn new(g: &'g GpuContext) -> Enc<'g> {
-        Enc { g, enc: None, pending: 0, missing: false }
+        let owner = g.pool.lock().map(|mut p| p.register()).unwrap_or(0);
+        Enc { g, enc: None, pending: 0, missing: false, owner }
     }
 
     /// Whether a dispatch since the last call named a kernel the device doesn't have (one
@@ -960,7 +1030,7 @@ impl<'g> Enc<'g> {
 
     /// 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);
+        let (img, reused) = self.g.pooled(w, h, self.owner);
         if reused {
             self.fill(&img, 0.0);
         }
@@ -970,7 +1040,7 @@ impl<'g> Enc<'g> {
     /// 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);
+        let (img, _) = self.g.pooled(w, h, self.owner);
         if cfg!(test) {
             self.fill(&img, f32::NAN);
         }

crates/gpu/src/tests.rs

@@ -857,3 +857,190 @@ fn readback_staging_stays_bounded_across_sizes() {
     let held = g.ctx.staging_bytes();
     assert!(held <= (256 << 20) + 2 * 1024 * 16 * 376, "staging buffers hold {held} bytes");
 }
+
+/// A static adjustment layer with a time-dependent effect (Noise) keeps its footprint key
+/// across frames (the footprint is the source through the masks, not the effects), so the GPU
+/// uploads nothing after the first frame; matches the CPU.
+#[test]
+fn adjustment_footprint_uploads_once_under_time_dependent_effects() {
+    let mut s = Scene::new(BitDepth::Bpc8);
+    let bg = s.solid([0.3, 0.5, 0.7], 96, 64);
+    s.push(bg);
+    let mut adj = s.solid([1.0, 1.0, 1.0], 80, 50);
+    adj.switches.adjustment = true;
+    s.effect(&mut adj, "ec.noise.noise", &[("amount", n(30.0))]);
+    s.push(adj);
+    // The key itself (no GPU needed): stable across time.
+    let comp = s.p.comp(s.cid).expect("the scene's comp");
+    let adj = comp.layers.last().expect("the adjustment layer");
+    let key = |k: f64| {
+        effectcraft_render::cache::footprint_key(&effectcraft_render::EvalCtx::new(&s.p, s.cid, comp, Tick::from_seconds_f64(k)), adj, 1.0, false, false)
+    };
+    assert!(key(0.0).is_some());
+    assert_eq!(key(0.0), key(0.5), "the footprint key changes with time");
+    let Some(g) = gpu() else { return };
+    let cache = effectcraft_render::cache::LayerCache::new(64 << 20);
+    let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
+    r.accel = Some(g);
+    r.cache = Some(&cache);
+    let t = |k: u64| Tick::from_seconds_f64(k as f64 / 30.0);
+    let mut prev = None;
+    for k in 0..4 {
+        let before = g.ctx.transfer_stats().upload_bytes;
+        let img = g.render(&r, s.cid, t(k)).expect("the GPU renders the frame");
+        let up = g.ctx.transfer_stats().upload_bytes - before;
+        if k >= 2 {
+            assert_eq!(up, 0, "frame {k} uploaded {up} bytes: the footprint was re-keyed");
+        }
+        // Noise changes every frame: the frames differ.
+        assert!(prev.as_ref() != Some(&img.data), "frame {k} repeats the previous one");
+        prev = Some(img.data);
+        check(&format!("noise adjustment frame {k}"), compare_at(&s, opts(), t(k)), 0.0);
+    }
+}
+
+/// An `Enc` reuses a working texture it released at once (`Pool::owned`); another `Enc` does
+/// not get it while the first is alive.
+#[test]
+fn enc_reuses_its_own_released_textures() {
+    let Some(g) = gpu() else { return };
+    let mut e = crate::context::Enc::new(&g.ctx);
+    let a = e.image(37, 23);
+    let ta = a.texture.clone();
+    drop(a);
+    let mut other = crate::context::Enc::new(&g.ctx);
+    let o = other.image(37, 23);
+    assert!(o.texture != ta, "another Enc got a texture the first may still be using");
+    let b = e.image(37, 23);
+    assert!(b.texture == ta, "the Enc allocated a new texture instead of reusing its own");
+}
+
+/// In-`Enc` texture reuse changes no pixel: a busy frame (layers, effects, adjustment layer,
+/// nested comp) renders bit-identically with reuse on and off, and matches the CPU.
+#[test]
+fn texture_reuse_does_not_change_output() {
+    let Some(g) = gpu() else { return };
+    let mut s = Scene::new(BitDepth::Bpc32);
+    for i in 0..4 {
+        let mut l = s.footage(90 - i * 7, 60 - i * 5);
+        s.effect(&mut l, "ec.blur.gaussian", &[("blurriness", n(2.0 + i as f64))]);
+        set(&mut l, "transform/rotation", Value::Scalar(10.0 * i as f64));
+        l.blend_mode = if i % 2 == 0 { BlendMode::Screen } else { BlendMode::Multiply };
+        s.push(l);
+    }
+    let mut adj = s.solid([1.0, 1.0, 1.0], 80, 50);
+    adj.switches.adjustment = true;
+    s.effect(&mut adj, "ec.color.levels", &[("gamma", n(1.2))]);
+    s.push(adj);
+    let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
+    r.accel = Some(g);
+    let t = Tick::from_seconds_f64(0.3);
+    let on = g.render(&r, s.cid, t).expect("the GPU renders the frame");
+    crate::context::TEXTURE_REUSE.store(false, std::sync::atomic::Ordering::Relaxed);
+    let off = g.render(&r, s.cid, t);
+    crate::context::TEXTURE_REUSE.store(true, std::sync::atomic::Ordering::Relaxed);
+    assert!(off.is_some_and(|off| off.data == on.data), "texture reuse changed the frame");
+    check("texture reuse vs CPU", compare_at(&s, opts(), t), 0.0);
+}
+
+/// Static layer buffers (here adjustment-layer footprints) upload once while they fit the upload budget, and upload again every frame
+/// when they do not (least recently used first); the frame is the same either way.
+#[test]
+fn upload_budget_bounds_the_layer_textures_kept() {
+    let Some(g) = gpu() else { return };
+    let mut s = Scene::new(BitDepth::Bpc8);
+    for (i, c) in [[0.8, 0.2, 0.1], [0.1, 0.3, 0.9]].into_iter().enumerate() {
+        let mut l = s.solid(c, 200, 200);
+        set(&mut l, "transform/position", v3(60.0 + 30.0 * i as f64, 40.0));
+        // Adjustment layers: their footprints are cached layer buffers the GPU uploads.
+        l.switches.adjustment = true;
+        s.effect(&mut l, "ec.color.levels", &[("gamma", n(1.1 + 0.2 * i as f64))]);
+        s.push(l);
+    }
+    let frame = |cache: &effectcraft_render::cache::LayerCache, k: u64| {
+        let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
+        r.accel = Some(g);
+        r.cache = Some(cache);
+        let before = g.ctx.transfer_stats().upload_bytes;
+        let img = g.render(&r, s.cid, Tick::from_seconds_f64(k as f64 / 30.0)).expect("the GPU renders the frame");
+        (img.data, g.ctx.transfer_stats().upload_bytes - before)
+    };
+    let fits = effectcraft_render::cache::LayerCache::new(64 << 20);
+    let (a, _) = frame(&fits, 0);
+    let (b, up) = frame(&fits, 1);
+    assert_eq!(up, 0, "static layers re-uploaded within the default budget");
+    // A budget smaller than the two 640 KB footprints: each upload evicts the other.
+    let default = g.ctx.upload_budget();
+    g.ctx.set_upload_budget(1 << 20);
+    let tight = effectcraft_render::cache::LayerCache::new(64 << 20);
+    let _ = frame(&tight, 0);
+    let (c, up) = frame(&tight, 1);
+    g.ctx.set_upload_budget(default);
+    assert!(up > 0, "over budget, the footprints must upload again");
+    assert!(a == b && b == c, "the upload budget changed the frame");
+}
+
+/// Layers with effect stacks run them resident on the GPU (`walk::draw_resident_fx`): GPU
+/// effects, a CPU-only effect in the middle (read back, run, upload), opacity and a blend mode,
+/// a shape layer whose effects run in effect space from its bounds' corner; 8 and 32 bpc, full
+/// and half resolution; all match the CPU.
+#[test]
+fn effect_stacks_render_resident() {
+    for depth in [BitDepth::Bpc8, BitDepth::Bpc32] {
+        let mut s = Scene::new(depth);
+        let bg = s.footage(97, 61);
+        s.push(bg);
+        let mut a = s.footage(70, 44);
+        s.effect(&mut a, "ec.blur.gaussian", &[("blurriness", n(4.0))]);
+        s.effect(&mut a, "ec.generate.fractal", &[]);
+        s.effect(&mut a, "ec.color.levels", &[("gamma", n(1.4))]);
+        set(&mut a, "transform/rotation", Value::Scalar(20.0));
+        set(&mut a, "transform/opacity", Value::Scalar(70.0));
+        a.blend_mode = BlendMode::Overlay;
+        s.push(a);
+        let mut l = build::layer(&mut s.p, &s.comp, "Shape", LayerSource::Shape, (97, 61), None);
+        let mut next = s.p.next_id;
+        let mut ids = Ids(&mut next);
+        let rect = build::shape_rect(&mut ids, [50.0, 30.0], [5.0, -4.0], 0.0);
+        let fill = build::shape_fill(&mut ids, [0.2, 0.8, 0.3, 1.0]);
+        let g = build::shape_group(&mut ids, "Rectangle 1", vec![rect, fill]);
+        s.p.next_id = next;
+        l.props.sub_mut("contents").expect("a shape layer has contents").children.push(g.into());
+        s.effect(&mut l, "ec.blur.gaussian", &[("blurriness", n(3.0))]);
+        s.effect(&mut l, "ec.stylize.mosaic", &[]);
+        s.push(l);
+        let t = Tick::from_seconds_f64(0.2);
+        check(&format!("effect stacks {depth:?}"), compare_at(&s, opts(), t), 0.0);
+        check(&format!("effect stacks half {depth:?}"), compare_at(&s, RenderOpts { scale: 0.5, ..opts() }, t), 0.0);
+    }
+}
+
+/// A layer whose effects give the same result as last frame draws from the layer cache (its
+/// processed buffer uploads once) instead of re-running the stack on the GPU every frame.
+#[test]
+fn static_effect_layers_draw_from_the_cache() {
+    let Some(g) = gpu() else { return };
+    let mut s = Scene::new(BitDepth::Bpc8);
+    let mut l = s.solid([0.6, 0.4, 0.2], 80, 50);
+    s.effect(&mut l, "ec.blur.gaussian", &[("blurriness", n(3.0))]);
+    s.push(l);
+    let cache = effectcraft_render::cache::LayerCache::new(64 << 20);
+    let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
+    r.accel = Some(g);
+    r.cache = Some(&cache);
+    let comp = s.p.comp(s.cid).expect("the scene's comp");
+    let layer = comp.layers.last().expect("the layer");
+    let t = |k: u64| Tick::from_seconds_f64(k as f64 / 30.0);
+    assert!(r.layer_is_static(&effectcraft_render::EvalCtx::new(&s.p, s.cid, comp, t(2)), layer));
+    let mut frames = Vec::new();
+    for k in 0..3 {
+        let before = g.ctx.transfer_stats().upload_bytes;
+        frames.push(g.render(&r, s.cid, t(k)).expect("the GPU renders the frame").data);
+        let up = g.ctx.transfer_stats().upload_bytes - before;
+        if k == 2 {
+            assert_eq!(up, 0, "the cached layer uploaded again");
+        }
+    }
+    assert!(frames[1] == frames[2]);
+    check("static effect layer vs CPU", compare_at(&s, opts(), t(2)), 0.0);
+}