Precompile CGUs while the main thread has the implicit job server token

Author: Zoxc

src/librustc_codegen_ssa/base.rs

@@ -609,54 +609,75 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
 
     let total_codegen_time = Lock::new(Duration::new(0, 0));
 
-    let cgu_reuse: Vec<_> = tcx.sess.time("find cgu reuse", || {
-        codegen_units.iter().map(|cgu| determine_cgu_reuse(tcx, &cgu)).collect()
-    });
-
-    let mut cgus: FxHashMap<usize, _> = if cfg!(parallel_compiler) {
-        tcx.sess.time("compile first CGUs", || {
-            // Try to find one CGU to compile per thread.
-            let cgus: Vec<_> = cgu_reuse
-                .iter()
-                .enumerate()
-                .filter(|&(_, reuse)| reuse == &CguReuse::No)
-                .take(tcx.sess.threads())
-                .collect();
-
-            // Compile the found CGUs in parallel.
-            par_iter(cgus)
-                .map(|(i, _)| {
-                    let start_time = Instant::now();
-                    let module = backend.compile_codegen_unit(tcx, codegen_units[i].name());
-                    let mut time = total_codegen_time.lock();
-                    *time += start_time.elapsed();
-                    (i, module)
-                })
-                .collect()
-        })
-    } else {
-        FxHashMap::default()
+    // The non-parallel compiler can only translate codegen units to LLVM IR
+    // on a single thread, leading to a staircase effect where the N LLVM
+    // threads have to wait on the single codegen threads to generate work
+    // for them. The parallel compiler does not have this restriction, so
+    // we can pre-load the LLVM queue in parallel before handing off
+    // coordination to the OnGoingCodegen scheduler.
+    //
+    // This likely is a temporary measure. Once we don't have to support the
+    // non-parallel compiler anymore, we can compile CGUs end-to-end in
+    // parallel and get rid of the complicated scheduling logic.
+    let pre_compile_cgus = |cgu_reuse: &[CguReuse]| {
+        if cfg!(parallel_compiler) {
+            tcx.sess.time("compile_first_CGU_batch", || {
+                // Try to find one CGU to compile per thread.
+                let cgus: Vec<_> = cgu_reuse
+                    .iter()
+                    .enumerate()
+                    .filter(|&(_, reuse)| reuse == &CguReuse::No)
+                    .take(tcx.sess.threads())
+                    .collect();
+
+                // Compile the found CGUs in parallel.
+                par_iter(cgus)
+                    .map(|(i, _)| {
+                        let start_time = Instant::now();
+                        let module = backend.compile_codegen_unit(tcx, codegen_units[i].name());
+                        let mut time = total_codegen_time.lock();
+                        *time += start_time.elapsed();
+                        (i, module)
+                    })
+                    .collect()
+            })
+        } else {
+            FxHashMap::default()
+        }
     };
 
-    let mut total_codegen_time = total_codegen_time.into_inner();
+    let mut cgu_reuse = Vec::new();
+    let mut pre_compiled_cgus: Option<FxHashMap<usize, _>> = None;
 
-    for (i, cgu) in codegen_units.into_iter().enumerate() {
+    for (i, cgu) in codegen_units.iter().enumerate() {
         ongoing_codegen.wait_for_signal_to_codegen_item();
         ongoing_codegen.check_for_errors(tcx.sess);
 
+        // Do some setup work in the first iteration
+        if pre_compiled_cgus.is_none() {
+            // Calculate the CGU reuse
+            cgu_reuse = tcx.sess.time("find_cgu_reuse", || {
+                codegen_units.iter().map(|cgu| determine_cgu_reuse(tcx, &cgu)).collect()
+            });
+            // Pre compile some CGUs
+            pre_compiled_cgus = Some(pre_compile_cgus(&cgu_reuse));
+        }
+
         let cgu_reuse = cgu_reuse[i];
         tcx.sess.cgu_reuse_tracker.set_actual_reuse(&cgu.name().as_str(), cgu_reuse);
 
         match cgu_reuse {
             CguReuse::No => {
-                let (module, cost) = if let Some(cgu) = cgus.remove(&i) {
-                    cgu
-                } else {
-                    let start_time = Instant::now();
-                    let module = backend.compile_codegen_unit(tcx, cgu.name());
-                    total_codegen_time += start_time.elapsed();
-                    module
-                };
+                let (module, cost) =
+                    if let Some(cgu) = pre_compiled_cgus.as_mut().unwrap().remove(&i) {
+                        cgu
+                    } else {
+                        let start_time = Instant::now();
+                        let module = backend.compile_codegen_unit(tcx, cgu.name());
+                        let mut time = total_codegen_time.lock();
+                        *time += start_time.elapsed();
+                        module
+                    };
                 submit_codegened_module_to_llvm(
                     &backend,
                     &ongoing_codegen.coordinator_send,
@@ -695,7 +716,11 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
 
     // Since the main thread is sometimes blocked during codegen, we keep track
     // -Ztime-passes output manually.
-    print_time_passes_entry(tcx.sess.time_passes(), "codegen_to_LLVM_IR", total_codegen_time);
+    print_time_passes_entry(
+        tcx.sess.time_passes(),
+        "codegen_to_LLVM_IR",
+        total_codegen_time.into_inner(),
+    );
 
     ::rustc_incremental::assert_module_sources::assert_module_sources(tcx);
 

src/librustc_codegen_ssa/traits/backend.rs

@@ -48,6 +48,8 @@ pub trait ExtraBackendMethods: CodegenBackend + WriteBackendMethods + Sized + Se
         mods: &mut Self::Module,
         kind: AllocatorKind,
     );
+    /// This generates the codegen unit and returns it along with
+    /// a `u64` giving an estimate of the unit's processing cost.
     fn compile_codegen_unit(
         &self,
         tcx: TyCtxt<'_>,