From 9569a6ea0cb2d8e1aef4643749cb7f1620a61627 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 12 Dec 2024 01:11:38 +0000 Subject: [PATCH] feat(zkvm): implement freeing allocator and reserved memory section for read_vec Co-Authored-By: John Guibas --- crates/zkvm/entrypoint/src/heap.rs | 59 +++++++++---------- crates/zkvm/entrypoint/src/syscalls/memory.rs | 28 +++++++++ crates/zkvm/lib/src/io.rs | 21 +------ crates/zkvm/lib/src/lib.rs | 3 + 4 files changed, 63 insertions(+), 48 deletions(-) diff --git a/crates/zkvm/entrypoint/src/heap.rs b/crates/zkvm/entrypoint/src/heap.rs index 0e084d8f19..357c621dd0 100644 --- a/crates/zkvm/entrypoint/src/heap.rs +++ b/crates/zkvm/entrypoint/src/heap.rs @@ -11,46 +11,44 @@ struct FreeBlock { next: Option>, } -/// A simple heap allocator with free list. -/// -/// Uses a first-fit strategy with coalescing of adjacent free blocks. -/// Designed for single-threaded embedded systems with a memory limit of 0x78000000. -pub struct SimpleAlloc { - head: AtomicUsize, // Stores the raw pointer value of the head FreeBlock -} +// Global free list head stored as raw pointer value +static FREE_LIST_HEAD: AtomicUsize = AtomicUsize::new(0); -// SAFETY: The allocator is thread-safe due to atomic operations -unsafe impl Sync for SimpleAlloc {} +/// A simple heap allocator that supports freeing memory. +/// +/// Uses a first-fit strategy for allocation and maintains a free list +/// for memory reuse. Designed for single-threaded embedded systems +/// with a memory limit of 0x78000000. +#[derive(Copy, Clone)] +pub struct SimpleAlloc; +// Implementation detail functions impl SimpleAlloc { - const fn new() -> Self { - Self { - head: AtomicUsize::new(0), - } - } - - unsafe fn add_free_block(&self, ptr: *mut u8, size: usize) { + unsafe fn add_free_block(ptr: *mut u8, size: usize) { let block = ptr as *mut FreeBlock; (*block).size = size; loop { - let current_head = self.head.load(Ordering::Relaxed); + let current_head = FREE_LIST_HEAD.load(Ordering::Relaxed); (*block).next = NonNull::new(current_head as *mut FreeBlock); - if self.head.compare_exchange( - current_head, - block as usize, - Ordering::Release, - Ordering::Relaxed, - ).is_ok() { + if FREE_LIST_HEAD + .compare_exchange( + current_head, + block as usize, + Ordering::Release, + Ordering::Relaxed, + ) + .is_ok() + { break; } } } - unsafe fn find_block(&self, size: usize, align: usize) -> Option<(*mut u8, usize)> { + unsafe fn find_block(size: usize, align: usize) -> Option<(*mut u8, usize)> { let mut prev: Option<*mut FreeBlock> = None; - let mut current_ptr = self.head.load(Ordering::Acquire) as *mut FreeBlock; + let mut current_ptr = FREE_LIST_HEAD.load(Ordering::Acquire) as *mut FreeBlock; while !current_ptr.is_null() { let addr = current_ptr as *mut u8; @@ -64,7 +62,7 @@ impl SimpleAlloc { match prev { Some(p) => (*p).next = next, None => { - self.head.store(next_raw, Ordering::Release); + FREE_LIST_HEAD.store(next_raw, Ordering::Release); } } return Some((aligned_addr, (*current_ptr).size)); @@ -82,7 +80,7 @@ unsafe impl GlobalAlloc for SimpleAlloc { let align = layout.align(); // Try to find a block in free list - if let Some((ptr, _)) = self.find_block(size, align) { + if let Some((ptr, _)) = Self::find_block(size, align) { return ptr; } @@ -92,9 +90,10 @@ unsafe impl GlobalAlloc for SimpleAlloc { unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { let size = layout.size().max(core::mem::size_of::()); - self.add_free_block(ptr, size); + Self::add_free_block(ptr, size); } } -#[global_allocator] -static ALLOCATOR: SimpleAlloc = SimpleAlloc::new(); +#[used] +#[no_mangle] +pub static HEAP_ALLOCATOR: SimpleAlloc = SimpleAlloc; diff --git a/crates/zkvm/entrypoint/src/syscalls/memory.rs b/crates/zkvm/entrypoint/src/syscalls/memory.rs index 28d22e22d6..7fa81e4b16 100644 --- a/crates/zkvm/entrypoint/src/syscalls/memory.rs +++ b/crates/zkvm/entrypoint/src/syscalls/memory.rs @@ -15,6 +15,10 @@ // Memory addresses must be lower than BabyBear prime. const MAX_MEMORY: usize = 0x78000000; +const RESERVED_SIZE: usize = 1024 * 1024; // 1MB reserved section +static mut RESERVED_POS: usize = 0; +static mut RESERVED_START: usize = 0; + #[allow(clippy::missing_safety_doc)] #[no_mangle] pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 { @@ -32,6 +36,11 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u if heap_pos == 0 { heap_pos = unsafe { (&_end) as *const u8 as usize }; + unsafe { + RESERVED_START = heap_pos; + RESERVED_POS = heap_pos; + heap_pos += RESERVED_SIZE; + } } let offset = heap_pos & (align - 1); @@ -49,3 +58,22 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u unsafe { HEAP_POS = heap_pos }; ptr } +#[allow(clippy::missing_safety_doc)] +#[no_mangle] +pub unsafe extern "C" fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8 { + let mut pos = RESERVED_POS; + + // Align the position + let offset = pos & (align - 1); + if offset != 0 { + pos += align - offset; + } + + let new_pos = pos + bytes; + if new_pos > RESERVED_START + RESERVED_SIZE { + panic!("Reserved memory section full"); + } + + RESERVED_POS = new_pos; + pos as *mut u8 +} diff --git a/crates/zkvm/lib/src/io.rs b/crates/zkvm/lib/src/io.rs index 94a92a602c..895bdf0cff 100644 --- a/crates/zkvm/lib/src/io.rs +++ b/crates/zkvm/lib/src/io.rs @@ -1,10 +1,7 @@ #![allow(unused_unsafe)] -use crate::{syscall_hint_len, syscall_hint_read, syscall_write}; +use crate::{sys_alloc_reserved, syscall_hint_len, syscall_hint_read, syscall_write}; use serde::{de::DeserializeOwned, Serialize}; -use std::{ - alloc::Layout, - io::{Result, Write}, -}; +use std::io::{Result, Write}; /// The file descriptor for public values. pub const FD_PUBLIC_VALUES: u32 = 3; @@ -43,24 +40,12 @@ impl Write for SyscallWriter { /// let data: Vec = sp1_zkvm::io::read_vec(); /// ``` pub fn read_vec() -> Vec { - // Round up to the nearest multiple of 4 so that the memory allocated is in whole words let len = unsafe { syscall_hint_len() }; let capacity = (len + 3) / 4 * 4; + let ptr = unsafe { sys_alloc_reserved(capacity, 4) }; - // Allocate a buffer of the required length that is 4 byte aligned - let layout = Layout::from_size_align(capacity, 4).expect("vec is too large"); - let ptr = unsafe { std::alloc::alloc(layout) }; - - // SAFETY: - // 1. `ptr` was allocated using alloc - // 2. We assuume that the VM global allocator doesn't dealloc - // 3/6. Size is correct from above - // 4/5. Length is 0 - // 7. Layout::from_size_align already checks this let mut vec = unsafe { Vec::from_raw_parts(ptr, 0, capacity) }; - // Read the vec into uninitialized memory. The syscall assumes the memory is uninitialized, - // which should be true because the allocator does not dealloc, so a new alloc should be fresh. unsafe { syscall_hint_read(ptr, len); vec.set_len(len); diff --git a/crates/zkvm/lib/src/lib.rs b/crates/zkvm/lib/src/lib.rs index e43c4d0d68..d106793f83 100644 --- a/crates/zkvm/lib/src/lib.rs +++ b/crates/zkvm/lib/src/lib.rs @@ -97,6 +97,9 @@ extern "C" { /// Allocates a buffer aligned to the given alignment. pub fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8; + /// Allocates from the reserved memory section + pub fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8; + /// Decompresses a BLS12-381 point. pub fn syscall_bls12381_decompress(point: &mut [u8; 96], is_odd: bool);