pbrt/src/utils/parallel.rs

74 lines
1.5 KiB
Rust

use crossbeam_channel::{Receiver, bounded};
use rayon::prelude::*;
pub fn init_parallel(n_threads: usize) {
let threads = if n_threads == 0 {
num_cpus::get()
} else {
n_threads
};
if let Err(e) = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build_global()
{
eprintln!("Warning: Rayon thread pool already initialized: {}", e);
}
}
pub fn num_system_cores() -> usize {
num_cpus::get()
}
pub fn max_concurrency() -> usize {
rayon::current_num_threads()
}
pub fn parallel_for<F>(start: i64, end: i64, func: F)
where
F: Fn(i64) + Sync + Send,
{
(start..end).into_par_iter().for_each(|i| func(i));
}
pub fn parallel_for_2d<F>(start_x: i64, end_x: i64, start_y: i64, end_y: i64, func: F)
where
F: Fn(i64, i64) + Sync + Send,
{
(start_y..end_y).into_par_iter().for_each(|y| {
(start_x..end_x).into_par_iter().for_each(|x| {
func(x, y);
});
});
}
pub struct AsyncJob<T> {
receiver: Receiver<T>,
}
impl<T> AsyncJob<T> {
pub fn wait(self) -> T {
self.receiver
.recv()
.expect("AsyncJob worker thread panicked or disconnected")
}
pub fn is_ready(&self) -> bool {
!self.receiver.is_empty()
}
}
pub fn run_async<F, T>(func: F) -> AsyncJob<T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let (tx, rx) = bounded(1);
rayon::spawn(move || {
let result = func();
let _ = tx.send(result);
});
AsyncJob { receiver: rx }
}