tenferro_cpu/
context.rs

1use std::env;
2use std::num::NonZeroUsize;
3use std::sync::Arc;
4
5use rayon::prelude::*;
6use thiserror::Error as ThisError;
7
8use crate::affinity::{CpuAffinityError, SystemThreadAffinity, ThreadAffinity};
9use crate::arbiter::{
10    current_execution_owner, register_worker_execution_scope, worker_execution_scope_matches,
11    ExecutionScopeState,
12};
13use crate::domain_executor::{
14    CpuDomainExecutor, CpuDomainExecutorCapabilities, CpuDomainExecutorError, CpuExecutorAffinity,
15    CpuExecutorReentrancy, CpuExecutorShutdown, CpuInnerParallelism, ScopedCpuJob, ScopedCpuJobs,
16};
17use crate::{CpuId, CpuSet, Error, ErrorKind, Result, ValidationKind};
18
19/// Failure to construct a CPU context with pinned Rayon workers.
20///
21/// # Examples
22///
23/// ```
24/// use tenferro_cpu::CpuContextError;
25///
26/// let error = CpuContextError::InvalidThreadCount;
27/// assert!(error.to_string().contains("thread count"));
28/// ```
29#[derive(Debug, ThisError)]
30pub enum CpuContextError {
31    /// A context must contain at least one worker.
32    #[error("thread count must be at least 1")]
33    InvalidThreadCount,
34    /// A pinned engine cannot create more workers than assigned CPUs.
35    #[error("requested {workers} workers for only {cpus} assigned CPUs")]
36    TooManyWorkers {
37        /// Requested Rayon worker count.
38        workers: usize,
39        /// Number of logical CPUs in the execution domain.
40        cpus: usize,
41    },
42    /// Rayon could not construct the custom thread pool.
43    #[error("failed to build pinned CPU thread pool: {source}")]
44    PoolBuild {
45        /// Rayon or OS thread-spawn error.
46        #[source]
47        source: rayon::ThreadPoolBuildError,
48    },
49    /// A worker could not set or verify its assigned CPU affinity.
50    #[error("failed to pin worker {worker} to CPU {cpu}: {source}")]
51    WorkerPinning {
52        /// Stable Rayon worker index.
53        worker: usize,
54        /// Assigned operating-system logical CPU.
55        cpu: CpuId,
56        /// OS or verification failure.
57        #[source]
58        source: CpuAffinityError,
59    },
60    /// A worker terminated before reporting startup affinity.
61    #[error("worker startup channel closed before all workers reported: {source}")]
62    WorkerStartupClosed {
63        /// Channel receive failure from the worker startup handshake.
64        #[source]
65        source: std::sync::mpsc::RecvError,
66    },
67}
68
69/// Reusable CPU execution context carrying CPU parallelism policy.
70///
71/// `CpuContext` stores the requested thread count as a kernel-level
72/// parallelism hint and owns the Rayon pool used by multi-threaded CPU work.
73///
74/// # Examples
75///
76/// ```
77/// use tenferro_cpu::CpuContext;
78///
79/// let ctx = CpuContext::with_threads(1).unwrap();
80/// let value = ctx.install(|| 1 + 1);
81/// assert_eq!(value, 2);
82/// assert_eq!(ctx.num_threads(), 1);
83/// ```
84#[derive(Clone, Debug)]
85pub struct CpuContext {
86    num_threads: usize,
87    pool: Option<Arc<rayon::ThreadPool>>,
88    pinned_cpus: Option<CpuSet>,
89    execution_scope: Arc<ExecutionScopeState>,
90    #[cfg(test)]
91    executor_install_calls: Arc<std::sync::atomic::AtomicUsize>,
92}
93
94impl CpuContext {
95    /// Create a CPU context from `RAYON_NUM_THREADS`, or fall back to a
96    /// single-threaded context with a stderr warning when validation fails.
97    ///
98    /// # Examples
99    ///
100    /// ```
101    /// use tenferro_cpu::CpuContext;
102    ///
103    /// let ctx = CpuContext::from_env();
104    /// assert!(ctx.num_threads() >= 1);
105    /// ```
106    pub fn from_env() -> Self {
107        Self::try_from_env().unwrap_or_else(|err| {
108            eprintln!(
109                "tenferro_cpu: falling back to single-threaded CPU context after configuration error: {err}"
110            );
111            Self::single_threaded()
112        })
113    }
114
115    /// Try to create a CPU context from `RAYON_NUM_THREADS`.
116    ///
117    /// # Examples
118    ///
119    /// ```
120    /// use tenferro_cpu::CpuContext;
121    ///
122    /// let ctx = CpuContext::try_from_env()
123    ///     .unwrap_or_else(|_| CpuContext::with_threads(1).unwrap());
124    /// assert!(ctx.num_threads() >= 1);
125    /// ```
126    ///
127    /// # Errors
128    ///
129    /// Returns [`CpuContextError`] when `RAYON_NUM_THREADS` is malformed or
130    /// requests an invalid worker count.
131    pub fn try_from_env() -> Result<Self> {
132        match env::var("RAYON_NUM_THREADS") {
133            Ok(value) => {
134                let num_threads = value.parse::<usize>().map_err(|err| {
135                    Error::extension(
136                        "CpuContext::try_from_env",
137                        "cpu",
138                        ErrorKind::Validation(ValidationKind::InvalidArgument),
139                        err,
140                    )
141                })?;
142                Self::with_threads(num_threads).map_err(|err| match err {
143                    Error::Validation { source, .. } => {
144                        Error::validation("CpuContext::try_from_env", source)
145                    }
146                    err => err,
147                })
148            }
149            Err(env::VarError::NotPresent) => {
150                Self::with_threads(super::affinity::available_parallelism())
151            }
152            Err(err) => Err(Error::extension(
153                "CpuContext::try_from_env",
154                "cpu",
155                ErrorKind::Validation(ValidationKind::InvalidArgument),
156                err,
157            )),
158        }
159    }
160
161    /// Create a CPU context with a fixed parallelism hint.
162    ///
163    /// # Examples
164    ///
165    /// ```
166    /// use tenferro_cpu::CpuContext;
167    ///
168    /// let ctx = CpuContext::with_threads(2).unwrap();
169    /// assert_eq!(ctx.num_threads(), 2);
170    /// ```
171    ///
172    /// # Errors
173    ///
174    /// Returns [`CpuContextError::InvalidThreadCount`] through
175    /// [`Error::Validation`] when `num_threads` is zero, or
176    /// [`Error::BackendSource`] when Rayon rejects the thread pool.
177    pub fn with_threads(num_threads: usize) -> Result<Self> {
178        if num_threads == 0 {
179            return Err(Error::invalid_argument(
180                "CpuContext::with_threads",
181                "configuration",
182                "thread count must be at least 1",
183            ));
184        }
185        let execution_scope = Arc::new(ExecutionScopeState::default());
186        let pool = if num_threads == 1 {
187            None
188        } else {
189            let (startup_tx, startup_rx) = std::sync::mpsc::channel();
190            let worker_scope = Arc::clone(&execution_scope);
191            let pool = rayon::ThreadPoolBuilder::new()
192                .num_threads(num_threads)
193                .start_handler(move |_| {
194                    register_worker_execution_scope(Arc::clone(&worker_scope));
195                    let _ = startup_tx.send(());
196                })
197                .build()
198                .map_err(|source| Error::backend_source("CpuContext::with_threads", source))?;
199            for _ in 0..num_threads {
200                startup_rx
201                    .recv()
202                    .map_err(|source| Error::backend_source("CpuContext::with_threads", source))?;
203            }
204            Some(Arc::new(pool))
205        };
206        Ok(Self {
207            num_threads,
208            pool,
209            pinned_cpus: None,
210            execution_scope,
211            #[cfg(test)]
212            executor_install_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
213        })
214    }
215
216    /// Create a Rayon context whose workers are pinned to assigned logical CPUs.
217    ///
218    /// A real Rayon pool is constructed even when `num_threads` is one. The
219    /// worker count cannot exceed the assigned CPU count.
220    ///
221    /// # Examples
222    ///
223    /// ```
224    /// use tenferro_cpu::{process_cpu_affinity, CpuContext};
225    ///
226    /// if let Some(allowed) = process_cpu_affinity() {
227    ///     let one_cpu = tenferro_cpu::CpuSet::new([allowed.as_slice()[0]])?;
228    ///     let context = CpuContext::with_pinned_cpus(one_cpu.clone(), 1)?;
229    ///     assert_eq!(context.pinned_cpus(), Some(&one_cpu));
230    /// }
231    /// # Ok::<(), Box<dyn std::error::Error>>(())
232    /// ```
233    ///
234    /// # Errors
235    ///
236    /// Returns [`CpuContextError::InvalidThreadCount`] for zero workers,
237    /// [`CpuContextError::TooManyWorkers`] when the request exceeds the CPU
238    /// set, or an affinity error when workers cannot be pinned.
239    pub fn with_pinned_cpus(
240        cpus: CpuSet,
241        num_threads: usize,
242    ) -> std::result::Result<Self, CpuContextError> {
243        Self::with_pinned_cpus_using(cpus, num_threads, SystemThreadAffinity)
244    }
245
246    pub(crate) fn with_pinned_cpus_using<A: ThreadAffinity>(
247        cpus: CpuSet,
248        num_threads: usize,
249        affinity: A,
250    ) -> std::result::Result<Self, CpuContextError> {
251        if num_threads == 0 {
252            return Err(CpuContextError::InvalidThreadCount);
253        }
254        if num_threads > cpus.len() {
255            return Err(CpuContextError::TooManyWorkers {
256                workers: num_threads,
257                cpus: cpus.len(),
258            });
259        }
260
261        let execution_scope = Arc::new(ExecutionScopeState::default());
262        let assigned_cpus = Arc::new(select_worker_cpus(&cpus, num_threads));
263        let (startup_tx, startup_rx) = std::sync::mpsc::channel();
264        let pool_assigned_cpus = Arc::clone(&assigned_cpus);
265        let worker_scope = Arc::clone(&execution_scope);
266        let pool = rayon::ThreadPoolBuilder::new()
267            .num_threads(num_threads)
268            .spawn_handler(move |thread| {
269                let worker = thread.index();
270                let cpu = pool_assigned_cpus[worker];
271                let startup_tx = startup_tx.clone();
272                let affinity = affinity.clone();
273                let worker_scope = Arc::clone(&worker_scope);
274                std::thread::Builder::new()
275                    .name(format!("tenferro-cpu-{cpu}"))
276                    .spawn(move || {
277                        register_worker_execution_scope(Arc::clone(&worker_scope));
278                        let result = affinity.pin_current(cpu).and_then(|observed| {
279                            (observed.len() == 1 && observed.contains(cpu))
280                                .then_some(())
281                                .ok_or_else(|| CpuAffinityError::Verification {
282                                    observed: observed.as_slice().to_vec(),
283                                })
284                        });
285                        let _ = startup_tx.send((worker, cpu, result));
286                        thread.run();
287                    })
288                    .map(|_| ())
289            })
290            .build()
291            .map_err(|source| CpuContextError::PoolBuild { source })?;
292        let pool = Arc::new(pool);
293        for _ in 0..num_threads {
294            let (worker, cpu, result) = startup_rx
295                .recv()
296                .map_err(|source| CpuContextError::WorkerStartupClosed { source })?;
297            if let Err(source) = result {
298                return Err(CpuContextError::WorkerPinning {
299                    worker,
300                    cpu,
301                    source,
302                });
303            }
304        }
305        Ok(Self {
306            num_threads,
307            pool: Some(pool),
308            pinned_cpus: Some(cpus),
309            execution_scope,
310            #[cfg(test)]
311            executor_install_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
312        })
313    }
314
315    fn single_threaded() -> Self {
316        Self {
317            num_threads: 1,
318            pool: None,
319            pinned_cpus: None,
320            execution_scope: Arc::new(ExecutionScopeState::default()),
321            #[cfg(test)]
322            executor_install_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
323        }
324    }
325
326    /// Return this context's CPU parallelism hint.
327    ///
328    /// # Examples
329    ///
330    /// ```
331    /// use tenferro_cpu::CpuContext;
332    ///
333    /// let ctx = CpuContext::with_threads(2).unwrap();
334    /// assert_eq!(ctx.num_threads(), 2);
335    /// ```
336    pub fn num_threads(&self) -> usize {
337        self.num_threads
338    }
339
340    /// Return the worker CPU domain for a pinned context.
341    ///
342    /// Legacy thread-count-only contexts return `None` because they do not own
343    /// worker affinity.
344    ///
345    /// # Examples
346    ///
347    /// ```
348    /// use tenferro_cpu::CpuContext;
349    ///
350    /// assert_eq!(CpuContext::with_threads(1)?.pinned_cpus(), None);
351    /// # Ok::<(), tenferro_tensor::Error>(())
352    /// ```
353    pub fn pinned_cpus(&self) -> Option<&CpuSet> {
354        self.pinned_cpus.as_ref()
355    }
356
357    /// Run a closure inside this context's CPU execution scope.
358    ///
359    /// # Examples
360    ///
361    /// ```
362    /// use tenferro_cpu::CpuContext;
363    ///
364    /// let ctx = CpuContext::with_threads(1).unwrap();
365    /// let value = ctx.install(|| 1 + 1);
366    /// assert_eq!(value, 2);
367    /// ```
368    pub fn install<R: Send>(&self, op: impl FnOnce() -> R + Send) -> R {
369        match &self.pool {
370            Some(pool) => pool.install(op),
371            None => op(),
372        }
373    }
374
375    pub(crate) fn install_if_needed<R: Send>(&self, op: impl FnOnce() -> R + Send) -> R {
376        if self.pool.is_some() && worker_execution_scope_matches(&self.execution_scope) {
377            op()
378        } else {
379            self.install(op)
380        }
381    }
382
383    #[cfg(test)]
384    pub(crate) fn owns_current_worker_for_test(&self) -> bool {
385        worker_execution_scope_matches(&self.execution_scope)
386    }
387
388    #[cfg(test)]
389    pub(crate) fn executor_install_calls_for_test(&self) -> usize {
390        self.executor_install_calls
391            .load(std::sync::atomic::Ordering::Relaxed)
392    }
393}
394
395impl CpuDomainExecutor for CpuContext {
396    fn capabilities(&self) -> CpuDomainExecutorCapabilities {
397        // INVARIANT: every CpuContext constructor rejects zero workers, and
398        // `num_threads` is private so it cannot be invalidated after creation.
399        let worker_count = match NonZeroUsize::new(self.num_threads) {
400            Some(worker_count) => worker_count,
401            None => unreachable!("CpuContext must contain at least one worker"),
402        };
403        CpuDomainExecutorCapabilities {
404            worker_count,
405            outer_parallelism: self.num_threads > 1,
406            inner_parallelism: if self.pool.is_some() {
407                CpuInnerParallelism::Rayon
408            } else {
409                CpuInnerParallelism::None
410            },
411            // This permits internal entry through the same executor. Public
412            // CpuBackend re-entry remains guarded by BACKEND_REENTRY_PANIC.
413            reentrancy: CpuExecutorReentrancy::SameExecutor,
414            affinity: if self.pinned_cpus.is_some() {
415                CpuExecutorAffinity::TenferroPinnedVerified
416            } else {
417                CpuExecutorAffinity::None
418            },
419            shutdown: CpuExecutorShutdown::TenferroOwned,
420        }
421    }
422
423    fn submit(&self, jobs: &dyn ScopedCpuJobs) -> std::result::Result<(), CpuDomainExecutorError> {
424        let _scope = current_execution_owner().map(|owner| self.execution_scope.enter(owner));
425        if self.pool.is_none() {
426            return (0..jobs.len()).try_for_each(|index| jobs.run(index));
427        }
428        self.install_if_needed(|| {
429            (0..jobs.len())
430                .into_par_iter()
431                .try_for_each(|index| jobs.run(index))
432        })
433    }
434
435    fn install(
436        &self,
437        job: &mut dyn ScopedCpuJob,
438    ) -> std::result::Result<(), CpuDomainExecutorError> {
439        #[cfg(test)]
440        self.executor_install_calls
441            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
442        let _scope = current_execution_owner().map(|owner| self.execution_scope.enter(owner));
443        self.install_if_needed(|| job.run())
444    }
445}
446
447fn select_worker_cpus(cpus: &CpuSet, num_threads: usize) -> Vec<CpuId> {
448    if num_threads == 1 {
449        return vec![cpus.as_slice()[cpus.len() / 2]];
450    }
451    (0..num_threads)
452        .map(|worker| {
453            let index = ((worker as u128) * ((cpus.len() - 1) as u128)
454                / ((num_threads - 1) as u128)) as usize;
455            cpus.as_slice()[index]
456        })
457        .collect()
458}
459
460#[cfg(test)]
461mod tests;
tenferro_cpu/context.rs

tenferro_cpu/
context.rs
