From 0621d5b4d5d04d35a6ff06879420275ba217cd35 Mon Sep 17 00:00:00 2001 From: Mrmaxmeier Date: Thu, 14 Apr 2022 20:37:41 +0200 Subject: [PATCH 1/6] cranelift-jit: add bump memory allocator, make configurable This adds a new JIT memory allocator which operates on an address range that is reserved up-front. This is useful in order to side-step issues where allocation behavior affect relocations, i.e. on x86 where we run into issues when the system allocator returns memory that is more than 4 GiB apart. The main drawback of this allocator option is that the memory range has to be reserved up-front, so allocation will fail if this space is exhausted. This commit: - adds a 'JITMemoryProvider' trait - moves the previous on-demand allocator to this trait - adds a new 'Arena' allocator --- cranelift/jit/src/backend.rs | 96 ++++---- cranelift/jit/src/lib.rs | 2 + cranelift/jit/src/memory/arena.rs | 218 ++++++++++++++++++ cranelift/jit/src/memory/mod.rs | 80 +++++++ .../jit/src/{memory.rs => memory/system.rs} | 116 +++++----- 5 files changed, 405 insertions(+), 107 deletions(-) create mode 100644 cranelift/jit/src/memory/arena.rs create mode 100644 cranelift/jit/src/memory/mod.rs rename cranelift/jit/src/{memory.rs => memory/system.rs} (73%) diff --git a/cranelift/jit/src/backend.rs b/cranelift/jit/src/backend.rs index 5a2f27d52a98..0f52b08361bb 100644 --- a/cranelift/jit/src/backend.rs +++ b/cranelift/jit/src/backend.rs @@ -1,6 +1,9 @@ //! Defines `JITModule`. -use crate::{compiled_blob::CompiledBlob, memory::BranchProtection, memory::Memory}; +use crate::{ + compiled_blob::CompiledBlob, + memory::{BranchProtection, JITMemoryProvider, SystemMemoryProvider}, +}; use cranelift_codegen::binemit::Reloc; use cranelift_codegen::isa::{OwnedTargetIsa, TargetIsa}; use cranelift_codegen::settings::Configurable; @@ -28,6 +31,7 @@ pub struct JITBuilder { symbols: HashMap>, lookup_symbols: Vec Option<*const u8> + Send>>, libcall_names: Box String + Send + Sync>, + memory: Option>, } impl JITBuilder { @@ -91,6 +95,7 @@ impl JITBuilder { symbols, lookup_symbols, libcall_names, + memory: None, } } @@ -141,6 +146,14 @@ impl JITBuilder { self.lookup_symbols.push(symbol_lookup_fn); self } + + /// Set the memory provider for the module. + /// + /// If unset, defaults to [`SystemMemoryProvider`]. + pub fn memory_provider(&mut self, provider: Box) -> &mut Self { + self.memory = Some(provider); + self + } } /// A wrapper that impls Send for the contents. @@ -159,7 +172,7 @@ pub struct JITModule { symbols: RefCell>>, lookup_symbols: Vec Option<*const u8> + Send>>, libcall_names: Box String + Send + Sync>, - memory: MemoryHandle, + memory: Box, declarations: ModuleDeclarations, compiled_functions: SecondaryMap>, compiled_data_objects: SecondaryMap>, @@ -167,13 +180,6 @@ pub struct JITModule { data_objects_to_finalize: Vec, } -/// A handle to allow freeing memory allocated by the `Module`. -struct MemoryHandle { - code: Memory, - readonly: Memory, - writable: Memory, -} - impl JITModule { /// Free memory allocated for code and data segments of compiled functions. /// @@ -184,9 +190,7 @@ impl JITModule { /// from that module are currently executing and none of the `fn` pointers /// are called afterwards. pub unsafe fn free_memory(mut self) { - self.memory.code.free_memory(); - self.memory.readonly.free_memory(); - self.memory.writable.free_memory(); + self.memory.free_memory(); } fn lookup_symbol(&self, name: &str) -> Option<*const u8> { @@ -325,8 +329,12 @@ impl JITModule { } // Now that we're done patching, prepare the memory for execution! - self.memory.readonly.set_readonly()?; - self.memory.code.set_readable_and_executable()?; + let branch_protection = if cfg!(target_arch = "aarch64") && use_bti(&self.isa.isa_flags()) { + BranchProtection::BTI + } else { + BranchProtection::None + }; + self.memory.finalize(branch_protection)?; Ok(()) } @@ -338,23 +346,15 @@ impl JITModule { "cranelift-jit needs is_pic=false" ); - let branch_protection = - if cfg!(target_arch = "aarch64") && use_bti(&builder.isa.isa_flags()) { - BranchProtection::BTI - } else { - BranchProtection::None - }; + let memory = builder + .memory + .unwrap_or_else(|| Box::new(SystemMemoryProvider::new())); Self { isa: builder.isa, symbols: RefCell::new(builder.symbols), lookup_symbols: builder.lookup_symbols, libcall_names: builder.libcall_names, - memory: MemoryHandle { - code: Memory::new(branch_protection), - // Branch protection is not applicable to non-executable memory. - readonly: Memory::new(BranchProtection::None), - writable: Memory::new(BranchProtection::None), - }, + memory, declarations: ModuleDeclarations::default(), compiled_functions: SecondaryMap::new(), compiled_data_objects: SecondaryMap::new(), @@ -436,15 +436,16 @@ impl Module for JITModule { let compiled_code = ctx.compiled_code().unwrap(); let size = compiled_code.code_info().total_size as usize; - let align = alignment.max(self.isa.symbol_alignment()); - let ptr = self - .memory - .code - .allocate(size, align) - .map_err(|e| ModuleError::Allocation { - message: "unable to alloc function", - err: e, - })?; + let align = alignment + .max(self.isa.function_alignment().minimum as u64) + .max(self.isa.symbol_alignment()); + let ptr = + self.memory + .allocate_readexec(size, align) + .map_err(|e| ModuleError::Allocation { + message: "unable to alloc function", + err: e, + })?; { let mem = unsafe { std::slice::from_raw_parts_mut(ptr, size) }; @@ -488,15 +489,16 @@ impl Module for JITModule { } let size = bytes.len(); - let align = alignment.max(self.isa.symbol_alignment()); - let ptr = self - .memory - .code - .allocate(size, align) - .map_err(|e| ModuleError::Allocation { - message: "unable to alloc function bytes", - err: e, - })?; + let align = alignment + .max(self.isa.function_alignment().minimum as u64) + .max(self.isa.symbol_alignment()); + let ptr = + self.memory + .allocate_readexec(size, align) + .map_err(|e| ModuleError::Allocation { + message: "unable to alloc function bytes", + err: e, + })?; unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), ptr, size); @@ -548,16 +550,14 @@ impl Module for JITModule { let ptr = if decl.writable { self.memory - .writable - .allocate(alloc_size, align.unwrap_or(WRITABLE_DATA_ALIGNMENT)) + .allocate_readwrite(alloc_size, align.unwrap_or(WRITABLE_DATA_ALIGNMENT)) .map_err(|e| ModuleError::Allocation { message: "unable to alloc writable data", err: e, })? } else { self.memory - .readonly - .allocate(alloc_size, align.unwrap_or(READONLY_DATA_ALIGNMENT)) + .allocate_readonly(alloc_size, align.unwrap_or(READONLY_DATA_ALIGNMENT)) .map_err(|e| ModuleError::Allocation { message: "unable to alloc readonly data", err: e, diff --git a/cranelift/jit/src/lib.rs b/cranelift/jit/src/lib.rs index a7e4190cd12f..ee0012270f17 100644 --- a/cranelift/jit/src/lib.rs +++ b/cranelift/jit/src/lib.rs @@ -9,6 +9,8 @@ mod backend; mod compiled_blob; mod memory; +pub use memory::*; + pub use crate::backend::{JITBuilder, JITModule}; /// Version number of this crate. diff --git a/cranelift/jit/src/memory/arena.rs b/cranelift/jit/src/memory/arena.rs new file mode 100644 index 000000000000..4bed9d9b324c --- /dev/null +++ b/cranelift/jit/src/memory/arena.rs @@ -0,0 +1,218 @@ +use std::io; +use std::mem::ManuallyDrop; +use std::ptr; + +use cranelift_module::ModuleResult; + +use super::{BranchProtection, JITMemoryProvider}; + +fn align_up(addr: usize, align: usize) -> usize { + debug_assert!(align.is_power_of_two()); + (addr + align - 1) & !(align - 1) +} + +#[derive(Debug)] +struct Segment { + ptr: *mut u8, + len: usize, + position: usize, + target_prot: region::Protection, + finalized: bool, +} + +impl Segment { + fn new(ptr: *mut u8, len: usize, target_prot: region::Protection) -> Self { + let mut segment = Segment { + ptr, + len, + target_prot, + position: 0, + finalized: false, + }; + // set segment to read-write for initialization + segment.set_rw(); + segment + } + + fn set_rw(&mut self) { + unsafe { + region::protect(self.ptr, self.len, region::Protection::READ_WRITE) + .expect("unable to change memory protection for jit memory segment"); + } + } + + fn finalize(&mut self, branch_protection: BranchProtection) { + if self.finalized { + return; + } + + if self.target_prot == region::Protection::READ_EXECUTE { + super::set_readable_and_executable(self.ptr, self.len, branch_protection) + .expect("unable to set memory protection for jit memory segment"); + } else { + unsafe { + region::protect(self.ptr, self.len, self.target_prot) + .expect("unable to change memory protection for jit memory segment"); + } + } + self.finalized = true; + } + + fn allocate(&mut self, size: usize, align: usize) -> *mut u8 { + assert!(self.has_space_for(size, align)); + self.position = align_up(self.position, align); + let ptr = unsafe { self.ptr.add(self.position) }; + self.position += size; + ptr + } + + fn has_space_for(&self, size: usize, align: usize) -> bool { + !self.finalized && align_up(self.position, align) + size <= self.len + } +} + +/// `ArenaMemoryProvider` allocates segments from a contiguous memory region +/// that is reserved up-front. +/// +/// The arena's memory is initially allocated with PROT_NONE and gradually +/// updated as the JIT requires more space. This approach allows for stable +/// addresses throughout the lifetime of the JIT. +/// +/// Note: Memory will be leaked by default unless +/// [`JITMemoryProvider::free_memory`] is called to ensure function pointers +/// remain valid for the remainder of the program's life. +pub struct ArenaMemoryProvider { + alloc: ManuallyDrop>, + ptr: *mut u8, + size: usize, + position: usize, + segments: Vec, +} + +impl ArenaMemoryProvider { + /// Create a new memory region with the given size. + pub fn new_with_size(reserve_size: usize) -> Result { + let size = align_up(reserve_size, region::page::size()); + let mut alloc = region::alloc(reserve_size, region::Protection::NONE)?; + let ptr = alloc.as_mut_ptr(); + + Ok(Self { + alloc: ManuallyDrop::new(Some(alloc)), + segments: Vec::new(), + ptr, + size, + position: 0, + }) + } + + fn allocate( + &mut self, + size: usize, + align: usize, + protection: region::Protection, + ) -> io::Result<*mut u8> { + // Note: Add a fast path without a linear scan over segments here? + + // can we fit this allocation into an existing segment + if let Some(segment) = self.segments.iter_mut().find(|seg| { + seg.target_prot == protection && !seg.finalized && seg.has_space_for(size, align) + }) { + return Ok(segment.allocate(size, align)); + } + + // can we resize the last segment? + if let Some(segment) = self.segments.iter_mut().last() { + if segment.target_prot == protection && !segment.finalized { + let align = align.max(region::page::size()); + let additional_size = align_up(size, align); + + // if our reserved arena can fit the additional size, extend the + // last segment + if self.position + additional_size <= self.size { + segment.len += additional_size; + segment.set_rw(); + self.position += additional_size; + return Ok(segment.allocate(size, align)); + } + } + } + + // allocate new segment for given size and alignment + self.allocate_segment(align_up(size, align), protection)?; + let i = self.segments.len() - 1; + Ok(self.segments[i].allocate(size, align)) + } + + fn allocate_segment( + &mut self, + size: usize, + target_prot: region::Protection, + ) -> Result<(), io::Error> { + let size = align_up(size, region::page::size()); + let ptr = unsafe { self.ptr.add(self.position) }; + if self.position + size > self.size { + return Err(io::Error::new( + io::ErrorKind::Other, + "pre-allocated jit memory region exhausted", + )); + } + self.position += size; + self.segments.push(Segment::new(ptr, size, target_prot)); + Ok(()) + } + + pub(crate) fn finalize(&mut self, branch_protection: BranchProtection) { + for segment in &mut self.segments { + segment.finalize(branch_protection); + } + } + + /// Frees the allocated memory region, which would be leaked otherwise. + /// Likely to invalidate existing function pointers, causing unsafety. + pub(crate) unsafe fn free_memory(&mut self) { + if self.ptr == ptr::null_mut() { + return; + } + self.segments.clear(); + // Drop the allocation, freeing memory + let _: Option = self.alloc.take(); + self.ptr = ptr::null_mut(); + } +} + +impl Drop for ArenaMemoryProvider { + fn drop(&mut self) { + if self.ptr == ptr::null_mut() { + return; + } + let is_live = self.segments.iter().any(|seg| seg.finalized); + if !is_live { + // Only free memory if it's not been finalized yet. + // Otherwise, leak it since JIT memory may still be in use. + unsafe { self.free_memory() }; + } + } +} + +impl JITMemoryProvider for ArenaMemoryProvider { + fn allocate_readexec(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { + self.allocate(size, align as usize, region::Protection::READ_EXECUTE) + } + + fn allocate_readwrite(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { + self.allocate(size, align as usize, region::Protection::READ_WRITE) + } + + fn allocate_readonly(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { + self.allocate(size, align as usize, region::Protection::READ) + } + + unsafe fn free_memory(&mut self) { + self.free_memory(); + } + + fn finalize(&mut self, branch_protection: BranchProtection) -> ModuleResult<()> { + self.finalize(branch_protection); + Ok(()) + } +} diff --git a/cranelift/jit/src/memory/mod.rs b/cranelift/jit/src/memory/mod.rs new file mode 100644 index 000000000000..0a042886622e --- /dev/null +++ b/cranelift/jit/src/memory/mod.rs @@ -0,0 +1,80 @@ +mod arena; +mod system; + +use std::io; + +use cranelift_module::ModuleResult; + +pub use arena::ArenaMemoryProvider; +pub use system::SystemMemoryProvider; + +/// Type of branch protection to apply to executable memory. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum BranchProtection { + /// No protection. + None, + /// Use the Branch Target Identification extension of the Arm architecture. + BTI, +} + +/// A provider of memory for the JIT. +pub trait JITMemoryProvider { + /// Allocate memory that will be executable once finalized. + fn allocate_readexec(&mut self, size: usize, align: u64) -> io::Result<*mut u8>; + /// Allocate writable memory. + fn allocate_readwrite(&mut self, size: usize, align: u64) -> io::Result<*mut u8>; + /// Allocate memory that will be read-only once finalized. + fn allocate_readonly(&mut self, size: usize, align: u64) -> io::Result<*mut u8>; + + /// Free the memory region. + unsafe fn free_memory(&mut self); + /// Finalize the memory region and apply memory protections. + fn finalize(&mut self, branch_protection: BranchProtection) -> ModuleResult<()>; +} + +/// Set all memory allocated in this `Memory` up to now as readable and executable. +/// This function deals with icache coherence, branch protection, and pipeline flushing. +pub(crate) fn set_readable_and_executable( + ptr: *mut u8, + len: usize, + branch_protection: BranchProtection, +) -> ModuleResult<()> { + // Clear all the newly allocated code from cache if the processor requires it + // + // Do this before marking the memory as R+X, technically we should be able to do it after + // but there are some CPU's that have had errata about doing this with read only memory. + unsafe { + wasmtime_jit_icache_coherence::clear_cache(ptr as *const libc::c_void, len) + .expect("Failed cache clear") + }; + + unsafe { + region::protect(ptr, len, region::Protection::READ_EXECUTE).map_err(|e| { + cranelift_module::ModuleError::Backend( + anyhow::Error::new(e).context("unable to make memory readable+executable"), + ) + })?; + } + + // If BTI is requested, and the architecture supports it, use mprotect to set the PROT_BTI flag + if branch_protection == BranchProtection::BTI { + #[cfg(all(target_arch = "aarch64", target_os = "linux"))] + if is_bti && std::arch::is_aarch64_feature_detected!("bti") { + let prot = libc::PROT_EXEC | libc::PROT_READ | /* PROT_BTI */ 0x10; + + unsafe { + if libc::mprotect(ptr as *mut libc::c_void, len, prot) < 0 { + return Err(ModuleError::Backend( + anyhow::Error::new(io::Error::last_os_error()) + .context("unable to make memory readable+executable"), + )); + } + } + } + } + + // Flush any in-flight instructions from the pipeline + wasmtime_jit_icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); + + Ok(()) +} diff --git a/cranelift/jit/src/memory.rs b/cranelift/jit/src/memory/system.rs similarity index 73% rename from cranelift/jit/src/memory.rs rename to cranelift/jit/src/memory/system.rs index 3e005716d26d..f25a03700aaa 100644 --- a/cranelift/jit/src/memory.rs +++ b/cranelift/jit/src/memory/system.rs @@ -5,11 +5,12 @@ use memmap2::MmapMut; #[cfg(not(any(feature = "selinux-fix", windows)))] use std::alloc; -use std::ffi::c_void; use std::io; use std::mem; use std::ptr; -use wasmtime_jit_icache_coherence as icache_coherence; + +use super::BranchProtection; +use super::JITMemoryProvider; /// A simple struct consisting of a pointer and length. struct PtrLen { @@ -111,15 +112,6 @@ impl Drop for PtrLen { // TODO: add a `Drop` impl for `cfg(target_os = "windows")` -/// Type of branch protection to apply to executable memory. -#[derive(Clone, Debug, PartialEq)] -pub(crate) enum BranchProtection { - /// No protection. - None, - /// Use the Branch Target Identification extension of the Arm architecture. - BTI, -} - /// JIT memory manager. This manages pages of suitably aligned and /// accessible memory. Memory will be leaked by default to have /// function pointers remain valid for the remainder of the @@ -129,19 +121,17 @@ pub(crate) struct Memory { already_protected: usize, current: PtrLen, position: usize, - branch_protection: BranchProtection, } unsafe impl Send for Memory {} impl Memory { - pub(crate) fn new(branch_protection: BranchProtection) -> Self { + pub(crate) fn new() -> Self { Self { allocations: Vec::new(), already_protected: 0, current: PtrLen::new(), position: 0, - branch_protection, } } @@ -175,56 +165,16 @@ impl Memory { } /// Set all memory allocated in this `Memory` up to now as readable and executable. - pub(crate) fn set_readable_and_executable(&mut self) -> ModuleResult<()> { + pub(crate) fn set_readable_and_executable( + &mut self, + branch_protection: BranchProtection, + ) -> ModuleResult<()> { self.finish_current(); - // Clear all the newly allocated code from cache if the processor requires it - // - // Do this before marking the memory as R+X, technically we should be able to do it after - // but there are some CPU's that have had errata about doing this with read only memory. - for &PtrLen { ptr, len, .. } in self.non_protected_allocations_iter() { - unsafe { - icache_coherence::clear_cache(ptr as *const c_void, len) - .expect("Failed cache clear") - }; - } - - let set_region_readable_and_executable = |ptr, len| -> ModuleResult<()> { - if self.branch_protection == BranchProtection::BTI { - #[cfg(all(target_arch = "aarch64", target_os = "linux"))] - if std::arch::is_aarch64_feature_detected!("bti") { - let prot = libc::PROT_EXEC | libc::PROT_READ | /* PROT_BTI */ 0x10; - - unsafe { - if libc::mprotect(ptr as *mut libc::c_void, len, prot) < 0 { - return Err(ModuleError::Backend( - anyhow::Error::new(io::Error::last_os_error()) - .context("unable to make memory readable+executable"), - )); - } - } - - return Ok(()); - } - } - - unsafe { - region::protect(ptr, len, region::Protection::READ_EXECUTE).map_err(|e| { - ModuleError::Backend( - anyhow::Error::new(e).context("unable to make memory readable+executable"), - ) - })?; - } - Ok(()) - }; - for &PtrLen { ptr, len, .. } in self.non_protected_allocations_iter() { - set_region_readable_and_executable(ptr, len)?; + super::set_readable_and_executable(ptr, len, branch_protection)?; } - // Flush any in-flight instructions from the pipeline - icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); - self.already_protected = self.allocations.len(); Ok(()) } @@ -274,3 +224,51 @@ impl Drop for Memory { .for_each(mem::forget); } } + +/// A memory provider that allocates memory on-demand using the system +/// allocator. +/// +/// Note: Memory will be leaked by default unless +/// [`JITMemoryProvider::free_memory`] is called to ensure function pointers +/// remain valid for the remainder of the program's life. +pub struct SystemMemoryProvider { + code: Memory, + readonly: Memory, + writable: Memory, +} + +impl SystemMemoryProvider { + /// Create a new memory handle with the given branch protection. + pub fn new() -> Self { + Self { + code: Memory::new(), + readonly: Memory::new(), + writable: Memory::new(), + } + } +} + +impl JITMemoryProvider for SystemMemoryProvider { + unsafe fn free_memory(&mut self) { + self.code.free_memory(); + self.readonly.free_memory(); + self.writable.free_memory(); + } + + fn finalize(&mut self, branch_protection: BranchProtection) -> ModuleResult<()> { + self.readonly.set_readonly()?; + self.code.set_readable_and_executable(branch_protection) + } + + fn allocate_readexec(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { + self.code.allocate(size, align) + } + + fn allocate_readwrite(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { + self.writable.allocate(size, align) + } + + fn allocate_readonly(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { + self.readonly.allocate(size, align) + } +} From 222d233e98590785681d5cc70a6f92fa5fe1daa2 Mon Sep 17 00:00:00 2001 From: Mrmaxmeier Date: Mon, 7 Apr 2025 11:28:14 +0200 Subject: [PATCH 2/6] ArenaMemoryProvider: add tests, simplify alignment logic --- cranelift/jit/src/memory/arena.rs | 79 +++++++++++++++++++++++++------ 1 file changed, 64 insertions(+), 15 deletions(-) diff --git a/cranelift/jit/src/memory/arena.rs b/cranelift/jit/src/memory/arena.rs index 4bed9d9b324c..dd4a036939fa 100644 --- a/cranelift/jit/src/memory/arena.rs +++ b/cranelift/jit/src/memory/arena.rs @@ -29,7 +29,8 @@ impl Segment { position: 0, finalized: false, }; - // set segment to read-write for initialization + // Set segment to read-write for initialization. The target permissions + // will be applied in `finalize`. segment.set_rw(); segment } @@ -46,6 +47,8 @@ impl Segment { return; } + // Executable regions are handled separately to correctly deal with + // branch protection and cache coherence. if self.target_prot == region::Protection::READ_EXECUTE { super::set_readable_and_executable(self.ptr, self.len, branch_protection) .expect("unable to set memory protection for jit memory segment"); @@ -58,6 +61,8 @@ impl Segment { self.finalized = true; } + // Note: We do pointer arithmetic on `ptr` passed to `Segment::new` here. + // This assumes that `ptr` is valid for `len` bytes, or will result in UB. fn allocate(&mut self, size: usize, align: usize) -> *mut u8 { assert!(self.has_space_for(size, align)); self.position = align_up(self.position, align); @@ -93,7 +98,7 @@ impl ArenaMemoryProvider { /// Create a new memory region with the given size. pub fn new_with_size(reserve_size: usize) -> Result { let size = align_up(reserve_size, region::page::size()); - let mut alloc = region::alloc(reserve_size, region::Protection::NONE)?; + let mut alloc = region::alloc(size, region::Protection::NONE)?; let ptr = alloc.as_mut_ptr(); Ok(Self { @@ -108,26 +113,31 @@ impl ArenaMemoryProvider { fn allocate( &mut self, size: usize, - align: usize, + align: u64, protection: region::Protection, ) -> io::Result<*mut u8> { + let align = usize::try_from(align).expect("alignment too big"); + debug_assert!( + align <= region::page::size(), + "alignment over page size is not supported" + ); + // Note: Add a fast path without a linear scan over segments here? - // can we fit this allocation into an existing segment + // Can we fit this allocation into an existing segment. if let Some(segment) = self.segments.iter_mut().find(|seg| { seg.target_prot == protection && !seg.finalized && seg.has_space_for(size, align) }) { return Ok(segment.allocate(size, align)); } - // can we resize the last segment? + // Can we resize the last segment? if let Some(segment) = self.segments.iter_mut().last() { if segment.target_prot == protection && !segment.finalized { - let align = align.max(region::page::size()); - let additional_size = align_up(size, align); + let additional_size = align_up(size, region::page::size()); - // if our reserved arena can fit the additional size, extend the - // last segment + // If our reserved arena can fit the additional size, extend the + // last segment. if self.position + additional_size <= self.size { segment.len += additional_size; segment.set_rw(); @@ -137,8 +147,8 @@ impl ArenaMemoryProvider { } } - // allocate new segment for given size and alignment - self.allocate_segment(align_up(size, align), protection)?; + // Allocate new segment for given size and alignment. + self.allocate_segment(size, protection)?; let i = self.segments.len() - 1; Ok(self.segments[i].allocate(size, align)) } @@ -174,7 +184,7 @@ impl ArenaMemoryProvider { return; } self.segments.clear(); - // Drop the allocation, freeing memory + // Drop the allocation, freeing memory. let _: Option = self.alloc.take(); self.ptr = ptr::null_mut(); } @@ -196,15 +206,15 @@ impl Drop for ArenaMemoryProvider { impl JITMemoryProvider for ArenaMemoryProvider { fn allocate_readexec(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { - self.allocate(size, align as usize, region::Protection::READ_EXECUTE) + self.allocate(size, align, region::Protection::READ_EXECUTE) } fn allocate_readwrite(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { - self.allocate(size, align as usize, region::Protection::READ_WRITE) + self.allocate(size, align, region::Protection::READ_WRITE) } fn allocate_readonly(&mut self, size: usize, align: u64) -> io::Result<*mut u8> { - self.allocate(size, align as usize, region::Protection::READ) + self.allocate(size, align, region::Protection::READ) } unsafe fn free_memory(&mut self) { @@ -216,3 +226,42 @@ impl JITMemoryProvider for ArenaMemoryProvider { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn alignment_ok() { + let mut arena = ArenaMemoryProvider::new_with_size(1 << 20).unwrap(); + + for align_log2 in 0..8 { + let align = 1usize << align_log2; + for size in 1..128 { + let ptr = arena.allocate_readwrite(size, align as u64).unwrap(); + // assert!(ptr.is_aligned_to(align)); + assert_eq!(ptr.addr() % align, 0); + } + } + } + + #[test] + #[cfg(target_pointer_width = "64")] + fn large_virtual_allocation() { + let reserve_size = 1 << 40; // Request 1TB of memory. This should be successful. + let mut arena = ArenaMemoryProvider::new_with_size(reserve_size).unwrap(); + let ptr = arena.allocate_readwrite(1, 1).unwrap(); + assert_eq!(ptr.addr(), arena.ptr.addr()); + arena.finalize(BranchProtection::None); + unsafe { ptr.write_volatile(42) }; + unsafe { arena.free_memory() }; + } + + #[test] + fn over_capacity() { + let mut arena = ArenaMemoryProvider::new_with_size(1 << 20).unwrap(); // 1 MB + + let _ = arena.allocate_readwrite(900_000, 1).unwrap(); + let _ = arena.allocate_readwrite(200_000, 1).unwrap_err(); + } +} From b665341f66fe99a2cc0825ede14d4018efda0170 Mon Sep 17 00:00:00 2001 From: Mrmaxmeier Date: Mon, 7 Apr 2025 11:44:04 +0200 Subject: [PATCH 3/6] ArenaMemoryProvider: add asserts --- cranelift/jit/src/memory/arena.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cranelift/jit/src/memory/arena.rs b/cranelift/jit/src/memory/arena.rs index dd4a036939fa..f4ee184ab403 100644 --- a/cranelift/jit/src/memory/arena.rs +++ b/cranelift/jit/src/memory/arena.rs @@ -22,6 +22,9 @@ struct Segment { impl Segment { fn new(ptr: *mut u8, len: usize, target_prot: region::Protection) -> Self { + // Segments are created on page boundaries. + debug_assert_eq!(ptr as usize % region::page::size(), 0); + debug_assert_eq!(len % region::page::size(), 0); let mut segment = Segment { ptr, len, @@ -117,14 +120,14 @@ impl ArenaMemoryProvider { protection: region::Protection, ) -> io::Result<*mut u8> { let align = usize::try_from(align).expect("alignment too big"); - debug_assert!( + assert!( align <= region::page::size(), "alignment over page size is not supported" ); // Note: Add a fast path without a linear scan over segments here? - // Can we fit this allocation into an existing segment. + // Can we fit this allocation into an existing segment? if let Some(segment) = self.segments.iter_mut().find(|seg| { seg.target_prot == protection && !seg.finalized && seg.has_space_for(size, align) }) { From c7984580d2ebb70dc9274fcfc5d3d553b60192ea Mon Sep 17 00:00:00 2001 From: Mrmaxmeier Date: Mon, 7 Apr 2025 23:56:40 +0200 Subject: [PATCH 4/6] Move pipeline flush back out of `set_readable_and_executable` --- cranelift/jit/src/lib.rs | 5 +++-- cranelift/jit/src/memory/arena.rs | 3 +++ cranelift/jit/src/memory/mod.rs | 18 +++++++++--------- cranelift/jit/src/memory/system.rs | 3 +++ 4 files changed, 18 insertions(+), 11 deletions(-) diff --git a/cranelift/jit/src/lib.rs b/cranelift/jit/src/lib.rs index ee0012270f17..838668e1a9e7 100644 --- a/cranelift/jit/src/lib.rs +++ b/cranelift/jit/src/lib.rs @@ -9,9 +9,10 @@ mod backend; mod compiled_blob; mod memory; -pub use memory::*; - pub use crate::backend::{JITBuilder, JITModule}; +pub use crate::memory::{ + ArenaMemoryProvider, BranchProtection, JITMemoryProvider, SystemMemoryProvider, +}; /// Version number of this crate. pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/cranelift/jit/src/memory/arena.rs b/cranelift/jit/src/memory/arena.rs index f4ee184ab403..f503587ec260 100644 --- a/cranelift/jit/src/memory/arena.rs +++ b/cranelift/jit/src/memory/arena.rs @@ -178,6 +178,9 @@ impl ArenaMemoryProvider { for segment in &mut self.segments { segment.finalize(branch_protection); } + + // Flush any in-flight instructions from the pipeline + wasmtime_jit_icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); } /// Frees the allocated memory region, which would be leaked otherwise. diff --git a/cranelift/jit/src/memory/mod.rs b/cranelift/jit/src/memory/mod.rs index 0a042886622e..40c85ebe7705 100644 --- a/cranelift/jit/src/memory/mod.rs +++ b/cranelift/jit/src/memory/mod.rs @@ -1,9 +1,8 @@ -mod arena; -mod system; - +use cranelift_module::ModuleResult; use std::io; -use cranelift_module::ModuleResult; +mod arena; +mod system; pub use arena::ArenaMemoryProvider; pub use system::SystemMemoryProvider; @@ -32,8 +31,12 @@ pub trait JITMemoryProvider { fn finalize(&mut self, branch_protection: BranchProtection) -> ModuleResult<()>; } -/// Set all memory allocated in this `Memory` up to now as readable and executable. -/// This function deals with icache coherence, branch protection, and pipeline flushing. +/// Marks the memory region as readable and executable. +/// +/// This function deals with applies branch protection and clears the icache, +/// but *doesn't* flush the pipeline. Callers have to ensure that +/// [`wasmtime_jit_icache_coherence::pipeline_flush_mt`] is called before the +/// mappings are used. pub(crate) fn set_readable_and_executable( ptr: *mut u8, len: usize, @@ -73,8 +76,5 @@ pub(crate) fn set_readable_and_executable( } } - // Flush any in-flight instructions from the pipeline - wasmtime_jit_icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); - Ok(()) } diff --git a/cranelift/jit/src/memory/system.rs b/cranelift/jit/src/memory/system.rs index f25a03700aaa..715c155f50f3 100644 --- a/cranelift/jit/src/memory/system.rs +++ b/cranelift/jit/src/memory/system.rs @@ -175,6 +175,9 @@ impl Memory { super::set_readable_and_executable(ptr, len, branch_protection)?; } + // Flush any in-flight instructions from the pipeline + wasmtime_jit_icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); + self.already_protected = self.allocations.len(); Ok(()) } From acc6a09bf53c28756e8bfc232f50727bb33acec7 Mon Sep 17 00:00:00 2001 From: Mrmaxmeier Date: Tue, 8 Apr 2025 16:34:04 +0200 Subject: [PATCH 5/6] Fix aarch64 build --- cranelift/jit/src/memory/mod.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cranelift/jit/src/memory/mod.rs b/cranelift/jit/src/memory/mod.rs index 40c85ebe7705..9052d2da56e6 100644 --- a/cranelift/jit/src/memory/mod.rs +++ b/cranelift/jit/src/memory/mod.rs @@ -1,4 +1,4 @@ -use cranelift_module::ModuleResult; +use cranelift_module::{ModuleError, ModuleResult}; use std::io; mod arena; @@ -53,16 +53,16 @@ pub(crate) fn set_readable_and_executable( unsafe { region::protect(ptr, len, region::Protection::READ_EXECUTE).map_err(|e| { - cranelift_module::ModuleError::Backend( + ModuleError::Backend( anyhow::Error::new(e).context("unable to make memory readable+executable"), ) })?; } - // If BTI is requested, and the architecture supports it, use mprotect to set the PROT_BTI flag + // If BTI is requested, and the architecture supports it, use mprotect to set the PROT_BTI flag. if branch_protection == BranchProtection::BTI { #[cfg(all(target_arch = "aarch64", target_os = "linux"))] - if is_bti && std::arch::is_aarch64_feature_detected!("bti") { + if std::arch::is_aarch64_feature_detected!("bti") { let prot = libc::PROT_EXEC | libc::PROT_READ | /* PROT_BTI */ 0x10; unsafe { From 3aca8ef4d44c6aed708817ce5392a93e170396ef Mon Sep 17 00:00:00 2001 From: Mrmaxmeier Date: Wed, 9 Apr 2025 00:04:07 +0200 Subject: [PATCH 6/6] Document Windows overcommit limitation --- cranelift/jit/src/memory/arena.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/cranelift/jit/src/memory/arena.rs b/cranelift/jit/src/memory/arena.rs index f503587ec260..5a598fcfccb9 100644 --- a/cranelift/jit/src/memory/arena.rs +++ b/cranelift/jit/src/memory/arena.rs @@ -86,6 +86,10 @@ impl Segment { /// updated as the JIT requires more space. This approach allows for stable /// addresses throughout the lifetime of the JIT. /// +/// Depending on the underlying platform, requesting large parts of the address +/// space to be allocated might fail. This implementation currently doesn't do +/// overcommit on Windows. +/// /// Note: Memory will be leaked by default unless /// [`JITMemoryProvider::free_memory`] is called to ensure function pointers /// remain valid for the remainder of the program's life. @@ -101,6 +105,10 @@ impl ArenaMemoryProvider { /// Create a new memory region with the given size. pub fn new_with_size(reserve_size: usize) -> Result { let size = align_up(reserve_size, region::page::size()); + // Note: The region crate uses `MEM_RESERVE | MEM_COMMIT` on Windows. + // This means that allocations that exceed the page file plus system + // memory will fail here. + // https://github.com/darfink/region-rs/pull/34 let mut alloc = region::alloc(size, region::Protection::NONE)?; let ptr = alloc.as_mut_ptr(); @@ -252,9 +260,12 @@ mod tests { } #[test] - #[cfg(target_pointer_width = "64")] + #[cfg(all(target_pointer_width = "64", not(target_os = "windows")))] + // Windows: See https://github.com/darfink/region-rs/pull/34 fn large_virtual_allocation() { - let reserve_size = 1 << 40; // Request 1TB of memory. This should be successful. + // We should be able to request 1TB of virtual address space on 64-bit + // platforms. Physical memory should be committed as we go. + let reserve_size = 1 << 40; let mut arena = ArenaMemoryProvider::new_with_size(reserve_size).unwrap(); let ptr = arena.allocate_readwrite(1, 1).unwrap(); assert_eq!(ptr.addr(), arena.ptr.addr());