feat: add rayon-compat crate

NthTensor · NthTensor · commit 3b682bbec64b · 2025-06-27T17:53:45.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,8 +8,7 @@ repository = "https://github.com/NthTensor/Forte"
 
 [workspace]
 resolver = "2"
-members = ["ci"]
-exclude = ["coz"]
+members = ["ci", "rayon-compat"]
 
 [dependencies]
 async-task = "4.7.1"
diff --git a/rayon-compat/Cargo.toml b/rayon-compat/Cargo.toml
@@ -0,0 +1,10 @@
+[package]
+name = "rayon-compat"
+version = "1.12.1"
+edition = "2024"
+
+[dependencies]
+forte = { path = ".." }
+
+[features]
+web_spin_lock = []
diff --git a/rayon-compat/README.md b/rayon-compat/README.md
@@ -0,0 +1,14 @@
+# Rayon Compat
+
+This is a way to run `rayon` on top of `forte`! The `rayon-compat` crate mocks the important bits of the api of `rayon_core` in a pretty simple and crude way, which is none-the-less enough to support most of what `rayon` needs.
+
+To use this crate, apply the following cargo patch like one of these:
+```
+// If you want to clone forte and use it locally
+[patch.crates-io]
+rayon-core = { path = "path to this repo", package = "rayon-compat" }
+
+// If you want to use the latest published version of forte
+[patch.crates-io]
+rayon-core = { path = "https://github.com/NthTensor/Forte", package = "rayon-compat" }
+```
diff --git a/rayon-compat/src/lib.rs b/rayon-compat/src/lib.rs
@@ -0,0 +1,169 @@
+use std::{
+    marker::PhantomData,
+    num::NonZero,
+    sync::atomic::{AtomicBool, Ordering},
+};
+
+pub static THREAD_POOL: forte::ThreadPool = forte::ThreadPool::new();
+
+pub static STARTED: AtomicBool = AtomicBool::new(false);
+
+#[inline]
+fn ensure_started() {
+    if !STARTED.load(Ordering::Relaxed) {
+        if !STARTED.swap(true, Ordering::Relaxed) {
+            THREAD_POOL.resize_to_available();
+        }
+    }
+}
+
+// -----------------------------------------------------------------------------
+// Join
+
+#[derive(Debug)]
+pub struct FnContext {
+    /// True if the task was migrated.
+    migrated: bool,
+}
+
+impl FnContext {
+    #[inline]
+    pub fn migrated(&self) -> bool {
+        self.migrated
+    }
+}
+
+pub fn join_context<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
+where
+    A: FnOnce(FnContext) -> RA + Send,
+    B: FnOnce(FnContext) -> RB + Send,
+    RA: Send,
+    RB: Send,
+{
+    ensure_started();
+    THREAD_POOL.join(
+        |worker| {
+            let ctx = FnContext {
+                migrated: worker.migrated(),
+            };
+            oper_a(ctx)
+        },
+        |worker| {
+            let ctx = FnContext {
+                migrated: worker.migrated(),
+            };
+            oper_b(ctx)
+        },
+    )
+}
+
+pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
+where
+    A: FnOnce() -> RA + Send,
+    B: FnOnce() -> RB + Send,
+    RA: Send,
+    RB: Send,
+{
+    ensure_started();
+    THREAD_POOL.join(|_| oper_a(), |_| oper_b())
+}
+
+pub fn current_num_threads() -> usize {
+    std::thread::available_parallelism()
+        .map(NonZero::get)
+        .unwrap_or(1)
+}
+
+pub fn current_thread_index() -> Option<usize> {
+    forte::Worker::map_current(|worker| worker.index())
+}
+
+// -----------------------------------------------------------------------------
+// Scope
+
+pub struct Scope<'scope> {
+    phantom: PhantomData<&'scope ()>,
+}
+
+impl<'scope> Scope<'scope> {
+    pub fn spawn<BODY>(&self, body: BODY)
+    where
+        BODY: FnOnce(&Scope) + Send + 'scope,
+    {
+        unimplemented!();
+    }
+}
+
+pub fn scope<'scope, OP, R>(op: OP) -> R
+where
+    OP: FnOnce(&Scope<'scope>) -> R + Send,
+    R: Send,
+{
+    unimplemented!();
+}
+
+pub fn in_place_scope<'scope, OP, R>(op: OP) -> R
+where
+    OP: FnOnce(&Scope<'scope>) -> R,
+{
+    unimplemented!();
+}
+
+pub fn in_place_scope_fifo() {
+    unimplemented!()
+}
+
+pub fn scope_fifo() {
+    unimplemented!()
+}
+
+// -----------------------------------------------------------------------------
+// Spawn
+
+pub fn spawn<F>(func: F)
+where
+    F: FnOnce() + Send + 'static,
+{
+    unimplemented!();
+}
+
+// -----------------------------------------------------------------------------
+// Fake stuff that dosn't work
+
+pub struct ThreadBuilder;
+
+pub struct ThreadPool;
+
+pub struct ThreadPoolBuildError;
+
+pub struct ThreadPoolBuilder;
+
+pub struct BroadcastContext;
+
+pub struct ScopeFifo;
+
+pub struct Yield;
+
+pub fn broadcast() {
+    unimplemented!()
+}
+
+pub fn spawn_broadcast() {
+    unimplemented!()
+}
+
+pub fn max_num_threads() {
+    unimplemented!()
+}
+
+pub fn spawn_fifo() {
+    unimplemented!()
+}
+
+pub fn yield_local() {
+    unimplemented!()
+}
+
+pub fn yield_now() {
+    unimplemented!()
+}
diff --git a/src/thread_pool.rs b/src/thread_pool.rs
@@ -647,6 +647,7 @@ thread_local! {
 /// Workers have one core memory-safety guarantee: Any jobs added to the worker
 /// will eventually be executed.
 pub struct Worker {
+    pub(crate) migrated: Cell<bool>,
     pub(crate) lease: Lease,
     pub(crate) queue: JobQueue,
 }
@@ -736,6 +737,7 @@ impl Worker {
         // problem that the same thread can occupy multiple workers on the same
         // thread. We many eventually need to design something to prevent this.
         let worker = Worker {
+            migrated: Cell::new(false),
             lease,
             queue: JobQueue::new(),
         };
@@ -749,7 +751,7 @@ impl Worker {
 
         // Execute the work queue until it's empty
         while let Some(job_ref) = worker.queue.pop_front() {
-            job_ref.execute(&worker);
+            worker.execute(job_ref, false);
         }
 
         // Swap back to pointing to the previous value (possibly null).
@@ -762,6 +764,10 @@ impl Worker {
         result
     }
 
+    pub fn index(&self) -> usize {
+        self.lease.index
+    }
+
     /// Tries to promote the oldest job in the local stack to a shared job. If
     /// the local job queue is empty, or if the shared queue is full, this does
     /// nothing. If the promotion is successful, it tries to wake another
@@ -814,12 +820,15 @@ impl Worker {
 
     /// Tries to find a job to execute, either in the local queue or shared on
     /// the threadpool.
+    ///
+    /// The second value is true if the job was shared, or false if it was spawned locally.
     #[inline]
-    pub fn find_work(&self) -> Option<JobRef> {
+    pub fn find_work(&self) -> Option<(JobRef, bool)> {
         // We give preference first to things in our local deque, then in other
         // workers deques, and finally to injected jobs from the outside. The
         // idea is to finish what we started before we take on something new.
-        self.queue.pop_back().or_else(|| self.claim_shared_job())
+        self.queue.pop_back().map(|job| (job, false))
+            .or_else(|| self.claim_shared_job().map(|job| (job, true)))
     }
 
     /// Claims a shared job from the thread pool.
@@ -842,7 +851,7 @@ impl Worker {
     pub fn yield_local(&self) -> Yield {
         match self.queue.pop_back() {
             Some(job_ref) => {
-                job_ref.execute(self);
+                self.execute(job_ref, false);
                 Yield::Executed
             }
             None => Yield::Idle,
@@ -860,13 +869,29 @@ impl Worker {
     #[inline]
     pub fn yield_now(&self) -> Yield {
         match self.find_work() {
-            Some(job_ref) => {
-                job_ref.execute(self);
+            Some((job_ref, migrated)) => {
+                self.execute(job_ref, migrated);
                 Yield::Executed
             }
             None => Yield::Idle,
         }
     }
+
+    /// Returns `true` if the current job is executing on a different thread
+    /// from the one on which it was created. Returns `false` if not executing a
+    /// job, or if the current job was created on the current thread.
+    #[inline]
+    pub fn migrated(&self) -> bool {
+        self.migrated.get()
+    }
+
+    /// Executes a job on a worker
+    #[inline]
+    pub fn execute(&self, job_ref: JobRef, migrated: bool) {
+        let migrated = self.migrated.replace(migrated);
+        job_ref.execute(self);
+        self.migrated.set(migrated);
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -1063,7 +1088,7 @@ impl Worker {
 
             // Even if it's not the droid we were looking for, we must still
             // execute the job.
-            job.execute(self);
+            self.execute(job, false);
         }
 
         // Wait for the job to complete, then return the result.
@@ -1217,7 +1242,7 @@ fn managed_worker(lease: Lease, halt: Arc<AtomicBool>, barrier: Arc<Barrier>) {
     Worker::occupy(lease, |worker| {
         while !halt.load(Ordering::Relaxed) {
             if let Some(job) = worker.queue.pop_back() {
-                job.execute(worker);
+                worker.execute(job, false);
                 continue;
             }
 
@@ -1226,7 +1251,7 @@ fn managed_worker(lease: Lease, halt: Arc<AtomicBool>, barrier: Arc<Barrier>) {
             while !halt.load(Ordering::Relaxed) {
                 if let Some(job) = state.claim_shared_job() {
                     drop(state);
-                    job.execute(worker);
+                    worker.execute(job, true);
                     break;
                 }
 
