Conditionals (https://github.com/Shopify/ruby/pull/323)

* CSEL on AArch64

* Implement various Op::CSel* instructions

3 files changed, 97 insertions, 0 deletions

diff --git a/yjit/src/asm/arm64/inst/conditional.rs b/yjit/src/asm/arm64/inst/conditional.rs
new file mode 100644
index 0000000000..e1950e95b4
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/conditional.rs
@@ -0,0 +1,73 @@
+use super::super::arg::Sf;
+
+/// The struct that represents an A64 conditional instruction that can be
+/// encoded.
+///
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
+/// |     0  0  1    1  0  1  0    1  0  0                                   0  0                                   |
+/// | sf                                   rm..............   cond.......         rn.............. rd.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct Conditional {
+    /// The number of the general-purpose destination register.
+    rd: u8,
+
+    /// The number of the first general-purpose source register.
+    rn: u8,
+
+    /// The condition to use for the conditional instruction.
+    cond: u8,
+
+    /// The number of the second general-purpose source register.
+    rm: u8,
+
+    /// The size of the registers of this instruction.
+    sf: Sf
+}
+
+impl Conditional {
+    /// CSEL
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/CSEL--Conditional-Select-?lang=en
+    pub fn csel(rd: u8, rn: u8, rm: u8, cond: u8, num_bits: u8) -> Self {
+        Self { rd, rn, cond, rm, sf: num_bits.into() }
+    }
+}
+
+/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Data-Processing----Register?lang=en#condsel
+const FAMILY: u32 = 0b101;
+
+impl From<Conditional> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: Conditional) -> Self {
+        0
+        | ((inst.sf as u32) << 31)
+        | (1 << 28)
+        | (FAMILY << 25)
+        | (1 << 23)
+        | ((inst.rm as u32) << 16)
+        | ((inst.cond as u32) << 12)
+        | ((inst.rn as u32) << 5)
+        | (inst.rd as u32)
+    }
+}
+
+impl From<Conditional> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: Conditional) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use super::super::super::arg::Condition;
+
+    #[test]
+    fn test_csel() {
+        let result: u32 = Conditional::csel(0, 1, 2, Condition::NE, 64).into();
+        assert_eq!(0x9a821020, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs
index 5d4d252d93..c69bb5e871 100644
--- a/yjit/src/asm/arm64/inst/mod.rs
+++ b/yjit/src/asm/arm64/inst/mod.rs
@@ -6,6 +6,7 @@ mod branch;
 mod branch_cond;
 mod breakpoint;
 mod call;
+mod conditional;
 mod data_imm;
 mod data_reg;
 mod load;
@@ -25,6 +26,7 @@ pub use branch::Branch;
 pub use branch_cond::BranchCond;
 pub use breakpoint::Breakpoint;
 pub use call::Call;
+pub use conditional::Conditional;
 pub use data_imm::DataImm;
 pub use data_reg::DataReg;
 pub use load::Load;
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index 6eebccaa61..3feee65197 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -272,6 +272,23 @@ pub fn cmp(cb: &mut CodeBlock, rn: A64Opnd, rm: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// CSEL - conditionally select between two registers
+pub fn csel(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd, cond: u8) {
+    let bytes: [u8; 4] = match (rd, rn, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(
+                rd.num_bits == rn.num_bits && rn.num_bits == rm.num_bits,
+                "All operands must be of the same size."
+            );
+
+            Conditional::csel(rd.reg_no, rn.reg_no, rm.reg_no, cond, rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to csel instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// LDADDAL - atomic add with acquire and release semantics
 pub fn ldaddal(cb: &mut CodeBlock, rs: A64Opnd, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rs, rt, rn) {
@@ -761,6 +778,11 @@ mod tests {
     }
 
     #[test]
+    fn test_csel() {
+        check_bytes("6a018c9a", |cb| csel(cb, X10, X11, X12, Condition::EQ));
+    }
+
+    #[test]
     fn test_ldaddal() {
         check_bytes("8b01eaf8", |cb| ldaddal(cb, X10, X11, X12));
     }