Simplify ThinLTO handling

compiler/rustc_codegen_gcc/src/lib.rs

@@ -86,9 +86,7 @@ use rustc_codegen_ssa::back::write::{
 };
 use rustc_codegen_ssa::base::codegen_crate;
 use rustc_codegen_ssa::target_features::cfg_target_feature;
-use rustc_codegen_ssa::traits::{
-    CodegenBackend, ExtraBackendMethods, ThinBufferMethods, WriteBackendMethods,
-};
+use rustc_codegen_ssa::traits::{CodegenBackend, ExtraBackendMethods, WriteBackendMethods};
 use rustc_codegen_ssa::{CodegenResults, CompiledModule, ModuleCodegen, TargetConfig};
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::profiling::SelfProfilerRef;
@@ -423,20 +421,11 @@ unsafe impl Send for SyncContext {}
 // FIXME(antoyo): that shouldn't be Sync. Parallel compilation is currently disabled with "CodegenBackend::supports_parallel()".
 unsafe impl Sync for SyncContext {}
 
-pub struct ThinBuffer;
-
-impl ThinBufferMethods for ThinBuffer {
-    fn data(&self) -> &[u8] {
-        &[]
-    }
-}
-
 impl WriteBackendMethods for GccCodegenBackend {
     type Module = GccContext;
     type TargetMachine = ();
     type ModuleBuffer = ModuleBuffer;
     type ThinData = ();
-    type ThinBuffer = ThinBuffer;
 
     fn run_and_optimize_fat_lto(
         cgcx: &CodegenContext,
@@ -458,7 +447,7 @@ impl WriteBackendMethods for GccCodegenBackend {
         // FIXME(bjorn3): Limit LTO exports to these symbols
         _exported_symbols_for_lto: &[String],
         _each_linked_rlib_for_lto: &[PathBuf],
-        _modules: Vec<(String, Self::ThinBuffer)>,
+        _modules: Vec<(String, Self::ModuleBuffer)>,
         _cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
     ) -> (Vec<ThinModule<Self>>, Vec<WorkProduct>) {
         unreachable!()
@@ -502,11 +491,7 @@ impl WriteBackendMethods for GccCodegenBackend {
         back::write::codegen(cgcx, prof, shared_emitter, module, config)
     }
 
-    fn prepare_thin(_module: ModuleCodegen<Self::Module>) -> (String, Self::ThinBuffer) {
-        unreachable!()
-    }
-
-    fn serialize_module(_module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer) {
+    fn serialize_module(_module: Self::Module, _is_thin: bool) -> Self::ModuleBuffer {
         unimplemented!();
     }
 }

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -2,7 +2,6 @@ use std::collections::BTreeMap;
 use std::ffi::{CStr, CString};
 use std::fs::File;
 use std::path::{Path, PathBuf};
-use std::ptr::NonNull;
 use std::sync::Arc;
 use std::{io, iter, slice};
 
@@ -187,7 +186,7 @@ pub(crate) fn run_thin(
     dcx: DiagCtxtHandle<'_>,
     exported_symbols_for_lto: &[String],
     each_linked_rlib_for_lto: &[PathBuf],
-    modules: Vec<(String, ThinBuffer)>,
+    modules: Vec<(String, ModuleBuffer)>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
 ) -> (Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>) {
     let (symbols_below_threshold, upstream_modules) =
@@ -203,12 +202,6 @@ pub(crate) fn run_thin(
     thin_lto(cgcx, prof, dcx, modules, upstream_modules, cached_modules, &symbols_below_threshold)
 }
 
-pub(crate) fn prepare_thin(module: ModuleCodegen<ModuleLlvm>) -> (String, ThinBuffer) {
-    let name = module.name;
-    let buffer = ThinBuffer::new(module.module_llvm.llmod(), true);
-    (name, buffer)
-}
-
 fn fat_lto(
     cgcx: &CodegenContext,
     prof: &SelfProfilerRef,
@@ -297,7 +290,7 @@ fn fat_lto(
         // way we know of to do that is to serialize them to a string and them parse
         // them later. Not great but hey, that's why it's "fat" LTO, right?
         for module in in_memory {
-            let buffer = ModuleBuffer::new(module.module_llvm.llmod());
+            let buffer = ModuleBuffer::new(module.module_llvm.llmod(), false);
             let llmod_id = CString::new(&module.name[..]).unwrap();
             serialized_modules.push((SerializedModule::Local(buffer), llmod_id));
         }
@@ -400,7 +393,7 @@ fn thin_lto(
     cgcx: &CodegenContext,
     prof: &SelfProfilerRef,
     dcx: DiagCtxtHandle<'_>,
-    modules: Vec<(String, ThinBuffer)>,
+    modules: Vec<(String, ModuleBuffer)>,
     serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
     symbols_below_threshold: &[*const libc::c_char],
@@ -634,7 +627,9 @@ pub(crate) fn run_pass_manager(
     };
 
     unsafe {
-        write::llvm_optimize(cgcx, prof, dcx, module, None, config, opt_level, opt_stage, stage);
+        write::llvm_optimize(
+            cgcx, prof, dcx, module, None, None, config, opt_level, opt_stage, stage,
+        );
     }
 
     if cfg!(feature = "llvm_enzyme") && enable_ad && !thin {
@@ -643,7 +638,7 @@ pub(crate) fn run_pass_manager(
         if !config.autodiff.contains(&config::AutoDiff::NoPostopt) {
             unsafe {
                 write::llvm_optimize(
-                    cgcx, prof, dcx, module, None, config, opt_level, opt_stage, stage,
+                    cgcx, prof, dcx, module, None, None, config, opt_level, opt_stage, stage,
                 );
             }
         }
@@ -658,31 +653,26 @@ pub(crate) fn run_pass_manager(
     debug!("lto done");
 }
 
-pub struct ModuleBuffer(&'static mut llvm::ModuleBuffer);
-
-unsafe impl Send for ModuleBuffer {}
-unsafe impl Sync for ModuleBuffer {}
+#[repr(transparent)]
+pub(crate) struct Buffer(&'static mut llvm::Buffer);
 
-impl ModuleBuffer {
-    pub(crate) fn new(m: &llvm::Module) -> ModuleBuffer {
-        ModuleBuffer(unsafe { llvm::LLVMRustModuleBufferCreate(m) })
-    }
-}
+unsafe impl Send for Buffer {}
+unsafe impl Sync for Buffer {}
 
-impl ModuleBufferMethods for ModuleBuffer {
-    fn data(&self) -> &[u8] {
+impl Buffer {
+    pub(crate) fn data(&self) -> &[u8] {
         unsafe {
-            let ptr = llvm::LLVMRustModuleBufferPtr(self.0);
-            let len = llvm::LLVMRustModuleBufferLen(self.0);
+            let ptr = llvm::LLVMRustBufferPtr(self.0);
+            let len = llvm::LLVMRustBufferLen(self.0);
             slice::from_raw_parts(ptr, len)
         }
     }
 }
 
-impl Drop for ModuleBuffer {
+impl Drop for Buffer {
     fn drop(&mut self) {
         unsafe {
-            llvm::LLVMRustModuleBufferFree(&mut *(self.0 as *mut _));
+            llvm::LLVMRustBufferFree(&mut *(self.0 as *mut _));
         }
     }
 }
@@ -700,48 +690,22 @@ impl Drop for ThinData {
     }
 }
 
-pub struct ThinBuffer(&'static mut llvm::ThinLTOBuffer);
-
-unsafe impl Send for ThinBuffer {}
-unsafe impl Sync for ThinBuffer {}
-
-impl ThinBuffer {
-    pub(crate) fn new(m: &llvm::Module, is_thin: bool) -> ThinBuffer {
-        unsafe {
-            let buffer = llvm::LLVMRustThinLTOBufferCreate(m, is_thin);
-            ThinBuffer(buffer)
-        }
-    }
-
-    pub(crate) unsafe fn from_raw_ptr(ptr: *mut llvm::ThinLTOBuffer) -> ThinBuffer {
-        let mut ptr = NonNull::new(ptr).unwrap();
-        ThinBuffer(unsafe { ptr.as_mut() })
-    }
-
-    pub(crate) fn thin_link_data(&self) -> &[u8] {
-        unsafe {
-            let ptr = llvm::LLVMRustThinLTOBufferThinLinkDataPtr(self.0) as *const _;
-            let len = llvm::LLVMRustThinLTOBufferThinLinkDataLen(self.0);
-            slice::from_raw_parts(ptr, len)
-        }
-    }
+pub struct ModuleBuffer {
+    data: Buffer,
 }
 
-impl ThinBufferMethods for ThinBuffer {
-    fn data(&self) -> &[u8] {
+impl ModuleBuffer {
+    pub(crate) fn new(m: &llvm::Module, is_thin: bool) -> ModuleBuffer {
         unsafe {
-            let ptr = llvm::LLVMRustThinLTOBufferPtr(self.0) as *const _;
-            let len = llvm::LLVMRustThinLTOBufferLen(self.0);
-            slice::from_raw_parts(ptr, len)
+            let buffer = llvm::LLVMRustModuleSerialize(m, is_thin);
+            ModuleBuffer { data: Buffer(buffer) }
         }
     }
 }
 
-impl Drop for ThinBuffer {
-    fn drop(&mut self) {
-        unsafe {
-            llvm::LLVMRustThinLTOBufferFree(&mut *(self.0 as *mut _));
-        }
+impl ModuleBufferMethods for ModuleBuffer {
+    fn data(&self) -> &[u8] {
+        self.data.data()
     }
 }
 

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -1,7 +1,6 @@
 use std::ffi::{CStr, CString};
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
-use std::ptr::null_mut;
 use std::sync::Arc;
 use std::{fs, slice, str};
 
@@ -29,7 +28,7 @@ use rustc_target::spec::{
 };
 use tracing::{debug, trace};
 
-use crate::back::lto::ThinBuffer;
+use crate::back::lto::{Buffer, ModuleBuffer};
 use crate::back::owned_target_machine::OwnedTargetMachine;
 use crate::back::profiling::{
     LlvmSelfProfiler, selfprofile_after_pass_callback, selfprofile_before_pass_callback,
@@ -563,7 +562,8 @@ pub(crate) unsafe fn llvm_optimize(
     prof: &SelfProfilerRef,
     dcx: DiagCtxtHandle<'_>,
     module: &ModuleCodegen<ModuleLlvm>,
-    thin_lto_buffer: Option<&mut *mut llvm::ThinLTOBuffer>,
+    thin_lto_buffer: Option<&mut Option<Buffer>>,
+    thin_lto_summary_buffer: Option<&mut Option<Buffer>>,
     config: &ModuleConfig,
     opt_level: config::OptLevel,
     opt_stage: llvm::OptStage,
@@ -786,8 +786,7 @@ pub(crate) unsafe fn llvm_optimize(
             config.verify_llvm_ir,
             config.lint_llvm_ir,
             thin_lto_buffer,
-            config.emit_thin_lto,
-            config.emit_thin_lto_summary,
+            thin_lto_summary_buffer,
             merge_functions,
             unroll_loops,
             vectorize_slp,
@@ -933,13 +932,14 @@ pub(crate) fn optimize(
         // The bitcode obtained during the `codegen` phase is no longer suitable for performing LTO.
         // It may have undergone LTO due to ThinLocal, so we need to obtain the embedded bitcode at
         // this point.
-        let mut thin_lto_buffer = if (module.kind == ModuleKind::Regular
+        let (mut thin_lto_buffer, mut thin_lto_summary_buffer) = if (module.kind
+            == ModuleKind::Regular
             && config.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full))
             || config.emit_thin_lto_summary
         {
-            Some(null_mut())
+            (Some(None), config.emit_thin_lto_summary.then_some(None))
         } else {
-            None
+            (None, None)
         };
         unsafe {
             llvm_optimize(
@@ -948,24 +948,26 @@ pub(crate) fn optimize(
                 dcx,
                 module,
                 thin_lto_buffer.as_mut(),
+                thin_lto_summary_buffer.as_mut(),
                 config,
                 opt_level,
                 opt_stage,
                 autodiff_stage,
             )
         };
         if let Some(thin_lto_buffer) = thin_lto_buffer {
-            let thin_lto_buffer = unsafe { ThinBuffer::from_raw_ptr(thin_lto_buffer) };
+            let thin_lto_buffer = thin_lto_buffer.unwrap();
             module.thin_lto_buffer = Some(thin_lto_buffer.data().to_vec());
             let bc_summary_out = cgcx.output_filenames.temp_path_for_cgu(
                 OutputType::ThinLinkBitcode,
                 &module.name,
                 cgcx.invocation_temp.as_deref(),
             );
-            if config.emit_thin_lto_summary
+            if let Some(thin_lto_summary_buffer) = thin_lto_summary_buffer
                 && let Some(thin_link_bitcode_filename) = bc_summary_out.file_name()
             {
-                let summary_data = thin_lto_buffer.thin_link_data();
+                let thin_lto_summary_buffer = thin_lto_summary_buffer.unwrap();
+                let summary_data = thin_lto_summary_buffer.data();
                 prof.artifact_size(
                     "llvm_bitcode_summary",
                     thin_link_bitcode_filename.to_string_lossy(),
@@ -1033,7 +1035,7 @@ pub(crate) fn codegen(
                         "LLVM_module_codegen_make_bitcode",
                         &*module.name,
                     );
-                    ThinBuffer::new(llmod, config.emit_thin_lto)
+                    ModuleBuffer::new(llmod, cgcx.lto != Lto::Fat)
                 };
                 let data = thin.data();
                 let _timer = prof

compiler/rustc_codegen_llvm/src/lib.rs

@@ -153,7 +153,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     type ModuleBuffer = back::lto::ModuleBuffer;
     type TargetMachine = OwnedTargetMachine;
     type ThinData = back::lto::ThinData;
-    type ThinBuffer = back::lto::ThinBuffer;
     fn print_pass_timings(&self) {
         let timings = llvm::build_string(|s| unsafe { llvm::LLVMRustPrintPassTimings(s) }).unwrap();
         print!("{timings}");
@@ -193,7 +192,7 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         dcx: DiagCtxtHandle<'_>,
         exported_symbols_for_lto: &[String],
         each_linked_rlib_for_lto: &[PathBuf],
-        modules: Vec<(String, Self::ThinBuffer)>,
+        modules: Vec<(String, Self::ModuleBuffer)>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
     ) -> (Vec<ThinModule<Self>>, Vec<WorkProduct>) {
         back::lto::run_thin(
@@ -233,11 +232,8 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     ) -> CompiledModule {
         back::write::codegen(cgcx, prof, shared_emitter, module, config)
     }
-    fn prepare_thin(module: ModuleCodegen<Self::Module>) -> (String, Self::ThinBuffer) {
-        back::lto::prepare_thin(module)
-    }
-    fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer) {
-        (module.name, back::lto::ModuleBuffer::new(module.module_llvm.llmod()))
+    fn serialize_module(module: Self::Module, is_thin: bool) -> Self::ModuleBuffer {
+        back::lto::ModuleBuffer::new(module.llmod(), is_thin)
     }
 }