diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 9b4656da5ad..730386ef755 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1,6 +1,7 @@ // spell-checker: disable use super::{JitCompileError, JitSig, JitType}; use alloc::collections::BTreeSet; +use alloc::rc::Rc; use cranelift::codegen::ir::FuncRef; use cranelift::prelude::*; use num_traits::cast::ToPrimitive; @@ -15,13 +16,23 @@ enum CustomTrapCode { /// Raised when shifting by a negative number NegativeShiftCount = 1, } - +#[derive(Debug)] +enum ObjectKind { + Opaque, + Tuple(Rc), +} +#[derive(Debug)] +enum ElementShape { + Scalar(JitType), + Object(Rc), +} +#[derive(Debug)] +struct TupleShape(Vec); #[derive(Clone)] -struct Local { - var: Variable, - ty: JitType, +enum Local { + Scalar { var: Variable, ty: JitType }, + Object { var: Variable, kind: Rc }, } - #[derive(Debug)] enum JitValue { Int(Value), @@ -29,7 +40,7 @@ enum JitValue { Bool(Value), None, Null, - Tuple(Vec), + Object(Value, Rc), FuncRef(FuncRef), } @@ -39,6 +50,7 @@ impl JitValue { JitType::Int => Self::Int(val), JitType::Float => Self::Float(val), JitType::Bool => Self::Bool(val), + JitType::Object => Self::Object(val, Rc::new(ObjectKind::Opaque)), } } @@ -47,14 +59,16 @@ impl JitValue { Self::Int(_) => Some(JitType::Int), Self::Float(_) => Some(JitType::Float), Self::Bool(_) => Some(JitType::Bool), - Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None, + Self::Object(_, _) => Some(JitType::Object), + Self::None | Self::Null | Self::FuncRef(_) => None, } } fn into_value(self) -> Option { match self { Self::Int(val) | Self::Float(val) | Self::Bool(val) => Some(val), - Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None, + Self::Object(val, _) => Some(val), + Self::None | Self::Null | Self::FuncRef(_) => None, } } } @@ -70,9 +84,9 @@ pub(crate) struct FunctionCompiler<'a, 'b> { stack: Vec, variables: Box<[Option]>, label_to_block: HashMap, + alloc_func: FuncRef, pub(crate) sig: JitSig, } - impl<'a, 'b> FunctionCompiler<'a, 'b> { pub(crate) fn new( builder: &'a mut FunctionBuilder<'b>, @@ -80,6 +94,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { arg_types: &[JitType], ret_type: Option, entry_block: Block, + alloc_func: FuncRef, ) -> Self { let mut compiler = Self { builder, @@ -90,6 +105,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { args: arg_types.to_vec(), ret: ret_type, }, + alloc_func, }; let params = compiler.builder.func.dfg.block_params(entry_block).to_vec(); for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() { @@ -103,30 +119,52 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { compiler } + /// Pops Multiple Values off the Stack, returning a Vector fn pop_multiple(&mut self, count: usize) -> Vec { let stack_len = self.stack.len(); self.stack.drain(stack_len - count..).collect() } - fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> { - #[expect(clippy::mut_mut, reason = "This seems like a false positive")] - let builder = &mut self.builder; - let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?; - let local = self.variables[idx].get_or_insert_with(|| { - let var = builder.declare_var(ty.to_cranelift()); - Local { - var, - ty: ty.clone(), + /// Generates a local out of a jitvalue + fn local_from_value( + builder: &mut FunctionBuilder<'_>, + value: JitValue, + ) -> Result { + match value { + JitValue::Object(ptr, kind) => { + let var = builder.declare_var(types::I64); + builder.def_var(var, ptr); + Ok(Local::Object { var, kind }) + } + + scalar => { + let ty = scalar.to_jit_type().ok_or(JitCompileError::NotSupported)?; + let var = builder.declare_var(ty.to_cranelift()); + builder.def_var(var, scalar.into_value().unwrap()); + + Ok(Local::Scalar { var, ty }) } - }); - if ty != local.ty { - Err(JitCompileError::NotSupported) - } else { - self.builder.def_var(local.var, val.into_value().unwrap()); - Ok(()) } } + /// Generates a JitValue out of a Local + fn local_to_value(builder: &mut FunctionBuilder<'_>, local: &Local) -> JitValue { + match local { + Local::Scalar { var, ty } => { + JitValue::from_type_and_value(ty.clone(), builder.use_var(*var)) + } + Local::Object { var, kind } => JitValue::Object(builder.use_var(*var), kind.clone()), + } + } + /// Stores a Local into the variables array + fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> { + #[expect(clippy::mut_mut, reason = "This seems like a false positive")] + let builder = &mut self.builder; + let local = Self::local_from_value(builder, val)?; + self.variables[idx] = Some(local); + Ok(()) + } + /// Generates a boolean value out of a JitValue, Nulls FuncRefs and Objects aren't supported fn boolean_val(&mut self, val: JitValue) -> Result { match val { JitValue::Float(val) => { @@ -141,7 +179,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } JitValue::Bool(val) => Ok(val), JitValue::None => Ok(self.builder.ins().iconst(types::I8, 0)), - JitValue::Null | JitValue::Tuple(_) | JitValue::FuncRef(_) => { + JitValue::Null | JitValue::FuncRef(_) | JitValue::Object(_, _) => { Err(JitCompileError::NotSupported) } } @@ -242,7 +280,6 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { if label_targets.contains(&label) { // Create or get the block for this label: let target_block = self.get_or_create_block(label); - // If the current block isn't terminated, add a fallthrough jump if let Some(cur) = self.builder.current_block() && cur != target_block @@ -321,34 +358,84 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let val = self.builder.ins().iconst(types::I8, value as i64); JitValue::Bool(val) } + BorrowedConstant::Tuple { elements } => { + let mut vals = Vec::new(); + for el in elements { + let b = el.borrow_constant(); + let res = self.prepare_const(b)?; + vals.push(res); + } + let (ptr, shape) = self.build_heap_tuple(vals)?; + JitValue::Object(ptr, Rc::new(ObjectKind::Tuple(shape))) + } BorrowedConstant::None => JitValue::None, _ => return Err(JitCompileError::NotSupported), }; Ok(value) } - - fn return_value(&mut self, val: JitValue) -> Result<(), JitCompileError> { - if let Some(ref ty) = self.sig.ret { - // If the signature has a return type, enforce it - if val.to_jit_type().as_ref() != Some(ty) { - return Err(JitCompileError::NotSupported); + fn build_heap_tuple( + &mut self, + elements: Vec, + ) -> Result<(Value, Rc), JitCompileError> { + let len_val = self.builder.ins().iconst(types::I64, elements.len() as i64); + let call = self.builder.ins().call(self.alloc_func, &[len_val]); + let ptr = self.builder.inst_results(call)[0]; + let mut shape = Vec::with_capacity(elements.len()); + for (i, element) in elements.into_iter().enumerate() { + let offset = 8 + 8 * i as i32; + match element { + JitValue::Object(val, kind) => { + self.builder.ins().store(MemFlags::new(), val, ptr, offset); + shape.push(ElementShape::Object(kind)); + } + scalar => { + let ty = scalar.to_jit_type().ok_or(JitCompileError::NotSupported)?; + let val = scalar.into_value().unwrap(); + self.builder.ins().store(MemFlags::new(), val, ptr, offset); + shape.push(ElementShape::Scalar(ty)); + } } - } else { - // First time we see a return, define it in the signature - let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?; - self.sig.ret = Some(ty.clone()); - self.builder - .func - .signature - .returns - .push(AbiParam::new(ty.to_cranelift())); } + Ok((ptr, Rc::new(TupleShape(shape)))) + } + fn unpack_heap_tuple(&mut self, ptr: Value, shape: &TupleShape) -> Vec { + let mut values = Vec::with_capacity(shape.0.len()); + for (i, element_shape) in shape.0.iter().enumerate() { + let offset = 8 + 8 * i as i32; + let value = match element_shape { + ElementShape::Scalar(ty) => { + let val = + self.builder + .ins() + .load(ty.to_cranelift(), MemFlags::new(), ptr, offset); + JitValue::from_type_and_value(ty.clone(), val) + } + ElementShape::Object(kind) => { + let val = self + .builder + .ins() + .load(types::I64, MemFlags::new(), ptr, offset); + JitValue::Object(val, kind.clone()) + } + }; + values.push(value); + } + values + } + fn return_value(&mut self, val: JitValue) -> Result<(), JitCompileError> { + let expected = self.sig.ret.as_ref().ok_or(JitCompileError::NotSupported)?; - // If this is e.g. an Int, Float, or Bool we have a Cranelift `Value`. - // If we have JitValue::None or .Tuple(...) but can't handle that, error out (or handle differently). - let cr_val = val.into_value().ok_or(JitCompileError::NotSupported)?; + let value = match (expected, val) { + (JitType::Int, JitValue::Int(v)) => v, + (JitType::Float, JitValue::Float(v)) => v, + (JitType::Float, JitValue::Int(v)) => self.builder.ins().fcvt_from_sint(types::F64, v), + + (JitType::Bool, JitValue::Bool(v)) => v, + (JitType::Object, JitValue::Object(v, _)) => v, + _ => return Err(JitCompileError::NotSupported), + }; - self.builder.ins().return_(&[cr_val]); + self.builder.ins().return_(&[value]); Ok(()) } @@ -524,7 +611,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } Instruction::BuildTuple { count } => { let elements = self.pop_multiple(count.get(arg) as usize); - self.stack.push(JitValue::Tuple(elements)); + let (ptr, shape) = self.build_heap_tuple(elements)?; + self.stack + .push(JitValue::Object(ptr, Rc::new(ObjectKind::Tuple(shape)))); Ok(()) } Instruction::Call { argc } => { @@ -656,10 +745,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let local = self.variables[var_num.get(arg)] .as_ref() .ok_or(JitCompileError::BadBytecode)?; - self.stack.push(JitValue::from_type_and_value( - local.ty.clone(), - self.builder.use_var(local.var), - )); + let value = Self::local_to_value(self.builder, local); + self.stack.push(value); Ok(()) } Instruction::LoadFastLoadFast { var_nums } @@ -675,10 +762,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let local = self.variables[idx] .as_ref() .ok_or(JitCompileError::BadBytecode)?; - self.stack.push(JitValue::from_type_and_value( - local.ty.clone(), - self.builder.use_var(local.var), - )); + let value = Self::local_to_value(self.builder, local); + self.stack.push(value); } Ok(()) } @@ -751,10 +836,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let local = self.variables[load_idx] .as_ref() .ok_or(JitCompileError::BadBytecode)?; - self.stack.push(JitValue::from_type_and_value( - local.ty.clone(), - self.builder.use_var(local.var), - )); + let value = Self::local_to_value(self.builder, local); + self.stack.push(value); Ok(()) } Instruction::StoreFastStoreFast { var_nums } => { @@ -801,23 +884,23 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } Instruction::UnpackSequence { count } => { let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; - - let elements = match val { - JitValue::Tuple(elements) => elements, + let (ptr, shape) = match val { + JitValue::Object(ptr, kind) => match &*kind { + ObjectKind::Tuple(shape) => (ptr, shape.clone()), + _ => return Err(JitCompileError::NotSupported), + }, _ => return Err(JitCompileError::NotSupported), }; - - if elements.len() != count.get(arg) as usize { + if shape.0.len() != count.get(arg) as usize { return Err(JitCompileError::NotSupported); } - + let elements = self.unpack_heap_tuple(ptr, &shape); self.stack.extend(elements.into_iter().rev()); Ok(()) } _ => Err(JitCompileError::NotSupported), } } - fn compile_sub(&mut self, a: Value, b: Value) -> Value { let (out, carry) = self.builder.ins().ssub_overflow(a, b); self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW); @@ -1273,6 +1356,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.builder.ins().jump(merge_block, &[nan_f.into()]); self.builder.switch_to_block(continue_neg_inf); + // b is an integer here; convert b_floor to an i64. let b_i64 = self.builder.ins().fcvt_to_sint(i64_ty, b_floor); let one_i = self.builder.ins().iconst(i64_ty, 1); diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index dbaa4a3eb26..8b25e088c70 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -1,7 +1,8 @@ +#![allow(clippy::alloc_instead_of_core)] mod instructions; extern crate alloc; - +use alloc::alloc::{Layout, alloc, handle_alloc_error}; use alloc::fmt; use core::mem::ManuallyDrop; use cranelift::prelude::*; @@ -36,20 +37,42 @@ pub enum JitArgumentError { WrongNumberOfArguments, } +/// Allocates a tuple in the heap and returns it's pointer, Unsafe function +extern "C" fn jit_alloc_tuple(len: i64) -> i64 { + let layout = Layout::array::(len as usize + 1).unwrap(); + let ptr = unsafe { alloc(layout) }; + if ptr.is_null() { + handle_alloc_error(layout); + } + unsafe { + *(ptr as *mut i64) = len; + } + ptr as i64 +} + struct Jit { builder_context: FunctionBuilderContext, ctx: codegen::Context, + alloc_func_id: FuncId, module: JITModule, } impl Jit { fn new() -> Self { - let builder = JITBuilder::new(cranelift_module::default_libcall_names()) + let mut builder = JITBuilder::new(cranelift_module::default_libcall_names()) .expect("Failed to build JITBuilder"); - let module = JITModule::new(builder); + builder.symbol("jit_alloc_tuple", jit_alloc_tuple as *const u8); + let mut module = JITModule::new(builder); + let mut alloc_sig = Signature::new(module.target_config().default_call_conv); + alloc_sig.params.push(AbiParam::new(types::I64)); + alloc_sig.returns.push(AbiParam::new(types::I64)); + let alloc_func_id = module + .declare_function("jit_alloc_tuple", Linkage::Import, &alloc_sig) + .expect("Failed to declare jit_alloc_tuple"); Self { builder_context: FunctionBuilderContext::new(), ctx: module.make_context(), + alloc_func_id, module, } } @@ -83,6 +106,9 @@ impl Jit { )?; let func_ref = self.module.declare_func_in_func(id, &mut self.ctx.func); + let alloc_func_ref = self + .module + .declare_func_in_func(self.alloc_func_id, &mut self.ctx.func); let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context); let entry_block = builder.create_block(); @@ -96,6 +122,7 @@ impl Jit { args, ret, entry_block, + alloc_func_ref, ); compiler.compile(func_ref, bytecode)?; @@ -193,6 +220,7 @@ pub enum JitType { Int, Float, Bool, + Object, } impl JitType { @@ -201,6 +229,7 @@ impl JitType { Self::Int => types::I64, Self::Float => types::F64, Self::Bool => types::I8, + Self::Object => types::I64, } } @@ -209,6 +238,7 @@ impl JitType { Self::Int => libffi::middle::Type::i64(), Self::Float => libffi::middle::Type::f64(), Self::Bool => libffi::middle::Type::u8(), + Self::Object => libffi::middle::Type::pointer(), } } } @@ -219,6 +249,7 @@ pub enum AbiValue { Float(f64), Int(i64), Bool(bool), + Object(usize), } impl AbiValue { @@ -227,6 +258,7 @@ impl AbiValue { Self::Int(i) => libffi::middle::Arg::new(i), Self::Float(f) => libffi::middle::Arg::new(f), Self::Bool(b) => libffi::middle::Arg::new(b), + Self::Object(p) => libffi::middle::Arg::new(p), } } } @@ -249,6 +281,12 @@ impl From for AbiValue { } } +impl From for AbiValue { + fn from(p: usize) -> Self { + Self::Object(p) + } +} + impl TryFrom for i64 { type Error = (); @@ -282,11 +320,23 @@ impl TryFrom for bool { } } +impl TryFrom for usize { + type Error = (); + + fn try_from(value: AbiValue) -> Result { + match value { + AbiValue::Object(p) => Ok(p), + _ => Err(()), + } + } +} + fn type_check(ty: &JitType, val: &AbiValue) -> Result<(), JitArgumentError> { match (ty, val) { (JitType::Int, AbiValue::Int(_)) | (JitType::Float, AbiValue::Float(_)) - | (JitType::Bool, AbiValue::Bool(_)) => Ok(()), + | (JitType::Bool, AbiValue::Bool(_)) + | (JitType::Object, AbiValue::Object(_)) => Ok(()), _ => Err(JitArgumentError::ArgumentTypeMismatch), } } @@ -296,6 +346,7 @@ union UnTypedAbiValue { float: f64, int: i64, boolean: u8, + object: usize, _void: (), } @@ -306,6 +357,7 @@ impl UnTypedAbiValue { JitType::Int => AbiValue::Int(self.int), JitType::Float => AbiValue::Float(self.float), JitType::Bool => AbiValue::Bool(self.boolean != 0), + JitType::Object => AbiValue::Object(self.object), } } } diff --git a/crates/jit/tests/bool_tests.rs b/crates/jit/tests/bool_tests.rs index 8a5f4ea9db3..e1e55452261 100644 --- a/crates/jit/tests/bool_tests.rs +++ b/crates/jit/tests/bool_tests.rs @@ -3,7 +3,7 @@ mod tests { #[test] fn basic_return() { let return_ = jit_function! { return_(a: bool) -> bool => r##" - def return_(a: bool): + def return_(a: bool) -> bool: return a "## }; @@ -14,13 +14,13 @@ mod tests { #[test] fn basic_const() { let const_true = jit_function! { const_true(a: i64) -> bool => r##" - def const_true(a: int): + def const_true(a: int) -> bool: return True "## }; assert_eq!(const_true(0), Ok(true)); let const_false = jit_function! { const_false(a: i64) -> bool => r##" - def const_false(a: int): + def const_false(a: int) -> bool: return False "## }; assert_eq!(const_false(0), Ok(false)); @@ -29,7 +29,7 @@ mod tests { #[test] fn basic_not() { let not_ = jit_function! { not_(a: bool) -> bool => r##" - def not_(a: bool): + def not_(a: bool) -> bool: return not a "## }; @@ -40,7 +40,7 @@ mod tests { #[test] fn basic_if_not() { let if_not = jit_function! { if_not(a: bool) -> i64 => r##" - def if_not(a: bool): + def if_not(a: bool) -> int: if not a: return 0 else: @@ -56,7 +56,7 @@ mod tests { #[test] fn basic_eq() { let eq = jit_function! { eq(a:bool, b:bool) -> i64 => r##" - def eq(a: bool, b: bool): + def eq(a: bool, b: bool) -> int: if a == b: return 1 return 0 @@ -71,7 +71,7 @@ mod tests { #[test] fn eq_with_integers() { let eq = jit_function! { eq(a:bool, b:i64) -> i64 => r##" - def eq(a: bool, b: int): + def eq(a: bool, b: int) -> int: if a == b: return 1 return 0 @@ -86,7 +86,7 @@ mod tests { #[test] fn basic_gt() { let gt = jit_function! { gt(a:bool, b:bool) -> i64 => r##" - def gt(a: bool, b: bool): + def gt(a: bool, b: bool) -> int: if a > b: return 1 return 0 @@ -101,7 +101,7 @@ mod tests { #[test] fn gt_with_integers() { let gt = jit_function! { gt(a:i64, b:bool) -> i64 => r##" - def gt(a: int, b: bool): + def gt(a: int, b: bool) -> int: if a > b: return 1 return 0 @@ -116,7 +116,7 @@ mod tests { #[test] fn basic_lt() { let lt = jit_function! { lt(a:bool, b:bool) -> i64 => r##" - def lt(a: bool, b: bool): + def lt(a: bool, b: bool) -> int: if a < b: return 1 return 0 @@ -131,7 +131,7 @@ mod tests { #[test] fn lt_with_integers() { let lt = jit_function! { lt(a:i64, b:bool) -> i64 => r##" - def lt(a: int, b: bool): + def lt(a: int, b: bool) -> int: if a < b: return 1 return 0 @@ -146,7 +146,7 @@ mod tests { #[test] fn basic_gte() { let gte = jit_function! { gte(a:bool, b:bool) -> i64 => r##" - def gte(a: bool, b: bool): + def gte(a: bool, b: bool) -> int: if a >= b: return 1 return 0 @@ -161,7 +161,7 @@ mod tests { #[test] fn gte_with_integers() { let gte = jit_function! { gte(a:bool, b:i64) -> i64 => r##" - def gte(a: bool, b: int): + def gte(a: bool, b: int) -> int: if a >= b: return 1 return 0 @@ -176,7 +176,7 @@ mod tests { #[test] fn basic_lte() { let lte = jit_function! { lte(a:bool, b:bool) -> i64 => r##" - def lte(a: bool, b: bool): + def lte(a: bool, b: bool) -> int: if a <= b: return 1 return 0 @@ -191,7 +191,7 @@ mod tests { #[test] fn lte_with_integers() { let lte = jit_function! { lte(a:bool, b:i64) -> i64 => r##" - def lte(a: bool, b: int): + def lte(a: bool, b: int) -> int: if a <= b: return 1 return 0 diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index b9bbb3ea63c..12a90e548be 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -21,7 +21,7 @@ mod tests { #[test] fn basic_add() { let add = jit_function! { add(a:f64, b:f64) -> f64 => r##" - def add(a: float, b: float): + def add(a: float, b: float) -> float: return a + b "## }; @@ -37,7 +37,7 @@ mod tests { #[test] fn add_with_integer() { let add = jit_function! { add(a:f64, b:i64) -> f64 => r##" - def add(a: float, b: int): + def add(a: float, b: int) -> float: return a + b "## }; @@ -49,7 +49,7 @@ mod tests { #[test] fn basic_sub() { let sub = jit_function! { sub(a:f64, b:f64) -> f64 => r##" - def sub(a: float, b: float): + def sub(a: float, b: float) -> float: return a - b "## }; @@ -66,7 +66,7 @@ mod tests { #[test] fn sub_with_integer() { let sub = jit_function! { sub(a:i64, b:f64) -> f64 => r##" - def sub(a: int, b: float): + def sub(a: int, b: float) -> float: return a - b "## }; @@ -79,7 +79,7 @@ mod tests { #[test] fn basic_mul() { let mul = jit_function! { mul(a:f64, b:f64) -> f64 => r##" - def mul(a: float, b: float): + def mul(a: float, b: float) -> float: return a * b "## }; @@ -100,7 +100,7 @@ mod tests { #[test] fn mul_with_integer() { let mul = jit_function! { mul(a:f64, b:i64) -> f64 => r##" - def mul(a: float, b: int): + def mul(a: float, b: int) -> float: return a * b "## }; @@ -115,7 +115,7 @@ mod tests { #[test] fn basic_power() { let pow = jit_function! { pow(a:f64, b:f64) -> f64 => r##" - def pow(a:float, b: float): + def pow(a:float, b: float) -> float: return a**b "##}; // Test base cases @@ -226,7 +226,7 @@ mod tests { #[test] fn basic_div() { let div = jit_function! { div(a:f64, b:f64) -> f64 => r##" - def div(a: float, b: float): + def div(a: float, b: float) -> float: return a / b "## }; @@ -247,7 +247,7 @@ mod tests { #[test] fn div_with_integer() { let div = jit_function! { div(a:f64, b:i64) -> f64 => r##" - def div(a: float, b: int): + def div(a: float, b: int) -> float: return a / b "## }; @@ -264,7 +264,7 @@ mod tests { #[test] fn basic_if_bool() { let if_bool = jit_function! { if_bool(a:f64) -> i64 => r##" - def if_bool(a: float): + def if_bool(a: float) -> int: if a: return 1 return 0 @@ -281,7 +281,7 @@ mod tests { #[test] fn basic_float_eq() { let float_eq = jit_function! { float_eq(a: f64, b: f64) -> bool => r##" - def float_eq(a: float, b: float): + def float_eq(a: float, b: float) -> bool: return a == b "## }; @@ -298,7 +298,7 @@ mod tests { #[test] fn basic_float_ne() { let float_ne = jit_function! { float_ne(a: f64, b: f64) -> bool => r##" - def float_ne(a: float, b: float): + def float_ne(a: float, b: float) -> bool: return a != b "## }; @@ -315,7 +315,7 @@ mod tests { #[test] fn basic_float_gt() { let float_gt = jit_function! { float_gt(a: f64, b: f64) -> bool => r##" - def float_gt(a: float, b: float): + def float_gt(a: float, b: float) -> bool: return a > b "## }; @@ -332,7 +332,7 @@ mod tests { #[test] fn basic_float_gte() { let float_gte = jit_function! { float_gte(a: f64, b: f64) -> bool => r##" - def float_gte(a: float, b: float): + def float_gte(a: float, b: float) -> bool: return a >= b "## }; @@ -349,7 +349,7 @@ mod tests { #[test] fn basic_float_lt() { let float_lt = jit_function! { float_lt(a: f64, b: f64) -> bool => r##" - def float_lt(a: float, b: float): + def float_lt(a: float, b: float) -> bool: return a < b "## }; @@ -366,7 +366,7 @@ mod tests { #[test] fn basic_float_lte() { let float_lte = jit_function! { float_lte(a: f64, b: f64) -> bool => r##" - def float_lte(a: float, b: float): + def float_lte(a: float, b: float) -> bool: return a <= b "## }; diff --git a/crates/jit/tests/int_tests.rs b/crates/jit/tests/int_tests.rs index 23cf98aafe1..290b8c5f1fd 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -5,7 +5,7 @@ mod tests { #[test] fn basic_add() { let add = jit_function! { add(a:i64, b:i64) -> i64 => r##" - def add(a: int, b: int): + def add(a: int, b: int) -> int: return a + b "## }; @@ -17,7 +17,7 @@ mod tests { #[test] fn basic_sub() { let sub = jit_function! { sub(a:i64, b:i64) -> i64 => r##" - def sub(a: int, b: int): + def sub(a: int, b: int) -> int: return a - b "## }; @@ -30,7 +30,7 @@ mod tests { #[test] fn basic_mul() { let mul = jit_function! { mul(a:i64, b:i64) -> i64 => r##" - def mul(a: int, b: int): + def mul(a: int, b: int) -> int: return a * b "## }; @@ -49,7 +49,7 @@ mod tests { #[test] fn basic_div() { let div = jit_function! { div(a:i64, b:i64) -> f64 => r##" - def div(a: int, b: int): + def div(a: int, b: int) -> float: return a / b "## }; @@ -74,7 +74,7 @@ mod tests { #[test] fn basic_floor_div() { let floor_div = jit_function! { floor_div(a:i64, b:i64) -> i64 => r##" - def floor_div(a: int, b: int): + def floor_div(a: int, b: int) -> int: return a // b "## }; @@ -89,7 +89,7 @@ mod tests { fn basic_exp() { let exp = jit_function! { exp(a: i64, b: i64) -> i64 => r##" - def exp(a: int, b: int): + def exp(a: int, b: int) -> int: return a ** b "## }; @@ -110,7 +110,7 @@ mod tests { #[test] fn basic_mod() { let modulo = jit_function! { modulo(a:i64, b:i64) -> i64 => r##" - def modulo(a: int, b: int): + def modulo(a: int, b: int) -> int: return a % b "## }; @@ -125,7 +125,7 @@ mod tests { #[test] fn basic_power() { let power = jit_function! { power(a:i64, b:i64) -> i64 => r##" - def power(a: int, b: int): + def power(a: int, b: int) -> int: return a ** b "## }; @@ -137,7 +137,7 @@ mod tests { #[test] fn basic_lshift() { let lshift = jit_function! { lshift(a:i64, b:i64) -> i64 => r##" - def lshift(a: int, b: int): + def lshift(a: int, b: int) -> int: return a << b "## }; @@ -152,7 +152,7 @@ mod tests { #[test] fn basic_rshift() { let rshift = jit_function! { rshift(a:i64, b:i64) -> i64 => r##" - def rshift(a: int, b: int): + def rshift(a: int, b: int) -> int: return a >> b "## }; @@ -167,7 +167,7 @@ mod tests { #[test] fn basic_and() { let bitand = jit_function! { bitand(a:i64, b:i64) -> i64 => r##" - def bitand(a: int, b: int): + def bitand(a: int, b: int) -> int: return a & b "## }; @@ -182,7 +182,7 @@ mod tests { #[test] fn basic_or() { let bitor = jit_function! { bitor(a:i64, b:i64) -> i64 => r##" - def bitor(a: int, b: int): + def bitor(a: int, b: int) -> int: return a | b "## }; @@ -197,7 +197,7 @@ mod tests { #[test] fn basic_xor() { let bitxor = jit_function! { bitxor(a:i64, b:i64) -> i64 => r##" - def bitxor(a: int, b: int): + def bitxor(a: int, b: int) -> int: return a ^ b "## }; @@ -212,7 +212,7 @@ mod tests { #[test] fn basic_eq() { let eq = jit_function! { eq(a:i64, b:i64) -> i64 => r##" - def eq(a: int, b: int): + def eq(a: int, b: int) -> int: if a == b: return 1 return 0 @@ -227,7 +227,7 @@ mod tests { #[test] fn basic_gt() { let gt = jit_function! { gt(a:i64, b:i64) -> i64 => r##" - def gt(a: int, b: int): + def gt(a: int, b: int) -> int: if a > b: return 1 return 0 @@ -244,7 +244,7 @@ mod tests { #[test] fn basic_lt() { let lt = jit_function! { lt(a:i64, b:i64) -> i64 => r##" - def lt(a: int, b: int): + def lt(a: int, b: int) -> int: if a < b: return 1 return 0 @@ -260,7 +260,7 @@ mod tests { #[test] fn basic_gte() { let gte = jit_function! { gte(a:i64, b:i64) -> i64 => r##" - def gte(a: int, b: int): + def gte(a: int, b: int) -> int: if a >= b: return 1 return 0 @@ -275,7 +275,7 @@ mod tests { #[test] fn basic_lte() { let lte = jit_function! { lte(a:i64, b:i64) -> i64 => r##" - def lte(a: int, b: int): + def lte(a: int, b: int) -> int: if a <= b: return 1 return 0 @@ -290,7 +290,7 @@ mod tests { #[test] fn basic_minus() { let minus = jit_function! { minus(a:i64) -> i64 => r##" - def minus(a: int): + def minus(a: int) -> int: return -a "## }; @@ -304,7 +304,7 @@ mod tests { #[test] fn basic_plus() { let plus = jit_function! { plus(a:i64) -> i64 => r##" - def plus(a: int): + def plus(a: int) -> int: return +a "## }; @@ -318,7 +318,7 @@ mod tests { #[test] fn basic_not() { let not_ = jit_function! { not_(a: i64) -> bool => r##" - def not_(a: int): + def not_(a: int) -> bool: return not a "## }; diff --git a/crates/jit/tests/misc_tests.rs b/crates/jit/tests/misc_tests.rs index b73100ad6ec..6ccf33357d1 100644 --- a/crates/jit/tests/misc_tests.rs +++ b/crates/jit/tests/misc_tests.rs @@ -16,7 +16,7 @@ mod tests { #[test] fn invoke() { let func = jit_function! { func => r##" - def func(a: int, b: float): + def func(a: int, b: float) -> int: return 1 "## }; @@ -41,7 +41,7 @@ mod tests { #[test] fn args_builder() { let func = jit_function! { func=> r##" - def func(a: int, b: float): + def func(a: int, b: float) -> int: return 1 "## }; @@ -71,7 +71,7 @@ mod tests { #[test] fn basic_if_else() { let if_else = jit_function! { if_else(a:i64) -> i64 => r##" - def if_else(a: int): + def if_else(a: int) -> int: if a: return 42 else: @@ -88,9 +88,11 @@ mod tests { } #[test] + #[ignore = "JIT loop lowering is currently broken / WIP"] fn basic_while_loop() { + // TODO: re-add this into the test suite when SSA phi nodes are implemented correctly let while_loop = jit_function! { while_loop(a:i64) -> i64 => r##" - def while_loop(a: int): + def while_loop(a: int) -> int: b = 0 while a > 0: b += 1 @@ -106,7 +108,7 @@ mod tests { #[test] fn basic_unpack_tuple() { let unpack_tuple = jit_function! { unpack_tuple(a:i64, b:i64) -> i64 => r##" - def unpack_tuple(a: int, b: int): + def unpack_tuple(a: int, b: int) -> int: a, b = b, a return a "## }; diff --git a/crates/jit/tests/none_tests.rs b/crates/jit/tests/none_tests.rs index a1d36af80b8..e02bfcbb59d 100644 --- a/crates/jit/tests/none_tests.rs +++ b/crates/jit/tests/none_tests.rs @@ -3,7 +3,7 @@ mod tests { #[test] fn basic_not() { let not_ = jit_function! { not_(x: i64) -> bool => r##" - def not_(x: int): + def not_(x: int) -> bool: return not None "## }; @@ -13,7 +13,7 @@ mod tests { #[test] fn basic_if_not() { let if_not = jit_function! { if_not(x: i64) -> i64 => r##" - def if_not(x: int): + def if_not(x: int) -> int: if not None: return 1 else: diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 4125ef4c4c6..3a690793188 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -1067,6 +1067,7 @@ impl PyFunction { let arg_types = jit::get_jit_arg_types(&zelf, vm)?; let ret_type = jit::jit_ret_type(&zelf, vm)?; let code: &Py = &zelf.code; + let compiled = rustpython_jit::compile(&code.code, &arg_types, ret_type) .map_err(|err| jit::new_jit_error(err.to_string(), vm))?; *jit_guard = Some(compiled); diff --git a/crates/vm/src/builtins/function/jit.rs b/crates/vm/src/builtins/function/jit.rs index 8432bb5369a..645109919ee 100644 --- a/crates/vm/src/builtins/function/jit.rs +++ b/crates/vm/src/builtins/function/jit.rs @@ -31,9 +31,9 @@ pub(super) enum ArgsError { impl ToPyObject for AbiValue { fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { match self { - AbiValue::Int(i) => i.to_pyobject(vm), - AbiValue::Float(f) => f.to_pyobject(vm), - AbiValue::Bool(b) => b.to_pyobject(vm), + Self::Int(i) => i.to_pyobject(vm), + Self::Float(f) => f.to_pyobject(vm), + Self::Bool(b) => b.to_pyobject(vm), _ => unimplemented!(), } } @@ -52,9 +52,11 @@ fn get_jit_arg_type(dict: &Py, name: &str, vm: &VirtualMachine) -> PyRes Ok(JitType::Float) } else if value.is(vm.ctx.types.bool_type) { Ok(JitType::Bool) + } else if value.is(vm.ctx.types.tuple_type) { + Ok(JitType::Object) } else { Err(new_jit_error( - "Jit requires argument to be either int, float or bool".to_owned(), + "Jit requires argument to be either int, float, bool or a tuple".to_owned(), vm, )) } diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 887cbb50e7e..e42ed2b0c44 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -1,29 +1,63 @@ -def foo(): +def foo() -> int: a = 5 return 10 + a -def bar(): +def bar() -> float: a = 1e6 return a / 5.0 -def baz(a: int, b: int): +def baz(a: int, b: float) -> float: return a + b + 12 +def tuple_identity(t: tuple) -> tuple: + return t + + +def mixed_args(a: int, b: float, c: tuple) -> tuple: + return a, b, c + + +def fib(n: int) -> int: + if n == 0 or n == 1: + return 1 + return fib(n - 1) + fib(n - 2) + + def tests(): - assert foo() == 15 - assert bar() == 2e5 - assert baz(17, 20) == 49 - assert baz(17, 22.5) == 51.5 + test_funcs = [foo, bar, baz, baz, tuple_identity, mixed_args, fib] + test_args = [ + [], + [], + [17, 20], + [17, 22.5], + [(1, 2)], + [1, 3.5, (1, 2)], + [5], + ] + test_expected = [ + 15, + 2e5, + 49, + 51.5, + (1, 2), + (1, 3.5, (1, 2)), + 8, + ] + for f, args, expected in zip(test_funcs, test_args, test_expected): + assert f(*args) == expected + print(f"Test {f.__name__}({', '.join(map(str, args))}) == {expected} PASSED") tests() if hasattr(foo, "__jit__"): - print("Has jit") + print("Has jit, JIT test start:") foo.__jit__() bar.__jit__() baz.__jit__() + tuple_identity.__jit__() + fib.__jit__() tests()