binius_utils/
rayon.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use cfg_if::cfg_if;
4
5/// In case when number of threads is set to 1, use rayon thread pool with
6/// `use_current_thread` set to true. This is solves two problems:
7/// 1. The performance is almost the same as if rayon wasn't used at all
8/// 2. Makes profiling and debugging results less noisy
9///
10/// NOTE: rayon doesn't allow initializing global thread pool several times, so
11/// in case when it was initialized before this function returns an error.
12/// The typical usage of the function is to place it's call in the beginning of the `main`.
13/// The function returns reference to the result because `ThreadPoolBuildError`
14/// doesn't implement `Clone`.
15pub fn adjust_thread_pool() -> &'static Result<(), binius_maybe_rayon::ThreadPoolBuildError> {
16	cfg_if! {
17		if #[cfg(feature = "rayon")] {
18			use std::sync::OnceLock;
19
20			static ONCE_GUARD: OnceLock<Result<(), binius_maybe_rayon::ThreadPoolBuildError>> = OnceLock::new();
21
22			ONCE_GUARD.get_or_init(|| {
23				// We cannot use `binius_maybe_rayon::get_current_threads` because it would force the global thread pool
24				// to initialize, so we won't be able to override it.
25				match std::env::var("RAYON_NUM_THREADS") {
26					Ok(v) if v == "1" => binius_maybe_rayon::ThreadPoolBuilder::new()
27						.num_threads(1)
28						.use_current_thread()
29						.build_global(),
30					_ => Ok(()),
31				}
32			})
33		}
34		else {
35			static RESULT: Result<(), binius_maybe_rayon::ThreadPoolBuildError> = Ok(());
36
37			&RESULT
38		}
39	}
40}
41
42/// Returns the base-2 logarithm of the number of threads that should be used for the task
43pub fn get_log_max_threads() -> usize {
44	(2 * binius_maybe_rayon::current_num_threads() - 1).ilog2() as _
45}