Allows customization of builtin functions under FatLTO

dianqk · dianqk · commit 49f62dfcd37b · 2023-12-27T22:36:07.000+08:00
diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -60,7 +60,7 @@ impl SerializedModuleInfo {
 fn prepare_lto(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     dcx: &DiagCtxt,
-) -> Result<(Vec<CString>, Vec<SerializedModuleInfo>), FatalError> {
+) -> Result<(Vec<CString>, Vec<SerializedModuleInfo>, Vec<CString>), FatalError> {
     let export_threshold = match cgcx.lto {
         // We're just doing LTO for our one crate
         Lto::ThinLocal => SymbolExportLevel::Rust,
@@ -85,6 +85,17 @@ fn prepare_lto(
     };
     info!("{} symbols to preserve in this crate", symbols_below_threshold.len());
 
+    let compiler_builtins_exported_symbols = match cgcx.compiler_builtins {
+        Some(crate_num) => {
+            if let Some(exported_symbols) = exported_symbols.get(&crate_num) {
+                exported_symbols.iter().filter_map(symbol_filter).collect::<Vec<CString>>()
+            } else {
+                Vec::new()
+            }
+        }
+        None => Vec::new(),
+    };
+
     // If we're performing LTO for the entire crate graph, then for each of our
     // upstream dependencies, find the corresponding rlib and load the bitcode
     // from the archive.
@@ -167,7 +178,7 @@ fn prepare_lto(
     // __llvm_profile_runtime, therefore we won't know until link time if this symbol
     // should have default visibility.
     symbols_below_threshold.push(CString::new("__llvm_profile_counter_bias").unwrap());
-    Ok((symbols_below_threshold, upstream_modules))
+    Ok((symbols_below_threshold, upstream_modules, compiler_builtins_exported_symbols))
 }
 
 fn get_bitcode_slice_from_object_data<'a>(
@@ -218,10 +229,21 @@ pub(crate) fn run_fat(
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
 ) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
     let dcx = cgcx.create_dcx();
-    let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, &dcx)?;
+    let (symbols_below_threshold, upstream_modules, compiler_builtins_exported_symbols) =
+        prepare_lto(cgcx, &dcx)?;
     let symbols_below_threshold =
         symbols_below_threshold.iter().map(|c| c.as_ptr()).collect::<Vec<_>>();
-    fat_lto(cgcx, &dcx, modules, cached_modules, upstream_modules, &symbols_below_threshold)
+    let compiler_builtins_exported_symbols =
+        compiler_builtins_exported_symbols.iter().map(|c| c.as_ptr()).collect::<Vec<_>>();
+    fat_lto(
+        cgcx,
+        &dcx,
+        modules,
+        cached_modules,
+        upstream_modules,
+        &symbols_below_threshold,
+        &compiler_builtins_exported_symbols,
+    )
 }
 
 /// Performs thin LTO by performing necessary global analysis and returning two
@@ -233,7 +255,7 @@ pub(crate) fn run_thin(
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
 ) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
     let dcx = cgcx.create_dcx();
-    let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, &dcx)?;
+    let (symbols_below_threshold, upstream_modules, _) = prepare_lto(cgcx, &dcx)?;
     let symbols_below_threshold =
         symbols_below_threshold.iter().map(|c| c.as_ptr()).collect::<Vec<_>>();
     if cgcx.opts.cg.linker_plugin_lto.enabled() {
@@ -258,6 +280,7 @@ fn fat_lto(
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
     mut serialized_modules: Vec<SerializedModuleInfo>,
     symbols_below_threshold: &[*const libc::c_char],
+    compiler_builtins_exported_symbols: &[*const libc::c_char],
 ) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
     let _timer = cgcx.prof.generic_activity("LLVM_fat_lto_build_monolithic_module");
     info!("going for a fat lto");
@@ -372,17 +395,19 @@ fn fat_lto(
         // above, this is all mostly handled in C++. Like above, though, we don't
         // know much about the memory management here so we err on the side of being
         // save and persist everything with the original module.
-        let mut linker = Linker::new(llmod);
+        let mut linker = Linker::new(llmod, compiler_builtins_exported_symbols);
         for serialized_module in serialized_modules {
-            let SerializedModuleInfo { module, name, .. } = serialized_module;
+            let SerializedModuleInfo { module, name, compiler_builtins } = serialized_module;
             let _timer = cgcx
                 .prof
                 .generic_activity_with_arg_recorder("LLVM_fat_lto_link_module", |recorder| {
                     recorder.record_arg(format!("{name:?}"))
                 });
             info!("linking {:?}", name);
             let data = module.data();
-            linker.add(data).map_err(|()| write::llvm_err(dcx, LlvmError::LoadBitcode { name }))?;
+            linker
+                .add(data, compiler_builtins)
+                .map_err(|()| write::llvm_err(dcx, LlvmError::LoadBitcode { name }))?;
             serialized_bitcode.push(module);
         }
         drop(linker);
@@ -406,16 +431,24 @@ fn fat_lto(
 pub(crate) struct Linker<'a>(&'a mut llvm::Linker<'a>);
 
 impl<'a> Linker<'a> {
-    pub(crate) fn new(llmod: &'a llvm::Module) -> Self {
-        unsafe { Linker(llvm::LLVMRustLinkerNew(llmod)) }
+    pub(crate) fn new(llmod: &'a llvm::Module, builtin_syms: &[*const libc::c_char]) -> Self {
+        let ptr = builtin_syms.as_ptr();
+        unsafe {
+            Linker(llvm::LLVMRustLinkerNew(
+                llmod,
+                ptr as *const *const libc::c_char,
+                builtin_syms.len() as libc::size_t,
+            ))
+        }
     }
 
-    pub(crate) fn add(&mut self, bytecode: &[u8]) -> Result<(), ()> {
+    pub(crate) fn add(&mut self, bytecode: &[u8], compiler_builtins: bool) -> Result<(), ()> {
         unsafe {
             if llvm::LLVMRustLinkerAdd(
                 self.0,
                 bytecode.as_ptr() as *const libc::c_char,
                 bytecode.len(),
+                compiler_builtins,
             ) {
                 Ok(())
             } else {
diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs
@@ -634,12 +634,12 @@ pub(crate) fn link(
     let (first, elements) =
         modules.split_first().expect("Bug! modules must contain at least one module.");
 
-    let mut linker = Linker::new(first.module_llvm.llmod());
+    let mut linker = Linker::new(first.module_llvm.llmod(), &[]);
     for module in elements {
         let _timer = cgcx.prof.generic_activity_with_arg("LLVM_link_module", &*module.name);
         let buffer = ModuleBuffer::new(module.module_llvm.llmod());
         linker
-            .add(buffer.data())
+            .add(buffer.data(), false)
             .map_err(|()| llvm_err(dcx, LlvmError::SerializeModule { name: &module.name }))?;
     }
     drop(linker);
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -2356,11 +2356,16 @@ extern "C" {
         out_len: &mut usize,
     ) -> *const u8;
 
-    pub fn LLVMRustLinkerNew(M: &Module) -> &mut Linker<'_>;
+    pub fn LLVMRustLinkerNew(
+        M: &Module,
+        builtin_syms: *const *const c_char,
+        len: size_t,
+    ) -> &mut Linker<'_>;
     pub fn LLVMRustLinkerAdd(
         linker: &Linker<'_>,
         bytecode: *const c_char,
         bytecode_len: usize,
+        compiler_builtins: bool,
     ) -> bool;
     pub fn LLVMRustLinkerFree<'a>(linker: &'a mut Linker<'a>);
     #[allow(improper_ctypes)]
diff --git a/compiler/rustc_llvm/llvm-wrapper/Linker.cpp b/compiler/rustc_llvm/llvm-wrapper/Linker.cpp
@@ -14,6 +14,7 @@
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/Linker/IRMover.h"
+#include "llvm/Object/ModuleSymbolTable.h"
 #include "llvm/Support/Error.h"
 
 #include "LLVMWrapper.h"
@@ -44,7 +45,10 @@ enum class LinkFrom { Dst, Src, Both };
 /// entrypoint for this file.
 class ModuleLinker {
   IRMover &Mover;
+  const StringSet<> &CompilerBuiltinsSymbols;
+  StringSet<> UserBuiltinsSymbols;
   std::unique_ptr<Module> SrcM;
+  bool SrcIsCompilerBuiltins;
 
   SetVector<GlobalValue *> ValuesToLink;
 
@@ -122,11 +126,14 @@ class ModuleLinker {
   bool linkIfNeeded(GlobalValue &GV, SmallVectorImpl<GlobalValue *> &GVToClone);
 
 public:
-  ModuleLinker(IRMover &Mover, std::unique_ptr<Module> SrcM, unsigned Flags,
+  ModuleLinker(IRMover &Mover, const StringSet<> &CompilerBuiltinsSymbols,
+               std::unique_ptr<Module> SrcM, bool SrcIsCompilerBuiltins,
+               unsigned Flags,
                std::function<void(Module &, const StringSet<> &)>
                    InternalizeCallback = {})
-      : Mover(Mover), SrcM(std::move(SrcM)), Flags(Flags),
-        InternalizeCallback(std::move(InternalizeCallback)) {}
+      : Mover(Mover), CompilerBuiltinsSymbols(CompilerBuiltinsSymbols),
+        SrcM(std::move(SrcM)), SrcIsCompilerBuiltins(SrcIsCompilerBuiltins),
+        Flags(Flags), InternalizeCallback(std::move(InternalizeCallback)) {}
 
   bool run();
 };
@@ -342,6 +349,10 @@ bool ModuleLinker::shouldLinkFromSource(bool &LinkFromSrc,
 
 bool ModuleLinker::linkIfNeeded(GlobalValue &GV,
                                 SmallVectorImpl<GlobalValue *> &GVToClone) {
+  // If a builtin symbol is defined in a non-compiler-builtins, the symbol of
+  // compiler-builtins is a non-prevailing symbol.
+  if (SrcIsCompilerBuiltins && UserBuiltinsSymbols.contains(GV.getName()))
+    return false;
   GlobalValue *DGV = getLinkedToGlobal(&GV);
 
   if (shouldLinkOnlyNeeded()) {
@@ -501,6 +512,27 @@ bool ModuleLinker::run() {
     ReplacedDstComdats.insert(DstC);
   }
 
+  if (SrcIsCompilerBuiltins) {
+    ModuleSymbolTable SymbolTable;
+    SymbolTable.addModule(&DstM);
+    for (auto &Sym : SymbolTable.symbols()) {
+      uint32_t Flags = SymbolTable.getSymbolFlags(Sym);
+      if ((Flags & object::BasicSymbolRef::SF_Weak) ||
+          !(Flags & object::BasicSymbolRef::SF_Global))
+        continue;
+      if (GlobalValue *GV = dyn_cast_if_present<GlobalValue *>(Sym)) {
+        if (CompilerBuiltinsSymbols.contains(GV->getName()))
+          UserBuiltinsSymbols.insert(GV->getName());
+      } else if (auto *AS =
+                     dyn_cast_if_present<ModuleSymbolTable::AsmSymbol *>(Sym)) {
+        if (CompilerBuiltinsSymbols.contains(AS->first))
+          UserBuiltinsSymbols.insert(AS->first);
+      } else {
+        llvm::report_fatal_error("unknown symbol type");
+      }
+    }
+  }
+
   // Alias have to go first, since we are not able to find their comdats
   // otherwise.
   for (GlobalAlias &GV : llvm::make_early_inc_range(DstM.aliases()))
@@ -617,6 +649,7 @@ namespace {
 struct RustLinker {
   IRMover Mover;
   LLVMContext &Ctx;
+  StringSet<> CompilerBuiltinsSymbols;
 
   enum Flags {
     None = 0,
@@ -634,37 +667,44 @@ struct RustLinker {
   /// callback.
   ///
   /// Returns true on error.
-  bool linkInModule(std::unique_ptr<Module> Src, unsigned Flags = Flags::None,
+  bool linkInModule(std::unique_ptr<Module> Src, bool SrcIsCompilerBuiltins,
+                    unsigned Flags = Flags::None,
                     std::function<void(Module &, const StringSet<> &)>
                         InternalizeCallback = {});
 
-  RustLinker(Module &M) : Mover(M), Ctx(M.getContext()) {}
+  RustLinker(Module &M, StringSet<> CompilerBuiltinsSymbols)
+      : Mover(M), Ctx(M.getContext()),
+        CompilerBuiltinsSymbols(CompilerBuiltinsSymbols) {}
 };
 
 } // namespace
 
 bool RustLinker::linkInModule(
-    std::unique_ptr<Module> Src, unsigned Flags,
+    std::unique_ptr<Module> Src, bool SrcIsCompilerBuiltins, unsigned Flags,
     std::function<void(Module &, const StringSet<> &)> InternalizeCallback) {
-  ModuleLinker ModLinker(Mover, std::move(Src), Flags,
+  ModuleLinker ModLinker(Mover, CompilerBuiltinsSymbols, std::move(Src),
+                         SrcIsCompilerBuiltins, Flags,
                          std::move(InternalizeCallback));
   return ModLinker.run();
 }
 
-extern "C" RustLinker*
-LLVMRustLinkerNew(LLVMModuleRef DstRef) {
+extern "C" RustLinker *LLVMRustLinkerNew(LLVMModuleRef DstRef, char **Symbols,
+                                         size_t Len) {
   Module *Dst = unwrap(DstRef);
-
-  return new RustLinker(*Dst);
+  StringSet<> CompilerBuiltinsSymbols;
+  for (size_t I = 0; I < Len; I++) {
+    CompilerBuiltinsSymbols.insert(Symbols[I]);
+  }
+  return new RustLinker(*Dst, CompilerBuiltinsSymbols);
 }
 
 extern "C" void
 LLVMRustLinkerFree(RustLinker *L) {
   delete L;
 }
 
-extern "C" bool
-LLVMRustLinkerAdd(RustLinker *L, char *BC, size_t Len) {
+extern "C" bool LLVMRustLinkerAdd(RustLinker *L, char *BC, size_t Len,
+                                  bool CompilerBuiltins) {
   std::unique_ptr<MemoryBuffer> Buf =
       MemoryBuffer::getMemBufferCopy(StringRef(BC, Len));
 
@@ -677,7 +717,7 @@ LLVMRustLinkerAdd(RustLinker *L, char *BC, size_t Len) {
 
   auto Src = std::move(*SrcOrError);
 
-  if (L->linkInModule(std::move(Src))) {
+  if (L->linkInModule(std::move(Src), CompilerBuiltins)) {
     LLVMRustSetLastError("");
     return false;
   }
diff --git a/tests/assembly/lto-custom-builtins.rs b/tests/assembly/lto-custom-builtins.rs
@@ -0,0 +1,45 @@
+// assembly-output: emit-asm
+// compile-flags: --crate-type cdylib -C lto=fat -C prefer-dynamic=no
+// only-x86_64-unknown-linux-gnu
+
+#![feature(lang_items, linkage)]
+#![no_std]
+#![no_main]
+
+#![crate_type = "bin"]
+
+// We want to use customized __subdf3.
+// CHECK: .globl __subdf3
+// CHECK-NEXT: __subdf3:
+// CHECK-NEXT: movq $2, %rax
+core::arch::global_asm!(".global __subdf3", "__subdf3:", "mov rax, 2");
+
+// We want to use __addsf3 of compiler-builtins.
+// CHECK: .globl __addsf3
+// CHECK: __addsf3:
+// CHECK: xorl %eax, %eax
+// CHECK-NEXT retq
+#[no_mangle]
+pub extern "C" fn __addsf3() -> i32 {
+    0
+}
+
+// We want to use customized __adddf3.
+// CHECK: .globl __adddf3
+// CHECK: __adddf3:
+// CHECK-NEXT: .cfi_startproc
+// CHECK-NOT: movl $1, %eax
+// CHECK: movq %xmm0, %rdx
+#[no_mangle]
+#[linkage = "weak"]
+pub extern "C" fn __adddf3() -> i32 {
+    1
+}
+
+#[panic_handler]
+fn panic(_panic: &core::panic::PanicInfo<'_>) -> ! {
+    loop {}
+}
+
+#[lang = "eh_personality"]
+extern "C" fn eh_personality() {}
