YJIT: A64: Use ADDS/SUBS/CMP (immediate) when possible (#10402)

* YJIT: A64: Use ADDS/SUBS/CMP (immediate) when possible

We were loading 1 into a register and then doing ADDS/SUBS previously.
That was particularly bad since those come up in fixnum operations.

```diff
   # integer left shift with rhs=1
-  mov x11, #1
-  subs x11, x1, x11
+  subs x11, x1, #1
   lsl x12, x11, #1
   asr x13, x12, #1
   cmp x13, x11
-  b.ne #0x106ab60f8
-  mov x11, #1
-  adds x12, x12, x11
+  b.ne #0x10903a0f8
+  adds x12, x12, #1
   mov x1, x12
```

Note that it's fine to cast between i64 and u64 since the bit pattern is
preserved, and the add/sub themselves don't care about the signedness of
the operands.

CMP is just another mnemonic for SUBS.

* YJIT: A64: Split asm.mul() with immediates properly

There is in fact no MUL on A64 that takes an immediate, so this
instruction was using the wrong split method. No current usages of this
form in YJIT.

---------

Co-authored-by: Maxime Chevalier-Boisvert <[email protected]>

2 files changed, 40 insertions, 2 deletions

diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index eb99c00ba7..8365c34955 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -276,6 +276,9 @@ pub fn cmp(cb: &mut CodeBlock, rn: A64Opnd, rm: A64Opnd) {
 
             DataReg::cmp(rn.reg_no, rm.reg_no, rn.num_bits).into()
         },
+        (A64Opnd::Reg(rn), A64Opnd::Imm(imm12)) => {
+            DataImm::cmp(rn.reg_no, (imm12 as u64).try_into().unwrap(), rn.num_bits).into()
+        },
         (A64Opnd::Reg(rn), A64Opnd::UImm(imm12)) => {
             DataImm::cmp(rn.reg_no, imm12.try_into().unwrap(), rn.num_bits).into()
         },
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index e9319b46e9..092df6326f 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -315,7 +315,11 @@ impl Assembler
             match opnd {
                 Opnd::Reg(_) | Opnd::CArg(_) | Opnd::InsnOut { .. } => opnd,
                 Opnd::Mem(_) => split_load_operand(asm, opnd),
-                Opnd::Imm(_) => asm.load(opnd),
+                Opnd::Imm(imm) => if ShiftedImmediate::try_from(imm as u64).is_ok() {
+                    opnd
+                } else {
+                    asm.load(opnd)
+                }
                 Opnd::UImm(uimm) => {
                     if ShiftedImmediate::try_from(uimm).is_ok() {
                         opnd
@@ -655,7 +659,7 @@ impl Assembler
                 },
                 Insn::Mul { left, right, .. } => {
                     let opnd0 = split_load_operand(asm, *left);
-                    let opnd1 = split_shifted_immediate(asm, *right);
+                    let opnd1 = split_load_operand(asm, *right);
                     asm.mul(opnd0, opnd1);
                 },
                 Insn::Test { left, right } => {
@@ -1704,4 +1708,35 @@ mod tests {
             0x8: csel x1, x11, x12, lt
         "});
     }
+
+    #[test]
+    fn test_add_with_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let out = asm.add(Opnd::Reg(TEMP_REGS[1]), 1.into());
+        let out = asm.add(out, 1_usize.into());
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), out);
+        asm.compile_with_num_regs(&mut cb, 2);
+
+        assert_disasm!(cb, "2b0500b16b0500b1e1030baa", {"
+            0x0: adds x11, x9, #1
+            0x4: adds x11, x11, #1
+            0x8: mov x1, x11
+        "});
+    }
+
+    #[test]
+    fn test_mul_with_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let out = asm.mul(Opnd::Reg(TEMP_REGS[1]), 3.into());
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), out);
+        asm.compile_with_num_regs(&mut cb, 2);
+
+        assert_disasm!(cb, "6b0080d22b7d0b9be1030baa", {"
+            0x0: mov x11, #3
+            0x4: mul x11, x9, x11
+            0x8: mov x1, x11
+        "});
+    }
 }