Native libraries use mimalloc as global allocator

.gitignore

@@ -15,3 +15,5 @@ __pycache__/
 /fuzz/artifacts
 /fuzz/coverage
 /fuzz/Cargo.lock
+.mypy_cache/
+.pytest_cache/

Cargo.toml

@@ -1,5 +1,7 @@
 [workspace]
 members = [
+    "allocator",
+    "allocator/mimalloc-sys",
     "compiler/qsc",
     "compiler/qsc_ast",
     "compiler/qsc_codegen",

README.md

@@ -23,12 +23,13 @@ Code from this repository powers the Q# development experience on <https://quant
 
 ## Building
 
-To build this repository there are 4 dependencies that need to be installed. These are:
+To build this repository there are dependencies that need to be installed. These are:
 
 - Python (<https://python.org>)
 - Rust (<https://www.rust-lang.org/tools/install>)
 - Node.js (<https://nodejs.org/>)
 - wasm-pack (<https://rustwasm.github.io/wasm-pack/installer/>)
+- cmake (<https://cmake.org/>) and a C compiler
 
 The build script will check these dependencies and their versions and fail if not met. (Or run
 `python ./prereqs.py` directly to check if the minimum required versions are installed).

allocator/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "allocator"
+authors.workspace = true
+homepage.workspace = true
+repository.workspace = true
+edition.workspace = true
+license.workspace = true
+version.workspace = true
+
+[dependencies]
+mimalloc-sys = { path = "./mimalloc-sys" }
+
+[lints]
+workspace = true
+

allocator/mimalloc-sys/CMakeLists.txt

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cmake_minimum_required(VERSION 3.10.0)
+
+
+project(allocator_external)
+include(ExternalProject)
+
+ExternalProject_Add(mimalloc
+  GIT_REPOSITORY    https://github.com/microsoft/mimalloc.git
+  GIT_TAG           $ENV{ALLOCATOR_MIMALLOC_TAG}
+  GIT_SHALLOW       TRUE
+  GIT_PROGRESS      TRUE
+  CONFIGURE_COMMAND ""
+  BUILD_COMMAND ""
+  INSTALL_COMMAND ""
+  TEST_COMMAND ""
+  USES_TERMINAL_DOWNLOAD  TRUE
+)

allocator/mimalloc-sys/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "mimalloc-sys"
+build = "build.rs"
+links = "mimalloc"
+authors.workspace = true
+homepage.workspace = true
+repository.workspace = true
+edition.workspace = true
+license.workspace = true
+version.workspace = true
+
+[dependencies]
+
+[lints]
+workspace = true
+
+[build-dependencies]
+cmake = "0.1"
+cc = "1.0"
+

allocator/mimalloc-sys/build.rs

@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use std::boxed::Box;
+use std::env;
+use std::error::Error;
+use std::fs;
+use std::path::{Path, PathBuf};
+
+use cmake::Config;
+
+// 1.8.2
+//static ALLOCATOR_MIMALLOC_TAG: &str = "b66e3214d8a104669c2ec05ae91ebc26a8f5ab78";
+// 2.1.2
+static ALLOCATOR_MIMALLOC_TAG: &str = "43ce4bd7fd34bcc730c1c7471c99995597415488";
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let dst = download_mimalloc()?;
+    compile_mimalloc(&dst);
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-changed=CMakeLists.txt");
+    Ok(())
+}
+
+// Compile mimalloc source code and link it to the crate.
+// The cc crate is used to compile the source code into a static library.
+// The cmake crate is used to download the source code and stage it in the build directory.
+// We don't use the cmake crate to compile the source code because the mimalloc build system
+// loads extra libraries, changes the name and path around, and does other things that are
+// difficult to handle. The cc crate is much simpler and more predictable.
+fn compile_mimalloc(dst: &Path) {
+    let src_dir = dst
+        .join("build")
+        .join("mimalloc-prefix")
+        .join("src")
+        .join("mimalloc");
+
+    let mut build = cc::Build::new();
+
+    build.include(src_dir.join("include"));
+    build.include(src_dir.join("src"));
+    build.file(src_dir.join("src/static.c"));
+
+    if build.get_compiler().is_like_msvc() {
+        build.cpp(true);
+        build.static_crt(true);
+    }
+    // turn off debug mode
+    build.define("MI_DEBUG", "0");
+
+    // turning on optimizations doesn't seem to make a difference
+    //build.opt_level(3);
+
+    build.compile("mimalloc");
+
+    println!(
+        "cargo:rustc-link-search=native={}",
+        dst.join("lib").display()
+    );
+    println!("cargo:rustc-link-lib=static=mimalloc");
+}
+
+// Use cmake to download mimallloc source code and stage
+// it in the build directory.
+fn download_mimalloc() -> Result<PathBuf, Box<dyn Error>> {
+    let build_dir = get_build_dir()?;
+    let mut config = Config::new(build_dir);
+
+    config
+        .no_build_target(true)
+        .env("ALLOCATOR_MIMALLOC_TAG", ALLOCATOR_MIMALLOC_TAG)
+        .very_verbose(true);
+
+    let dst = config.build();
+
+    Ok(dst)
+}
+
+fn get_build_dir() -> Result<PathBuf, Box<dyn Error>> {
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR")?;
+    let build_dir = PathBuf::from(manifest_dir.as_str());
+    let normalized_build_dir = fs::canonicalize(build_dir)?;
+    Ok(normalized_build_dir)
+}

allocator/mimalloc-sys/src/lib.rs

@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use core::ffi::c_void;
+pub static MI_ALIGNMENT_MAX: usize = 1024 * 1024; // 1 MiB
+
+extern "C" {
+    /// Allocate size bytes aligned by alignment.
+    /// size: the number of bytes to allocate
+    /// alignment: the minimal alignment of the allocated memory. Must be less than MI_ALIGNMENT_MAX
+    /// returns: a pointer to the allocated memory, or null if out of memory. The returned pointer is aligned by alignment
+    pub fn mi_malloc_aligned(size: usize, alignment: usize) -> *mut c_void;
+    pub fn mi_zalloc_aligned(size: usize, alignment: usize) -> *mut c_void;
+
+    /// Free previously allocated memory.
+    /// The pointer p must have been allocated before (or be nullptr).
+    /// p: the pointer to the memory to free or nullptr
+    pub fn mi_free(p: *mut c_void);
+    pub fn mi_realloc_aligned(p: *mut c_void, newsize: usize, alignment: usize) -> *mut c_void;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn memory_can_be_allocated_and_freed() {
+        let ptr = unsafe { mi_malloc_aligned(8, 8) }.cast::<u8>();
+        assert!(!ptr.cast::<c_void>().is_null());
+        unsafe { mi_free(ptr.cast::<c_void>()) };
+    }
+
+    #[test]
+    fn memory_can_be_allocated_zeroed_and_freed() {
+        let ptr = unsafe { mi_zalloc_aligned(8, 8) }.cast::<u8>();
+        assert!(!ptr.cast::<c_void>().is_null());
+        unsafe { mi_free(ptr.cast::<c_void>()) };
+    }
+
+    #[test]
+    fn memory_can_be_reallocated_and_freed() {
+        let ptr = unsafe { mi_malloc_aligned(8, 8) }.cast::<u8>();
+        assert!(!ptr.cast::<c_void>().is_null());
+        let realloc_ptr = unsafe { mi_realloc_aligned(ptr.cast::<c_void>(), 8, 8) }.cast::<u8>();
+        assert!(!realloc_ptr.cast::<c_void>().is_null());
+        unsafe { mi_free(ptr.cast::<c_void>()) };
+    }
+}

allocator/src/lib.rs

@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#[cfg(not(target_family = "wasm"))]
+pub mod mimalloc;
+
+/// Declare a global allocator if the platform supports it.
+#[macro_export]
+macro_rules! assign_global {
+    () => {
+        #[cfg(not(target_family = "wasm"))]
+        #[global_allocator]
+        static GLOBAL: allocator::mimalloc::Mimalloc = allocator::mimalloc::Mimalloc;
+    };
+}

allocator/src/mimalloc.rs

@@ -0,0 +1,79 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use core::alloc::{GlobalAlloc, Layout};
+use core::ffi::c_void;
+
+use mimalloc_sys::{mi_free, mi_malloc_aligned, mi_realloc_aligned, mi_zalloc_aligned};
+
+pub struct Mimalloc;
+
+unsafe impl GlobalAlloc for Mimalloc {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        debug_assert!(layout.align() < mimalloc_sys::MI_ALIGNMENT_MAX);
+        mi_malloc_aligned(layout.size(), layout.align()).cast::<u8>()
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+        mi_free(ptr.cast::<c_void>());
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        debug_assert!(layout.align() < mimalloc_sys::MI_ALIGNMENT_MAX);
+        mi_zalloc_aligned(layout.size(), layout.align()).cast::<u8>()
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        debug_assert!(layout.align() < mimalloc_sys::MI_ALIGNMENT_MAX);
+        mi_realloc_aligned(ptr.cast::<c_void>(), new_size, layout.align()).cast::<u8>()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::error::Error;
+
+    #[test]
+    fn memory_can_be_allocated_and_freed() -> Result<(), Box<dyn Error>> {
+        let layout = Layout::from_size_align(8, 8)?;
+        let alloc = Mimalloc;
+
+        unsafe {
+            let ptr = alloc.alloc(layout);
+            assert!(!ptr.cast::<c_void>().is_null());
+            alloc.dealloc(ptr, layout);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn memory_can_be_alloc_zeroed_and_freed() -> Result<(), Box<dyn Error>> {
+        let layout = Layout::from_size_align(8, 8)?;
+        let alloc = Mimalloc;
+
+        unsafe {
+            let ptr = alloc.alloc_zeroed(layout);
+            assert!(!ptr.cast::<c_void>().is_null());
+            alloc.dealloc(ptr, layout);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn large_chunks_of_memory_can_be_allocated_and_freed() -> Result<(), Box<dyn Error>> {
+        let layout = Layout::from_size_align(2 * 1024 * 1024 * 1024, 8)?;
+        let alloc = Mimalloc;
+
+        unsafe {
+            let ptr = alloc.alloc(layout);
+            assert!(!ptr.cast::<c_void>().is_null());
+            alloc.dealloc(ptr, layout);
+        }
+        Ok(())
+    }
+}

compiler/qsc/Cargo.toml

@@ -29,6 +29,9 @@ qsc_project = { path = "../qsc_project", features = ["fs"] }
 rustc-hash = { workspace = true }
 thiserror = { workspace = true }
 
+[target.'cfg(not(any(target_family = "wasm")))'.dependencies]
+allocator = { path = "../../allocator" }
+
 [dev-dependencies]
 criterion = { workspace = true, features = ["cargo_bench_support"] }
 expect-test = { workspace = true }