[Bug #22224] Ractors, YJIT: Skip re-reporting recorded EP escapes to the JITs (#18176)

* [22224] Ractors, YJIT: Skip re-reporting recorded EP escapes to the JITs

vm_make_env_each reports every EP escape, and both
rb_yjit_invalidate_ep_is_bp and rb_zjit_invalidate_no_ep_escape take
the VM lock with a stop-all-Ractors barrier on every call. For a given
iseq only the first report does anything: afterwards YJIT's
no_ep_escape_iseqs entry is permanently an empty set and ZJIT's
ep_escape_iseqs insert is a no-op, so the barrier is paid for nothing.
Under multiple Ractors this serializes the whole process on every Proc
created from an already-escaped iseq: a workload creating one lambda
per iteration runs 150x slower on 8 Ractors than on 1 with YJIT
enabled.

Record the report in a jit_ep_escape_recorded flag on the iseq body,
set by the active JIT while it holds the VM lock, and skip the calls
in vm_make_env_each once it is set. The flag sits in existing padding
and the body is ZALLOC'd, so it starts false and dies with the iseq.
Setting it in the JIT rather than in the caller keeps mid-process
RubyVM::YJIT.enable correct: escapes reported while no JIT is
initialized are recorded by nobody and must not set the flag. The
unlocked read is benign: a stale false causes one redundant,
idempotent report.

On the micro-benchmark from the bug, --yjit at 8 Ractors goes from
0.04M to 11.07M iterations/s (277x).

* [22224] Deduplicate the EP escape record onto the iseq body

The jit_ep_escape_recorded flag introduced in the previous commit is
now the JITs' only record of an EP escape: YJIT's no_ep_escape_iseqs
keeps only iseqs with live assumptions (the empty-set escape markers
are gone), and ZJIT's ep_escape_iseqs set is removed together with its
update_ep_escape_iseqs compaction pass, since the flag lives in the
iseq body, which the GC does not move. Both JITs read the flag back at
compile time through a new rb_jit_iseq_ep_escape_recorded_p helper.

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

jit.c

@@ -244,6 +244,20 @@ rb_jit_iseq_builtin_attrs(const rb_iseq_t *iseq)
     return iseq->body->builtin_attrs;
 }
 
+// Relaxed memory ordering, but called by the JIT with VM lock and barrier.
+void
+rb_jit_iseq_mark_ep_escape_recorded(const rb_iseq_t *iseq)
+{
+    rbimpl_atomic_store(&iseq->body->jit_ep_escape_recorded, 1, RBIMPL_ATOMIC_RELAXED);
+}
+
+// Whether an EP escape of this iseq has been reported to the enabled JIT.
+bool
+rb_jit_iseq_ep_escape_recorded_p(const rb_iseq_t *iseq)
+{
+    return rbimpl_atomic_load(&iseq->body->jit_ep_escape_recorded, RBIMPL_ATOMIC_RELAXED) != 0;
+}
+
 int
 rb_get_mct_argc(const rb_method_cfunc_t *mct)
 {

vm.c

@@ -1136,7 +1136,10 @@ 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).
     // This is done before creating the imemo_env because VM_STACK_ENV_WRITE
     // below leaves the on-stack ep in a state that is unsafe to GC.
-    if (VM_FRAME_RUBYFRAME_P(cfp)) {
+    // Once the enabled JIT has recorded this iseq's escape, the invalidations
+    // are no longer useful and can slow down Ractors.
+    if (VM_FRAME_RUBYFRAME_P(cfp) &&
+        !rbimpl_atomic_load(&ISEQ_BODY(iseq)->jit_ep_escape_recorded, RBIMPL_ATOMIC_RELAXED)) {
         rb_yjit_invalidate_ep_is_bp(iseq);
         rb_zjit_invalidate_no_ep_escape(iseq);
     }

vm_core.h

@@ -543,6 +543,9 @@ struct rb_iseq_constant_body {
 
     bool prism; // ISEQ was generated from prism compiler
 
+    // Set once an EP escape of this iseq has been reported to the enabled JIT.
+    rb_atomic_t jit_ep_escape_recorded;
+
     union {
         iseq_bits_t * list; /* Find references for GC */
         iseq_bits_t single;

yjit/src/cruby.rs

@@ -118,6 +118,9 @@ extern "C" {
         ci: *const rb_callinfo,
     ) -> *const rb_callable_method_entry_t;
 
+    pub fn rb_jit_iseq_mark_ep_escape_recorded(iseq: IseqPtr);
+    pub fn rb_jit_iseq_ep_escape_recorded_p(iseq: IseqPtr) -> bool;
+
     // Floats within range will be encoded without creating objects in the heap.
     // (Range is 0x3000000000000001 to 0x4fffffffffffffff (1.7272337110188893E-77 to 2.3158417847463237E+77).
     pub fn rb_float_new(d: f64) -> VALUE;

yjit/src/invariants.rs

@@ -170,10 +170,7 @@ pub fn track_no_ep_escape_assumption(uninit_block: BlockRef, iseq: IseqPtr) {
 
 /// 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)
-        .map_or(false, |blocks| blocks.is_empty())
+    unsafe { rb_jit_iseq_ep_escape_recorded_p(iseq) }
 }
 
 /// Forget an ISEQ remembered in invariants
@@ -583,21 +580,16 @@ pub extern "C" fn rb_yjit_invalidate_ep_is_bp(iseq: IseqPtr) {
 
     with_vm_lock(src_loc!(), || {
         // If an EP escape for this ISEQ is detected for the first time, invalidate all blocks
-        // associated to the ISEQ.
-        let no_ep_escape_iseqs = &mut Invariants::get_instance().no_ep_escape_iseqs;
-        match no_ep_escape_iseqs.get_mut(&iseq) {
-            Some(blocks) => {
-                // Invalidate existing blocks and make jit.ep_is_bp() return false
-                for block in mem::take(blocks) {
-                    invalidate_block_version(&block);
-                    incr_counter!(invalidate_ep_escape);
-                }
-            }
-            None => {
-                // Let jit.ep_is_bp() return false for this ISEQ
-                no_ep_escape_iseqs.insert(iseq, HashSet::new());
+        // associated to the ISEQ. The iseq flag records the escape, so the map keeps only
+        // iseqs with live assumptions.
+        if let Some(blocks) = Invariants::get_instance().no_ep_escape_iseqs.remove(&iseq) {
+            for block in blocks {
+                invalidate_block_version(&block);
+                incr_counter!(invalidate_ep_escape);
             }
         }
+
+        unsafe { rb_jit_iseq_mark_ep_escape_recorded(iseq) };
     });
 }
 

zjit/src/cruby.rs

@@ -122,6 +122,9 @@ unsafe extern "C" {
         ci: *const rb_callinfo,
     ) -> *const rb_callable_method_entry_t;
 
+    pub fn rb_jit_iseq_mark_ep_escape_recorded(iseq: IseqPtr);
+    pub fn rb_jit_iseq_ep_escape_recorded_p(iseq: IseqPtr) -> bool;
+
     // Floats within range will be encoded without creating objects in the heap.
     // (Range is 0x3000000000000001 to 0x4fffffffffffffff (1.7272337110188893E-77 to 2.3158417847463237E+77).
     pub fn rb_float_new(d: f64) -> VALUE;

zjit/src/invariants.rs

@@ -5,7 +5,7 @@ use std::{collections::{HashMap, HashSet}, mem};
 use crate::{backend::lir::{Assembler, asm_comment}, cruby::{ID, IseqPtr, RedefinitionFlag, VALUE, iseq_name, rb_callable_method_entry_t, rb_gc_location, ruby_basic_operators, src_loc, with_vm_lock}, hir::Invariant, options::debug, state::{ZJITState, zjit_enabled_p, trace_invalidation}, virtualmem::CodePtr};
 use crate::payload::{IseqVersionRef, get_or_create_iseq_payload};
 use crate::codegen::invalidate_iseq_version;
-use crate::cruby::rb_iseq_reset_jit_func;
+use crate::cruby::{rb_iseq_reset_jit_func, rb_jit_iseq_ep_escape_recorded_p, rb_jit_iseq_mark_ep_escape_recorded};
 use crate::stats::with_time_stat;
 use crate::stats::Counter::invalidation_time_ns;
 use crate::gc::remove_gc_offsets;
@@ -70,9 +70,6 @@ impl PatchPoint {
 /// 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>,
-
     /// Map from ISEQ that's assumed to not escape EP to a set of patch points
     no_ep_escape_iseq_patch_points: HashMap<IseqPtr, HashSet<PatchPoint>>,
 
@@ -108,7 +105,6 @@ pub struct Invariants {
 impl Invariants {
     /// Update object references in Invariants
     pub fn update_references(&mut self) {
-        self.update_ep_escape_iseqs();
         self.update_no_ep_escape_iseq_patch_points();
         self.update_cme_patch_points();
         self.update_no_singleton_class_patch_points();
@@ -121,7 +117,6 @@ impl Invariants {
         // Why not patch the patch points? If the ISEQ is dead then the GC also proved that all
         // generated code referencing the ISEQ are unreachable. We mark the ISEQs baked into
         // generated code.
-        self.ep_escape_iseqs.remove(&iseq);
         self.no_ep_escape_iseq_patch_points.remove(&iseq);
     }
 
@@ -135,15 +130,6 @@ impl Invariants {
         self.no_singleton_class_patch_points.remove(&klass);
     }
 
-    /// Update ISEQ references in Invariants::ep_escape_iseqs
-    fn update_ep_escape_iseqs(&mut self) {
-        let updated = std::mem::take(&mut self.ep_escape_iseqs)
-            .into_iter()
-            .map(|iseq| unsafe { rb_gc_location(iseq.into()) }.as_iseq())
-            .collect();
-        self.ep_escape_iseqs = updated;
-    }
-
     /// Update ISEQ references in Invariants::no_ep_escape_iseq_patch_points
     fn update_no_ep_escape_iseq_patch_points(&mut self) {
         let updated = std::mem::take(&mut self.no_ep_escape_iseq_patch_points)
@@ -215,8 +201,8 @@ pub extern "C" fn rb_zjit_invalidate_no_ep_escape(iseq: IseqPtr) {
 
     with_vm_lock(src_loc!(), || {
         // Remember that this ISEQ may escape EP
+        unsafe { rb_jit_iseq_mark_ep_escape_recorded(iseq) };
         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 let Some(patch_points) = invariants.no_ep_escape_iseq_patch_points.remove(&iseq) {
@@ -291,7 +277,7 @@ pub fn track_no_ep_escape_assumption(
 
 /// Returns true if a given ISEQ has previously escaped environment pointer.
 pub fn iseq_seen_ep_escape(iseq: IseqPtr) -> bool {
-    ZJITState::get_invariants().ep_escape_iseqs.contains(&iseq)
+    unsafe { rb_jit_iseq_ep_escape_recorded_p(iseq) }
 }
 
 /// Track a patch point for a basic operator in a given class.