ZJIT: Use LoadField for Class/Module ivars (ruby#16252)

Assume only one box (root box) and invalidate otherwise. Drops C calls to `rb_ivar_get_at_no_ractor_check`.

Before:
```
Top-20 calls to C functions from JIT code (77.3% of total 64,311,573):
                             rb_vm_opt_send_without_block: 11,939,854 (18.6%)
                                             rb_hash_aref:  5,400,091 ( 8.4%)
                                        rb_vm_invokeblock:  4,453,357 ( 6.9%)
                     rb_zjit_writebarrier_check_immediate:  4,279,890 ( 6.7%)
                                rb_vm_getinstancevariable:  3,504,908 ( 5.4%)
                                               rb_vm_send:  3,103,424 ( 4.8%)
                           rb_ivar_get_at_no_ractor_check:  2,864,766 ( 4.5%)
                                        rb_obj_is_kind_of:  2,313,479 ( 3.6%)
                                             rb_hash_aset:  1,903,359 ( 3.0%)
                                               Hash#fetch:  1,639,937 ( 2.5%)
                                rb_vm_setinstancevariable:  1,596,791 ( 2.5%)
                               rb_vm_opt_getconstant_path:  1,328,761 ( 2.1%)
                                          rb_jit_ary_push:    960,563 ( 1.5%)
                                rb_ec_ary_new_from_values:    722,913 ( 1.1%)
                               rb_class_allocate_instance:    721,483 ( 1.1%)
                                                    fetch:    713,134 ( 1.1%)
                                        rb_str_buf_append:    667,545 ( 1.0%)
                                              rb_ivar_get:    585,817 ( 0.9%)
                                    rb_hash_new_with_size:    520,347 ( 0.8%)
                                        rb_vm_sendforward:    479,029 ( 0.7%)
```

After:
```
Top-20 calls to C functions from JIT code (76.5% of total 62,282,359):
                             rb_vm_opt_send_without_block: 11,939,850 (19.2%)
                                             rb_hash_aref:  5,400,092 ( 8.7%)
                                        rb_vm_invokeblock:  4,453,357 ( 7.2%)
                     rb_zjit_writebarrier_check_immediate:  4,279,893 ( 6.9%)
                                rb_vm_getinstancevariable:  3,504,920 ( 5.6%)
                                               rb_vm_send:  3,103,441 ( 5.0%)
                                        rb_obj_is_kind_of:  2,313,510 ( 3.7%)
                                             rb_hash_aset:  1,903,359 ( 3.1%)
                                               Hash#fetch:  1,639,937 ( 2.6%)
                                rb_vm_setinstancevariable:  1,596,797 ( 2.6%)
                               rb_vm_opt_getconstant_path:  1,328,761 ( 2.1%)
                                          rb_jit_ary_push:    960,563 ( 1.5%)
                           rb_ivar_get_at_no_ractor_check:    835,498 ( 1.3%)
                                rb_ec_ary_new_from_values:    722,921 ( 1.2%)
                               rb_class_allocate_instance:    721,492 ( 1.2%)
                                                    fetch:    713,135 ( 1.1%)
                                        rb_str_buf_append:    667,545 ( 1.1%)
                                              rb_ivar_get:    585,815 ( 0.9%)
                                    rb_hash_new_with_size:    520,348 ( 0.8%)
                                        rb_vm_sendforward:    479,029 ( 0.8%)
```

The remaining `rb_ivar_get_at_no_ractor_check` are due to TypedData access, mostly on `Thread`.

Author: tekknolagi

box.c

@@ -18,6 +18,7 @@
 #include "ruby/util.h"
 #include "vm_core.h"
 #include "darray.h"
+#include "zjit.h"
 
 #include <stdio.h>
 
@@ -390,6 +391,9 @@ box_initialize(VALUE box_value)
 
     rb_ivar_set(box_value, id_box_entry, entry);
 
+    // Invalidate ZJIT code that assumes only the root box is active
+    rb_zjit_invalidate_root_box();
+
     return box_value;
 }
 

depend

@@ -980,6 +980,7 @@ box.$(OBJEXT): {$(VPATH)}vm_core.h
 box.$(OBJEXT): {$(VPATH)}vm_debug.h
 box.$(OBJEXT): {$(VPATH)}vm_opts.h
 box.$(OBJEXT): {$(VPATH)}vm_sync.h
+box.$(OBJEXT): {$(VPATH)}zjit.h
 builtin.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 builtin.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 builtin.$(OBJEXT): $(CCAN_DIR)/list/list.h

jit.c

@@ -16,12 +16,17 @@
 #include "vm_sync.h"
 #include "internal/fixnum.h"
 #include "internal/string.h"
+#include "internal/class.h"
+#include "internal/imemo.h"
 
 enum jit_bindgen_constants {
     // Field offsets for the RObject struct
     ROBJECT_OFFSET_AS_HEAP_FIELDS = offsetof(struct RObject, as.heap.fields),
     ROBJECT_OFFSET_AS_ARY = offsetof(struct RObject, as.ary),
 
+    // Field offset for prime classext's fields_obj from a class pointer
+    RCLASS_OFFSET_PRIME_FIELDS_OBJ = offsetof(struct RClass_and_rb_classext_t, classext.fields_obj),
+
     // Field offsets for the RString struct
     RUBY_OFFSET_RSTRING_LEN = offsetof(struct RString, len),
 
@@ -529,6 +534,13 @@ rb_jit_multi_ractor_p(void)
     return rb_multi_ractor_p();
 }
 
+bool
+rb_jit_class_fields_embedded_p(VALUE klass)
+{
+    VALUE fields_obj = RCLASS_EXT_PRIME(klass)->fields_obj;
+    return !fields_obj || !FL_TEST_RAW(fields_obj, OBJ_FIELD_HEAP);
+}
+
 // Acquire the VM lock and then signal all other Ruby threads (ractors) to
 // contend for the VM lock, putting them to sleep. ZJIT and YJIT use this to
 // evict threads running inside generated code so among other things, it can

yjit/src/cruby_bindings.inc.rs

@@ -28,6 +28,7 @@ void rb_zjit_iseq_free(const rb_iseq_t *iseq);
 void rb_zjit_before_ractor_spawn(void);
 void rb_zjit_tracing_invalidate_all(void);
 void rb_zjit_invalidate_no_singleton_class(VALUE klass);
+void rb_zjit_invalidate_root_box(void);
 #else
 #define rb_zjit_entry 0
 static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception) {}
@@ -40,6 +41,7 @@ static inline void rb_zjit_constant_state_changed(ID id) {}
 static inline void rb_zjit_before_ractor_spawn(void) {}
 static inline void rb_zjit_tracing_invalidate_all(void) {}
 static inline void rb_zjit_invalidate_no_singleton_class(VALUE klass) {}
+static inline void rb_zjit_invalidate_root_box(void) {}
 #endif // #if USE_ZJIT
 
 #define rb_zjit_enabled_p (rb_zjit_entry != 0)

zjit.h

@@ -309,6 +309,7 @@ fn main() {
         .allowlist_function("rb_assert_holding_vm_lock")
         .allowlist_function("rb_jit_shape_too_complex_p")
         .allowlist_function("rb_jit_multi_ractor_p")
+        .allowlist_function("rb_jit_class_fields_embedded_p")
         .allowlist_function("rb_jit_vm_lock_then_barrier")
         .allowlist_function("rb_jit_vm_unlock")
         .allowlist_function("rb_jit_for_each_iseq")

zjit/bindgen/src/main.rs

@@ -10,7 +10,8 @@ use std::slice;
 use crate::backend::current::ALLOC_REGS;
 use crate::invariants::{
     track_bop_assumption, track_cme_assumption, track_no_ep_escape_assumption, track_no_trace_point_assumption,
-    track_single_ractor_assumption, track_stable_constant_names_assumption, track_no_singleton_class_assumption
+    track_single_ractor_assumption, track_stable_constant_names_assumption, track_no_singleton_class_assumption,
+    track_root_box_assumption
 };
 use crate::gc::append_gc_offsets;
 use crate::payload::{get_or_create_iseq_payload, IseqCodePtrs, IseqVersion, IseqVersionRef, IseqStatus};
@@ -874,6 +875,9 @@ pub fn split_patch_point(asm: &mut Assembler, target: &Target, invariant: Invari
             Invariant::NoSingletonClass { klass } => {
                 track_no_singleton_class_assumption(klass, code_ptr, side_exit_ptr, version);
             }
+            Invariant::RootBoxOnly => {
+                track_root_box_assumption(code_ptr, side_exit_ptr, version);
+            }
         }
     });
 }

zjit/src/codegen.rs

@@ -581,6 +581,10 @@ impl VALUE {
         }
     }
 
+    pub fn class_fields_embedded_p(self) -> bool {
+        unsafe { rb_jit_class_fields_embedded_p(self) }
+    }
+
     pub fn as_fixnum(self) -> i64 {
         assert!(self.fixnum_p());
         (self.0 as i64) >> 1
@@ -1491,6 +1495,7 @@ pub(crate) mod ids {
         name: aref               content: b"[]"
         name: len
         name: _as_heap
+        name: _fields_obj
         name: thread_ptr
         name: self_              content: b"self"
         name: rb_ivar_get_at_no_ractor_check

zjit/src/cruby.rs

@@ -7,7 +7,7 @@
 #![allow(clippy::match_like_matches_macro)]
 use crate::{
     backend::lir::C_ARG_OPNDS,
-    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, invariants::has_singleton_class_of, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json
+    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, invariants::{self, has_singleton_class_of}, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json
 };
 use std::{
     cell::RefCell, collections::{BTreeSet, HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, c_int, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
@@ -151,6 +151,9 @@ pub enum Invariant {
     NoSingletonClass {
         klass: VALUE,
     },
+    /// Only the root box is active, so we can safely read from the prime classext.
+    /// Invalidated if a non-root box duplicates any classext.
+    RootBoxOnly,
 }
 
 impl Invariant {
@@ -290,6 +293,7 @@ impl<'a> std::fmt::Display for InvariantPrinter<'a> {
                     class_name,
                     self.ptr_map.map_ptr(klass.as_ptr::<VALUE>()))
             }
+            Invariant::RootBoxOnly => write!(f, "RootBoxOnly"),
         }
     }
 }
@@ -2182,6 +2186,17 @@ impl Function {
         }
     }
 
+    /// Assume that only the root box is active, so we can safely read from the prime classext.
+    /// Returns true if safe to assume so and emits a PatchPoint.
+    pub fn assume_root_box(&mut self, block: BlockId, state: InsnId) -> bool {
+        if invariants::non_root_box_created() {
+            false
+        } else {
+            self.push_insn(block, Insn::PatchPoint { invariant: Invariant::RootBoxOnly, state });
+            true
+        }
+    }
+
     /// Assume that objects of a given class will have no singleton class.
     /// Returns true if safe to assume so and emits a PatchPoint.
     /// Returns false if we've already seen a singleton class for this class,
@@ -3899,10 +3914,52 @@ impl Function {
                             // entered the compiler.  That means we can just return nil for this
                             // shape + iv name
                             self.push_insn(block, Insn::Const { val: Const::Value(Qnil) })
+                        } else if recv_type.flags().is_t_class_or_module() {
+                            // Class/module ivar: load from prime classext's fields_obj
+                            if self.assume_root_box(block, state) {
+                                // Root box only: load directly from prime classext
+                                let fields_obj = self.push_insn(block, Insn::LoadField {
+                                    recv: self_val, id: ID!(_fields_obj),
+                                    offset: RCLASS_OFFSET_PRIME_FIELDS_OBJ as i32,
+                                    return_type: types::RubyValue,
+                                });
+                                if recv_type.flags().is_fields_embedded() {
+                                    let offset = ROBJECT_OFFSET_AS_ARY as i32
+                                        + (SIZEOF_VALUE * ivar_index.to_usize()) as i32;
+                                    self.push_insn(block, Insn::LoadField {
+                                        recv: fields_obj, id, offset,
+                                        return_type: types::BasicObject,
+                                    })
+                                } else {
+                                    let ptr = self.push_insn(block, Insn::LoadField {
+                                        recv: fields_obj, id: ID!(_as_heap),
+                                        offset: ROBJECT_OFFSET_AS_HEAP_FIELDS as i32,
+                                        return_type: types::CPtr,
+                                    });
+                                    let offset = SIZEOF_VALUE_I32 * ivar_index as i32;
+                                    self.push_insn(block, Insn::LoadField {
+                                        recv: ptr, id, offset,
+                                        return_type: types::BasicObject,
+                                    })
+                                }
+                            } else {
+                                // Non-root box active: fall back to C call
+                                // NOTE: it's fine to use rb_ivar_get_at_no_ractor_check because
+                                // getinstancevariable does assume_single_ractor_mode()
+                                let ivar_index_insn = self.push_insn(block, Insn::Const { val: Const::CUInt16(ivar_index as u16) });
+                                self.push_insn(block, Insn::CCall {
+                                    cfunc: rb_ivar_get_at_no_ractor_check as *const u8,
+                                    recv: self_val,
+                                    args: vec![ivar_index_insn],
+                                    name: ID!(rb_ivar_get_at_no_ractor_check),
+                                    return_type: types::BasicObject,
+                                    elidable: true })
+                            }
                         } else if !recv_type.flags().is_t_object() {
+                            // Non-T_OBJECT, non-class/module (e.g. T_DATA): fall back to C call
                             // NOTE: it's fine to use rb_ivar_get_at_no_ractor_check because
                             // getinstancevariable does assume_single_ractor_mode()
-                            let ivar_index_insn: InsnId = self.push_insn(block, Insn::Const { val: Const::CUInt16(ivar_index as u16) });
+                            let ivar_index_insn = self.push_insn(block, Insn::Const { val: Const::CUInt16(ivar_index as u16) });
                             self.push_insn(block, Insn::CCall {
                                 cfunc: rb_ivar_get_at_no_ractor_check as *const u8,
                                 recv: self_val,