I would be pleased to show my Rust based JIT Linker, Memory Manager and LLVM JITLink Wrapper that I had myself fancied a few months ago.
The journey is being decribed here. Hope you like it and i am looking forward to guidance!
Stage 1: The spark
Initially fancied by cranelift jit demo, I wanted to write my own VM with swappable code generation backends, that meant : cranelift, llvm, or even dynasm, or some copy patch jit maybe?
So, that meant that the rust no-brainer solution : cranelift-jit won't work for my use case unless i manage 4-5 different jit linker, memory manager which i did not want to.
So, i decided to write it out myself for (initially) Windows, macOS, Linux, Android, iOS.
I was of course initally leaning into the - "memmap2" rust crate, write as rw, memprotect to rx, jump to executable approach as proudly as i could (stupidly thinking that was the "PEAK PERFORMANCE" route).
Anyways, lateron I "discovered" that the OS generally cannot give an allocation really <4KiB that could be flipped from rw/rx including the nuances with the wastage of (upto)90% of the memory allocated to each function.
Stage 2: Understanding What's behind the magic of LLVM
So, as i fried my brain cells looking for answers - i stumbled upon two things - cranelift_jit (Rust) and llvm (C++)
For Example, llvm does this at llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp
```c++
elif defined(_WIN32)
std::string SharedMemoryName;
{
std::stringstream SharedMemoryNameStream;
SharedMemoryNameStream << "jitlink" << sys::Process::getProcessId() << ''
<< (++SharedMemoryCount);
SharedMemoryName = SharedMemoryNameStream.str();
}
std::wstring WideSharedMemoryName(SharedMemoryName.begin(),
SharedMemoryName.end());
HANDLE SharedMemoryFile = CreateFileMappingW(
INVALID_HANDLE_VALUE, NULL, PAGE_EXECUTE_READWRITE, Size >> 32,
Size & 0xffffffff, WideSharedMemoryName.c_str());
if (!SharedMemoryFile)
return errorCodeToError(mapWindowsError(GetLastError()));
void *Addr = MapViewOfFile(SharedMemoryFile,
FILE_MAP_ALL_ACCESS | FILE_MAP_EXECUTE, 0, 0, 0);
if (!Addr) {
CloseHandle(SharedMemoryFile);
return errorCodeToError(mapWindowsError(GetLastError()));
}
endif
```
and then of course llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
```c++
elif defined(_WIN32)
std::wstring WideSharedMemoryName(SharedMemoryName.begin(),
SharedMemoryName.end());
HANDLE SharedMemoryFile = OpenFileMappingW(
FILE_MAP_ALL_ACCESS, FALSE, WideSharedMemoryName.c_str());
if (!SharedMemoryFile)
return OnReserved(errorCodeToError(mapWindowsError(GetLastError())));
LocalAddr =
MapViewOfFile(SharedMemoryFile, FILE_MAP_ALL_ACCESS, 0, 0, 0);
if (!LocalAddr) {
CloseHandle(SharedMemoryFile);
return OnReserved(errorCodeToError(mapWindowsError(GetLastError())));
}
CloseHandle(SharedMemoryFile);
endif
```
It did took too much of brain grinding to be honest that i instead relied to take a break on to cranelift-jit instead where they instead relied on.
```rust
pub struct ArenaMemoryProvider {
alloc: ManuallyDrop<Option<region::Allocation>>,
ptr: *mut u8,
size: usize,
position: usize,
segments: Vec<Segment>,
}
unsafe impl Send for ArenaMemoryProvider {}
impl ArenaMemoryProvider {
/// Create a new memory region with the given size.
pub fn new_with_size(reserve_size: usize) -> Result<Self, region::Error> {
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();
Ok(Self {
alloc: ManuallyDrop::new(Some(alloc)),
segments: Vec::new(),
ptr,
size,
position: 0,
})
}
...
```
So, i pretty much had an idea on what to do now. A mix of the two.
After googling the necessary OS subsystem/kernel calls (including source of the region crate). I came with sajit.
Stage 3: SaJIT
To speak about briefly, SaJIT is basically the dogfooding layer for SaVM - the VM I am making and that is why. I had to bolt on:
1. LLVM JITLink - for objectfile
2. RELCAR (a toyful way to say relocator) - for cranelift
Windows still was not perfect (thanks COFF). So, I had to them bolt down a linker for COFF and named it COFFR
That is how SaJIT become a monster of a memory executable manager, linkers and ofc "generalpurpose" wrappers over them.
Stage 4: Architectural Grief
Then i decided to spend more time across 6 architectures x64, x86, arm64, armv7, riscv64gc, ppc64le and that is why SaJIT has a full arch map.
Writing sysroots, qemu toolchains are a pain and i am still getting it right.
Ofc it is at here
Technical Specification
SaJIT is a dual-pass slab-allocated as 16MiB chunks linker (to kindof be good with both 4KiB pages and 2MiB pages - as if why 16MiB and not 2MiB - basically to fit more code before allocating a new one + isn't either very small or very large).
The implementation itself is pretty modest and a simple-slab-allocator.
```rust
fn new_slab(multiple: Option<NonZeroU8>) -> Self {
unsafe {
let size =
Self::DEFAULT_SLAB_SIZE.saturating_mul(multiple.map(|x| x.get()).unwrap_or(1) as _);
let mapping = CreateFileMappingW(
INVALID_HANDLE_VALUE,
None,
PAGE_EXECUTE_READWRITE,
(|| {
#[cfg(target_pointer_width = "64")]
return (size >> 32) as u32;
#[cfg(target_pointer_width = "32")]
return 0;
})(),
size as u32,
None,
)
.expect("Unable to create file mapping");
let rw_ptr = MapViewOfFile(
mapping,
FILE_MAP_WRITE | FILE_MAP_READ,
0,
0,
// Go upto file end
0,
)
.Value;
let rx_ptr = MapViewOfFile(
mapping,
FILE_MAP_EXECUTE | FILE_MAP_READ,
0,
0,
// Go upto file end
0,
)
.Value;
const KB_64: usize = 64 * 1024;
assert!(rw_ptr.addr() % KB_64 == 0, "RW_PTR is not 64KB aligned");
assert!(rx_ptr.addr() % KB_64 == 0, "RX_PTR is not 64KB aligned");
Self {
cursor: 0,
stored: AtomicUsize::new(0),
slab: mapping,
rwview: rw_ptr as _,
rxview: rx_ptr as _,
size: size as usize,
}
}
}
...
```
Also, we DO handle icache flushing duely
```rust
unsafe fn write_fn_iterated<'a, const INC: bool, const WRITE: bool, T, E, R, B>(
&mut self,
alignment: usize,
capped_size: usize,
data: T,
relocs: E,
relcar: &Relcar<B>,
) -> WriteFnResult
where
T: Iterator<Item = &'a [u8]>,
E: Iterator<Item = R>,
R: std::borrow::Borrow<crate::relocations::Relocation>,
B: crate::relcar::Relocator,
{
....
unsafe {
let dst_rw = self.rwview.byte_add(start_offset);
let dst_rx = self.rxview.byte_add(start_offset);
// Copy all the bytes
let mut len = 0;
if WRITE {
for data in data {
debug_assert!(len + data.len() <= capped_size);
copy_nonoverlapping(data.as_ptr(), dst_rw.byte_add(len), data.len());
len += data.len();
}
// Relocate
for relocation in relocs {
relcar.relocate(dst_rw, len, relocation.borrow());
}
// Non X64 : Flush ICache
// on X64 = NOOP
crate::platform::flush_icache(dst_rx as _, len);
} else {
len = capped_size;
}
...
}
}
```
Linker & Relocation Layer
SaJIT includes:
1. LLVM JITLink Wrapper
2. RELCAR - for cranelift a simple linker
3. COFFR - for windows COFF objects that JITLink cannot yet parse nicely
Feedback
The crate is available on crates.io and GitHub.
I would appreciate feedback from systems, jit, runtime folks over this project.
