Skip to content

Commit c7c5e67

Browse files
ShaharNavehShaharNaveh
committed
Inplace ops support in JIT
1 parent 0c9e00b commit c7c5e67

File tree

1 file changed

+87
-51
lines changed

1 file changed

+87
-51
lines changed

‎crates/jit/src/instructions.rs‎

Lines changed: 87 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -446,18 +446,30 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
446446
let b_type = b.to_jit_type();
447447

448448
let val = match (op, a, b) {
449-
(BinaryOperator::Add, JitValue::Int(a), JitValue::Int(b)) => {
449+
(
450+
BinaryOperator::Add | BinaryOperator::InplaceAdd,
451+
JitValue::Int(a),
452+
JitValue::Int(b),
453+
) => {
450454
let (out, carry) = self.builder.ins().sadd_overflow(a, b);
451455
self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW);
452456
JitValue::Int(out)
453457
}
454-
(BinaryOperator::Subtract, JitValue::Int(a), JitValue::Int(b)) => {
455-
JitValue::Int(self.compile_sub(a, b))
456-
}
457-
(BinaryOperator::FloorDivide, JitValue::Int(a), JitValue::Int(b)) => {
458-
JitValue::Int(self.builder.ins().sdiv(a, b))
459-
}
460-
(BinaryOperator::TrueDivide, JitValue::Int(a), JitValue::Int(b)) => {
458+
(
459+
BinaryOperator::Subtract | BinaryOperator::InplaceSubtract,
460+
JitValue::Int(a),
461+
JitValue::Int(b),
462+
) => JitValue::Int(self.compile_sub(a, b)),
463+
(
464+
BinaryOperator::FloorDivide | BinaryOperator::InplaceFloorDivide,
465+
JitValue::Int(a),
466+
JitValue::Int(b),
467+
) => JitValue::Int(self.builder.ins().sdiv(a, b)),
468+
(
469+
BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide,
470+
JitValue::Int(a),
471+
JitValue::Int(b),
472+
) => {
461473
// Check if b == 0, If so trap with a division by zero error
462474
self.builder
463475
.ins()
@@ -467,15 +479,21 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
467479
let b_float = self.builder.ins().fcvt_from_sint(types::F64, b);
468480
JitValue::Float(self.builder.ins().fdiv(a_float, b_float))
469481
}
470-
(BinaryOperator::Multiply, JitValue::Int(a), JitValue::Int(b)) => {
471-
JitValue::Int(self.builder.ins().imul(a, b))
472-
}
473-
(BinaryOperator::Remainder, JitValue::Int(a), JitValue::Int(b)) => {
474-
JitValue::Int(self.builder.ins().srem(a, b))
475-
}
476-
(BinaryOperator::Power, JitValue::Int(a), JitValue::Int(b)) => {
477-
JitValue::Int(self.compile_ipow(a, b))
478-
}
482+
(
483+
BinaryOperator::Multiply | BinaryOperator::InplaceMultiply,
484+
JitValue::Int(a),
485+
JitValue::Int(b),
486+
) => JitValue::Int(self.builder.ins().imul(a, b)),
487+
(
488+
BinaryOperator::Remainder | BinaryOperator::InplaceRemainder,
489+
JitValue::Int(a),
490+
JitValue::Int(b),
491+
) => JitValue::Int(self.builder.ins().srem(a, b)),
492+
(
493+
BinaryOperator::Power | BinaryOperator::InplacePower,
494+
JitValue::Int(a),
495+
JitValue::Int(b),
496+
) => JitValue::Int(self.compile_ipow(a, b)),
479497
(
480498
BinaryOperator::Lshift | BinaryOperator::Rshift,
481499
JitValue::Int(a),
@@ -489,39 +507,57 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
489507
TrapCode::user(CustomTrapCode::NegativeShiftCount as u8).unwrap(),
490508
);
491509

492-
let out = if op == BinaryOperator::Lshift {
493-
self.builder.ins().ishl(a, b)
494-
} else {
495-
self.builder.ins().sshr(a, b)
496-
};
510+
let out =
511+
if matches!(op, BinaryOperator::Lshift | BinaryOperator::InplaceLshift)
512+
{
513+
self.builder.ins().ishl(a, b)
514+
} else {
515+
self.builder.ins().sshr(a, b)
516+
};
497517
JitValue::Int(out)
498518
}
499-
(BinaryOperator::And, JitValue::Int(a), JitValue::Int(b)) => {
500-
JitValue::Int(self.builder.ins().band(a, b))
501-
}
502-
(BinaryOperator::Or, JitValue::Int(a), JitValue::Int(b)) => {
503-
JitValue::Int(self.builder.ins().bor(a, b))
504-
}
505-
(BinaryOperator::Xor, JitValue::Int(a), JitValue::Int(b)) => {
506-
JitValue::Int(self.builder.ins().bxor(a, b))
507-
}
519+
(
520+
BinaryOperator::And | BinaryOperator::InplaceAnd,
521+
JitValue::Int(a),
522+
JitValue::Int(b),
523+
) => JitValue::Int(self.builder.ins().band(a, b)),
524+
(
525+
BinaryOperator::Or | BinaryOperator::InplaceOr,
526+
JitValue::Int(a),
527+
JitValue::Int(b),
528+
) => JitValue::Int(self.builder.ins().bor(a, b)),
529+
(
530+
BinaryOperator::Xor | BinaryOperator::InplaceXor,
531+
JitValue::Int(a),
532+
JitValue::Int(b),
533+
) => JitValue::Int(self.builder.ins().bxor(a, b)),
508534

509535
// Floats
510-
(BinaryOperator::Add, JitValue::Float(a), JitValue::Float(b)) => {
511-
JitValue::Float(self.builder.ins().fadd(a, b))
512-
}
513-
(BinaryOperator::Subtract, JitValue::Float(a), JitValue::Float(b)) => {
514-
JitValue::Float(self.builder.ins().fsub(a, b))
515-
}
516-
(BinaryOperator::Multiply, JitValue::Float(a), JitValue::Float(b)) => {
517-
JitValue::Float(self.builder.ins().fmul(a, b))
518-
}
519-
(BinaryOperator::TrueDivide, JitValue::Float(a), JitValue::Float(b)) => {
520-
JitValue::Float(self.builder.ins().fdiv(a, b))
521-
}
522-
(BinaryOperator::Power, JitValue::Float(a), JitValue::Float(b)) => {
523-
JitValue::Float(self.compile_fpow(a, b))
524-
}
536+
(
537+
BinaryOperator::Add | BinaryOperator::InplaceAdd,
538+
JitValue::Float(a),
539+
JitValue::Float(b),
540+
) => JitValue::Float(self.builder.ins().fadd(a, b)),
541+
(
542+
BinaryOperator::Subtract | BinaryOperator::InplaceSubtract,
543+
JitValue::Float(a),
544+
JitValue::Float(b),
545+
) => JitValue::Float(self.builder.ins().fsub(a, b)),
546+
(
547+
BinaryOperator::Multiply | BinaryOperator::InplaceMultiply,
548+
JitValue::Float(a),
549+
JitValue::Float(b),
550+
) => JitValue::Float(self.builder.ins().fmul(a, b)),
551+
(
552+
BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide,
553+
JitValue::Float(a),
554+
JitValue::Float(b),
555+
) => JitValue::Float(self.builder.ins().fdiv(a, b)),
556+
(
557+
BinaryOperator::Power | BinaryOperator::InplacePower,
558+
JitValue::Float(a),
559+
JitValue::Float(b),
560+
) => JitValue::Float(self.compile_fpow(a, b)),
525561

526562
// Floats and Integers
527563
(_, JitValue::Int(a), JitValue::Float(b))
@@ -537,19 +573,19 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
537573
};
538574

539575
match op {
540-
BinaryOperator::Add => {
576+
BinaryOperator::Add | BinaryOperator::InplaceAdd => {
541577
JitValue::Float(self.builder.ins().fadd(operand_one, operand_two))
542578
}
543-
BinaryOperator::Subtract => {
579+
BinaryOperator::Subtract | BinaryOperator::InplaceSubtract => {
544580
JitValue::Float(self.builder.ins().fsub(operand_one, operand_two))
545581
}
546-
BinaryOperator::Multiply => {
582+
BinaryOperator::Multiply | BinaryOperator::InplaceMultiply => {
547583
JitValue::Float(self.builder.ins().fmul(operand_one, operand_two))
548584
}
549-
BinaryOperator::TrueDivide => {
585+
BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide => {
550586
JitValue::Float(self.builder.ins().fdiv(operand_one, operand_two))
551587
}
552-
BinaryOperator::Power => {
588+
BinaryOperator::Power | BinaryOperator::InplacePower => {
553589
JitValue::Float(self.compile_fpow(operand_one, operand_two))
554590
}
555591
_ => return Err(JitCompileError::NotSupported),

0 commit comments

Comments
 (0)