diff --git a/Cargo.lock b/Cargo.lock index 8ac10d7d4fb53..494535542f3a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -756,7 +756,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.2", ] [[package]] @@ -3837,6 +3837,8 @@ dependencies = [ "rustc_target", "serde_json", "shlex", + "tikv-jemalloc-sys", + "tikv-jemallocator", "tracing", "windows 0.61.3", ] @@ -5551,6 +5553,16 @@ dependencies = [ "libc", ] +[[package]] +name = "tikv-jemallocator" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "tinystr" version = "0.8.2" diff --git a/compiler/rustc/Cargo.toml b/compiler/rustc/Cargo.toml index 8cef9e0644bb0..1886ff7c27e7a 100644 --- a/compiler/rustc/Cargo.toml +++ b/compiler/rustc/Cargo.toml @@ -30,12 +30,11 @@ wasi = "=0.14.2" [dependencies.tikv-jemalloc-sys] version = "0.6.1" optional = true -features = ['override_allocator_on_supported_platforms'] [features] # tidy-alphabetical-start check_only = ['rustc_driver_impl/check_only'] -jemalloc = ['dep:tikv-jemalloc-sys'] +jemalloc = ['dep:tikv-jemalloc-sys', 'rustc_driver_impl/jemalloc'] llvm = ['rustc_driver_impl/llvm'] llvm_enzyme = ['rustc_driver_impl/llvm_enzyme'] llvm_offload = ['rustc_driver_impl/llvm_offload'] diff --git a/compiler/rustc/src/main.rs b/compiler/rustc/src/main.rs index 30d64b05cfde9..ab1918a428901 100644 --- a/compiler/rustc/src/main.rs +++ b/compiler/rustc/src/main.rs @@ -37,8 +37,68 @@ use std::process::ExitCode; // to compare their performance, see // https://github.com/rust-lang/rust/commit/b90cfc887c31c3e7a9e6d462e2464db1fe506175#diff-43914724af6e464c1da2171e4a9b6c7e607d5bc1203fa95c0ab85be4122605ef // for an example of how to do so. + #[cfg(feature = "jemalloc")] -use tikv_jemalloc_sys as _; +mod c_alloc { + use std::os::raw::{c_int, c_void}; + #[used] + static _F1: unsafe extern "C" fn(usize, usize) -> *mut c_void = calloc; + #[used] + static _F2: unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c_int = posix_memalign; + #[used] + static _F3: unsafe extern "C" fn(usize, usize) -> *mut c_void = aligned_alloc; + #[used] + static _F4: unsafe extern "C" fn(usize) -> *mut c_void = malloc; + #[used] + static _F5: unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void = realloc; + #[used] + static _F6: unsafe extern "C" fn(*mut c_void) = free; + + // On OSX, jemalloc doesn't directly override malloc/free, but instead + // registers itself with the allocator's zone APIs in a ctor. However, + // the linker doesn't seem to consider ctors as "used" when statically + // linking, so we need to explicitly depend on the function. + #[cfg(target_os = "macos")] + unsafe extern "C" { + fn _rjem_je_zone_register(); + } + + #[used] + #[cfg(target_os = "macos")] + static _F7: unsafe extern "C" fn() = _rjem_je_zone_register; + + #[unsafe(no_mangle)] + unsafe extern "C" fn calloc(items: usize, size: usize) -> *mut c_void { + unsafe { rustc_driver::jemalloc::calloc(items, size) } + } + + #[unsafe(no_mangle)] + unsafe extern "C" fn posix_memalign(ptr: *mut *mut c_void, size: usize, align: usize) -> c_int { + unsafe { rustc_driver::jemalloc::posix_memalign(ptr, size, align) } + } + + #[unsafe(no_mangle)] + unsafe extern "C" fn aligned_alloc(size: usize, align: usize) -> *mut c_void { + unsafe { rustc_driver::jemalloc::aligned_alloc(size, align) } + } + + #[unsafe(no_mangle)] + unsafe extern "C" fn malloc(size: usize) -> *mut c_void { + unsafe { rustc_driver::jemalloc::malloc(size) } + } + + #[unsafe(no_mangle)] + unsafe extern "C" fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void { + unsafe { rustc_driver::jemalloc::realloc(ptr, size) } + } + + #[unsafe(no_mangle)] + unsafe extern "C" fn free(ptr: *mut c_void) { + unsafe { + rustc_driver::jemalloc::free(ptr); + } + } +} fn main() -> ExitCode { rustc_driver::main() diff --git a/compiler/rustc_driver_impl/Cargo.toml b/compiler/rustc_driver_impl/Cargo.toml index 10e1b1987194a..fe9243ce62b5b 100644 --- a/compiler/rustc_driver_impl/Cargo.toml +++ b/compiler/rustc_driver_impl/Cargo.toml @@ -55,9 +55,18 @@ features = [ ctrlc = "3.4.4" # tidy-alphabetical-end +[dependencies.tikv-jemallocator] +version = "0.6.1" +optional = true + +[dependencies.tikv-jemalloc-sys] +version = "0.6.1" +optional = true + [features] # tidy-alphabetical-start check_only = ['rustc_interface/check_only'] +jemalloc = ['dep:tikv-jemallocator', 'dep:tikv-jemalloc-sys'] llvm = ['rustc_interface/llvm'] llvm_enzyme = ['rustc_interface/llvm_enzyme'] llvm_offload = ['rustc_interface/llvm_offload'] diff --git a/compiler/rustc_driver_impl/src/lib.rs b/compiler/rustc_driver_impl/src/lib.rs index c15c3c229398c..dda7397e3496b 100644 --- a/compiler/rustc_driver_impl/src/lib.rs +++ b/compiler/rustc_driver_impl/src/lib.rs @@ -63,6 +63,44 @@ use rustc_target::json::ToJson; use rustc_target::spec::{Target, TargetTuple}; use tracing::trace; +#[cfg(feature = "jemalloc")] +pub mod jemalloc { + use std::os::raw::{c_int, c_void}; + + #[global_allocator] + static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + + pub unsafe extern "C" fn calloc(items: usize, size: usize) -> *mut c_void { + unsafe { tikv_jemalloc_sys::calloc(items, size) } + } + + pub unsafe extern "C" fn posix_memalign( + ptr: *mut *mut c_void, + size: usize, + align: usize, + ) -> c_int { + unsafe { tikv_jemalloc_sys::posix_memalign(ptr, size, align) } + } + + pub unsafe extern "C" fn aligned_alloc(size: usize, align: usize) -> *mut c_void { + unsafe { tikv_jemalloc_sys::aligned_alloc(size, align) } + } + + pub unsafe extern "C" fn malloc(size: usize) -> *mut c_void { + unsafe { tikv_jemalloc_sys::malloc(size) } + } + + pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void { + unsafe { tikv_jemalloc_sys::realloc(ptr, size) } + } + + pub unsafe extern "C" fn free(ptr: *mut c_void) { + unsafe { + tikv_jemalloc_sys::free(ptr); + } + } +} + #[allow(unused_macros)] macro do_not_use_print($($t:tt)*) { std::compile_error!( diff --git a/src/tools/tidy/src/deps.rs b/src/tools/tidy/src/deps.rs index 1d12456d2bb23..60a3de06f7272 100644 --- a/src/tools/tidy/src/deps.rs +++ b/src/tools/tidy/src/deps.rs @@ -442,6 +442,7 @@ const PERMITTED_RUSTC_DEPENDENCIES: &[&str] = &[ "thorin-dwp", "thread_local", "tikv-jemalloc-sys", + "tikv-jemallocator", "tinystr", "tinyvec", "tinyvec_macros",