imxyy1soope1/nixjit
1
0
Fork 0
Code Issues Pull Requests Actions Projects 1 Releases Activity

feat(jit): fib!

Browse Source
This commit is contained in:
imxyy_soope_
2025-05-19 19:29:25 +08:00
parent 6d26716412
commit 9e172bf013
10 changed files with 363 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
src/jit/helpers.rs
View File

@@ -1,5 +1,3 @@
use std::rc::Rc;
use inkwell::AddressSpace;
use inkwell::context::Context;
use inkwell::execution_engine::ExecutionEngine;
@@ -8,8 +6,8 @@ use inkwell::types::{FloatType, FunctionType, IntType, PointerType, StructType};
use inkwell::values::{BasicValueEnum, FunctionValue};
use crate::jit::JITValueData;
use crate::ty::internal::Thunk;
use crate::vm::{Env, VM};
use crate::ty::internal::{Thunk, Value};
use crate::vm::{LetEnv, VM};
use super::{JITFunc, JITValue, ValueTag};
@@ -25,10 +23,14 @@ pub struct Helpers<'ctx> {
pub debug: FunctionValue<'ctx>,
pub capture_env: FunctionValue<'ctx>,
pub neg: FunctionValue<'ctx>,
pub not: FunctionValue<'ctx>,
pub add: FunctionValue<'ctx>,
pub sub: FunctionValue<'ctx>,
pub eq: FunctionValue<'ctx>,
pub or: FunctionValue<'ctx>,
pub call: FunctionValue<'ctx>,
pub lookup: FunctionValue<'ctx>,
pub force: FunctionValue<'ctx>,
}
impl<'ctx> Helpers<'ctx> {
@@ -61,6 +63,11 @@ impl<'ctx> Helpers<'ctx> {
value_type.fn_type(&[value_type.into(), ptr_type.into()], false),
None,
);
let not = module.add_function(
"not",
value_type.fn_type(&[value_type.into(), ptr_type.into()], false),
None,
);
let add = module.add_function(
"add",
value_type.fn_type(&[value_type.into(), value_type.into()], false),
@@ -71,15 +78,23 @@ impl<'ctx> Helpers<'ctx> {
value_type.fn_type(&[value_type.into(), value_type.into()], false),
None,
);
// Assuming a single argument for now based on the test case
let eq = module.add_function(
"eq",
value_type.fn_type(&[value_type.into(), value_type.into()], false),
None,
);
let or = module.add_function(
"or",
value_type.fn_type(&[value_type.into(), value_type.into()], false),
None,
);
let call = module.add_function(
"call",
value_type.fn_type(
&[
ptr_type.into(),
ptr_type.into(),
ptr_type.into(),
value_type.into(),
value_type.into(),
ptr_type.into(),
],
false,
),
@@ -90,14 +105,23 @@ impl<'ctx> Helpers<'ctx> {
value_type.fn_type(&[ptr_int_type.into(), ptr_type.into()], false),
None,
);
let force = module.add_function(
"force",
value_type.fn_type(&[value_type.into(), ptr_type.into()], false),
None,
);
execution_engine.add_global_mapping(&debug, helper_debug as _);
execution_engine.add_global_mapping(&capture_env, helper_capture_env as _);
execution_engine.add_global_mapping(&neg, helper_neg as _);
execution_engine.add_global_mapping(&not, helper_not as _);
execution_engine.add_global_mapping(&add, helper_add as _);
execution_engine.add_global_mapping(&sub, helper_sub as _);
execution_engine.add_global_mapping(&eq, helper_eq as _);
execution_engine.add_global_mapping(&or, helper_or as _);
execution_engine.add_global_mapping(&call, helper_call as _);
execution_engine.add_global_mapping(&lookup, helper_lookup as _);
execution_engine.add_global_mapping(&force, helper_force as _);
Helpers {
int_type,
@@ -111,10 +135,14 @@ impl<'ctx> Helpers<'ctx> {
debug,
capture_env,
neg,
not,
add,
sub,
eq,
or,
call,
lookup,
force
}
}
@@ -179,13 +207,13 @@ extern "C" fn helper_debug(value: JITValue) {
}
#[unsafe(no_mangle)]
extern "C" fn helper_capture_env(thunk: JITValue, env: *const Env) {
extern "C" fn helper_capture_env(thunk: JITValue, env: *const LetEnv) {
let thunk: &Thunk = unsafe { std::mem::transmute(thunk.data.ptr.as_ref().unwrap()) };
thunk.capture(unsafe { env.as_ref().unwrap() }.clone());
}
#[unsafe(no_mangle)]
extern "C" fn helper_neg(rhs: JITValue, _env: *const Env) -> JITValue {
extern "C" fn helper_neg(rhs: JITValue, _env: *const LetEnv) -> JITValue {
use ValueTag::*;
match rhs.tag {
Int => JITValue {
@@ -198,6 +226,20 @@ extern "C" fn helper_neg(rhs: JITValue, _env: *const Env) -> JITValue {
}
}
#[unsafe(no_mangle)]
extern "C" fn helper_not(rhs: JITValue, _env: *const LetEnv) -> JITValue {
use ValueTag::*;
match rhs.tag {
Bool => JITValue {
tag: Bool,
data: JITValueData {
bool: !unsafe { rhs.data.bool },
},
},
_ => todo!(),
}
}
#[unsafe(no_mangle)]
extern "C" fn helper_add(lhs: JITValue, rhs: JITValue) -> JITValue {
use ValueTag::*;
@@ -235,20 +277,68 @@ extern "C" fn helper_sub(lhs: JITValue, rhs: JITValue) -> JITValue {
}
#[unsafe(no_mangle)]
extern "C" fn helper_call<'jit, 'vm>(
vm: *const VM<'jit>,
env: *const Env<'jit, 'vm>,
func_ptr: *const (),
arg: JITValue,
) -> JITValue {
let func: JITFunc = unsafe { std::mem::transmute(func_ptr) };
todo!();
unsafe { func(vm, env) }
extern "C" fn helper_eq(lhs: JITValue, rhs: JITValue) -> JITValue {
use ValueTag::*;
match (lhs.tag, rhs.tag) {
(Int, Int) => JITValue {
tag: Bool,
data: JITValueData {
bool: unsafe { lhs.data.int == rhs.data.int }
},
},
_ => todo!(
"Equation not implemented for {:?} and {:?}",
lhs.tag,
rhs.tag
),
}
}
#[unsafe(no_mangle)]
extern "C" fn helper_lookup<'jit, 'vm>(sym: usize, env: *const Env<'jit, 'vm>) -> JITValue {
extern "C" fn helper_or(lhs: JITValue, rhs: JITValue) -> JITValue {
use ValueTag::*;
match (lhs.tag, rhs.tag) {
(Bool, Bool) => JITValue {
tag: Bool,
data: JITValueData {
bool: unsafe { lhs.data.bool || rhs.data.bool },
},
},
_ => todo!(
"Substruction not implemented for {:?} and {:?}",
lhs.tag,
rhs.tag
),
}
}
#[unsafe(no_mangle)]
extern "C" fn helper_call<'jit, 'vm>(
func: JITValue,
arg: JITValue,
vm: *const VM<'jit>,
) -> JITValue {
use ValueTag::*;
match func.tag {
Function => {
let func: Value = func.into();
func.call(unsafe { vm.as_ref() }.unwrap(), vec![arg.into()]).unwrap().into()
}
_ => todo!(),
}
}
#[unsafe(no_mangle)]
extern "C" fn helper_lookup<'jit, 'vm>(sym: usize, env: *const LetEnv<'jit, 'vm>) -> JITValue {
let env = unsafe { env.as_ref() }.unwrap();
let val = env.lookup(sym);
dbg!(val.as_ref().unwrap().typename());
val.unwrap().into()
}
#[unsafe(no_mangle)]
extern "C" fn helper_force<'jit, 'vm>(thunk: JITValue, vm: *const VM<'jit>) -> JITValue {
let mut val = Value::from(thunk);
val.force(unsafe { vm.as_ref() }.unwrap()).unwrap();
val.into()
}

218
src/jit/mod.rs
View File

@@ -1,18 +1,19 @@
use std::rc::Rc;
use inkwell::OptimizationLevel;
use inkwell::basic_block::BasicBlock;
use inkwell::builder::Builder;
use inkwell::context::Context;
use inkwell::execution_engine::{ExecutionEngine, JitFunction};
use inkwell::module::Module;
use inkwell::values::{BasicValueEnum, PointerValue};
use inkwell::values::{BasicValueEnum, FunctionValue, PointerValue};
use crate::bytecode::{Func, OpCode, UnOp};
use crate::error::*;
use crate::stack::Stack;
use crate::ty::common::Const;
use crate::ty::internal::{Thunk, Value};
use crate::vm::{Env, VM};
use crate::vm::{LetEnv, VM};
mod helpers;
@@ -48,17 +49,19 @@ pub struct JITValue {
pub union JITValueData {
int: i64,
float: f64,
boolean: bool,
bool: bool,
ptr: *const (),
}
impl<'jit: 'vm, 'vm> Into<Value<'jit, 'vm>> for JITValue {
fn into(self) -> Value<'jit, 'vm> {
impl<'jit: 'vm, 'vm> From<JITValue> for Value<'jit, 'vm> {
fn from(value: JITValue) -> Self {
use ValueTag::*;
match self.tag {
Int => Value::Const(Const::Int(unsafe { self.data.int })),
match value.tag {
Int => Value::Const(Const::Int(unsafe { value.data.int })),
Null => Value::Const(Const::Null),
_ => todo!("not implemented for {:?}", self.tag),
Function => Value::Func(unsafe { Rc::from_raw(value.data.ptr as *const _) }),
Thunk => Value::Thunk(unsafe { Rc::from_raw(value.data.ptr as *const _) }),
_ => todo!("not implemented for {:?}", value.tag),
}
}
}
@@ -70,13 +73,25 @@ impl From<Value<'_, '_>> for JITValue {
tag: ValueTag::Int,
data: JITValueData { int },
},
Value::Func(func) => JITValue {
tag: ValueTag::Function,
data: JITValueData {
ptr: Rc::into_raw(func) as *const _
}
},
Value::Thunk(thunk) => JITValue {
tag: ValueTag::Thunk,
data: JITValueData {
ptr: Rc::into_raw(thunk) as *const _
}
},
_ => todo!(),
}
}
}
pub type JITFunc<'jit, 'vm> =
unsafe extern "C" fn(*const VM<'jit>, *const Env<'jit, 'vm>) -> JITValue;
unsafe extern "C" fn(*const VM<'jit>, *const LetEnv<'jit, 'vm>) -> JITValue;
pub struct JITContext<'ctx> {
context: &'ctx Context,
@@ -125,17 +140,25 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
let func_ = self
.module
.add_function("nixjit_function", self.helpers.func_type, None);
let entry = self.context.append_basic_block(func_, "entry");
self.builder.position_at_end(entry);
let env = func_.get_nth_param(1).unwrap().into_pointer_value();
while let Some(opcode) = iter.next() {
self.single_op(opcode, vm, env, &mut stack)?;
}
let entry = self.context.append_basic_block(func_, "entry");
self.build_expr(
&mut iter,
vm,
env,
&mut stack,
func_,
entry,
func.opcodes.len(),
)?;
assert_eq!(stack.len(), 1);
let value = stack.pop();
let exit = self.context.append_basic_block(func_, "exit");
self.builder.build_unconditional_branch(exit)?;
self.builder.position_at_end(exit);
self.builder.build_return(Some(&value))?;
if func_.verify(false) {
func_.print_to_stderr();
if func_.verify(true) {
unsafe {
let name = func_.get_name().to_str().unwrap();
let addr = self.execution_engine.get_function(name).unwrap();
@@ -146,13 +169,82 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
}
}
fn build_expr<const CAP: usize>(
&self,
iter: &mut impl Iterator<Item = OpCode>,
vm: &'vm VM<'_>,
env: PointerValue<'ctx>,
stack: &mut Stack<BasicValueEnum<'ctx>, CAP>,
func: FunctionValue<'ctx>,
bb: BasicBlock<'ctx>,
mut length: usize,
) -> Result<usize> {
self.builder.position_at_end(bb);
while length > 1 {
let opcode = iter.next().unwrap();
let br = self.single_op(opcode, vm, env, stack)?;
length -= 1;
if br > 0 {
let consq = self.context.append_basic_block(func, "consq");
let alter = self.context.append_basic_block(func, "alter");
let cont = self.context.append_basic_block(func, "cont");
let cond = self
.builder
.build_alloca(self.helpers.value_type, "cond_alloca")?;
let result = self
.builder
.build_alloca(self.helpers.value_type, "result_alloca")?;
self.builder.build_store(cond, stack.pop())?;
self.builder.build_conditional_branch(
self.builder
.build_load(
self.context.bool_type(),
self.builder.build_struct_gep(
self.helpers.value_type,
cond,
1,
"gep_cond",
)?,
"load_cond",
)?
.into_int_value(),
consq,
alter,
)?;
length -= br;
let br = self.build_expr(iter, vm, env, stack, func, consq, br)?;
self.builder.build_store(result, stack.pop())?;
self.builder.build_unconditional_branch(cont)?;
length -= br;
self.build_expr(iter, vm, env, stack, func, alter, br)?;
self.builder.build_store(result, stack.pop())?;
self.builder.build_unconditional_branch(cont)?;
self.builder.position_at_end(cont);
stack.push(self.builder.build_load(
self.helpers.value_type,
result,
"load_result",
)?)?;
}
}
if length > 0 {
self.single_op(iter.next().unwrap(), vm, env, stack)
} else {
Ok(0)
}
}
#[inline(always)]
fn single_op<const CAP: usize>(
&self,
opcode: OpCode,
vm: &'vm VM<'_>,
env: PointerValue<'ctx>,
stack: &mut Stack<BasicValueEnum<'ctx>, CAP>,
) -> Result<()> {
) -> Result<usize> {
match opcode {
OpCode::Const { idx } => {
use Const::*;
@@ -186,54 +278,14 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
.try_as_basic_value()
.left()
.unwrap(),
_ => todo!(),
Not => self
.builder
.build_direct_call(self.helpers.not, &[rhs.into(), env.into()], "call_neg")?
.try_as_basic_value()
.left()
.unwrap(),
})?
}
OpCode::Func { idx } => {
let func = vm.get_func(idx);
let jit_func_ptr = self.compile_function(func, vm)?;
let jit_value = self
.helpers
.value_type
.const_named_struct(&[
self.helpers
.int_type
.const_int(ValueTag::Function as _, false)
.into(),
self.helpers
.ptr_int_type
.const_int(unsafe { jit_func_ptr.as_raw() } as _, false)
.into(),
])
.into();
stack.push(jit_value)?;
}
OpCode::Call { arity } => {
// Assuming arity is 1 for the test case
assert_eq!(arity, 1);
let arg = stack.pop();
let func_value = stack.pop();
let func_ptr = self
.builder
.build_extract_value(func_value.into_struct_value(), 1, "func_ptr")?
.into_pointer_value();
let result = self
.builder
.build_direct_call(
self.helpers.call,
&[
self.new_ptr(vm).into(),
env.into(),
func_ptr.into(),
arg.into(),
],
"call_func",
)?
.try_as_basic_value()
.left()
.unwrap();
stack.push(result)?;
}
OpCode::BinOp { op } => {
use crate::bytecode::BinOp;
let rhs = stack.pop();
@@ -269,9 +321,22 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
let result = self
.builder
.build_direct_call(
self.helpers.add,
self.helpers.eq,
&[lhs.into(), rhs.into()],
"call_add",
"call_eq",
)?
.try_as_basic_value()
.left()
.unwrap();
stack.push(result)?;
}
BinOp::Or => {
let result = self
.builder
.build_direct_call(
self.helpers.or,
&[lhs.into(), rhs.into()],
"call_or",
)?
.try_as_basic_value()
.left()
@@ -299,8 +364,29 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
.left()
.unwrap(),
)?,
OpCode::Call { arity } => {
// TODO:
assert_eq!(arity, 1);
let mut args = Vec::with_capacity(arity);
for _ in 0..arity {
args.insert(0, stack.pop());
}
let func = self.builder
.build_direct_call(self.helpers.force, &[stack.pop().into(), self.new_ptr(vm).into()], "force")?
.try_as_basic_value()
.left()
.unwrap();
let ret = self.builder
.build_direct_call(self.helpers.call, &[func.into(), args[0].into(), self.new_ptr(vm).into()], "call")?
.try_as_basic_value()
.left()
.unwrap();
stack.push(ret)?;
}
OpCode::JmpIfFalse { step } => return Ok(step),
OpCode::Jmp { step } => return Ok(step),
_ => todo!("{opcode:?} not implemented"),
}
Ok(())
Ok(0)
}
}