YJIT: Fix not reading locals from cfp->ep after YJIT.enable and exceptional entry

[Backport #21941]

In case of `--yjit-disable`, YJIT only starts to record environment
escapes after `RubyVM::YJIT.enable`. Previously we falsely assumed that
we always have a full history all the way back to VM boot. This had YJIT
install and run code that assume EP=BP when EP≠BP for some exceptional
entry into the middle of a running frame, if the environment escaped
before `YJIT.enable`.

The fix is to reject exceptional entry with an escaped environment.
Rename things and explain in more detail how the predicate for deciding
to assume EP=BP works. It's quite subtle since it reasons about all
parties in the system that push a control frame and then run JIT code.

Note that while can_assume_on_stack_env() checks the currently running
environment if it so happens to be the one YJIT is compiling against, it
can return true for any ISEQ. The check isn't necessary for fixing the
bug, and the load bearing part of this patch is the change to
exceptional entries.

This fix is flat on speed and space on ruby-bench headline benchmarks.

Many thanks for the community effort to create a small test case for
this bug.

Author: XrXr

test/ruby/test_yjit.rb

@@ -1784,6 +1784,27 @@ def bar(w:, x:, y:, z:, **kwrest) = yield kwrest
     RUBY
   end
 
+  def test_exceptional_entry_into_env_escaped_before_yjit_enablement
+    threshold = 2
+    assert_separately(["--disable-all", "--yjit-disable", "--yjit-call-threshold=#{threshold}"], <<~RUBY)
+      def run
+        @captured_env = ->{}
+        RubyVM::YJIT.enable
+
+        i = 0
+        while i < #{threshold}
+          next_i = i + 1
+          from_break = tap { break i + 1 } # break from the block generates an exceptional entry
+          assert_equal(from_break, next_i, '[Bug #21941]')
+          i = next_i
+        end
+      end
+
+      run
+      assert_equal(#{threshold}, @captured_env.binding.local_variable_get(:i))
+    RUBY
+  end
+
   private
 
   def code_gc_helpers

yjit.rb

@@ -353,6 +353,9 @@ def _print_stats(out: $stderr) # :nodoc:
       # Number of failed compiler invocations
       compilation_failure = stats[:compilation_failure]
 
+      # Number of refused exceptional entries with an escaped environment
+      exceptional_entry_escaped_env = stats[:exceptional_entry_escaped_env]
+
       code_region_overhead = stats[:code_region_size] - (stats[:inline_code_size] + stats[:outlined_code_size])
 
       out.puts "num_send:              " + format_number(13, stats[:num_send])
@@ -389,6 +392,7 @@ def _print_stats(out: $stderr) # :nodoc:
       out.puts "bindings_allocations:  " + format_number(13, stats[:binding_allocations])
       out.puts "bindings_set:          " + format_number(13, stats[:binding_set])
       out.puts "compilation_failure:   " + format_number(13, compilation_failure) if compilation_failure != 0
+      out.puts "exceptional_entry_escaped_env:" + format_number(6, exceptional_entry_escaped_env) if exceptional_entry_escaped_env != 0
       out.puts "live_iseq_count:       " + format_number(13, stats[:live_iseq_count])
       out.puts "iseq_alloc_count:      " + format_number(13, stats[:iseq_alloc_count])
       out.puts "compiled_iseq_entry:   " + format_number(13, stats[:compiled_iseq_entry])

yjit/src/codegen.rs

@@ -234,20 +234,36 @@ impl<'a> JITState<'a> {
         result
     }
 
-    /// Return true if the current ISEQ could escape an environment.
+    /// Return true if the JIT code can use [`Self::assume_no_ep_escape`]
+    /// and will run with an on-stack (`!VM_ENV_ESCAPED_P`) environment.
     ///
-    /// As of vm_push_frame(), EP is always equal to BP. However, after pushing
-    /// a frame, some ISEQ setups call vm_bind_update_env(), which redirects EP.
-    /// Also, some method calls escape the environment to the heap.
-    fn escapes_ep(&self) -> bool {
+    /// ## Reasoning about ISEQs that are not currently running
+    ///
+    /// As of vm_push_frame() and its JIT code equivalent, EP is always equal to BP (the
+    /// environment is on-stack and has not escaped). We can usually assume this is the starting
+    /// condition upon entry into JIT code. However, after pushing a frame and before entry into
+    /// JIT code, some ISEQ setups call vm_bind_update_env(), which redirects EP.
+    ///
+    /// ## After making the assumption
+    ///
+    /// After JIT code entry, many ruby operations can have the environment escape to heap. These
+    /// are handled by [`crate::invariants`].
+    ///
+    /// Exceptional entry through jit_exec_exception() is an extreme case of the environment state
+    /// changing between vm_push_frame() and entry into JIT code. The frame could have been pushed
+    /// before YJIT is enabled. The exception entry point refuses entry with an escaped environment.
+    fn can_assume_on_stack_env(&self) -> bool {
         match unsafe { get_iseq_body_type(self.iseq) } {
             // <main> frame is always associated to TOPLEVEL_BINDING.
             ISEQ_TYPE_MAIN |
             // Kernel#eval uses a heap EP when a Binding argument is not nil.
-            ISEQ_TYPE_EVAL => true,
-            // If this ISEQ has previously escaped EP, give up the optimization.
-            _ if iseq_escapes_ep(self.iseq) => true,
-            _ => false,
+            ISEQ_TYPE_EVAL => false,
+            // Check the running environment if compiling for it
+            _ if unsafe { self.iseq == get_cfp_iseq(self.get_cfp()) && cfp_env_has_escaped(self.get_cfp()) } => false,
+            // If we've seen this ISEQ run with an escaped environment, give up the optimization
+            // to avoid excessive invalidations (even though it may be fine for soundness).
+            _ if seen_escaped_env(self.iseq) => false,
+            _ => true,
         }
     }
 
@@ -376,8 +392,8 @@ impl<'a> JITState<'a> {
         if jit_ensure_block_entry_exit(self, asm).is_none() {
             return false; // out of space, give up
         }
-        if self.escapes_ep() {
-            return false; // EP has been escaped in this ISEQ. disable the optimization to avoid an invalidation loop.
+        if !self.can_assume_on_stack_env() {
+            return false; // Unsound or unprofitable to make the assumption
         }
         self.no_ep_escape = true;
         true
@@ -2448,7 +2464,7 @@ fn gen_getlocal_generic(
     level: u32,
 ) -> Option<CodegenStatus> {
     // Split the block if we need to invalidate this instruction when EP escapes
-    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+    if level == 0 && jit.can_assume_on_stack_env() && !jit.at_compile_target() {
         return jit.defer_compilation(asm);
     }
 
@@ -2549,7 +2565,7 @@ fn gen_setlocal_generic(
     }
 
     // Split the block if we need to invalidate this instruction when EP escapes
-    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+    if level == 0 && jit.can_assume_on_stack_env() && !jit.at_compile_target() {
         return jit.defer_compilation(asm);
     }
 

yjit/src/cruby.rs

@@ -598,6 +598,13 @@ impl From<VALUE> for u16 {
     }
 }
 
+/// Check whether a control frame has an escaped environment
+pub unsafe fn cfp_env_has_escaped(cfp: *mut rb_control_frame_struct) -> bool {
+    use crate::utils::IntoUsize;
+    let ep = get_cfp_ep(cfp);
+    0 != ep.offset(VM_ENV_DATA_INDEX_FLAGS as isize).read().0 & VM_ENV_FLAG_ESCAPED.as_usize()
+}
+
 /// Produce a Ruby string from a Rust string slice
 pub fn rust_str_to_ruby(str: &str) -> VALUE {
     unsafe { rb_utf8_str_new(str.as_ptr() as *const _, str.len() as i64) }

yjit/src/invariants.rs

@@ -168,8 +168,8 @@ pub fn track_no_ep_escape_assumption(uninit_block: BlockRef, iseq: IseqPtr) {
         .insert(uninit_block);
 }
 
-/// Returns true if a given ISEQ has previously escaped an environment.
-pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
+/// Returns true if a given ISEQ has escaped an environment since YJIT boot.
+pub fn seen_escaped_env(iseq: IseqPtr) -> bool {
     Invariants::get_instance()
         .no_ep_escape_iseqs
         .get(&iseq)

yjit/src/stats.rs

@@ -248,6 +248,7 @@ pub const DEFAULT_COUNTERS: &'static [Counter] = &[
     Counter::compiled_blockid_count,
     Counter::compiled_block_count,
     Counter::deleted_defer_block_count,
+    Counter::exceptional_entry_escaped_env,
     Counter::compiled_branch_count,
     Counter::compile_time_ns,
     Counter::compilation_failure,
@@ -599,6 +600,8 @@ make_counters! {
     iseq_stack_too_large,
     iseq_too_long,
 
+    exceptional_entry_escaped_env,
+
     temp_reg_opnd,
     temp_mem_opnd,
     temp_spill,

yjit/src/yjit.rs

@@ -157,6 +157,13 @@ pub extern "C" fn rb_yjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr, jit_exc
         return std::ptr::null();
     }
 
+    // In case of exceptional entry, reject escaped environment.
+    // This allows us to use the fact that new frames generally start with an on-stack environment.
+    if jit_exception && unsafe { cfp_env_has_escaped(get_ec_cfp(ec)) } {
+        incr_counter!(exceptional_entry_escaped_env);
+        return std::ptr::null();
+    }
+
     // If a custom call threshold was not specified on the command-line and
     // this is a large application (has very many ISEQs), switch to
     // using the call threshold for large applications after this entry point