Add support for using InsnOut as memory operand base

2 files changed, 88 insertions, 24 deletions

diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index d14e3485aa..1f6307db9e 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -105,12 +105,20 @@ pub enum Op
     Breakpoint,
 }
 
+// Memory operand base
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum MemBase
+{
+    Reg(u8),
+    InsnOut(usize),
+}
+
 // Memory location
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Mem
 {
-    // Base register
-    pub(super) base_reg: Reg,
+    // Base register number or instruction index
+    pub(super) base: MemBase,
 
     // Offset relative to the base pointer
     pub(super) disp: i32,
@@ -148,11 +156,20 @@ impl Opnd
             Opnd::Reg(base_reg) => {
                 assert!(base_reg.num_bits == 64);
                 Opnd::Mem(Mem {
+                    base: MemBase::Reg(base_reg.reg_no),
+                    disp: disp,
                     num_bits: num_bits,
-                    base_reg: base_reg,
+                })
+            },
+
+            Opnd::InsnOut(idx) => {
+                Opnd::Mem(Mem {
+                    base: MemBase::InsnOut(idx),
                     disp: disp,
+                    num_bits: num_bits,
                 })
             },
+
             _ => unreachable!("memory operand with non-register base")
         }
     }
@@ -161,6 +178,13 @@ impl Opnd
     pub fn const_ptr(ptr: *const u8) -> Self {
         Opnd::UImm(ptr as u64)
     }
+
+    pub fn unwrap_reg(&self) -> Reg {
+        match self {
+            Opnd::Reg(reg) => *reg,
+            _ => unreachable!("trying to unwrap {:?} into reg", self)
+        }
+    }
 }
 
 impl From<usize> for Opnd {
@@ -264,8 +288,14 @@ impl Assembler
         // one.
         let insn_idx = self.insns.len();
         for opnd in &opnds {
-            if let Opnd::InsnOut(idx) = opnd {
-                self.live_ranges[*idx] = insn_idx;
+            match opnd {
+                Opnd::InsnOut(idx) => {
+                    self.live_ranges[*idx] = insn_idx;
+                }
+                Opnd::Mem( Mem { base: MemBase::InsnOut(idx), .. }) => {
+                    self.live_ranges[*idx] = insn_idx;
+                }
+                _ => {}
             }
         }
 
@@ -483,22 +513,26 @@ impl Assembler
             // spans more than one instruction. In that case, return the
             // allocated register to the pool.
             for opnd in &opnds {
-                if let Opnd::InsnOut(idx) = opnd {
-                    // Since we have an InsnOut, we know it spans more that one
-                    // instruction.
-                    let start_index = *idx;
-                    assert!(start_index < index);
-
-                    // We're going to check if this is the last instruction that
-                    // uses this operand. If it is, we can return the allocated
-                    // register to the pool.
-                    if live_ranges[start_index] == index {
-                        if let Opnd::Reg(reg) = asm.insns[start_index].out {
-                            dealloc_reg(&mut pool, &regs, &reg);
-                        } else {
-                            unreachable!("no register allocated for insn");
+                match opnd {
+                    Opnd::InsnOut(idx) | Opnd::Mem( Mem { base: MemBase::InsnOut(idx), .. }) => {
+                        // Since we have an InsnOut, we know it spans more that one
+                        // instruction.
+                        let start_index = *idx;
+                        assert!(start_index < index);
+
+                        // We're going to check if this is the last instruction that
+                        // uses this operand. If it is, we can return the allocated
+                        // register to the pool.
+                        if live_ranges[start_index] == index {
+                            if let Opnd::Reg(reg) = asm.insns[start_index].out {
+                                dealloc_reg(&mut pool, &regs, &reg);
+                            } else {
+                                unreachable!("no register allocated for insn");
+                            }
                         }
                     }
+
+                    _ => {}
                 }
             }
 
@@ -541,7 +575,15 @@ impl Assembler
             // Replace InsnOut operands by their corresponding register
             let reg_opnds = opnds.into_iter().map(|opnd|
                 match opnd {
-                     Opnd::InsnOut(idx) => asm.insns[idx].out,
+                    Opnd::InsnOut(idx) => asm.insns[idx].out,
+                    Opnd::Mem(Mem { base: MemBase::InsnOut(idx), disp, num_bits }) => {
+                        let out_reg = asm.insns[idx].out.unwrap_reg();
+                        Opnd::Mem(Mem {
+                            base: MemBase::Reg(out_reg.reg_no),
+                            disp,
+                            num_bits
+                        })
+                    }
                      _ => opnd,
                 }
             ).collect();
@@ -864,6 +906,22 @@ mod tests {
         asm.compile_with_regs(&mut cb, regs);
     }
 
+    // Use instruction output as base register for memory operand
+    #[test]
+    fn test_base_insn_out()
+    {
+        let (mut asm, mut cb, regs) = setup_asm(1);
+
+        // Load the pointer into a register
+        let ptr_reg = asm.load(Opnd::const_ptr(0 as *const u8));
+        let counter_opnd = Opnd::mem(64, ptr_reg, 0);
+
+        // Increment and store the updated value
+        asm.incr_counter(counter_opnd, 1.into() );
+
+        asm.compile_with_regs(&mut cb, regs);
+    }
+
     #[test]
     fn test_c_call()
     {
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index daa8005088..0c23781e20 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -6,7 +6,7 @@ use crate::asm::{CodeBlock};
 use crate::asm::x86_64::*;
 use crate::codegen::{JITState};
 use crate::cruby::*;
-use crate::backend::ir::{Assembler, Opnd, Target, Op, Mem};
+use crate::backend::ir::{Assembler, Opnd, Target, Op, MemBase, Mem};
 
 // Use the x86 register type for this platform
 pub type Reg = X86Reg;
@@ -49,8 +49,14 @@ impl From<Opnd> for X86Opnd {
             Opnd::Reg(reg) => X86Opnd::Reg(reg),
 
             // Memory operand with displacement
-            Opnd::Mem(Mem{ num_bits, base_reg, disp }) => {
-                mem_opnd(num_bits, X86Opnd::Reg(base_reg), disp)
+            Opnd::Mem(Mem{ base: MemBase::Reg(reg_no), num_bits, disp }) => {
+                let reg = X86Reg {
+                    reg_no,
+                    num_bits: 64,
+                    reg_type: RegType::GP
+                };
+
+                mem_opnd(num_bits, X86Opnd::Reg(reg), disp)
             }
 
             _ => panic!("unsupported x86 operand type")
@@ -186,7 +192,7 @@ impl Assembler
                 // Atomically increment a counter at a given memory location
                 Op::IncrCounter => {
                     assert!(matches!(insn.opnds[0], Opnd::Mem(_)));
-                    assert!(matches!(insn.opnds[0], Opnd::UImm(_)));
+                    assert!(matches!(insn.opnds[1], Opnd::UImm(_) | Opnd::Imm(_) ) );
                     write_lock_prefix(cb);
                     add(cb, insn.opnds[0].into(), insn.opnds[1].into());
                 },