ZJIT: Mark profiled objects when marking ISEQ (#13784)

common.mk

@@ -9376,6 +9376,7 @@ iseq.$(OBJEXT): {$(VPATH)}vm_debug.h
 iseq.$(OBJEXT): {$(VPATH)}vm_opts.h
 iseq.$(OBJEXT): {$(VPATH)}vm_sync.h
 iseq.$(OBJEXT): {$(VPATH)}yjit.h
+iseq.$(OBJEXT): {$(VPATH)}zjit.h
 jit.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 jit.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 jit.$(OBJEXT): $(CCAN_DIR)/list/list.h

iseq.c

@@ -44,6 +44,7 @@
 #include "builtin.h"
 #include "insns.inc"
 #include "insns_info.inc"
+#include "zjit.h"
 
 VALUE rb_cISeq;
 static VALUE iseqw_new(const rb_iseq_t *iseq);
@@ -401,11 +402,17 @@ rb_iseq_mark_and_move(rb_iseq_t *iseq, bool reference_updating)
         if (reference_updating) {
 #if USE_YJIT
             rb_yjit_iseq_update_references(iseq);
+#endif
+#if USE_ZJIT
+            rb_zjit_iseq_update_references(body->zjit_payload);
 #endif
         }
         else {
 #if USE_YJIT
             rb_yjit_iseq_mark(body->yjit_payload);
+#endif
+#if USE_ZJIT
+            rb_zjit_iseq_mark(body->zjit_payload);
 #endif
         }
     }

jit.c

@@ -415,6 +415,14 @@ rb_assert_iseq_handle(VALUE handle)
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_iseq));
 }
 
+// Assert that we have the VM lock. Relevant mostly for multi ractor situations.
+// The GC takes the lock before calling us, and this asserts that it indeed happens.
+void
+rb_assert_holding_vm_lock(void)
+{
+    ASSERT_vm_locking();
+}
+
 int
 rb_IMEMO_TYPE_P(VALUE imemo, enum imemo_type imemo_type)
 {

test/ruby/test_zjit.rb

@@ -951,6 +951,12 @@ class TestZJIT < Test::Unit::TestCase
     end
   end
 
+  def test_require_rubygems
+    assert_runs 'true', %q{
+      require 'rubygems'
+    }, call_threshold: 2
+  end
+
   def test_module_name_with_guard_passes
     assert_compiles '"Integer"', %q{
       def test(mod)

yjit.c

@@ -792,14 +792,6 @@ rb_yjit_shape_index(shape_id_t shape_id)
     return RSHAPE_INDEX(shape_id);
 }
 
-// Assert that we have the VM lock. Relevant mostly for multi ractor situations.
-// The GC takes the lock before calling us, and this asserts that it indeed happens.
-void
-rb_yjit_assert_holding_vm_lock(void)
-{
-    ASSERT_vm_locking();
-}
-
 // The number of stack slots that vm_sendish() pops for send and invokesuper.
 size_t
 rb_yjit_sendish_sp_pops(const struct rb_callinfo *ci)

yjit/bindgen/src/main.rs

@@ -341,7 +341,6 @@ fn main() {
         .allowlist_function("rb_yjit_exit_locations_dict")
         .allowlist_function("rb_yjit_icache_invalidate")
         .allowlist_function("rb_optimized_call")
-        .allowlist_function("rb_yjit_assert_holding_vm_lock")
         .allowlist_function("rb_yjit_sendish_sp_pops")
         .allowlist_function("rb_yjit_invokeblock_sp_pops")
         .allowlist_function("rb_yjit_set_exception_return")
@@ -349,6 +348,9 @@ fn main() {
         .allowlist_type("robject_offsets")
         .allowlist_type("rstring_offsets")
 
+        // From jit.c
+        .allowlist_function("rb_assert_holding_vm_lock")
+
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")
 

yjit/src/core.rs

@@ -1920,7 +1920,7 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
         // For aliasing, having the VM lock hopefully also implies that no one
         // else has an overlapping &mut IseqPayload.
         unsafe {
-            rb_yjit_assert_holding_vm_lock();
+            rb_assert_holding_vm_lock();
             &*(payload as *const IseqPayload)
         }
     };
@@ -2009,7 +2009,7 @@ pub extern "C" fn rb_yjit_iseq_update_references(iseq: IseqPtr) {
         // For aliasing, having the VM lock hopefully also implies that no one
         // else has an overlapping &mut IseqPayload.
         unsafe {
-            rb_yjit_assert_holding_vm_lock();
+            rb_assert_holding_vm_lock();
             &*(payload as *const IseqPayload)
         }
     };

yjit/src/cruby_bindings.inc.rs

@@ -1244,7 +1244,6 @@ extern "C" {
     pub fn rb_yjit_shape_obj_too_complex_p(obj: VALUE) -> bool;
     pub fn rb_yjit_shape_capacity(shape_id: shape_id_t) -> attr_index_t;
     pub fn rb_yjit_shape_index(shape_id: shape_id_t) -> attr_index_t;
-    pub fn rb_yjit_assert_holding_vm_lock();
     pub fn rb_yjit_sendish_sp_pops(ci: *const rb_callinfo) -> usize;
     pub fn rb_yjit_invokeblock_sp_pops(ci: *const rb_callinfo) -> usize;
     pub fn rb_yjit_set_exception_return(
@@ -1325,6 +1324,7 @@ extern "C" {
     pub fn rb_BASIC_OP_UNREDEFINED_P(bop: ruby_basic_operators, klass: u32) -> bool;
     pub fn rb_RCLASS_ORIGIN(c: VALUE) -> VALUE;
     pub fn rb_assert_iseq_handle(handle: VALUE);
+    pub fn rb_assert_holding_vm_lock();
     pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int;
     pub fn rb_assert_cme_handle(handle: VALUE);
     pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE;

zjit.h

@@ -13,6 +13,8 @@ void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec
 void rb_zjit_profile_enable(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);
+void rb_zjit_iseq_mark(void *payload);
+void rb_zjit_iseq_update_references(void *payload);
 #else
 #define rb_zjit_enabled_p false
 static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}

zjit/bindgen/src/main.rs

@@ -352,6 +352,9 @@ fn main() {
         .allowlist_type("robject_offsets")
         .allowlist_type("rstring_offsets")
 
+        // From jit.c
+        .allowlist_function("rb_assert_holding_vm_lock")
+
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")
 

zjit/src/codegen.rs

@@ -4,7 +4,7 @@ use std::num::NonZeroU32;
 
 use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::invariants::track_bop_assumption;
-use crate::profile::get_or_create_iseq_payload;
+use crate::gc::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, SP};

zjit/src/cruby_bindings.inc.rs

@@ -1006,6 +1006,7 @@ unsafe extern "C" {
     pub fn rb_BASIC_OP_UNREDEFINED_P(bop: ruby_basic_operators, klass: u32) -> bool;
     pub fn rb_RCLASS_ORIGIN(c: VALUE) -> VALUE;
     pub fn rb_assert_iseq_handle(handle: VALUE);
+    pub fn rb_assert_holding_vm_lock();
     pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int;
     pub fn rb_assert_cme_handle(handle: VALUE);
     pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE;

zjit/src/gc.rs

@@ -0,0 +1,75 @@
+// This module is responsible for marking/moving objects on GC.
+
+use std::ffi::c_void;
+use crate::{cruby::*, profile::IseqProfile, virtualmem::CodePtr};
+
+/// This is all the data ZJIT stores on an ISEQ. We mark objects in this struct on GC.
+#[derive(Default, Debug)]
+pub struct IseqPayload {
+    /// Type information of YARV instruction operands
+    pub profile: IseqProfile,
+
+    /// JIT code address of the first block
+    pub start_ptr: Option<CodePtr>,
+
+    // TODO: Add references to GC offsets in JIT code
+}
+
+/// Get the payload object associated with an iseq. Create one if none exists.
+pub fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
+    type VoidPtr = *mut c_void;
+
+    let payload_non_null = unsafe {
+        let payload = rb_iseq_get_zjit_payload(iseq);
+        if payload.is_null() {
+            // Allocate a new payload with Box and transfer ownership to the GC.
+            // We drop the payload with Box::from_raw when the GC frees the iseq and calls us.
+            // NOTE(alan): Sometimes we read from an iseq without ever writing to it.
+            // We allocate in those cases anyways.
+            let new_payload = IseqPayload::default();
+            let new_payload = Box::into_raw(Box::new(new_payload));
+            rb_iseq_set_zjit_payload(iseq, new_payload as VoidPtr);
+
+            new_payload
+        } else {
+            payload as *mut IseqPayload
+        }
+    };
+
+    // SAFETY: we should have the VM lock and all other Ruby threads should be asleep. So we have
+    // exclusive mutable access.
+    // Hmm, nothing seems to stop calling this on the same
+    // iseq twice, though, which violates aliasing rules.
+    unsafe { payload_non_null.as_mut() }.unwrap()
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_iseq_mark(payload: *mut c_void) {
+    let payload = if payload.is_null() {
+        return; // nothing to mark
+    } else {
+        // SAFETY: The GC takes the VM lock while marking, which
+        // we assert, so we should be synchronized and data race free.
+        //
+        // For aliasing, having the VM lock hopefully also implies that no one
+        // else has an overlapping &mut IseqPayload.
+        unsafe {
+            rb_assert_holding_vm_lock();
+            &*(payload as *const IseqPayload)
+        }
+    };
+
+    payload.profile.each_object(|object| {
+        // TODO: Implement `rb_zjit_iseq_update_references` and use `rb_gc_mark_movable`
+        unsafe { rb_gc_mark(object); }
+    });
+
+    // TODO: Mark objects in JIT code
+}
+
+/// GC callback for updating GC objects in the per-iseq payload.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_iseq_update_references(_payload: *mut c_void) {
+    // TODO: let `rb_zjit_iseq_mark` use `rb_gc_mark_movable`
+    // and update references using `rb_gc_location` here.
+}

zjit/src/hir.rs

@@ -4,7 +4,7 @@
 #![allow(non_upper_case_globals)]
 
 use crate::{
-    cast::IntoUsize, cruby::*, options::{get_option, DumpHIR}, profile::{get_or_create_iseq_payload, IseqPayload}, state::ZJITState
+    cast::IntoUsize, cruby::*, options::{get_option, DumpHIR}, gc::{get_or_create_iseq_payload, IseqPayload}, state::ZJITState
 };
 use std::{
     cell::RefCell,
@@ -2367,7 +2367,7 @@ impl ProfileOracle {
     /// Map the interpreter-recorded types of the stack onto the HIR operands on our compile-time virtual stack
     fn profile_stack(&mut self, state: &FrameState) {
         let iseq_insn_idx = state.insn_idx;
-        let Some(operand_types) = self.payload.get_operand_types(iseq_insn_idx) else { return };
+        let Some(operand_types) = self.payload.profile.get_operand_types(iseq_insn_idx) else { return };
         let entry = self.types.entry(iseq_insn_idx).or_insert_with(|| vec![]);
         // operand_types is always going to be <= stack size (otherwise it would have an underflow
         // at run-time) so use that to drive iteration.

zjit/src/lib.rs

@@ -24,3 +24,4 @@ mod invariants;
 #[cfg(test)]
 mod assertions;
 mod bitset;
+mod gc;

zjit/src/profile.rs

@@ -1,10 +1,9 @@
 // We use the YARV bytecode constants which have a CRuby-style name
 #![allow(non_upper_case_globals)]
 
-use core::ffi::c_void;
 use std::collections::HashMap;
 
-use crate::{cruby::*, hir_type::{types::{Empty, Fixnum}, Type}, virtualmem::CodePtr};
+use crate::{cruby::*, gc::get_or_create_iseq_payload, hir_type::{types::{Empty, Fixnum}, Type}};
 
 /// Ephemeral state for profiling runtime information
 struct Profiler {
@@ -77,8 +76,8 @@ fn profile_insn(profiler: &mut Profiler, opcode: ruby_vminsn_type) {
 
 /// Profile the Type of top-`n` stack operands
 fn profile_operands(profiler: &mut Profiler, n: usize) {
-    let payload = get_or_create_iseq_payload(profiler.iseq);
-    let mut types = if let Some(types) = payload.opnd_types.get(&profiler.insn_idx) {
+    let profile = &mut get_or_create_iseq_payload(profiler.iseq).profile;
+    let mut types = if let Some(types) = profile.opnd_types.get(&profiler.insn_idx) {
         types.clone()
     } else {
         vec![Empty; n]
@@ -89,21 +88,16 @@ fn profile_operands(profiler: &mut Profiler, n: usize) {
         types[i] = types[i].union(opnd_type);
     }
 
-    payload.opnd_types.insert(profiler.insn_idx, types);
+    profile.opnd_types.insert(profiler.insn_idx, types);
 }
 
-/// This is all the data ZJIT stores on an iseq. This will be dynamically allocated by C code
-/// C code should pass an &mut IseqPayload to us when calling into ZJIT.
 #[derive(Default, Debug)]
-pub struct IseqPayload {
+pub struct IseqProfile {
     /// Type information of YARV instruction operands, indexed by the instruction index
     opnd_types: HashMap<usize, Vec<Type>>,
-
-    /// JIT code address of the first block
-    pub start_ptr: Option<CodePtr>,
 }
 
-impl IseqPayload {
+impl IseqProfile {
     /// Get profiled operand types for a given instruction index
     pub fn get_operand_types(&self, insn_idx: usize) -> Option<&[Type]> {
         self.opnd_types.get(&insn_idx).map(|types| types.as_slice())
@@ -116,40 +110,15 @@ impl IseqPayload {
             _ => false,
         }
     }
-}
 
-/// Get the payload for an iseq. For safety it's up to the caller to ensure the returned `&mut`
-/// upholds aliasing rules and that the argument is a valid iseq.
-pub fn get_iseq_payload(iseq: IseqPtr) -> Option<&'static mut IseqPayload> {
-    let payload = unsafe { rb_iseq_get_zjit_payload(iseq) };
-    let payload: *mut IseqPayload = payload.cast();
-    unsafe { payload.as_mut() }
-}
-
-/// Get the payload object associated with an iseq. Create one if none exists.
-pub fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
-    type VoidPtr = *mut c_void;
-
-    let payload_non_null = unsafe {
-        let payload = rb_iseq_get_zjit_payload(iseq);
-        if payload.is_null() {
-            // Allocate a new payload with Box and transfer ownership to the GC.
-            // We drop the payload with Box::from_raw when the GC frees the iseq and calls us.
-            // NOTE(alan): Sometimes we read from an iseq without ever writing to it.
-            // We allocate in those cases anyways.
-            let new_payload = IseqPayload::default();
-            let new_payload = Box::into_raw(Box::new(new_payload));
-            rb_iseq_set_zjit_payload(iseq, new_payload as VoidPtr);
-
-            new_payload
-        } else {
-            payload as *mut IseqPayload
+    /// Run a given callback with every object in IseqProfile
+    pub fn each_object(&self, callback: impl Fn(VALUE)) {
+        for types in self.opnd_types.values() {
+            for opnd_type in types {
+                if let Some(object) = opnd_type.ruby_object() {
+                    callback(object);
+                }
+            }
         }
-    };
-
-    // SAFETY: we should have the VM lock and all other Ruby threads should be asleep. So we have
-    // exclusive mutable access.
-    // Hmm, nothing seems to stop calling this on the same
-    // iseq twice, though, which violates aliasing rules.
-    unsafe { payload_non_null.as_mut() }.unwrap()
+    }
 }