Only work-steal in the main loop for join and scope

rust-lang · Aug 27, 2023 · 7fdf1fd · 7fdf1fd
1 parent f192a48
commit 7fdf1fd
Show file tree

Hide file tree

Showing 8 changed files with 127 additions and 89 deletions.
diff --git a/rayon-core/Cargo.toml b/rayon-core/Cargo.toml
@@ -22,6 +22,7 @@ num_cpus = "1.2"
 crossbeam-channel = "0.5.0"
 crossbeam-deque = "0.8.1"
 crossbeam-utils = "0.8.0"
+smallvec = "1.11.0"
 
 [dev-dependencies]
 rand = "0.8"

diff --git a/rayon-core/src/broadcast/mod.rs b/rayon-core/src/broadcast/mod.rs
@@ -117,7 +117,7 @@ where
     registry.inject_broadcast(job_refs);
 
     // Wait for all jobs to complete, then collect the results, maybe propagating a panic.
-    latch.wait(current_thread);
+    latch.wait(current_thread, None);
     jobs.into_iter().map(|job| job.into_result()).collect()
 }
 

diff --git a/rayon-core/src/job.rs b/rayon-core/src/job.rs
@@ -26,6 +26,11 @@ pub(super) trait Job {
     unsafe fn execute(this: *const ());
 }
 
+#[derive(PartialEq, Eq, Hash, Copy, Clone)]
+pub(super) struct JobRefId {
+    pointer: usize,
+}
+
 /// Effectively a Job trait object. Each JobRef **must** be executed
 /// exactly once, or else data may leak.
 ///
@@ -54,11 +59,11 @@ impl JobRef {
         }
     }
 
-    /// Returns an opaque handle that can be saved and compared,
-    /// without making `JobRef` itself `Copy + Eq`.
     #[inline]
-    pub(super) fn id(&self) -> impl Eq {
-        (self.pointer, self.execute_fn)
+    pub(super) fn id(&self) -> JobRefId {
+        JobRefId {
+            pointer: self.pointer as usize,
+        }
     }
 
     #[inline]
@@ -102,10 +107,6 @@ where
         JobRef::new(self)
     }
 
-    pub(super) unsafe fn run_inline(self, stolen: bool) -> R {
-        self.func.into_inner().unwrap()(stolen)
-    }
-
     pub(super) unsafe fn into_result(self) -> R {
         self.result.into_inner().into_return_value()
     }
@@ -136,15 +137,15 @@ where
 /// (Probably `StackJob` should be refactored in a similar fashion.)
 pub(super) struct HeapJob<BODY>
 where
-    BODY: FnOnce() + Send,
+    BODY: FnOnce(JobRefId) + Send,
 {
     job: BODY,
     tlv: Tlv,
 }
 
 impl<BODY> HeapJob<BODY>
 where
-    BODY: FnOnce() + Send,
+    BODY: FnOnce(JobRefId) + Send,
 {
     pub(super) fn new(tlv: Tlv, job: BODY) -> Box<Self> {
         Box::new(HeapJob { job, tlv })
@@ -168,12 +169,13 @@ where
 
 impl<BODY> Job for HeapJob<BODY>
 where
-    BODY: FnOnce() + Send,
+    BODY: FnOnce(JobRefId) + Send,
 {
     unsafe fn execute(this: *const ()) {
+        let pointer = this as usize;
         let this = Box::from_raw(this as *mut Self);
         tlv::set(this.tlv);
-        (this.job)();
+        (this.job)(JobRefId { pointer });
     }
 }
 

diff --git a/rayon-core/src/join/mod.rs b/rayon-core/src/join/mod.rs
@@ -1,9 +1,9 @@
 use crate::job::StackJob;
 use crate::latch::SpinLatch;
-use crate::registry::{self, WorkerThread};
-use crate::tlv::{self, Tlv};
+use crate::registry;
+use crate::tlv;
 use crate::unwind;
-use std::any::Any;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use crate::FnContext;
 
@@ -135,68 +135,37 @@ where
         // Create virtual wrapper for task b; this all has to be
         // done here so that the stack frame can keep it all live
         // long enough.
-        let job_b = StackJob::new(tlv, call_b(oper_b), SpinLatch::new(worker_thread));
+        let job_b_started = AtomicBool::new(false);
+        let job_b = StackJob::new(
+            tlv,
+            |migrated| {
+                job_b_started.store(true, Ordering::Relaxed);
+                call_b(oper_b)(migrated)
+            },
+            SpinLatch::new(worker_thread),
+        );
         let job_b_ref = job_b.as_job_ref();
         let job_b_id = job_b_ref.id();
         worker_thread.push(job_b_ref);
 
         // Execute task a; hopefully b gets stolen in the meantime.
         let status_a = unwind::halt_unwinding(call_a(oper_a, injected));
-        let result_a = match status_a {
-            Ok(v) => v,
-            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err, tlv),
-        };
-
-        // Now that task A has finished, try to pop job B from the
-        // local stack.  It may already have been popped by job A; it
-        // may also have been stolen. There may also be some tasks
-        // pushed on top of it in the stack, and we will have to pop
-        // those off to get to it.
-        while !job_b.latch.probe() {
-            if let Some(job) = worker_thread.take_local_job() {
-                if job_b_id == job.id() {
-                    // Found it! Let's run it.
-                    //
-                    // Note that this could panic, but it's ok if we unwind here.
 
-                    // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
-                    tlv::set(tlv);
-
-                    let result_b = job_b.run_inline(injected);
-                    return (result_a, result_b);
-                } else {
-                    worker_thread.execute(job);
-                }
-            } else {
-                // Local deque is empty. Time to steal from other
-                // threads.
-                worker_thread.wait_until(&job_b.latch);
-                debug_assert!(job_b.latch.probe());
-                break;
-            }
-        }
+        // Wait for job B or execute it if it's in the local queue.
+        worker_thread.wait_for_jobs(
+            &job_b.latch,
+            || job_b_started.load(Ordering::Relaxed),
+            |job| job.id() == job_b_id,
+        );
 
         // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
         tlv::set(tlv);
 
+        let result_a = match status_a {
+            Ok(v) => v,
+            Err(err) => unwind::resume_unwinding(err),
+        };
+
         (result_a, job_b.into_result())
     })
 }
-
-/// If job A panics, we still cannot return until we are sure that job
-/// B is complete. This is because it may contain references into the
-/// enclosing stack frame(s).
-#[cold] // cold path
-unsafe fn join_recover_from_panic(
-    worker_thread: &WorkerThread,
-    job_b_latch: &SpinLatch<'_>,
-    err: Box<dyn Any + Send>,
-    tlv: Tlv,
-) -> ! {
-    worker_thread.wait_until(job_b_latch);
-
-    // Restore the TLV since we might have run some jobs overwriting it when waiting for job b.
-    tlv::set(tlv);
-
-    unwind::resume_unwinding(err)
-}
diff --git a/rayon-core/src/latch.rs b/rayon-core/src/latch.rs
@@ -177,11 +177,6 @@ impl<'r> SpinLatch<'r> {
             ..SpinLatch::new(thread)
         }
     }
-
-    #[inline]
-    pub(super) fn probe(&self) -> bool {
-        self.core_latch.probe()
-    }
 }
 
 impl<'r> AsCoreLatch for SpinLatch<'r> {

diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs
@@ -10,6 +10,7 @@ use crate::{
     ReleaseThreadHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder, Yield,
 };
 use crossbeam_deque::{Injector, Steal, Stealer, Worker};
+use smallvec::SmallVec;
 use std::cell::Cell;
 use std::collections::hash_map::DefaultHasher;
 use std::fmt;
@@ -840,14 +841,58 @@ impl WorkerThread {
     /// stealing tasks as necessary.
     #[inline]
     pub(super) unsafe fn wait_until<L: AsCoreLatch + ?Sized>(&self, latch: &L) {
+        self.wait_or_steal_until(latch, false)
+    }
+
+    /// Wait until the latch is set. Executes local jobs if `is_job` is true for them and
+    /// `all_jobs_started` still returns false.
+    #[inline]
+    pub(super) unsafe fn wait_for_jobs<L: AsCoreLatch + ?Sized>(
+        &self,
+        latch: &L,
+        mut all_jobs_started: impl FnMut() -> bool,
+        mut is_job: impl FnMut(&JobRef) -> bool,
+    ) {
+        let mut jobs = SmallVec::<[JobRef; 8]>::new();
+
+        // Make sure all jobs have started.
+        while !all_jobs_started() {
+            if let Some(job) = self.worker.pop() {
+                if is_job(&job) {
+                    // Found a job, let's run it.
+                    self.execute(job);
+                } else {
+                    jobs.push(job);
+                }
+            } else {
+                break;
+            }
+        }
+
+        // Restore the jobs that we weren't looking for.
+        for job in jobs.into_iter().rev() {
+            self.worker.push(job);
+        }
+
+        // Wait for the jobs to finish.
+        self.wait_until(latch);
+        debug_assert!(latch.as_core_latch().probe());
+    }
+
+    #[inline]
+    pub(super) unsafe fn wait_or_steal_until<L: AsCoreLatch + ?Sized>(
+        &self,
+        latch: &L,
+        steal: bool,
+    ) {
         let latch = latch.as_core_latch();
         if !latch.probe() {
-            self.wait_until_cold(latch);
+            self.wait_until_cold(latch, steal);
         }
     }
 
     #[cold]
-    unsafe fn wait_until_cold(&self, latch: &CoreLatch) {
+    unsafe fn wait_until_cold(&self, latch: &CoreLatch, steal: bool) {
         // the code below should swallow all panics and hence never
         // unwind; but if something does wrong, we want to abort,
         // because otherwise other code in rayon may assume that the
@@ -857,10 +902,16 @@ impl WorkerThread {
 
         let mut idle_state = self.registry.sleep.start_looking(self.index, latch);
         while !latch.probe() {
-            if let Some(job) = self.find_work() {
-                self.registry.sleep.work_found(idle_state);
-                self.execute(job);
-                idle_state = self.registry.sleep.start_looking(self.index, latch);
+            if steal {
+                if let Some(job) = self.find_work() {
+                    self.registry.sleep.work_found(idle_state);
+                    self.execute(job);
+                    idle_state = self.registry.sleep.start_looking(self.index, latch);
+                } else {
+                    self.registry
+                        .sleep
+                        .no_work_found(&mut idle_state, latch, &self)
+                }
             } else {
                 self.registry
                     .sleep
@@ -988,7 +1039,7 @@ unsafe fn main_loop(thread: ThreadBuilder) {
         terminate_addr: my_terminate_latch.as_core_latch().addr(),
     });
     registry.acquire_thread();
-    worker_thread.wait_until(my_terminate_latch);
+    worker_thread.wait_or_steal_until(my_terminate_latch, true);
 
     // Should not be any work left in our queue.
     debug_assert!(worker_thread.take_local_job().is_none());