Run linking and incremental saving / finalizing in parallel

Author: Zoxc

Cargo.lock

@@ -3545,6 +3545,7 @@ dependencies = [
  "portable-atomic",
  "rustc-hash 2.1.1",
  "rustc-rayon",
+ "rustc-rayon-core",
  "rustc-stable-hash",
  "rustc_arena",
  "rustc_graphviz",

compiler/rustc_codegen_cranelift/src/lib.rs

@@ -7,6 +7,7 @@
 // Note: please avoid adding other feature gates where possible
 #![feature(rustc_private)]
 // Note: please avoid adding other feature gates where possible
+#![recursion_limit = "256"]
 #![warn(rust_2018_idioms)]
 #![warn(unreachable_pub)]
 #![warn(unused_lifetimes)]
@@ -43,6 +44,7 @@ use cranelift_codegen::isa::TargetIsa;
 use cranelift_codegen::settings::{self, Configurable};
 use rustc_codegen_ssa::CodegenResults;
 use rustc_codegen_ssa::traits::CodegenBackend;
+use rustc_data_structures::sync::{DynSend, downcast_box_any_dyn_send};
 use rustc_metadata::EncodedMetadata;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_session::Session;
@@ -207,7 +209,7 @@ impl CodegenBackend for CraneliftCodegenBackend {
         tcx: TyCtxt<'_>,
         metadata: EncodedMetadata,
         need_metadata_module: bool,
-    ) -> Box<dyn Any> {
+    ) -> Box<dyn Any + DynSend> {
         info!("codegen crate {}", tcx.crate_name(LOCAL_CRATE));
         let config = self.config.clone().unwrap_or_else(|| {
             BackendConfig::from_opts(&tcx.sess.opts.cg.llvm_args)
@@ -226,11 +228,13 @@ impl CodegenBackend for CraneliftCodegenBackend {
 
     fn join_codegen(
         &self,
-        ongoing_codegen: Box<dyn Any>,
+        ongoing_codegen: Box<dyn Any + DynSend>,
         sess: &Session,
         outputs: &OutputFilenames,
     ) -> (CodegenResults, FxIndexMap<WorkProductId, WorkProduct>) {
-        ongoing_codegen.downcast::<driver::aot::OngoingCodegen>().unwrap().join(sess, outputs)
+        downcast_box_any_dyn_send::<driver::aot::OngoingCodegen>(ongoing_codegen)
+            .unwrap()
+            .join(sess, outputs)
     }
 }
 

compiler/rustc_codegen_gcc/src/lib.rs

@@ -104,7 +104,7 @@ use rustc_codegen_ssa::base::codegen_crate;
 use rustc_codegen_ssa::traits::{CodegenBackend, ExtraBackendMethods, WriteBackendMethods};
 use rustc_codegen_ssa::{CodegenResults, CompiledModule, ModuleCodegen};
 use rustc_data_structures::fx::FxIndexMap;
-use rustc_data_structures::sync::IntoDynSyncSend;
+use rustc_data_structures::sync::{DynSend, IntoDynSyncSend, downcast_box_any_dyn_send};
 use rustc_errors::DiagCtxtHandle;
 use rustc_metadata::EncodedMetadata;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
@@ -230,12 +230,12 @@ impl CodegenBackend for GccCodegenBackend {
         providers.global_backend_features = |tcx, ()| gcc_util::global_gcc_features(tcx.sess, true)
     }
 
-    fn codegen_crate(
+    fn codegen_crate<'tcx>(
         &self,
-        tcx: TyCtxt<'_>,
+        tcx: TyCtxt<'tcx>,
         metadata: EncodedMetadata,
         need_metadata_module: bool,
-    ) -> Box<dyn Any> {
+    ) -> Box<dyn Any + DynSend> {
         let target_cpu = target_cpu(tcx.sess);
         let res = codegen_crate(
             self.clone(),
@@ -250,12 +250,11 @@ impl CodegenBackend for GccCodegenBackend {
 
     fn join_codegen(
         &self,
-        ongoing_codegen: Box<dyn Any>,
+        ongoing_codegen: Box<dyn Any + DynSend>,
         sess: &Session,
         _outputs: &OutputFilenames,
     ) -> (CodegenResults, FxIndexMap<WorkProductId, WorkProduct>) {
-        ongoing_codegen
-            .downcast::<rustc_codegen_ssa::back::write::OngoingCodegen<GccCodegenBackend>>()
+        downcast_box_any_dyn_send::<rustc_codegen_ssa::back::write::OngoingCodegen<GccCodegenBackend>>( ongoing_codegen)
             .expect("Expected GccCodegenBackend's OngoingCodegen, found Box<Any>")
             .join(sess)
     }

compiler/rustc_codegen_llvm/src/lib.rs

@@ -19,6 +19,7 @@
 #![feature(rustdoc_internals)]
 #![feature(slice_as_array)]
 #![feature(try_blocks)]
+#![recursion_limit = "256"]
 // tidy-alphabetical-end
 
 use std::any::Any;
@@ -39,6 +40,7 @@ use rustc_codegen_ssa::back::write::{
 use rustc_codegen_ssa::traits::*;
 use rustc_codegen_ssa::{CodegenResults, CompiledModule, ModuleCodegen};
 use rustc_data_structures::fx::FxIndexMap;
+use rustc_data_structures::sync::{DynSend, downcast_box_any_dyn_send};
 use rustc_errors::{DiagCtxtHandle, FatalError};
 use rustc_metadata::EncodedMetadata;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
@@ -347,7 +349,7 @@ impl CodegenBackend for LlvmCodegenBackend {
         tcx: TyCtxt<'tcx>,
         metadata: EncodedMetadata,
         need_metadata_module: bool,
-    ) -> Box<dyn Any> {
+    ) -> Box<dyn Any + DynSend> {
         Box::new(rustc_codegen_ssa::base::codegen_crate(
             LlvmCodegenBackend(()),
             tcx,
@@ -359,14 +361,15 @@ impl CodegenBackend for LlvmCodegenBackend {
 
     fn join_codegen(
         &self,
-        ongoing_codegen: Box<dyn Any>,
+        ongoing_codegen: Box<dyn Any + DynSend>,
         sess: &Session,
         outputs: &OutputFilenames,
     ) -> (CodegenResults, FxIndexMap<WorkProductId, WorkProduct>) {
-        let (codegen_results, work_products) = ongoing_codegen
-            .downcast::<rustc_codegen_ssa::back::write::OngoingCodegen<LlvmCodegenBackend>>()
-            .expect("Expected LlvmCodegenBackend's OngoingCodegen, found Box<Any>")
-            .join(sess);
+        let (codegen_results, work_products) = downcast_box_any_dyn_send::<
+            rustc_codegen_ssa::back::write::OngoingCodegen<LlvmCodegenBackend>,
+        >(ongoing_codegen)
+        .expect("Expected LlvmCodegenBackend's OngoingCodegen, found Box<dyn Any + DynSend>")
+        .join(sess);
 
         if sess.opts.unstable_opts.llvm_time_trace {
             sess.time("llvm_dump_timing_file", || {

compiler/rustc_codegen_ssa/src/traits/backend.rs

@@ -36,7 +36,7 @@ pub trait BackendTypes {
     type DIVariable: Copy;
 }
 
-pub trait CodegenBackend {
+pub trait CodegenBackend: DynSync + DynSend {
     /// Locale resources for diagnostic messages - a string the content of the Fluent resource.
     /// Called before `init` so that all other functions are able to emit translatable diagnostics.
     fn locale_resource(&self) -> &'static str;
@@ -73,16 +73,16 @@ pub trait CodegenBackend {
         tcx: TyCtxt<'tcx>,
         metadata: EncodedMetadata,
         need_metadata_module: bool,
-    ) -> Box<dyn Any>;
+    ) -> Box<dyn Any + DynSend>;
 
-    /// This is called on the returned `Box<dyn Any>` from [`codegen_crate`](Self::codegen_crate)
+    /// This is called on the returned `Box<dyn Any + DynSend>` from [`codegen_crate`](Self::codegen_crate)
     ///
     /// # Panics
     ///
-    /// Panics when the passed `Box<dyn Any>` was not returned by [`codegen_crate`](Self::codegen_crate).
+    /// Panics when the passed `Box<dyn Any + DynSend>` was not returned by [`codegen_crate`](Self::codegen_crate).
     fn join_codegen(
         &self,
-        ongoing_codegen: Box<dyn Any>,
+        ongoing_codegen: Box<dyn Any + DynSend>,
         sess: &Session,
         outputs: &OutputFilenames,
     ) -> (CodegenResults, FxIndexMap<WorkProductId, WorkProduct>);

compiler/rustc_data_structures/Cargo.toml

@@ -15,6 +15,7 @@ jobserver_crate = { version = "0.1.28", package = "jobserver" }
 measureme = "12.0.1"
 rustc-hash = "2.0.0"
 rustc-rayon = { version = "0.5.1", features = ["indexmap"] }
+rustc-rayon-core = { version = "0.5.0" }
 rustc-stable-hash = { version = "0.1.0", features = ["nightly"] }
 rustc_arena = { path = "../rustc_arena" }
 rustc_graphviz = { path = "../rustc_graphviz" }

compiler/rustc_data_structures/src/marker.rs

@@ -1,4 +1,5 @@
 use std::alloc::Allocator;
+use std::any::Any;
 
 #[rustc_on_unimplemented(message = "`{Self}` doesn't implement `DynSend`. \
             Add it to `rustc_data_structures::marker` or use `IntoDynSyncSend` if it's already `Send`")]
@@ -51,9 +52,20 @@ macro_rules! already_send {
 
 // These structures are already `Send`.
 already_send!(
-    [std::backtrace::Backtrace][std::io::Stdout][std::io::Stderr][std::io::Error][std::fs::File]
-        [rustc_arena::DroplessArena][crate::memmap::Mmap][crate::profiling::SelfProfiler]
-        [crate::owned_slice::OwnedSlice]
+    [std::backtrace::Backtrace]
+    [std::io::Stdout]
+    [std::io::Stderr]
+    [std::io::Error]
+    [std::fs::File]
+    [std::sync::Condvar]
+    [jobserver_crate::Client]
+    [jobserver_crate::HelperThread]
+    [jobserver_crate::Acquired]
+    [Box<dyn Any + Send>]
+    [rustc_arena::DroplessArena]
+    [crate::memmap::Mmap]
+    [crate::profiling::SelfProfiler]
+    [crate::owned_slice::OwnedSlice]
 );
 
 macro_rules! impl_dyn_send {
@@ -64,10 +76,14 @@ macro_rules! impl_dyn_send {
 
 impl_dyn_send!(
     [std::sync::atomic::AtomicPtr<T> where T]
-    [std::sync::Mutex<T> where T: ?Sized+ DynSend]
+            [std::sync::Mutex<T> where T: ?Sized + DynSend]
+            [std::sync::RwLock<T> where T: ?Sized + DynSend]
     [std::sync::mpsc::Sender<T> where T: DynSend]
+            [std::sync::mpsc::Receiver<T> where T: DynSend]
     [std::sync::Arc<T> where T: ?Sized + DynSync + DynSend]
+            [std::sync::OnceLock<T> where T: DynSend]
     [std::sync::LazyLock<T, F> where T: DynSend, F: DynSend]
+            [std::thread::JoinHandle<T> where T]
     [std::collections::HashSet<K, S> where K: DynSend, S: DynSend]
     [std::collections::HashMap<K, V, S> where K: DynSend, V: DynSend, S: DynSend]
     [std::collections::BTreeMap<K, V, A> where K: DynSend, V: DynSend, A: std::alloc::Allocator + Clone + DynSend]
@@ -119,9 +135,9 @@ macro_rules! already_sync {
 // These structures are already `Sync`.
 already_sync!(
     [std::sync::atomic::AtomicBool][std::sync::atomic::AtomicUsize][std::sync::atomic::AtomicU8]
-        [std::sync::atomic::AtomicU32][std::backtrace::Backtrace][std::io::Error][std::fs::File]
-        [jobserver_crate::Client][crate::memmap::Mmap][crate::profiling::SelfProfiler]
-        [crate::owned_slice::OwnedSlice]
+        [std::sync::atomic::AtomicU32][std::backtrace::Backtrace][std::sync::Condvar]
+        [std::io::Error][std::fs::File][jobserver_crate::Client][crate::memmap::Mmap]
+        [crate::profiling::SelfProfiler][crate::owned_slice::OwnedSlice]
 );
 
 // Use portable AtomicU64 for targets without native 64-bit atomics
@@ -142,7 +158,9 @@ impl_dyn_sync!(
     [std::sync::OnceLock<T> where T: DynSend + DynSync]
     [std::sync::Mutex<T> where T: ?Sized + DynSend]
     [std::sync::Arc<T> where T: ?Sized + DynSync + DynSend]
+            [std::sync::RwLock<T> where T: ?Sized + DynSend + DynSync]
     [std::sync::LazyLock<T, F> where T: DynSend + DynSync, F: DynSend]
+    [std::sync::mpsc::SyncSender<T> where T: DynSend]
     [std::collections::HashSet<K, S> where K: DynSync, S: DynSync]
     [std::collections::HashMap<K, V, S> where K: DynSync, V: DynSync, S: DynSync]
     [std::collections::BTreeMap<K, V, A> where K: DynSync, V: DynSync, A: std::alloc::Allocator + Clone + DynSync]
@@ -224,3 +242,15 @@ impl<T> std::ops::DerefMut for IntoDynSyncSend<T> {
         &mut self.0
     }
 }
+
+#[inline]
+pub fn downcast_box_any_dyn_send<T: Any>(this: Box<dyn Any + DynSend>) -> Result<Box<T>, ()> {
+    if <dyn Any>::is::<T>(&*this) {
+        unsafe {
+            let (raw, alloc): (*mut (dyn Any + DynSend), _) = Box::into_raw_with_allocator(this);
+            Ok(Box::from_raw_in(raw as *mut T, alloc))
+        }
+    } else {
+        Err(())
+    }
+}

compiler/rustc_data_structures/src/sync.rs

@@ -50,6 +50,10 @@ pub use self::worker_local::{Registry, WorkerLocal};
 pub use crate::marker::*;
 
 mod freeze;
+
+mod task;
+pub use task::{Task, task};
+
 mod lock;
 mod parallel;
 mod vec;

compiler/rustc_data_structures/src/sync/task.rs

@@ -0,0 +1,101 @@
+use std::any::Any;
+use std::mem;
+use std::panic::{self, AssertUnwindSafe};
+use std::sync::Arc;
+
+use parking_lot::{Condvar, Mutex};
+
+use crate::jobserver;
+use crate::sync::{DynSend, FromDyn, IntoDynSyncSend, mode};
+
+enum TaskState<T: DynSend + 'static> {
+    Unexecuted(Box<dyn FnOnce() -> T + DynSend>),
+    Running,
+    Joined,
+    Result(Result<T, IntoDynSyncSend<Box<dyn Any + Send + 'static>>>),
+}
+
+struct TaskData<T: DynSend + 'static> {
+    state: Mutex<TaskState<T>>,
+    waiter: Condvar,
+}
+
+#[must_use]
+pub struct Task<T: DynSend + 'static> {
+    data: Arc<TaskData<T>>,
+}
+
+/// This attempts to run a closure in a background thread. It returns a `Task` type which
+/// you must call `join` on to ensure that the closure runs.
+pub fn task<T: DynSend>(f: impl FnOnce() -> T + DynSend + 'static) -> Task<T> {
+    let task = Task {
+        data: Arc::new(TaskData {
+            state: Mutex::new(TaskState::Unexecuted(Box::new(f))),
+            waiter: Condvar::new(),
+        }),
+    };
+
+    if mode::is_dyn_thread_safe() {
+        let data = FromDyn::from(Arc::clone(&task.data));
+
+        // Try to execute the task on a separate thread.
+        rayon::spawn(move || {
+            let data = data.into_inner();
+            let mut state = data.state.lock();
+            if matches!(*state, TaskState::Unexecuted(..)) {
+                if let TaskState::Unexecuted(f) = mem::replace(&mut *state, TaskState::Running) {
+                    drop(state);
+                    let result = panic::catch_unwind(AssertUnwindSafe(f));
+
+                    let unblock = {
+                        let mut state = data.state.lock();
+                        let unblock = matches!(*state, TaskState::Joined);
+                        *state = TaskState::Result(result.map_err(|e| IntoDynSyncSend(e)));
+                        unblock
+                    };
+
+                    if unblock {
+                        rayon_core::mark_unblocked(&rayon_core::Registry::current());
+                    }
+
+                    data.waiter.notify_one();
+                }
+            }
+        });
+    }
+
+    task
+}
+
+#[inline]
+fn unwind<T>(result: Result<T, IntoDynSyncSend<Box<dyn Any + Send + 'static>>>) -> T {
+    match result {
+        Ok(r) => r,
+        Err(err) => panic::resume_unwind(err.0),
+    }
+}
+
+impl<T: DynSend> Task<T> {
+    pub fn join(self) -> T {
+        let mut state_guard = self.data.state.lock();
+
+        match mem::replace(&mut *state_guard, TaskState::Joined) {
+            TaskState::Unexecuted(f) => f(),
+            TaskState::Result(result) => unwind(result),
+            TaskState::Running => {
+                rayon_core::mark_blocked();
+                jobserver::release_thread();
+
+                self.data.waiter.wait(&mut state_guard);
+
+                jobserver::acquire_thread();
+
+                match mem::replace(&mut *state_guard, TaskState::Joined) {
+                    TaskState::Result(result) => unwind(result),
+                    _ => panic!(),
+                }
+            }
+            TaskState::Joined => panic!(),
+        }
+    }
+}

compiler/rustc_driver_impl/src/lib.rs

@@ -380,13 +380,14 @@ pub fn run_compiler(at_args: &[String], callbacks: &mut (dyn Callbacks + Send))
                 }
             }
 
-            Some(Linker::codegen_and_build_linker(tcx, &*compiler.codegen_backend))
+            Some(Linker::codegen_and_build_linker(tcx, compiler))
         });
 
         // Linking is done outside the `compiler.enter()` so that the
         // `GlobalCtxt` within `Queries` can be freed as early as possible.
         if let Some(linker) = linker {
-            linker.link(sess, codegen_backend);
+            let _timer = sess.timer("waiting for linking");
+            linker.join();
         }
     })
 }