ffai_core/
threads.rs

1//! How many threads this process may actually run on.
2//!
3//! `std::thread::available_parallelism` ignores the process affinity mask on
4//! Windows: under `start /affinity 4` (one core) it still reports every
5//! logical CPU. rayon sizes its global pool the same way, and candle hands
6//! gemm `Parallelism::Rayon(physical cores)` on every matmul, so a process
7//! pinned to one core ran two dozen threads on it.
8//!
9//! `PdfCraft` measured what sizing the pool to the mask is worth (ABBA, output
10//! identical): on 8 cores, 0.927x the time (10/10 rounds) and peak memory
11//! 1 959 -> 1 240 MB; on 1 core, time neutral and memory 1 285 -> 747 MB. It
12//! could not do it itself, because reading the mask is an FFI call and it
13//! forbids `unsafe`. This module is that one call, behind a safe function.
14//!
15//! Use [`affinity_pool`] and run the engine inside it:
16//!
17//! ```no_run
18//! let pool = ffai_core::threads::affinity_pool().expect("thread pool");
19//! pool.install(|| {
20//!     // engine.recognize(&image, &options) — every rayon task the engine and
21//!     // candle spawn now runs on this pool's threads.
22//! });
23//! ```
24
25/// The number of CPUs this process may run on: the affinity mask on Windows,
26/// `available_parallelism` elsewhere (which already honours affinity and
27/// cgroup quotas on Linux). At least 1.
28#[must_use]
29pub fn available_threads() -> usize {
30    #[cfg(windows)]
31    if let Some(n) = windows_affinity() {
32        return n;
33    }
34    std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get)
35}
36
37#[cfg(windows)]
38fn windows_affinity() -> Option<usize> {
39    use windows_sys::Win32::System::Threading::{GetCurrentProcess, GetProcessAffinityMask};
40    let (mut process, mut system) = (0usize, 0usize);
41    // SAFETY: `GetCurrentProcess` returns a pseudo-handle that is always valid
42    // for the calling process and needs no closing; both out-pointers point at
43    // live, writable `usize`s for the duration of the call.
44    #[allow(unsafe_code)]
45    let ok = unsafe { GetProcessAffinityMask(GetCurrentProcess(), &raw mut process, &raw mut system) };
46    // The mask covers the process's processor group (at most 64 CPUs).
47    (ok != 0 && process != 0).then(|| process.count_ones() as usize)
48}
49
50/// A rayon pool with [`available_threads`] threads.
51///
52/// # Errors
53/// When rayon cannot spawn the threads.
54pub fn affinity_pool() -> Result<rayon::ThreadPool, rayon::ThreadPoolBuildError> {
55    rayon::ThreadPoolBuilder::new().num_threads(available_threads()).build()
56}
57
58#[cfg(test)]
59mod tests {
60    #[test]
61    fn at_least_one_and_at_most_the_machine() {
62        let n = super::available_threads();
63        assert!(n >= 1);
64        // On Windows the mask can only narrow the count, never widen it.
65        #[cfg(windows)]
66        assert!(n <= std::thread::available_parallelism().map_or(usize::MAX, std::num::NonZeroUsize::get));
67    }
68}
ffai_core/threads.rs

ffai_core/
threads.rs
