From 776a3667a796e2bd33756ccbc182268ec132857d Mon Sep 17 00:00:00 2001 From: darmie Date: Mon, 24 Nov 2025 14:56:53 +0000 Subject: [PATCH 1/2] fix(jit): Add MAP_JIT and icache support for ARM64 macOS - Use mmap with MAP_JIT flag for JIT memory allocation - Call pthread_jit_write_protect_np(1) to switch to execute mode - Add sys_icache_invalidate for proper icache coherence - Custom Drop using munmap for MAP_JIT memory Fixes non-deterministic JIT execution failures on Apple Silicon. --- cranelift/jit/src/memory/mod.rs | 11 ++++ cranelift/jit/src/memory/system.rs | 71 +++++++++++++++++++++++-- crates/jit-icache-coherence/src/libc.rs | 19 ++++--- 3 files changed, 92 insertions(+), 9 deletions(-) diff --git a/cranelift/jit/src/memory/mod.rs b/cranelift/jit/src/memory/mod.rs index 9052d2da56e6..c3b6e59bd5b1 100644 --- a/cranelift/jit/src/memory/mod.rs +++ b/cranelift/jit/src/memory/mod.rs @@ -76,5 +76,16 @@ pub(crate) fn set_readable_and_executable( } } + // ARM64 macOS: Switch to execute mode for W^X compliance. + #[cfg(all(target_arch = "aarch64", target_os = "macos"))] + { + unsafe extern "C" { + fn pthread_jit_write_protect_np(enabled: libc::c_int); + } + unsafe { + pthread_jit_write_protect_np(1); + } + } + Ok(()) } diff --git a/cranelift/jit/src/memory/system.rs b/cranelift/jit/src/memory/system.rs index 08f8f0fd0a4a..94e34df69495 100644 --- a/cranelift/jit/src/memory/system.rs +++ b/cranelift/jit/src/memory/system.rs @@ -3,7 +3,11 @@ use cranelift_module::{ModuleError, ModuleResult}; #[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))] use memmap2::MmapMut; -#[cfg(not(any(feature = "selinux-fix", windows)))] +#[cfg(not(any( + feature = "selinux-fix", + windows, + all(target_arch = "aarch64", target_os = "macos") +)))] use std::alloc; use std::io; use std::mem; @@ -49,7 +53,45 @@ impl PtrLen { }) } - #[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))] + /// macOS ARM64: Use mmap with MAP_JIT for W^X policy compliance. + #[cfg(all( + target_arch = "aarch64", + target_os = "macos", + not(feature = "selinux-fix") + ))] + fn with_size(size: usize) -> io::Result { + assert_ne!(size, 0); + let alloc_size = region::page::ceil(size as *const ()) as usize; + + const MAP_JIT: libc::c_int = 0x0800; + + let ptr = unsafe { + libc::mmap( + ptr::null_mut(), + alloc_size, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_PRIVATE | libc::MAP_ANON | MAP_JIT, + -1, + 0, + ) + }; + + if ptr == libc::MAP_FAILED { + Err(io::Error::last_os_error()) + } else { + Ok(Self { + ptr: ptr as *mut u8, + len: alloc_size, + }) + } + } + + /// Non-macOS ARM64: Use standard allocator + #[cfg(all( + not(target_os = "windows"), + not(feature = "selinux-fix"), + not(all(target_arch = "aarch64", target_os = "macos")) + ))] fn with_size(size: usize) -> io::Result { assert_ne!(size, 0); let page_size = region::page::size(); @@ -95,7 +137,30 @@ impl PtrLen { } // `MMapMut` from `cfg(feature = "selinux-fix")` already deallocates properly. -#[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))] + +/// macOS ARM64: Free MAP_JIT memory with munmap. +#[cfg(all( + target_arch = "aarch64", + target_os = "macos", + not(feature = "selinux-fix") +))] +impl Drop for PtrLen { + fn drop(&mut self) { + if !self.ptr.is_null() { + unsafe { + let _ = region::protect(self.ptr, self.len, region::Protection::READ_WRITE); + libc::munmap(self.ptr as *mut libc::c_void, self.len); + } + } + } +} + +/// Other Unix platforms: Use standard allocator dealloc. +#[cfg(all( + not(target_os = "windows"), + not(feature = "selinux-fix"), + not(all(target_arch = "aarch64", target_os = "macos")) +))] impl Drop for PtrLen { fn drop(&mut self) { if !self.ptr.is_null() { diff --git a/crates/jit-icache-coherence/src/libc.rs b/crates/jit-icache-coherence/src/libc.rs index bac0a68d1933..614e7f6708df 100644 --- a/crates/jit-icache-coherence/src/libc.rs +++ b/crates/jit-icache-coherence/src/libc.rs @@ -141,16 +141,23 @@ fn riscv_flush_icache(start: u64, end: u64) -> Result<()> { pub(crate) use details::*; +// macOS ARM64: Use sys_icache_invalidate for icache coherence. +#[cfg(all(target_arch = "aarch64", target_os = "macos"))] +unsafe extern "C" { + fn sys_icache_invalidate(start: *const c_void, len: usize); +} + /// See docs on [crate::clear_cache] for a description of what this function is trying to do. #[inline] pub(crate) fn clear_cache(_ptr: *const c_void, _len: usize) -> Result<()> { - // TODO: On AArch64 we currently rely on the `mprotect` call that switches the memory from W+R - // to R+X to do this for us, however that is an implementation detail and should not be relied - // upon. - // We should call some implementation of `clear_cache` here. - // - // See: https://github.com/bytecodealliance/wasmtime/issues/3310 + // macOS ARM64: Use sys_icache_invalidate for icache coherence. + #[cfg(all(target_arch = "aarch64", target_os = "macos"))] + unsafe { + sys_icache_invalidate(_ptr, _len); + } + #[cfg(all(target_arch = "riscv64", target_os = "linux"))] riscv_flush_icache(_ptr as u64, (_ptr as u64) + (_len as u64))?; + Ok(()) } From 6942622684d7e205585466ba3152f790b6092d23 Mon Sep 17 00:00:00 2001 From: darmie Date: Mon, 24 Nov 2025 16:20:17 +0000 Subject: [PATCH 2/2] fix(jit): Address PR review feedback for ARM64 macOS JIT - Use libc::MAP_JIT instead of custom constant - Add pthread_jit_write_protect_np(0) after mmap to enable write mode - Change doc comments to regular comments for implementation details - Add TODO comment about AArch64 icache reliance on mprotect --- cranelift/jit/src/memory/mod.rs | 2 +- cranelift/jit/src/memory/system.rs | 32 +++++++++++++++---------- crates/jit-icache-coherence/src/libc.rs | 7 ++++++ 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/cranelift/jit/src/memory/mod.rs b/cranelift/jit/src/memory/mod.rs index c3b6e59bd5b1..925ff535b435 100644 --- a/cranelift/jit/src/memory/mod.rs +++ b/cranelift/jit/src/memory/mod.rs @@ -76,7 +76,7 @@ pub(crate) fn set_readable_and_executable( } } - // ARM64 macOS: Switch to execute mode for W^X compliance. + // ARM64 macOS: Switch to execute mode (W^X: disable writing, enable execution). #[cfg(all(target_arch = "aarch64", target_os = "macos"))] { unsafe extern "C" { diff --git a/cranelift/jit/src/memory/system.rs b/cranelift/jit/src/memory/system.rs index 94e34df69495..ff4f0e931e3c 100644 --- a/cranelift/jit/src/memory/system.rs +++ b/cranelift/jit/src/memory/system.rs @@ -53,7 +53,7 @@ impl PtrLen { }) } - /// macOS ARM64: Use mmap with MAP_JIT for W^X policy compliance. + // macOS ARM64: Use mmap with MAP_JIT for W^X policy compliance. #[cfg(all( target_arch = "aarch64", target_os = "macos", @@ -63,30 +63,36 @@ impl PtrLen { assert_ne!(size, 0); let alloc_size = region::page::ceil(size as *const ()) as usize; - const MAP_JIT: libc::c_int = 0x0800; - let ptr = unsafe { libc::mmap( ptr::null_mut(), alloc_size, libc::PROT_READ | libc::PROT_WRITE, - libc::MAP_PRIVATE | libc::MAP_ANON | MAP_JIT, + libc::MAP_PRIVATE | libc::MAP_ANON | libc::MAP_JIT, -1, 0, ) }; if ptr == libc::MAP_FAILED { - Err(io::Error::last_os_error()) - } else { - Ok(Self { - ptr: ptr as *mut u8, - len: alloc_size, - }) + return Err(io::Error::last_os_error()); + } + + // Enable writing to MAP_JIT memory (W^X: start in write mode) + unsafe extern "C" { + fn pthread_jit_write_protect_np(enabled: libc::c_int); } + unsafe { + pthread_jit_write_protect_np(0); + } + + Ok(Self { + ptr: ptr as *mut u8, + len: alloc_size, + }) } - /// Non-macOS ARM64: Use standard allocator + // Non-macOS ARM64: Use standard allocator. #[cfg(all( not(target_os = "windows"), not(feature = "selinux-fix"), @@ -138,7 +144,7 @@ impl PtrLen { // `MMapMut` from `cfg(feature = "selinux-fix")` already deallocates properly. -/// macOS ARM64: Free MAP_JIT memory with munmap. +// macOS ARM64: Free MAP_JIT memory with munmap. #[cfg(all( target_arch = "aarch64", target_os = "macos", @@ -155,7 +161,7 @@ impl Drop for PtrLen { } } -/// Other Unix platforms: Use standard allocator dealloc. +// Other Unix platforms: Use standard allocator dealloc. #[cfg(all( not(target_os = "windows"), not(feature = "selinux-fix"), diff --git a/crates/jit-icache-coherence/src/libc.rs b/crates/jit-icache-coherence/src/libc.rs index 614e7f6708df..137c239bfffa 100644 --- a/crates/jit-icache-coherence/src/libc.rs +++ b/crates/jit-icache-coherence/src/libc.rs @@ -150,6 +150,13 @@ unsafe extern "C" { /// See docs on [crate::clear_cache] for a description of what this function is trying to do. #[inline] pub(crate) fn clear_cache(_ptr: *const c_void, _len: usize) -> Result<()> { + // TODO: On AArch64 we currently rely on the `mprotect` call that switches the memory from W+R + // to R+X to do this for us, however that is an implementation detail and should not be relied + // upon. + // We should call some implementation of `clear_cache` here. + // + // See: https://github.com/bytecodealliance/wasmtime/issues/3310 + // macOS ARM64: Use sys_icache_invalidate for icache coherence. #[cfg(all(target_arch = "aarch64", target_os = "macos"))] unsafe {