Support arguments in jitted functions by BenLewis-Seequent · Pull Request #2164 · RustPython/RustPython

BenLewis-Seequent · 2020-08-29T04:08:24Z

To jit functions that have arguments, the type of the arguments needs to be known when it's compiled. this PR uses the annotations to determine the type of the arguments.

coolreader18

Nice work!

coolreader18 · 2020-08-29T06:20:51Z

+                    self.values,
+                )
+            };
+            let cif_args = values.iter().map(|v| libffi::middle::Arg::new(v)).collect();


If you're going to recollect the vec anyway, you may as well just v.assume_init() in the mapping function and avoid the transmute.

I can't use assume_init as I need the reference into the values vector, cif_args is effectively a vector of pointers that are passed to libffi. This would be nicer if we had either rust-lang/rust#63291 or rust-lang/rust#63568.

coolreader18 · 2020-08-29T06:22:48Z

+#[cfg(feature = "jit")]
+use super::objfloat;
+#[cfg(feature = "jit")]
+use super::objint;


nit; you could condense this into one #[cfg jit] use super::{objint, objfloat}

coolreader18 · 2020-08-29T06:27:35Z

            .get_or_try_init(|| {
-                rustpython_jit::compile(&self.code.code)
+                let arg_types = PyFunction::get_jit_arg_types(&zelf, vm)?;
+                rustpython_jit::compile(&zelf.code.code, &arg_types)


This is a lot of code in objfunction, maybe there could be a submodule obj/objfunction/jitfunc.rs or something to store all the logic for py<->jit conversion?

youknowone · 2020-08-29T09:35:05Z

+                ))
+            }
+        } else {
+            Err(vm.new_runtime_error(format!("argument {} needs annotation", name)))


maybe we need new_jit_error

…itfunc.

coolreader18 · 2020-08-30T00:02:24Z

+    }
+}
+
+#[derive(Copy, Clone, PartialEq)]


The reason I didn't make this Copy is because I was thinking of eventually having composite types, like List(Box<JitType>) or Tuple(Vec<JitType>) or maybe some sort of product or sum type.

coolreader18 · 2020-08-30T00:07:17Z

+impl<'a> ArgsBuilder<'a> {
+    fn new(code: &'a CompiledCode) -> ArgsBuilder<'a> {
+        ArgsBuilder {
+            values: vec![MaybeUninit::uninit(); code.sig.args.len()],


Couldn't this be Vec<Option<UntypedAbiArg>>? and you could like do self.values.iter().map(|val| val.as_ref().map(|val| middle::Abi::new(val)).collect::<Option<_>>()

And you could just store the Vec<Option<>> in the Args struct, since it just matters that it's not dropped

ArniDagur · 2020-08-30T17:44:02Z

An example like the following can cause UB, where the function is given incorrect input:

def add_two(a: int, b: int) -> int:
    return a + b

add_two.__jit__()
result = add_two(2, 3.0)  # cpython would return 5.0 as a float

Here, I give a function that should take an int a float instead. CPython doesn't complain because it doesn't actually look at the type annotations. This branch understandably gives a nonsense result:

[arni][~/src/RustPython][jit-args]% target/release/rustpython
Welcome to the magnificent Rust Python 0.1.2 interpreter 😱 🖖
>>>>> def add_two(a: int, b: int) -> int:
..... 	return a + b
.....
>>>>> add_two.__jit__()
>>>>> add_two(2, 3.0)
4613937818241073154
>>>>> add_two(2, 3)
5
>>>>> def return_float() -> float:
..... 	return 3.0
.....
>>>>> add_two(2, return_float())
4613937818241073154

I was wondering if we're okay with this sort of undefined behaviour when incorrect input is given to jitted functions?

If not, we could either

Do a runtime check of some sort; or
In the native code we output, insert calls to the same VM methods that would have been executed if it was bytecode. These can handle data of unknown type. It would still improve performance because we wouldn't have to match on the instruction every time, right? You'd only go through the match statements in https://github.com/RustPython/RustPython/blob/master/vm/src/frame.rs#L319 once. When working with data whose type we can determine statically, such as the output of in-built functions such as range(...) we could then output more specialized native code as we do in this PR. Perhaps this requires a change in the bytecode format to include type information when it's known?

I'm not sure if 2. makes sense since I don't quite know how RustPython or interpreters in general work.

darleybarreto · 2020-08-30T18:46:54Z

I was wondering if we're okay with this sort of undefined behaviour when incorrect input is given to jitted functions?

If not, we could either

Do a runtime check of some sort; or

In the native code we output, insert calls to the same VM methods that would have been executed if it was bytecode. These can handle data of unknown type. It would still improve performance because we wouldn't have to match on the instruction every time, right? You'd only go through the match statements in https://github.com/RustPython/RustPython/blob/master/vm/src/frame.rs#L319 once. When working with data whose type we can determine statically, such as the output of in-built functions such as range(...) we could then output more specialized native code as we do in this PR. Perhaps this requires a change in the bytecode format to include type information when it's known?

Another option is to use some sort of multiple dispatch, as Julia does. Which also implies introducing some sort of type inference system and other complex things... Perhaps it would be easier to just throw an error for now, and eventually introduce more complicated machinery.

coolreader18 · 2020-08-30T21:35:01Z

I think ideally we're trying to make a jitted function transparent -- it should be no different in behavior from a non-jitted function, and it shouldn't error because of a jit limitation (e.g. i64 ints instead of BigInts). So yeah, I think we should definitely type-check and then fall back to non-jit behavior to either error or get whatever dynamic behavior would happen with that.

coolreader18 · 2020-08-30T21:37:49Z

Doing something similar to Julia's multiple dispatch definitely sounds cool though -- with

def add(a, b):
    return a + b

We could have like a hashmap of Vec<JitType> -> CompiledCode, and when it's called with e.g. (int, int), we try compiling with a as an int and b as an int, and if that succeeds we populate the hashmap with that signature. The same thing could happen for (float, float), or (float, int), or (int, float), etc

BenLewis-Seequent · 2020-08-31T07:29:32Z

@ArniDagur Good catch, this PR did previously do a runtime check of the arguments types, and if it didn't match it would fallback to interpreting the bytecode instead, but a subsequent refactoring broke that.

@darleybarreto @coolreader18 I'm not familiar with Julia at all, but does sound like an interesting approach.

youknowone · 2020-08-31T07:48:39Z

By looking PyPy practice, it make guards with typecheck and falling back to original code when the jit assumption failed.

skinnyBat added 2 commits August 29, 2020 12:53

Allow passing arguments to jitted functions.

0651175

Fix bugs with passing arguments to jitted functions.

751b549

BenLewis-Seequent requested a review from coolreader18 August 29, 2020 04:08

Fix clippy warnings.

3151d89

coolreader18 approved these changes Aug 29, 2020

View reviewed changes

youknowone reviewed Aug 29, 2020

View reviewed changes

skinnyBat added 3 commits August 30, 2020 08:19

Move jit specific functions out of objfunction.rs into a sub module j…

588facf

…itfunc.

Add JitError.

a83c6c5

Fix test_inheritance test.

b83cd27

coolreader18 reviewed Aug 30, 2020

View reviewed changes

Remove the use of MaybeUninit from jit args handling.

bd3c4fc

BenLewis-Seequent force-pushed the jit-args branch from 9a15a26 to bd3c4fc Compare August 30, 2020 08:12

Fix not type checking arguments to jit function.

b79d7ed

BenLewis-Seequent requested a review from coolreader18 August 31, 2020 14:59

coolreader18 approved these changes Sep 1, 2020

View reviewed changes

coolreader18 merged commit 3b2a16f into RustPython:master Sep 1, 2020

Conversation

BenLewis-Seequent commented Aug 29, 2020

Uh oh!

coolreader18 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArniDagur commented Aug 30, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

darleybarreto commented Aug 30, 2020

Uh oh!

coolreader18 commented Aug 30, 2020

Uh oh!

coolreader18 commented Aug 30, 2020

Uh oh!

BenLewis-Seequent commented Aug 31, 2020

Uh oh!

youknowone commented Aug 31, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ArniDagur commented Aug 30, 2020 •

edited

Loading