From 3c4de946c9a782c3d993cfa9886f09dba46ece1d Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Tue, 2 Apr 2024 12:29:14 -0400
Subject: [PATCH] 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 <maxime.chevalierboisvert@shopify.com>
---
 yjit/src/asm/arm64/mod.rs     |  3 +++
 yjit/src/backend/arm64/mod.rs | 39 +++++++++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index eb99c00ba7cba1..8365c34955a2b3 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 e9319b46e9ea37..092df6326f2e05 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
+        "});
+    }
 }