GPU: an Enc reuses the working textures it released at once instead of allocating new ones for every intermediate of a frame (#687)

crates/gpu/src/context.rs

@@ -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 {
@@ -817,12 +860,12 @@ 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)
     }
 
@@ -923,22 +966,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 +1007,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 +1017,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

@@ -898,3 +898,47 @@ fn adjustment_footprint_uploads_once_under_time_dependent_effects() {
         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);
+}