perf: close the gap with CPython across the benchmark suite

Cargo.toml

@@ -141,6 +141,7 @@ unicode_names2        = "2.0"
 memmap2               = "0.9"
 memchr                = "2.8"
 libc                  = "0.2"
+libmimalloc-sys       = { version = "0.1.49", default-features = false }
 rustls                = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
 rustls-pki-types      = "1.7"
 # PKCS#8 "ENCRYPTED PRIVATE KEY" (PBES2) decryption for password-protected
@@ -317,7 +318,9 @@ opt-level = 1
 
 [profile.release]
 opt-level     = 3
-lto           = "thin"
+# Fat LTO inlines across the VM, compiler, and JIT crates: 1-10% faster
+# on the benchmark fixtures for about a third more build time.
+lto           = "fat"
 codegen-units = 1
 debug         = "line-tables-only"
 strip         = "debuginfo"

crates/weavepy-capi/src/memoryview.rs

@@ -329,7 +329,9 @@ pub unsafe extern "C" fn PyMemoryView_FromObjectAndFlags(
                 Some(obj.clone())
             };
             PyMemoryView {
-                buffer: MemoryViewBuffer::Shared(weavepy_vm::sync::Rc::new(region)),
+                buffer: MemoryViewBuffer::Shared(
+                    weavepy_vm::rc_unsize!(weavepy_vm::sync::Rc::new(region) => dyn SharedMemBuffer),
+                ),
                 start: Cell::new(start),
                 len: Cell::new(view_len),
                 readonly: Cell::new(readonly),
@@ -403,7 +405,9 @@ pub unsafe extern "C" fn PyMemoryView_FromMemory(
         readonly,
     };
     let mv = PyMemoryView::contiguous_1d(
-        MemoryViewBuffer::Shared(weavepy_vm::sync::Rc::new(region)),
+        MemoryViewBuffer::Shared(
+            weavepy_vm::rc_unsize!(weavepy_vm::sync::Rc::new(region) => dyn SharedMemBuffer),
+        ),
         len,
         readonly,
         "B".to_owned(),
@@ -486,7 +490,9 @@ pub unsafe extern "C" fn PyMemoryView_FromBuffer(view: *const Py_buffer) -> *mut
         readonly,
     };
     let mv = PyMemoryView {
-        buffer: MemoryViewBuffer::Shared(weavepy_vm::sync::Rc::new(region)),
+        buffer: MemoryViewBuffer::Shared(
+            weavepy_vm::rc_unsize!(weavepy_vm::sync::Rc::new(region) => dyn SharedMemBuffer),
+        ),
         start: Cell::new(0),
         len: Cell::new(len),
         readonly: Cell::new(readonly),

crates/weavepy-cli/Cargo.toml

@@ -39,6 +39,7 @@ weavepy-compiler   = { workspace = true }
 weavepy-conformance = { workspace = true }
 weavepy-version    = { workspace = true }
 libc               = { workspace = true }
+libmimalloc-sys    = { workspace = true }
 anyhow             = { workspace = true }
 clap               = { workspace = true }
 serde_json         = { workspace = true }
@@ -59,6 +60,9 @@ windows-sys = { workspace = true }
 default = ["jit"]
 # RFC 0032 — build the `weavepy` binary with the tier-2 JIT compiled in.
 jit = ["weavepy/jit", "weavepy-vm/jit"]
+# Sample live allocations by call stack (`WEAVEPY_ALLOC_PROFILE=<file>`,
+# macOS only); see `src/alloc_profile.rs`.
+alloc-profile = []
 
 [lints]
 workspace = true

crates/weavepy-cli/src/alloc.rs

@@ -0,0 +1,81 @@
+//! The process allocator: mimalloc, entered through its plain allocation
+//! functions whenever they already satisfy the requested alignment.
+//!
+//! Every mimalloc block is at least word aligned, and every block of 16
+//! bytes or more is 16-byte aligned (`MI_MAX_ALIGN_SIZE`), so a layout
+//! aligned to that much needs no aligned entry point. The aligned entry points
+//! take their fast path only when the size class's free list happens to
+//! offer a suitably aligned block, and otherwise fall to a generic path
+//! that may over-allocate. Rust asks for alignment on every allocation,
+//! so going through them would put every allocation on that path.
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::ffi::c_void;
+
+use libmimalloc_sys::{
+    mi_free, mi_malloc, mi_malloc_aligned, mi_realloc, mi_realloc_aligned, mi_zalloc,
+    mi_zalloc_aligned,
+};
+
+/// The largest alignment every mimalloc block already has.
+const WORD_ALIGN: usize = std::mem::size_of::<usize>();
+
+/// mimalloc's `MI_MAX_ALIGN_SIZE`: the alignment of every block at least
+/// this large.
+const MAX_ALIGN: usize = 16;
+
+/// Whether the plain entry points already satisfy `align` for `size`.
+#[inline(always)]
+fn plain(size: usize, align: usize) -> bool {
+    align <= WORD_ALIGN || (align <= MAX_ALIGN && size >= align)
+}
+
+/// mimalloc as the global allocator (see the module docs).
+pub(crate) struct Mimalloc;
+
+// SAFETY: every block comes from mimalloc with at least the layout's
+// alignment (word-aligned blocks, or the aligned entry points for larger
+// alignments), and every block is released or resized through mimalloc.
+unsafe impl GlobalAlloc for Mimalloc {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: plain FFI allocation calls.
+        unsafe {
+            if plain(layout.size(), layout.align()) {
+                mi_malloc(layout.size()).cast()
+            } else {
+                mi_malloc_aligned(layout.size(), layout.align()).cast()
+            }
+        }
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: as in `alloc`.
+        unsafe {
+            if plain(layout.size(), layout.align()) {
+                mi_zalloc(layout.size()).cast()
+            } else {
+                mi_zalloc_aligned(layout.size(), layout.align()).cast()
+            }
+        }
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+        // SAFETY: `ptr` came from this allocator (the caller's contract).
+        unsafe { mi_free(ptr.cast::<c_void>()) }
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        // SAFETY: `ptr` came from this allocator with `layout`.
+        unsafe {
+            if plain(new_size, layout.align()) {
+                mi_realloc(ptr.cast::<c_void>(), new_size).cast()
+            } else {
+                mi_realloc_aligned(ptr.cast::<c_void>(), new_size, layout.align()).cast()
+            }
+        }
+    }
+}

crates/weavepy-cli/src/alloc_profile.rs

@@ -0,0 +1,211 @@
+//! Opt-in allocation-site profiler (`--features alloc-profile`, macOS).
+//!
+//! The global allocator samples one allocation per [`SAMPLE`] bytes and
+//! records its call stack. A sampled block is charged [`SAMPLE`] bytes to
+//! its stack while it lives, so at exit the table estimates the *live* heap
+//! by allocation site. With `WEAVEPY_ALLOC_PROFILE=<file>` set, the live
+//! samples are written to `<file>`: one line per sample, the charged bytes
+//! then the slide-adjusted return addresses (resolve them with `atos`).
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+/// Bytes between samples.
+const SAMPLE: usize = 16 * 1024;
+/// Return addresses kept per sample.
+const DEPTH: usize = 24;
+/// Sampled blocks tracked at once (open addressing, power of two).
+const SLOTS: usize = 1 << 18;
+
+struct Sample {
+    ptr: usize,
+    stack: [usize; DEPTH],
+}
+
+static ENABLED: AtomicBool = AtomicBool::new(false);
+static LOCK: AtomicBool = AtomicBool::new(false);
+static TABLE: AtomicUsize = AtomicUsize::new(0);
+
+thread_local! {
+    static IN_HOOK: Cell<bool> = const { Cell::new(false) };
+    static UNTIL_SAMPLE: Cell<usize> = const { Cell::new(SAMPLE) };
+}
+
+extern "C" {
+    fn _dyld_get_image_vmaddr_slide(image_index: u32) -> isize;
+}
+
+fn lock() {
+    while LOCK
+        .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
+        .is_err()
+    {
+        std::hint::spin_loop();
+    }
+}
+
+fn unlock() {
+    LOCK.store(false, Ordering::Release);
+}
+
+fn table() -> *mut Sample {
+    TABLE.load(Ordering::Acquire) as *mut Sample
+}
+
+fn slot_of(ptr: usize) -> usize {
+    (ptr.wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 20) & (SLOTS - 1)
+}
+
+/// Start sampling (allocates the table from the system allocator).
+pub(crate) fn start() {
+    if std::env::var_os("WEAVEPY_ALLOC_PROFILE").is_none() {
+        return;
+    }
+    let bytes = SLOTS * std::mem::size_of::<Sample>();
+    // SAFETY: a fresh zeroed mapping owned for the rest of the process.
+    let p = unsafe {
+        libc::mmap(
+            std::ptr::null_mut(),
+            bytes,
+            libc::PROT_READ | libc::PROT_WRITE,
+            libc::MAP_PRIVATE | libc::MAP_ANON,
+            -1,
+            0,
+        )
+    };
+    if p == libc::MAP_FAILED {
+        return;
+    }
+    TABLE.store(p as usize, Ordering::Release);
+    ENABLED.store(true, Ordering::Release);
+}
+
+fn record(ptr: usize) {
+    let mut stack = [0usize; DEPTH];
+    // SAFETY: `backtrace` fills at most `DEPTH` entries of the buffer.
+    let n = unsafe { libc::backtrace(stack.as_mut_ptr().cast(), DEPTH as libc::c_int) };
+    let _ = n;
+    let t = table();
+    lock();
+    let mut i = slot_of(ptr);
+    for _ in 0..SLOTS {
+        // SAFETY: `i < SLOTS`, inside the mapping.
+        let s = unsafe { &mut *t.add(i) };
+        if s.ptr == 0 || s.ptr == ptr {
+            s.ptr = ptr;
+            s.stack = stack;
+            break;
+        }
+        i = (i + 1) & (SLOTS - 1);
+    }
+    unlock();
+}
+
+fn forget(ptr: usize) {
+    let t = table();
+    lock();
+    let mut i = slot_of(ptr);
+    for _ in 0..SLOTS {
+        // SAFETY: as in `record`.
+        let s = unsafe { &mut *t.add(i) };
+        if s.ptr == 0 {
+            break;
+        }
+        if s.ptr == ptr {
+            // A tombstone keeps later probes of this chain reachable.
+            s.ptr = usize::MAX;
+            break;
+        }
+        i = (i + 1) & (SLOTS - 1);
+    }
+    unlock();
+}
+
+/// The profiling allocator: mimalloc plus the sampler.
+pub(crate) struct Profiled;
+
+// SAFETY: every allocation is mimalloc's; the sampler only records
+// addresses and never touches the blocks.
+unsafe impl GlobalAlloc for Profiled {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: forwarded unchanged.
+        let p = unsafe { crate::alloc::Mimalloc.alloc(layout) };
+        if ENABLED.load(Ordering::Relaxed) && !p.is_null() {
+            sample(p as usize, layout.size());
+        }
+        p
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        if ENABLED.load(Ordering::Relaxed) {
+            forget(ptr as usize);
+        }
+        // SAFETY: forwarded unchanged.
+        unsafe { crate::alloc::Mimalloc.dealloc(ptr, layout) }
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        if ENABLED.load(Ordering::Relaxed) {
+            forget(ptr as usize);
+        }
+        // SAFETY: forwarded unchanged.
+        let p = unsafe { crate::alloc::Mimalloc.realloc(ptr, layout, new_size) };
+        if ENABLED.load(Ordering::Relaxed) && !p.is_null() {
+            sample(p as usize, new_size);
+        }
+        p
+    }
+}
+
+fn sample(ptr: usize, size: usize) {
+    let _ = IN_HOOK.try_with(|hook| {
+        if hook.get() {
+            return;
+        }
+        let due = UNTIL_SAMPLE.with(|left| {
+            let l = left.get();
+            if size >= l {
+                left.set(SAMPLE);
+                true
+            } else {
+                left.set(l - size);
+                false
+            }
+        });
+        if due {
+            hook.set(true);
+            record(ptr);
+            hook.set(false);
+        }
+    });
+}
+
+/// Write the live samples (see the module docs).
+pub(crate) fn finish() {
+    let Some(path) = std::env::var_os("WEAVEPY_ALLOC_PROFILE") else {
+        return;
+    };
+    if !ENABLED.swap(false, Ordering::AcqRel) {
+        return;
+    }
+    let t = table();
+    // SAFETY: the main executable is image 0.
+    let slide = unsafe { _dyld_get_image_vmaddr_slide(0) } as usize;
+    let mut out = String::new();
+    lock();
+    for i in 0..SLOTS {
+        // SAFETY: `i < SLOTS`.
+        let s = unsafe { &*t.add(i) };
+        if s.ptr == 0 || s.ptr == usize::MAX {
+            continue;
+        }
+        out.push_str(&SAMPLE.to_string());
+        for &a in s.stack.iter().take_while(|a| **a != 0) {
+            out.push_str(&format!(" {:x}", a.wrapping_sub(slide)));
+        }
+        out.push('\n');
+    }
+    unlock();
+    let _ = std::fs::write(path, out);
+}