binius_utils/rayon/
config.rs

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

binius_utils/rayon/
config.rs
