Skip to content

Commit 58ebf04

Browse files
dohearNicholasPaulickNickJoeLoparcon4t3rusch
authored
Add JIT compilation support for integer multiplication, division, and exponents (RustPython#5561)
* Initial commit for power function * Float power jit developed * Addded support for Floats and Ints * Integration Testing for JITPower implementation. * Update instructions.rs cranelift more like painlift * Update instructions.rs * Update instructions.rs * initial commit for making stable PR ready features * fixed final edge case for compile_ipow * fixed final edge case for compile_ipow * commenting out compile_ipow * fixed spelling errors * removed unused tests * forgot to run clippy --------- Co-authored-by: Nicholas Paulick <paulicknicholas@gmail.com> Co-authored-by: Nick <nick@Samanthas-MBP.wi.rr.com> Co-authored-by: JoeLoparco <loparcojoseph@gmail.com> Co-authored-by: Nathan Rusch <nathan.rusch@icloud.com> Co-authored-by: dohear <daniel.ohear@marquette.edu>
1 parent 7fea1e1 commit 58ebf04

2 files changed

Lines changed: 232 additions & 1 deletion

File tree

‎jit/src/instructions.rs‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,12 +425,25 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
425425
(BinaryOperator::Subtract, JitValue::Int(a), JitValue::Int(b)) => {
426426
JitValue::Int(self.compile_sub(a, b))
427427
}
428+
(BinaryOperator::Multiply, JitValue::Int(a), JitValue::Int(b)) => {
429+
JitValue::Int(self.builder.ins().imul(a, b))
430+
}
428431
(BinaryOperator::FloorDivide, JitValue::Int(a), JitValue::Int(b)) => {
429432
JitValue::Int(self.builder.ins().sdiv(a, b))
430433
}
434+
(BinaryOperator::Divide, JitValue::Int(a), JitValue::Int(b)) => {
435+
// Convert to float for regular division
436+
let a_float = self.builder.ins().fcvt_from_sint(types::F64, a);
437+
let b_float = self.builder.ins().fcvt_from_sint(types::F64, b);
438+
JitValue::Float(self.builder.ins().fdiv(a_float, b_float))
439+
}
431440
(BinaryOperator::Modulo, JitValue::Int(a), JitValue::Int(b)) => {
432441
JitValue::Int(self.builder.ins().srem(a, b))
433442
}
443+
// Todo: This should return int when possible
444+
(BinaryOperator::Power, JitValue::Int(a), JitValue::Int(b)) => {
445+
JitValue::Float(self.compile_ipow(a, b))
446+
}
434447
(
435448
BinaryOperator::Lshift | BinaryOperator::Rshift,
436449
JitValue::Int(a),
@@ -562,4 +575,154 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
562575
.trapif(IntCC::Overflow, carry, TrapCode::IntegerOverflow);
563576
out
564577
}
578+
fn compile_ipow(&mut self, a: Value, b: Value) -> Value {
579+
// Convert base to float since result might not always be a Int
580+
let float_base = self.builder.ins().fcvt_from_sint(types::F64, a);
581+
582+
// Create code blocks
583+
let check_block1 = self.builder.create_block();
584+
let check_block2 = self.builder.create_block();
585+
let check_block3 = self.builder.create_block();
586+
let handle_neg_exp = self.builder.create_block();
587+
let loop_block = self.builder.create_block();
588+
let continue_block = self.builder.create_block();
589+
let exit_block = self.builder.create_block();
590+
591+
// Set code block params
592+
// Set code block params
593+
self.builder.append_block_param(check_block1, types::F64);
594+
self.builder.append_block_param(check_block1, types::I64);
595+
596+
self.builder.append_block_param(check_block2, types::F64);
597+
self.builder.append_block_param(check_block2, types::I64);
598+
599+
self.builder.append_block_param(check_block3, types::F64);
600+
self.builder.append_block_param(check_block3, types::I64);
601+
602+
self.builder.append_block_param(handle_neg_exp, types::F64);
603+
self.builder.append_block_param(handle_neg_exp, types::I64);
604+
605+
self.builder.append_block_param(loop_block, types::F64); //base
606+
self.builder.append_block_param(loop_block, types::F64); //result
607+
self.builder.append_block_param(loop_block, types::I64); //exponent
608+
609+
self.builder.append_block_param(continue_block, types::F64); //base
610+
self.builder.append_block_param(continue_block, types::F64); //result
611+
self.builder.append_block_param(continue_block, types::I64); //exponent
612+
613+
self.builder.append_block_param(exit_block, types::F64);
614+
615+
// Begin evaluating by jumping to first check block
616+
self.builder.ins().jump(check_block1, &[float_base, b]);
617+
618+
// Check block one:
619+
// Checks if input is O ** n where n > 0
620+
// Jumps to exit_block as 0 if true
621+
self.builder.switch_to_block(check_block1);
622+
let paramsc1 = self.builder.block_params(check_block1);
623+
let basec1 = paramsc1[0];
624+
let expc1 = paramsc1[1];
625+
let zero_f64 = self.builder.ins().f64const(0.0);
626+
let zero_i64 = self.builder.ins().iconst(types::I64, 0);
627+
let is_base_zero = self.builder.ins().fcmp(FloatCC::Equal, zero_f64, basec1);
628+
let is_exp_positive = self
629+
.builder
630+
.ins()
631+
.icmp(IntCC::SignedGreaterThan, expc1, zero_i64);
632+
let is_zero_to_positive = self.builder.ins().band(is_base_zero, is_exp_positive);
633+
self.builder
634+
.ins()
635+
.brnz(is_zero_to_positive, exit_block, &[zero_f64]);
636+
self.builder.ins().jump(check_block2, &[basec1, expc1]);
637+
638+
// Check block two:
639+
// Checks if exponent is negative
640+
// Jumps to a special handle_neg_exponent block if true
641+
self.builder.switch_to_block(check_block2);
642+
let paramsc2 = self.builder.block_params(check_block2);
643+
let basec2 = paramsc2[0];
644+
let expc2 = paramsc2[1];
645+
let zero_i64 = self.builder.ins().iconst(types::I64, 0);
646+
let is_neg = self
647+
.builder
648+
.ins()
649+
.icmp(IntCC::SignedLessThan, expc2, zero_i64);
650+
self.builder
651+
.ins()
652+
.brnz(is_neg, handle_neg_exp, &[basec2, expc2]);
653+
self.builder.ins().jump(check_block3, &[basec2, expc2]);
654+
655+
// Check block three:
656+
// Checks if exponent is one
657+
// jumps to exit block with the base of the exponents value
658+
self.builder.switch_to_block(check_block3);
659+
let paramsc3 = self.builder.block_params(check_block3);
660+
let basec3 = paramsc3[0];
661+
let expc3 = paramsc3[1];
662+
let resc3 = self.builder.ins().f64const(1.0);
663+
let one_i64 = self.builder.ins().iconst(types::I64, 1);
664+
let is_one = self.builder.ins().icmp(IntCC::Equal, expc3, one_i64);
665+
self.builder.ins().brnz(is_one, exit_block, &[basec3]);
666+
self.builder.ins().jump(loop_block, &[basec3, resc3, expc3]);
667+
668+
// Handles negative Exponents
669+
// calculates x^(-n) = (1/x)^n
670+
// then proceeds to the loop to evaluate
671+
self.builder.switch_to_block(handle_neg_exp);
672+
let paramshn = self.builder.block_params(handle_neg_exp);
673+
let basehn = paramshn[0];
674+
let exphn = paramshn[1];
675+
let one_f64 = self.builder.ins().f64const(1.0);
676+
let base_inverse = self.builder.ins().fdiv(one_f64, basehn);
677+
let pos_exp = self.builder.ins().ineg(exphn);
678+
self.builder
679+
.ins()
680+
.jump(loop_block, &[base_inverse, one_f64, pos_exp]);
681+
682+
// Main loop block
683+
// checks loop condition (exp > 0)
684+
// Jumps to continue block if true, exit block if false
685+
self.builder.switch_to_block(loop_block);
686+
let paramslb = self.builder.block_params(loop_block);
687+
let baselb = paramslb[0];
688+
let reslb = paramslb[1];
689+
let explb = paramslb[2];
690+
let zero = self.builder.ins().iconst(types::I64, 0);
691+
let is_zero = self.builder.ins().icmp(IntCC::Equal, explb, zero);
692+
self.builder.ins().brnz(is_zero, exit_block, &[reslb]);
693+
self.builder
694+
.ins()
695+
.jump(continue_block, &[baselb, reslb, explb]);
696+
697+
// Continue block
698+
// Main math logic
699+
// Always jumps back to loob_block
700+
self.builder.switch_to_block(continue_block);
701+
let paramscb = self.builder.block_params(continue_block);
702+
let basecb = paramscb[0];
703+
let rescb = paramscb[1];
704+
let expcb = paramscb[2];
705+
let is_odd = self.builder.ins().band_imm(expcb, 1);
706+
let is_odd = self.builder.ins().icmp_imm(IntCC::Equal, is_odd, 1);
707+
let mul_result = self.builder.ins().fmul(rescb, basecb);
708+
let new_result = self.builder.ins().select(is_odd, mul_result, rescb);
709+
let squared_base = self.builder.ins().fmul(basecb, basecb);
710+
let new_exp = self.builder.ins().sshr_imm(expcb, 1);
711+
self.builder
712+
.ins()
713+
.jump(loop_block, &[squared_base, new_result, new_exp]);
714+
715+
self.builder.switch_to_block(exit_block);
716+
let result = self.builder.block_params(exit_block)[0];
717+
718+
self.builder.seal_block(check_block1);
719+
self.builder.seal_block(check_block2);
720+
self.builder.seal_block(check_block3);
721+
self.builder.seal_block(handle_neg_exp);
722+
self.builder.seal_block(loop_block);
723+
self.builder.seal_block(continue_block);
724+
self.builder.seal_block(exit_block);
725+
726+
result
727+
}
565728
}

‎jit/tests/int_tests.rs‎

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
use core::f64;
2+
13
#[test]
24
fn test_add() {
35
let add = jit_function! { add(a:i64, b:i64) -> i64 => r##"
@@ -23,6 +25,51 @@ fn test_sub() {
2325
assert_eq!(sub(-3, -10), Ok(7));
2426
}
2527

28+
#[test]
29+
fn test_mul() {
30+
let mul = jit_function! { mul(a:i64, b:i64) -> i64 => r##"
31+
def mul(a: int, b: int):
32+
return a * b
33+
"## };
34+
35+
assert_eq!(mul(5, 10), Ok(50));
36+
assert_eq!(mul(0, 5), Ok(0));
37+
assert_eq!(mul(5, 0), Ok(0));
38+
assert_eq!(mul(0, 0), Ok(0));
39+
assert_eq!(mul(-5, 10), Ok(-50));
40+
assert_eq!(mul(5, -10), Ok(-50));
41+
assert_eq!(mul(-5, -10), Ok(50));
42+
assert_eq!(mul(999999, 999999), Ok(999998000001));
43+
assert_eq!(mul(i64::MAX, 1), Ok(i64::MAX));
44+
assert_eq!(mul(1, i64::MAX), Ok(i64::MAX));
45+
}
46+
47+
#[test]
48+
49+
fn test_div() {
50+
let div = jit_function! { div(a:i64, b:i64) -> f64 => r##"
51+
def div(a: int, b: int):
52+
return a / b
53+
"## };
54+
55+
assert_eq!(div(0, 1), Ok(0.0));
56+
assert_eq!(div(5, 1), Ok(5.0));
57+
assert_eq!(div(5, 10), Ok(0.5));
58+
assert_eq!(div(5, 2), Ok(2.5));
59+
assert_eq!(div(12, 10), Ok(1.2));
60+
assert_eq!(div(7, 10), Ok(0.7));
61+
assert_eq!(div(-3, -1), Ok(3.0));
62+
assert_eq!(div(-3, 1), Ok(-3.0));
63+
assert_eq!(div(1, 1000), Ok(0.001));
64+
assert_eq!(div(1, 100000), Ok(0.00001));
65+
assert_eq!(div(2, 3), Ok(0.6666666666666666));
66+
assert_eq!(div(1, 3), Ok(0.3333333333333333));
67+
assert_eq!(div(i64::MAX, 2), Ok(4611686018427387904.0));
68+
assert_eq!(div(i64::MIN, 2), Ok(-4611686018427387904.0));
69+
assert_eq!(div(i64::MIN, -1), Ok(9223372036854775808.0)); // Overflow case
70+
assert_eq!(div(i64::MIN, i64::MAX), Ok(-1.0));
71+
}
72+
2673
#[test]
2774
fn test_floor_div() {
2875
let floor_div = jit_function! { floor_div(a:i64, b:i64) -> i64 => r##"
@@ -35,7 +82,28 @@ fn test_floor_div() {
3582
assert_eq!(floor_div(12, 10), Ok(1));
3683
assert_eq!(floor_div(7, 10), Ok(0));
3784
assert_eq!(floor_div(-3, -1), Ok(3));
38-
assert_eq!(floor_div(-3, 1), Ok(-3));
85+
}
86+
87+
#[test]
88+
89+
fn test_exp() {
90+
let exp = jit_function! { exp(a: i64, b: i64) -> f64 => r##"
91+
def exp(a: int, b: int):
92+
return a ** b
93+
"## };
94+
95+
assert_eq!(exp(2, 3), Ok(8.0));
96+
assert_eq!(exp(3, 2), Ok(9.0));
97+
assert_eq!(exp(5, 0), Ok(1.0));
98+
assert_eq!(exp(0, 0), Ok(1.0));
99+
assert_eq!(exp(-5, 0), Ok(1.0));
100+
assert_eq!(exp(0, 1), Ok(0.0));
101+
assert_eq!(exp(0, 5), Ok(0.0));
102+
assert_eq!(exp(-2, 2), Ok(4.0));
103+
assert_eq!(exp(-3, 4), Ok(81.0));
104+
assert_eq!(exp(-2, 3), Ok(-8.0));
105+
assert_eq!(exp(-3, 3), Ok(-27.0));
106+
assert_eq!(exp(1000, 2), Ok(1000000.0));
39107
}
40108

41109
#[test]

0 commit comments

Comments
 (0)