forte_rayon_compat/
lib.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2
3pub static THREAD_POOL: forte::ThreadPool = const { forte::ThreadPool::new() };
4
5pub static STARTED: AtomicBool = const { AtomicBool::new(false) };
6
7#[inline(always)]
8fn ensure_started() {
9    if !STARTED.load(Ordering::Relaxed) && !STARTED.swap(true, Ordering::Relaxed) {
10        THREAD_POOL.resize_to_available();
11    }
12}
13
14#[inline(always)]
15pub fn current_num_threads() -> usize {
16    64 // Forte prefers smaller tasks, so it's better to lie to rayon about the size of the pool
17}
18
19#[inline(always)]
20pub fn current_thread_index() -> Option<usize> {
21    forte::Worker::map_current(|worker| worker.index())
22}
23
24#[inline(always)]
25pub fn max_num_threads() -> usize {
26    usize::MAX // The number of forte workers is only bounded by the size of a vector.
27}
28
29// -----------------------------------------------------------------------------
30// Join
31
32#[derive(Debug)]
33pub struct FnContext {
34    /// True if the task was migrated.
35    migrated: bool,
36}
37
38impl FnContext {
39    #[inline(always)]
40    pub fn migrated(&self) -> bool {
41        self.migrated
42    }
43}
44
45#[inline(always)]
46pub fn join_context<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
47where
48    A: FnOnce(FnContext) -> RA + Send,
49    B: FnOnce(FnContext) -> RB + Send,
50    RA: Send,
51    RB: Send,
52{
53    ensure_started();
54    THREAD_POOL.join(
55        |worker| {
56            let migrated = worker.migrated();
57            let ctx = FnContext { migrated };
58            oper_a(ctx)
59        },
60        |worker| {
61            let migrated = worker.migrated();
62            let ctx = FnContext { migrated };
63            oper_b(ctx)
64        },
65    )
66}
67
68#[inline(always)]
69pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
70where
71    A: FnOnce() -> RA + Send,
72    B: FnOnce() -> RB + Send,
73    RA: Send,
74    RB: Send,
75{
76    ensure_started();
77    THREAD_POOL.join(|_| oper_a(), |_| oper_b())
78}
79
80// -----------------------------------------------------------------------------
81// Scope
82
83pub use forte::Scope;
84
85#[inline(always)]
86pub fn scope<'scope, OP, R>(op: OP) -> R
87where
88    OP: FnOnce(&Scope<'scope>) -> R + Send,
89    R: Send,
90{
91    ensure_started();
92    forte::scope(op)
93}
94
95#[inline(always)]
96pub fn in_place_scope<'scope, OP, R>(op: OP) -> R
97where
98    OP: FnOnce(&Scope<'scope>) -> R,
99{
100    ensure_started();
101    forte::scope(op)
102}
103
104// -----------------------------------------------------------------------------
105// Spawn
106
107#[inline(always)]
108pub fn spawn<F>(func: F)
109where
110    F: FnOnce() + Send + 'static,
111{
112    ensure_started();
113    THREAD_POOL.spawn(|_| func())
114}
115
116// -----------------------------------------------------------------------------
117// Yield
118
119pub use forte::Yield;
120
121pub fn yield_local() -> Yield {
122    let result = forte::Worker::map_current(forte::Worker::yield_local);
123    match result {
124        Some(status) => status,
125        _ => Yield::Idle,
126    }
127}
128
129pub fn yield_now() -> Yield {
130    let result = forte::Worker::map_current(forte::Worker::yield_now);
131    match result {
132        Some(status) => status,
133        _ => Yield::Idle,
134    }
135}
136
137// -----------------------------------------------------------------------------
138// Fake stuff that dosn't work. These are here only so so that rayon can export
139// them.
140
141pub struct ThreadBuilder;
142
143pub struct ThreadPool;
144
145pub struct ThreadPoolBuildError;
146
147pub struct ThreadPoolBuilder;
148
149pub struct BroadcastContext;
150
151pub struct ScopeFifo;
152
153pub fn broadcast() {
154    unimplemented!()
155}
156
157pub fn spawn_broadcast() {
158    unimplemented!()
159}
160
161pub fn scope_fifo() {
162    unimplemented!()
163}
164
165pub fn in_place_scope_fifo() {
166    unimplemented!()
167}
168
169pub fn spawn_fifo() {
170    unimplemented!()
171}
forte_rayon_compat/lib.rs

forte_rayon_compat/
lib.rs
