ZJIT: Use LoadField for TypedData ivars (ruby#16259)

tekknolagi · XrXr · web-flow · commit a81c3eb795b9 · 2026-03-02T16:06:03.000-05:00
Drops C calls to `rb_ivar_get_at_no_ractor_check` out of the stats completely.

Co-authored-by: Alan Wu &lt;XrXr@users.noreply.github.com&gt;
diff --git a/depend b/depend
@@ -7809,6 +7809,7 @@ jit.$(OBJEXT): {$(VPATH)}internal/error.h
 jit.$(OBJEXT): {$(VPATH)}internal/eval.h
 jit.$(OBJEXT): {$(VPATH)}internal/event.h
 jit.$(OBJEXT): {$(VPATH)}internal/fl_type.h
+jit.$(OBJEXT): {$(VPATH)}internal/core/rtypeddata.h
 jit.$(OBJEXT): {$(VPATH)}internal/gc.h
 jit.$(OBJEXT): {$(VPATH)}internal/glob.h
 jit.$(OBJEXT): {$(VPATH)}internal/globals.h
diff --git a/jit.c b/jit.c
@@ -18,6 +18,7 @@
 #include "internal/string.h"
 #include "internal/class.h"
 #include "internal/imemo.h"
+#include "ruby/internal/core/rtypeddata.h"
 
 enum jit_bindgen_constants {
     // Field offsets for the RObject struct
@@ -27,6 +28,9 @@ enum jit_bindgen_constants {
     // 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 offset for fields_obj in RTypedData
+    RTYPEDDATA_OFFSET_FIELDS_OBJ = offsetof(struct RTypedData, fields_obj),
+
     // Field offsets for the RString struct
     RUBY_OFFSET_RSTRING_LEN = offsetof(struct RString, len),
 
@@ -541,6 +545,13 @@ rb_jit_class_fields_embedded_p(VALUE klass)
     return !fields_obj || !FL_TEST_RAW(fields_obj, OBJ_FIELD_HEAP);
 }
 
+bool
+rb_jit_typed_data_fields_embedded_p(VALUE obj)
+{
+    VALUE fields_obj = RTYPEDDATA(obj)->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
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
@@ -81,6 +81,7 @@ fn main() {
         .allowlist_type("RBasic")
 
         .allowlist_type("ruby_rstring_flags")
+        .allowlist_type("rbimpl_typeddata_flags")
 
         // This function prints info about a value and is useful for debugging
         .allowlist_function("rb_raw_obj_info")
@@ -310,6 +311,8 @@ fn main() {
         .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_typed_data_p")
+        .allowlist_function("rb_jit_typed_data_fields_embedded_p")
         .allowlist_function("rb_jit_vm_lock_then_barrier")
         .allowlist_function("rb_jit_vm_unlock")
         .allowlist_function("rb_jit_for_each_iseq")
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
@@ -94,6 +94,8 @@ use std::fmt::{Debug, Display, Formatter};
 use std::os::raw::{c_char, c_int, c_uint};
 use std::panic::{catch_unwind, UnwindSafe};
 
+use crate::cast::IntoUsize as _;
+
 // We check that we can do this with the configure script and a couple of
 // static asserts. u64 and not usize to play nice with lowering to x86.
 pub type size_t = u64;
@@ -602,6 +604,16 @@ impl VALUE {
         unsafe { rb_jit_class_fields_embedded_p(self) }
     }
 
+    pub fn typed_data_p(self) -> bool {
+        !self.special_const_p() &&
+            self.builtin_type() == RUBY_T_DATA &&
+            0 != (self.builtin_flags() & RUBY_TYPED_FL_IS_TYPED_DATA.to_usize())
+    }
+
+    pub fn typed_data_fields_embedded_p(self) -> bool {
+        unsafe { rb_jit_typed_data_fields_embedded_p(self) }
+    }
+
     pub fn as_fixnum(self) -> i64 {
         assert!(self.fixnum_p());
         (self.0 as i64) >> 1
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -3949,8 +3949,34 @@ impl Function {
                                     return_type: types::BasicObject,
                                     elidable: true })
                             }
+                        } else if recv_type.flags().is_typed_data() {
+                            // Typed T_DATA: load from fields_obj at fixed offset in RTypedData
+                            let fields_obj = self.push_insn(block, Insn::LoadField {
+                                recv: self_val, id: ID!(_fields_obj),
+                                offset: RTYPEDDATA_OFFSET_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 if !recv_type.flags().is_t_object() {
-                            // Non-T_OBJECT, non-class/module (e.g. T_DATA): fall back to C call
+                            // Non-T_OBJECT, non-class/module, non-typed-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 = self.push_insn(block, Insn::Const { val: Const::CUInt16(ivar_index as u16) });
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -7253,7 +7253,8 @@ mod hir_opt_tests {
     }
 
     #[test]
-    fn test_optimize_getivar_on_t_data() {
+    fn test_optimize_getivar_on_t_struct() {
+        // Range is T_STRUCT (not T_DATA): falls back to CCall
         eval("
             class C < Range
               def test = @a
@@ -7285,6 +7286,78 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_optimize_getivar_on_typed_data() {
+        // Thread is typed T_DATA: uses LoadField chain via RTypedData fields_obj
+        eval("
+            class C < Thread
+              def test = @a
+            end
+            obj = C.new { }
+            obj.join
+            obj.instance_variable_set(:@a, 1)
+            obj.test
+            TEST = C.instance_method(:test)
+        ");
+        assert_snapshot!(hir_string_proc("TEST"), @r"
+        fn test@<compiled>:3:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          PatchPoint SingleRactorMode
+          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v18:CShape = LoadField v17, :_shape_id@0x1000
+          v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
+          v20:RubyValue = LoadField v17, :_fields_obj@0x1002
+          v21:BasicObject = LoadField v20, :@a@0x1002
+          CheckInterrupts
+          Return v21
+        ");
+    }
+
+    #[test]
+    fn test_optimize_getivar_on_typed_data_heap_fields() {
+        // Typed T_DATA with enough ivars to force heap field storage
+        eval("
+            class C < Thread
+              def test = @var1000
+            end
+            obj = C.new { }
+            obj.join
+            1000.times { |i| obj.instance_variable_set(:\"@var#{i}\", 1) }
+            obj.instance_variable_set(:@var1000, 42)
+            obj.test
+            TEST = C.instance_method(:test)
+        ");
+        assert_snapshot!(hir_string_proc("TEST"), @r"
+        fn test@<compiled>:3:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          PatchPoint SingleRactorMode
+          v17:HeapBasicObject = GuardType v6, HeapBasicObject
+          v18:CShape = LoadField v17, :_shape_id@0x1000
+          v19:CShape[0x1001] = GuardBitEquals v18, CShape(0x1001)
+          v20:RubyValue = LoadField v17, :_fields_obj@0x1002
+          v21:CPtr = LoadField v20, :_as_heap@0x1002
+          v22:BasicObject = LoadField v21, :@var1000@0x1003
+          CheckInterrupts
+          Return v22
+        ");
+    }
+
     #[test]
     fn test_optimize_getivar_on_module_multi_ractor() {
         eval("
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
@@ -207,6 +207,8 @@ impl Flags {
     const IS_FIELDS_EMBEDDED: u32 = 1 << 5;
     /// Object is a T_CLASS or T_MODULE
     const IS_T_CLASS_OR_MODULE: u32 = 1 << 6;
+    /// Object is a typed T_DATA (RTYPEDDATA_P)
+    const IS_TYPED_DATA: u32 = 1 << 7;
 
     pub fn none() -> Self { Self(Self::NONE) }
 
@@ -218,6 +220,7 @@ impl Flags {
     pub fn is_object_profiling(self) -> bool { (self.0 & Self::IS_OBJECT_PROFILING) != 0 }
     pub fn is_fields_embedded(self) -> bool { (self.0 & Self::IS_FIELDS_EMBEDDED) != 0 }
     pub fn is_t_class_or_module(self) -> bool { (self.0 & Self::IS_T_CLASS_OR_MODULE) != 0 }
+    pub fn is_typed_data(self) -> bool { (self.0 & Self::IS_TYPED_DATA) != 0 }
 }
 
 /// opt_send_without_block/opt_plus/... should store:
@@ -300,6 +303,12 @@ impl ProfiledType {
                 flags.0 |= Flags::IS_FIELDS_EMBEDDED;
             }
         }
+        if obj.typed_data_p() {
+            flags.0 |= Flags::IS_TYPED_DATA;
+            if obj.typed_data_fields_embedded_p() {
+                flags.0 |= Flags::IS_FIELDS_EMBEDDED;
+            }
+        }
         Self { class: obj.class_of(), shape: obj.shape_id_of(), flags }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -207,6 +207,8 @@ impl Flags {`
`207`	`207`	`const IS_FIELDS_EMBEDDED: u32 = 1 << 5;`
`208`	`208`	`/// Object is a T_CLASS or T_MODULE`
`209`	`209`	`const IS_T_CLASS_OR_MODULE: u32 = 1 << 6;`
	`210`	`+ /// Object is a typed T_DATA (RTYPEDDATA_P)`
	`211`	`+ const IS_TYPED_DATA: u32 = 1 << 7;`
`210`	`212`
`211`	`213`	`pub fn none() -> Self { Self(Self::NONE) }`
`212`	`214`
`@@ -218,6 +220,7 @@ impl Flags {`
`218`	`220`	`pub fn is_object_profiling(self) -> bool { (self.0 & Self::IS_OBJECT_PROFILING) != 0 }`
`219`	`221`	`pub fn is_fields_embedded(self) -> bool { (self.0 & Self::IS_FIELDS_EMBEDDED) != 0 }`
`220`	`222`	`pub fn is_t_class_or_module(self) -> bool { (self.0 & Self::IS_T_CLASS_OR_MODULE) != 0 }`
	`223`	`+ pub fn is_typed_data(self) -> bool { (self.0 & Self::IS_TYPED_DATA) != 0 }`
`221`	`224`	`}`
`222`	`225`
`223`	`226`	`/// opt_send_without_block/opt_plus/... should store:`
`@@ -300,6 +303,12 @@ impl ProfiledType {`
`300`	`303`	`flags.0 \|= Flags::IS_FIELDS_EMBEDDED;`
`301`	`304`	`}`
`302`	`305`	`}`
	`306`	`+ if obj.typed_data_p() {`
	`307`	`+ flags.0 \|= Flags::IS_TYPED_DATA;`
	`308`	`+ if obj.typed_data_fields_embedded_p() {`
	`309`	`+ flags.0 \|= Flags::IS_FIELDS_EMBEDDED;`
	`310`	`+ }`
	`311`	`+ }`
`303`	`312`	`Self { class: obj.class_of(), shape: obj.shape_id_of(), flags }`
`304`	`313`	`}`
`305`	`314`