diff --git a/.cspell.json b/.cspell.json index 69ace83483..c04d3c0b92 100644 --- a/.cspell.json +++ b/.cspell.json @@ -69,6 +69,7 @@ "internable", "lossily", "makeunicodedata", + "mimalloc", "miri", "notrace", "openat", diff --git a/Cargo.lock b/Cargo.lock index 60f0c44ce1..2e4d73d171 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1224,6 +1224,16 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" +[[package]] +name = "libmimalloc-sys" +version = "0.1.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec9d6fac27761dabcd4ee73571cdb06b7022dc99089acbe5435691edffaac0f4" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "libredox" version = "0.1.3" @@ -1410,6 +1420,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "mimalloc" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "995942f432bbb4822a7e9c3faa87a695185b0d09273ba85f097b54f4e458f2af" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -2450,6 +2469,7 @@ dependencies = [ "malachite-bigint", "memchr", "memoffset", + "mimalloc", "nix", "num-complex", "num-integer", diff --git a/src/main.rs b/src/main.rs index e88ea40f3d..05ea300009 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,8 @@ +use rustpython_vm::RustPythonAllocator; + +#[global_allocator] +static ALLOCATOR: RustPythonAllocator = RustPythonAllocator::new(); + pub fn main() -> std::process::ExitCode { rustpython::run(|_vm| {}) } diff --git a/vm/Cargo.toml b/vm/Cargo.toml index 125c263da9..23a7ffc619 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -108,6 +108,7 @@ widestring = { workspace = true } [target.'cfg(all(any(target_os = "linux", target_os = "macos", target_os = "windows"), not(any(target_env = "musl", target_env = "sgx"))))'.dependencies] libffi = { workspace = true, features = ["system"] } libloading = "0.8" +mimalloc = "0.1" [target.'cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))'.dependencies] num_cpus = "1.13.1" diff --git a/vm/src/alloc.rs b/vm/src/alloc.rs new file mode 100644 index 0000000000..dce668f056 --- /dev/null +++ b/vm/src/alloc.rs @@ -0,0 +1,132 @@ +use std::alloc::{Layout, GlobalAlloc}; +use std::cell::UnsafeCell; +use std::ffi::CStr; + +// Pretend these are the actual allocator impls + +type GA_Alloc = unsafe fn(layout: Layout) -> *mut u8; +type GA_Dealloc = unsafe fn(ptr: *mut u8, layout: Layout); +type GA_AllocZeroed = unsafe fn(layout: Layout) -> *mut u8; +type GA_Realloc = unsafe fn(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8; + +mod mimalloc_functions { + use std::alloc::Layout; + + use mimalloc::MiMalloc; + + pub unsafe fn alloc(layout: Layout) -> *mut u8 { + MiMalloc::alloc(layout) + } + + pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) { + MiMalloc::dealloc(ptr, layout) + } + + pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 { + MiMalloc::alloc_zeroed(layout) + } + + pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + MiMalloc::realloc(ptr, layout, new_size) + } +} + +mod malloc_functions { + use std::alloc::System as Malloc; + use std::alloc::Layout; + + pub unsafe fn alloc(layout: Layout) -> *mut u8 { + Malloc::alloc(layout) + } + + pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) { + Malloc::dealloc(ptr, layout) + } + + pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 { + Malloc::alloc_zeroed(layout) + } + + pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + Malloc::realloc(ptr, layout, new_size) + } +} + +struct ConfigurableAllocator { + alloc: GA_Alloc, + dealloc: GA_Dealloc, + alloc_zeroed: GA_AllocZeroed, + realloc: GA_Realloc, +} + +struct MakeMutable { + inner: UnsafeCell, +} + +impl MakeMutable { + const fn new(inner: T) -> Self { + MakeMutable { + inner: UnsafeCell::new(inner), + } + } + + fn get(&self) -> &T { + unsafe { &*self.inner.get() } + } + + fn get_mut(&self) -> &mut T { + unsafe { &mut *self.inner.get() } + } +} + +unsafe impl Sync for MakeMutable {} + +unsafe impl GlobalAlloc for MakeMutable { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ((*self.inner.get()).alloc)(layout) + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + ((*self.inner.get()).dealloc)(ptr, layout) + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + ((*self.inner.get()).alloc_zeroed)(layout) + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + ((*self.inner.get()).realloc)(ptr, layout, new_size) + } +} + +impl ConfigurableAllocator { + const fn default() -> Self { + ConfigurableAllocator { + alloc: mimalloc_functions::alloc, + dealloc: mimalloc_functions::dealloc, + alloc_zeroed: mimalloc_functions::alloc_zeroed, + realloc: mimalloc_functions::realloc, + } + } +} + +#[global_allocator] +static ALLOC: MakeMutable = MakeMutable::new(ConfigurableAllocator::default()); + +fn switch_to_malloc() { + unsafe { + (*ALLOC.inner.get()).alloc = malloc_functions::alloc; + (*ALLOC.inner.get()).dealloc = malloc_functions::dealloc; + (*ALLOC.inner.get()).alloc_zeroed = malloc_functions::alloc_zeroed; + (*ALLOC.inner.get()).realloc = malloc_functions::realloc; + } +} + +fn switch_to_mimalloc() { + unsafe { + (*ALLOC.inner.get()).alloc = mimalloc_functions::alloc; + (*ALLOC.inner.get()).dealloc = mimalloc_functions::dealloc; + (*ALLOC.inner.get()).alloc_zeroed = mimalloc_functions::alloc_zeroed; + (*ALLOC.inner.get()).realloc = mimalloc_functions::realloc; + } +} diff --git a/vm/src/lib.rs b/vm/src/lib.rs index de0042c619..c36ac5b1cc 100644 --- a/vm/src/lib.rs +++ b/vm/src/lib.rs @@ -40,6 +40,7 @@ pub use rustpython_derive::*; #[macro_use] pub(crate) mod macros; +mod alloc; mod anystr; pub mod buffer; pub mod builtins; @@ -85,6 +86,7 @@ pub mod warn; #[cfg(windows)] pub mod windows; +pub use self::alloc::RustPythonAllocator; pub use self::compiler::source; pub use self::convert::{TryFromBorrowedObject, TryFromObject}; pub use self::object::{