temp

yjit/src/backend/arm64/mod.rs

@@ -2,6 +2,7 @@ use std::mem::take;
 
 use crate::asm::{CodeBlock, OutlinedCb};
 use crate::asm::arm64::*;
+use crate::core::MAX_TEMP_REGS;
 use crate::cruby::*;
 use crate::backend::ir::*;
 use crate::virtualmem::CodePtr;
@@ -178,7 +179,7 @@ fn emit_load_value(cb: &mut CodeBlock, rd: A64Opnd, value: u64) -> usize {
 
 /// List of registers that can be used for stack temps.
 /// These are caller-saved registers.
-pub static TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
+pub static TEMP_REGS: [Reg; MAX_TEMP_REGS] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
 
 #[derive(Debug, PartialEq)]
 enum EmitError {

yjit/src/backend/ir.rs

@@ -6,7 +6,7 @@ use crate::codegen::{gen_outlined_exit, gen_counted_exit};
 use crate::cruby::{vm_stack_canary, SIZEOF_VALUE_I32, VALUE};
 use crate::virtualmem::CodePtr;
 use crate::asm::{CodeBlock, OutlinedCb};
-use crate::core::{Context, RegTemps, MAX_REG_TEMPS};
+use crate::core::{Context, RegTemps, RegTemp, RegTemps2, MAX_REG_TEMPS};
 use crate::options::*;
 use crate::stats::*;
 
@@ -1203,14 +1203,20 @@ impl Assembler
             return;
         }
 
-        // Allocate a register if there's no conflict.
+        // Allocate a register if there's no conflict. (old)
         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);
         }
+
+        // Allocate a register if there's an unused register.
+        let mut reg_temps2 = self.ctx.get_reg_temps2();
+        if reg_temps2.set(RegTemp::Stack(stack_idx)) {
+            self.set_reg_temps2(reg_temps2);
+        }
     }
 
     /// Erase local variable type information
@@ -1268,7 +1274,7 @@ impl Assembler
         incr_counter!(temp_spill);
     }
 
-    /// Update which stack temps are in a register
+    /// Update which stack temps are in a register (old)
     pub fn set_reg_temps(&mut self, reg_temps: RegTemps) {
         if self.ctx.get_reg_temps() != reg_temps {
             asm_comment!(self, "reg_temps: {:08b} -> {:08b}", self.ctx.get_reg_temps().as_u8(), reg_temps.as_u8());
@@ -1286,6 +1292,17 @@ impl Assembler
         }
     }
 
+    /// Update which stack values are in a register
+    pub fn set_reg_temps2(&mut self, temp_regs: RegTemps2) {
+        if self.ctx.get_reg_temps2() != temp_regs {
+            // TODO: implement this
+            // asm_comment!(self, "temp_regs: {:08b} -> {:08b}", self.ctx.get_reg_temps().as_u8(), reg_temps.as_u8());
+            self.ctx.set_reg_temps2(temp_regs);
+            // TODO: implement this
+            // self.verify_reg_temps();
+        }
+    }
+
     /// 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.

yjit/src/backend/x86_64/mod.rs

@@ -3,6 +3,7 @@ use std::mem::take;
 use crate::asm::*;
 use crate::asm::x86_64::*;
 use crate::codegen::CodePtr;
+use crate::core::MAX_TEMP_REGS;
 use crate::cruby::*;
 use crate::backend::ir::*;
 use crate::options::*;
@@ -80,7 +81,7 @@ impl From<&Opnd> for X86Opnd {
 }
 
 /// List of registers that can be used for stack temps.
-pub static TEMP_REGS: [Reg; 5] = [RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG];
+pub static TEMP_REGS: [Reg; MAX_TEMP_REGS] = [RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG];
 
 impl Assembler
 {

yjit/src/core.rs

@@ -411,6 +411,9 @@ impl From<Opnd> for YARVOpnd {
     }
 }
 
+/// Number of registers that can be used for stack temps
+pub const MAX_TEMP_REGS: usize = 5;
+
 /// Maximum index of stack temps that could be in a register
 pub const MAX_REG_TEMPS: u8 = 8;
 
@@ -450,6 +453,60 @@ impl RegTemps {
     }
 }
 
+/// A stack slot or a local variable. u8 represents the index of it.
+#[derive(Copy, Clone, Eq, Hash, PartialEq, Debug)]
+pub enum RegTemp {
+    Stack(u8),
+    Local(u8),
+}
+
+/// RegTemps manages a set of registers used for temporary values on the stack.
+/// Each element of the array represents each of the registers.
+#[derive(Copy, Clone, Default, Eq, Hash, PartialEq, Debug)]
+pub struct RegTemps2([Option<RegTemp>; MAX_TEMP_REGS]);
+
+impl RegTemps2 {
+    /// Allocate a register for a given stack value if available.
+    /// Return true if self is updated.
+    pub fn set(&mut self, new_temp: RegTemp) -> bool {
+        // If a given value is already set, skip allocation.
+        for &temp in self.0.iter() {
+            if temp == Some(new_temp) {
+                return false;
+            }
+        }
+
+        // Allocate a register if available.
+        if let Some(reg_index) = self.alloc_reg(new_temp) {
+            self.0[reg_index] = Some(new_temp);
+            return true;
+        }
+        false
+    }
+
+    /// Pick a free register for a given stack temp.
+    /// Use lower-index registers for Stack, and higher-index registers for Local.
+    fn alloc_reg(&self, new_temp: RegTemp) -> Option<usize> {
+        // If the default index for the stack value is available, use that to minimize
+        // discrepancies among Contexts.
+        let default_index = match new_temp {
+            RegTemp::Stack(index) => index.as_usize() % MAX_TEMP_REGS,
+            RegTemp::Local(index) => MAX_TEMP_REGS - (index.as_usize() % MAX_TEMP_REGS) - 1,
+        };
+        if self.0[default_index].is_none() {
+            return Some(default_index);
+        }
+
+        // If not, pick any other available register. Like default indexes, prefer
+        // lower indexes for Stack, and higher indexes for Local.
+        let temps = self.0.iter();
+        match new_temp {
+            RegTemp::Stack(_) => temps.enumerate().find(|(_, temp)| temp.is_none()),
+            RegTemp::Local(_) => temps.rev().enumerate().find(|(_, temp)| temp.is_none()),
+        }.map(|(index, _)| index)
+    }
+}
+
 /// Bits for chain_depth_return_landing_defer
 const RETURN_LANDING_BIT: u8 = 0b10000000;
 const DEFER_BIT: u8          = 0b01000000;
@@ -475,6 +532,7 @@ pub struct Context {
 
     /// Bitmap of which stack temps are in a register
     reg_temps: RegTemps,
+    reg_temps2: RegTemps2,
 
     /// Fields packed into u8
     /// - 1st bit from the left: Whether this code is the target of a JIT-to-JIT Ruby return ([Self::is_return_landing])
@@ -2363,6 +2421,14 @@ impl Context {
         self.reg_temps = reg_temps;
     }
 
+    pub fn get_reg_temps2(&self) -> RegTemps2 {
+        self.reg_temps2
+    }
+
+    pub fn set_reg_temps2(&mut self, temp_regs: RegTemps2) {
+        self.reg_temps2 = temp_regs;
+    }
+
     pub fn get_chain_depth(&self) -> u8 {
         self.chain_depth_and_flags & CHAIN_DEPTH_MASK
     }

yjit/src/options.rs

@@ -1,5 +1,5 @@
 use std::{ffi::{CStr, CString}, ptr::null, fs::File};
-use crate::{backend::current::TEMP_REGS, stats::Counter};
+use crate::{core::MAX_TEMP_REGS, stats::Counter};
 use std::os::raw::{c_char, c_int, c_uint};
 
 // Call threshold for small deployments and command-line apps
@@ -84,7 +84,7 @@ pub static mut OPTIONS: Options = Options {
     exec_mem_size: 48 * 1024 * 1024,
     no_type_prop: false,
     max_versions: 4,
-    num_temp_regs: 5,
+    num_temp_regs: MAX_TEMP_REGS,
     gen_stats: false,
     trace_exits: None,
     print_stats: true,
@@ -223,7 +223,7 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
 
         ("temp-regs", _) => match opt_val.parse() {
             Ok(n) => {
-                assert!(n <= TEMP_REGS.len(), "--yjit-temp-regs must be <= {}", TEMP_REGS.len());
+                assert!(n <= MAX_TEMP_REGS, "--yjit-temp-regs must be <= {}", MAX_TEMP_REGS);
                 unsafe { OPTIONS.num_temp_regs = n }
             }
             Err(_) => {