ZJIT: Add MulHighBits LIR instruction, eliminate implicit X16

tekknolagi · tekknolagi · commit e7b5202fdda6 · 2026-03-11T02:11:38.000-07:00
Replace the speculative smulh emission (Mul always writing X16) with an explicit MulHighBits instruction that the register allocator handles like any other value. JoMul now takes both the high and low operands, and on ARM64 emits CMP high, low, ASR Shopify#62 using the barrel shifter. No implicit register contracts between instructions. The full multiply overflow sequence is: MulHighBits → Mul → JoMul, each independently emitted, each with explicit register-allocated operands. On x86, MulHighBits emits nothing (imul sets OF directly).
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
@@ -664,7 +664,8 @@ impl Assembler {
                     *opnd = split_load_operand(asm, *opnd);
                     asm.push_insn(insn);
                 },
-                Insn::Mul { left, right, .. } => {
+                Insn::Mul { left, right, .. } |
+                Insn::MulHighBits { left, right, .. } => {
                     *left = split_load_operand(asm, *left);
                     *right = split_load_operand(asm, *right);
                     asm.push_insn(insn);
@@ -793,7 +794,8 @@ impl Assembler {
                         asm.store(mem_out, SCRATCH0_OPND);
                     }
                 }
-                Insn::Mul { left, right, out } => {
+                Insn::Mul { left, right, out } |
+                Insn::MulHighBits { left, right, out } => {
                     *left = split_memory_read(asm, *left, SCRATCH0_OPND);
                     *right = split_memory_read(asm, *right, SCRATCH1_OPND);
                     let mem_out = split_memory_write(out, SCRATCH0_OPND);
@@ -907,8 +909,9 @@ impl Assembler {
                         }
                     }
                 }
-                Insn::JoMul(opnd, _) => {
-                    *opnd = split_memory_read(asm, *opnd, SCRATCH0_OPND);
+                Insn::JoMul(high, low, _) => {
+                    *high = split_memory_read(asm, *high, SCRATCH0_OPND);
+                    *low = split_memory_read(asm, *low, SCRATCH1_OPND);
                     asm.push_insn(insn);
                 }
                 &mut Insn::PatchPoint { ref target, invariant, version } => {
@@ -1235,13 +1238,11 @@ impl Assembler {
                     }
                 },
                 Insn::Mul { left, right, out } => {
-                    // Speculatively emit smulh into EMIT_OPND (X16) for a
-                    // potential following JoMul. If no JoMul follows, X16 is
-                    // simply overwritten later. Must come before mul since mul
-                    // may clobber an input register.
-                    smulh(cb, Self::EMIT_OPND, left.into(), right.into());
                     mul(cb, out.into(), left.into(), right.into());
                 },
+                Insn::MulHighBits { left, right, out } => {
+                    smulh(cb, out.into(), left.into(), right.into());
+                },
                 Insn::And { left, right, out } => {
                     and(cb, out.into(), left.into(), right.into());
                 },
@@ -1507,11 +1508,11 @@ impl Assembler {
                 Insn::Jne(target) | Insn::Jnz(target) => {
                     emit_conditional_jump::<{Condition::NE}>(self, cb, target.clone());
                 },
-                Insn::JoMul(val, target) => {
-                    // Compare smulh result (in EMIT_OPND/X16 from preceding Mul)
-                    // with the mul output sign-extended from bit 62. Uses the
-                    // barrel shifter built into CMP for a single instruction.
-                    cmp_shifted(cb, Self::EMIT_OPND, val.into(), 0b10, 62); // ASR #62
+                Insn::JoMul(high, low, target) => {
+                    // Compare MulHighBits result with the Mul output
+                    // sign-extended from bit 62. Uses the barrel shifter
+                    // built into CMP for a single instruction.
+                    cmp_shifted(cb, high.into(), low.into(), 0b10, 62); // ASR #62
                     emit_conditional_jump::<{Condition::NE}>(self, cb, target.clone());
                 },
                 Insn::Jl(target) => {
@@ -1763,11 +1764,10 @@ mod tests {
 
         assert_disasm_snapshot!(cb.disasm(), @"
         0x0: mov x0, #3
-        0x4: smulh x16, x9, x0
-        0x8: mul x0, x9, x0
-        0xc: mov x1, x0
+        0x4: mul x0, x9, x0
+        0x8: mov x1, x0
         ");
-        assert_snapshot!(cb.hexdump(), @"600080d2307d409b207d009be10300aa");
+        assert_snapshot!(cb.hexdump(), @"600080d2207d009be10300aa");
     }
 
     #[test]
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
@@ -622,8 +622,9 @@ pub enum Insn {
     Jo(Target),
 
     /// Jump if overflow in multiplication.
-    /// The operand is the Mul output, used on ARM64 for the barrel-shifted compare.
-    JoMul(Opnd, Target),
+    /// First operand is MulHighBits output, second is Mul output.
+    /// On ARM64, emits CMP high, low, ASR #62. On x86, emits jo (ignores both).
+    JoMul(Opnd, Opnd, Target),
 
     /// Jump if zero
     Jz(Target),
@@ -698,6 +699,9 @@ pub enum Insn {
     // Integer multiplication
     Mul { left: Opnd, right: Opnd, out: Opnd },
 
+    // High 64 bits of a signed 128-bit multiply (smulh on ARM64, no-op on x86)
+    MulHighBits { left: Opnd, right: Opnd, out: Opnd },
+
     // Bitwise AND test instruction
     Test { left: Opnd, right: Opnd },
 
@@ -735,7 +739,7 @@ impl Insn {
             Insn::Jne(target) |
             Insn::Jnz(target) |
             Insn::Jo(target) |
-            Insn::JoMul(_, target) |
+            Insn::JoMul(_, _, target) |
             Insn::Jz(target) |
             Insn::Joz(_, target) |
             Insn::Jonz(_, target) |
@@ -810,6 +814,7 @@ impl Insn {
             Insn::Store { .. } => "Store",
             Insn::Sub { .. } => "Sub",
             Insn::Mul { .. } => "Mul",
+            Insn::MulHighBits { .. } => "MulHighBits",
             Insn::Test { .. } => "Test",
             Insn::URShift { .. } => "URShift",
             Insn::Xor { .. } => "Xor"
@@ -843,6 +848,7 @@ impl Insn {
             Insn::RShift { out, .. } |
             Insn::Sub { out, .. } |
             Insn::Mul { out, .. } |
+            Insn::MulHighBits { out, .. } |
             Insn::URShift { out, .. } |
             Insn::Xor { out, .. } => Some(out),
             _ => None
@@ -876,6 +882,7 @@ impl Insn {
             Insn::RShift { out, .. } |
             Insn::Sub { out, .. } |
             Insn::Mul { out, .. } |
+            Insn::MulHighBits { out, .. } |
             Insn::URShift { out, .. } |
             Insn::Xor { out, .. } => Some(out),
             _ => None
@@ -895,7 +902,7 @@ impl Insn {
             Insn::Jne(target) |
             Insn::Jnz(target) |
             Insn::Jo(target) |
-            Insn::JoMul(_, target) |
+            Insn::JoMul(_, _, target) |
             Insn::Jz(target) |
             Insn::Joz(_, target) |
             Insn::Jonz(_, target) |
@@ -967,7 +974,7 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::Jne(target) |
             Insn::Jnz(target) |
             Insn::Jo(target) |
-            Insn::JoMul(_, target) |
+            Insn::JoMul(_, _, target) |
             Insn::Jz(target) |
             Insn::Label(target) |
             Insn::LeaJumpTarget { target, .. } |
@@ -1084,6 +1091,7 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::Store { dest: opnd0, src: opnd1 } |
             Insn::Sub { left: opnd0, right: opnd1, .. } |
             Insn::Mul { left: opnd0, right: opnd1, .. } |
+            Insn::MulHighBits { left: opnd0, right: opnd1, .. } |
             Insn::Test { left: opnd0, right: opnd1 } |
             Insn::URShift { opnd: opnd0, shift: opnd1, .. } |
             Insn::Xor { left: opnd0, right: opnd1, .. } => {
@@ -1192,7 +1200,43 @@ impl<'a> InsnOpndMutIterator<'a> {
                 }
             }
 
-            Insn::JoMul(opnd, target) |
+            Insn::JoMul(high, low, target) => {
+                match self.idx {
+                    0 => { self.idx += 1; return Some(high); }
+                    1 => { self.idx += 1; return Some(low); }
+                    _ => {}
+                }
+
+                match target {
+                    Target::SideExit { exit: SideExit { stack, locals, .. }, .. } => {
+                        let stack_idx = self.idx - 2;
+                        if stack_idx < stack.len() {
+                            let opnd = &mut stack[stack_idx];
+                            self.idx += 1;
+                            return Some(opnd);
+                        }
+
+                        let local_idx = stack_idx - stack.len();
+                        if local_idx < locals.len() {
+                            let opnd = &mut locals[local_idx];
+                            self.idx += 1;
+                            return Some(opnd);
+                        }
+                        None
+                    }
+                    Target::Block(edge) => {
+                        let arg_idx = self.idx - 2;
+                        if arg_idx < edge.args.len() {
+                            let opnd = &mut edge.args[arg_idx];
+                            self.idx += 1;
+                            return Some(opnd);
+                        }
+                        None
+                    }
+                    _ => None
+                }
+            }
+
             Insn::Joz(opnd, target) |
             Insn::Jonz(opnd, target) => {
                 if self.idx == 0 {
@@ -1278,6 +1322,7 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::Store { dest: opnd0, src: opnd1 } |
             Insn::Sub { left: opnd0, right: opnd1, .. } |
             Insn::Mul { left: opnd0, right: opnd1, .. } |
+            Insn::MulHighBits { left: opnd0, right: opnd1, .. } |
             Insn::Test { left: opnd0, right: opnd1 } |
             Insn::URShift { opnd: opnd0, shift: opnd1, .. } |
             Insn::Xor { left: opnd0, right: opnd1, .. } => {
@@ -1799,7 +1844,7 @@ impl Assembler
             Insn::Jbe(Target::Block(edge)) => Insn::Jbe(Target::Label(process_edge(edge))),
             Insn::Jb(Target::Block(edge)) => Insn::Jb(Target::Label(process_edge(edge))),
             Insn::Jo(Target::Block(edge)) => Insn::Jo(Target::Label(process_edge(edge))),
-            Insn::JoMul(opnd, Target::Block(edge)) => Insn::JoMul(*opnd, Target::Label(process_edge(edge))),
+            Insn::JoMul(high, low, Target::Block(edge)) => Insn::JoMul(*high, *low, Target::Label(process_edge(edge))),
             Insn::Joz(opnd, Target::Block(edge)) => Insn::Joz(*opnd, Target::Label(process_edge(edge))),
             Insn::Jonz(opnd, Target::Block(edge)) => Insn::Jonz(*opnd, Target::Label(process_edge(edge))),
             _ => insn.clone()
@@ -2453,7 +2498,9 @@ impl fmt::Display for Assembler {
                         // If the instruction has a SideExit, avoid using opnd_iter(), which has stack/locals.
                         // Here, only handle instructions that have both Opnd and Target.
                         match insn {
-                            Insn::JoMul(opnd, _) |
+                            Insn::JoMul(high, low, _) => {
+                                write!(f, ", {high}, {low}")?;
+                            }
                             Insn::Joz(opnd, _) |
                             Insn::Jonz(opnd, _) |
                             Insn::LeaJumpTarget { out: opnd, target: _ } => {
@@ -2465,7 +2512,9 @@ impl fmt::Display for Assembler {
                         // If the instruction has a Block target, avoid using opnd_iter() for branch args
                         // since they're already printed inline with the target. Only print non-target operands.
                         match insn {
-                            Insn::JoMul(opnd, _) |
+                            Insn::JoMul(high, low, _) => {
+                                write!(f, ", {high}, {low}")?;
+                            }
                             Insn::Joz(opnd, _) |
                             Insn::Jonz(opnd, _) |
                             Insn::LeaJumpTarget { out: opnd, target: _ } => {
@@ -2790,8 +2839,8 @@ impl Assembler {
         self.push_insn(Insn::Jo(target));
     }
 
-    pub fn jo_mul(&mut self, val: Opnd, target: Target) {
-        self.push_insn(Insn::JoMul(val, target));
+    pub fn jo_mul(&mut self, high: Opnd, low: Opnd, target: Target) {
+        self.push_insn(Insn::JoMul(high, low, target));
     }
 
     pub fn jz(&mut self, target: Target) {
@@ -2919,6 +2968,15 @@ impl Assembler {
         out
     }
 
+    /// High 64 bits of signed 128-bit multiply. On ARM64 this emits smulh;
+    /// on x86 it's a no-op since imul sets the overflow flag directly.
+    #[must_use]
+    pub fn mul_high_bits(&mut self, left: Opnd, right: Opnd) -> Opnd {
+        let out = self.new_vreg(Opnd::match_num_bits(&[left, right]));
+        self.push_insn(Insn::MulHighBits { left, right, out });
+        out
+    }
+
     pub fn test(&mut self, left: Opnd, right: Opnd) {
         self.push_insn(Insn::Test { left, right });
     }
diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
@@ -785,6 +785,8 @@ impl Assembler {
                 Insn::Mul { left, right, .. } => {
                     imul(cb, left.into(), right.into());
                 },
+                // x86 imul sets OF directly, no high bits needed
+                Insn::MulHighBits { .. } => {},
 
                 Insn::And { left, right, .. } => {
                     and(cb, left.into(), right.into());
@@ -1007,7 +1009,7 @@ impl Assembler {
                 }
 
                 Insn::Jo(target) |
-                Insn::JoMul(_, target) => {
+                Insn::JoMul(_, _, target) => {
                     match *target {
                         Target::CodePtr(code_ptr) => jo_ptr(cb, code_ptr),
                         Target::Label(label) => jo_label(cb, label),
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -1984,8 +1984,10 @@ fn gen_fixnum_mult(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, rig
     let right_untag = asm.sub(right, Opnd::UImm(1));
     let out_val = asm.mul(left_untag, right_untag);
 
-    // Test for overflow (on ARM64, JoMul uses out_val for barrel-shifted cmp)
-    asm.jo_mul(out_val, side_exit(jit, state, FixnumMultOverflow));
+    // Test for overflow: on ARM64, compare smulh (high bits) with sign-extended
+    // low bits. On x86, mul_high_bits is a no-op and jo_mul just reads OF.
+    let high = asm.mul_high_bits(left_untag, right_untag);
+    asm.jo_mul(high, out_val, side_exit(jit, state, FixnumMultOverflow));
     asm.add(out_val, Opnd::UImm(1))
 }
 
