strided_basic/
exec_context.rs

1//! Execution policy passed through erased kernel replay boundaries.
2//!
3//! `ExecContext` is explicit even for kernels that are currently serial so
4//! downstream runtimes do not accidentally bake ambient thread-pool state into
5//! their prepared-kernel ABI.
6
7use core::num::NonZeroUsize;
8
9use crate::{with_execution_policy, ExecutionPolicy, Result, StridedError};
10
11/// Caller-selected execution policy for prepared kernel replay.
12///
13/// This type intentionally hides its representation so future replay families
14/// can add provider-owned pools or scheduling scopes without forcing downstream
15/// crates to pattern-match a closed enum.
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub struct ExecContext {
18    kind: ExecContextKind,
19}
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22enum ExecContextKind {
23    Serial,
24    MaxThreads(NonZeroUsize),
25    Ambient,
26}
27
28impl ExecContext {
29    /// Execute without entering a parallel worker pool.
30    #[inline]
31    pub const fn serial() -> Self {
32        Self {
33            kind: ExecContextKind::Serial,
34        }
35    }
36
37    /// Execute with an operation-local upper bound on worker threads.
38    #[inline]
39    pub fn max_threads(max_threads: usize) -> Result<Self> {
40        match NonZeroUsize::new(max_threads) {
41            Some(max_threads) => Ok(Self {
42                kind: ExecContextKind::MaxThreads(max_threads),
43            }),
44            None => Err(StridedError::InvalidThreadBudget { max_threads }),
45        }
46    }
47
48    /// Execute using the ambient runtime policy.
49    ///
50    /// This is useful for direct `strided-kernel` users. Runtime crates that
51    /// own CPU resources should prefer [`ExecContext::serial`] or
52    /// [`ExecContext::max_threads`] so thread ownership remains explicit.
53    #[inline]
54    pub const fn ambient() -> Self {
55        Self {
56            kind: ExecContextKind::Ambient,
57        }
58    }
59
60    /// Returns `true` when this context requires serial execution.
61    #[inline]
62    pub fn is_serial(&self) -> bool {
63        matches!(self.kind, ExecContextKind::Serial)
64    }
65
66    /// Returns `true` when this context delegates to ambient runtime policy.
67    #[inline]
68    pub fn is_ambient(&self) -> bool {
69        matches!(self.kind, ExecContextKind::Ambient)
70    }
71
72    /// Returns the configured worker-thread upper bound, if any.
73    #[inline]
74    pub fn max_threads_limit(&self) -> Option<NonZeroUsize> {
75        match self.kind {
76            ExecContextKind::MaxThreads(max_threads) => Some(max_threads),
77            ExecContextKind::Serial | ExecContextKind::Ambient => None,
78        }
79    }
80
81    #[inline]
82    /// Run an operation under this context, preserving enclosing execution limits.
83    ///
84    /// # Examples
85    ///
86    /// ```
87    /// use strided_basic::ExecContext;
88    /// assert_eq!(ExecContext::serial().run(|| 42), 42);
89    /// ```
90    pub fn run<R>(&self, operation: impl FnOnce() -> R) -> R {
91        match self.kind {
92            ExecContextKind::Serial => {
93                with_execution_policy(ExecutionPolicy::Sequential, operation)
94            }
95            ExecContextKind::MaxThreads(max_threads) => {
96                with_execution_policy(ExecutionPolicy::Rayon { max_threads }, operation)
97            }
98            ExecContextKind::Ambient => operation(),
99        }
100    }
101}
102
103#[cfg(feature = "parallel")]
104impl ExecContext {
105    /// Enter `pool` for `operation`, bounding fanout to `max_threads` workers.
106    ///
107    /// Placement belongs to the caller: this never creates a pool, and the
108    /// operation runs on the pool it is handed rather than on whichever pool
109    /// happens to be installed, so a host that already owns a pool shares that
110    /// one with these kernels instead of risking a second team.
111    ///
112    /// The width is clamped to the pool's worker count. A width of one enters no
113    /// Rayon machinery at all and `operation` sees [`ExecContext::serial`]. When
114    /// the calling thread is already a worker of `pool`, the operation runs in
115    /// place rather than nesting a second team.
116    ///
117    /// # Errors
118    ///
119    /// [`StridedError::InvalidThreadBudget`] when `max_threads` is zero.
120    ///
121    /// # Examples
122    ///
123    /// ```
124    /// use strided_basic::ExecContext;
125    ///
126    /// let pool = rayon::ThreadPoolBuilder::new()
127    ///     .num_threads(2)
128    ///     .build()
129    ///     .unwrap();
130    /// assert!(ExecContext::run_on(&pool, 1, |ctx| ctx.is_serial()).unwrap());
131    /// let limit = ExecContext::run_on(&pool, 2, |ctx| ctx.max_threads_limit()).unwrap();
132    /// assert_eq!(limit.map(|value| value.get()), Some(2));
133    /// ```
134    pub fn run_on<R: Send>(
135        pool: &rayon::ThreadPool,
136        max_threads: usize,
137        operation: impl FnOnce(ExecContext) -> R + Send,
138    ) -> Result<R> {
139        let budget = NonZeroUsize::new(max_threads)
140            .ok_or(StridedError::InvalidThreadBudget { max_threads })?;
141        let width = budget.get().min(pool.current_num_threads());
142        if width <= 1 {
143            let ctx = Self::serial();
144            return Ok(ctx.run(|| operation(ctx)));
145        }
146        let ctx = Self::max_threads(width)?;
147        let run = move || ctx.run(|| operation(ctx));
148        if pool.current_thread_index().is_some() {
149            Ok(run())
150        } else {
151            Ok(pool.install(run))
152        }
153    }
154}
155
156impl Default for ExecContext {
157    #[inline]
158    fn default() -> Self {
159        Self::serial()
160    }
161}
162
163#[cfg(test)]
164#[path = "exec_context/tests/tests.rs"]
165mod tests;
strided_basic/exec_context.rs

strided_basic/
exec_context.rs
