Chapter 5: JIT with LLVM

Chapter 4 lowered Kaleidoscope IR into the LLVM dialect. In this chapter we take the final step: convert that module to LLVM IR, JIT compile it, and invoke the generated machine code directly from Rust.

The implementation for this chapter lives in examples/kaleidoscope/jit.rs.

Design

The flow is intentionally small and linear:

1. Parse Kaleidoscope source into AST (parse_program).
2. Lower each AST function into Kaleidoscope dialect (lower_function).
3. Lower the module to LLVM dialect (lower_module).
4. Convert LLVM dialect to LLVM IR (to_llvm_ir::convert_module).
5. Build an LLVM ORC JIT instance, add the module, look up a symbol, execute.

The one extra safeguard in this implementation is a runtime signature check. Before JIT invocation, we verify that the target function really has signature fn(i64) -> i64. This limitation is intentionally set for simplicity.

Step 1-4: Lower source all the way to LLVM IR

lower_to_llvm_ir is a helper that runs the complete frontend and lowering pipeline and returns an LLVMModule:

#![allow(unused)]
fn main() {
fn lower_to_llvm_ir(src: &str, llvm_ctx: &LLVMContext) -> Result<LLVMModule> {
    let funcs =
        parse_program(src).map_err(|e| input_error_noloc!("Failed to parse program: {}", e))?;
    let ctx = &mut Context::new();
    let module = ModuleOp::new(ctx, "test".try_into().expect("valid module name"));
    for func in &funcs {
        let func_op = lower_function(ctx, func)?;
        module.append_operation(ctx, func_op.get_operation(), 0);
    }
    lower_module(ctx, module)?;
    verify_operation(module.get_operation(), ctx)?;
    // Convert from LLVM dialect to LLVM IR
    let llvm_module = to_llvm_ir::convert_module(ctx, llvm_ctx, module)?;
    llvm_module
        .verify()
        .map_err(|e| input_error_noloc!("Generated LLVM module is invalid: {}", e))?;
    Ok(llvm_module)
}
}

Notes:

• We build a fresh Context and ModuleOp per call.
• Every parsed function is lowered and appended to the module.
• lower_module mutates the module in place from Kaleidoscope dialect to LLVM dialect.
• to_llvm_ir::convert_module produces real LLVM IR (LLVMModule) consumable by JIT.

Step 5: JIT compile and execute

The public API for execution is exec_fn:

#![allow(unused)]
fn main() {
pub fn exec_fn(src: &str, name: &str, arg: i64) -> Result<i64> {
    initialize_native()
        .map_err(|e| input_error_noloc!("Failed to initialize native target: {}", e))?;
    let llvm_ctx = LLVMContext::default();
    let llvm_module = lower_to_llvm_ir(src, &llvm_ctx)?;

    let Some(f) = llvm_get_named_function(&llvm_module, name) else {
        return Err(input_error_noloc!(
            "Function '{}' not found in generated LLVM module",
            name
        ));
    };
    let f_ty = llvm_global_get_value_type(f);
    let param_types = llvm_get_param_types(f_ty);
    let ret_type = llvm_get_return_type(f_ty);
    let llvm_int64_ty = llvm_int_type_in_context(&llvm_ctx, 64);
    if param_types.len() != 1 || param_types[0] != llvm_int64_ty || ret_type != llvm_int64_ty {
        return Err(input_error_noloc!(
            "Expected function '{}' to have exactly one parameter of type i64 and return type i64, but found different signature",
            name
        ));
    }

    // println!("Generated LLVM IR:\n{}", llvm_module.to_string());

    // JIT compile and execute the main function
    let lljit = pliron_llvm::llvm_sys::lljit::LLVMLLJIT::new_with_default_builder()
        .map_err(|e| input_error_noloc!("Failed to create JIT execution engine: {}", e))?;
    lljit
        .add_module(llvm_module)
        .map_err(|e| input_error_noloc!("Failed to add module to JIT: {}", e))?;
    let main_fn = lljit
        .lookup_symbol(name)
        .map_err(|e| input_error_noloc!("Failed to find main function in JIT: {}", e))?;

    let main_fn: extern "C" fn(i64) -> i64 = unsafe { std::mem::transmute(main_fn) };
    Ok(main_fn(arg))
}
}

Key parts of exec_fn:

• initialize_native() sets up the host target backend once per process.
• llvm_get_named_function checks that the requested symbol exists in the generated module.
• llvm_get_param_types and llvm_get_return_type are used to validate the function ABI as exactly one i64 argument and an i64 return value.
• LLVMLLJIT::new_with_default_builder() creates an ORC JIT instance.
• add_module loads the generated LLVM module into JIT.
• lookup_symbol(name) resolves the function entry address.
• std::mem::transmute casts the raw symbol address to extern "C" fn(i64) -> i64, then calls it.

If the function is missing, or has a different signature, exec_fn returns a structured pliron::result::Error instead of panicking.

Tests

The file includes three focused tests that exercise the end-to-end JIT path.

Fibonacci from source file:

#![allow(unused)]
fn main() {
    #[test]
    fn fibonacci_jit() {
        let src = std::fs::read_to_string("examples/kaleidoscope/fibonacci.kal")
            .expect("failed to read fibonacci.kal");
        let result = exec_fn(&src, "main", 5).expect("failed to execute main function");
        assert_eq!(result, 5);
    }
}

Factorial from source file:

#![allow(unused)]
fn main() {
    #[test]
    fn factorial_jit() {
        let src = std::fs::read_to_string("examples/kaleidoscope/factorial.kal")
            .expect("failed to read factorial.kal");
        let result = exec_fn(&src, "main", 5).expect("failed to execute main function");
        assert_eq!(result, 120);
    }
}

Inline if/else program:

#![allow(unused)]
fn main() {
    #[test]
    fn if_else_jit() {
        let src = "
            def abs(x) {
                var result = 0;
                if x < 0 {
                    result = 0 - x;
                } else {
                    result = x;
                }
                return result;
            }
        ";
        let result = exec_fn(src, "abs", 42).expect("failed to execute main function");
        assert_eq!(result, 42);
        let result = exec_fn(src, "abs", -42).expect("failed to execute main function");
        assert_eq!(result, 42);
        let result = exec_fn(src, "abs", 0).expect("failed to execute main function");
        assert_eq!(result, 0);
    }
}

These tests verify that:

• file-backed programs run correctly through JIT,
• recursive calls work (fib, factorial),
• structured control flow (if/else) lowers and executes correctly.

Try it out

Run individual JIT tests:

cargo test --example kaleidoscope -- --show-output fibonacci_jit
cargo test --example kaleidoscope -- --show-output factorial_jit
cargo test --example kaleidoscope -- --show-output if_else_jit

You can also run the example binary in this workspace and execute a chosen function from a .kal source file.

Running via the example CLI

The example binary in examples/kaleidoscope/main.rs is a thin wrapper around jit::exec_fn: it reads source from --input, selects a function with --fn, passes an integer argument via --arg, and prints the result.

Run Fibonacci’s main(5):

cargo run --example kaleidoscope -- --input examples/kaleidoscope/fibonacci.kal --fn main --arg 5

Run Factorial’s main(5):

cargo run --example kaleidoscope -- --input examples/kaleidoscope/factorial.kal --fn main --arg 5

If the function name is missing from the module, or does not match the expected fn(i64) -> i64 signature, the CLI reports the error returned by exec_fn.

Recap

At this point, the tutorial pipeline is complete:

1. Parse source text into AST.
2. Lower AST to a custom high-level dialect.
3. Lower to LLVM dialect.
4. Convert to LLVM IR.
5. JIT and execute native code.

This is the minimal but complete compiler path from source language to running machine code using pliron + pliron-llvm.