strided_kernel/
execution_policy.rs

1use core::cell::Cell;
2use core::num::NonZeroUsize;
3
4/// Controls how a strided operation may use CPU threads.
5///
6/// This policy controls fanout only. It does not create a Rayon pool or choose
7/// CPU placement; callers that need placement control should install the
8/// operation in their chosen executor and then select [`Self::Sequential`] or
9/// [`Self::Rayon`].
10///
11/// The policy applies to fanout owned by `strided-kernel`. A bounded worker
12/// partition runs nested strided operations sequentially so nested operations
13/// cannot multiply the outer budget. The scope is worker-local, not a
14/// task-local Rayon context, and is not propagated to threads or Rayon tasks
15/// created by user callbacks.
16///
17/// User callbacks that rely on policy isolation must not enter their own Rayon
18/// scheduling or yield boundary (`join`, `scope`, `spawn`, and similar APIs).
19/// At such a callback-owned boundary, Rayon may execute unrelated work on the
20/// waiting worker while its worker-local policy is active. `strided-kernel`
21/// suspends the policy at every scheduler boundary it owns, but cannot do so
22/// for arbitrary scheduling performed inside a callback.
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub enum ExecutionPolicy {
25    /// Preserve the compatibility behavior of using all threads in the current
26    /// installed Rayon pool, or the global pool when no pool is installed.
27    ///
28    /// Explicit runtimes should prefer [`Self::Sequential`] or [`Self::Rayon`].
29    AmbientRayon,
30    /// Execute entirely on the calling thread with zero Rayon fanout.
31    Sequential,
32    /// Use the currently installed Rayon pool while limiting the operation to
33    /// at most `max_threads` concurrent partitions.
34    ///
35    /// Without the `parallel` feature, operations remain sequential.
36    Rayon { max_threads: NonZeroUsize },
37}
38
39thread_local! {
40    static ACTIVE_POLICY: Cell<ExecutionPolicy> = const { Cell::new(ExecutionPolicy::AmbientRayon) };
41    static ACTIVE_FANOUT: Cell<bool> = const { Cell::new(false) };
42}
43
44fn restrict(outer: ExecutionPolicy, inner: ExecutionPolicy) -> ExecutionPolicy {
45    match (outer, inner) {
46        (ExecutionPolicy::Sequential, _) | (_, ExecutionPolicy::Sequential) => {
47            ExecutionPolicy::Sequential
48        }
49        (ExecutionPolicy::AmbientRayon, policy) | (policy, ExecutionPolicy::AmbientRayon) => policy,
50        (
51            ExecutionPolicy::Rayon { max_threads: outer },
52            ExecutionPolicy::Rayon { max_threads: inner },
53        ) => ExecutionPolicy::Rayon {
54            max_threads: outer.min(inner),
55        },
56    }
57}
58
59#[derive(Clone, Copy)]
60struct ExecutionState {
61    policy: ExecutionPolicy,
62    fanout_active: bool,
63}
64
65struct StateGuard {
66    previous: ExecutionState,
67}
68
69impl Drop for StateGuard {
70    fn drop(&mut self) {
71        set_state(self.previous);
72    }
73}
74
75fn state() -> ExecutionState {
76    ExecutionState {
77        policy: ACTIVE_POLICY.with(Cell::get),
78        fanout_active: ACTIVE_FANOUT.with(Cell::get),
79    }
80}
81
82fn set_state(state: ExecutionState) {
83    ACTIVE_POLICY.with(|active| active.set(state.policy));
84    ACTIVE_FANOUT.with(|active| active.set(state.fanout_active));
85}
86
87#[cfg(feature = "parallel")]
88fn with_state<R>(next: ExecutionState, operation: impl FnOnce() -> R) -> R {
89    let previous = state();
90    set_state(next);
91    let _guard = StateGuard { previous };
92    operation()
93}
94
95/// Execute `operation` under an explicit strided-kernel execution policy.
96///
97/// Nested scopes combine conservatively: sequential execution dominates and
98/// nested Rayon budgets use the smaller limit. The previous policy is restored
99/// even if `operation` panics.
100///
101/// The contract covers strided operations and their library-owned fanout. It
102/// does not turn Rayon worker-local state into task-local state. In particular,
103/// callbacks should not invoke Rayon scheduling or yielding APIs if they depend
104/// on isolation from unrelated work; see [`ExecutionPolicy`] for details.
105///
106/// # Examples
107///
108/// ```rust
109/// use strided_kernel::{
110///     map_into, with_execution_policy, ExecutionPolicy, StridedArray,
111/// };
112///
113/// let source = StridedArray::<f64>::from_fn_col_major(&[3], |index| index[0] as f64);
114/// let mut destination = StridedArray::<f64>::col_major(&[3]);
115/// with_execution_policy(ExecutionPolicy::Sequential, || {
116///     map_into(&mut destination.view_mut(), &source.view(), |value| value + 1.0)
117///         .unwrap();
118/// });
119/// assert_eq!(destination.into_data(), vec![1.0, 2.0, 3.0]);
120/// ```
121#[inline]
122pub fn with_execution_policy<R>(policy: ExecutionPolicy, operation: impl FnOnce() -> R) -> R {
123    let policy = match policy {
124        ExecutionPolicy::AmbientRayon => return operation(),
125        policy => policy,
126    };
127    let previous = state();
128    set_state(ExecutionState {
129        policy: restrict(previous.policy, policy),
130        fanout_active: previous.fanout_active,
131    });
132    let _guard = StateGuard { previous };
133    operation()
134}
135
136#[cfg(feature = "parallel")]
137pub(crate) fn active_policy() -> ExecutionPolicy {
138    ACTIVE_POLICY.with(Cell::get)
139}
140
141#[cfg(feature = "parallel")]
142pub(crate) fn fanout_active() -> bool {
143    ACTIVE_FANOUT.with(Cell::get)
144}
145
146#[cfg(feature = "parallel")]
147#[inline(always)]
148pub(crate) fn with_owned_execution<R>(
149    policy: ExecutionPolicy,
150    fanout_active: bool,
151    operation: impl FnOnce() -> R,
152) -> R {
153    match policy {
154        ExecutionPolicy::AmbientRayon => operation(),
155        ExecutionPolicy::Sequential | ExecutionPolicy::Rayon { .. } => {
156            let previous = state();
157            with_state(
158                ExecutionState {
159                    policy: restrict(previous.policy, policy),
160                    fanout_active: previous.fanout_active || fanout_active,
161                },
162                operation,
163            )
164        }
165    }
166}
167
168#[cfg(feature = "parallel")]
169pub(crate) fn with_scheduler_suspended<R>(operation: impl FnOnce() -> R) -> R {
170    with_state(
171        ExecutionState {
172            policy: ExecutionPolicy::AmbientRayon,
173            fanout_active: false,
174        },
175        operation,
176    )
177}
178
179#[cfg(feature = "parallel")]
180pub(crate) fn permutation_copy_parallel_eligible(
181    policy: ExecutionPolicy,
182    fanout_active: bool,
183    current_pool_threads: usize,
184) -> bool {
185    if fanout_active || current_pool_threads <= 1 {
186        return false;
187    }
188    match policy {
189        ExecutionPolicy::AmbientRayon => true,
190        ExecutionPolicy::Sequential => false,
191        ExecutionPolicy::Rayon { max_threads } => current_pool_threads <= max_threads.get(),
192    }
193}
194
195#[cfg(feature = "parallel")]
196pub(crate) fn rayon_threads() -> usize {
197    if fanout_active() {
198        return 1;
199    }
200    match active_policy() {
201        ExecutionPolicy::Sequential => 1,
202        ExecutionPolicy::AmbientRayon => crate::threading::current_pool_threads(),
203        ExecutionPolicy::Rayon { max_threads } => {
204            crate::threading::current_pool_threads().min(max_threads.get())
205        }
206    }
207}
208
209#[cfg(test)]
210mod default_tests {
211    use super::*;
212    use std::panic::{catch_unwind, AssertUnwindSafe};
213
214    #[test]
215    fn ambient_scope_returns_result_without_changing_state() {
216        let before = state();
217
218        let value = with_execution_policy(ExecutionPolicy::AmbientRayon, || {
219            let active = state();
220            assert_eq!(active.policy, before.policy);
221            assert_eq!(active.fanout_active, before.fanout_active);
222            17usize
223        });
224
225        let after = state();
226        assert_eq!(value, 17);
227        assert_eq!(after.policy, before.policy);
228        assert_eq!(after.fanout_active, before.fanout_active);
229    }
230
231    #[test]
232    fn explicit_scope_restores_state_after_return_and_panic() {
233        let before = state();
234        let two = NonZeroUsize::new(2).unwrap();
235
236        let value = with_execution_policy(ExecutionPolicy::Rayon { max_threads: two }, || {
237            assert_eq!(state().policy, ExecutionPolicy::Rayon { max_threads: two });
238            23usize
239        });
240        assert_eq!(value, 23);
241        assert_eq!(state().policy, before.policy);
242        assert_eq!(state().fanout_active, before.fanout_active);
243
244        let panic = catch_unwind(AssertUnwindSafe(|| {
245            with_execution_policy(ExecutionPolicy::Sequential, || panic!("policy scope panic"));
246        }));
247        assert!(panic.is_err());
248        assert_eq!(state().policy, before.policy);
249        assert_eq!(state().fanout_active, before.fanout_active);
250    }
251
252    #[test]
253    fn nested_explicit_scopes_combine_conservatively() {
254        let two = NonZeroUsize::new(2).unwrap();
255        let four = NonZeroUsize::new(4).unwrap();
256
257        with_execution_policy(ExecutionPolicy::Rayon { max_threads: four }, || {
258            assert_eq!(state().policy, ExecutionPolicy::Rayon { max_threads: four });
259
260            with_execution_policy(ExecutionPolicy::Rayon { max_threads: two }, || {
261                assert_eq!(state().policy, ExecutionPolicy::Rayon { max_threads: two });
262            });
263            assert_eq!(state().policy, ExecutionPolicy::Rayon { max_threads: four });
264
265            with_execution_policy(ExecutionPolicy::Sequential, || {
266                assert_eq!(state().policy, ExecutionPolicy::Sequential);
267            });
268            assert_eq!(state().policy, ExecutionPolicy::Rayon { max_threads: four });
269        });
270
271        with_execution_policy(ExecutionPolicy::Sequential, || {
272            with_execution_policy(ExecutionPolicy::Rayon { max_threads: four }, || {
273                assert_eq!(state().policy, ExecutionPolicy::Sequential);
274            });
275        });
276        assert_eq!(state().policy, ExecutionPolicy::AmbientRayon);
277        assert!(!state().fanout_active);
278    }
279}
280
281#[cfg(all(test, feature = "parallel"))]
282mod tests {
283    use super::*;
284    use std::panic::{catch_unwind, AssertUnwindSafe};
285
286    #[test]
287    fn permutation_copy_parallel_eligibility_is_deterministic() {
288        let two = NonZeroUsize::new(2).unwrap();
289        let four = NonZeroUsize::new(4).unwrap();
290
291        assert!(permutation_copy_parallel_eligible(
292            ExecutionPolicy::Rayon { max_threads: two },
293            false,
294            2,
295        ));
296        assert!(permutation_copy_parallel_eligible(
297            ExecutionPolicy::Rayon { max_threads: four },
298            false,
299            2,
300        ));
301        assert!(!permutation_copy_parallel_eligible(
302            ExecutionPolicy::Rayon { max_threads: two },
303            false,
304            4,
305        ));
306        assert!(!permutation_copy_parallel_eligible(
307            ExecutionPolicy::Rayon { max_threads: two },
308            true,
309            2,
310        ));
311        assert!(!permutation_copy_parallel_eligible(
312            ExecutionPolicy::Sequential,
313            false,
314            2,
315        ));
316        assert!(permutation_copy_parallel_eligible(
317            ExecutionPolicy::AmbientRayon,
318            false,
319            2,
320        ));
321    }
322
323    #[test]
324    fn scheduler_panic_restores_owned_policy_and_fanout_state() {
325        let two = NonZeroUsize::new(2).unwrap();
326        let policy = ExecutionPolicy::Rayon { max_threads: two };
327
328        with_execution_policy(policy, || {
329            with_owned_execution(policy, true, || {
330                let panic = catch_unwind(AssertUnwindSafe(|| {
331                    with_scheduler_suspended(|| panic!("scheduler boundary panic"));
332                }));
333                assert!(panic.is_err());
334                assert_eq!(active_policy(), policy);
335                assert!(fanout_active());
336            });
337        });
338        assert_eq!(active_policy(), ExecutionPolicy::AmbientRayon);
339        assert!(!fanout_active());
340    }
341
342    #[test]
343    fn leaf_panic_restores_ambient_policy_and_inactive_fanout() {
344        let two = NonZeroUsize::new(2).unwrap();
345        let policy = ExecutionPolicy::Rayon { max_threads: two };
346
347        let panic = catch_unwind(AssertUnwindSafe(|| {
348            with_owned_execution(policy, true, || panic!("owned leaf panic"));
349        }));
350        assert!(panic.is_err());
351        assert_eq!(active_policy(), ExecutionPolicy::AmbientRayon);
352        assert!(!fanout_active());
353    }
354}
strided_kernel/execution_policy.rs

strided_kernel/
execution_policy.rs
