Rust の Cranelift で逆ポーランド記法を JIT コンパイル

Cranelift のコンパイラ基盤を使ったハローワールドメモ。

Cranelift is a low-level retargetable code generator. It translates a target-independent intermediate representation into executable machine code.

逆ポーランド記法をパースして JIT で計算させる

Cargo.toml

[package]
name = "cranelift-example"
version = "0.1.0"
edition = "2021"

[dependencies]
cranelift = "0.117.0"
cranelift-codegen = "0.117.0"
cranelift-frontend = "0.117.0"
cranelift-jit = "0.117.0"
cranelift-module = "0.117.0"
libc = "0.2"

bin/rpn.rs

use cranelift::codegen::ir::{types, AbiParam, Function, InstBuilder, Value};
use cranelift::frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_codegen::ir::UserFuncName;
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::Module;
use std::collections::VecDeque;

fn main() {
    // RPN expression input
    let rpn_expression = "5 1 2 + 4 * + 3 -";

    // Create a JIT builder and module.
    let jit_builder = JITBuilder::new(cranelift_module::default_libcall_names()).expect("Failed to create JITBuilder");
    let mut module = JITModule::new(jit_builder);

    // Create the main function signature.
    let mut ctx = module.make_context();
    let mut sig = module.make_signature();
    sig.returns.push(AbiParam::new(types::I32));
    let main_func = module.declare_function("main", cranelift_module::Linkage::Export, &sig).unwrap();

    ctx.func = Function::with_name_signature(UserFuncName::user(0, 1), sig);
    let mut builder_ctx = FunctionBuilderContext::new();
    let mut builder = FunctionBuilder::new(&mut ctx.func, &mut builder_ctx);

    // Create the entry block.
    let entry_block = builder.create_block();
    builder.switch_to_block(entry_block);
    builder.seal_block(entry_block);

    // Stack to hold values
    let mut stack: Vec<Value> = Vec::new();

    // Parse and evaluate the RPN expression
    let tokens: VecDeque<&str> = rpn_expression.split_whitespace().collect();
    for token in tokens {
        match token {
            "+" => {
                let b = stack.pop().expect("Stack underflow");
                let a = stack.pop().expect("Stack underflow");
                let result = builder.ins().iadd(a, b);
                stack.push(result);
            }
            "-" => {
                let b = stack.pop().expect("Stack underflow");
                let a = stack.pop().expect("Stack underflow");
                let result = builder.ins().isub(a, b);
                stack.push(result);
            }
            "*" => {
                let b = stack.pop().expect("Stack underflow");
                let a = stack.pop().expect("Stack underflow");
                let result = builder.ins().imul(a, b);
                stack.push(result);
            }
            "/" => {
                let b = stack.pop().expect("Stack underflow");
                let a = stack.pop().expect("Stack underflow");
                let result = builder.ins().sdiv(a, b);
                stack.push(result);
            }
            _ => {
                let value: i32 = token.parse().expect("Invalid token");
                let value = builder.ins().iconst(types::I32, value as i64);
                stack.push(value);
            }
        }
    }

    // The final result should be the only value left on the stack
    let result = stack.pop().expect("No result on stack");
    builder.ins().return_(&[result]);

    // Finalize the function.
    builder.finalize();

    println!("rpn_expression: {}", rpn_expression);
    println!("Compiled function: ");
    println!("{}", ctx.func.display());

    // Compile and run the function.
    module.define_function(main_func, &mut ctx).unwrap();
    module.clear_context(&mut ctx);
    module.finalize_definitions().expect("Failed to finalize definitions");

    let code_ptr = module.get_finalized_function(main_func);
    let code_fn = unsafe { std::mem::transmute::<_, fn() -> i32>(code_ptr) };
    let result = code_fn();

    println!("Result: {}", result); // Expected output: Result: 14
}

生成される Cranelift IR と実行ログ:

$ cargo run --bin rpn
   Compiling cranelift-example v0.1.0 (/Users/hk2a/devel/rust/cranelift-example)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 1.20s
     Running `target/debug/rpn`
rpn_expression: 5 1 2 + 4 * + 3 -
Compiled function:
function u0:1() -> i32 system_v {
block0:
    v0 = iconst.i32 5
    v1 = iconst.i32 1
    v2 = iconst.i32 2
    v3 = iadd v1, v2  ; v1 = 1, v2 = 2
    v4 = iconst.i32 4
    v5 = imul v3, v4  ; v4 = 4
    v6 = iadd v0, v5  ; v0 = 5
    v7 = iconst.i32 3
    v8 = isub v6, v7  ; v7 = 3
    return v8
}

Result: 14

逆ポーランド記法をパースして JIT で計算させる(スタックスロット)

bin/rpn.rs

use cranelift::codegen::ir::{types, AbiParam, Function, InstBuilder, StackSlot, StackSlotData, StackSlotKind, Value};
use cranelift::frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_codegen::ir::UserFuncName;
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::Module;
use std::collections::VecDeque;

fn main() {
    // RPN expression input
    let rpn_expression = "5 1 2 + 4 * + 3 -";

    // Create a JIT builder and module.
    let jit_builder = JITBuilder::new(cranelift_module::default_libcall_names()).expect("Failed to create JITBuilder");
    let mut module = JITModule::new(jit_builder);

    // Create the main function signature.
    let mut ctx = module.make_context();
    let mut sig = module.make_signature();
    sig.returns.push(AbiParam::new(types::I32));
    let main_func = module.declare_function("main", cranelift_module::Linkage::Export, &sig).unwrap();

    ctx.func = Function::with_name_signature(UserFuncName::user(0, 1), sig);
    let mut builder_ctx = FunctionBuilderContext::new();
    let mut builder = FunctionBuilder::new(&mut ctx.func, &mut builder_ctx);

    // Create the entry block.
    let entry_block = builder.create_block();
    builder.switch_to_block(entry_block);
    builder.seal_block(entry_block);

    // Stack slot to hold values
    let stack_slot = builder.create_sized_stack_slot(StackSlotData::new(StackSlotKind::ExplicitSlot, 16, 0));
    let mut offset = 0;

    // Parse and evaluate the RPN expression
    let tokens: VecDeque<&str> = rpn_expression.split_whitespace().collect();
    for token in tokens {
        match token {
            "+" => {
                let b = pop(&mut builder, stack_slot, &mut offset);
                let a = pop(&mut builder, stack_slot, &mut offset);
                let result = builder.ins().iadd(a, b);
                push(&mut builder, stack_slot, &mut offset, result);
            }
            "-" => {
                let b = pop(&mut builder, stack_slot, &mut offset);
                let a = pop(&mut builder, stack_slot, &mut offset);
                let result = builder.ins().isub(a, b);
                push(&mut builder, stack_slot, &mut offset, result);
            }
            "*" => {
                let b = pop(&mut builder, stack_slot, &mut offset);
                let a = pop(&mut builder, stack_slot, &mut offset);
                let result = builder.ins().imul(a, b);
                push(&mut builder, stack_slot, &mut offset, result);
            }
            "/" => {
                let b = pop(&mut builder, stack_slot, &mut offset);
                let a = pop(&mut builder, stack_slot, &mut offset);
                let result = builder.ins().sdiv(a, b);
                push(&mut builder, stack_slot, &mut offset, result);
            }
            _ => {
                let value: i32 = token.parse().expect("Invalid token");
                let value = builder.ins().iconst(types::I32, value as i64);
                push(&mut builder, stack_slot, &mut offset, value);
            }
        }
    }

    // The final result should be the only value left on the stack
    let result = pop(&mut builder, stack_slot, &mut offset);
    builder.ins().return_(&[result]);

    // Finalize the function.
    builder.finalize();

    println!("rpn_expression: {}", rpn_expression);
    println!("Compiled function: ");
    println!("{}", ctx.func.display());

    // Compile and run the function.
    module.define_function(main_func, &mut ctx).unwrap();
    module.clear_context(&mut ctx);
    module.finalize_definitions().expect("Failed to finalize definitions");

    let code_ptr = module.get_finalized_function(main_func);
    let code_fn = unsafe { std::mem::transmute::<_, fn() -> i32>(code_ptr) };
    let result = code_fn();

    println!("Result: {}", result); // Expected output: Result: 14
}

// Helper functions for stack operations
fn push(builder: &mut FunctionBuilder, stack_slot: StackSlot, offset: &mut i32, value: Value) {
    builder.ins().stack_store(value, stack_slot, *offset);
    *offset += 4;
}

fn pop(builder: &mut FunctionBuilder, stack_slot: StackSlot, offset: &mut i32) -> Value {
    *offset -= 4;
    builder.ins().stack_load(types::I32, stack_slot, *offset)
}

生成される Cranelift IR と実行ログ:

$ cargo run --bin rpn
   Compiling cranelift-example v0.1.0 (/Users/hk2a/devel/rust/cranelift-example)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.74s
     Running `target/debug/rpn`
rpn_expression: 5 1 2 + 4 * + 3 -
Compiled function:
function u0:1() -> i32 system_v {
    ss0 = explicit_slot 16

block0:
    v0 = iconst.i32 5
    stack_store v0, ss0  ; v0 = 5
    v1 = iconst.i32 1
    stack_store v1, ss0+4  ; v1 = 1
    v2 = iconst.i32 2
    stack_store v2, ss0+8  ; v2 = 2
    v3 = stack_load.i32 ss0+8
    v4 = stack_load.i32 ss0+4
    v5 = iadd v4, v3
    stack_store v5, ss0+4
    v6 = iconst.i32 4
    stack_store v6, ss0+8  ; v6 = 4
    v7 = stack_load.i32 ss0+8
    v8 = stack_load.i32 ss0+4
    v9 = imul v8, v7
    stack_store v9, ss0+4
    v10 = stack_load.i32 ss0+4
    v11 = stack_load.i32 ss0
    v12 = iadd v11, v10
    stack_store v12, ss0
    v13 = iconst.i32 3
    stack_store v13, ss0+4  ; v13 = 3
    v14 = stack_load.i32 ss0+4
    v15 = stack_load.i32 ss0
    v16 = isub v15, v14
    stack_store v16, ss0
    v17 = stack_load.i32 ss0
    return v17
}

Result: 14

システムコールする例

libc の printf をコールして Hello, World を出力。

bin/hello.rs

use cranelift::codegen::ir::{types, AbiParam, Function, InstBuilder};
use cranelift::frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_codegen::ir::UserFuncName;
use cranelift_jit::{JITBuilder, JITModule};
use std::ffi::CString;
use cranelift_module::{Module, Linkage};
use libc;

fn main() {
    // Create a JIT builder and module.
    let mut jit_builder = JITBuilder::new(cranelift_module::default_libcall_names()).expect("Failed to create JITBuilder");
    jit_builder.symbol("printf", printf as *const u8);
    let mut module = JITModule::new(jit_builder);

    // Create a function signature for `printf`.
    let mut ctx = module.make_context();
    let mut sig = module.make_signature();
    let pointer_type = module.target_config().pointer_type();
    sig.params.push(AbiParam::new(pointer_type)); // フォーマット文字列の引数
    sig.params.push(AbiParam::new(pointer_type)); // 可変引数のためのダミー引数
    sig.returns.push(AbiParam::new(types::I32));
    let printf = module.declare_function("printf", Linkage::Import, &sig).unwrap();

    // Create the main function signature.
    let mut sig = module.make_signature();
    sig.returns.push(AbiParam::new(types::I32));
    let main_func = module.declare_function("main", Linkage::Export, &sig).unwrap();

    ctx.func = Function::with_name_signature(UserFuncName::user(0, 1), sig);
    let mut builder_ctx = FunctionBuilderContext::new();
    let mut builder = FunctionBuilder::new(&mut ctx.func, &mut builder_ctx);

    // Create the entry block.
    let entry_block = builder.create_block();
    builder.append_block_params_for_function_params(entry_block);
    builder.switch_to_block(entry_block);
    builder.seal_block(entry_block);

    // Create the string data.
    let hello_world = CString::new("Hello, World!\n").unwrap();
    let hello_world_ptr = hello_world.as_ptr() as i64;

    // Call the `printf` function.
    let printf_func = module.declare_func_in_func(printf, builder.func);
    let format_str = builder.ins().iconst(types::I64, hello_world_ptr);
    let zero = builder.ins().iconst(types::I64, 0);
    let call = builder.ins().call(printf_func, &[format_str, zero]);

    // Return the result of the call.
    let result = builder.inst_results(call)[0];
    builder.ins().return_(&[result]);

    builder.finalize();
    println!("{}", ctx.func.display());

    // Compile and run the function.
    module.define_function(main_func, &mut ctx).unwrap();
    module.clear_context(&mut ctx);
    module.finalize_definitions().expect("Failed to finalize definitions");

    let code_ptr = module.get_finalized_function(main_func);
    let code_fn = unsafe { std::mem::transmute::<_, fn() -> i32>(code_ptr) };
    code_fn();
}

// Dummy printf function to link with.
extern "C" fn printf(fmt: *const i8, _dummy: i64) -> i32 {
    unsafe {
        libc::printf(fmt)
    }
}

生成される Cranelift IR と実行ログ:

$ cargo run --bin hello
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.47s
     Running `target/debug/hello`
function u0:1() -> i32 system_v {
    sig0 = (i64, i64) -> i32 system_v
    fn0 = u0:0 sig0

block0:
    v0 = iconst.i64 0x6000_0258_0180
    v1 = iconst.i64 0
    v2 = call fn0(v0, v1)  ; v0 = 0x6000_0258_0180, v1 = 0
    return v2
}

Hello, World!

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です