feat(jit): lookup

This commit is contained in:
2025-05-18 21:52:36 +08:00
parent af5a312e1e
commit 4124156d52
4 changed files with 88 additions and 41 deletions

View File

@@ -72,7 +72,7 @@ impl From<Value<'_, '_>> for JITValue {
}
}
pub type JITFunc<'jit, 'vm> = unsafe extern "C" fn(*const VM<'jit>, *const Env<'jit, 'vm>, JITValue) -> JITValue;
pub type JITFunc<'jit, 'vm> = unsafe extern "C" fn(*const VM<'jit>, *const Env<'jit, 'vm>) -> JITValue;
pub struct JITContext<'ctx> {
context: &'ctx Context,
@@ -89,8 +89,6 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
let execution_engine = module
.create_jit_execution_engine(OptimizationLevel::Default)
.unwrap();
let helpers = Helpers::new(context, &module, &execution_engine);
JITContext {
@@ -99,7 +97,6 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
context,
module,
helpers
}
}
@@ -114,12 +111,12 @@ 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, &mut stack)?;
self.single_op(opcode, vm, env, &mut stack)?;
}
assert_eq!(stack.len(), 1);
let value = stack.pop();
self.builder.build_direct_call(self.helpers.debug, &[value.into()], "call_debug").unwrap();
self.builder.build_return(Some(&value))?;
if func_.verify(false) {
func_.print_to_stderr();
@@ -137,6 +134,7 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
&self,
opcode: OpCode,
vm: &'vm VM<'_>,
env: PointerValue<'ctx>,
stack: &mut Stack<BasicValueEnum<'ctx>, CAP>,
) -> Result<()> {
match opcode {
@@ -151,48 +149,59 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
}
}
OpCode::LoadThunk { idx } => stack.push(self.helpers.new_thunk(Rc::into_raw(Rc::new(Thunk::new(vm.get_thunk(idx))))))?,
/* OpCode::CaptureEnv => {
OpCode::CaptureEnv => {
let thunk = *stack.tos()?;
self.builder.build_direct_call(self.helpers.captured_env, &[thunk.into(), self.new_ptr(env).into()], "call_capture_env")?;
} */
self.builder.build_direct_call(self.helpers.capture_env, &[thunk.into(), env.into()], "call_capture_env")?;
}
OpCode::UnOp { op } => {
use UnOp::*;
let rhs = stack.pop();
stack.push(match op {
Neg => self.builder.build_direct_call(self.helpers.neg, &[rhs.into(), self.new_ptr(std::ptr::null::<Env>()).into()], "call_neg")?.try_as_basic_value().left().unwrap(),
Neg => self.builder.build_direct_call(self.helpers.neg, &[rhs.into(), env.into()], "call_neg")?.try_as_basic_value().left().unwrap(),
_ => todo!()
})?
}
/* OpCode::Func { idx } => {
OpCode::Func { idx } => {
let func = vm.get_func(idx);
let jit_func_ptr = self.compile_function(&Func::new(func, unsafe { env.as_ref() }.unwrap().clone(), OnceCell::new(), Cell::new(1)), vm)?;
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(jit_func_ptr as *const _ 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(), self.new_ptr(std::ptr::null::<Env>()).into(), func_ptr.into(), arg.into()], "call_func")?.try_as_basic_value().left().unwrap();
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();
let lhs = stack.pop();
match op {
BinOp::Add => {
let rhs = stack.pop();
let lhs = stack.pop();
let result = self.builder.build_direct_call(self.helpers.add, &[lhs.into(), rhs.into()], "call_add")?.try_as_basic_value().left().unwrap();
stack.push(result)?;
}
BinOp::Sub => {
let result = self.builder.build_direct_call(self.helpers.sub, &[lhs.into(), rhs.into()], "call_add")?.try_as_basic_value().left().unwrap();
stack.push(result)?;
}
BinOp::Eq => {
let result = self.builder.build_direct_call(self.helpers.add, &[lhs.into(), rhs.into()], "call_add")?.try_as_basic_value().left().unwrap();
stack.push(result)?;
}
_ => todo!("BinOp::{:?} not implemented in JIT", op),
}
}
OpCode::LookUp { sym } => {
stack.push(self.builder.build_direct_call(self.helpers.lookup, &[self.helpers.ptr_int_type.const_int(sym as u64, false).into(), env.into()], "call_lookup").unwrap().try_as_basic_value().left().unwrap())?
}
_ => todo!("{opcode:?} not implemented")
}
Ok(())