Skip to content

Commit c3f18a3

Browse files
skinnyBatskinnyBat
committed
Use libffi to call jitted functions.
1 parent 343c40c commit c3f18a3

5 files changed

Lines changed: 184 additions & 16 deletions

File tree

‎Cargo.lock‎

Lines changed: 139 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎jit/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,5 +12,6 @@ cranelift = "0.66.0"
1212
cranelift-module = "0.66.0"
1313
cranelift-simplejit = "0.66.0"
1414
num-traits = "0.2"
15+
libffi = "0.9.0"
1516
rustpython-bytecode = { path = "../bytecode", version = "0.1.1" }
1617
thiserror = "1.0"

‎jit/src/instructions.rs‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,16 @@ pub struct JitSig {
1010
pub ret: Option<JitType>,
1111
}
1212

13+
impl JitSig {
14+
pub fn to_cif(&self) -> libffi::middle::Cif {
15+
let ret = match self.ret {
16+
Some(ref ty) => ty.to_libffi(),
17+
None => libffi::middle::Type::void(),
18+
};
19+
libffi::middle::Cif::new(Vec::new(), ret)
20+
}
21+
}
22+
1323
#[derive(Clone, PartialEq)]
1424
pub enum JitType {
1525
Int,
@@ -23,6 +33,13 @@ impl JitType {
2333
Self::Float => types::F64,
2434
}
2535
}
36+
37+
fn to_libffi(&self) -> libffi::middle::Type {
38+
match self {
39+
Self::Int => libffi::middle::Type::i64(),
40+
Self::Float => libffi::middle::Type::f64(),
41+
}
42+
}
2643
}
2744

2845
#[derive(Clone)]
@@ -132,8 +149,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
132149
Ok(())
133150
}
134151
Instruction::BinaryOperation { op, .. } => {
135-
let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
152+
// the rhs is popped off first
136153
let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
154+
let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
137155
match (a.ty, b.ty) {
138156
(JitType::Int, JitType::Int) => match op {
139157
BinaryOperator::Add => {

‎jit/src/lib.rs‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use std::fmt;
2-
use std::mem;
32

43
use cranelift::prelude::*;
54
use cranelift_module::{Backend, FuncId, Linkage, Module, ModuleError};
@@ -98,20 +97,13 @@ pub struct CompiledCode {
9897

9998
impl CompiledCode {
10099
pub fn invoke(&self) -> Option<AbiValue> {
101-
match self.sig.ret {
102-
Some(JitType::Int) => {
103-
let func = unsafe { mem::transmute::<_, fn() -> i64>(self.code) };
104-
Some(AbiValue::Int(func()))
105-
}
106-
Some(JitType::Float) => {
107-
let func = unsafe { mem::transmute::<_, fn() -> f64>(self.code) };
108-
Some(AbiValue::Float(func()))
109-
}
110-
None => {
111-
let func = unsafe { mem::transmute::<_, fn()>(self.code) };
112-
func();
113-
None
114-
}
100+
let cif = self.sig.to_cif();
101+
unsafe {
102+
let value = cif.call::<UnTypedAbiValue>(
103+
libffi::middle::CodePtr::from_ptr(self.code as *const _),
104+
&[],
105+
);
106+
self.sig.ret.as_ref().map(|ty| value.to_typed(ty))
115107
}
116108
}
117109
}
@@ -121,6 +113,21 @@ pub enum AbiValue {
121113
Int(i64),
122114
}
123115

116+
union UnTypedAbiValue {
117+
float: f64,
118+
int: i64,
119+
_void: (),
120+
}
121+
122+
impl UnTypedAbiValue {
123+
unsafe fn to_typed(&self, ty: &JitType) -> AbiValue {
124+
match ty {
125+
JitType::Int => AbiValue::Int(self.int),
126+
JitType::Float => AbiValue::Float(self.float),
127+
}
128+
}
129+
}
130+
124131
unsafe impl Send for CompiledCode {}
125132
unsafe impl Sync for CompiledCode {}
126133

‎tests/snippets/jit.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ def foo():
33
a = 5
44
return 10 + a
55

6+
67
def bar():
78
a = 1e6
89
return a / 5.0
@@ -12,9 +13,11 @@ def tests():
1213
assert foo() == 15
1314
assert bar() == 2e5
1415

16+
1517
tests()
1618

1719
if hasattr(foo, "__jit__"):
1820
print("Has jit")
1921
foo.__jit__()
22+
bar.__jit__()
2023
tests()

0 commit comments

Comments
 (0)