Add more jit tests, and utility to make testing it easier.

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",
 ]
 

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"

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<AbiValue> {
-        debug_assert_eq!(self as *const _, args.code as *const _);
-        let cif = self.sig.to_cif();
-        unsafe {
-            let value = cif.call::<UnTypedAbiValue>(
-                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<Option<AbiValue>, 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::<Result<Vec<_>, _>>()?;
+        Ok(unsafe { self.invoke_raw(&cif_args) })
+    }
+
+    unsafe fn invoke_raw(&self, cif_args: &[libffi::middle::Arg]) -> Option<AbiValue> {
+        let cif = self.sig.to_cif();
+        let value = cif.call::<UnTypedAbiValue>(
+            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<i64> for AbiValue {
+    fn from(i: i64) -> Self {
+        AbiValue::Int(i)
+    }
+}
+
+impl From<f64> for AbiValue {
+    fn from(f: f64) -> Self {
+        AbiValue::Float(f)
+    }
+}
+
+impl TryFrom<AbiValue> for i64 {
+    type Error = ();
+
+    fn try_from(value: AbiValue) -> Result<Self, Self::Error> {
+        match value {
+            AbiValue::Int(i) => Ok(i),
+            AbiValue::Float(_) => Err(()),
+        }
+    }
+}
+
+impl TryFrom<AbiValue> for f64 {
+    type Error = ();
+
+    fn try_from(value: AbiValue) -> Result<Self, Self::Error> {
+        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<Args<'a>> {
         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::<Option<_>>()
             .map(|cif_args| Args {
                 _values: self.values,
@@ -258,3 +315,9 @@ pub struct Args<'a> {
     cif_args: Vec<libffi::middle::Arg>,
     code: &'a CompiledCode,
 }
+
+impl<'a> Args<'a> {
+    pub fn invoke(&self) -> Option<AbiValue> {
+        unsafe { self.code.invoke_raw(&self.cif_args) }
+    }
+}

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<CodeObject>,
+    name: String,
+    annotations: HashMap<String, StackValue>,
+}
+
+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<String, StackValue>),
+    Code(Box<CodeObject>),
+    Function(Function),
+}
+
+impl From<Constant> 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<StackValue>,
+    locals: HashMap<String, StackValue>,
+}
+
+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")
+                    })
+            }
+        }
+    };
+}

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));
+}

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));
+}

jit/tests/lib.rs

@@ -0,0 +1,4 @@
+#[macro_use]
+mod common;
+mod float_tests;
+mod int_tests;

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 \