YJIT: Remove Insn::RegTemps (#7741)

* YJIT: Remove Insn::RegTemps

* Update a comment

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

---------

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

5 files changed, 83 insertions, 97 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 3795062da8..ee7e0edfa6 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -1113,14 +1113,13 @@ impl Assembler
                 Insn::CSelGE { truthy, falsy, out } => {
                     csel(cb, out.into(), truthy.into(), falsy.into(), Condition::GE);
                 }
-                Insn::LiveReg { .. } |
-                Insn::RegTemps(_) |
-                Insn::SpillTemp(_) => (), // just a reg alloc signal, no code
+                Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
                 Insn::PadInvalPatch => {
                     while (cb.get_write_pos().saturating_sub(std::cmp::max(start_write_pos, cb.page_start_pos()))) < cb.jmp_ptr_bytes() && !cb.has_dropped_bytes() {
                         nop(cb);
                     }
                 }
+                Insn::SpillTemp(_) => unreachable!("Insn::SpillTemp should have been lowered by lower_stack"),
             };
 
             // On failure, jump to the next page and retry the current insn
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 2e85a660ec..9a578189d9 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -75,8 +75,19 @@ pub enum Opnd
     // Output of a preceding instruction in this block
     InsnOut{ idx: usize, num_bits: u8 },
 
-    // Pointer to a slot on the VM stack
-    Stack { idx: i32, stack_size: u8, sp_offset: i8, num_bits: u8 },
+    /// Pointer to a slot on the VM stack
+    Stack {
+        /// Index from stack top. Used for conversion to StackOpnd.
+        idx: i32,
+        /// Number of bits for Opnd::Reg and Opnd::Mem.
+        num_bits: u8,
+        /// ctx.stack_size when this operand is made. Used with idx for Opnd::Reg.
+        stack_size: u8,
+        /// ctx.sp_offset when this operand is made. Used with idx for Opnd::Mem.
+        sp_offset: i8,
+        /// ctx.reg_temps when this operand is read. Used for register allocation.
+        reg_temps: Option<RegTemps>
+    },
 
     // Low-level operands, for lowering
     Imm(i64),           // Raw signed immediate
@@ -165,7 +176,7 @@ impl Opnd
             Opnd::Reg(reg) => Some(Opnd::Reg(reg.with_num_bits(num_bits))),
             Opnd::Mem(Mem { base, disp, .. }) => Some(Opnd::Mem(Mem { base, disp, num_bits })),
             Opnd::InsnOut { idx, .. } => Some(Opnd::InsnOut { idx, num_bits }),
-            Opnd::Stack { idx, stack_size, sp_offset, .. } => Some(Opnd::Stack { idx, stack_size, sp_offset, num_bits }),
+            Opnd::Stack { idx, stack_size, sp_offset, reg_temps, .. } => Some(Opnd::Stack { idx, num_bits, stack_size, sp_offset, reg_temps }),
             _ => None,
         }
     }
@@ -440,9 +451,6 @@ pub enum Insn {
     /// Take a specific register. Signal the register allocator to not use it.
     LiveReg { opnd: Opnd, out: Opnd },
 
-    /// Update live stack temps without spill
-    RegTemps(RegTemps),
-
     // A low-level instruction that loads a value into a register.
     Load { opnd: Opnd, out: Opnd },
 
@@ -571,7 +579,6 @@ impl Insn {
             Insn::LeaLabel { .. } => "LeaLabel",
             Insn::Lea { .. } => "Lea",
             Insn::LiveReg { .. } => "LiveReg",
-            Insn::RegTemps(_) => "RegTemps",
             Insn::Load { .. } => "Load",
             Insn::LoadInto { .. } => "LoadInto",
             Insn::LoadSExt { .. } => "LoadSExt",
@@ -717,7 +724,6 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::Jz(_) |
             Insn::Label(_) |
             Insn::LeaLabel { .. } |
-            Insn::RegTemps(_) |
             Insn::PadInvalPatch |
             Insn::PosMarker(_) => None,
             Insn::CPopInto(opnd) |
@@ -816,7 +822,6 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::Jz(_) |
             Insn::Label(_) |
             Insn::LeaLabel { .. } |
-            Insn::RegTemps(_) |
             Insn::PadInvalPatch |
             Insn::PosMarker(_) => None,
             Insn::CPopInto(opnd) |
@@ -929,10 +934,6 @@ pub struct Assembler {
     /// Index of the last insn using the output of this insn
     pub(super) live_ranges: Vec<usize>,
 
-    /// Parallel vec with insns
-    /// Bitmap of which temps are in a register for this insn
-    pub(super) reg_temps: Vec<RegTemps>,
-
     /// Names of labels
     pub(super) label_names: Vec<String>,
 
@@ -959,7 +960,6 @@ impl Assembler
         Self {
             insns: Vec::default(),
             live_ranges: Vec::default(),
-            reg_temps: Vec::default(),
             label_names,
             ctx: Context::default(),
             side_exits,
@@ -994,11 +994,12 @@ impl Assembler
         // Index of this instruction
         let insn_idx = self.insns.len();
 
-        // If we find any InsnOut from previous instructions, we're going to
-        // update the live range of the previous instruction to point to this
-        // one.
-        for opnd in insn.opnd_iter() {
+        let mut insn = insn;
+        let mut opnd_iter = insn.opnd_iter_mut();
+        while let Some(opnd) = opnd_iter.next() {
             match opnd {
+                // If we find any InsnOut from previous instructions, we're going to update
+                // the live range of the previous instruction to point to this one.
                 Opnd::InsnOut { idx, .. } => {
                     assert!(*idx < self.insns.len());
                     self.live_ranges[*idx] = insn_idx;
@@ -1007,29 +1008,20 @@ impl Assembler
                     assert!(*idx < self.insns.len());
                     self.live_ranges[*idx] = insn_idx;
                 }
+                // Set current ctx.reg_temps to Opnd::Stack.
+                Opnd::Stack { idx, num_bits, stack_size, sp_offset, reg_temps: None } => {
+                    *opnd = Opnd::Stack {
+                        idx: *idx,
+                        num_bits: *num_bits,
+                        stack_size: *stack_size,
+                        sp_offset: *sp_offset,
+                        reg_temps: Some(self.ctx.get_reg_temps()),
+                    };
+                }
                 _ => {}
             }
         }
 
-        // Update live stack temps for this instruction
-        let mut reg_temps = self.get_reg_temps();
-        match insn {
-            Insn::RegTemps(next_temps) => {
-                reg_temps = next_temps;
-            }
-            Insn::SpillTemp(opnd) => {
-                assert_eq!(reg_temps.get(opnd.stack_idx()), true);
-                reg_temps.set(opnd.stack_idx(), false);
-            }
-            _ => {}
-        }
-        // Assert no conflict
-        for stack_idx in 0..MAX_REG_TEMPS {
-            if reg_temps.get(stack_idx) {
-                assert!(!reg_temps.conflicts_with(stack_idx));
-            }
-        }
-
         // Set a side exit context to Target::SideExit
         let mut insn = insn;
         if let Some(Target::SideExit { context, .. }) = insn.target_mut() {
@@ -1044,12 +1036,6 @@ impl Assembler
 
         self.insns.push(insn);
         self.live_ranges.push(insn_idx);
-        self.reg_temps.push(reg_temps);
-    }
-
-    /// Get stack temps that are currently in a register
-    pub fn get_reg_temps(&self) -> RegTemps {
-        *self.reg_temps.last().unwrap_or(&RegTemps::default())
     }
 
     /// Get a cached side exit, wrapping a counter if specified
@@ -1083,8 +1069,7 @@ impl Assembler
     }
 
     /// Convert Stack operands to memory operands
-    pub fn lower_stack(mut self) -> Assembler
-    {
+    pub fn lower_stack(mut self) -> Assembler {
         // Convert Opnd::Stack to Opnd::Mem
         fn mem_opnd(opnd: &Opnd) -> Opnd {
             if let Opnd::Stack { idx, sp_offset, num_bits, .. } = *opnd {
@@ -1107,13 +1092,11 @@ impl Assembler
 
         let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), take(&mut self.side_exits));
         let regs = Assembler::get_temp_regs();
-        let reg_temps = take(&mut self.reg_temps);
         let mut iterator = self.into_draining_iter();
 
         while let Some((index, mut insn)) = iterator.next_mapped() {
             match &insn {
                 // The original insn is pushed to the new asm to satisfy ccall's reg_temps assertion.
-                Insn::RegTemps(_) => {} // noop
                 Insn::SpillTemp(opnd) => {
                     incr_counter!(temp_spill);
                     asm.mov(mem_opnd(opnd), reg_opnd(opnd, &regs));
@@ -1129,17 +1112,17 @@ impl Assembler
                     // Lower Opnd::Stack to Opnd::Reg or Opnd::Mem
                     let mut opnd_iter = insn.opnd_iter_mut();
                     while let Some(opnd) = opnd_iter.next() {
-                        if let Opnd::Stack { idx, stack_size, sp_offset, num_bits } = *opnd {
-                            *opnd = if opnd.stack_idx() < MAX_REG_TEMPS && reg_temps[index].get(opnd.stack_idx()) {
+                        if let Opnd::Stack { idx, num_bits, stack_size, sp_offset, reg_temps } = *opnd {
+                            *opnd = if opnd.stack_idx() < MAX_REG_TEMPS && reg_temps.unwrap().get(opnd.stack_idx()) {
                                 reg_opnd(opnd, &regs)
                             } else {
                                 mem_opnd(opnd)
                             };
                         }
                     }
+                    asm.push_insn(insn);
                 }
             }
-            asm.push_insn(insn);
             iterator.map_insn_index(&mut asm);
         }
 
@@ -1152,46 +1135,60 @@ impl Assembler
             return;
         }
 
-        assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());
-        let mut reg_temps = self.get_reg_temps();
-
         // Allocate a register if there's no conflict.
+        let mut reg_temps = self.ctx.get_reg_temps();
         if reg_temps.conflicts_with(stack_idx) {
             assert!(!reg_temps.get(stack_idx));
         } else {
             reg_temps.set(stack_idx, true);
             self.set_reg_temps(reg_temps);
-            self.ctx.set_reg_temps(reg_temps);
         }
     }
 
     /// Spill all live stack temps from registers to the stack
     pub fn spill_temps(&mut self) {
-        assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());
-
         // Forget registers above the stack top
-        let mut reg_temps = self.get_reg_temps();
+        let mut reg_temps = self.ctx.get_reg_temps();
         for stack_idx in self.ctx.get_stack_size()..MAX_REG_TEMPS {
             reg_temps.set(stack_idx, false);
         }
         self.set_reg_temps(reg_temps);
 
         // Spill live stack temps
-        if self.get_reg_temps() != RegTemps::default() {
-            self.comment(&format!("spill_temps: {:08b} -> {:08b}", self.get_reg_temps().as_u8(), RegTemps::default().as_u8()));
+        if self.ctx.get_reg_temps() != RegTemps::default() {
+            self.comment(&format!("spill_temps: {:08b} -> {:08b}", self.ctx.get_reg_temps().as_u8(), RegTemps::default().as_u8()));
             for stack_idx in 0..u8::min(MAX_REG_TEMPS, self.ctx.get_stack_size()) {
-                if self.get_reg_temps().get(stack_idx) {
+                if self.ctx.get_reg_temps().get(stack_idx) {
                     let idx = self.ctx.get_stack_size() - 1 - stack_idx;
                     self.spill_temp(self.stack_opnd(idx.into()));
+                    reg_temps.set(stack_idx, false);
                 }
             }
+            self.ctx.set_reg_temps(reg_temps);
         }
 
         // Every stack temp should have been spilled
-        assert_eq!(self.get_reg_temps(), RegTemps::default());
         assert_eq!(self.ctx.get_reg_temps(), RegTemps::default());
     }
 
+    /// Update which stack temps are in a register
+    pub fn set_reg_temps(&mut self, reg_temps: RegTemps) {
+        if self.ctx.get_reg_temps() != reg_temps {
+            self.comment(&format!("reg_temps: {:08b} -> {:08b}", self.ctx.get_reg_temps().as_u8(), reg_temps.as_u8()));
+            self.ctx.set_reg_temps(reg_temps);
+            self.verify_reg_temps();
+        }
+    }
+
+    /// Assert there's no conflict in stack temp register allocation
+    fn verify_reg_temps(&self) {
+        for stack_idx in 0..MAX_REG_TEMPS {
+            if self.ctx.get_reg_temps().get(stack_idx) {
+                assert!(!self.ctx.get_reg_temps().conflicts_with(stack_idx));
+            }
+        }
+    }
+
     /// Sets the out field on the various instructions that require allocated
     /// registers because their output is used as the operand on a subsequent
     /// instruction. This is our implementation of the linear scan algorithm.
@@ -1572,7 +1569,7 @@ impl Assembler {
     }
 
     pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd {
-        assert_eq!(self.get_reg_temps(), RegTemps::default(), "temps must be spilled before ccall");
+        assert_eq!(self.ctx.get_reg_temps(), RegTemps::default(), "temps must be spilled before ccall");
         let out = self.next_opnd_out(Opnd::match_num_bits(&opnds));
         self.push_insn(Insn::CCall { fptr, opnds, out });
         out
@@ -1810,23 +1807,10 @@ impl Assembler {
         out
     }
 
-    /// Update which stack temps are in a register
-    pub fn set_reg_temps(&mut self, reg_temps: RegTemps) {
-        if self.get_reg_temps() != reg_temps {
-            self.comment(&format!("reg_temps: {:08b} -> {:08b}", self.get_reg_temps().as_u8(), reg_temps.as_u8()));
-            self.push_insn(Insn::RegTemps(reg_temps));
-            self.ctx.set_reg_temps(self.get_reg_temps());
-        }
-    }
-
     /// Spill a stack temp from a register to the stack
-    pub fn spill_temp(&mut self, opnd: Opnd) {
-        assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());
-
-        if opnd.stack_idx() < MAX_REG_TEMPS && self.get_reg_temps().get(opnd.stack_idx()) {
-            self.push_insn(Insn::SpillTemp(opnd));
-            self.ctx.set_reg_temps(self.get_reg_temps());
-        }
+    fn spill_temp(&mut self, opnd: Opnd) {
+        assert!(self.ctx.get_reg_temps().get(opnd.stack_idx()));
+        self.push_insn(Insn::SpillTemp(opnd));
     }
 
     pub fn store(&mut self, dest: Opnd, src: Opnd) {
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 0fa947d1ba..69e8d24ed7 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -737,15 +737,14 @@ impl Assembler
                 Insn::CSelGE { truthy, falsy, out } => {
                     emit_csel(cb, *truthy, *falsy, *out, cmovl);
                 }
-                Insn::LiveReg { .. } |
-                Insn::RegTemps(_) |
-                Insn::SpillTemp(_) => (), // just a reg alloc signal, no code
+                Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
                 Insn::PadInvalPatch => {
                     let code_size = cb.get_write_pos().saturating_sub(std::cmp::max(start_write_pos, cb.page_start_pos()));
                     if code_size < cb.jmp_ptr_bytes() {
                         nop(cb, (cb.jmp_ptr_bytes() - code_size) as u32);
                     }
                 }
+                Insn::SpillTemp(_) => unreachable!("Insn::SpillTemp should have been lowered by lower_stack"),
             };
 
             // On failure, jump to the next page and retry the current insn
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 640940777c..599de226a4 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -805,7 +805,6 @@ pub fn gen_single_block(
         let chain_depth = if asm.ctx.get_chain_depth() > 0 { format!(", chain_depth: {}", asm.ctx.get_chain_depth()) } else { "".to_string() };
         asm.comment(&format!("Block: {} (ISEQ offset: {}{})", iseq_get_location(blockid.iseq, blockid_idx), blockid_idx, chain_depth));
     }
-    asm.set_reg_temps(asm.ctx.get_reg_temps());
 
     // For each instruction to compile
     // NOTE: could rewrite this loop with a std::iter::Iterator
@@ -834,11 +833,9 @@ pub fn gen_single_block(
 
         // stack_pop doesn't immediately deallocate a register for stack temps,
         // but it's safe to do so at this instruction boundary.
-        assert_eq!(asm.get_reg_temps(), asm.ctx.get_reg_temps());
         for stack_idx in asm.ctx.get_stack_size()..MAX_REG_TEMPS {
             asm.ctx.dealloc_temp_reg(stack_idx);
         }
-        asm.set_reg_temps(asm.ctx.get_reg_temps());
 
         // If previous instruction requested to record the boundary
         if jit.record_boundary_patch_point {
@@ -6113,7 +6110,7 @@ fn gen_send_iseq(
         // pushed onto the stack that represents the parameters that weren't
         // explicitly given a value and have a non-constant default.
         let unspec_opnd = VALUE::fixnum_from_usize(unspecified_bits).as_u64();
-        asm.spill_temp(asm.stack_opnd(-1)); // avoid using a register for unspecified_bits
+        asm.ctx.dealloc_temp_reg(asm.stack_opnd(-1).stack_idx()); // avoid using a register for unspecified_bits
         asm.mov(asm.stack_opnd(-1), unspec_opnd.into());
     }
 
@@ -6243,10 +6240,8 @@ fn gen_send_iseq(
     return_asm.ctx = asm.ctx.clone();
     return_asm.stack_pop(sp_offset.try_into().unwrap());
     let return_val = return_asm.stack_push(Type::Unknown);
-    if return_val.stack_idx() < MAX_REG_TEMPS {
-        // The callee writes a return value on stack. Update reg_temps accordingly.
-        return_asm.ctx.dealloc_temp_reg(return_val.stack_idx());
-    }
+    // The callee writes a return value on stack. Update reg_temps accordingly.
+    return_asm.ctx.dealloc_temp_reg(return_val.stack_idx());
     return_asm.ctx.set_sp_offset(1);
     return_asm.ctx.reset_chain_depth();
 
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 5362618a28..4dd0a387d5 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -1622,10 +1622,14 @@ impl Context {
     }
 
     /// Stop using a register for a given stack temp.
+    /// This allows us to reuse the register for a value that we know is dead
+    /// and will no longer be used (e.g. popped stack temp).
     pub fn dealloc_temp_reg(&mut self, stack_idx: u8) {
-        let mut reg_temps = self.get_reg_temps();
-        reg_temps.set(stack_idx, false);
-        self.set_reg_temps(reg_temps);
+        if stack_idx < MAX_REG_TEMPS {
+            let mut reg_temps = self.get_reg_temps();
+            reg_temps.set(stack_idx, false);
+            self.set_reg_temps(reg_temps);
+        }
     }
 
     /// Get the type of an instruction operand
@@ -1911,7 +1915,6 @@ impl Assembler {
         }
 
         // Allocate a register to the stack operand
-        assert_eq!(self.ctx.reg_temps, self.get_reg_temps());
         if self.ctx.stack_size < MAX_REG_TEMPS {
             self.alloc_temp_reg(self.ctx.stack_size);
         }
@@ -1982,7 +1985,13 @@ impl Assembler {
 
     /// Get an operand pointing to a slot on the temp stack
     pub fn stack_opnd(&self, idx: i32) -> Opnd {
-        Opnd::Stack { idx, stack_size: self.ctx.stack_size, sp_offset: self.ctx.sp_offset, num_bits: 64 }
+        Opnd::Stack {
+            idx,
+            num_bits: 64,
+            stack_size: self.ctx.stack_size,
+            sp_offset: self.ctx.sp_offset,
+            reg_temps: None, // push_insn will set this
+        }
     }
 }