strided_basic/
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_basic::{
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    let policy = ACTIVE_POLICY.with(Cell::get);
139    let Some(width) = crate::threading::target_pool_width() else {
140        return policy;
141    };
142    if width <= 1 {
143        return ExecutionPolicy::Sequential;
144    }
145    let Some(max_threads) = NonZeroUsize::new(width) else {
146        return policy;
147    };
148    restrict(policy, ExecutionPolicy::Rayon { max_threads })
149}
150
151#[cfg(feature = "parallel")]
152pub(crate) fn fanout_active() -> bool {
153    ACTIVE_FANOUT.with(Cell::get)
154}
155
156#[cfg(feature = "parallel")]
157#[inline(always)]
158pub(crate) fn with_owned_execution<R>(
159    policy: ExecutionPolicy,
160    fanout_active: bool,
161    operation: impl FnOnce() -> R,
162) -> R {
163    match policy {
164        ExecutionPolicy::AmbientRayon => operation(),
165        ExecutionPolicy::Sequential | ExecutionPolicy::Rayon { .. } => {
166            let previous = state();
167            with_state(
168                ExecutionState {
169                    policy: restrict(previous.policy, policy),
170                    fanout_active: previous.fanout_active || fanout_active,
171                },
172                operation,
173            )
174        }
175    }
176}
177
178#[cfg(feature = "parallel")]
179pub(crate) fn with_scheduler_suspended<R>(operation: impl FnOnce() -> R) -> R {
180    crate::threading::with_target_pool_suspended(|| {
181        with_state(
182            ExecutionState {
183                policy: ExecutionPolicy::AmbientRayon,
184                fanout_active: false,
185            },
186            operation,
187        )
188    })
189}
190
191/// Clamp a caller-supplied split budget to the effective policy width.
192#[cfg(feature = "parallel")]
193#[inline]
194pub(crate) fn clamp_fanout_width(requested: usize, policy: ExecutionPolicy) -> usize {
195    match policy {
196        ExecutionPolicy::Sequential => requested.min(1),
197        ExecutionPolicy::AmbientRayon => requested,
198        ExecutionPolicy::Rayon { max_threads } => requested.min(max_threads.get()),
199    }
200}
201
202#[cfg(feature = "parallel")]
203pub(crate) fn permutation_copy_parallel_eligible(
204    policy: ExecutionPolicy,
205    fanout_active: bool,
206    current_pool_threads: usize,
207) -> bool {
208    if fanout_active || current_pool_threads <= 1 {
209        return false;
210    }
211    match policy {
212        ExecutionPolicy::AmbientRayon => true,
213        ExecutionPolicy::Sequential => false,
214        ExecutionPolicy::Rayon { max_threads } => current_pool_threads <= max_threads.get(),
215    }
216}
217
218#[cfg(feature = "parallel")]
219pub fn rayon_threads() -> usize {
220    if fanout_active() {
221        return 1;
222    }
223    let target_width = crate::threading::target_pool_width();
224    match ACTIVE_POLICY.with(Cell::get) {
225        ExecutionPolicy::Sequential => 1,
226        ExecutionPolicy::AmbientRayon => {
227            target_width.unwrap_or_else(crate::threading::current_pool_threads)
228        }
229        ExecutionPolicy::Rayon { max_threads } => target_width
230            .unwrap_or_else(crate::threading::current_pool_threads)
231            .min(max_threads.get()),
232    }
233}
234
235#[cfg(test)]
236#[path = "execution_policy/tests/default_tests.rs"]
237mod default_tests;
238
239#[cfg(all(test, feature = "parallel"))]
240#[path = "execution_policy/tests/tests.rs"]
241mod tests;
strided_basic/execution_policy.rs

strided_basic/
execution_policy.rs
