Merge pull request RustPython#2164 from skinny121/jit-args

Support arguments in jitted functions

Author: coolreader18

Cargo.lock

@@ -77,6 +77,10 @@ def test_inheritance(self):
                 last_depth = depth
         finally:
             inheritance_tree.close()
+
+        # RUSTPYTHON specific
+        exc_set.discard("JitError")
+
         self.assertEqual(len(exc_set), 0, "%s not accounted for" % exc_set)
 
     interface_tests = ("length", "args", "str", "repr")

Lib/test/test_baseexception.py

@@ -3,44 +3,7 @@ use num_traits::cast::ToPrimitive;
 use rustpython_bytecode::bytecode::{BinaryOperator, Constant, Instruction, NameScope};
 use std::collections::HashMap;
 
-use super::JitCompileError;
-
-#[derive(Default)]
-pub struct JitSig {
-    pub ret: Option<JitType>,
-}
-
-impl JitSig {
-    pub fn to_cif(&self) -> libffi::middle::Cif {
-        let ret = match self.ret {
-            Some(ref ty) => ty.to_libffi(),
-            None => libffi::middle::Type::void(),
-        };
-        libffi::middle::Cif::new(Vec::new(), ret)
-    }
-}
-
-#[derive(Clone, PartialEq)]
-pub enum JitType {
-    Int,
-    Float,
-}
-
-impl JitType {
-    fn to_cranelift(&self) -> types::Type {
-        match self {
-            Self::Int => types::I64,
-            Self::Float => types::F64,
-        }
-    }
-
-    fn to_libffi(&self) -> libffi::middle::Type {
-        match self {
-            Self::Int => libffi::middle::Type::i64(),
-            Self::Float => libffi::middle::Type::f64(),
-        }
-    }
-}
+use super::{JitCompileError, JitSig, JitType};
 
 #[derive(Clone)]
 struct Local {
@@ -53,20 +16,63 @@ struct JitValue {
     ty: JitType,
 }
 
+impl JitValue {
+    fn new(val: Value, ty: JitType) -> JitValue {
+        JitValue { val, ty }
+    }
+}
+
 pub struct FunctionCompiler<'a, 'b> {
     builder: &'a mut FunctionBuilder<'b>,
     stack: Vec<JitValue>,
     variables: HashMap<String, Local>,
-    pub sig: JitSig,
+    pub(crate) sig: JitSig,
 }
 
 impl<'a, 'b> FunctionCompiler<'a, 'b> {
-    pub fn new(builder: &'a mut FunctionBuilder<'b>) -> FunctionCompiler<'a, 'b> {
-        FunctionCompiler {
+    pub fn new(
+        builder: &'a mut FunctionBuilder<'b>,
+        arg_names: &[String],
+        arg_types: &[JitType],
+        entry_block: Block,
+    ) -> FunctionCompiler<'a, 'b> {
+        let mut compiler = FunctionCompiler {
             builder,
             stack: Vec::new(),
             variables: HashMap::new(),
-            sig: JitSig::default(),
+            sig: JitSig {
+                args: arg_types.to_vec(),
+                ret: None,
+            },
+        };
+        let params = compiler.builder.func.dfg.block_params(entry_block).to_vec();
+        debug_assert_eq!(arg_names.len(), arg_types.len());
+        debug_assert_eq!(arg_names.len(), params.len());
+        for ((name, ty), val) in arg_names.iter().zip(arg_types).zip(params) {
+            compiler
+                .store_variable(name.clone(), JitValue::new(val, ty.clone()))
+                .unwrap();
+        }
+        compiler
+    }
+
+    fn store_variable(&mut self, name: String, val: JitValue) -> Result<(), JitCompileError> {
+        let len = self.variables.len();
+        let builder = &mut self.builder;
+        let local = self.variables.entry(name).or_insert_with(|| {
+            let var = Variable::new(len);
+            let local = Local {
+                var,
+                ty: val.ty.clone(),
+            };
+            builder.declare_var(var, val.ty.to_cranelift());
+            local
+        });
+        if val.ty != local.ty {
+            Err(JitCompileError::NotSupported)
+        } else {
+            self.builder.def_var(local.var, val.val);
+            Ok(())
         }
     }
 
@@ -91,22 +97,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 scope: NameScope::Local,
             } => {
                 let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
-                let len = self.variables.len();
-                let builder = &mut self.builder;
-                let local = self.variables.entry(name.clone()).or_insert_with(|| {
-                    let var = Variable::new(len);
-                    let local = Local {
-                        var,
-                        ty: val.ty.clone(),
-                    };
-                    builder.declare_var(var, val.ty.to_cranelift());
-                    local
-                });
-                if val.ty != local.ty {
-                    return Err(JitCompileError::NotSupported);
-                }
-                self.builder.def_var(local.var, val.val);
-                Ok(())
+                self.store_variable(name.clone(), val)
             }
             Instruction::LoadConst {
                 value: Constant::Integer { value },

jit/src/instructions.rs

@@ -8,7 +8,7 @@ use rustpython_bytecode::bytecode;
 
 mod instructions;
 
-use instructions::{FunctionCompiler, JitSig, JitType};
+use instructions::FunctionCompiler;
 
 #[derive(Debug, thiserror::Error)]
 pub enum JitCompileError {
@@ -20,6 +20,12 @@ pub enum JitCompileError {
     CraneliftError(#[from] ModuleError),
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum JitArgumentError {
+    #[error("argument is of wrong type")]
+    ArgumentTypeMismatch,
+}
+
 struct Jit {
     builder_context: FunctionBuilderContext,
     ctx: codegen::Context,
@@ -40,15 +46,26 @@ impl Jit {
     fn build_function(
         &mut self,
         bytecode: &bytecode::CodeObject,
+        args: &[JitType],
     ) -> Result<(FuncId, JitSig), JitCompileError> {
+        for arg in args {
+            self.ctx
+                .func
+                .signature
+                .params
+                .push(AbiParam::new(arg.to_cranelift()));
+        }
+
         let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context);
         let entry_block = builder.create_block();
-        // builder.append_block_params_for_function_params(entry_block);
+        builder.append_block_params_for_function_params(entry_block);
         builder.switch_to_block(entry_block);
         builder.seal_block(entry_block);
 
         let sig = {
-            let mut compiler = FunctionCompiler::new(&mut builder);
+            let mut arg_names = bytecode.arg_names.clone();
+            arg_names.extend(bytecode.kwonlyarg_names.iter().cloned());
+            let mut compiler = FunctionCompiler::new(&mut builder, &arg_names, args, entry_block);
 
             for instruction in &bytecode.instructions {
                 compiler.add_instruction(instruction)?;
@@ -74,10 +91,13 @@ impl Jit {
     }
 }
 
-pub fn compile(bytecode: &bytecode::CodeObject) -> Result<CompiledCode, JitCompileError> {
+pub fn compile(
+    bytecode: &bytecode::CodeObject,
+    args: &[JitType],
+) -> Result<CompiledCode, JitCompileError> {
     let mut jit = Jit::new();
 
-    let (id, sig) = jit.build_function(bytecode)?;
+    let (id, sig) = jit.build_function(bytecode, args)?;
 
     jit.module.finalize_definitions();
 
@@ -96,23 +116,67 @@ pub struct CompiledCode {
 }
 
 impl CompiledCode {
-    pub fn invoke(&self) -> Option<AbiValue> {
+    pub fn args_builder(&self) -> ArgsBuilder<'_> {
+        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))
         }
     }
 }
 
+struct JitSig {
+    args: Vec<JitType>,
+    ret: Option<JitType>,
+}
+
+impl JitSig {
+    fn to_cif(&self) -> libffi::middle::Cif {
+        let ret = match self.ret {
+            Some(ref ty) => ty.to_libffi(),
+            None => libffi::middle::Type::void(),
+        };
+        libffi::middle::Cif::new(self.args.iter().map(JitType::to_libffi), ret)
+    }
+}
+
+#[derive(Clone, PartialEq)]
+pub enum JitType {
+    Int,
+    Float,
+}
+
+impl JitType {
+    fn to_cranelift(&self) -> types::Type {
+        match self {
+            Self::Int => types::I64,
+            Self::Float => types::F64,
+        }
+    }
+
+    fn to_libffi(&self) -> libffi::middle::Type {
+        match self {
+            Self::Int => libffi::middle::Type::i64(),
+            Self::Float => libffi::middle::Type::f64(),
+        }
+    }
+}
+
+#[derive(Clone)]
 pub enum AbiValue {
     Float(f64),
     Int(i64),
 }
 
+#[derive(Copy, Clone)]
 union UnTypedAbiValue {
     float: f64,
     int: i64,
@@ -143,3 +207,54 @@ impl fmt::Debug for CompiledCode {
         f.write_str("[compiled code]")
     }
 }
+
+pub struct ArgsBuilder<'a> {
+    values: Vec<Option<AbiValue>>,
+    code: &'a CompiledCode,
+}
+
+impl<'a> ArgsBuilder<'a> {
+    fn new(code: &'a CompiledCode) -> ArgsBuilder<'a> {
+        ArgsBuilder {
+            values: vec![None; code.sig.args.len()],
+            code,
+        }
+    }
+
+    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),
+        }
+    }
+
+    pub fn is_set(&self, idx: usize) -> bool {
+        self.values[idx].is_some()
+    }
+
+    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),
+                })
+            })
+            .collect::<Option<_>>()
+            .map(|cif_args| Args {
+                _values: self.values,
+                cif_args,
+                code: self.code,
+            })
+    }
+}
+
+pub struct Args<'a> {
+    _values: Vec<Option<AbiValue>>,
+    cif_args: Vec<libffi::middle::Arg>,
+    code: &'a CompiledCode,
+}

jit/src/lib.rs

@@ -9,9 +9,15 @@ def bar():
     return a / 5.0
 
 
+def baz(a: int, b: int):
+    return a + b + 12
+
+
 def tests():
     assert foo() == 15
     assert bar() == 2e5
+    assert baz(17, 20) == 49
+    assert baz(17, 22.5) == 51.5
 
 
 tests()
@@ -20,4 +26,5 @@ def tests():
     print("Has jit")
     foo.__jit__()
     bar.__jit__()
+    baz.__jit__()
     tests()

tests/snippets/jit.py

@@ -77,6 +77,7 @@ parking_lot = "0.11"
 thread_local = "1.0"
 cfg-if = "0.1.10"
 timsort = "0.1"
+thiserror = "1.0"
 
 ## unicode stuff
 unicode_names2 = "0.4"

vm/Cargo.toml

@@ -977,4 +977,9 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
         "BytesWarning" => ctx.exceptions.bytes_warning.clone(),
         "ResourceWarning" => ctx.exceptions.resource_warning.clone(),
     });
+
+    #[cfg(feature = "jit")]
+    extend_module!(vm, module, {
+        "JitError" => ctx.exceptions.jit_error.clone(),
+    });
 }