Parallelize across subsample frames and dump intermediates.

rtiow/src/bin/tracer.rs

@@ -173,7 +173,7 @@ pub struct Opt {
     #[structopt(short = "s", long = "subsample", default_value = "10")]
     pub subsamples: usize,
 
-    /// Output file
+    /// Output directory
     #[structopt(parse(from_os_str))]
     pub output: PathBuf,
 }
@@ -182,7 +182,7 @@ fn main() -> Result<(), std::io::Error> {
     let start = Instant::now();
     let opt = Opt::from_args();
     let scene = build_scene(&opt);
-    let img = render(&scene);
+    let res = render(scene, &opt.output);
     let runtime = start.elapsed();
     eprintln!(
         "Render time {}.{} seconds",
@@ -190,9 +190,5 @@ fn main() -> Result<(), std::io::Error> {
         runtime.subsec_millis()
     );
 
-    let path = "/tmp/test.png";
-    // Write the contents of this image to the Writer in PNG format.
-    img.save(path).unwrap();
-    eprintln!("Saved {}", path);
-    Ok(())
+    res
 }

rtiow/src/lib.rs

@@ -7,6 +7,8 @@ pub mod renderer;
 pub mod sphere;
 pub mod vec3;
 
+extern crate crossbeam_channel;
 extern crate image;
+extern crate num_cpus;
 extern crate rand;
 extern crate rayon;

rtiow/src/renderer.rs

@@ -1,10 +1,14 @@
 use std;
+use std::path::Path;
+use std::sync;
+use std::thread;
 
+use crossbeam_channel as channel;
 use image;
 use image::RgbImage;
+use num_cpus;
 use rand;
 use rand::Rng;
-use rayon::prelude::*;
 
 use camera::Camera;
 use hitable::Hit;
@@ -29,49 +33,93 @@ fn color(r: Ray, world: &Hit, depth: usize) -> Vec3 {
         }
         return Default::default();
     }
-
     // No hit, choose color from background.
     let unit_direction = r.direction().unit_vector();
     let t = 0.5 * (unit_direction.y + 1.);
     Vec3::new(1., 1., 1.) * (1. - t) + Vec3::new(0.5, 0.7, 1.) * t
 }
 
-fn trace_pixel(x: usize, y: usize, scene: &Scene) -> [u8; 3] {
+fn trace_pixel(x: usize, y: usize, scene: &Scene) -> Vec3 {
     let mut rng = rand::thread_rng();
-    let mut col: Vec3 = Default::default();
-    for _ in 0..scene.subsamples {
-        let u = (rng.gen_range::<f32>(0., 1.) + x as f32) / scene.width as f32;
-        let v = (rng.gen_range::<f32>(0., 1.) + y as f32) / scene.height as f32;
-        let ray = scene.camera.get_ray(u, v);
-        col = col + color(ray, &scene.world, 0);
+    let u = (rng.gen_range::<f32>(0., 1.) + x as f32) / scene.width as f32;
+    let v = (rng.gen_range::<f32>(0., 1.) + y as f32) / scene.height as f32;
+    let ray = scene.camera.get_ray(u, v);
+    color(ray, &scene.world, 0)
+}
+
+fn render_worker(
+    scene: &Scene,
+    input_chan: channel::Receiver<usize>,
+    output_chan: channel::Sender<(usize, Vec<Vec3>)>,
+) {
+    for subsample in input_chan {
+        let mut pixel_data: Vec<Vec3> = Vec::with_capacity(scene.width * scene.height);
+        for y in 0..scene.height {
+            for x in 0..scene.width {
+                let p = trace_pixel(x, y, scene);
+                pixel_data.push(p);
+            }
+        }
+        output_chan.send((subsample, pixel_data));
     }
-    col = col / scene.subsamples as f32;
-    // Gamma correct, use gamma 2 correction, which is 1/gamma where gamma=2 which is 1/2
-    // or sqrt.
-    col = Vec3::new(col[0].sqrt(), col[1].sqrt(), col[2].sqrt());
-    let ir = (255.99 * col[0]) as u8;
-    let ig = (255.99 * col[1]) as u8;
-    let ib = (255.99 * col[2]) as u8;
-    [ir, ig, ib]
+    eprintln!("Shutting down worker");
 }
 
-pub fn render(scene: &Scene) -> image::RgbImage {
+pub fn render(scene: Scene, output_dir: &Path) -> std::result::Result<(), std::io::Error> {
+    let (seq_tx, seq_rx) = channel::unbounded();
+    let (pixel_data_tx, pixel_data_rx) = channel::unbounded();
+
+    let scene = sync::Arc::new(scene);
+    for _ in 0..num_cpus::get() {
+        let s = sync::Arc::clone(&scene);
+        let seq_rx = seq_rx.clone();
+        let pixel_data_tx = pixel_data_tx.clone();
+        thread::spawn(move || {
+            render_worker(&s, seq_rx, pixel_data_tx);
+        });
+    }
+    drop(seq_rx);
+    drop(pixel_data_tx);
+
+    (1..=scene.subsamples).for_each(|idx| seq_tx.send(idx));
+    drop(seq_tx);
+
+    let mut acc_count = 0;
+    let mut acc: Vec<Vec3> = Vec::with_capacity(scene.width * scene.height);
+    for _ in 0..(scene.width * scene.height) {
+        acc.push(Default::default());
+    }
+
     let mut img = RgbImage::new(scene.width as u32, scene.height as u32);
-    let coords: Vec<_> = (0..scene.height)
-        .flat_map(|j| (0..scene.width).map(move |i| (i, j)))
-        .collect();
+    for (_subsample, pixel_data) in pixel_data_rx {
+        acc_count += 1;
+        pixel_data.iter().enumerate().for_each(|(idx, p)| {
+            let x = idx % scene.width;
+            let y = idx / scene.width;
+            let y_inv = scene.height - y - 1;
+            acc[idx] = acc[idx] + *p;
 
-    let pixels = coords
-        .par_iter()
-        .map(|(i, j)| {
-            let p = trace_pixel(*i, *j, &scene);
-            // height-j is to flip y-axis.
-            (*i as u32, (scene.height - *j - 1) as u32, image::Rgb(p))
-        })
-        .collect::<Vec<_>>();
-
-    pixels
-        .iter()
-        .for_each(|(x, y, p)| img.put_pixel(*x, *y, *p));
-    img
+            // Gamma correct, use gamma 2 correction, which is 1/gamma where gamma=2 which is 1/2 or
+            // sqrt.
+            let col = acc[idx] / acc_count as f32;
+            let col = Vec3::new(col[0].sqrt(), col[1].sqrt(), col[2].sqrt());
+            img.put_pixel(
+                x as u32,
+                y_inv as u32,
+                image::Rgb([
+                    (col[0] * 255.).min(255.) as u8,
+                    (col[1] * 255.).min(255.) as u8,
+                    (col[2] * 255.).min(255.) as u8,
+                ]),
+            );
+        });
+        let path = output_dir.join(format!("iteration{:05}.png", acc_count));
+        eprintln!("saving {}", path.to_string_lossy());
+        img.save(&path)
+            .unwrap_or_else(|_| panic!("Failed save {}", path.to_string_lossy()));
+    }
+    let path = output_dir.join("final.png");
+    // Write the contents of this image to the Writer in PNG format.
+    eprintln!("Saving {}", path.to_string_lossy());
+    img.save(path)
 }