YJIT: Avoid splitting mov for small values on arm64 (#7745)

* YJIT: Avoid splitting mov for small values on arm64

* Fix a comment

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

* YJIT: Test the 0xffff boundary

---------

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

2 files changed, 32 insertions, 3 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 928406a956..3eb6419054 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -605,7 +605,11 @@ impl Assembler
                         },
                         // Otherwise we'll use the normal mov instruction.
                         (Opnd::Reg(_), _) => {
-                            let value = split_bitmask_immediate(asm, src, dest.rm_num_bits());
+                            let value = match src {
+                                // Unlike other instructions, we can avoid splitting this case, using movz.
+                                Opnd::UImm(uimm) if uimm <= 0xffff => src,
+                                _ => split_bitmask_immediate(asm, src, dest.rm_num_bits()),
+                            };
                             asm.mov(dest, value);
                         },
                         _ => unreachable!()
@@ -956,7 +960,18 @@ impl Assembler
                     };
                 },
                 Insn::Mov { dest, src } => {
-                    mov(cb, dest.into(), src.into());
+                    // This supports the following two kinds of immediates:
+                    //   * The value fits into a single movz instruction
+                    //   * It can be encoded with the special bitmask immediate encoding
+                    // arm64_split() should have split other immediates that require multiple instructions.
+                    match src {
+                        Opnd::UImm(uimm) if *uimm <= 0xffff => {
+                            movz(cb, dest.into(), A64Opnd::new_uimm(*uimm), 0);
+                        },
+                        _ => {
+                            mov(cb, dest.into(), src.into());
+                        }
+                    }
                 },
                 Insn::Lea { opnd, out } => {
                     let opnd: A64Opnd = opnd.into();
@@ -1588,4 +1603,18 @@ mod tests {
             0x0: ldur x1, [x19, #8]
         "});
     }
+
+    #[test]
+    fn test_not_split_mov() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.mov(Opnd::Reg(Assembler::TEMP_REGS[0]), Opnd::UImm(0xffff));
+        asm.mov(Opnd::Reg(Assembler::TEMP_REGS[0]), Opnd::UImm(0x10000));
+        asm.compile_with_num_regs(&mut cb, 1);
+
+        assert_disasm!(cb, "e1ff9fd2e10370b2", {"
+            0x0: mov x1, #0xffff
+            0x4: orr x1, xzr, #0x10000
+        "});
+    }
 }
diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
index f9a5744979..57806ccc30 100644
--- a/yjit/src/disasm.rs
+++ b/yjit/src/disasm.rs
@@ -207,7 +207,6 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
 #[cfg(test)]
 macro_rules! assert_disasm {
     ($cb:expr, $hex:expr, $disasm:expr) => {
-        assert_eq!(format!("{:x}", $cb), $hex);
         #[cfg(feature = "disasm")]
         {
             let disasm = disasm_addr_range(
@@ -217,6 +216,7 @@ macro_rules! assert_disasm {
             );
             assert_eq!(unindent(&disasm, false), unindent(&$disasm, true));
         }
+        assert_eq!(format!("{:x}", $cb), $hex);
     };
 }
 #[cfg(test)]