Skip to content

Allows customization of builtin functions under FatLTO #119348

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 94 additions & 25 deletions compiler/rustc_codegen_llvm/src/back/lto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,22 @@ pub fn crate_type_allows_lto(crate_type: CrateType) -> bool {
}
}

struct SerializedModuleInfo {
module: SerializedModule<ModuleBuffer>,
name: CString,
compiler_builtins: bool,
}

impl SerializedModuleInfo {
fn new(module: SerializedModule<ModuleBuffer>, name: CString, compiler_builtins: bool) -> Self {
Self { module, name, compiler_builtins }
}
}

fn prepare_lto(
cgcx: &CodegenContext<LlvmCodegenBackend>,
dcx: &DiagCtxt,
) -> Result<(Vec<CString>, Vec<(SerializedModule<ModuleBuffer>, CString)>), 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,
Expand All @@ -73,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.
Expand Down Expand Up @@ -127,6 +150,7 @@ fn prepare_lto(
})
})
.filter(|&(name, _)| looks_like_rust_object_file(name));
let compiler_builtins = cgcx.compiler_builtins == Some(cnum);
for (name, child) in obj_files {
info!("adding bitcode from {}", name);
match get_bitcode_slice_from_object_data(
Expand All @@ -135,7 +159,11 @@ fn prepare_lto(
) {
Ok(data) => {
let module = SerializedModule::FromRlib(data.to_vec());
upstream_modules.push((module, CString::new(name).unwrap()));
upstream_modules.push(SerializedModuleInfo::new(
module,
CString::new(name).unwrap(),
compiler_builtins,
));
}
Err(e) => {
dcx.emit_err(e);
Expand All @@ -150,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>(
Expand Down Expand Up @@ -201,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
Expand All @@ -216,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() {
Expand All @@ -239,8 +278,9 @@ fn fat_lto(
dcx: &DiagCtxt,
modules: Vec<FatLtoInput<LlvmCodegenBackend>>,
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
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");
Expand All @@ -258,15 +298,19 @@ fn fat_lto(
let mut in_memory = Vec::new();
serialized_modules.extend(cached_modules.into_iter().map(|(buffer, wp)| {
info!("pushing cached module {:?}", wp.cgu_name);
(buffer, CString::new(wp.cgu_name).unwrap())
SerializedModuleInfo::new(buffer, CString::new(wp.cgu_name).unwrap(), false)
}));
for module in modules {
match module {
FatLtoInput::InMemory(m) => in_memory.push(m),
FatLtoInput::Serialized { name, buffer } => {
info!("pushing serialized module {:?}", name);
let buffer = SerializedModule::Local(buffer);
serialized_modules.push((buffer, CString::new(name).unwrap()));
serialized_modules.push(SerializedModuleInfo::new(
buffer,
CString::new(name).unwrap(),
false,
));
}
}
}
Expand Down Expand Up @@ -299,10 +343,10 @@ fn fat_lto(
Some((_cost, i)) => in_memory.remove(i),
None => {
assert!(!serialized_modules.is_empty(), "must have at least one serialized module");
let (buffer, name) = serialized_modules.remove(0);
let SerializedModuleInfo { module, name, .. } = serialized_modules.remove(0);
info!("no in-memory regular modules to choose from, parsing {:?}", name);
ModuleCodegen {
module_llvm: ModuleLlvm::parse(cgcx, &name, buffer.data(), dcx)?,
module_llvm: ModuleLlvm::parse(cgcx, &name, module.data(), dcx)?,
name: name.into_string().unwrap(),
kind: ModuleKind::Regular,
}
Expand Down Expand Up @@ -330,26 +374,41 @@ fn fat_lto(
for module in in_memory {
let buffer = ModuleBuffer::new(module.module_llvm.llmod());
let llmod_id = CString::new(&module.name[..]).unwrap();
serialized_modules.push((SerializedModule::Local(buffer), llmod_id));
serialized_modules.push(SerializedModuleInfo::new(
SerializedModule::Local(buffer),
llmod_id,
false,
));
}
// Sort the modules to ensure we produce deterministic results.
serialized_modules.sort_by(|module1, module2| module1.1.cmp(&module2.1));
// Place compiler_builtins at the end. After that we check if any crate wants to
// customize the builtin functions. Remove the function of compiler_builtins, if any.
serialized_modules.sort_by(|module1, module2| {
let compiler_builtins_cmp = module1.compiler_builtins.cmp(&module2.compiler_builtins);
if compiler_builtins_cmp.is_ne() {
return compiler_builtins_cmp;
}
module1.name.cmp(&module2.name)
});

// For all serialized bitcode files we parse them and link them in as we did
// 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);
for (bc_decoded, name) in serialized_modules {
let mut linker = Linker::new(llmod, compiler_builtins_exported_symbols);
for serialized_module in serialized_modules {
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 = bc_decoded.data();
linker.add(data).map_err(|()| write::llvm_err(dcx, LlvmError::LoadBitcode { name }))?;
serialized_bitcode.push(bc_decoded);
let data = module.data();
linker
.add(data, compiler_builtins)
.map_err(|()| write::llvm_err(dcx, LlvmError::LoadBitcode { name }))?;
serialized_bitcode.push(module);
}
drop(linker);
save_temp_bitcode(cgcx, &module, "lto.input");
Expand All @@ -372,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 {
Expand Down Expand Up @@ -433,7 +500,7 @@ fn thin_lto(
cgcx: &CodegenContext<LlvmCodegenBackend>,
dcx: &DiagCtxt,
modules: Vec<(String, ThinBuffer)>,
serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
serialized_modules: Vec<SerializedModuleInfo>,
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
symbols_below_threshold: &[*const libc::c_char],
) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
Expand Down Expand Up @@ -479,10 +546,12 @@ fn thin_lto(
// we must always unconditionally look at the index).
let mut serialized = Vec::with_capacity(serialized_modules.len() + cached_modules.len());

let cached_modules =
cached_modules.into_iter().map(|(sm, wp)| (sm, CString::new(wp.cgu_name).unwrap()));
let cached_modules = cached_modules.into_iter().map(|(sm, wp)| {
SerializedModuleInfo::new(sm, CString::new(wp.cgu_name).unwrap(), false)
});

for (module, name) in serialized_modules.into_iter().chain(cached_modules) {
for serialized_module in serialized_modules.into_iter().chain(cached_modules) {
let SerializedModuleInfo { module, name, .. } = serialized_module;
info!("upstream or cached module {:?}", name);
thin_modules.push(llvm::ThinLTOModule {
identifier: name.as_ptr(),
Expand Down
4 changes: 2 additions & 2 deletions compiler/rustc_codegen_llvm/src/back/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
7 changes: 6 additions & 1 deletion compiler/rustc_codegen_llvm/src/llvm/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
2 changes: 2 additions & 0 deletions compiler/rustc_codegen_ssa/src/back/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,7 @@ pub struct CodegenContext<B: WriteBackendMethods> {
pub opts: Arc<config::Options>,
pub crate_types: Vec<CrateType>,
pub each_linked_rlib_for_lto: Vec<(CrateNum, PathBuf)>,
pub compiler_builtins: Option<CrateNum>,
pub output_filenames: Arc<OutputFilenames>,
pub regular_module_config: Arc<ModuleConfig>,
pub metadata_module_config: Arc<ModuleConfig>,
Expand Down Expand Up @@ -1089,6 +1090,7 @@ fn start_executing_work<B: ExtraBackendMethods>(
let cgcx = CodegenContext::<B> {
crate_types: tcx.crate_types().to_vec(),
each_linked_rlib_for_lto,
compiler_builtins: crate_info.compiler_builtins,
lto: sess.lto(),
fewer_names: sess.fewer_names(),
save_temps: sess.opts.cg.save_temps,
Expand Down
Loading