Skip to content

Commit

Permalink
[WIP] Use weak linkage instead of compiler generated shims
Browse files Browse the repository at this point in the history
This is still keeping the allocator shim for the oom handler and the
__rust_no_alloc_shim_is_unstable symbol for now.

TODO: Update comments everywhere and test on macOS and Windows
  • Loading branch information
bjorn3 committed Dec 20, 2024
1 parent 3bf62cc commit 1a51663
Show file tree
Hide file tree
Showing 15 changed files with 107 additions and 245 deletions.
4 changes: 0 additions & 4 deletions compiler/rustc_ast/src/expand/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,6 @@ pub fn global_fn_name(base: Symbol) -> String {
format!("__rust_{base}")
}

pub fn default_fn_name(base: Symbol) -> String {
format!("__rdl_{base}")
}

pub fn alloc_error_handler_name(alloc_error_handler_kind: AllocatorKind) -> &'static str {
match alloc_error_handler_kind {
AllocatorKind::Global => "__rg_oom",
Expand Down
62 changes: 12 additions & 50 deletions compiler/rustc_codegen_cranelift/src/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,72 +2,34 @@
// Adapted from rustc

use rustc_ast::expand::allocator::{
ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
alloc_error_handler_name, default_fn_name, global_fn_name,
AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, alloc_error_handler_name,
};
use rustc_codegen_ssa::base::allocator_kind_for_codegen;
use rustc_codegen_ssa::base::needs_allocator_shim;
use rustc_session::config::OomStrategy;

use crate::prelude::*;

/// Returns whether an allocator shim was created
pub(crate) fn codegen(tcx: TyCtxt<'_>, module: &mut dyn Module) -> bool {
let Some(kind) = allocator_kind_for_codegen(tcx) else { return false };
codegen_inner(
module,
kind,
tcx.alloc_error_handler_kind(()).unwrap(),
tcx.sess.opts.unstable_opts.oom,
);
true
if needs_allocator_shim(tcx) {
codegen_inner(
module,
tcx.alloc_error_handler_kind(()).unwrap(),
tcx.sess.opts.unstable_opts.oom,
);
true
} else {
false
}
}

fn codegen_inner(
module: &mut dyn Module,
kind: AllocatorKind,
alloc_error_handler_kind: AllocatorKind,
oom_strategy: OomStrategy,
) {
let usize_ty = module.target_config().pointer_type();

if kind == AllocatorKind::Default {
for method in ALLOCATOR_METHODS {
let mut arg_tys = Vec::with_capacity(method.inputs.len());
for input in method.inputs.iter() {
match input.ty {
AllocatorTy::Layout => {
arg_tys.push(usize_ty); // size
arg_tys.push(usize_ty); // align
}
AllocatorTy::Ptr => arg_tys.push(usize_ty),
AllocatorTy::Usize => arg_tys.push(usize_ty),

AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
}
}
let output = match method.output {
AllocatorTy::ResultPtr => Some(usize_ty),
AllocatorTy::Unit => None,

AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
panic!("invalid allocator output")
}
};

let sig = Signature {
call_conv: module.target_config().default_call_conv,
params: arg_tys.iter().cloned().map(AbiParam::new).collect(),
returns: output.into_iter().map(AbiParam::new).collect(),
};
crate::common::create_wrapper_function(
module,
sig,
&global_fn_name(method.name),
&default_fn_name(method.name),
);
}
}

let sig = Signature {
call_conv: module.target_config().default_call_conv,
params: vec![AbiParam::new(usize_ty), AbiParam::new(usize_ty)],
Expand Down
35 changes: 1 addition & 34 deletions compiler/rustc_codegen_gcc/src/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@
use gccjit::FnAttribute;
use gccjit::{Context, FunctionType, GlobalKind, ToRValue, Type};
use rustc_ast::expand::allocator::{
ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
alloc_error_handler_name, default_fn_name, global_fn_name,
AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, alloc_error_handler_name,
};
use rustc_middle::bug;
use rustc_middle::ty::TyCtxt;
Expand All @@ -15,7 +14,6 @@ pub(crate) unsafe fn codegen(
tcx: TyCtxt<'_>,
mods: &mut GccContext,
_module_name: &str,
kind: AllocatorKind,
alloc_error_handler_kind: AllocatorKind,
) {
let context = &mods.context;
Expand All @@ -26,37 +24,6 @@ pub(crate) unsafe fn codegen(
tws => bug!("Unsupported target word size for int: {}", tws),
};
let i8 = context.new_type::<i8>();
let i8p = i8.make_pointer();

if kind == AllocatorKind::Default {
for method in ALLOCATOR_METHODS {
let mut types = Vec::with_capacity(method.inputs.len());
for input in method.inputs.iter() {
match input.ty {
AllocatorTy::Layout => {
types.push(usize);
types.push(usize);
}
AllocatorTy::Ptr => types.push(i8p),
AllocatorTy::Usize => types.push(usize),

AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
}
}
let output = match method.output {
AllocatorTy::ResultPtr => Some(i8p),
AllocatorTy::Unit => None,

AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
panic!("invalid allocator output")
}
};
let from_name = global_fn_name(method.name);
let to_name = default_fn_name(method.name);

create_wrapper_function(tcx, context, &from_name, &to_name, &types, output);
}
}

// FIXME(bjorn3): Add noreturn attribute
create_wrapper_function(
Expand Down
3 changes: 1 addition & 2 deletions compiler/rustc_codegen_gcc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,6 @@ impl ExtraBackendMethods for GccCodegenBackend {
&self,
tcx: TyCtxt<'_>,
module_name: &str,
kind: AllocatorKind,
alloc_error_handler_kind: AllocatorKind,
) -> Self::Module {
let mut mods = GccContext {
Expand All @@ -301,7 +300,7 @@ impl ExtraBackendMethods for GccCodegenBackend {
};

unsafe {
allocator::codegen(tcx, &mut mods, module_name, kind, alloc_error_handler_kind);
allocator::codegen(tcx, &mut mods, module_name, alloc_error_handler_kind);
}
mods
}
Expand Down
36 changes: 1 addition & 35 deletions compiler/rustc_codegen_llvm/src/allocator.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use libc::c_uint;
use rustc_ast::expand::allocator::{
ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
alloc_error_handler_name, default_fn_name, global_fn_name,
AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, alloc_error_handler_name,
};
use rustc_middle::bug;
use rustc_middle::ty::TyCtxt;
Expand All @@ -15,7 +14,6 @@ pub(crate) unsafe fn codegen(
tcx: TyCtxt<'_>,
module_llvm: &mut ModuleLlvm,
module_name: &str,
kind: AllocatorKind,
alloc_error_handler_kind: AllocatorKind,
) {
let llcx = &*module_llvm.llcx;
Expand All @@ -29,38 +27,6 @@ pub(crate) unsafe fn codegen(
}
};
let i8 = unsafe { llvm::LLVMInt8TypeInContext(llcx) };
let i8p = unsafe { llvm::LLVMPointerTypeInContext(llcx, 0) };

if kind == AllocatorKind::Default {
for method in ALLOCATOR_METHODS {
let mut args = Vec::with_capacity(method.inputs.len());
for input in method.inputs.iter() {
match input.ty {
AllocatorTy::Layout => {
args.push(usize); // size
args.push(usize); // align
}
AllocatorTy::Ptr => args.push(i8p),
AllocatorTy::Usize => args.push(usize),

AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
}
}
let output = match method.output {
AllocatorTy::ResultPtr => Some(i8p),
AllocatorTy::Unit => None,

AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
panic!("invalid allocator output")
}
};

let from_name = global_fn_name(method.name);
let to_name = default_fn_name(method.name);

create_wrapper_function(tcx, llcx, llmod, &from_name, &to_name, &args, output, false);
}
}

// rust alloc error handler
create_wrapper_function(
Expand Down
3 changes: 1 addition & 2 deletions compiler/rustc_codegen_llvm/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,12 +114,11 @@ impl ExtraBackendMethods for LlvmCodegenBackend {
&self,
tcx: TyCtxt<'tcx>,
module_name: &str,
kind: AllocatorKind,
alloc_error_handler_kind: AllocatorKind,
) -> ModuleLlvm {
let mut module_llvm = ModuleLlvm::new_metadata(tcx, module_name);
unsafe {
allocator::codegen(tcx, &mut module_llvm, module_name, kind, alloc_error_handler_kind);
allocator::codegen(tcx, &mut module_llvm, module_name, alloc_error_handler_kind);
}
module_llvm
}
Expand Down
4 changes: 2 additions & 2 deletions compiler/rustc_codegen_ssa/src/back/symbol_export.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use rustc_session::config::{CrateType, OomStrategy};
use rustc_target::spec::{SanitizerSet, TlsModel};
use tracing::debug;

use crate::base::allocator_kind_for_codegen;
use crate::base::needs_allocator_shim;

fn threshold(tcx: TyCtxt<'_>) -> SymbolExportLevel {
crates_export_threshold(tcx.crate_types())
Expand Down Expand Up @@ -206,7 +206,7 @@ fn exported_symbols_provider_local(
}

// Mark allocator shim symbols as exported only if they were generated.
if allocator_kind_for_codegen(tcx).is_some() {
if needs_allocator_shim(tcx) {
for symbol_name in ALLOCATOR_METHODS
.iter()
.map(|method| format!("__rust_{}", method.name))
Expand Down
9 changes: 4 additions & 5 deletions compiler/rustc_codegen_ssa/src/base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use std::time::{Duration, Instant};

use itertools::Itertools;
use rustc_abi::FIRST_VARIANT;
use rustc_ast::expand::allocator::{ALLOCATOR_METHODS, AllocatorKind, global_fn_name};
use rustc_ast::expand::allocator::{ALLOCATOR_METHODS, global_fn_name};
use rustc_data_structures::fx::{FxHashMap, FxIndexSet};
use rustc_data_structures::profiling::{get_resident_set_size, print_time_passes_entry};
use rustc_data_structures::sync::{Lrc, par_map};
Expand Down Expand Up @@ -584,7 +584,7 @@ pub fn collect_debugger_visualizers_transitive(
/// Decide allocator kind to codegen. If `Some(_)` this will be the same as
/// `tcx.allocator_kind`, but it may be `None` in more cases (e.g. if using
/// allocator definitions from a dylib dependency).
pub fn allocator_kind_for_codegen(tcx: TyCtxt<'_>) -> Option<AllocatorKind> {
pub fn needs_allocator_shim(tcx: TyCtxt<'_>) -> bool {
// If the crate doesn't have an `allocator_kind` set then there's definitely
// no shim to generate. Otherwise we also check our dependency graph for all
// our output crate types. If anything there looks like its a `Dynamic`
Expand All @@ -595,7 +595,7 @@ pub fn allocator_kind_for_codegen(tcx: TyCtxt<'_>) -> Option<AllocatorKind> {
use rustc_middle::middle::dependency_format::Linkage;
list.iter().any(|&linkage| linkage == Linkage::Dynamic)
});
if any_dynamic_crate { None } else { tcx.allocator_kind(()) }
if any_dynamic_crate { false } else { tcx.allocator_kind(()).is_some() }
}

pub fn codegen_crate<B: ExtraBackendMethods>(
Expand Down Expand Up @@ -664,14 +664,13 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
start_async_codegen(backend.clone(), tcx, target_cpu, metadata, metadata_module);

// Codegen an allocator shim, if necessary.
if let Some(kind) = allocator_kind_for_codegen(tcx) {
if needs_allocator_shim(tcx) {
let llmod_id =
cgu_name_builder.build_cgu_name(LOCAL_CRATE, &["crate"], Some("allocator")).to_string();
let module_llvm = tcx.sess.time("write_allocator_module", || {
backend.codegen_allocator(
tcx,
&llmod_id,
kind,
// If allocator_kind is Some then alloc_error_handler_kind must
// also be Some.
tcx.alloc_error_handler_kind(()).unwrap(),
Expand Down
1 change: 0 additions & 1 deletion compiler/rustc_codegen_ssa/src/traits/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,6 @@ pub trait ExtraBackendMethods:
&self,
tcx: TyCtxt<'tcx>,
module_name: &str,
kind: AllocatorKind,
alloc_error_handler_kind: AllocatorKind,
) -> Self::Module;

Expand Down
2 changes: 2 additions & 0 deletions compiler/rustc_metadata/src/creader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1127,6 +1127,8 @@ fn global_allocator_spans(krate: &ast::Crate) -> Vec<Span> {
fn visit_item(&mut self, item: &'ast ast::Item) {
if item.ident.name == self.name
&& attr::contains_name(&item.attrs, sym::rustc_std_internal_symbol)
// Ignore the default allocator in libstd with weak linkage
&& attr::find_by_name(&item.attrs, sym::linkage).is_none()
{
self.spans.push(item.span);
}
Expand Down
1 change: 1 addition & 0 deletions compiler/rustc_monomorphize/src/partitioning.rs
Original file line number Diff line number Diff line change
Expand Up @@ -887,6 +887,7 @@ fn mono_item_visibility<'tcx>(
// Removal of these functions can't be done by LLVM but rather must be
// done by the linker as it's a non-local decision.
//
// FIXME update comment
// * Second is "std internal symbols". Currently this is primarily used
// for allocator symbols. Allocators are a little weird in their
// implementation, but the idea is that the compiler, at the last
Expand Down
2 changes: 1 addition & 1 deletion library/alloc/src/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ extern "Rust" {
// These are the magic symbols to call the global allocator. rustc generates
// them to call `__rg_alloc` etc. if there is a `#[global_allocator]` attribute
// (the code expanding that attribute macro generates those functions), or to call
// the default implementations in std (`__rdl_alloc` etc. in `library/std/src/alloc.rs`)
// the default implementations in std (weak symbols in `library/std/src/alloc.rs`)
// otherwise.
// The rustc fork of LLVM 14 and earlier also special-cases these function names to be able to optimize them
// like `malloc`, `realloc`, and `free`, respectively.
Expand Down
17 changes: 10 additions & 7 deletions library/std/src/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -382,9 +382,8 @@ pub fn rust_oom(layout: Layout) -> ! {
#[unstable(feature = "alloc_internals", issue = "none")]
pub mod __default_lib_allocator {
use super::{GlobalAlloc, Layout, System};
// These magic symbol names are used as a fallback for implementing the
// `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs`) when there is
// no `#[global_allocator]` attribute.
// These are used as a fallback for implementing the `__rust_alloc`, etc symbols
// (see `src/liballoc/alloc.rs`) when there is no `#[global_allocator]` attribute.

// for symbol names src/librustc_ast/expand/allocator.rs
// for signatures src/librustc_allocator/lib.rs
Expand All @@ -393,7 +392,8 @@ pub mod __default_lib_allocator {
// ABI

#[rustc_std_internal_symbol]
pub unsafe extern "C" fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
#[linkage = "weak"]
pub unsafe extern "Rust" fn __rust_alloc(size: usize, align: usize) -> *mut u8 {
// SAFETY: see the guarantees expected by `Layout::from_size_align` and
// `GlobalAlloc::alloc`.
unsafe {
Expand All @@ -403,14 +403,16 @@ pub mod __default_lib_allocator {
}

#[rustc_std_internal_symbol]
pub unsafe extern "C" fn __rdl_dealloc(ptr: *mut u8, size: usize, align: usize) {
#[linkage = "weak"]
pub unsafe extern "Rust" fn __rust_dealloc(ptr: *mut u8, size: usize, align: usize) {
// SAFETY: see the guarantees expected by `Layout::from_size_align` and
// `GlobalAlloc::dealloc`.
unsafe { System.dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
}

#[rustc_std_internal_symbol]
pub unsafe extern "C" fn __rdl_realloc(
#[linkage = "weak"]
pub unsafe extern "Rust" fn __rust_realloc(
ptr: *mut u8,
old_size: usize,
align: usize,
Expand All @@ -425,7 +427,8 @@ pub mod __default_lib_allocator {
}

#[rustc_std_internal_symbol]
pub unsafe extern "C" fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
#[linkage = "weak"]
pub unsafe extern "Rust" fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
// SAFETY: see the guarantees expected by `Layout::from_size_align` and
// `GlobalAlloc::alloc_zeroed`.
unsafe {
Expand Down
Loading

0 comments on commit 1a51663

Please sign in to comment.