Implement Insn::Param using the SP register (https://github.com/Shopify/zjit/pull/39)

Merged: https://github.com/ruby/ruby/pull/13131

vm.c

@@ -1048,6 +1048,7 @@ vm_make_env_each(const rb_execution_context_t * const ec, rb_control_frame_t *co
     // Invalidate JIT code that assumes cfp->ep == vm_base_ptr(cfp).
     if (env->iseq) {
         rb_yjit_invalidate_ep_is_bp(env->iseq);
+        rb_zjit_invalidate_ep_is_bp(env->iseq);
     }
 
     return (VALUE)env;

zjit.h

@@ -11,11 +11,13 @@ void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, boo
 void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec);
 void rb_zjit_profile_iseq(const rb_iseq_t *iseq);
 void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
+void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq);
 #else
-void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
-void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec) {}
-void rb_zjit_profile_iseq(const rb_iseq_t *iseq) {}
-void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
+static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
+static inline void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec) {}
+static inline void rb_zjit_profile_iseq(const rb_iseq_t *iseq) {}
+static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
+static inline void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq) {}
 #endif // #if USE_YJIT
 
 #endif // #ifndef ZJIT_H

zjit/src/codegen.rs

@@ -1,6 +1,8 @@
-use crate::{
-    asm::CodeBlock, backend::lir, backend::lir::{asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, EC, SP}, cruby::*, debug, hir::{Const, FrameState, Function, Insn, InsnId}, hir_type::{types::Fixnum, Type}, virtualmem::CodePtr
-};
+use crate::{asm::CodeBlock, cruby::*, debug, virtualmem::CodePtr};
+use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
+use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, EC, SP};
+use crate::hir::{Const, FrameState, Function, Insn, InsnId};
+use crate::hir_type::{types::Fixnum, Type};
 
 /// Ephemeral code generation state
 struct JITState {
@@ -28,6 +30,15 @@ impl JITState {
         }
         opnd
     }
+
+    /// Assume that this ISEQ doesn't escape EP. Return false if it's known to escape EP.
+    fn assume_no_ep_escape(&mut self) -> bool {
+        if iseq_escapes_ep(self.iseq) {
+            return false;
+        }
+        track_no_ep_escape_assumption(self.iseq);
+        true
+    }
 }
 
 /// Compile High-level IR into machine code
@@ -104,17 +115,23 @@ fn gen_const(val: VALUE) -> Opnd {
 }
 
 /// Compile a method/block paramter read. For now, it only supports method parameters.
-fn gen_param(jit: &JITState, asm: &mut Assembler, local_idx: usize) -> Option<lir::Opnd> {
-    // Get the EP of the current CFP
-    // TODO: Use the SP register and invalidate on EP escape
-    let ep_opnd = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_EP);
-    let ep_reg = asm.load(ep_opnd);
-
-    // Load the local variable
-    // val = *(vm_get_ep(GET_EP(), level) - idx);
+fn gen_param(jit: &mut JITState, asm: &mut Assembler, local_idx: usize) -> Option<lir::Opnd> {
     let ep_offset = local_idx_to_ep_offset(jit.iseq, local_idx);
-    let offs = -(SIZEOF_VALUE_I32 * ep_offset);
-    let local_opnd = Opnd::mem(64, ep_reg, offs);
+
+    let local_opnd = if jit.assume_no_ep_escape() {
+        // Create a reference to the local variable using the SP register. We assume EP == BP.
+        // TODO: Implement the invalidation in rb_zjit_invalidate_ep_is_bp()
+        let offs = -(SIZEOF_VALUE_I32 * (ep_offset + 1));
+        Opnd::mem(64, SP, offs)
+    } else {
+        // Get the EP of the current CFP
+        let ep_opnd = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_EP);
+        let ep_reg = asm.load(ep_opnd);
+
+        // Create a reference to the local variable using cfp->ep
+        let offs = -(SIZEOF_VALUE_I32 * ep_offset);
+        Opnd::mem(64, ep_reg, offs)
+    };
 
     Some(local_opnd)
 }

zjit/src/invariants.rs

@@ -1,4 +1,17 @@
-use crate::{cruby::{ruby_basic_operators, RedefinitionFlag}, zjit_enabled_p};
+use std::collections::HashSet;
+
+use crate::{cruby::{ruby_basic_operators, IseqPtr, RedefinitionFlag}, state::ZJITState, zjit_enabled_p};
+
+/// Used to track all of the various block references that contain assumptions
+/// about the state of the virtual machine.
+#[derive(Default)]
+pub struct Invariants {
+    /// Set of ISEQs that are known to escape EP
+    ep_escape_iseqs: HashSet<IseqPtr>,
+
+    /// Set of ISEQs whose JIT code assumes that it doesn't escape EP
+    no_ep_escape_iseqs: HashSet<IseqPtr>,
+}
 
 /// Called when a basic operator is redefined. Note that all the blocks assuming
 /// the stability of different operators are invalidated together and we don't
@@ -12,3 +25,33 @@ pub extern "C" fn rb_zjit_bop_redefined(_klass: RedefinitionFlag, _bop: ruby_bas
 
     unimplemented!("Invalidation on BOP redefinition is not implemented yet");
 }
+
+/// Invalidate blocks for a given ISEQ that assumes environment pointer is
+/// equal to base pointer.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_invalidate_ep_is_bp(iseq: IseqPtr) {
+    // Skip tracking EP escapes on boot. We don't need to invalidate anything during boot.
+    if !ZJITState::has_instance() {
+        return;
+    }
+
+    // Remember that this ISEQ may escape EP
+    let invariants = ZJITState::get_invariants();
+    invariants.ep_escape_iseqs.insert(iseq);
+
+    // If the ISEQ has been compiled assuming it doesn't escape EP, invalidate the JIT code.
+    if invariants.no_ep_escape_iseqs.contains(&iseq) {
+        unimplemented!("Invalidation on EP escape is not implemented yet");
+    }
+}
+
+/// Track that JIT code for a ISEQ will assume that base pointer is equal to environment pointer.
+pub fn track_no_ep_escape_assumption(iseq: IseqPtr) {
+    let invariants = ZJITState::get_invariants();
+    invariants.no_ep_escape_iseqs.insert(iseq);
+}
+
+/// Returns true if a given ISEQ has previously escaped environment pointer.
+pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
+    ZJITState::get_invariants().ep_escape_iseqs.contains(&iseq)
+}

zjit/src/profile.rs

@@ -92,6 +92,7 @@ impl IseqPayload {
         self.opnd_types.get(&insn_idx).map(|types| types.as_slice())
     }
 
+    /// Return true if top-two stack operands are Fixnums
     pub fn have_two_fixnums(&self, insn_idx: usize) -> bool {
         match self.get_operand_types(insn_idx) {
             Some([left, right]) => left.is_subtype(Fixnum) && right.is_subtype(Fixnum),

zjit/src/state.rs

@@ -1,3 +1,4 @@
+use crate::invariants::Invariants;
 use crate::options::Options;
 use crate::asm::CodeBlock;
 
@@ -8,6 +9,9 @@ pub struct ZJITState {
 
     /// ZJIT command-line options
     options: Options,
+
+    /// Assumptions that require invalidation
+    invariants: Invariants,
 }
 
 /// Private singleton instance of the codegen globals
@@ -59,10 +63,16 @@ impl ZJITState {
         let zjit_state = ZJITState {
             code_block: cb,
             options,
+            invariants: Invariants::default(),
         };
         unsafe { ZJIT_STATE = Some(zjit_state); }
     }
 
+    /// Return true if zjit_state has been initialized
+    pub fn has_instance() -> bool {
+        unsafe { ZJIT_STATE.as_mut().is_some() }
+    }
+
     /// Get a mutable reference to the codegen globals instance
     fn get_instance() -> &'static mut ZJITState {
         unsafe { ZJIT_STATE.as_mut().unwrap() }
@@ -73,8 +83,13 @@ impl ZJITState {
         &mut ZJITState::get_instance().code_block
     }
 
-    // Get a mutable reference to the options
+    /// Get a mutable reference to the options
     pub fn get_options() -> &'static mut Options {
         &mut ZJITState::get_instance().options
     }
+
+    /// Get a mutable reference to the invariants
+    pub fn get_invariants() -> &'static mut Invariants {
+        &mut ZJITState::get_instance().invariants
+    }
 }