diff --git a/Cargo.lock b/Cargo.lock index 374231ad8..700a0dd1d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -48,6 +48,15 @@ version = "1.0.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b602bfe940d21c130f3895acd65221e8a61270debe89d628b9cb4e3ccb8569b" +[[package]] +name = "approx" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0e60b75072ecd4168020818c0107f2857bb6c4e64252d8d3983f6263b40a5c3" +dependencies = [ + "num-traits", +] + [[package]] name = "arr_macro" version = "0.1.3" @@ -1656,12 +1665,14 @@ dependencies = [ name = "rustpython-jit" version = "0.1.2" dependencies = [ + "approx", "cranelift", "cranelift-module", "cranelift-simplejit", "libffi", "num-traits", "rustpython-bytecode", + "rustpython-derive", "thiserror", ] diff --git a/jit/Cargo.toml b/jit/Cargo.toml index 3f33538fe..7b32035be 100644 --- a/jit/Cargo.toml +++ b/jit/Cargo.toml @@ -7,11 +7,21 @@ repository = "https://github.com/RustPython/RustPython" license = "MIT" edition = "2018" +autotests = false + [dependencies] cranelift = "0.66.0" cranelift-module = "0.66.0" cranelift-simplejit = "0.66.0" num-traits = "0.2" libffi = "0.9.0" -rustpython-bytecode = { path = "../bytecode", version = "0.1.1" } +rustpython-bytecode = { path = "../bytecode", version = "0.1.2" } thiserror = "1.0" + +[dev-dependencies] +approx = "0.3.2" +rustpython-derive = { path = "../derive", version = "0.1.2" } + +[[test]] +name = "integration" +path = "tests/lib.rs" diff --git a/jit/src/lib.rs b/jit/src/lib.rs index a484e15fa..32667bbfe 100644 --- a/jit/src/lib.rs +++ b/jit/src/lib.rs @@ -9,6 +9,7 @@ use rustpython_bytecode::bytecode; mod instructions; use instructions::FunctionCompiler; +use std::convert::TryFrom; #[derive(Debug, thiserror::Error)] pub enum JitCompileError { @@ -20,10 +21,12 @@ pub enum JitCompileError { CraneliftError(#[from] ModuleError), } -#[derive(Debug, thiserror::Error)] +#[derive(Debug, thiserror::Error, Eq, PartialEq)] pub enum JitArgumentError { #[error("argument is of wrong type")] ArgumentTypeMismatch, + #[error("wrong number of arguments")] + WrongNumberOfArguments, } struct Jit { @@ -120,16 +123,29 @@ impl CompiledCode { ArgsBuilder::new(self) } - pub fn invoke<'a>(&self, args: &Args<'a>) -> Option { - debug_assert_eq!(self as *const _, args.code as *const _); - let cif = self.sig.to_cif(); - unsafe { - let value = cif.call::( - libffi::middle::CodePtr::from_ptr(self.code as *const _), - &args.cif_args, - ); - self.sig.ret.as_ref().map(|ty| value.to_typed(ty)) + pub fn invoke(&self, args: &[AbiValue]) -> Result, JitArgumentError> { + if self.sig.args.len() != args.len() { + return Err(JitArgumentError::WrongNumberOfArguments); } + + let cif_args = self + .sig + .args + .iter() + .zip(args.iter()) + .map(|(ty, val)| type_check(ty, val).map(|_| val)) + .map(|v| v.map(AbiValue::to_libffi_arg)) + .collect::, _>>()?; + Ok(unsafe { self.invoke_raw(&cif_args) }) + } + + unsafe fn invoke_raw(&self, cif_args: &[libffi::middle::Arg]) -> Option { + let cif = self.sig.to_cif(); + let value = cif.call::( + libffi::middle::CodePtr::from_ptr(self.code as *const _), + cif_args, + ); + self.sig.ret.as_ref().map(|ty| value.to_typed(ty)) } } @@ -170,12 +186,62 @@ impl JitType { } } -#[derive(Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum AbiValue { Float(f64), Int(i64), } +impl AbiValue { + fn to_libffi_arg(&self) -> libffi::middle::Arg { + match self { + AbiValue::Int(ref i) => libffi::middle::Arg::new(i), + AbiValue::Float(ref f) => libffi::middle::Arg::new(f), + } + } +} + +impl From for AbiValue { + fn from(i: i64) -> Self { + AbiValue::Int(i) + } +} + +impl From for AbiValue { + fn from(f: f64) -> Self { + AbiValue::Float(f) + } +} + +impl TryFrom for i64 { + type Error = (); + + fn try_from(value: AbiValue) -> Result { + match value { + AbiValue::Int(i) => Ok(i), + AbiValue::Float(_) => Err(()), + } + } +} + +impl TryFrom for f64 { + type Error = (); + + fn try_from(value: AbiValue) -> Result { + match value { + AbiValue::Int(_) => Err(()), + AbiValue::Float(f) => Ok(f), + } + } +} + +fn type_check(ty: &JitType, val: &AbiValue) -> Result<(), JitArgumentError> { + match (ty, val) { + (JitType::Int, AbiValue::Int(_)) | (JitType::Float, AbiValue::Float(_)) => Ok(()), + _ => Err(JitArgumentError::ArgumentTypeMismatch), + } +} + #[derive(Copy, Clone)] union UnTypedAbiValue { float: f64, @@ -222,13 +288,9 @@ impl<'a> ArgsBuilder<'a> { } pub fn set(&mut self, idx: usize, value: AbiValue) -> Result<(), JitArgumentError> { - match (&self.code.sig.args[idx], &value) { - (JitType::Int, AbiValue::Int(_)) | (JitType::Float, AbiValue::Float(_)) => { - self.values[idx] = Some(value); - Ok(()) - } - _ => Err(JitArgumentError::ArgumentTypeMismatch), - } + type_check(&self.code.sig.args[idx], &value).map(|_| { + self.values[idx] = Some(value); + }) } pub fn is_set(&self, idx: usize) -> bool { @@ -238,12 +300,7 @@ impl<'a> ArgsBuilder<'a> { pub fn into_args(self) -> Option> { self.values .iter() - .map(|v| { - v.as_ref().map(|v| match v { - AbiValue::Int(ref i) => libffi::middle::Arg::new(i), - AbiValue::Float(ref f) => libffi::middle::Arg::new(f), - }) - }) + .map(|v| v.as_ref().map(AbiValue::to_libffi_arg)) .collect::>() .map(|cif_args| Args { _values: self.values, @@ -258,3 +315,9 @@ pub struct Args<'a> { cif_args: Vec, code: &'a CompiledCode, } + +impl<'a> Args<'a> { + pub fn invoke(&self) -> Option { + unsafe { self.code.invoke_raw(&self.cif_args) } + } +} diff --git a/jit/tests/common.rs b/jit/tests/common.rs new file mode 100644 index 000000000..cd3d6bc30 --- /dev/null +++ b/jit/tests/common.rs @@ -0,0 +1,184 @@ +use std::collections::HashMap; + +use rustpython_bytecode::bytecode::{CodeObject, Constant, Instruction, NameScope}; +use rustpython_jit::{CompiledCode, JitType}; + +#[derive(Debug, Clone)] +pub struct Function { + code: Box, + name: String, + annotations: HashMap, +} + +impl Function { + pub fn compile(self) -> CompiledCode { + let mut arg_types = Vec::new(); + for arg in self.code.arg_names.iter() { + let arg_type = match self.annotations.get(arg) { + Some(StackValue::String(annotation)) => match annotation.as_str() { + "int" => JitType::Int, + "float" => JitType::Float, + _ => panic!("Unrecognised jit type"), + }, + _ => panic!("Argument have annotation"), + }; + arg_types.push(arg_type); + } + + rustpython_jit::compile(&self.code, &arg_types).expect("Compile failure") + } +} + +#[derive(Debug, Clone)] +enum StackValue { + String(String), + None, + Map(HashMap), + Code(Box), + Function(Function), +} + +impl From for StackValue { + fn from(value: Constant) -> Self { + match value { + Constant::String { value } => StackValue::String(value), + Constant::None => StackValue::None, + Constant::Code { code } => StackValue::Code(code), + c => unimplemented!("constant {:?} isn't yet supported in py_function!", c), + } + } +} + +pub struct StackMachine { + stack: Vec, + locals: HashMap, +} + +impl StackMachine { + pub fn new() -> StackMachine { + StackMachine { + stack: Vec::new(), + locals: HashMap::new(), + } + } + + pub fn run(&mut self, code: CodeObject) { + for instruction in code.instructions { + if self.process_instruction(instruction) { + break; + } + } + } + + fn process_instruction(&mut self, instruction: Instruction) -> bool { + match instruction { + Instruction::LoadConst { value } => self.stack.push(value.into()), + Instruction::LoadName { + name, + scope: NameScope::Free, + } => self.stack.push(StackValue::String(name)), + Instruction::StoreName { name, .. } => { + self.locals.insert(name, self.stack.pop().unwrap()); + } + Instruction::StoreAttr { .. } => { + // Do nothing except throw away the stack values + self.stack.pop().unwrap(); + self.stack.pop().unwrap(); + } + Instruction::BuildMap { size, .. } => { + let mut map = HashMap::new(); + for _ in 0..size { + let value = self.stack.pop().unwrap(); + let name = if let Some(StackValue::String(name)) = self.stack.pop() { + name + } else { + unimplemented!("no string keys isn't yet supported in py_function!") + }; + map.insert(name, value); + } + self.stack.push(StackValue::Map(map)); + } + Instruction::MakeFunction => { + let name = if let Some(StackValue::String(name)) = self.stack.pop() { + name + } else { + panic!("Expected function name") + }; + let code = if let Some(StackValue::Code(code)) = self.stack.pop() { + code + } else { + panic!("Expected function code") + }; + let annotations = if let Some(StackValue::Map(map)) = self.stack.pop() { + map + } else { + panic!("Expected function annotations") + }; + self.stack.push(StackValue::Function(Function { + name, + code, + annotations, + })); + } + Instruction::Duplicate => { + let value = self.stack.last().unwrap().clone(); + self.stack.push(value); + } + Instruction::Rotate { amount } => { + let mut values = Vec::new(); + + // Pop all values from stack: + for _ in 0..amount { + values.push(self.stack.pop().unwrap()); + } + + // Push top of stack back first: + self.stack.push(values.remove(0)); + + // Push other value back in order: + for value in values.into_iter().rev() { + self.stack.push(value); + } + } + Instruction::ReturnValue => return true, + _ => unimplemented!( + "instruction {:?} isn't yet supported in py_function!", + instruction + ), + } + return false; + } + + pub fn get_function(&self, name: &str) -> Function { + if let Some(StackValue::Function(function)) = self.locals.get(name) { + function.clone() + } else { + panic!("There was no function named {}", name) + } + } +} + +macro_rules! jit_function { + ($func_name:ident($($arg_name:ident:$arg_type:ty),*) -> $ret_type:ty => $($t:tt)*) => { + { + use std::convert::TryInto; + + let code = rustpython_derive::py_compile!( + crate_name = "rustpython_bytecode", + source = $($t)* + ); + let mut machine = $crate::common::StackMachine::new(); + machine.run(code); + let jit_code = machine.get_function(stringify!($func_name)).compile(); + + move |$($arg_name:$arg_type),*| -> Result<$ret_type, rustpython_jit::JitArgumentError> { + jit_code + .invoke(&[$($arg_name.into()),*]) + .map(|ret| match ret { + Some(ret) => ret.try_into().expect("jit function returned unexpected type"), + None => panic!("jit function unexpectedly returned None") + }) + } + } + }; +} diff --git a/jit/tests/float_tests.rs b/jit/tests/float_tests.rs new file mode 100644 index 000000000..7140fd6dc --- /dev/null +++ b/jit/tests/float_tests.rs @@ -0,0 +1,92 @@ +macro_rules! assert_approx_eq { + ($left:expr, $right:expr) => { + match ($left, $right) { + (Ok(lhs), Ok(rhs)) => approx::assert_relative_eq!(lhs, rhs), + (lhs, rhs) => assert_eq!(lhs, rhs), + } + }; +} + +macro_rules! assert_bits_eq { + ($left:expr, $right:expr) => { + match ($left, $right) { + (Ok(lhs), Ok(rhs)) => assert!(lhs.to_bits() == rhs.to_bits()), + (lhs, rhs) => assert_eq!(lhs, rhs), + } + }; +} + +#[test] +fn test_add() { + let add = jit_function! { add(a:f64, b:f64) -> f64 => r##" + def add(a: float, b: float): + return a + b + "## }; + + assert_approx_eq!(add(5.5, 10.2), Ok(15.7)); + assert_approx_eq!(add(-4.5, 7.6), Ok(3.1)); + assert_approx_eq!(add(-5.2, -3.9), Ok(-9.1)); + assert_bits_eq!(add(-5.2, f64::NAN), Ok(f64::NAN)); + assert_eq!(add(2.0, f64::INFINITY), Ok(f64::INFINITY)); + assert_eq!(add(-2.0, f64::NEG_INFINITY), Ok(f64::NEG_INFINITY)); + assert_eq!(add(1.0, f64::NEG_INFINITY), Ok(f64::NEG_INFINITY)); +} + +#[test] +fn test_sub() { + let sub = jit_function! { sub(a:f64, b:f64) -> f64 => r##" + def sub(a: float, b: float): + return a - b + "## }; + + assert_approx_eq!(sub(5.2, 3.6), Ok(1.6)); + assert_approx_eq!(sub(3.4, 4.2), Ok(-0.8)); + assert_approx_eq!(sub(-2.1, 1.3), Ok(-3.4)); + assert_approx_eq!(sub(3.1, -1.3), Ok(4.4)); + assert_bits_eq!(sub(-5.2, f64::NAN), Ok(f64::NAN)); + assert_eq!(sub(f64::INFINITY, 2.0), Ok(f64::INFINITY)); + assert_eq!(sub(-2.0, f64::NEG_INFINITY), Ok(f64::INFINITY)); + assert_eq!(sub(1.0, f64::INFINITY), Ok(f64::NEG_INFINITY)); +} + +#[test] +fn test_mul() { + let mul = jit_function! { mul(a:f64, b:f64) -> f64 => r##" + def mul(a: float, b: float): + return a * b + "## }; + + assert_approx_eq!(mul(5.2, 2.0), Ok(10.4)); + assert_approx_eq!(mul(3.4, -1.7), Ok(-5.779999999999999)); + assert_bits_eq!(mul(1.0, 0.0), Ok(0.0f64)); + assert_bits_eq!(mul(1.0, -0.0), Ok(-0.0f64)); + assert_bits_eq!(mul(-1.0, 0.0), Ok(-0.0f64)); + assert_bits_eq!(mul(-1.0, -0.0), Ok(0.0f64)); + assert_bits_eq!(mul(-5.2, f64::NAN), Ok(f64::NAN)); + assert_eq!(mul(1.0, f64::INFINITY), Ok(f64::INFINITY)); + assert_eq!(mul(1.0, f64::NEG_INFINITY), Ok(f64::NEG_INFINITY)); + assert_eq!(mul(-1.0, f64::INFINITY), Ok(f64::NEG_INFINITY)); + assert!(mul(0.0, f64::INFINITY).unwrap().is_nan()); + assert_eq!(mul(f64::NEG_INFINITY, f64::INFINITY), Ok(f64::NEG_INFINITY)); +} + +#[test] +fn test_div() { + let div = jit_function! { div(a:f64, b:f64) -> f64 => r##" + def div(a: float, b: float): + return a / b + "## }; + + assert_approx_eq!(div(5.2, 2.0), Ok(2.6)); + assert_approx_eq!(div(3.4, -1.7), Ok(-2.0)); + assert_eq!(div(1.0, 0.0), Ok(f64::INFINITY)); + assert_eq!(div(1.0, -0.0), Ok(f64::NEG_INFINITY)); + assert_eq!(div(-1.0, 0.0), Ok(f64::NEG_INFINITY)); + assert_eq!(div(-1.0, -0.0), Ok(f64::INFINITY)); + assert_bits_eq!(div(-5.2, f64::NAN), Ok(f64::NAN)); + assert_eq!(div(f64::INFINITY, 2.0), Ok(f64::INFINITY)); + assert_bits_eq!(div(-2.0, f64::NEG_INFINITY), Ok(0.0f64)); + assert_bits_eq!(div(1.0, f64::INFINITY), Ok(0.0f64)); + assert_bits_eq!(div(2.0, f64::NEG_INFINITY), Ok(-0.0f64)); + assert_bits_eq!(div(-1.0, f64::INFINITY), Ok(-0.0f64)); +} diff --git a/jit/tests/int_tests.rs b/jit/tests/int_tests.rs new file mode 100644 index 000000000..191229b57 --- /dev/null +++ b/jit/tests/int_tests.rs @@ -0,0 +1,24 @@ +#[test] +fn test_add() { + let add = jit_function! { add(a:i64, b:i64) -> i64 => r##" + def add(a: int, b: int): + return a + b + "## }; + + assert_eq!(add(5, 10), Ok(15)); + assert_eq!(add(-5, 12), Ok(7)); + assert_eq!(add(-5, -3), Ok(-8)); +} + +#[test] +fn test_sub() { + let sub = jit_function! { sub(a:i64, b:i64) -> i64 => r##" + def sub(a: int, b: int): + return a - b + "## }; + + assert_eq!(sub(5, 10), Ok(-5)); + assert_eq!(sub(12, 10), Ok(2)); + assert_eq!(sub(7, 10), Ok(-3)); + assert_eq!(sub(-3, -10), Ok(7)); +} diff --git a/jit/tests/lib.rs b/jit/tests/lib.rs new file mode 100644 index 000000000..c1805deef --- /dev/null +++ b/jit/tests/lib.rs @@ -0,0 +1,4 @@ +#[macro_use] +mod common; +mod float_tests; +mod int_tests; diff --git a/vm/src/obj/objfunction.rs b/vm/src/obj/objfunction.rs index 308466534..4e42f2e5b 100644 --- a/vm/src/obj/objfunction.rs +++ b/vm/src/obj/objfunction.rs @@ -240,7 +240,7 @@ impl PyFunction { if let Some(jitted_code) = self.jitted_code.get() { match jitfunc::get_jit_args(self, &func_args, jitted_code, vm) { Ok(args) => { - return Ok(jitted_code.invoke(&args).into_pyobject(vm)); + return Ok(args.invoke().into_pyobject(vm)); } Err(err) => info!( "jit: function `{}` is falling back to being interpreted because of the \