Add spawn_fifo

Author: cuviper

rayon-core/src/lib.rs

@@ -63,7 +63,7 @@ pub mod internal;
 pub use join::{join, join_context};
 pub use scope::{scope, Scope};
 pub use scope::{scope_fifo, ScopeFifo};
-pub use spawn::spawn;
+pub use spawn::{spawn, spawn_fifo};
 pub use thread_pool::current_thread_has_pending_tasks;
 pub use thread_pool::current_thread_index;
 pub use thread_pool::ThreadPool;

rayon-core/src/registry.rs

@@ -4,7 +4,7 @@ use crossbeam_deque::{Deque, Steal, Stealer};
 use internal::task::Task;
 #[cfg(rayon_unstable)]
 use job::Job;
-use job::{JobRef, StackJob};
+use job::{JobFifo, JobRef, StackJob};
 use latch::{CountLatch, Latch, LatchProbe, LockLatch, SpinLatch, TickleLatch};
 use log::Event::*;
 use sleep::Sleep;
@@ -467,6 +467,9 @@ pub struct WorkerThread {
     /// the "worker" half of our local deque
     worker: Deque<JobRef>,
 
+    /// local queue used for `spawn_fifo` indirection
+    fifo: JobFifo,
+
     index: usize,
 
     /// are these workers configured to steal breadth-first or not?
@@ -523,6 +526,11 @@ impl WorkerThread {
         self.registry.sleep.tickle(self.index);
     }
 
+    #[inline]
+    pub unsafe fn push_fifo(&self, job: JobRef) {
+        self.push(self.fifo.push(job));
+    }
+
     #[inline]
     pub fn local_deque_is_empty(&self) -> bool {
         self.worker.len() == 0
@@ -652,6 +660,7 @@ unsafe fn main_loop(
 ) {
     let worker_thread = WorkerThread {
         worker: worker,
+        fifo: JobFifo::new(),
         breadth_first: breadth_first,
         index: index,
         rng: XorShift64Star::new(),

rayon-core/src/spawn/mod.rs

@@ -61,14 +61,29 @@ where
 ///
 /// Not a public API, but used elsewhere in Rayon.
 pub unsafe fn spawn_in<F>(func: F, registry: &Arc<Registry>)
+where
+    F: FnOnce() + Send + 'static,
+{
+    // We assert that this does not hold any references (we know
+    // this because of the `'static` bound in the inferface);
+    // moreover, we assert that the code below is not supposed to
+    // be able to panic, and hence the data won't leak but will be
+    // enqueued into some deque for later execution.
+    let abort_guard = unwind::AbortIfPanic; // just in case we are wrong, and code CAN panic
+    let job_ref = spawn_job(func, registry);
+    registry.inject_or_push(job_ref);
+    mem::forget(abort_guard);
+}
+
+unsafe fn spawn_job<F>(func: F, registry: &Arc<Registry>) -> JobRef
 where
     F: FnOnce() + Send + 'static,
 {
     // Ensure that registry cannot terminate until this job has
     // executed. This ref is decremented at the (*) below.
     registry.increment_terminate_count();
 
-    let async_job = Box::new(HeapJob::new({
+    Box::new(HeapJob::new({
         let registry = registry.clone();
         move || {
             match unwind::halt_unwinding(func) {
@@ -79,16 +94,42 @@ where
             }
             registry.terminate(); // (*) permit registry to terminate now
         }
-    }));
+    }))
+    .as_job_ref()
+}
 
+/// TODO: like `spawn`, but FIFO
+pub fn spawn_fifo<F>(func: F)
+where
+    F: FnOnce() + Send + 'static,
+{
+    // We assert that current registry has not terminated.
+    unsafe { spawn_fifo_in(func, &Registry::current()) }
+}
+
+/// Spawn an asynchronous FIFO job in `registry.`
+///
+/// Unsafe because `registry` must not yet have terminated.
+///
+/// Not a public API, but used elsewhere in Rayon.
+pub unsafe fn spawn_fifo_in<F>(func: F, registry: &Arc<Registry>)
+where
+    F: FnOnce() + Send + 'static,
+{
     // We assert that this does not hold any references (we know
     // this because of the `'static` bound in the inferface);
     // moreover, we assert that the code below is not supposed to
     // be able to panic, and hence the data won't leak but will be
     // enqueued into some deque for later execution.
     let abort_guard = unwind::AbortIfPanic; // just in case we are wrong, and code CAN panic
-    let job_ref = HeapJob::as_job_ref(async_job);
-    registry.inject_or_push(job_ref);
+    let job_ref = spawn_job(func, registry);
+
+    // If we're in the pool, use our thread's private fifo for this thread to execute
+    // in a locally-FIFO order.  Otherwise, just use the pool's global injector.
+    match registry.current_thread() {
+        Some(worker) => worker.push_fifo(job_ref),
+        None => registry.inject(&[job_ref]),
+    }
     mem::forget(abort_guard);
 }
 

rayon-core/src/spawn/test.rs

@@ -3,7 +3,7 @@ use std::any::Any;
 use std::sync::mpsc::channel;
 use std::sync::Mutex;
 
-use super::spawn;
+use super::{spawn, spawn_fifo};
 use ThreadPoolBuilder;
 
 #[test]
@@ -141,3 +141,62 @@ fn custom_panic_handler_and_nested_spawn() {
         }
     }
 }
+
+macro_rules! test_order {
+    ($outer_spawn:ident, $inner_spawn:ident) => {{
+        let builder = ThreadPoolBuilder::new().num_threads(1);
+        let pool = builder.build().unwrap();
+        let (tx, rx) = channel();
+        pool.install(move || {
+            for i in 0..10 {
+                let tx = tx.clone();
+                $outer_spawn(move || {
+                    for j in 0..10 {
+                        let tx = tx.clone();
+                        $inner_spawn(move || {
+                            tx.send(i * 10 + j).unwrap();
+                        });
+                    }
+                });
+            }
+        });
+        rx.iter().collect::<Vec<i32>>()
+    }};
+}
+
+#[test]
+fn lifo_order() {
+    // In the absense of stealing, `spawn()` jobs on a thread will run in LIFO order.
+    let vec = test_order!(spawn, spawn);
+    let expected: Vec<i32> = (0..100).rev().collect(); // LIFO -> reversed
+    assert_eq!(vec, expected);
+}
+
+#[test]
+fn fifo_order() {
+    // In the absense of stealing, `spawn_fifo()` jobs on a thread will run in FIFO order.
+    let vec = test_order!(spawn_fifo, spawn_fifo);
+    let expected: Vec<i32> = (0..100).collect(); // FIFO -> natural order
+    assert_eq!(vec, expected);
+}
+
+#[test]
+fn lifo_fifo_order() {
+    // LIFO on the outside, FIFO on the inside
+    let vec = test_order!(spawn, spawn_fifo);
+    let expected: Vec<i32> = (0..10)
+        .rev()
+        .flat_map(|i| (0..10).map(move |j| i * 10 + j))
+        .collect();
+    assert_eq!(vec, expected);
+}
+
+#[test]
+fn fifo_lifo_order() {
+    // FIFO on the outside, LIFO on the inside
+    let vec = test_order!(spawn_fifo, spawn);
+    let expected: Vec<i32> = (0..10)
+        .flat_map(|i| (0..10).rev().map(move |j| i * 10 + j))
+        .collect();
+    assert_eq!(vec, expected);
+}

rayon-core/src/thread_pool/mod.rs

@@ -269,6 +269,15 @@ impl ThreadPool {
         // We assert that `self.registry` has not terminated.
         unsafe { spawn::spawn_in(op, &self.registry) }
     }
+
+    /// TODO: like `spawn`, but FIFO
+    pub fn spawn_fifo<OP>(&self, op: OP)
+    where
+        OP: FnOnce() + Send + 'static,
+    {
+        // We assert that `self.registry` has not terminated.
+        unsafe { spawn::spawn_in(op, &self.registry) }
+    }
 }
 
 impl Drop for ThreadPool {