Rollup merge of #143388 - bjorn3:lto_refactors, r=compiler-errors
Various refactors to the LTO handling code In particular reducing the sharing of code paths between fat and thin-LTO and making the fat LTO implementation more self-contained. This also moves some autodiff handling out of cg_ssa into cg_llvm given that Enzyme only works with LLVM anyway and an implementation for another backend may do things entirely differently. This will also make it a bit easier to split LTO handling out of the coordinator thread main loop into a separate loop, which should reduce the complexity of the coordinator thread.
This commit is contained in:
@@ -7,7 +7,7 @@ use std::sync::Arc;
|
||||
use std::{io, iter, slice};
|
||||
|
||||
use object::read::archive::ArchiveFile;
|
||||
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
|
||||
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
|
||||
use rustc_codegen_ssa::back::symbol_export;
|
||||
use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
|
||||
use rustc_codegen_ssa::traits::*;
|
||||
@@ -201,7 +201,7 @@ pub(crate) fn run_fat(
|
||||
cgcx: &CodegenContext<LlvmCodegenBackend>,
|
||||
modules: Vec<FatLtoInput<LlvmCodegenBackend>>,
|
||||
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
|
||||
) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
|
||||
) -> Result<ModuleCodegen<ModuleLlvm>, FatalError> {
|
||||
let dcx = cgcx.create_dcx();
|
||||
let dcx = dcx.handle();
|
||||
let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, dcx)?;
|
||||
@@ -217,7 +217,7 @@ pub(crate) fn run_thin(
|
||||
cgcx: &CodegenContext<LlvmCodegenBackend>,
|
||||
modules: Vec<(String, ThinBuffer)>,
|
||||
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
|
||||
) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
|
||||
) -> Result<(Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
|
||||
let dcx = cgcx.create_dcx();
|
||||
let dcx = dcx.handle();
|
||||
let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, dcx)?;
|
||||
@@ -248,7 +248,7 @@ fn fat_lto(
|
||||
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
|
||||
mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
|
||||
symbols_below_threshold: &[*const libc::c_char],
|
||||
) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
|
||||
) -> Result<ModuleCodegen<ModuleLlvm>, FatalError> {
|
||||
let _timer = cgcx.prof.generic_activity("LLVM_fat_lto_build_monolithic_module");
|
||||
info!("going for a fat lto");
|
||||
|
||||
@@ -366,7 +366,7 @@ fn fat_lto(
|
||||
save_temp_bitcode(cgcx, &module, "lto.after-restriction");
|
||||
}
|
||||
|
||||
Ok(LtoModuleCodegen::Fat(module))
|
||||
Ok(module)
|
||||
}
|
||||
|
||||
pub(crate) struct Linker<'a>(&'a mut llvm::Linker<'a>);
|
||||
@@ -436,7 +436,7 @@ fn thin_lto(
|
||||
serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
|
||||
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
|
||||
symbols_below_threshold: &[*const libc::c_char],
|
||||
) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
|
||||
) -> Result<(Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
|
||||
let _timer = cgcx.prof.generic_activity("LLVM_thin_lto_global_analysis");
|
||||
unsafe {
|
||||
info!("going for that thin, thin LTO");
|
||||
@@ -568,10 +568,7 @@ fn thin_lto(
|
||||
}
|
||||
|
||||
info!(" - {}: re-compiled", module_name);
|
||||
opt_jobs.push(LtoModuleCodegen::Thin(ThinModule {
|
||||
shared: Arc::clone(&shared),
|
||||
idx: module_index,
|
||||
}));
|
||||
opt_jobs.push(ThinModule { shared: Arc::clone(&shared), idx: module_index });
|
||||
}
|
||||
|
||||
// Save the current ThinLTO import information for the next compilation
|
||||
|
||||
@@ -817,10 +817,12 @@ pub(crate) fn link(
|
||||
|
||||
pub(crate) fn codegen(
|
||||
cgcx: &CodegenContext<LlvmCodegenBackend>,
|
||||
dcx: DiagCtxtHandle<'_>,
|
||||
module: ModuleCodegen<ModuleLlvm>,
|
||||
config: &ModuleConfig,
|
||||
) -> Result<CompiledModule, FatalError> {
|
||||
let dcx = cgcx.create_dcx();
|
||||
let dcx = dcx.handle();
|
||||
|
||||
let _timer = cgcx.prof.generic_activity_with_arg("LLVM_module_codegen", &*module.name);
|
||||
{
|
||||
let llmod = module.module_llvm.llmod();
|
||||
|
||||
@@ -2,7 +2,6 @@ use std::ptr;
|
||||
|
||||
use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
|
||||
use rustc_codegen_ssa::ModuleCodegen;
|
||||
use rustc_codegen_ssa::back::write::ModuleConfig;
|
||||
use rustc_codegen_ssa::common::TypeKind;
|
||||
use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
|
||||
use rustc_errors::FatalError;
|
||||
@@ -461,7 +460,6 @@ pub(crate) fn differentiate<'ll>(
|
||||
module: &'ll ModuleCodegen<ModuleLlvm>,
|
||||
cgcx: &CodegenContext<LlvmCodegenBackend>,
|
||||
diff_items: Vec<AutoDiffItem>,
|
||||
_config: &ModuleConfig,
|
||||
) -> Result<(), FatalError> {
|
||||
for item in &diff_items {
|
||||
trace!("{}", item);
|
||||
|
||||
@@ -37,10 +37,6 @@ impl<G: EmissionGuarantee> Diagnostic<'_, G> for ParseTargetMachineConfig<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Diagnostic)]
|
||||
#[diag(codegen_llvm_autodiff_without_lto)]
|
||||
pub(crate) struct AutoDiffWithoutLTO;
|
||||
|
||||
#[derive(Diagnostic)]
|
||||
#[diag(codegen_llvm_autodiff_without_enable)]
|
||||
pub(crate) struct AutoDiffWithoutEnable;
|
||||
|
||||
@@ -26,11 +26,11 @@ use std::mem::ManuallyDrop;
|
||||
use back::owned_target_machine::OwnedTargetMachine;
|
||||
use back::write::{create_informational_target_machine, create_target_machine};
|
||||
use context::SimpleCx;
|
||||
use errors::{AutoDiffWithoutLTO, ParseTargetMachineConfig};
|
||||
use errors::ParseTargetMachineConfig;
|
||||
use llvm_util::target_config;
|
||||
use rustc_ast::expand::allocator::AllocatorKind;
|
||||
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
|
||||
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
|
||||
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
|
||||
use rustc_codegen_ssa::back::write::{
|
||||
CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryConfig, TargetMachineFactoryFn,
|
||||
};
|
||||
@@ -43,7 +43,7 @@ use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
|
||||
use rustc_middle::ty::TyCtxt;
|
||||
use rustc_middle::util::Providers;
|
||||
use rustc_session::Session;
|
||||
use rustc_session::config::{Lto, OptLevel, OutputFilenames, PrintKind, PrintRequest};
|
||||
use rustc_session::config::{OptLevel, OutputFilenames, PrintKind, PrintRequest};
|
||||
use rustc_span::Symbol;
|
||||
|
||||
mod back {
|
||||
@@ -174,18 +174,29 @@ impl WriteBackendMethods for LlvmCodegenBackend {
|
||||
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
|
||||
back::write::link(cgcx, dcx, modules)
|
||||
}
|
||||
fn run_fat_lto(
|
||||
fn run_and_optimize_fat_lto(
|
||||
cgcx: &CodegenContext<Self>,
|
||||
modules: Vec<FatLtoInput<Self>>,
|
||||
cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
|
||||
) -> Result<LtoModuleCodegen<Self>, FatalError> {
|
||||
back::lto::run_fat(cgcx, modules, cached_modules)
|
||||
diff_fncs: Vec<AutoDiffItem>,
|
||||
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
|
||||
let mut module = back::lto::run_fat(cgcx, modules, cached_modules)?;
|
||||
|
||||
if !diff_fncs.is_empty() {
|
||||
builder::autodiff::differentiate(&module, cgcx, diff_fncs)?;
|
||||
}
|
||||
|
||||
let dcx = cgcx.create_dcx();
|
||||
let dcx = dcx.handle();
|
||||
back::lto::run_pass_manager(cgcx, dcx, &mut module, false)?;
|
||||
|
||||
Ok(module)
|
||||
}
|
||||
fn run_thin_lto(
|
||||
cgcx: &CodegenContext<Self>,
|
||||
modules: Vec<(String, Self::ThinBuffer)>,
|
||||
cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
|
||||
) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError> {
|
||||
) -> Result<(Vec<ThinModule<Self>>, Vec<WorkProduct>), FatalError> {
|
||||
back::lto::run_thin(cgcx, modules, cached_modules)
|
||||
}
|
||||
fn optimize(
|
||||
@@ -196,14 +207,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
|
||||
) -> Result<(), FatalError> {
|
||||
back::write::optimize(cgcx, dcx, module, config)
|
||||
}
|
||||
fn optimize_fat(
|
||||
cgcx: &CodegenContext<Self>,
|
||||
module: &mut ModuleCodegen<Self::Module>,
|
||||
) -> Result<(), FatalError> {
|
||||
let dcx = cgcx.create_dcx();
|
||||
let dcx = dcx.handle();
|
||||
back::lto::run_pass_manager(cgcx, dcx, module, false)
|
||||
}
|
||||
fn optimize_thin(
|
||||
cgcx: &CodegenContext<Self>,
|
||||
thin: ThinModule<Self>,
|
||||
@@ -212,11 +215,10 @@ impl WriteBackendMethods for LlvmCodegenBackend {
|
||||
}
|
||||
fn codegen(
|
||||
cgcx: &CodegenContext<Self>,
|
||||
dcx: DiagCtxtHandle<'_>,
|
||||
module: ModuleCodegen<Self::Module>,
|
||||
config: &ModuleConfig,
|
||||
) -> Result<CompiledModule, FatalError> {
|
||||
back::write::codegen(cgcx, dcx, module, config)
|
||||
back::write::codegen(cgcx, module, config)
|
||||
}
|
||||
fn prepare_thin(
|
||||
module: ModuleCodegen<Self::Module>,
|
||||
@@ -227,19 +229,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
|
||||
fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer) {
|
||||
(module.name, back::lto::ModuleBuffer::new(module.module_llvm.llmod()))
|
||||
}
|
||||
/// Generate autodiff rules
|
||||
fn autodiff(
|
||||
cgcx: &CodegenContext<Self>,
|
||||
module: &ModuleCodegen<Self::Module>,
|
||||
diff_fncs: Vec<AutoDiffItem>,
|
||||
config: &ModuleConfig,
|
||||
) -> Result<(), FatalError> {
|
||||
if cgcx.lto != Lto::Fat {
|
||||
let dcx = cgcx.create_dcx();
|
||||
return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO));
|
||||
}
|
||||
builder::autodiff::differentiate(module, cgcx, diff_fncs, config)
|
||||
}
|
||||
}
|
||||
|
||||
impl LlvmCodegenBackend {
|
||||
|
||||
Reference in New Issue
Block a user