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}