How to have seedable RNG in parallel in rust

Question

I am learning rust by implementing a raytracer. I have a working prototype that is single threaded and I am trying to make it multithreaded.

In my code, I have a sampler which is basically a wrapper around StdRng::seed_from_u64(123) (this will change when I will add different types of samplers) that is mutable because of StdRNG. I need to have a repeatable behaviour that is why i am seeding the random number generator.

In my rendering loop I use the sampler in the following way

        let mut sampler = create_sampler(&self.sampler_value);
        let sample_count = sampler.sample_count();

        println!("Rendering ...");
        let progress_bar = get_progress_bar(image.size());

        // Generate multiple rays for each pixel in the image
        for y in 0..image.size_y {
            for x in 0..image.size_x {
                image[(x, y)] = (0..sample_count)
                    .into_iter()
                    .map(|_| {
                        let pixel = Vec2::new(x as f32, y as f32) + sampler.next2f();
                        let ray = self.camera.generate_ray(&pixel);
                        self.integrator.li(self, &mut sampler, &ray)
                    })
                    .sum::<Vec3>()
                    / (sample_count as f32);

                progress_bar.inc(1);
            }
        }

When I replace into_iter by par_into_iter the compiler tells me cannot borrow sampler as mutable, as it is a captured variable in a Fn closure

What should I do in this situation?

P.s. If it is of any use, this is the repo : https://github.com/jgsimard/rustrt

Answer 1

The Rust Rand Book has a whole chapter dedicated to random generators and parallelism.

Use a single master seed. For each work unit, seed an RNG using the master seed and set the RNG's stream to the work unit number.

Note that this is not supported by all RNGs.

Here the example they gave that is fully deterministic:

use rand::distributions::{Distribution, Uniform};
use rand_chacha::{rand_core::SeedableRng, ChaCha8Rng};
use rayon::prelude::*;

static SEED: u64 = 0;
static BATCH_SIZE: u64 = 10_000;
static BATCHES: u64 = 1000;

fn main() {
    let range = Uniform::new(-1.0f64, 1.0);

    let in_circle = (0..BATCHES)
        .into_par_iter()
        .map(|i| {
            let mut rng = ChaCha8Rng::seed_from_u64(SEED);
            rng.set_stream(i);
            let mut count = 0;
            for _ in 0..BATCH_SIZE {
                let a = range.sample(&mut rng);
                let b = range.sample(&mut rng);
                if a * a + b * b <= 1.0 {
                    count += 1;
                }
            }
            count
        })
        .reduce(|| 0usize, |a, b| a + b);

    // prints 3.1409052 (determinstic and reproducible result)
    println!(
        "π is approximately {}",
        4. * (in_circle as f64) / ((BATCH_SIZE * BATCHES) as f64)
    );
}

Answer 2

This is the way I did it. I used this ressource : rust-random.github.io/book/guide-parallel.html. So I used ChaCha8Rng with the set_stream function to get seedable PRNG in parallel. I had to put the image[(x, y)] outside of the iterator because into_par_iter does not allow mutable borrow inside a closure. If you see something dumb in my solution, please tell me!

let size_x = image.size_x;
let img: Vec<Vec<Vec3>> = (0..image.size_y)
    .into_par_iter()
    .map(|y| {
         (0..image.size_x)
             .into_par_iter()
             .map(|x| {
                 let mut rng = ChaCha8Rng::seed_from_u64(sampler.seed());
                 rng.set_stream((y * size_x + x) as u64);
                 let v = (0..sample_count)
                     .into_iter()
                     .map(|_| {
                         let pixel = Vec2::new(x as f32, y as f32) + sampler.next2f(&mut rng);
                         let ray = self.camera.generate_ray(&pixel);
                         self.integrator.li(self, &sampler, &mut rng, &ray)
                     })
                     .sum::<Vec3>()
                     / (sample_count as f32);
                  progress_bar.inc(1);
                  v
              }).collect()
      }).collect();

for (y, row) in img.into_iter().enumerate() {
    for (x, p) in row.into_iter().enumerate() {
         image[(x, y)] = p;
    }
}

Answer 3

Even if Rust wasn't stopping you, you cannot just use a seeded PRNG with parallelism and get a reproducible result out.

Think about it this way: a PRNG with a certain seed/state produces a certain sequence of numbers. Reproducibility (determinism) requires not just that the numbers are the same, but that the way they are taken from the sequence is the same. But if you have multiple threads computing different pixels (different uses) which are racing with each other to fetch numbers from the single PRNG, then the pixels will fetch different numbers on different runs.

In order to get the determinism you want, you must deterministically choose which random number is used for which purpose.

One way to do this would be to make up an “image” of random numbers, computed sequentially, and pass that to the parallel loop. Then each ray has its own random number, which it can use as its seed for another PRNG that only that ray uses.

Another way that can be much more efficient and usable (because it doesn't require any sequentiality at all) is to use hash functions instead of PRNGs. Whenever you want a random number, use a hash function (like those which implement the std::hash::Hasher trait in Rust, but not necessarily the particular one std provides since it's not the fastest) to combine a bunch of information, like

• the seed value
• the pixel x and y location
• which bounce or secondary ray of this pixel you're computing

into a single value which you can use as a pseudorandom number. This way, the “random” results are the same for the same circumstances (because you explicitly specfied that it should be computed from them) even if some other part of the program execution changes (whether that's a code change or a thread scheduling decision by the OS).

Answer 4

Your sampler is not thread-safe, if only because it is a &mut Sampler and mutable references cannot be shared between threads, obviously.

The easy thing would be to wrap it into an Arc<Mutex<Sampler>> and clone it to every closure. Something like (untested):

let sampler = Arc::new(Mutex::new(create_sampler(&self.sampler_value)));
//...
for y in 0..image.size_y {
            for x in 0..image.size_x {
                image[(x, y)] = (0..sample_count)
                    .par_into_iter()
                    .map({
                       let sampler = Arc::clone(sampler);
                       move |_| {
                           let mut sampler = sampler.lock().unwrap();
                           // use the sampler
                       }
                    })
                    .sum::<Vec3>() //...

But that may not be very efficient way because the mutex will be locked most of the time, and you will kill the paralellism. You may try locking/unlocking the mutex during the ray tracing and see if it improves.

The ideal solution would be to make the Sampler thread-safe and inner mutable, so that the next2f and friends do not need the &mut self part (Sampler::next2f(&self)). Again, the easiest way is having an internal mutex.

Or you can try going lock-less! I mean, your current implementation of that function is:

    fn next2f(&mut self) -> Vec2 {
        self.current_dimension += 2;
        Vec2::new(self.rng.gen(), self.rng.gen())
    }

You could replace the current_dimension with an AtomicI32 and the rng with a rand::thread_rng (also untested):

    fn next2f(&self) -> Vec2 {
        self.current_dimension.fetch_add(2, Ordering::SeqCst);
        let mut rng = rand::thread_rng();
        Vec2::new(rng.gen(), rng.gen())
    }