Ensure fat LTO doesn't merge everything into the allocator module

Author: bjorn3

compiler/rustc_codegen_cranelift/src/driver/aot.rs

@@ -12,7 +12,9 @@ use cranelift_object::{ObjectBuilder, ObjectModule};
 use rustc_codegen_ssa::assert_module_sources::CguReuse;
 use rustc_codegen_ssa::back::link::ensure_removed;
 use rustc_codegen_ssa::base::determine_cgu_reuse;
-use rustc_codegen_ssa::{CodegenResults, CompiledModule, CrateInfo, errors as ssa_errors};
+use rustc_codegen_ssa::{
+ CodegenResults, CompiledModule, CrateInfo, ModuleKind, errors as ssa_errors,
+};
 use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{IntoDynSyncSend, par_map};
@@ -361,6 +363,7 @@ fn emit_cgu(
  invocation_temp,
  prof,
  product.object,
+ ModuleKind::Regular,
  name.clone(),
  producer,
  )?;
@@ -369,6 +372,7 @@ fn emit_cgu(
  module_regular,
  module_global_asm: global_asm_object_file.map(|global_asm_object_file| CompiledModule {
  name: format!("{name}.asm"),
+ kind: ModuleKind::Regular,
  object: Some(global_asm_object_file),
  dwarf_object: None,
  bytecode: None,
@@ -385,6 +389,7 @@ fn emit_module(
  invocation_temp: Option<&str>,
  prof: &SelfProfilerRef,
  mut object: cranelift_object::object::write::Object<'_>,
+ kind: ModuleKind,
  name: String,
  producer_str: &str,
 ) -> Result<CompiledModule, String> {
@@ -425,6 +430,7 @@ fn emit_module(
 
  Ok(CompiledModule {
  name,
+ kind,
  object: Some(tmp_file),
  dwarf_object: None,
  bytecode: None,
@@ -479,6 +485,7 @@ fn reuse_workproduct_for_cgu(
  Ok(ModuleCodegenResult {
  module_regular: CompiledModule {
  name: cgu.name().to_string(),
+ kind: ModuleKind::Regular,
  object: Some(obj_out_regular),
  dwarf_object: None,
  bytecode: None,
@@ -488,6 +495,7 @@ fn reuse_workproduct_for_cgu(
  },
  module_global_asm: source_file_global_asm.map(|source_file| CompiledModule {
  name: cgu.name().to_string(),
+ kind: ModuleKind::Regular,
  object: Some(obj_out_global_asm),
  dwarf_object: None,
  bytecode: None,
@@ -643,6 +651,7 @@ fn emit_allocator_module(tcx: TyCtxt<'_>) -> Option<CompiledModule> {
  tcx.sess.invocation_temp.as_deref(),
  &tcx.sess.prof,
  product.object,
+ ModuleKind::Allocator,
  "allocator_shim".to_owned(),
  &crate::debuginfo::producer(tcx.sess),
  ) {

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -11,7 +11,7 @@ use object::{Object, ObjectSection};
 use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
 use rustc_codegen_ssa::traits::*;
-use rustc_codegen_ssa::{ModuleCodegen, looks_like_rust_object_file};
+use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::memmap::Mmap;
 use rustc_errors::DiagCtxtHandle;
@@ -225,9 +225,15 @@ fn fat_lto(
  // All the other modules will be serialized and reparsed into the new
  // context, so this hopefully avoids serializing and parsing the largest
  // codegen unit.
+ //
+ // Additionally use a regular module as the base here to ensure that various
+ // file copy operations in the backend work correctly. The only other kind
+ // of module here should be an allocator one, and if your crate is smaller
+ // than the allocator module then the size doesn't really matter anyway.
  let costliest_module = in_memory
  .iter()
  .enumerate()
+ .filter(|&(_, module)| module.kind == ModuleKind::Regular)
  .map(|(i, module)| {
  let cost = unsafe { llvm::LLVMRustModuleCost(module.module_llvm.llmod()) };
  (cost, i)

compiler/rustc_codegen_ssa/src/back/write.rs

@@ -490,7 +490,7 @@ fn copy_all_cgu_workproducts_to_incr_comp_cache_dir(
 
  let _timer = sess.timer("copy_all_cgu_workproducts_to_incr_comp_cache_dir");
 
- for module in &compiled_modules.modules {
+ for module in compiled_modules.modules.iter().filter(|m| m.kind == ModuleKind::Regular) {
  let mut files = Vec::new();
  if let Some(object_file_path) = &module.object {
  files.push((OutputType::Object.extension(), object_file_path.as_path()));
@@ -960,6 +960,7 @@ fn execute_copy_from_cache_work_item<B: ExtraBackendMethods>(
 
  WorkItemResult::Finished(CompiledModule {
  links_from_incr_cache,
+ kind: ModuleKind::Regular,
  name: module.name,
  object,
  dwarf_object,

compiler/rustc_codegen_ssa/src/lib.rs

@@ -120,6 +120,7 @@ impl<M> ModuleCodegen<M> {
 
  CompiledModule {
  name: self.name,
+ kind: self.kind,
  object,
  dwarf_object,
  bytecode,
@@ -133,6 +134,7 @@ impl<M> ModuleCodegen<M> {
 #[derive(Debug, Encodable, Decodable)]
 pub struct CompiledModule {
  pub name: String,
+ pub kind: ModuleKind,
  pub object: Option<PathBuf>,
  pub dwarf_object: Option<PathBuf>,
  pub bytecode: Option<PathBuf>,