YJIT: Chain-guard opt_mult overflow (#8554)

* YJIT: Chain-guard opt_mult overflow

* YJIT: Support regenerating Jo after Mul

5 files changed, 29 insertions, 11 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 79af4f5858..23bcb18a06 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -824,7 +824,6 @@ impl Assembler
         let start_write_pos = cb.get_write_pos();
         let mut insn_idx: usize = 0;
         while let Some(insn) = self.insns.get(insn_idx) {
-            let mut next_insn_idx = insn_idx + 1;
             let src_ptr = cb.get_write_ptr();
             let had_dropped_bytes = cb.has_dropped_bytes();
             let old_label_state = cb.get_label_state();
@@ -878,8 +877,9 @@ impl Assembler
                 },
                 Insn::Mul { left, right, out } => {
                     // If the next instruction is jo (jump on overflow)
-                    match self.insns.get(insn_idx + 1) {
-                        Some(Insn::Jo(target)) => {
+                    match (self.insns.get(insn_idx + 1), self.insns.get(insn_idx + 2)) {
+                        (Some(Insn::JoMul(target)), _) |
+                        (Some(Insn::PosMarker(_)), Some(Insn::JoMul(target))) => {
                             // Compute the high 64 bits
                             smulh(cb, Self::SCRATCH0, left.into(), right.into());
 
@@ -895,9 +895,7 @@ impl Assembler
                             // If the high 64-bits are not all zeros or all ones,
                             // matching the sign bit, then we have an overflow
                             cmp(cb, Self::SCRATCH0, Self::SCRATCH1);
-                            emit_conditional_jump::<{Condition::NE}>(cb, compile_side_exit(*target, self, ocb));
-
-                            next_insn_idx += 1;
+                            // Insn::JoMul will emit_conditional_jump::<{Condition::NE}>
                         }
                         _ => {
                             mul(cb, out.into(), left.into(), right.into());
@@ -1120,7 +1118,7 @@ impl Assembler
                 Insn::Je(target) | Insn::Jz(target) => {
                     emit_conditional_jump::<{Condition::EQ}>(cb, compile_side_exit(*target, self, ocb));
                 },
-                Insn::Jne(target) | Insn::Jnz(target) => {
+                Insn::Jne(target) | Insn::Jnz(target) | Insn::JoMul(target) => {
                     emit_conditional_jump::<{Condition::NE}>(cb, compile_side_exit(*target, self, ocb));
                 },
                 Insn::Jl(target) => {
@@ -1197,7 +1195,7 @@ impl Assembler
                     return Err(());
                 }
             } else {
-                insn_idx = next_insn_idx;
+                insn_idx += 1;
                 gc_offsets.append(&mut insn_gc_offsets);
             }
         }
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 479a4df3db..10e463885b 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -447,6 +447,9 @@ pub enum Insn {
     /// Jump if overflow
     Jo(Target),
 
+    /// Jump if overflow in multiplication
+    JoMul(Target),
+
     /// Jump if zero
     Jz(Target),
 
@@ -590,6 +593,7 @@ impl Insn {
             Insn::Jne(_) => "Jne",
             Insn::Jnz(_) => "Jnz",
             Insn::Jo(_) => "Jo",
+            Insn::JoMul(_) => "JoMul",
             Insn::Jz(_) => "Jz",
             Insn::Label(_) => "Label",
             Insn::LeaLabel { .. } => "LeaLabel",
@@ -743,6 +747,7 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::Jne(_) |
             Insn::Jnz(_) |
             Insn::Jo(_) |
+            Insn::JoMul(_) |
             Insn::Jz(_) |
             Insn::Label(_) |
             Insn::LeaLabel { .. } |
@@ -843,6 +848,7 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::Jne(_) |
             Insn::Jnz(_) |
             Insn::Jo(_) |
+            Insn::JoMul(_) |
             Insn::Jz(_) |
             Insn::Label(_) |
             Insn::LeaLabel { .. } |
@@ -1861,6 +1867,10 @@ impl Assembler {
         self.push_insn(Insn::Jo(target));
     }
 
+    pub fn jo_mul(&mut self, target: Target) {
+        self.push_insn(Insn::JoMul(target));
+    }
+
     pub fn jz(&mut self, target: Target) {
         self.push_insn(Insn::Jz(target));
     }
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 718e95188e..7a67429488 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -740,7 +740,8 @@ impl Assembler
                     }
                 }
 
-                Insn::Jo(target) => {
+                Insn::Jo(target) |
+                Insn::JoMul(target) => {
                     match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jo_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jo_label(cb, label_idx),
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 972933e855..5c06b08b47 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -245,6 +245,7 @@ pub enum JCCKinds {
     JCC_JBE,
     JCC_JNA,
     JCC_JNAE,
+    JCC_JO_MUL,
 }
 
 #[inline(always)]
@@ -2025,6 +2026,7 @@ fn jit_chain_guard(
         JCC_JZ | JCC_JE => BranchGenFn::JZToTarget0,
         JCC_JBE | JCC_JNA => BranchGenFn::JBEToTarget0,
         JCC_JB | JCC_JNAE => BranchGenFn::JBToTarget0,
+        JCC_JO_MUL => BranchGenFn::JOMulToTarget0,
     };
 
     if (asm.ctx.get_chain_depth() as i32) < depth_limit {
@@ -3415,7 +3417,8 @@ fn gen_opt_mult(
         }
     };
 
-    if two_fixnums {
+    // Fallback to a method call if it overflows
+    if two_fixnums && asm.ctx.get_chain_depth() == 0 {
         if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_MULT) {
             return None;
         }
@@ -3432,7 +3435,7 @@ fn gen_opt_mult(
         let arg0_untag = asm.rshift(arg0, Opnd::UImm(1));
         let arg1_untag = asm.sub(arg1, Opnd::UImm(1));
         let out_val = asm.mul(arg0_untag, arg1_untag);
-        asm.jo(Target::side_exit(Counter::opt_mult_overflow));
+        jit_chain_guard(JCC_JO_MUL, jit, asm, ocb, 1, Counter::opt_mult_overflow);
         let out_val = asm.add(out_val, Opnd::UImm(1));
 
         // Push the output on the stack
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index acea71179e..aa1b12483d 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -493,6 +493,7 @@ pub enum BranchGenFn {
     JZToTarget0,
     JBEToTarget0,
     JBToTarget0,
+    JOMulToTarget0,
     JITReturn,
 }
 
@@ -549,6 +550,9 @@ impl BranchGenFn {
             BranchGenFn::JBToTarget0 => {
                 asm.jb(target0)
             }
+            BranchGenFn::JOMulToTarget0 => {
+                asm.jo_mul(target0)
+            }
             BranchGenFn::JITReturn => {
                 asm_comment!(asm, "update cfp->jit_return");
                 asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_JIT_RETURN), Opnd::const_ptr(target0.unwrap_code_ptr().raw_ptr()));
@@ -566,6 +570,7 @@ impl BranchGenFn {
             BranchGenFn::JZToTarget0 |
             BranchGenFn::JBEToTarget0 |
             BranchGenFn::JBToTarget0 |
+            BranchGenFn::JOMulToTarget0 |
             BranchGenFn::JITReturn => BranchShape::Default,
         }
     }
@@ -587,6 +592,7 @@ impl BranchGenFn {
             BranchGenFn::JZToTarget0 |
             BranchGenFn::JBEToTarget0 |
             BranchGenFn::JBToTarget0 |
+            BranchGenFn::JOMulToTarget0 |
             BranchGenFn::JITReturn => {
                 assert_eq!(new_shape, BranchShape::Default);
             }