Fuse codegen into LTO optimize methods

Author: bjorn3

compiler/rustc_codegen_gcc/src/back/lto.rs

@@ -26,15 +26,15 @@ use object::read::archive::ArchiveFile;
 use rustc_codegen_ssa::back::lto::SerializedModule;
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput, SharedEmitter};
 use rustc_codegen_ssa::traits::*;
-use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
+use rustc_codegen_ssa::{CompiledModule, ModuleCodegen, ModuleKind, looks_like_rust_object_file};
 use rustc_data_structures::memmap::Mmap;
 use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_errors::{DiagCtxt, DiagCtxtHandle};
 use rustc_log::tracing::info;
 use rustc_session::config::Lto;
 use tempfile::{TempDir, tempdir};
 
-use crate::back::write::save_temp_bitcode;
+use crate::back::write::{codegen, save_temp_bitcode};
 use crate::errors::LtoBitcodeFromRlib;
 use crate::{GccCodegenBackend, GccContext, LtoMode, to_gcc_opt_level};
 
@@ -112,7 +112,7 @@ pub(crate) fn run_fat(
     shared_emitter: &SharedEmitter,
     each_linked_rlib_for_lto: &[PathBuf],
     modules: Vec<FatLtoInput<GccCodegenBackend>>,
-) -> ModuleCodegen<GccContext> {
+) -> CompiledModule {
     let dcx = DiagCtxt::new(Box::new(shared_emitter.clone()));
     let dcx = dcx.handle();
     let lto_data = prepare_lto(cgcx, each_linked_rlib_for_lto, dcx);
@@ -132,12 +132,12 @@ pub(crate) fn run_fat(
 fn fat_lto(
     cgcx: &CodegenContext,
     prof: &SelfProfilerRef,
-    _dcx: DiagCtxtHandle<'_>,
+    dcx: DiagCtxtHandle<'_>,
     modules: Vec<FatLtoInput<GccCodegenBackend>>,
     mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
     tmp_path: TempDir,
     //symbols_below_threshold: &[String],
-) -> ModuleCodegen<GccContext> {
+) -> CompiledModule {
     let _timer = prof.generic_activity("GCC_fat_lto_build_monolithic_module");
     info!("going for a fat lto");
 
@@ -260,7 +260,7 @@ fn fat_lto(
     // of now.
     module.module_llvm.temp_dir = Some(tmp_path);
 
-    module
+    codegen(cgcx, prof, dcx, module, &cgcx.module_config)
 }
 
 pub struct ModuleBuffer(PathBuf);

compiler/rustc_codegen_gcc/src/back/write.rs

@@ -2,12 +2,10 @@ use std::{env, fs};
 
 use gccjit::{Context, OutputKind};
 use rustc_codegen_ssa::back::link::ensure_removed;
-use rustc_codegen_ssa::back::write::{
-    BitcodeSection, CodegenContext, EmitObj, ModuleConfig, SharedEmitter,
-};
+use rustc_codegen_ssa::back::write::{BitcodeSection, CodegenContext, EmitObj, ModuleConfig};
 use rustc_codegen_ssa::{CompiledModule, ModuleCodegen};
 use rustc_data_structures::profiling::SelfProfilerRef;
-use rustc_errors::DiagCtxt;
+use rustc_errors::DiagCtxtHandle;
 use rustc_fs_util::link_or_copy;
 use rustc_log::tracing::debug;
 use rustc_session::config::OutputType;
@@ -20,13 +18,10 @@ use crate::{GccContext, LtoMode};
 pub(crate) fn codegen(
     cgcx: &CodegenContext,
     prof: &SelfProfilerRef,
-    shared_emitter: &SharedEmitter,
+    dcx: DiagCtxtHandle<'_>,
     module: ModuleCodegen<GccContext>,
     config: &ModuleConfig,
 ) -> CompiledModule {
-    let dcx = DiagCtxt::new(Box::new(shared_emitter.clone()));
-    let dcx = dcx.handle();
-
     let _timer = prof.generic_activity_with_arg("GCC_module_codegen", &*module.name);
     {
         let context = &module.module_llvm.context;

compiler/rustc_codegen_gcc/src/lib.rs

@@ -93,7 +93,7 @@ use rustc_codegen_ssa::{CodegenResults, CompiledModule, ModuleCodegen, TargetCon
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_data_structures::sync::IntoDynSyncSend;
-use rustc_errors::DiagCtxtHandle;
+use rustc_errors::{DiagCtxt, DiagCtxtHandle};
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::util::Providers;
@@ -438,7 +438,7 @@ impl WriteBackendMethods for GccCodegenBackend {
     type ThinData = ();
     type ThinBuffer = ThinBuffer;
 
-    fn run_and_optimize_fat_lto(
+    fn optimize_and_codegen_fat_lto(
         cgcx: &CodegenContext,
         prof: &SelfProfilerRef,
         shared_emitter: &SharedEmitter,
@@ -447,7 +447,7 @@ impl WriteBackendMethods for GccCodegenBackend {
         _exported_symbols_for_lto: &[String],
         each_linked_rlib_for_lto: &[PathBuf],
         modules: Vec<FatLtoInput<Self>>,
-    ) -> ModuleCodegen<Self::Module> {
+    ) -> CompiledModule {
         back::lto::run_fat(cgcx, prof, shared_emitter, each_linked_rlib_for_lto, modules)
     }
 
@@ -474,13 +474,13 @@ impl WriteBackendMethods for GccCodegenBackend {
         module.module_llvm.context.set_optimization_level(to_gcc_opt_level(config.opt_level));
     }
 
-    fn optimize_thin(
+    fn optimize_and_codegen_thin(
         _cgcx: &CodegenContext,
         _prof: &SelfProfilerRef,
         _shared_emitter: &SharedEmitter,
         _tm_factory: TargetMachineFactoryFn<Self>,
         _thin: ThinModule<Self>,
-    ) -> ModuleCodegen<Self::Module> {
+    ) -> CompiledModule {
         unreachable!()
     }
 
@@ -491,7 +491,9 @@ impl WriteBackendMethods for GccCodegenBackend {
         module: ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> CompiledModule {
-        back::write::codegen(cgcx, prof, shared_emitter, module, config)
+        let dcx = DiagCtxt::new(Box::new(shared_emitter.clone()));
+        let dcx = dcx.handle();
+        back::write::codegen(cgcx, prof, dcx, module, config)
     }
 
     fn prepare_thin(_module: ModuleCodegen<Self::Module>) -> (String, Self::ThinBuffer) {

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -13,7 +13,7 @@ use rustc_codegen_ssa::back::write::{
     CodegenContext, FatLtoInput, SharedEmitter, TargetMachineFactoryFn,
 };
 use rustc_codegen_ssa::traits::*;
-use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
+use rustc_codegen_ssa::{CompiledModule, ModuleCodegen, ModuleKind, looks_like_rust_object_file};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::memmap::Mmap;
 use rustc_data_structures::profiling::SelfProfilerRef;
@@ -25,7 +25,8 @@ use rustc_session::config::{self, Lto};
 use tracing::{debug, info};
 
 use crate::back::write::{
-    self, CodegenDiagnosticsStage, DiagnosticHandlers, bitcode_section_name, save_temp_bitcode,
+    self, CodegenDiagnosticsStage, DiagnosticHandlers, bitcode_section_name, codegen,
+    save_temp_bitcode,
 };
 use crate::errors::{LlvmError, LtoBitcodeFromRlib};
 use crate::llvm::{self, build_string};
@@ -745,13 +746,13 @@ impl Drop for ThinBuffer {
     }
 }
 
-pub(crate) fn optimize_thin_module(
+pub(crate) fn optimize_and_codegen_thin_module(
     cgcx: &CodegenContext,
     prof: &SelfProfilerRef,
     shared_emitter: &SharedEmitter,
     tm_factory: TargetMachineFactoryFn<LlvmCodegenBackend>,
     thin_module: ThinModule<LlvmCodegenBackend>,
-) -> ModuleCodegen<ModuleLlvm> {
+) -> CompiledModule {
     let dcx = DiagCtxt::new(Box::new(shared_emitter.clone()));
     let dcx = dcx.handle();
 
@@ -830,7 +831,7 @@ pub(crate) fn optimize_thin_module(
             save_temp_bitcode(cgcx, &module, "thin-lto-after-pm");
         }
     }
-    module
+    codegen(cgcx, prof, shared_emitter, module, &cgcx.module_config)
 }
 
 /// Maps LLVM module identifiers to their corresponding LLVM LTO cache keys

compiler/rustc_codegen_llvm/src/lib.rs

@@ -136,15 +136,15 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     type TargetMachine = OwnedTargetMachine;
     type ThinData = back::lto::ThinData;
     type ThinBuffer = back::lto::ThinBuffer;
-    fn run_and_optimize_fat_lto(
+    fn optimize_and_codegen_fat_lto(
         cgcx: &CodegenContext,
         prof: &SelfProfilerRef,
         shared_emitter: &SharedEmitter,
         tm_factory: TargetMachineFactoryFn<LlvmCodegenBackend>,
         exported_symbols_for_lto: &[String],
         each_linked_rlib_for_lto: &[PathBuf],
         modules: Vec<FatLtoInput<Self>>,
-    ) -> ModuleCodegen<Self::Module> {
+    ) -> CompiledModule {
         let mut module = back::lto::run_fat(
             cgcx,
             prof,
@@ -159,7 +159,7 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         let dcx = dcx.handle();
         back::lto::run_pass_manager(cgcx, prof, dcx, &mut module, false);
 
-        module
+        back::write::codegen(cgcx, prof, shared_emitter, module, &cgcx.module_config)
     }
     fn run_thin_lto(
         cgcx: &CodegenContext,
@@ -189,14 +189,14 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     ) {
         back::write::optimize(cgcx, prof, shared_emitter, module, config)
     }
-    fn optimize_thin(
+    fn optimize_and_codegen_thin(
         cgcx: &CodegenContext,
         prof: &SelfProfilerRef,
         shared_emitter: &SharedEmitter,
         tm_factory: TargetMachineFactoryFn<LlvmCodegenBackend>,
         thin: ThinModule<Self>,
-    ) -> ModuleCodegen<Self::Module> {
-        back::lto::optimize_thin_module(cgcx, prof, shared_emitter, tm_factory, thin)
+    ) -> CompiledModule {
+        back::lto::optimize_and_codegen_thin_module(cgcx, prof, shared_emitter, tm_factory, thin)
     }
     fn codegen(
         cgcx: &CodegenContext,

compiler/rustc_codegen_ssa/src/back/write.rs

@@ -1006,16 +1006,15 @@ fn do_fat_lto<B: ExtraBackendMethods>(
         needs_fat_lto.push(FatLtoInput::Serialized { name: wp.cgu_name, buffer: module })
     }
 
-    let module = B::run_and_optimize_fat_lto(
+    B::optimize_and_codegen_fat_lto(
         cgcx,
         prof,
         &shared_emitter,
         tm_factory,
         exported_symbols_for_lto,
         each_linked_rlib_for_lto,
         needs_fat_lto,
-    );
-    B::codegen(cgcx, prof, &shared_emitter, module, &cgcx.module_config)
+    )
 }
 
 fn do_thin_lto<B: ExtraBackendMethods>(
@@ -1180,8 +1179,7 @@ fn execute_thin_lto_work_item<B: ExtraBackendMethods>(
 ) -> CompiledModule {
     let _timer = prof.generic_activity_with_arg("codegen_module_perform_lto", module.name());
 
-    let module = B::optimize_thin(cgcx, prof, &shared_emitter, tm_factory, module);
-    B::codegen(cgcx, prof, &shared_emitter, module, &cgcx.module_config)
+    B::optimize_and_codegen_thin(cgcx, prof, &shared_emitter, tm_factory, module)
 }
 
 /// Messages sent to the coordinator.

compiler/rustc_codegen_ssa/src/traits/write.rs

@@ -19,15 +19,15 @@ pub trait WriteBackendMethods: Clone + 'static {
 
     /// Performs fat LTO by merging all modules into a single one, running autodiff
     /// if necessary and running any further optimizations
-    fn run_and_optimize_fat_lto(
+    fn optimize_and_codegen_fat_lto(
         cgcx: &CodegenContext,
         prof: &SelfProfilerRef,
         shared_emitter: &SharedEmitter,
         tm_factory: TargetMachineFactoryFn<Self>,
         exported_symbols_for_lto: &[String],
         each_linked_rlib_for_lto: &[PathBuf],
         modules: Vec<FatLtoInput<Self>>,
-    ) -> ModuleCodegen<Self::Module>;
+    ) -> CompiledModule;
     /// Performs thin LTO by performing necessary global analysis and returning two
     /// lists, one of the modules that need optimization and another for modules that
     /// can simply be copied over from the incr. comp. cache.
@@ -47,13 +47,13 @@ pub trait WriteBackendMethods: Clone + 'static {
         module: &mut ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     );
-    fn optimize_thin(
+    fn optimize_and_codegen_thin(
         cgcx: &CodegenContext,
         prof: &SelfProfilerRef,
         shared_emitter: &SharedEmitter,
         tm_factory: TargetMachineFactoryFn<Self>,
         thin: ThinModule<Self>,
-    ) -> ModuleCodegen<Self::Module>;
+    ) -> CompiledModule;
     fn codegen(
         cgcx: &CodegenContext,
         prof: &SelfProfilerRef,