From 023bf6c3ba82e131dbf5bbef2463676eeb94a81d Mon Sep 17 00:00:00 2001 From: Tom Schuster Date: Sun, 11 Oct 2020 02:37:25 +0200 Subject: [PATCH] Basic block support in JIT --- compiler/src/compile.rs | 16 ++++--- jit/src/instructions.rs | 98 ++++++++++++++++++++++++++++++++++++++++- jit/src/lib.rs | 6 +-- jit/tests/int_tests.rs | 32 ++++++++++++++ jit/tests/misc_tests.rs | 19 ++++++++ 5 files changed, 159 insertions(+), 12 deletions(-) diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs index f92eadf30..13eac5e22 100644 --- a/compiler/src/compile.rs +++ b/compiler/src/compile.rs @@ -1388,14 +1388,18 @@ impl Compiler { self.emit(Instruction::CompareOperation { op: to_operator(ops.last().unwrap()), }); - self.emit(Instruction::Jump { target: last_label }); - // early exit left us with stack: `rhs, comparison_result`. We need to clean up rhs. - self.set_label(break_label); - self.emit(Instruction::Rotate { amount: 2 }); - self.emit(Instruction::Pop); + if vals.len() > 2 { + self.emit(Instruction::Jump { target: last_label }); + + // early exit left us with stack: `rhs, comparison_result`. We need to clean up rhs. + self.set_label(break_label); + self.emit(Instruction::Rotate { amount: 2 }); + self.emit(Instruction::Pop); + + self.set_label(last_label); + } - self.set_label(last_label); Ok(()) } diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs index 8abfa9b6f..cd75b5e4f 100644 --- a/jit/src/instructions.rs +++ b/jit/src/instructions.rs @@ -1,6 +1,8 @@ use cranelift::prelude::*; use num_traits::cast::ToPrimitive; -use rustpython_bytecode::bytecode::{BinaryOperator, Constant, Instruction, NameScope}; +use rustpython_bytecode::bytecode::{ + BinaryOperator, CodeObject, ComparisonOperator, Constant, Instruction, Label, NameScope, +}; use std::collections::HashMap; use super::{JitCompileError, JitSig, JitType}; @@ -26,6 +28,7 @@ pub struct FunctionCompiler<'a, 'b> { builder: &'a mut FunctionBuilder<'b>, stack: Vec, variables: HashMap, + label_to_block: HashMap, pub(crate) sig: JitSig, } @@ -40,6 +43,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { builder, stack: Vec::new(), variables: HashMap::new(), + label_to_block: HashMap::new(), sig: JitSig { args: arg_types.to_vec(), ret: None, @@ -76,8 +80,70 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } } - pub fn add_instruction(&mut self, instruction: &Instruction) -> Result<(), JitCompileError> { + fn boolean_val(&mut self, val: JitValue) -> Result { + match val.ty { + JitType::Float => Err(JitCompileError::NotSupported), + JitType::Int => Ok(val.val), + } + } + + pub fn compile(&mut self, bytecode: &CodeObject) -> Result<(), JitCompileError> { + let offset_to_label: HashMap<&usize, &Label> = + bytecode.label_map.iter().map(|(k, v)| (v, k)).collect(); + + for (offset, instruction) in bytecode.instructions.iter().enumerate() { + if let Some(&label) = offset_to_label.get(&offset) { + let builder = &mut self.builder; + let block = self + .label_to_block + .entry(*label) + .or_insert_with(|| builder.create_block()); + + // If the current block is not terminated/filled just jump + // into the new block. + if !self.builder.is_filled() { + self.builder.ins().jump(*block, &[]); + } + + self.builder.switch_to_block(*block); + } + + // Sometimes the bytecode contains instructions after a return + // just ignore those until we are at the next label + if self.builder.is_filled() { + continue; + } + + self.add_instruction(&instruction)?; + } + + Ok(()) + } + + fn add_instruction(&mut self, instruction: &Instruction) -> Result<(), JitCompileError> { match instruction { + Instruction::JumpIfFalse { target } => { + let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; + + let then_block = self.builder.create_block(); + self.label_to_block.insert(*target, then_block); + + let val = self.boolean_val(cond)?; + self.builder.ins().brz(val, then_block, &[]); + + let block = self.builder.create_block(); + self.builder.ins().fallthrough(block, &[]); + self.builder.switch_to_block(block); + + Ok(()) + } + Instruction::Jump { target } => { + let target_block = self.builder.create_block(); + self.label_to_block.insert(*target, target_block); + self.builder.ins().jump(target_block, &[]); + + Ok(()) + } Instruction::LoadName { name, scope: NameScope::Local, @@ -139,6 +205,34 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.builder.ins().return_(&[val.val]); Ok(()) } + Instruction::CompareOperation { op, .. } => { + // the rhs is popped off first + let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; + let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; + + match (a.ty, b.ty) { + (JitType::Int, JitType::Int) => { + let cond = match op { + ComparisonOperator::Equal => IntCC::Equal, + ComparisonOperator::NotEqual => IntCC::NotEqual, + ComparisonOperator::Less => IntCC::SignedLessThan, + ComparisonOperator::LessOrEqual => IntCC::SignedLessThanOrEqual, + ComparisonOperator::Greater => IntCC::SignedGreaterThan, + ComparisonOperator::GreaterOrEqual => IntCC::SignedLessThanOrEqual, + _ => return Err(JitCompileError::NotSupported), + }; + + let val = self.builder.ins().icmp(cond, a.val, b.val); + self.stack.push(JitValue { + val, + ty: JitType::Int, // TODO: Boolean + }); + + Ok(()) + } + _ => Err(JitCompileError::NotSupported), + } + } Instruction::BinaryOperation { op, .. } => { // the rhs is popped off first let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; diff --git a/jit/src/lib.rs b/jit/src/lib.rs index 32667bbfe..1defaf7f9 100644 --- a/jit/src/lib.rs +++ b/jit/src/lib.rs @@ -63,20 +63,18 @@ impl Jit { let entry_block = builder.create_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 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)?; - } + compiler.compile(bytecode)?; compiler.sig }; + builder.seal_all_blocks(); builder.finalize(); let id = self.module.declare_function( diff --git a/jit/tests/int_tests.rs b/jit/tests/int_tests.rs index 191229b57..0725ed064 100644 --- a/jit/tests/int_tests.rs +++ b/jit/tests/int_tests.rs @@ -22,3 +22,35 @@ fn test_sub() { assert_eq!(sub(7, 10), Ok(-3)); assert_eq!(sub(-3, -10), Ok(7)); } + +#[test] +fn test_eq() { + let eq = jit_function! { eq(a:i64, b:i64) -> i64 => r##" + def eq(a: int, b: int): + if a == b: + return 1 + return 0 + "## }; + + assert_eq!(eq(0, 0), Ok(1)); + assert_eq!(eq(1, 1), Ok(1)); + assert_eq!(eq(0, 1), Ok(0)); + assert_eq!(eq(-200, 200), Ok(0)); +} + +#[test] +fn test_gt() { + let gt = jit_function! { gt(a:i64, b:i64) -> i64 => r##" + def gt(a: int, b: int): + if a > b: + return 1 + return 0 + "## }; + + assert_eq!(gt(5, 2), Ok(1)); + assert_eq!(gt(2, 5), Ok(0)); + assert_eq!(gt(2, 2), Ok(0)); + assert_eq!(gt(5, 5), Ok(0)); + assert_eq!(gt(-1, -10), Ok(1)); + assert_eq!(gt(1, -1), Ok(1)); +} diff --git a/jit/tests/misc_tests.rs b/jit/tests/misc_tests.rs index 26f907b06..45674c12c 100644 --- a/jit/tests/misc_tests.rs +++ b/jit/tests/misc_tests.rs @@ -65,3 +65,22 @@ fn test_args_builder() { assert!(args.is_some()); assert_eq!(args.unwrap().invoke(), Some(AbiValue::Int(1))); } + +#[test] +fn test_if_else() { + let if_else = jit_function! { if_else(a:i64) -> i64 => r##" + def if_else(a: int): + if a: + return 42 + else: + return 0 + + # Prevent type failure from implicit `return None` + return 0 + "## }; + + assert_eq!(if_else(0), Ok(0)); + assert_eq!(if_else(1), Ok(42)); + assert_eq!(if_else(-1), Ok(42)); + assert_eq!(if_else(100), Ok(42)); +}