From 18364ce19ff3595a321968c93cdc87025ede5527 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Mon, 27 Jul 2026 10:45:48 -0300 Subject: [PATCH 1/2] Add opt-in dlmalloc guest allocator --- syscalls/Cargo.lock | 28 +++++++ syscalls/Cargo.toml | 13 +++ syscalls/src/allocator.rs | 170 ++++++++++++++++++++++++++++++++++---- 3 files changed, 196 insertions(+), 15 deletions(-) diff --git a/syscalls/Cargo.lock b/syscalls/Cargo.lock index 34e481dd8..1921ff54a 100644 --- a/syscalls/Cargo.lock +++ b/syscalls/Cargo.lock @@ -64,6 +64,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dlmalloc" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad5208a115eaba24916f7456929832e310a81518c641f93fee4f89aa93aa3675" +dependencies = [ + "cfg-if", + "libc", + "windows-sys", +] + [[package]] name = "embedded-alloc" version = "0.6.0" @@ -128,6 +139,8 @@ dependencies = [ name = "lambda-vm-syscalls" version = "0.1.0" dependencies = [ + "critical-section", + "dlmalloc", "embedded-alloc", "getrandom 0.2.17", "getrandom 0.3.4", @@ -379,6 +392,21 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/syscalls/Cargo.toml b/syscalls/Cargo.toml index 0460a2435..e8c499d7e 100644 --- a/syscalls/Cargo.toml +++ b/syscalls/Cargo.toml @@ -11,6 +11,19 @@ getrandom = { version = "0.3.4", default-features = false } getrandom_v2 = {version = "0.2.15", features = ["custom"], package = "getrandom"} lazy_static = "1.5.0" rand = "0.9.2" +# `dlmalloc` feature only (optional so the default TLSF build never pulls them). +# `critical-section` gives the Sync a #[global_allocator] static needs; its single-hart +# impl already comes from `riscv` above (same as embedded-alloc's TLSF heap uses). +dlmalloc = { version = "0.2.14", default-features = false, optional = true } +critical-section = { version = "1.2", optional = true } + +[features] +# BENCH: swap the guest global allocator from embedded-alloc TLSF to Doug Lea's malloc +# (dlmalloc), backed by a bump "system" provider over the guest heap. Unlike a raw bump +# allocator it reclaims freed memory (no OOM on churny workloads) while measuring cheaper +# than the default TLSF on zkVM guests (Bencik's ZisK comparison). Off by default; see +# `src/allocator.rs`. +dlmalloc = ["dep:dlmalloc", "dep:critical-section"] [dev-dependencies] keccak = "0.1" diff --git a/syscalls/src/allocator.rs b/syscalls/src/allocator.rs index 78b2933e5..a95ac6e73 100644 --- a/syscalls/src/allocator.rs +++ b/syscalls/src/allocator.rs @@ -1,23 +1,163 @@ -use embedded_alloc::TlsfHeap as Heap; use riscv as _; -// Only the guest routes Rust allocations through this heap; on host (e.g. -// `cargo test` for the sponge's differential tests) the attribute would hijack -// the test harness's allocator with a never-initialized heap and abort. -#[cfg_attr(target_arch = "riscv64", global_allocator)] -static HEAP: Heap = Heap::empty(); - const MAX_MEMORY_SIZE: usize = 0xC000_0000; const WORD_SIZE: usize = 4; -pub fn init_allocator() { - { - unsafe extern "C" { - static _end: u8; +// Guest global allocator, selectable at build time: +// - default: embedded-alloc TLSF — reclaims freed memory (safe for arbitrary churn). +// - `dlmalloc` feature (BENCH): Doug Lea's malloc, backed by a bump "system" +// provider that hands dlmalloc page-aligned chunks from the guest heap. dlmalloc +// does all sub-allocation churn itself, so — unlike a raw bump allocator — it +// reuses freed memory (no OOM) while measuring cheaper than TLSF on zkVM guests +// (Bencik's ZisK allocator comparison). +// +// Only the guest installs a #[global_allocator]; on host (e.g. `cargo test` for the +// sponge's differential tests) the attribute would hijack the test harness's +// allocator with a never-initialized heap and abort. + +#[cfg(not(feature = "dlmalloc"))] +mod imp { + use embedded_alloc::TlsfHeap as Heap; + + #[cfg_attr(target_arch = "riscv64", global_allocator)] + static HEAP: Heap = Heap::empty(); + + pub fn init(heap_start: usize, heap_end: usize) { + unsafe { HEAP.init(heap_start, heap_end - heap_start) } + } +} + +#[cfg(feature = "dlmalloc")] +mod imp { + use core::alloc::{GlobalAlloc, Layout}; + use core::cell::RefCell; + use core::sync::atomic::{AtomicUsize, Ordering}; + use critical_section::Mutex; + use dlmalloc::{Allocator, Dlmalloc}; + + // Page granularity dlmalloc requests memory in. Must be a power of two; the guest + // heap region is 3 GiB so the value only affects the segment rounding below. + const PAGE_SIZE: usize = 4096; + + // The "system" side of dlmalloc: instead of mmap/sbrk (absent on the guest) it + // bump-allocates page-aligned segments from the single contiguous heap region + // [_end, MAX_MEMORY_SIZE). It never releases a segment (`free`/`free_part`/ + // `remap` all decline) — dlmalloc itself owns all reuse of freed *user* + // allocations against this fixed backing store, which is what keeps churny + // workloads OOM-free unlike a raw bump allocator. + struct BumpSystem; + + // Single-hart guest → `Relaxed` atomics are contention-free. + static HEAP_POS: AtomicUsize = AtomicUsize::new(0); + static HEAP_END: AtomicUsize = AtomicUsize::new(0); + + unsafe impl Allocator for BumpSystem { + fn alloc(&self, size: usize) -> (*mut u8, usize, u32) { + // Round up to a page so consecutive segments stay page-aligned. + let size = size.wrapping_add(PAGE_SIZE - 1) & !(PAGE_SIZE - 1); + let pos = HEAP_POS.load(Ordering::Relaxed); + match pos.checked_add(size) { + Some(new_pos) if new_pos <= HEAP_END.load(Ordering::Relaxed) => { + HEAP_POS.store(new_pos, Ordering::Relaxed); + // flags = 0: a plain external segment (never partially released). + (pos as *mut u8, size, 0) + } + // Out of heap → null makes dlmalloc return null → handle_alloc_error. + _ => (core::ptr::null_mut(), 0, 0), + } + } + + fn remap(&self, _ptr: *mut u8, _old: usize, _new: usize, _can_move: bool) -> *mut u8 { + core::ptr::null_mut() + } + + fn free_part(&self, _ptr: *mut u8, _old: usize, _new: usize) -> bool { + false + } + + fn free(&self, _ptr: *mut u8, _size: usize) -> bool { + false + } + + fn can_release_part(&self, _flags: u32) -> bool { + false + } + + fn allocates_zeros(&self) -> bool { + // Guest memory is zero-initialized and this provider never reuses a + // segment, so system-fresh bytes read as 0 → dlmalloc's calloc skips the + // memset for system-fresh memory (it still zeroes recycled blocks itself). + true + } + + fn page_size(&self) -> usize { + PAGE_SIZE + } + } + + // Dlmalloc is Send but !Sync, so it can't sit in a static directly. A single-hart + // critical section serializes access and supplies the Sync a #[global_allocator] + // static requires — the same primitive embedded-alloc's TLSF heap uses. + static DLMALLOC: Mutex>> = + Mutex::new(RefCell::new(Dlmalloc::new_with_allocator(BumpSystem))); + + struct DlGlobal; + + #[cfg_attr(target_arch = "riscv64", global_allocator)] + static ALLOC: DlGlobal = DlGlobal; + + pub fn init(heap_start: usize, heap_end: usize) { + HEAP_POS.store(heap_start, Ordering::Relaxed); + HEAP_END.store(heap_end, Ordering::Relaxed); + } + + unsafe impl GlobalAlloc for DlGlobal { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + critical_section::with(|cs| unsafe { + DLMALLOC + .borrow(cs) + .borrow_mut() + .malloc(layout.size(), layout.align()) + }) } - let heap_pos: usize = unsafe { (&_end) as *const u8 as usize }; - unsafe { HEAP.init(heap_pos, MAX_MEMORY_SIZE - heap_pos) } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + critical_section::with(|cs| unsafe { + DLMALLOC + .borrow(cs) + .borrow_mut() + .free(ptr, layout.size(), layout.align()) + }) + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + critical_section::with(|cs| unsafe { + DLMALLOC + .borrow(cs) + .borrow_mut() + .calloc(layout.size(), layout.align()) + }) + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + critical_section::with(|cs| unsafe { + DLMALLOC.borrow(cs).borrow_mut().realloc( + ptr, + layout.size(), + layout.align(), + new_size, + ) + }) + } + } +} + +pub fn init_allocator() { + unsafe extern "C" { + static _end: u8; } + let heap_pos: usize = unsafe { (&_end) as *const u8 as usize }; + imp::init(heap_pos, MAX_MEMORY_SIZE); } /// # Safety @@ -26,8 +166,8 @@ pub fn init_allocator() { /// It is only for rust std internal uses #[unsafe(no_mangle)] pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 { - use core::alloc::GlobalAlloc; - unsafe { HEAP.alloc(core::alloc::Layout::from_size_align(bytes, align).unwrap()) } + // Route through whichever `#[global_allocator]` is installed (TLSF or dlmalloc). + unsafe { std::alloc::alloc(core::alloc::Layout::from_size_align(bytes, align).unwrap()) } } /// # Safety From 9f28a234a8a9df87d619fa9294ca22268892d4ba Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Mon, 27 Jul 2026 11:57:30 -0300 Subject: [PATCH 2/2] Use dlmalloc as the default guest allocator --- Cargo.lock | 13 +++++++++++++ syscalls/Cargo.toml | 24 ++++++++++++----------- syscalls/src/allocator.rs | 41 ++++++++++++++++++++------------------- 3 files changed, 47 insertions(+), 31 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 74986dcc9..8ca304d11 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -496,6 +496,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dlmalloc" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad5208a115eaba24916f7456929832e310a81518c641f93fee4f89aa93aa3675" +dependencies = [ + "cfg-if", + "libc", + "windows-sys", +] + [[package]] name = "ecsm" version = "0.1.0" @@ -846,6 +857,8 @@ dependencies = [ name = "lambda-vm-syscalls" version = "0.1.0" dependencies = [ + "critical-section", + "dlmalloc", "embedded-alloc", "getrandom 0.2.16", "getrandom 0.3.4", diff --git a/syscalls/Cargo.toml b/syscalls/Cargo.toml index e8c499d7e..f1601ee35 100644 --- a/syscalls/Cargo.toml +++ b/syscalls/Cargo.toml @@ -4,6 +4,10 @@ version = "0.1.0" edition = "2024" [dependencies] +# embedded-alloc TLSF heap: the previous default guest allocator, now selectable via +# the `tlsf-alloc` feature. Kept a hard dep (like the codebase's other allocator +# toggles) so switching needs no dependency edits; it's only linked into the guest and +# dropped when dlmalloc (the default) is selected. embedded-alloc = "0.6" riscv = { version = "0.15", features = ["critical-section-single-hart"] } thiserror = "1.0" @@ -11,19 +15,17 @@ getrandom = { version = "0.3.4", default-features = false } getrandom_v2 = {version = "0.2.15", features = ["custom"], package = "getrandom"} lazy_static = "1.5.0" rand = "0.9.2" -# `dlmalloc` feature only (optional so the default TLSF build never pulls them). -# `critical-section` gives the Sync a #[global_allocator] static needs; its single-hart -# impl already comes from `riscv` above (same as embedded-alloc's TLSF heap uses). -dlmalloc = { version = "0.2.14", default-features = false, optional = true } -critical-section = { version = "1.2", optional = true } +# Default guest allocator: Doug Lea's malloc. Cheaper per alloc/free than TLSF on zkVM +# guests (fewer guest cycles) while still reclaiming freed memory (no OOM). `critical- +# section` gives the Sync a #[global_allocator] static needs; its single-hart impl comes +# from `riscv` above (same as embedded-alloc's TLSF heap uses). See `src/allocator.rs`. +dlmalloc = { version = "0.2.14", default-features = false } +critical-section = "1.2" [features] -# BENCH: swap the guest global allocator from embedded-alloc TLSF to Doug Lea's malloc -# (dlmalloc), backed by a bump "system" provider over the guest heap. Unlike a raw bump -# allocator it reclaims freed memory (no OOM on churny workloads) while measuring cheaper -# than the default TLSF on zkVM guests (Bencik's ZisK comparison). Off by default; see -# `src/allocator.rs`. -dlmalloc = ["dep:dlmalloc", "dep:critical-section"] +# Fallback/BENCH: use the embedded-alloc TLSF heap instead of the default dlmalloc guest +# allocator (to A/B the two or fall back). Off by default. See `src/allocator.rs`. +tlsf-alloc = [] [dev-dependencies] keccak = "0.1" diff --git a/syscalls/src/allocator.rs b/syscalls/src/allocator.rs index a95ac6e73..c99482b45 100644 --- a/syscalls/src/allocator.rs +++ b/syscalls/src/allocator.rs @@ -4,30 +4,19 @@ const MAX_MEMORY_SIZE: usize = 0xC000_0000; const WORD_SIZE: usize = 4; // Guest global allocator, selectable at build time: -// - default: embedded-alloc TLSF — reclaims freed memory (safe for arbitrary churn). -// - `dlmalloc` feature (BENCH): Doug Lea's malloc, backed by a bump "system" -// provider that hands dlmalloc page-aligned chunks from the guest heap. dlmalloc -// does all sub-allocation churn itself, so — unlike a raw bump allocator — it -// reuses freed memory (no OOM) while measuring cheaper than TLSF on zkVM guests -// (Bencik's ZisK allocator comparison). +// - default: Doug Lea's malloc (dlmalloc), backed by a bump "system" provider that +// hands dlmalloc page-aligned chunks from the guest heap. dlmalloc does all +// sub-allocation churn itself, so — unlike a raw bump allocator — it reuses freed +// memory (no OOM) while executing fewer guest instructions per alloc/free than TLSF +// (Bencik's ZisK allocator comparison; measured cheaper on our ethrex proving too). +// - `tlsf-alloc` feature: embedded-alloc's TLSF heap — the previous default, kept one +// flag away for A/B comparison or fallback. // // Only the guest installs a #[global_allocator]; on host (e.g. `cargo test` for the // sponge's differential tests) the attribute would hijack the test harness's // allocator with a never-initialized heap and abort. -#[cfg(not(feature = "dlmalloc"))] -mod imp { - use embedded_alloc::TlsfHeap as Heap; - - #[cfg_attr(target_arch = "riscv64", global_allocator)] - static HEAP: Heap = Heap::empty(); - - pub fn init(heap_start: usize, heap_end: usize) { - unsafe { HEAP.init(heap_start, heap_end - heap_start) } - } -} - -#[cfg(feature = "dlmalloc")] +#[cfg(not(feature = "tlsf-alloc"))] mod imp { use core::alloc::{GlobalAlloc, Layout}; use core::cell::RefCell; @@ -152,6 +141,18 @@ mod imp { } } +#[cfg(feature = "tlsf-alloc")] +mod imp { + use embedded_alloc::TlsfHeap as Heap; + + #[cfg_attr(target_arch = "riscv64", global_allocator)] + static HEAP: Heap = Heap::empty(); + + pub fn init(heap_start: usize, heap_end: usize) { + unsafe { HEAP.init(heap_start, heap_end - heap_start) } + } +} + pub fn init_allocator() { unsafe extern "C" { static _end: u8; @@ -166,7 +167,7 @@ pub fn init_allocator() { /// It is only for rust std internal uses #[unsafe(no_mangle)] pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 { - // Route through whichever `#[global_allocator]` is installed (TLSF or dlmalloc). + // Route through whichever `#[global_allocator]` is installed (dlmalloc or TLSF). unsafe { std::alloc::alloc(core::alloc::Layout::from_size_align(bytes, align).unwrap()) } }