Switch to jemalloc as default?

jguhlin · Jan 7, 2025 · 4f21176 · 4f21176
1 parent cd0a981
commit 4f21176
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 5 deletions.
diff --git a/minimappers2/Cargo.toml b/minimappers2/Cargo.toml
@@ -10,12 +10,20 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 minimap2 = { version = "0.1.23", features = ["simde"], path = ".." }
 crossbeam = "0.8.4"
-mimalloc = {version = "0.1", default-features = false }
-
-pyo3 = { version = "0.22" }
+pyo3 = { version = "0.22", features = ["abi3-py39", "chrono", "extension-module"] }
 polars = "0.45"
 pyo3-polars = "0.19"
 
+[target.'cfg(all(any(not(target_family = "unix"), target_os = "emscripten", allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
+mimalloc = { version = "0.1", default-features = false }
+
+# Feature background_threads is unsupported on MacOS (https://github.com/jemalloc/jemalloc/issues/843).
+[target.'cfg(all(target_family = "unix", not(target_os = "macos"), not(target_os = "emscripten"), not(allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls", "background_threads"] }
+
+[target.'cfg(all(target_family = "unix", target_os = "macos", not(allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
+
 [profile.release]
 opt-level = 3
 lto = "fat"

diff --git a/minimappers2/build.rs b/minimappers2/build.rs
@@ -0,0 +1,7 @@
+fn main() {
+    println!("cargo::rustc-check-cfg=cfg(allocator, values(\"default\", \"mimalloc\"))");
+    println!(
+        "cargo:rustc-env=TARGET={}",
+        std::env::var("TARGET").unwrap()
+    );
+}
diff --git a/minimappers2/src/lib.rs b/minimappers2/src/lib.rs
@@ -2,15 +2,52 @@ use std::num::NonZeroI32;
 use std::sync::{Arc, Mutex};
 
 use crossbeam::queue::ArrayQueue;
-use mimalloc::MiMalloc;
 use minimap2::*;
 
 use polars::{df, prelude::*};
 use pyo3::prelude::*;
 use pyo3_polars::{error::PyPolarsErr, PyDataFrame};
 
+#[cfg(all(
+    target_family = "unix",
+    not(target_os = "emscripten"),
+    not(allocator = "default"),
+    not(allocator = "mimalloc"),
+))]
+use jemallocator::Jemalloc;
+#[cfg(all(
+    not(debug_assertions),
+    not(allocator = "default"),
+    any(
+        not(target_family = "unix"),
+        target_os = "emscripten",
+        allocator = "mimalloc"
+    ),
+))]
+use mimalloc::MiMalloc;
+
 #[global_allocator]
-static GLOBAL: MiMalloc = MiMalloc;
+#[cfg(all(
+    not(debug_assertions),
+    not(allocator = "mimalloc"),
+    not(allocator = "default"),
+    target_family = "unix",
+    not(target_os = "emscripten"),
+))]
+static ALLOC: Jemalloc = Jemalloc;
+
+#[global_allocator]
+#[cfg(all(
+    not(debug_assertions),
+    not(allocator = "default"),
+    any(
+        not(target_family = "unix"),
+        target_os = "emscripten",
+        allocator = "mimalloc"
+    ),
+))]
+static ALLOC: MiMalloc = MiMalloc;
+
 
 mod multithreading;