strided_kernel/
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    pub(crate) fn run<R>(&self, operation: impl FnOnce() -> R) -> R {
83        match self.kind {
84            ExecContextKind::Serial => {
85                with_execution_policy(ExecutionPolicy::Sequential, operation)
86            }
87            ExecContextKind::MaxThreads(max_threads) => {
88                with_execution_policy(ExecutionPolicy::Rayon { max_threads }, operation)
89            }
90            ExecContextKind::Ambient => operation(),
91        }
92    }
93}
94
95impl Default for ExecContext {
96    #[inline]
97    fn default() -> Self {
98        Self::serial()
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::ExecContext;
105    use crate::StridedError;
106
107    #[test]
108    fn serial_context_has_no_thread_limit() {
109        let ctx = ExecContext::serial();
110
111        assert!(ctx.is_serial());
112        assert!(!ctx.is_ambient());
113        assert_eq!(ctx.max_threads_limit(), None);
114        assert_eq!(ExecContext::default(), ctx);
115    }
116
117    #[test]
118    fn bounded_context_rejects_zero_and_exposes_limit() {
119        let ctx = ExecContext::max_threads(4).unwrap();
120
121        assert!(!ctx.is_serial());
122        assert!(!ctx.is_ambient());
123        assert_eq!(ctx.max_threads_limit().map(|value| value.get()), Some(4));
124        assert!(matches!(
125            ExecContext::max_threads(0).unwrap_err(),
126            StridedError::InvalidThreadBudget { max_threads: 0 }
127        ));
128    }
129
130    #[test]
131    fn ambient_context_has_no_thread_limit() {
132        let ctx = ExecContext::ambient();
133
134        assert!(!ctx.is_serial());
135        assert!(ctx.is_ambient());
136        assert_eq!(ctx.max_threads_limit(), None);
137    }
138
139    #[cfg(feature = "parallel")]
140    #[test]
141    fn run_installs_execution_policy_and_restores_previous_policy() {
142        use core::num::NonZeroUsize;
143
144        use crate::execution_policy::{active_policy, with_execution_policy, ExecutionPolicy};
145
146        let two = NonZeroUsize::new(2).unwrap();
147        let four = NonZeroUsize::new(4).unwrap();
148        let bounded = ExecContext::max_threads(2).unwrap();
149
150        let observed_bounded =
151            with_execution_policy(ExecutionPolicy::Rayon { max_threads: four }, || {
152                bounded.run(active_policy)
153            });
154        let observed_serial =
155            with_execution_policy(ExecutionPolicy::Rayon { max_threads: four }, || {
156                ExecContext::serial().run(active_policy)
157            });
158        let observed_ambient =
159            with_execution_policy(ExecutionPolicy::Rayon { max_threads: four }, || {
160                ExecContext::ambient().run(active_policy)
161            });
162
163        assert_eq!(
164            observed_bounded,
165            ExecutionPolicy::Rayon { max_threads: two }
166        );
167        assert_eq!(observed_serial, ExecutionPolicy::Sequential);
168        assert_eq!(
169            observed_ambient,
170            ExecutionPolicy::Rayon { max_threads: four }
171        );
172        assert_eq!(active_policy(), ExecutionPolicy::AmbientRayon);
173    }
174}
strided_kernel/exec_context.rs

strided_kernel/
exec_context.rs
