YJIT: Stack temp register allocation for arm64 (#7659)

* YJIT: Stack temp register allocation for arm64

* Update a comment

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

* Update comments about assertion

* Update a comment

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

---------

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

6 files changed, 66 insertions, 38 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index b87068df11..b85c032076 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -5,9 +5,11 @@
 use crate::asm::{CodeBlock};
 use crate::asm::arm64::*;
 use crate::codegen::{JITState, CodegenGlobals};
+use crate::core::Context;
 use crate::cruby::*;
 use crate::backend::ir::*;
 use crate::virtualmem::CodePtr;
+use crate::options::*;
 
 // Use the arm64 register type for this platform
 pub type Reg = A64Reg;
@@ -167,6 +169,10 @@ impl Assembler
     const SCRATCH0: A64Opnd = A64Opnd::Reg(X16_REG);
     const SCRATCH1: A64Opnd = A64Opnd::Reg(X17_REG);
 
+    /// List of registers that can be used for stack temps.
+    /// These are caller-saved registers.
+    pub const TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
+
     /// Get the list of registers from which we will allocate on this platform
     /// These are caller-saved registers
     /// Note: we intentionally exclude C_RET_REG (X0) from this list
@@ -175,16 +181,9 @@ impl Assembler
         vec![X11_REG, X12_REG, X13_REG]
     }
 
-    /// Get the list of registers that can be used for stack temps.
-    pub fn get_temp_regs() -> Vec<Reg> {
-        // FIXME: arm64 is not supported yet. Insn::Store doesn't support registers
-        // in its dest operand. Currently crashing at split_memory_address.
-        vec![]
-    }
-
     /// Get a list of all of the caller-saved registers
     pub fn get_caller_save_regs() -> Vec<Reg> {
-        vec![X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
+        vec![X1_REG, X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
     }
 
     /// Split platform-specific instructions
@@ -595,11 +594,6 @@ impl Assembler
                     asm.not(opnd0);
                 },
                 Insn::Store { dest, src } => {
-                    // The displacement for the STUR instruction can't be more
-                    // than 9 bits long. If it's longer, we need to load the
-                    // memory address into a register first.
-                    let opnd0 = split_memory_address(asm, dest);
-
                     // The value being stored must be in a register, so if it's
                     // not already one we'll load it first.
                     let opnd1 = match src {
@@ -610,7 +604,19 @@ impl Assembler
                         _ => split_load_operand(asm, src)
                     };
 
-                    asm.store(opnd0, opnd1);
+                    match dest {
+                        Opnd::Reg(_) => {
+                            // Store does not support a register as a dest operand.
+                            asm.mov(dest, opnd1);
+                        }
+                        _ => {
+                            // The displacement for the STUR instruction can't be more
+                            // than 9 bits long. If it's longer, we need to load the
+                            // memory address into a register first.
+                            let opnd0 = split_memory_address(asm, dest);
+                            asm.store(opnd0, opnd1);
+                        }
+                    }
                 },
                 Insn::Sub { left, right, .. } => {
                     let opnd0 = split_load_operand(asm, left);
@@ -1085,7 +1091,9 @@ impl Assembler
     /// Optimize and compile the stored instructions
     pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Vec<u32>
     {
-        let mut asm = self.lower_stack().arm64_split().alloc_regs(regs);
+        let asm = self.lower_stack();
+        let asm = asm.arm64_split();
+        let mut asm = asm.alloc_regs(regs);
 
         // Create label instances in the code block
         for (idx, name) in asm.label_names.iter().enumerate() {
@@ -1170,7 +1178,7 @@ mod tests {
     fn test_emit_cpop_all() {
         let (mut asm, mut cb) = setup_asm();
 
-        asm.cpop_all();
+        asm.cpop_all(&Context::default());
         asm.compile_with_num_regs(&mut cb, 0);
     }
 
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 37fb378905..a32750eae2 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -913,6 +913,13 @@ impl Assembler
         }
     }
 
+    /// Get the list of registers that can be used for stack temps.
+    pub fn get_temp_regs() -> Vec<Reg> {
+        let num_regs = get_option!(num_temp_regs);
+        let mut regs = Self::TEMP_REGS.to_vec();
+        regs.drain(0..num_regs).collect()
+    }
+
     /// Build an Opnd::InsnOut from the current index of the assembler and the
     /// given number of bits.
     pub(super) fn next_opnd_out(&self, num_bits: u8) -> Opnd {
@@ -1493,8 +1500,12 @@ impl Assembler {
         out
     }
 
-    pub fn cpop_all(&mut self) {
+    pub fn cpop_all(&mut self, ctx: &Context) {
         self.push_insn(Insn::CPopAll);
+
+        // Re-enable ccall's RegTemps assertion disabled by cpush_all.
+        // cpush_all + cpop_all preserve all stack temp registers, so it's safe.
+        self.set_reg_temps(ctx.get_reg_temps());
     }
 
     pub fn cpop_into(&mut self, opnd: Opnd) {
@@ -1507,6 +1518,12 @@ impl Assembler {
 
     pub fn cpush_all(&mut self) {
         self.push_insn(Insn::CPushAll);
+
+        // Mark all temps as not being in registers.
+        // Temps will be marked back as being in registers by cpop_all.
+        // We assume that cpush_all + cpop_all are used for C functions in utils.rs
+        // that don't require spill_temps for GC.
+        self.set_reg_temps(RegTemps::default());
     }
 
     pub fn cret(&mut self, opnd: Opnd) {
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 0121e19142..dffac1758d 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -88,6 +88,9 @@ impl Assembler
     // a closure and we don't want it to have to capture anything.
     const SCRATCH0: X86Opnd = X86Opnd::Reg(R11_REG);
 
+    /// List of registers that can be used for stack temps.
+    pub const TEMP_REGS: [Reg; 5] = [RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG];
+
     /// Get the list of registers from which we can allocate on this platform
     pub fn get_alloc_regs() -> Vec<Reg>
     {
@@ -98,13 +101,6 @@ impl Assembler
         ]
     }
 
-    /// Get the list of registers that can be used for stack temps.
-    pub fn get_temp_regs() -> Vec<Reg> {
-        let num_regs = get_option!(num_temp_regs);
-        let mut regs = vec![RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG];
-        regs.drain(0..num_regs).collect()
-    }
-
     /// Get a list of all of the caller-save registers
     pub fn get_caller_save_regs() -> Vec<Reg> {
         vec![RAX_REG, RCX_REG, RDX_REG, RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG, R11_REG]
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index fdee17c166..e54b2e2752 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -933,7 +933,7 @@ pub fn gen_single_block(
             // If requested, dump instructions for debugging
             if get_option!(dump_insns) {
                 println!("compiling {}", insn_name(opcode));
-                print_str(&mut asm, &format!("executing {}", insn_name(opcode)));
+                print_str(&mut asm, &ctx, &format!("executing {}", insn_name(opcode)));
             }
 
             // Call the code generation function
@@ -8402,7 +8402,9 @@ mod tests {
         let status = gen_pop(&mut jit, &mut context, &mut asm, &mut ocb);
 
         assert_eq!(status, KeepCompiling);
-        assert_eq!(context.diff(&Context::default()), TypeDiff::Compatible(0));
+        let mut default = Context::default();
+        default.set_reg_temps(context.get_reg_temps());
+        assert_eq!(context.diff(&default), TypeDiff::Compatible(0));
     }
 
     #[test]
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 759ed16205..03acc7bbe0 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -1,4 +1,5 @@
 use std::ffi::CStr;
+use crate::backend::ir::Assembler;
 
 // Command-line options
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -55,7 +56,7 @@ pub static mut OPTIONS: Options = Options {
     greedy_versioning: false,
     no_type_prop: false,
     max_versions: 4,
-    num_temp_regs: 0,
+    num_temp_regs: 5,
     gen_stats: false,
     gen_trace_exits: false,
     pause: false,
@@ -146,7 +147,10 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         },
 
         ("temp-regs", _) => match opt_val.parse() {
-            Ok(n) => unsafe { OPTIONS.num_temp_regs = n },
+            Ok(n) => {
+                assert!(n <= Assembler::TEMP_REGS.len(), "--yjit-temp-regs must be <= {}", Assembler::TEMP_REGS.len());
+                unsafe { OPTIONS.num_temp_regs = n }
+            }
             Err(_) => {
                 return None;
             }
diff --git a/yjit/src/utils.rs b/yjit/src/utils.rs
index d9a75b5302..5291e05c96 100644
--- a/yjit/src/utils.rs
+++ b/yjit/src/utils.rs
@@ -1,6 +1,7 @@
 #![allow(dead_code)] // Some functions for print debugging in here
 
 use crate::backend::ir::*;
+use crate::core::Context;
 use crate::cruby::*;
 use std::slice;
 
@@ -141,7 +142,7 @@ macro_rules! c_callable {
 }
 pub(crate) use c_callable;
 
-pub fn print_int(asm: &mut Assembler, opnd: Opnd) {
+pub fn print_int(asm: &mut Assembler, ctx: &Context, opnd: Opnd) {
     c_callable!{
         fn print_int_fn(val: i64) {
             println!("{}", val);
@@ -164,11 +165,11 @@ pub fn print_int(asm: &mut Assembler, opnd: Opnd) {
     };
 
     asm.ccall(print_int_fn as *const u8, vec![argument]);
-    asm.cpop_all();
+    asm.cpop_all(ctx);
 }
 
 /// Generate code to print a pointer
-pub fn print_ptr(asm: &mut Assembler, opnd: Opnd) {
+pub fn print_ptr(asm: &mut Assembler, ctx: &Context, opnd: Opnd) {
     c_callable!{
         fn print_ptr_fn(ptr: *const u8) {
             println!("{:p}", ptr);
@@ -179,11 +180,11 @@ pub fn print_ptr(asm: &mut Assembler, opnd: Opnd) {
 
     asm.cpush_all();
     asm.ccall(print_ptr_fn as *const u8, vec![opnd]);
-    asm.cpop_all();
+    asm.cpop_all(ctx);
 }
 
 /// Generate code to print a value
-pub fn print_value(asm: &mut Assembler, opnd: Opnd) {
+pub fn print_value(asm: &mut Assembler, ctx: &Context, opnd: Opnd) {
     c_callable!{
         fn print_value_fn(val: VALUE) {
             unsafe { rb_obj_info_dump(val) }
@@ -194,11 +195,11 @@ pub fn print_value(asm: &mut Assembler, opnd: Opnd) {
 
     asm.cpush_all();
     asm.ccall(print_value_fn as *const u8, vec![opnd]);
-    asm.cpop_all();
+    asm.cpop_all(ctx);
 }
 
 /// Generate code to print constant string to stdout
-pub fn print_str(asm: &mut Assembler, str: &str) {
+pub fn print_str(asm: &mut Assembler, ctx: &Context, str: &str) {
     c_callable!{
         fn print_str_cfun(ptr: *const u8, num_bytes: usize) {
             unsafe {
@@ -222,7 +223,7 @@ pub fn print_str(asm: &mut Assembler, str: &str) {
     let opnd = asm.lea_label(string_data);
     asm.ccall(print_str_cfun as *const u8, vec![opnd, Opnd::UImm(str.len() as u64)]);
 
-    asm.cpop_all();
+    asm.cpop_all(ctx);
 }
 
 #[cfg(test)]
@@ -262,7 +263,7 @@ mod tests {
         let mut asm = Assembler::new();
         let mut cb = CodeBlock::new_dummy(1024);
 
-        print_int(&mut asm, Opnd::Imm(42));
+        print_int(&mut asm, &Context::default(), Opnd::Imm(42));
         asm.compile(&mut cb);
     }
 
@@ -271,7 +272,7 @@ mod tests {
         let mut asm = Assembler::new();
         let mut cb = CodeBlock::new_dummy(1024);
 
-        print_str(&mut asm, "Hello, world!");
+        print_str(&mut asm, &Context::default(), "Hello, world!");
         asm.compile(&mut cb);
     }
 }