Merge pull request #1511 from MatejKozlovsky/fallback-parallel

Export: render in parallel once the GPU stops for good

crates/export/src/pipeline.rs

@@ -48,22 +48,43 @@ pub(crate) fn gpu_inflight(var: Option<&str>) -> usize {
 /// one window of rendered images is held in memory. Chunks hold no lock or slot, so a failing
 /// or panicking frame cannot leak one, and nested parallelism (or a single thread) cannot
 /// deadlock.
+#[cfg(test)]
 pub(crate) fn map_frames<I: Send, T: Send>(
     ks: Vec<u64>,
     cap: Option<usize>,
     render: impl Fn(u64) -> I + Sync + Send,
     encode: impl Fn(u64, I) -> T + Sync + Send,
+) -> Vec<T> {
+    map_frames_while(ks, cap, || true, render, encode)
+}
+
+/// [`map_frames`] whose `cap` holds only while `capped()`: once it turns false (the GPU stopped
+/// for good, so every frame falls back to the CPU), the remaining frames render in parallel
+/// like the uncapped CPU path instead of one chunk of `cap` at a time.
+pub(crate) fn map_frames_while<I: Send, T: Send>(
+    ks: Vec<u64>,
+    cap: Option<usize>,
+    capped: impl Fn() -> bool,
+    render: impl Fn(u64) -> I + Sync + Send,
+    encode: impl Fn(u64, I) -> T + Sync + Send,
 ) -> Vec<T> {
     use rayon::prelude::*;
-    let Some(n) = cap.map(|n| n.max(1)) else {
+    let Some(n) = cap.map(|n| n.max(1)).filter(|_| capped()) else {
         return ks.into_par_iter().map(|k| encode(k, render(k))).collect();
     };
     let window = n.max(rayon::current_num_threads());
     let mut out = Vec::with_capacity(ks.len());
-    for win in ks.chunks(window) {
+    for (w, win) in ks.chunks(window).enumerate() {
+        if !capped() {
+            out.par_extend(ks[w * window..].par_iter().map(|&k| encode(k, render(k))));
+            break;
+        }
         let mut imgs = Vec::with_capacity(win.len());
-        for chunk in win.chunks(n) {
-            imgs.par_extend(chunk.par_iter().map(|&k| render(k)));
+        let mut i = 0;
+        while i < win.len() {
+            let end = if capped() { (i + n).min(win.len()) } else { win.len() };
+            imgs.par_extend(win[i..end].par_iter().map(|&k| render(k)));
+            i = end;
         }
         out.par_extend(win.par_iter().copied().zip(imgs.into_par_iter()).map(|(k, img)| encode(k, img)));
     }
@@ -187,7 +208,10 @@ impl<'a> Cx<'a> {
             }
             return out;
         }
-        map_frames(ks, self.gpu_cap(), |k| self.frame(comp, k), f)
+        // The cap is for the GPU's memory: once it has stopped for good, frames render on the
+        // CPU in parallel.
+        let gpu_live = || self.job.accel.is_none_or(|a| a.available());
+        map_frames_while(ks, self.gpu_cap(), gpu_live, |k| self.frame(comp, k), f)
     }
 
     /// GPU frames rendered at once (#879), decided once per batch: with a GPU accelerator every
@@ -453,6 +477,36 @@ mod tests {
         (out, max.load(Ordering::SeqCst))
     }
 
+    #[test]
+    fn frames_render_in_parallel_once_the_gpu_stops() {
+        // Every frame after the GPU stopped falls back to the CPU: the cap no longer applies.
+        use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+        let pool = rayon::ThreadPoolBuilder::new().num_threads(4).build().unwrap();
+        let (live, now, max) = (AtomicBool::new(true), AtomicUsize::new(0), AtomicUsize::new(0));
+        let out = pool.install(|| {
+            map_frames_while(
+                (0..24).collect(),
+                Some(1),
+                || live.load(Ordering::SeqCst),
+                |k| {
+                    if k == 1 {
+                        live.store(false, Ordering::SeqCst);
+                    }
+                    let c = now.fetch_add(1, Ordering::SeqCst) + 1;
+                    if k >= 2 {
+                        max.fetch_max(c, Ordering::SeqCst);
+                    }
+                    std::thread::sleep(std::time::Duration::from_millis(20));
+                    now.fetch_sub(1, Ordering::SeqCst);
+                    k * 10
+                },
+                |_, v| v,
+            )
+        });
+        assert_eq!(out, (0..24).map(|k| k * 10).collect::<Vec<_>>());
+        assert!(max.load(Ordering::SeqCst) > 1, "frames after the GPU stopped ran one at a time");
+    }
+
     #[test]
     fn capped_frames_keep_their_order() {
         let want: Vec<u64> = (0..40).map(|k| k * 10).collect();

crates/gpu/src/lib.rs

@@ -317,6 +317,10 @@ impl Accelerator for Gpu {
         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() {

crates/render/src/lib.rs

@@ -455,6 +455,12 @@ pub trait Accelerator: Send + Sync {
     fn auto_pick(&self) -> Option<&AutoPick> {
         None
     }
+    /// Whether the accelerator can still render frames. `false` once it has stopped for good
+    /// (its device was lost or retired): every frame it is asked for falls back to the CPU, so
+    /// callers that limit frames in flight for it can stop doing so.
+    fn available(&self) -> bool {
+        true
+    }
     /// Deferred readbacks: where a GPU readback cannot be waited for (WebGPU in a browser
     /// worker), a frame renders in *passes*. A readback the accelerator does not have yet is
     /// started, the step returns a cheap placeholder and the pass is marked missed