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

Zoxc · Zoxc · commit 7fdf1fd47b5a · 2023-08-27T19:39:57.000+02:00
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());
diff --git a/rayon-core/src/scope/mod.rs b/rayon-core/src/scope/mod.rs
diff --git a/rayon-core/src/spawn/mod.rs b/rayon-core/src/spawn/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -117,7 +117,7 @@ where`
`117`	`117`	`registry.inject_broadcast(job_refs);`
`118`	`118`
`119`	`119`	`// Wait for all jobs to complete, then collect the results, maybe propagating a panic.`
`120`		`- latch.wait(current_thread);`
	`120`	`+ latch.wait(current_thread, None);`
`121`	`121`	`jobs.into_iter().map(\|job\| job.into_result()).collect()`
`122`	`122`	`}`
`123`	`123`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,11 @@ pub(super) trait Job {`
`26`	`26`	`unsafe fn execute(this: *const ());`
`27`	`27`	`}`
`28`	`28`
	`29`	`+#[derive(PartialEq, Eq, Hash, Copy, Clone)]`
	`30`	`+pub(super) struct JobRefId {`
	`31`	`+ pointer: usize,`
	`32`	`+}`
	`33`	`+`
`29`	`34`	`/// Effectively a Job trait object. Each JobRef must be executed`
`30`	`35`	`/// exactly once, or else data may leak.`
`31`	`36`	`///`
`@@ -54,11 +59,11 @@ impl JobRef {`
`54`	`59`	`}`
`55`	`60`	`}`
`56`	`61`
`57`		`- /// Returns an opaque handle that can be saved and compared,`
`58`		- /// without making `JobRef` itself `Copy + Eq`.
`59`	`62`	`#[inline]`
`60`		`- pub(super) fn id(&self) -> impl Eq {`
`61`		`- (self.pointer, self.execute_fn)`
	`63`	`+ pub(super) fn id(&self) -> JobRefId {`
	`64`	`+ JobRefId {`
	`65`	`+ pointer: self.pointer as usize,`
	`66`	`+ }`
`62`	`67`	`}`
`63`	`68`
`64`	`69`	`#[inline]`
`@@ -102,10 +107,6 @@ where`
`102`	`107`	`JobRef::new(self)`
`103`	`108`	`}`
`104`	`109`
`105`		`- pub(super) unsafe fn run_inline(self, stolen: bool) -> R {`
`106`		`- self.func.into_inner().unwrap()(stolen)`
`107`		`- }`
`108`		`-`
`109`	`110`	`pub(super) unsafe fn into_result(self) -> R {`
`110`	`111`	`self.result.into_inner().into_return_value()`
`111`	`112`	`}`
`@@ -136,15 +137,15 @@ where`
`136`	`137`	/// (Probably `StackJob` should be refactored in a similar fashion.)
`137`	`138`	`pub(super) struct HeapJob<BODY>`
`138`	`139`	`where`
`139`		`- BODY: FnOnce() + Send,`
	`140`	`+ BODY: FnOnce(JobRefId) + Send,`
`140`	`141`	`{`
`141`	`142`	`job: BODY,`
`142`	`143`	`tlv: Tlv,`
`143`	`144`	`}`
`144`	`145`
`145`	`146`	`impl<BODY> HeapJob<BODY>`
`146`	`147`	`where`
`147`		`- BODY: FnOnce() + Send,`
	`148`	`+ BODY: FnOnce(JobRefId) + Send,`
`148`	`149`	`{`
`149`	`150`	`pub(super) fn new(tlv: Tlv, job: BODY) -> Box<Self> {`
`150`	`151`	`Box::new(HeapJob { job, tlv })`
`@@ -168,12 +169,13 @@ where`
`168`	`169`
`169`	`170`	`impl<BODY> Job for HeapJob<BODY>`
`170`	`171`	`where`
`171`		`- BODY: FnOnce() + Send,`
	`172`	`+ BODY: FnOnce(JobRefId) + Send,`
`172`	`173`	`{`
`173`	`174`	`unsafe fn execute(this: *const ()) {`
	`175`	`+ let pointer = this as usize;`
`174`	`176`	`let this = Box::from_raw(this as *mut Self);`
`175`	`177`	`tlv::set(this.tlv);`
`176`		`- (this.job)();`
	`178`	`+ (this.job)(JobRefId { pointer });`
`177`	`179`	`}`
`178`	`180`	`}`
`179`	`181`
Original file line number	Diff line number	Diff line change
`@@ -177,11 +177,6 @@ impl<'r> SpinLatch<'r> {`
`177`	`177`	`..SpinLatch::new(thread)`
`178`	`178`	`}`
`179`	`179`	`}`
`180`		`-`
`181`		`- #[inline]`
`182`		`- pub(super) fn probe(&self) -> bool {`
`183`		`- self.core_latch.probe()`
`184`		`- }`
`185`	`180`	`}`
`186`	`181`
`187`	`182`	`impl<'r> AsCoreLatch for SpinLatch<'r> {`