YJIT: add mul() instruction to backend IR (#8195)

3 files changed, 36 insertions, 7 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 01485d1df6..a991a4b215 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -642,6 +642,11 @@ impl Assembler
                     let opnd1 = split_shifted_immediate(asm, *right);
                     asm.sub(opnd0, opnd1);
                 },
+                Insn::Mul { left, right, .. } => {
+                    let opnd0 = split_load_operand(asm, *left);
+                    let opnd1 = split_shifted_immediate(asm, *right);
+                    asm.mul(opnd0, opnd1);
+                },
                 Insn::Test { left, right } => {
                     // The value being tested must be in a register, so if it's
                     // not already one we'll load it first.
@@ -839,9 +844,6 @@ impl Assembler
                         cb.write_byte(0);
                     }
                 },
-                Insn::Add { left, right, out } => {
-                    adds(cb, out.into(), left.into(), right.into());
-                },
                 Insn::FrameSetup => {
                     stp_pre(cb, X29, X30, A64Opnd::new_mem(128, C_SP_REG, -16));
 
@@ -854,9 +856,15 @@ impl Assembler
 
                     ldp_post(cb, X29, X30, A64Opnd::new_mem(128, C_SP_REG, 16));
                 },
+                Insn::Add { left, right, out } => {
+                    adds(cb, out.into(), left.into(), right.into());
+                },
                 Insn::Sub { left, right, out } => {
                     subs(cb, out.into(), left.into(), right.into());
                 },
+                Insn::Mul { left, right, out } => {
+                    mul(cb, out.into(), left.into(), right.into());
+                },
                 Insn::And { left, right, out } => {
                     and(cb, out.into(), left.into(), right.into());
                 },
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 51fc24677f..d29199a390 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -507,9 +507,12 @@ pub enum Insn {
     // Low-level instruction to store a value to memory.
     Store { dest: Opnd, src: Opnd },
 
-    // This is the same as the OP_ADD instruction, except for subtraction.
+    // This is the same as the add instruction, except for subtraction.
     Sub { left: Opnd, right: Opnd, out: Opnd },
 
+    // Integer multiplication
+    Mul { left: Opnd, right: Opnd, out: Opnd },
+
     // Bitwise AND test instruction
     Test { left: Opnd, right: Opnd },
 
@@ -609,6 +612,7 @@ impl Insn {
             Insn::RShift { .. } => "RShift",
             Insn::Store { .. } => "Store",
             Insn::Sub { .. } => "Sub",
+            Insn::Mul { .. } => "Mul",
             Insn::Test { .. } => "Test",
             Insn::URShift { .. } => "URShift",
             Insn::Xor { .. } => "Xor"
@@ -641,6 +645,7 @@ impl Insn {
             Insn::Or { out, .. } |
             Insn::RShift { out, .. } |
             Insn::Sub { out, .. } |
+            Insn::Mul { out, .. } |
             Insn::URShift { out, .. } |
             Insn::Xor { out, .. } => Some(out),
             _ => None
@@ -673,6 +678,7 @@ impl Insn {
             Insn::Or { out, .. } |
             Insn::RShift { out, .. } |
             Insn::Sub { out, .. } |
+            Insn::Mul { out, .. } |
             Insn::URShift { out, .. } |
             Insn::Xor { out, .. } => Some(out),
             _ => None
@@ -782,6 +788,7 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::RShift { opnd: opnd0, shift: opnd1, .. } |
             Insn::Store { dest: opnd0, src: opnd1 } |
             Insn::Sub { left: opnd0, right: opnd1, .. } |
+            Insn::Mul { left: opnd0, right: opnd1, .. } |
             Insn::Test { left: opnd0, right: opnd1 } |
             Insn::URShift { opnd: opnd0, shift: opnd1, .. } |
             Insn::Xor { left: opnd0, right: opnd1, .. } => {
@@ -881,6 +888,7 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::RShift { opnd: opnd0, shift: opnd1, .. } |
             Insn::Store { dest: opnd0, src: opnd1 } |
             Insn::Sub { left: opnd0, right: opnd1, .. } |
+            Insn::Mul { left: opnd0, right: opnd1, .. } |
             Insn::Test { left: opnd0, right: opnd1 } |
             Insn::URShift { opnd: opnd0, shift: opnd1, .. } |
             Insn::Xor { left: opnd0, right: opnd1, .. } => {
@@ -1953,6 +1961,13 @@ impl Assembler {
         out
     }
 
+    #[must_use]
+    pub fn mul(&mut self, left: Opnd, right: Opnd) -> Opnd {
+        let out = self.next_opnd_out(Opnd::match_num_bits(&[left, right]));
+        self.push_insn(Insn::Mul { left, right, out });
+        out
+    }
+
     pub fn test(&mut self, left: Opnd, right: Opnd) {
         self.push_insn(Insn::Test { left, right });
     }
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index f60a31a9c3..c5f458cb20 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -172,6 +172,7 @@ impl Assembler
             match &mut insn {
                 Insn::Add { left, right, out } |
                 Insn::Sub { left, right, out } |
+                Insn::Mul { left, right, out } |
                 Insn::And { left, right, out } |
                 Insn::Or { left, right, out } |
                 Insn::Xor { left, right, out } => {
@@ -496,19 +497,24 @@ impl Assembler
                     cb.write_byte(0);
                 },
 
+                Insn::FrameSetup => {},
+                Insn::FrameTeardown => {},
+
                 Insn::Add { left, right, .. } => {
                     let opnd1 = emit_64bit_immediate(cb, right);
                     add(cb, left.into(), opnd1);
                 },
 
-                Insn::FrameSetup => {},
-                Insn::FrameTeardown => {},
-
                 Insn::Sub { left, right, .. } => {
                     let opnd1 = emit_64bit_immediate(cb, right);
                     sub(cb, left.into(), opnd1);
                 },
 
+                Insn::Mul { left, right, .. } => {
+                    let opnd1 = emit_64bit_immediate(cb, right);
+                    imul(cb, left.into(), opnd1);
+                },
+
                 Insn::And { left, right, .. } => {
                     let opnd1 = emit_64bit_immediate(cb, right);
                     and(cb, left.into(), opnd1);