YJIT: Optimize cmp REG, 0 into test REG, REG

yjit/src/backend/x86_64/mod.rs

@@ -105,9 +105,40 @@ impl Assembler
     // These are the callee-saved registers in the x86-64 SysV ABI
     // RBX, RSP, RBP, and R12–R15
 
-    /// Merge IR instructions for the x86 platform. As of x86_split, all `out` operands
-    /// are Opnd::Out, but you sometimes want to use Opnd::Reg for example to shorten the
-    /// generated code, which is what this pass does.
+    /// Replace IR instructions for the x86 platform. This function optimizes
+    /// IR, looking at IR *before* x86_split. Optimization in this function
+    /// depends on IR to NOT interleave inserted by x86_split, e.g. Load.
+    fn x86_replace(mut self) -> Assembler {
+        let live_ranges: Vec<usize> = take(&mut self.live_ranges);
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut iterator = self.into_draining_iter();
+
+        while let Some((index, mut insn)) = iterator.next_unmapped() {
+            match (&insn, iterator.peek()) {
+                // Replace `cmp REG, 0` (4 bytes) with `test REG, REG` (3 bytes)
+                // when next IR is `je`, `jne`, `csel_e`, or `csel_ne`
+                (Insn::Cmp { left, right: Opnd::UImm(0) | Opnd::Imm(0) },
+                 Some(Insn::Je(_) | Insn::Jne(_) | Insn::CSelE { .. } | Insn::CSelNE { .. }))
+                if matches!(left, Opnd::InsnOut { .. }) => {
+                    let left = iterator.map_opnd(*left);
+                    asm.push_insn(Insn::Test { left, right: left });
+                }
+                _ => {
+                    let mut opnd_iter = insn.opnd_iter_mut();
+                    while let Some(opnd) = opnd_iter.next() {
+                        *opnd = iterator.map_opnd(*opnd);
+                    }
+                    asm.push_insn(insn);
+                }
+            }
+            iterator.map_insn_index(&mut asm);
+        }
+        asm
+    }
+
+    /// Merge IR instructions for the x86 platform. This function optimizes IR,
+    /// looking at IR *after* x86_split. x86_split overwrites all out operands
+    /// to Opnd::Out, but this function may emit Opnd::Reg operands on out operands.
     fn x86_merge(mut self) -> Assembler {
         let live_ranges: Vec<usize> = take(&mut self.live_ranges);
         let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
@@ -732,6 +763,7 @@ impl Assembler
     pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Vec<u32>
     {
         let asm = self.lower_stack();
+        let asm = asm.x86_replace();
         let asm = asm.x86_split();
         let asm = asm.x86_merge();
         let mut asm = asm.alloc_regs(regs);
@@ -943,4 +975,17 @@ mod tests {
         asm.compile_with_num_regs(&mut cb, 1);
         assert_eq!(format!("{:x}", cb), "488d4308488903");
     }
+
+    #[test]
+    fn test_replace_cmp_0() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let val = asm.load(Opnd::mem(64, SP, 8));
+        asm.cmp(val, 0.into());
+        let result = asm.csel_e(Qtrue.into(), Qfalse.into());
+        asm.mov(Opnd::Reg(RAX_REG), result);
+        asm.compile_with_num_regs(&mut cb, 2);
+
+        assert_eq!(format!("{:x}", cb), "488b43084885c0b814000000b900000000480f45c14889c0");
+    }
 }