Add create_mcjit_execution_engine_with_memory_manager for custom MCJIT memory management #566

ktanaka101 · 2025-01-25T15:30:53Z

Description

This PR adds a new function, create_mcjit_execution_engine_with_memory_manager, which allows users to supply their own JIT memory manager (McjitMemoryManager trait). This enables custom allocation, finalization, and teardown logic for code and data sections under MCJIT.

Related Issue

Close #296

How This Has Been Tested

Tested on a Mac M1 MAX host running Docker (Debian bookworm),
with LLVM 18 installed inside the container.

Manually and testcode verified on AArch64 Linux with a small JIT-compiled function.

f64 test_fn() {
entry:
  call void @llvm.experimental.stackmap(i64 12345, i32 0)
  ret f64 64.0
}

Confirmed that user-defined allocate_code_section, allocate_data_section, finalize_memory, and destroy callbacks are invoked by MCJIT.
No regressions observed in existing tests (the standard test suite still passes).

Checklist

I have read the Contributing Guide
My code follows the style and guidelines of this project
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas
I have added tests that prove my fix is effective or that my feature works

Changes

New Public API

pub fn create_mcjit_execution_engine_with_memory_manager(
    &self,
    memory_manager: impl McjitMemoryManager + 'static,
    opt_level: OptimizationLevel,
    code_model: CodeModel,
    no_frame_pointer_elim: bool,
    enable_fast_isel: bool,
) -> Result<ExecutionEngine<'ctx>, LLVMString> {
    // ...
    // Implementation that boxes the memory manager, creates an LLVMMCJITMemoryManager, 
    // and finalizes the MCJIT options before returning an `ExecutionEngine`.
    // ...
}

McjitMemoryManager Trait
- Allows the user to override how code/data sections are allocated and finalized.
- Includes allocate_code_section, allocate_data_section, finalize_memory, and destroy.
MemoryManagerAdapter
- Internal struct used to connect user-defined trait implementations to LLVM’s C callbacks (allocate_code_section_adapter, allocate_data_section_adapter, etc.).
- Automatically cleans up the memory manager upon destroy().

By merging this PR, Inkwell will support custom JIT memory management for MCJIT, addressing use cases like:

Allocating executable code in custom memory pools.
Intercepting .llvm_stackmaps for garbage collection.
Applying security restrictions or sandboxing JIT-compiled code.

Thank you for reviewing!

… custom MemoryManager

…ts across LLVM versions

TheDan64 · 2025-01-25T16:22:24Z

src/memory_manager.rs

+///
+/// The caller must ensure `ptr` points to a valid, null-terminated UTF-8 string.
+/// If the string is invalid UTF-8 or `ptr` is null, an empty string is returned.
+fn c_str_to_str<'a>(ptr: *const libc::c_char) -> &'a str {


We should mark this fn as unsafe; it's too easy to abuse with an arbitrary lifetime

TheDan64 · 2025-01-28T01:58:09Z

src/memory_manager.rs

+
+    // Re-box to drop the adapter and its contents
+    unsafe {
+        let _ = Box::from_raw(adapter);


I think it'd be safer to create the box upfront from the casted pointer; this currently leaves a moment where both the box and &mut exist pointing to the same data at the same time which would certainly be UB

TheDan64 · 2025-01-28T02:02:57Z

src/module.rs

+        }
+
+        // 4) Build LLVMMCJITCompilerOptions
+        let mut options: llvm_sys::execution_engine::LLVMMCJITCompilerOptions = unsafe { std::mem::zeroed() };


Please use MaybeUninit::zeroed() instead

TheDan64 · 2025-01-28T02:04:18Z

src/module.rs

+        options.OptLevel = opt_level as u32;
+        options.CodeModel = code_model.into();
+        options.NoFramePointerElim = if no_frame_pointer_elim { 1 } else { 0 };
+        options.EnableFastISel = if enable_fast_isel { 1 } else { 0 };


Could cast bool as int

TheDan64 · 2025-01-28T02:05:04Z

tests/all/test_execution_engine.rs

+    );
+    assert!(
+        second_result.is_err(),
+        "Expected an error when creating a second ExecutionEngine on the same module"


These tests are great!

TheDan64 self-requested a review January 25, 2025 16:06

TheDan64 assigned ktanaka101 Jan 25, 2025

ktanaka101 added 2 commits January 26, 2025 11:32

Add a function to the Module that creates an MCJIT capable of using a…

4e6d86a

… custom MemoryManager

Fix test failures caused by differing allocate_data_section call coun…

ba33e2b

…ts across LLVM versions

ktanaka101 force-pushed the feature-296-mcjit-memory-manager-access branch from 731165a to ba33e2b Compare January 26, 2025 02:33

TheDan64 reviewed Jan 28, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add create_mcjit_execution_engine_with_memory_manager for custom MCJIT memory management #566

Add create_mcjit_execution_engine_with_memory_manager for custom MCJIT memory management #566

ktanaka101 commented Jan 25, 2025 •

edited

Loading

TheDan64 Jan 25, 2025

TheDan64 Jan 28, 2025

TheDan64 Jan 28, 2025

TheDan64 Jan 28, 2025

TheDan64 Jan 28, 2025

Add create_mcjit_execution_engine_with_memory_manager for custom MCJIT memory management #566

Are you sure you want to change the base?

Add create_mcjit_execution_engine_with_memory_manager for custom MCJIT memory management #566

Conversation

ktanaka101 commented Jan 25, 2025 • edited Loading

Description

Related Issue

How This Has Been Tested

Checklist

Changes

TheDan64 Jan 25, 2025

Choose a reason for hiding this comment

TheDan64 Jan 28, 2025

Choose a reason for hiding this comment

TheDan64 Jan 28, 2025

Choose a reason for hiding this comment

TheDan64 Jan 28, 2025

Choose a reason for hiding this comment

TheDan64 Jan 28, 2025

Choose a reason for hiding this comment

ktanaka101 commented Jan 25, 2025 •

edited

Loading