YJIT: Support regenerating Jo after Mul

Shopify · Sep 29, 2023 · 639cecd · 639cecd
1 parent d701d0f
commit 639cecd
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 17 deletions.
diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
@@ -4181,3 +4181,6 @@ def foo(a, b)
 
   [r1, r2, r3, r4, r5]
 }
+
+# Integer multiplication and overflow (minimized regression test from test-basic)
+assert_equal '8515157028618240000', %q{2128789257154560000 * 4}
diff --git 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
@@ -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
@@ -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
@@ -245,7 +245,7 @@ pub enum JCCKinds {
     JCC_JBE,
     JCC_JNA,
     JCC_JNAE,
-    JCC_JO,
+    JCC_JO_MUL,
 }
 
 #[inline(always)]
@@ -2026,7 +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 => BranchGenFn::JOToTarget0,
+        JCC_JO_MUL => BranchGenFn::JOMulToTarget0,
     };
 
     if (asm.ctx.get_chain_depth() as i32) < depth_limit {
@@ -3435,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);
-        jit_chain_guard(JCC_JO, jit, asm, ocb, 1, 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
@@ -493,7 +493,7 @@ pub enum BranchGenFn {
     JZToTarget0,
     JBEToTarget0,
     JBToTarget0,
-    JOToTarget0,
+    JOMulToTarget0,
     JITReturn,
 }
 
@@ -550,8 +550,8 @@ impl BranchGenFn {
             BranchGenFn::JBToTarget0 => {
                 asm.jb(target0)
             }
-            BranchGenFn::JOToTarget0 => {
-                asm.jo(target0)
+            BranchGenFn::JOMulToTarget0 => {
+                asm.jo_mul(target0)
             }
             BranchGenFn::JITReturn => {
                 asm_comment!(asm, "update cfp->jit_return");
@@ -570,7 +570,7 @@ impl BranchGenFn {
             BranchGenFn::JZToTarget0 |
             BranchGenFn::JBEToTarget0 |
             BranchGenFn::JBToTarget0 |
-            BranchGenFn::JOToTarget0 |
+            BranchGenFn::JOMulToTarget0 |
             BranchGenFn::JITReturn => BranchShape::Default,
         }
     }
@@ -592,7 +592,7 @@ impl BranchGenFn {
             BranchGenFn::JZToTarget0 |
             BranchGenFn::JBEToTarget0 |
             BranchGenFn::JBToTarget0 |
-            BranchGenFn::JOToTarget0 |
+            BranchGenFn::JOMulToTarget0 |
             BranchGenFn::JITReturn => {
                 assert_eq!(new_shape, BranchShape::Default);
             }