From 0b8115706f2547d256961656fcca5c9b06e90b8a Mon Sep 17 00:00:00 2001 From: Ashwin Naren Date: Sun, 6 Apr 2025 18:51:06 -0700 Subject: [PATCH 1/6] use mimalloc for better performance --- Cargo.lock | 20 ++++++++++++++++++++ Cargo.toml | 3 +++ src/main.rs | 5 +++++ 3 files changed, 28 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 60f0c44ce1..285e84aa3b 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.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b20daca3a4ac14dbdc753c5e90fc7b490a48a9131daed3c9a9ced7b2defd37b" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "libredox" version = "0.1.3" @@ -1410,6 +1420,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "mimalloc" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03cb1f88093fe50061ca1195d336ffec131347c7b833db31f9ab62a2d1b7925f" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -2163,6 +2182,7 @@ dependencies = [ "lexopt", "libc", "log", + "mimalloc", "pyo3", "ruff_python_parser", "rustpython-compiler", diff --git a/Cargo.toml b/Cargo.toml index 271c11e782..54793df639 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,8 @@ dirs = { package = "dirs-next", version = "2.0.0" } env_logger = { version = "0.9.0", default-features = false, features = ["atty", "termcolor"] } flamescope = { version = "0.1.2", optional = true } +mimalloc = "0.1" + [target.'cfg(windows)'.dependencies] libc = { workspace = true } @@ -77,6 +79,7 @@ opt-level = 3 [profile.release] lto = "thin" +debug = true [patch.crates-io] # REDOX START, Uncomment when you want to compile/check with redoxer diff --git a/src/main.rs b/src/main.rs index e88ea40f3d..49cfcaee8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,8 @@ +use mimalloc::MiMalloc; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + pub fn main() -> std::process::ExitCode { rustpython::run(|_vm| {}) } From 645dce20b41520460b0e7c4f15b7ae024f28de40 Mon Sep 17 00:00:00 2001 From: Ashwin Naren Date: Mon, 7 Apr 2025 11:08:17 -0700 Subject: [PATCH 2/6] fix CI --- Cargo.toml | 2 +- src/main.rs | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 54793df639..a59e0f08ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,7 @@ dirs = { package = "dirs-next", version = "2.0.0" } env_logger = { version = "0.9.0", default-features = false, features = ["atty", "termcolor"] } flamescope = { version = "0.1.2", optional = true } +[target.'cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))'.dependencies] mimalloc = "0.1" [target.'cfg(windows)'.dependencies] @@ -79,7 +80,6 @@ opt-level = 3 [profile.release] lto = "thin" -debug = true [patch.crates-io] # REDOX START, Uncomment when you want to compile/check with redoxer diff --git a/src/main.rs b/src/main.rs index 49cfcaee8a..6d8e8a8a9c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,10 @@ -use mimalloc::MiMalloc; +#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] +mod _alloc { + use mimalloc::MiMalloc; -#[global_allocator] -static GLOBAL: MiMalloc = MiMalloc; + #[global_allocator] + static GLOBAL: MiMalloc = MiMalloc; +} pub fn main() -> std::process::ExitCode { rustpython::run(|_vm| {}) From 7d25353f708860125e89f1468f4b246a7ec322cb Mon Sep 17 00:00:00 2001 From: Ashwin Naren Date: Mon, 7 Apr 2025 20:47:23 -0700 Subject: [PATCH 3/6] update cspell --- .cspell.json | 1 + 1 file changed, 1 insertion(+) 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", From 32bfe9e37a87a689b685e644754c1b7cfb4403fa Mon Sep 17 00:00:00 2001 From: Ashwin Naren Date: Mon, 7 Apr 2025 20:50:33 -0700 Subject: [PATCH 4/6] fix exotic targets build --- Cargo.toml | 2 +- src/main.rs | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index a59e0f08ed..1c1f48c641 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,7 +40,7 @@ dirs = { package = "dirs-next", version = "2.0.0" } env_logger = { version = "0.9.0", default-features = false, features = ["atty", "termcolor"] } flamescope = { version = "0.1.2", optional = true } -[target.'cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))'.dependencies] +[target.'cfg(all(any(target_os = "linux", target_os = "macos", target_os = "windows"), not(any(target_env = "musl", target_env = "sgx"))))'.dependencies] mimalloc = "0.1" [target.'cfg(windows)'.dependencies] diff --git a/src/main.rs b/src/main.rs index 6d8e8a8a9c..59064b0d84 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,7 @@ -#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] +#[cfg(all( + any(target_os = "linux", target_os = "macos", target_os = "windows"), + not(any(target_env = "musl", target_env = "sgx")) +))] mod _alloc { use mimalloc::MiMalloc; From 226c1d544730b990d6da0d18d2d1b1a91aaf029e Mon Sep 17 00:00:00 2001 From: Ashwin Naren Date: Tue, 8 Apr 2025 10:13:40 -0700 Subject: [PATCH 5/6] eh --- Cargo.lock | 10 ++-- Cargo.toml | 3 -- src/main.rs | 12 ++--- vm/Cargo.toml | 1 + vm/src/alloc.rs | 129 ++++++++++++++++++++++++++++++++++++++++++++++++ vm/src/lib.rs | 2 + 6 files changed, 140 insertions(+), 17 deletions(-) create mode 100644 vm/src/alloc.rs diff --git a/Cargo.lock b/Cargo.lock index 285e84aa3b..2e4d73d171 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1226,9 +1226,9 @@ checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" [[package]] name = "libmimalloc-sys" -version = "0.1.41" +version = "0.1.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b20daca3a4ac14dbdc753c5e90fc7b490a48a9131daed3c9a9ced7b2defd37b" +checksum = "ec9d6fac27761dabcd4ee73571cdb06b7022dc99089acbe5435691edffaac0f4" dependencies = [ "cc", "libc", @@ -1422,9 +1422,9 @@ dependencies = [ [[package]] name = "mimalloc" -version = "0.1.45" +version = "0.1.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03cb1f88093fe50061ca1195d336ffec131347c7b833db31f9ab62a2d1b7925f" +checksum = "995942f432bbb4822a7e9c3faa87a695185b0d09273ba85f097b54f4e458f2af" dependencies = [ "libmimalloc-sys", ] @@ -2182,7 +2182,6 @@ dependencies = [ "lexopt", "libc", "log", - "mimalloc", "pyo3", "ruff_python_parser", "rustpython-compiler", @@ -2470,6 +2469,7 @@ dependencies = [ "malachite-bigint", "memchr", "memoffset", + "mimalloc", "nix", "num-complex", "num-integer", diff --git a/Cargo.toml b/Cargo.toml index 1c1f48c641..271c11e782 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,9 +40,6 @@ dirs = { package = "dirs-next", version = "2.0.0" } env_logger = { version = "0.9.0", default-features = false, features = ["atty", "termcolor"] } flamescope = { version = "0.1.2", optional = true } -[target.'cfg(all(any(target_os = "linux", target_os = "macos", target_os = "windows"), not(any(target_env = "musl", target_env = "sgx"))))'.dependencies] -mimalloc = "0.1" - [target.'cfg(windows)'.dependencies] libc = { workspace = true } diff --git a/src/main.rs b/src/main.rs index 59064b0d84..05ea300009 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,7 @@ -#[cfg(all( - any(target_os = "linux", target_os = "macos", target_os = "windows"), - not(any(target_env = "musl", target_env = "sgx")) -))] -mod _alloc { - use mimalloc::MiMalloc; +use rustpython_vm::RustPythonAllocator; - #[global_allocator] - static GLOBAL: MiMalloc = MiMalloc; -} +#[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..27981aaa91 --- /dev/null +++ b/vm/src/alloc.rs @@ -0,0 +1,129 @@ +//! Configurable allocator for RustPython. +//! Currently it supports `mimalloc` and `system` allocators, +//! whereas cpython uses `pymalloc`` for most operations. + +#[cfg(all( + any(target_os = "linux", target_os = "macos", target_os = "windows"), + not(any(target_env = "musl", target_env = "sgx")) +))] +mod inner { + use std::alloc::{GlobalAlloc, Layout, System}; + + pub enum RustPythonAllocator { + System(System), + Mimalloc(mimalloc::MiMalloc), + } + + impl RustPythonAllocator { + pub fn new(allocator: &str) -> Self { + match allocator { + "system" | "malloc" => RustPythonAllocator::System(System), + "pymalloc" | "mimalloc" | "default" => { + RustPythonAllocator::Mimalloc(mimalloc::MiMalloc) + } + _ => RustPythonAllocator::System(System), + } + } + } + + unsafe impl GlobalAlloc for RustPythonAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + unsafe { + match self { + RustPythonAllocator::System(system) => system.alloc(layout), + RustPythonAllocator::Mimalloc(mimalloc) => mimalloc.alloc(layout), + } + } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { + match self { + RustPythonAllocator::System(system) => system.dealloc(ptr, layout), + RustPythonAllocator::Mimalloc(mimalloc) => mimalloc.dealloc(ptr, layout), + } + } + } + } +} + +#[cfg(not(all( + any(target_os = "linux", target_os = "macos", target_os = "windows"), + not(any(target_env = "musl", target_env = "sgx")) +)))] +mod inner { + use std::alloc::{GlobalAlloc, Layout, System}; + + pub enum RustPythonAllocator { + System(System), + } + + impl RustPythonAllocator { + pub fn new(_allocator: &str) -> Self { + RustPythonAllocator::System(System) + } + } + + impl GlobalAlloc for RustPythonAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + unsafe { self.0.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { self.0.dealloc(ptr, layout) } + } + } +} + +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::UnsafeCell; + +pub use inner::RustPythonAllocator as InternalAllocator; + +pub struct RustPythonAllocator { + inner: UnsafeCell, +} + +unsafe impl Send for RustPythonAllocator {} +unsafe impl Sync for RustPythonAllocator {} + +impl RustPythonAllocator { + /// Create a new allocator based on the PYTHONMALLOC environment variable + /// or the default allocator if not set. + /// If this is not intended, use [`InternalAllocator::new`] directly. + pub const fn new() -> Self { + Self { + inner: UnsafeCell::new(None), + } + } +} + +impl RustPythonAllocator { + unsafe fn get_or_init(&self) -> &InternalAllocator { + unsafe { + let inner = self.inner.get(); + if *inner.is_none() { + let env = std::env::var("PYTHONMALLOC").unwrap_or_default(); + let allocator = InternalAllocator::new(&env); + *inner = allocator; + } + inner as *const InternalAllocator as _ + } + } +} + +unsafe impl GlobalAlloc for RustPythonAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + unsafe { + let inner = self.get_or_init(); + inner.alloc(layout) + } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { + let inner = self.get_or_init(); + inner.dealloc(ptr, layout) + } + } +} 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::{ From 1418dcd01a631481f494c1bc0aa3a8c1205ccea9 Mon Sep 17 00:00:00 2001 From: Ashwin Naren Date: Fri, 11 Apr 2025 22:53:32 -0700 Subject: [PATCH 6/6] fix alloc --- vm/src/alloc.rs | 201 ++++++++++++++++++++++++------------------------ 1 file changed, 102 insertions(+), 99 deletions(-) diff --git a/vm/src/alloc.rs b/vm/src/alloc.rs index 27981aaa91..dce668f056 100644 --- a/vm/src/alloc.rs +++ b/vm/src/alloc.rs @@ -1,129 +1,132 @@ -//! Configurable allocator for RustPython. -//! Currently it supports `mimalloc` and `system` allocators, -//! whereas cpython uses `pymalloc`` for most operations. - -#[cfg(all( - any(target_os = "linux", target_os = "macos", target_os = "windows"), - not(any(target_env = "musl", target_env = "sgx")) -))] -mod inner { - use std::alloc::{GlobalAlloc, Layout, System}; - - pub enum RustPythonAllocator { - System(System), - Mimalloc(mimalloc::MiMalloc), - } - - impl RustPythonAllocator { - pub fn new(allocator: &str) -> Self { - match allocator { - "system" | "malloc" => RustPythonAllocator::System(System), - "pymalloc" | "mimalloc" | "default" => { - RustPythonAllocator::Mimalloc(mimalloc::MiMalloc) - } - _ => RustPythonAllocator::System(System), - } - } - } +use std::alloc::{Layout, GlobalAlloc}; +use std::cell::UnsafeCell; +use std::ffi::CStr; - unsafe impl GlobalAlloc for RustPythonAllocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - unsafe { - match self { - RustPythonAllocator::System(system) => system.alloc(layout), - RustPythonAllocator::Mimalloc(mimalloc) => mimalloc.alloc(layout), - } - } - } +// Pretend these are the actual allocator impls - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - unsafe { - match self { - RustPythonAllocator::System(system) => system.dealloc(ptr, layout), - RustPythonAllocator::Mimalloc(mimalloc) => mimalloc.dealloc(ptr, layout), - } - } - } +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) } } -#[cfg(not(all( - any(target_os = "linux", target_os = "macos", target_os = "windows"), - not(any(target_env = "musl", target_env = "sgx")) -)))] -mod inner { - use std::alloc::{GlobalAlloc, Layout, System}; +mod malloc_functions { + use std::alloc::System as Malloc; + use std::alloc::Layout; - pub enum RustPythonAllocator { - System(System), + 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 RustPythonAllocator { - pub fn new(_allocator: &str) -> Self { - RustPythonAllocator::System(System) +impl MakeMutable { + const fn new(inner: T) -> Self { + MakeMutable { + inner: UnsafeCell::new(inner), } } - impl GlobalAlloc for RustPythonAllocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - unsafe { self.0.alloc(layout) } - } + fn get(&self) -> &T { + unsafe { &*self.inner.get() } + } - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - unsafe { self.0.dealloc(ptr, layout) } - } + fn get_mut(&self) -> &mut T { + unsafe { &mut *self.inner.get() } } } -use std::alloc::{GlobalAlloc, Layout}; -use std::cell::UnsafeCell; +unsafe impl Sync for MakeMutable {} -pub use inner::RustPythonAllocator as InternalAllocator; +unsafe impl GlobalAlloc for MakeMutable { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ((*self.inner.get()).alloc)(layout) + } -pub struct RustPythonAllocator { - inner: UnsafeCell, -} + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + ((*self.inner.get()).dealloc)(ptr, layout) + } -unsafe impl Send for RustPythonAllocator {} -unsafe impl Sync for RustPythonAllocator {} + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + ((*self.inner.get()).alloc_zeroed)(layout) + } -impl RustPythonAllocator { - /// Create a new allocator based on the PYTHONMALLOC environment variable - /// or the default allocator if not set. - /// If this is not intended, use [`InternalAllocator::new`] directly. - pub const fn new() -> Self { - Self { - inner: UnsafeCell::new(None), - } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + ((*self.inner.get()).realloc)(ptr, layout, new_size) } } -impl RustPythonAllocator { - unsafe fn get_or_init(&self) -> &InternalAllocator { - unsafe { - let inner = self.inner.get(); - if *inner.is_none() { - let env = std::env::var("PYTHONMALLOC").unwrap_or_default(); - let allocator = InternalAllocator::new(&env); - *inner = allocator; - } - inner as *const InternalAllocator as _ +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, } } } -unsafe impl GlobalAlloc for RustPythonAllocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - unsafe { - let inner = self.get_or_init(); - inner.alloc(layout) - } +#[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; } +} - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - unsafe { - let inner = self.get_or_init(); - inner.dealloc(ptr, layout) - } +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; } }