binius_utils/rayon/
config.rs

1// Copyright 2024-2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4use std::{num::NonZero, sync::OnceLock, thread::available_parallelism};
5
6use super::ThreadPoolBuildError;
7
8/// Builds the global rayon pool, sized to the machine's physical cores.
9///
10/// Rayon's own default is one worker per logical CPU.
11/// On a machine with simultaneous multithreading that puts two workers on every core.
12///
13/// The prover's hot loops are carry-less-multiply and bandwidth bound.
14/// A second worker on the same core competes for those ports and adds no throughput.
15///
16/// The environment still chooses the width when it says so:
17///
18/// ```text
19///     RAYON_NUM_THREADS=1   one worker, running on the calling thread
20///     RAYON_NUM_THREADS=n   left to rayon, which reads the variable itself
21///     unset                 one worker per physical core
22/// ```
23///
24/// A single-worker pool on the calling thread keeps worker frames out of stack traces.
25///
26/// Rayon builds the global pool on first use and refuses to build it twice.
27/// So this must run before anything else touches the pool.
28/// A later call reports that earlier use as an error, which a caller may treat as advisory.
29///
30/// The result is computed once and cached, so calling more than once is harmless.
31///
32/// # Returns
33///
34/// A reference, because the error type cannot be cloned.
35pub fn adjust_thread_pool() -> &'static Result<(), ThreadPoolBuildError> {
36	static ONCE_GUARD: OnceLock<Result<(), ThreadPoolBuildError>> = OnceLock::new();
37
38	ONCE_GUARD.get_or_init(|| {
39		// Reading the environment avoids asking rayon for its current width.
40		// Asking would build the global pool, leaving nothing left to override.
41		match std::env::var("RAYON_NUM_THREADS") {
42			// One worker on the calling thread, so no worker frames appear in a stack trace.
43			Ok(v) if v == "1" => super::ThreadPoolBuilder::new()
44				.num_threads(1)
45				.use_current_thread()
46				.build_global(),
47			// Rayon reads this variable itself.
48			// Leaving the pool unbuilt is what lets it apply the requested width.
49			Ok(_) => Ok(()),
50			// Unset: size the pool to the cores that can actually run in parallel.
51			// An unreadable topology falls back to rayon's default rather than guessing.
52			Err(_) => physical_core_count().map_or(Ok(()), |n| {
53				super::ThreadPoolBuilder::new()
54					.num_threads(n.get())
55					.build_global()
56			}),
57		}
58	})
59}
60
61/// The number of physical cores this process may run on.
62///
63/// A core is one package-and-core pair in the kernel's topology.
64/// Distinct pairs collapse the sibling threads that share a core into one.
65///
66/// Returns nothing when the topology is unreadable, leaving the choice to the caller.
67fn physical_core_count() -> Option<NonZero<usize>> {
68	let logical = available_parallelism().ok()?;
69	let physical = platform_physical_cores()?;
70	// A cgroup quota or an affinity mask can hide part of the machine.
71	// Clamping keeps the width within what this process is allowed to use.
72	Some(physical.min(logical))
73}
74
75/// Counts the distinct package-and-core pairs the kernel reports under `/sys`.
76///
77/// The kernel exposes one topology directory per online logical CPU.
78/// Every sibling of a core repeats that core's pair, so the set of pairs is the set of cores.
79#[cfg(target_os = "linux")]
80fn platform_physical_cores() -> Option<NonZero<usize>> {
81	use std::{collections::HashSet, fs, path::Path};
82
83	let mut cores = HashSet::new();
84	for entry in fs::read_dir("/sys/devices/system/cpu").ok()? {
85		let path = entry.ok()?.path();
86
87		// Only the per-CPU directories carry a topology.
88		// Anything else in this directory is unrelated, so skip it instead of failing.
89		//
90		//     cpu0, cpu1, ... cpu31   -> read
91		//     cpufreq, power, online  -> skip
92		let is_cpu_dir = path
93			.file_name()
94			.and_then(|name| name.to_str())
95			.is_some_and(|name| {
96				name.starts_with("cpu") && name[3..].bytes().all(|b| b.is_ascii_digit())
97			});
98		if !is_cpu_dir {
99			continue;
100		}
101
102		// Each identifier is a single decimal number in its own file.
103		let read_id = |field: &str| -> Option<u32> {
104			let file: &Path = &path.join("topology").join(field);
105			fs::read_to_string(file).ok()?.trim().parse().ok()
106		};
107
108		// A CPU missing either identifier is offline or unsupported, so skip it.
109		if let (Some(package), Some(core)) = (read_id("physical_package_id"), read_id("core_id")) {
110			cores.insert((package, core));
111		}
112	}
113
114	NonZero::new(cores.len())
115}
116
117/// Reports no core count on platforms whose topology this does not read.
118#[cfg(not(target_os = "linux"))]
119fn platform_physical_cores() -> Option<NonZero<usize>> {
120	None
121}
122
123#[cfg(test)]
124mod tests {
125	use super::*;
126
127	#[test]
128	fn physical_core_count_is_within_available_parallelism() {
129		// Invariant: the count is a thread-pool width, so it can never exceed the parallelism
130		// this process is permitted to use.
131		let Some(physical) = physical_core_count() else {
132			// A platform without a readable topology has nothing to check.
133			return;
134		};
135
136		// On this host: 16 physical cores against 32 logical CPUs.
137		let logical =
138			available_parallelism().expect("available parallelism is known on test hosts");
139		assert!(physical <= logical);
140	}
141}
binius_utils/rayon/config.rs

binius_utils/rayon/
config.rs
