Better b.cond usage on AArch64 (#6305)

* Better b.cond usage on AArch64

When we're lowering a conditional jump, we previously had a bit of
a complicated setup where we could emit a conditional jump to skip
over a jump that was the next instruction, and then write out the
destination and use a branch register.

Now instead we use the b.cond instruction if our offset fits (not
common, but not unused either) and if it doesn't we write out an
inverse condition to jump past loading the destination and
branching directly.

* Added an inverse fn for Condition (#443)

Prevents the need to pass two params and potentially reduces errors.

Co-authored-by: Jimmy Miller <jimmyhmiller@jimmys-mbp.lan>

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>
Co-authored-by: Jimmy Miller <jimmyhmiller@jimmys-mbp.lan>

3 files changed, 42 insertions, 11 deletions

diff --git a/yjit/src/asm/arm64/arg/condition.rs b/yjit/src/asm/arm64/arg/condition.rs
index e791e4b078..bb9ce570c3 100644
--- a/yjit/src/asm/arm64/arg/condition.rs
+++ b/yjit/src/asm/arm64/arg/condition.rs
@@ -19,4 +19,34 @@ impl Condition {
     pub const GT: u8 = 0b1100; // greater than (signed)
     pub const LE: u8 = 0b1101; // less than or equal to (signed)
     pub const AL: u8 = 0b1110; // always
-}
+
+    pub const fn inverse(condition: u8) -> u8 {
+        match condition {
+            Condition::EQ => Condition::NE,
+            Condition::NE => Condition::EQ,
+
+            Condition::CS => Condition::CC,
+            Condition::CC => Condition::CS,
+
+            Condition::MI => Condition::PL,
+            Condition::PL => Condition::MI,
+
+            Condition::VS => Condition::VC,
+            Condition::VC => Condition::VS,
+
+            Condition::HI => Condition::LS,
+            Condition::LS => Condition::HI,
+
+            Condition::LT => Condition::GE,
+            Condition::GE => Condition::LT,
+
+            Condition::GT => Condition::LE,
+            Condition::LE => Condition::GT,
+
+            Condition::AL => Condition::AL,
+
+            _ => panic!("Unknown condition")
+
+        }
+    }
+}
\ No newline at end of file
diff --git a/yjit/src/asm/arm64/inst/branch_cond.rs b/yjit/src/asm/arm64/inst/branch_cond.rs
index a6bc79dffe..c489bacef0 100644
--- a/yjit/src/asm/arm64/inst/branch_cond.rs
+++ b/yjit/src/asm/arm64/inst/branch_cond.rs
@@ -20,8 +20,8 @@ pub struct BranchCond {
 impl BranchCond {
     /// B.cond
     /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/B-cond--Branch-conditionally-
-    pub fn bcond(cond: u8, byte_offset: i32) -> Self {
-        Self { cond, imm19: byte_offset >> 2 }
+    pub fn bcond(cond: u8, imm19: i32) -> Self {
+        Self { cond, imm19 }
     }
 }
 
@@ -53,25 +53,25 @@ mod tests {
 
     #[test]
     fn test_b_eq() {
-        let result: u32 = BranchCond::bcond(Condition::EQ, 128).into();
+        let result: u32 = BranchCond::bcond(Condition::EQ, 32).into();
         assert_eq!(0x54000400, result);
     }
 
     #[test]
     fn test_b_vs() {
-        let result: u32 = BranchCond::bcond(Condition::VS, 128).into();
+        let result: u32 = BranchCond::bcond(Condition::VS, 32).into();
         assert_eq!(0x54000406, result);
     }
 
     #[test]
     fn test_b_eq_max() {
-        let result: u32 = BranchCond::bcond(Condition::EQ, (1 << 20) - 4).into();
+        let result: u32 = BranchCond::bcond(Condition::EQ, (1 << 18) - 1).into();
         assert_eq!(0x547fffe0, result);
     }
 
     #[test]
     fn test_b_eq_min() {
-        let result: u32 = BranchCond::bcond(Condition::EQ, -(1 << 20)).into();
+        let result: u32 = BranchCond::bcond(Condition::EQ, -(1 << 18)).into();
         assert_eq!(0x54800000, result);
     }
 }
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index a6aa8ffcbb..b73b3125e2 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -203,9 +203,10 @@ pub fn b(cb: &mut CodeBlock, imm26: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
-/// Whether or not the offset between two instructions fits into the b.cond
-/// instruction. If it doesn't, then we have to load the value into a register
-/// first, then use the b.cond instruction to skip past a direct jump.
+/// Whether or not the offset in number of instructions between two instructions
+/// fits into the b.cond instruction. If it doesn't, then we have to load the
+/// value into a register first, then use the b.cond instruction to skip past a
+/// direct jump.
 pub const fn bcond_offset_fits_bits(offset: i64) -> bool {
     imm_fits_bits(offset, 21) && (offset & 0b11 == 0)
 }
@@ -216,7 +217,7 @@ pub fn bcond(cb: &mut CodeBlock, cond: u8, byte_offset: A64Opnd) {
         A64Opnd::Imm(imm) => {
             assert!(bcond_offset_fits_bits(imm), "The immediate operand must be 21 bits or less and be aligned to a 2-bit boundary.");
 
-            BranchCond::bcond(cond, imm as i32).into()
+            BranchCond::bcond(cond, (imm / 4) as i32).into()
         },
         _ => panic!("Invalid operand combination to bcond instruction."),
     };