Implement MUL instruction for aarch64 (#8193)

3 files changed, 94 insertions, 0 deletions

diff --git a/yjit/src/asm/arm64/inst/madd.rs b/yjit/src/asm/arm64/inst/madd.rs
new file mode 100644
index 0000000000..683e643189
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/madd.rs
@@ -0,0 +1,73 @@
+use super::super::arg::Sf;
+
+/// The struct that represents an A64 multiply-add 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  1    0  0  0                     0                                                    |
+/// | sf                                   rm..............      ra.............. rn.............. rd.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct MAdd {
+    /// The number of the general-purpose destination register.
+    rd: u8,
+
+    /// The number of the first general-purpose source register.
+    rn: u8,
+
+    /// The number of the third general-purpose source register.
+    ra: u8,
+
+    /// The number of the second general-purpose source register.
+    rm: u8,
+
+    /// The size of the registers of this instruction.
+    sf: Sf
+}
+
+impl MAdd {
+    /// MUL
+    /// https://developer.arm.com/documentation/ddi0602/2023-06/Base-Instructions/MUL--Multiply--an-alias-of-MADD-
+    pub fn mul(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn, ra: 0b11111, rm, sf: num_bits.into() }
+    }
+}
+
+impl From<MAdd> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: MAdd) -> Self {
+        0
+        | ((inst.sf as u32) << 31)
+        | (0b11011 << 24)
+        | ((inst.rm as u32) << 16)
+        | ((inst.ra as u32) << 10)
+        | ((inst.rn as u32) << 5)
+        | (inst.rd as u32)
+    }
+}
+
+impl From<MAdd> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: MAdd) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_mul_32() {
+        let result: u32 = MAdd::mul(0, 1, 2, 32).into();
+        assert_eq!(0x1B027C20, result);
+    }
+
+    #[test]
+    fn test_mul_64() {
+        let result: u32 = MAdd::mul(0, 1, 2, 64).into();
+        assert_eq!(0x9B027C20, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs
index 9821e6a334..665ebef57c 100644
--- a/yjit/src/asm/arm64/inst/mod.rs
+++ b/yjit/src/asm/arm64/inst/mod.rs
@@ -16,6 +16,7 @@ mod load_store;
 mod load_store_exclusive;
 mod logical_imm;
 mod logical_reg;
+mod madd;
 mod mov;
 mod nop;
 mod pc_rel;
@@ -40,6 +41,7 @@ pub use load_store::LoadStore;
 pub use load_store_exclusive::LoadStoreExclusive;
 pub use logical_imm::LogicalImm;
 pub use logical_reg::LogicalReg;
+pub use madd::MAdd;
 pub use mov::Mov;
 pub use nop::Nop;
 pub use pc_rel::PCRelative;
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index 9d4c4d639b..bcdbda8dc0 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -699,6 +699,20 @@ pub fn msr(cb: &mut CodeBlock, systemregister: SystemRegister, rt: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// MUL - multiply two registers, put the result in a third register
+pub fn mul(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
+    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, "Expected registers to be the same size");
+
+            MAdd::mul(rd.reg_no, rn.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to mul instruction")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// MVN - move a value in a register to another register, negating it
 pub fn mvn(cb: &mut CodeBlock, rd: A64Opnd, rm: A64Opnd) {
     let bytes: [u8; 4] = match (rd, rm) {
@@ -1414,6 +1428,11 @@ mod tests {
     }
 
     #[test]
+    fn test_mul() {
+        check_bytes("6a7d0c9b", |cb| mul(cb, X10, X11, X12));
+    }
+
+    #[test]
     fn test_mvn() {
         check_bytes("ea032baa", |cb| mvn(cb, X10, X11));
     }