YJIT: Stack temp register allocation (#7651)

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

.github/workflows/yjit-ubuntu.yml

@@ -78,6 +78,10 @@ jobs:
             configure: "--enable-yjit=dev"
             yjit_opts: "--yjit-call-threshold=1 --yjit-verify-ctx"
 
+          - test_task: "check"
+            configure: "--enable-yjit=dev"
+            yjit_opts: "--yjit-call-threshold=1 --yjit-temp-regs=5"
+
           - test_task: "test-all TESTS=--repeat-count=2"
             configure: "--enable-yjit=dev"
 

yjit.rb

@@ -268,6 +268,9 @@ module RubyVM::YJIT
 
       $stderr.puts "iseq_stack_too_large:  " + format_number(13, stats[:iseq_stack_too_large])
       $stderr.puts "iseq_too_long:         " + format_number(13, stats[:iseq_too_long])
+      $stderr.puts "temp_reg_opnd:         " + format_number(13, stats[:temp_reg_opnd])
+      $stderr.puts "temp_mem_opnd:         " + format_number(13, stats[:temp_mem_opnd])
+      $stderr.puts "temp_spill:            " + format_number(13, stats[:temp_spill])
       $stderr.puts "bindings_allocations:  " + format_number(13, stats[:binding_allocations])
       $stderr.puts "bindings_set:          " + format_number(13, stats[:binding_set])
       $stderr.puts "compilation_failure:   " + format_number(13, compilation_failure) if compilation_failure != 0

yjit/src/backend/arm64/mod.rs

@@ -175,6 +175,13 @@ 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]
@@ -1046,7 +1053,9 @@ impl Assembler
                 Insn::CSelGE { truthy, falsy, out } => {
                     csel(cb, out.into(), truthy.into(), falsy.into(), Condition::GE);
                 }
-                Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
+                Insn::LiveReg { .. } |
+                Insn::RegTemps(_) |
+                Insn::SpillTemp(_) => (), // 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()))) < JMP_PTR_BYTES && !cb.has_dropped_bytes() {
                         nop(cb);

yjit/src/backend/ir.rs

@@ -10,8 +10,9 @@ use std::mem::take;
 use crate::cruby::{VALUE, SIZEOF_VALUE_I32};
 use crate::virtualmem::{CodePtr};
 use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits};
-use crate::core::{Context, Type, TempMapping};
+use crate::core::{Context, Type, TempMapping, RegTemps, MAX_REG_TEMPS, MAX_TEMP_TYPES};
 use crate::options::*;
+use crate::stats::*;
 
 #[cfg(target_arch = "x86_64")]
 use crate::backend::x86_64::*;
@@ -73,7 +74,7 @@ pub enum Opnd
     InsnOut{ idx: usize, num_bits: u8 },
 
     // Pointer to a slot on the VM stack
-    Stack { idx: i32, sp_offset: i8, num_bits: u8 },
+    Stack { idx: i32, stack_size: u8, sp_offset: i8, num_bits: u8 },
 
     // Low-level operands, for lowering
     Imm(i64),           // Raw signed immediate
@@ -162,7 +163,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, sp_offset, .. } => Some(Opnd::Stack { idx, sp_offset, num_bits }),
+            Opnd::Stack { idx, stack_size, sp_offset, .. } => Some(Opnd::Stack { idx, stack_size, sp_offset, num_bits }),
             _ => None,
         }
     }
@@ -216,6 +217,26 @@ impl Opnd
     pub fn match_num_bits(opnds: &[Opnd]) -> u8 {
         Self::match_num_bits_iter(opnds.iter())
     }
+
+    /// Calculate Opnd::Stack's index from the stack bottom.
+    pub fn stack_idx(&self) -> u8 {
+        match self {
+            Opnd::Stack { idx, stack_size, .. } => {
+                (*stack_size as isize - *idx as isize - 1) as u8
+            },
+            _ => unreachable!(),
+        }
+    }
+
+    /// Get the index for stack temp registers.
+    pub fn reg_idx(&self) -> usize {
+        match self {
+            Opnd::Stack { .. } => {
+                self.stack_idx() as usize % get_option!(num_temp_regs)
+            },
+            _ => unreachable!(),
+        }
+    }
 }
 
 impl From<usize> for Opnd {
@@ -408,6 +429,9 @@ 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 },
 
@@ -443,6 +467,9 @@ pub enum Insn {
     /// Shift a value right by a certain amount (signed).
     RShift { opnd: Opnd, shift: Opnd, out: Opnd },
 
+    /// Spill a stack temp from a register into memory
+    SpillTemp(Opnd),
+
     // Low-level instruction to store a value to memory.
     Store { dest: Opnd, src: Opnd },
 
@@ -514,6 +541,7 @@ impl Insn {
             Insn::LeaLabel { .. } => "LeaLabel",
             Insn::Lea { .. } => "Lea",
             Insn::LiveReg { .. } => "LiveReg",
+            Insn::RegTemps(_) => "RegTemps",
             Insn::Load { .. } => "Load",
             Insn::LoadInto { .. } => "LoadInto",
             Insn::LoadSExt { .. } => "LoadSExt",
@@ -524,6 +552,7 @@ impl Insn {
             Insn::PadInvalPatch => "PadEntryExit",
             Insn::PosMarker(_) => "PosMarker",
             Insn::RShift { .. } => "RShift",
+            Insn::SpillTemp(_) => "SpillTemp",
             Insn::Store { .. } => "Store",
             Insn::Sub { .. } => "Sub",
             Insn::Test { .. } => "Test",
@@ -658,6 +687,7 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::Jz(_) |
             Insn::Label(_) |
             Insn::LeaLabel { .. } |
+            Insn::RegTemps(_) |
             Insn::PadInvalPatch |
             Insn::PosMarker(_) => None,
             Insn::CPopInto(opnd) |
@@ -668,7 +698,8 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::LiveReg { opnd, .. } |
             Insn::Load { opnd, .. } |
             Insn::LoadSExt { opnd, .. } |
-            Insn::Not { opnd, .. } => {
+            Insn::Not { opnd, .. } |
+            Insn::SpillTemp(opnd) => {
                 match self.idx {
                     0 => {
                         self.idx += 1;
@@ -755,6 +786,7 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::Jz(_) |
             Insn::Label(_) |
             Insn::LeaLabel { .. } |
+            Insn::RegTemps(_) |
             Insn::PadInvalPatch |
             Insn::PosMarker(_) => None,
             Insn::CPopInto(opnd) |
@@ -765,7 +797,8 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::LiveReg { opnd, .. } |
             Insn::Load { opnd, .. } |
             Insn::LoadSExt { opnd, .. } |
-            Insn::Not { opnd, .. } => {
+            Insn::Not { opnd, .. } |
+            Insn::SpillTemp(opnd) => {
                 match self.idx {
                     0 => {
                         self.idx += 1;
@@ -857,6 +890,10 @@ 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>,
 }
@@ -871,6 +908,7 @@ impl Assembler
         Self {
             insns: Vec::default(),
             live_ranges: Vec::default(),
+            reg_temps: Vec::default(),
             label_names
         }
     }
@@ -905,8 +943,33 @@ impl Assembler
             }
         }
 
+        // 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));
+            }
+        }
+
         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())
     }
 
     /// Create a new label instance that we can jump to
@@ -922,22 +985,113 @@ impl Assembler
     /// Convert Stack operands to memory operands
     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 {
+                incr_counter!(temp_mem_opnd);
+                Opnd::mem(num_bits, SP, (sp_offset as i32 - idx - 1) * SIZEOF_VALUE_I32)
+            } else {
+                unreachable!()
+            }
+        }
+
+        // Convert Opnd::Stack to Opnd::Reg
+        fn reg_opnd(opnd: &Opnd, regs: &Vec<Reg>) -> Opnd {
+            if let Opnd::Stack { num_bits, .. } = *opnd {
+                incr_counter!(temp_reg_opnd);
+                Opnd::Reg(regs[opnd.reg_idx()]).with_num_bits(num_bits).unwrap()
+            } else {
+                unreachable!()
+            }
+        }
+
         let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        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_unmapped() {
+        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));
+                }
+                _ => {
+                    // next_mapped() doesn't map out_opnd. So we need to map it here.
+                    if insn.out_opnd().is_some() {
+                        let out_num_bits = Opnd::match_num_bits_iter(insn.opnd_iter());
+                        let out = insn.out_opnd_mut().unwrap();
+                        *out = asm.next_opnd_out(out_num_bits);
+                    }
+
+                    // 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, sp_offset, num_bits } = *opnd {
+                        if let Opnd::Stack { idx, stack_size, sp_offset, num_bits } = *opnd {
-                    *opnd = Opnd::mem(num_bits, SP, (sp_offset as i32 - idx - 1) * SIZEOF_VALUE_I32);
+                            *opnd = if opnd.stack_idx() < MAX_REG_TEMPS && reg_temps[index].get(opnd.stack_idx()) {
+                                reg_opnd(opnd, &regs)
+                            } else {
+                                mem_opnd(opnd)
+                            };
+                        }
+                    }
                 }
             }
             asm.push_insn(insn);
+            iterator.map_insn_index(&mut asm);
         }
 
         asm
     }
 
+    /// Allocate a register to a stack temp if available.
+    pub fn alloc_temp_reg(&mut self, ctx: &mut Context, stack_idx: u8) {
+        if get_option!(num_temp_regs) == 0 {
+            return;
+        }
+
+        assert_eq!(self.get_reg_temps(), ctx.get_reg_temps());
+        let mut reg_temps = self.get_reg_temps();
+
+        // Allocate a register if there's no conflict.
+        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);
+            ctx.set_reg_temps(reg_temps);
+        }
+    }
+
+    /// Spill all live stack temps from registers to the stack
+    pub fn spill_temps(&mut self, ctx: &mut Context) {
+        assert_eq!(self.get_reg_temps(), ctx.get_reg_temps());
+
+        // Forget registers above the stack top
+        let mut reg_temps = self.get_reg_temps();
+        for stack_idx in 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()));
+            for stack_idx in 0..u8::min(MAX_REG_TEMPS, ctx.get_stack_size()) {
+                if self.get_reg_temps().get(stack_idx) {
+                    let idx = ctx.get_stack_size() - 1 - stack_idx;
+                    self.spill_temp(ctx.stack_opnd(idx.into()));
+                }
+            }
+        }
+
+        // Every stack temp should have been spilled
+        assert_eq!(self.get_reg_temps(), RegTemps::default());
+        ctx.set_reg_temps(self.get_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.
@@ -1318,6 +1472,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");
         let out = self.next_opnd_out(Opnd::match_num_bits(&opnds));
         self.push_insn(Insn::CCall { fptr, opnds, out });
         out
@@ -1545,6 +1700,20 @@ 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));
+        }
+    }
+
+    /// Spill a stack temp from a register to the stack
+    pub fn spill_temp(&mut self, opnd: Opnd) {
+        assert!(self.get_reg_temps().get(opnd.stack_idx()));
+        self.push_insn(Insn::SpillTemp(opnd));
+    }
+
     pub fn store(&mut self, dest: Opnd, src: Opnd) {
         self.push_insn(Insn::Store { dest, src });
     }

yjit/src/backend/x86_64/mod.rs

@@ -10,6 +10,7 @@ use crate::codegen::{JITState};
 use crate::cruby::*;
 use crate::backend::ir::*;
 use crate::codegen::CodegenGlobals;
+use crate::options::*;
 
 // Use the x86 register type for this platform
 pub type Reg = X86Reg;
@@ -97,6 +98,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 = 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]
@@ -709,7 +717,9 @@ impl Assembler
                 Insn::CSelGE { truthy, falsy, out } => {
                     emit_csel(cb, *truthy, *falsy, *out, cmovl);
                 }
-                Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
+                Insn::LiveReg { .. } |
+                Insn::RegTemps(_) |
+                Insn::SpillTemp(_) => (), // 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 < JMP_PTR_BYTES {

yjit/src/core.rs

@@ -371,6 +371,45 @@ impl From<Opnd> for YARVOpnd {
     }
 }
 
+/// Maximum index of stack temps that could be in a register
+pub const MAX_REG_TEMPS: u8 = 8;
+
+/// Bitmap of which stack temps are in a register
+#[derive(Copy, Clone, Default, Eq, Hash, PartialEq, Debug)]
+pub struct RegTemps(u8);
+
+impl RegTemps {
+    pub fn get(&self, index: u8) -> bool {
+        assert!(index < MAX_REG_TEMPS);
+        (self.0 >> index) & 1 == 1
+    }
+
+    pub fn set(&mut self, index: u8, value: bool) {
+        assert!(index < MAX_REG_TEMPS);
+        if value {
+            self.0 = self.0 | (1 << index);
+        } else {
+            self.0 = self.0 & !(1 << index);
+        }
+    }
+
+    pub fn as_u8(&self) -> u8 {
+        self.0
+    }
+
+    /// Return true if there's a register that conflicts with a given stack_idx.
+    pub fn conflicts_with(&self, stack_idx: u8) -> bool {
+        let mut other_idx = stack_idx as isize - get_option!(num_temp_regs) as isize;
+        while other_idx >= 0 {
+            if self.get(other_idx as u8) {
+                return true;
+            }
+            other_idx -= get_option!(num_temp_regs) as isize;
+        }
+        false
+    }
+}
+
 /// Code generation context
 /// Contains information we can use to specialize/optimize code
 /// There are a lot of context objects so we try to keep the size small.
@@ -383,6 +422,9 @@ pub struct Context {
     // This represents how far the JIT's SP is from the "real" SP
     sp_offset: i8,
 
+    /// Bitmap of which stack temps are in a register
+    reg_temps: RegTemps,
+
     // Depth of this block in the sidechain (eg: inline-cache chain)
     chain_depth: u8,
 
@@ -698,7 +740,7 @@ impl PendingBranch {
         // The branch struct is uninitialized right now but as a stable address.
         // We make sure the stub runs after the branch is initialized.
         let branch_struct_addr = self.uninit_branch.as_ptr() as usize;
-        let stub_addr = gen_branch_stub(ocb, branch_struct_addr, target_idx);
+        let stub_addr = gen_branch_stub(ctx, ocb, branch_struct_addr, target_idx);
 
         if let Some(stub_addr) = stub_addr {
             // Fill the branch target with a stub
@@ -1333,6 +1375,7 @@ pub fn limit_block_versions(blockid: BlockId, ctx: &Context) -> Context {
         let mut generic_ctx = Context::default();
         generic_ctx.stack_size = ctx.stack_size;
         generic_ctx.sp_offset = ctx.sp_offset;
+        generic_ctx.reg_temps = ctx.reg_temps;
 
         debug_assert_ne!(
             TypeDiff::Incompatible,
@@ -1534,6 +1577,14 @@ impl Context {
         self.sp_offset = offset;
     }
 
+    pub fn get_reg_temps(&self) -> RegTemps {
+        self.reg_temps
+    }
+
+    pub fn set_reg_temps(&mut self, reg_temps: RegTemps) {
+        self.reg_temps = reg_temps;
+    }
+
     pub fn get_chain_depth(&self) -> u8 {
         self.chain_depth
     }
@@ -1553,12 +1604,19 @@ impl Context {
         return Opnd::mem(64, SP, offset);
     }
 
+    /// Stop using a register for a given 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);
+    }
+
     /// Push one new value on the temp stack with an explicit mapping
     /// Return a pointer to the new stack top
-    pub fn stack_push_mapping(&mut self, (mapping, temp_type): (TempMapping, Type)) -> Opnd {
+    pub fn stack_push_mapping(&mut self, asm: &mut Assembler, (mapping, temp_type): (TempMapping, Type)) -> Opnd {
         // If type propagation is disabled, store no types
         if get_option!(no_type_prop) {
-            return self.stack_push_mapping((mapping, Type::Unknown));
+            return self.stack_push_mapping(asm, (mapping, Type::Unknown));
         }
 
         let stack_size: usize = self.stack_size.into();
@@ -1573,6 +1631,12 @@ impl Context {
             }
         }
 
+        // Allocate a register to the stack operand
+        assert_eq!(self.reg_temps, asm.get_reg_temps());
+        if self.stack_size < MAX_REG_TEMPS {
+            asm.alloc_temp_reg(self, self.stack_size);
+        }
+
         self.stack_size += 1;
         self.sp_offset += 1;
 
@@ -1581,22 +1645,22 @@ impl Context {
 
     /// Push one new value on the temp stack
     /// Return a pointer to the new stack top
-    pub fn stack_push(&mut self, val_type: Type) -> Opnd {
+    pub fn stack_push(&mut self, asm: &mut Assembler, val_type: Type) -> Opnd {
-        return self.stack_push_mapping((MapToStack, val_type));
+        return self.stack_push_mapping(asm, (MapToStack, val_type));
     }
 
     /// Push the self value on the stack
-    pub fn stack_push_self(&mut self) -> Opnd {
+    pub fn stack_push_self(&mut self, asm: &mut Assembler) -> Opnd {
-        return self.stack_push_mapping((MapToSelf, Type::Unknown));
+        return self.stack_push_mapping(asm, (MapToSelf, Type::Unknown));
     }
 
     /// Push a local variable on the stack
-    pub fn stack_push_local(&mut self, local_idx: usize) -> Opnd {
+    pub fn stack_push_local(&mut self, asm: &mut Assembler, local_idx: usize) -> Opnd {
         if local_idx >= MAX_LOCAL_TYPES {
-            return self.stack_push(Type::Unknown);
+            return self.stack_push(asm, Type::Unknown);
         }
 
-        return self.stack_push_mapping((MapToLocal((local_idx as u8).into()), Type::Unknown));
+        return self.stack_push_mapping(asm, (MapToLocal((local_idx as u8).into()), Type::Unknown));
     }
 
     // Pop N values off the stack
@@ -1639,7 +1703,7 @@ impl Context {
 
     /// Get an operand pointing to a slot on the temp stack
     pub fn stack_opnd(&self, idx: i32) -> Opnd {
-        Opnd::Stack { idx, sp_offset: self.sp_offset, num_bits: 64 }
+        Opnd::Stack { idx, stack_size: self.stack_size, sp_offset: self.sp_offset, num_bits: 64 }
     }
 
     /// Get the type of an instruction operand
@@ -1838,6 +1902,10 @@ impl Context {
             return TypeDiff::Incompatible;
         }
 
+        if dst.reg_temps != src.reg_temps {
+            return TypeDiff::Incompatible;
+        }
+
         // Difference sum
         let mut diff = 0;
 
@@ -2466,6 +2534,7 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
 /// Generate a "stub", a piece of code that calls the compiler back when run.
 /// A piece of code that redeems for more code; a thunk for code.
 fn gen_branch_stub(
+    ctx: &Context,
     ocb: &mut OutlinedCb,
     branch_struct_address: usize,
     target_idx: u32,
@@ -2476,8 +2545,18 @@ fn gen_branch_stub(
     let stub_addr = ocb.get_write_ptr();
 
     let mut asm = Assembler::new();
+    asm.set_reg_temps(ctx.reg_temps);
     asm.comment("branch stub hit");
 
+    // Save caller-saved registers before C_ARG_OPNDS get clobbered.
+    // Spill all registers for consistency with the trampoline.
+    for &reg in caller_saved_temp_regs().iter() {
+        asm.cpush(reg);
+    }
+
+    // Spill temps to the VM stack as well for jit.peek_at_stack()
+    asm.spill_temps(&mut ctx.clone());
+
     // Set up the arguments unique to this stub for:
     //
     //    branch_stub_hit(branch_ptr, target_idx, ec)
@@ -2522,6 +2601,11 @@ pub fn gen_branch_stub_hit_trampoline(ocb: &mut OutlinedCb) -> CodePtr {
         ]
     );
 
+    // Restore caller-saved registers for stack temps
+    for &reg in caller_saved_temp_regs().iter().rev() {
+        asm.cpop_into(reg);
+    }
+
     // Jump to the address returned by the branch_stub_hit() call
     asm.jmp_opnd(jump_addr);
 
@@ -2530,6 +2614,16 @@ pub fn gen_branch_stub_hit_trampoline(ocb: &mut OutlinedCb) -> CodePtr {
     code_ptr
 }
 
+/// Return registers to be pushed and popped on branch_stub_hit.
+/// The return value may include an extra register for x86 alignment.
+fn caller_saved_temp_regs() -> Vec<Opnd> {
+    let mut regs = Assembler::get_temp_regs();
+    if regs.len() % 2 == 1 {
+        regs.push(*regs.last().unwrap()); // x86 alignment
+    }
+    regs.iter().map(|&reg| Opnd::Reg(reg)).collect()
+}
+
 impl Assembler
 {
     /// Mark the start position of a patchable entry point in the machine code
@@ -2888,7 +2982,7 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
         }
 
         // Create a stub for this branch target
-        let stub_addr = gen_branch_stub(ocb, branchref.as_ptr() as usize, target_idx as u32);
+        let stub_addr = gen_branch_stub(&block.ctx, ocb, branchref.as_ptr() as usize, target_idx as u32);
 
         // In case we were unable to generate a stub (e.g. OOM). Use the block's
         // exit instead of a stub for the block. It's important that we
@@ -3031,6 +3125,33 @@ mod tests {
         assert_eq!(Type::Fixnum.diff(Type::UnknownHeap), TypeDiff::Incompatible);
     }
 
+    #[test]
+    fn reg_temps() {
+        let mut reg_temps = RegTemps(0);
+
+        // 0 means every slot is not spilled
+        for stack_idx in 0..MAX_REG_TEMPS {
+            assert_eq!(reg_temps.get(stack_idx), false);
+        }
+
+        // Set 0, 2, 7
+        reg_temps.set(0, true);
+        reg_temps.set(2, true);
+        reg_temps.set(3, true);
+        reg_temps.set(3, false);
+        reg_temps.set(7, true);
+
+        // Get 0..8
+        assert_eq!(reg_temps.get(0), true);
+        assert_eq!(reg_temps.get(1), false);
+        assert_eq!(reg_temps.get(2), true);
+        assert_eq!(reg_temps.get(3), false);
+        assert_eq!(reg_temps.get(4), false);
+        assert_eq!(reg_temps.get(5), false);
+        assert_eq!(reg_temps.get(6), false);
+        assert_eq!(reg_temps.get(7), true);
+    }
+
     #[test]
     fn context() {
         // Valid src => dst
@@ -3038,7 +3159,7 @@ mod tests {
 
         // Try pushing an operand and getting its type
         let mut ctx = Context::default();
-        ctx.stack_push(Type::Fixnum);
+        ctx.stack_push(&mut Assembler::new(), Type::Fixnum);
         let top_type = ctx.get_opnd_type(StackOpnd(0));
         assert!(top_type == Type::Fixnum);
 

yjit/src/options.rs

@@ -22,6 +22,9 @@ pub struct Options {
     // 1 means always create generic versions
     pub max_versions: usize,
 
+    // The number of registers allocated for stack temps
+    pub num_temp_regs: usize,
+
     // Capture and print out stats
     pub gen_stats: bool,
 
@@ -52,6 +55,7 @@ pub static mut OPTIONS: Options = Options {
     greedy_versioning: false,
     no_type_prop: false,
     max_versions: 4,
+    num_temp_regs: 0,
     gen_stats: false,
     gen_trace_exits: false,
     pause: false,
@@ -141,6 +145,13 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             OPTIONS.pause = true;
         },
 
+        ("temp-regs", _) => match opt_val.parse() {
+            Ok(n) => unsafe { OPTIONS.num_temp_regs = n },
+            Err(_) => {
+                return None;
+            }
+        },
+
         ("dump-disasm", _) => match opt_val.to_string().as_str() {
             "" => unsafe { OPTIONS.dump_disasm = Some(DumpDisasm::Stdout) },
             directory => {

yjit/src/stats.rs

@@ -352,6 +352,11 @@ make_counters! {
 
     iseq_stack_too_large,
     iseq_too_long,
+
+    temp_reg_opnd,
+    temp_mem_opnd,
+    temp_spill,
+    temp_reload,
 }
 
 //===========================================================================

yjit/yjit.mk

@@ -92,8 +92,7 @@ yjit-smoke-test:
 ifneq ($(strip $(CARGO)),)
 	$(CARGO) test --all-features -q --manifest-path='$(top_srcdir)/yjit/Cargo.toml'
 endif
-	$(MAKE) btest RUN_OPTS='--yjit-call-threshold=1' BTESTS=-j
+	$(MAKE) btest RUN_OPTS='--yjit-call-threshold=1 --yjit-temp-regs=5' BTESTS=-j
-	$(MAKE) test-all TESTS='$(top_srcdir)/test/ruby/test_yjit.rb'
 
 # Generate Rust bindings. See source for details.
 # Needs `./configure --enable-yjit=dev` and Clang.