ZJIT: Specialize monomorphic GetIvar (ruby#14388)

Specialize monomorphic `GetIvar` into:

* `GuardType(HeapObject)`
* `GuardShape`
* `LoadIvarEmbedded` or `LoadIvarExtended`

This requires profiling self for `getinstancevariable` (it's not on the operand
stack).

This also optimizes `GetIvar`s that happen as a result of inlining
`attr_reader` and `attr_accessor`.

Also move some (newly) shared JIT helpers into jit.c.

Author: tekknolagi

insns.def

@@ -212,6 +212,7 @@ getinstancevariable
 (VALUE val)
 /* Ractor crashes when it accesses class/module-level instances variables. */
 // attr bool leaf = false; /* has IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR() */
+// attr bool zjit_profile = true;
 {
     val = vm_getinstancevariable(GET_ISEQ(), GET_SELF(), id, ic);
 }

jit.c

@@ -14,6 +14,12 @@
 #include "iseq.h"
 #include "internal/gc.h"
 
+// Field offsets for the RObject struct
+enum robject_offsets {
+    ROBJECT_OFFSET_AS_HEAP_FIELDS = offsetof(struct RObject, as.heap.fields),
+    ROBJECT_OFFSET_AS_ARY = offsetof(struct RObject, as.ary),
+};
+
 unsigned int
 rb_iseq_encoded_size(const rb_iseq_t *iseq)
 {
@@ -454,3 +460,9 @@ rb_set_cfp_sp(struct rb_control_frame_struct *cfp, VALUE *sp)
 {
     cfp->sp = sp;
 }
+
+bool
+rb_jit_shape_too_complex_p(shape_id_t shape_id)
+{
+    return rb_shape_too_complex_p(shape_id);
+}

test/ruby/test_zjit.rb

@@ -1162,6 +1162,36 @@ def test() = @foo
     }
   end
 
+  def test_getinstancevariable_miss
+    assert_compiles '[1, 1, 4]', %q{
+      class C
+        def foo
+          @foo
+        end
+
+        def foo_then_bar
+          @foo = 1
+          @bar = 2
+        end
+
+        def bar_then_foo
+          @bar = 3
+          @foo = 4
+        end
+      end
+
+      o1 = C.new
+      o1.foo_then_bar
+      result = []
+      result << o1.foo
+      result << o1.foo
+      o2 = C.new
+      o2.bar_then_foo
+      result << o2.foo
+      result
+    }
+  end
+
   def test_setinstancevariable
     assert_compiles '1', %q{
       def test() = @foo = 1

yjit.c

@@ -39,12 +39,6 @@
 
 #include <errno.h>
 
-// Field offsets for the RObject struct
-enum robject_offsets {
-    ROBJECT_OFFSET_AS_HEAP_FIELDS = offsetof(struct RObject, as.heap.fields),
-    ROBJECT_OFFSET_AS_ARY = offsetof(struct RObject, as.ary),
-};
-
 // Field offsets for the RString struct
 enum rstring_offsets {
     RUBY_OFFSET_RSTRING_LEN = offsetof(struct RString, len)
@@ -758,12 +752,6 @@ rb_object_shape_count(void)
     return ULONG2NUM((unsigned long)rb_shapes_count());
 }
 
-bool
-rb_yjit_shape_too_complex_p(shape_id_t shape_id)
-{
-    return rb_shape_too_complex_p(shape_id);
-}
-
 bool
 rb_yjit_shape_obj_too_complex_p(VALUE obj)
 {

yjit/bindgen/src/main.rs

@@ -99,7 +99,6 @@ fn main() {
         .allowlist_function("rb_shape_get_iv_index")
         .allowlist_function("rb_shape_transition_add_ivar_no_warnings")
         .allowlist_function("rb_yjit_shape_obj_too_complex_p")
-        .allowlist_function("rb_yjit_shape_too_complex_p")
         .allowlist_function("rb_yjit_shape_capacity")
         .allowlist_function("rb_yjit_shape_index")
         .allowlist_var("SHAPE_ID_NUM_BITS")
@@ -351,11 +350,12 @@ fn main() {
         .allowlist_function("rb_yjit_invokeblock_sp_pops")
         .allowlist_function("rb_yjit_set_exception_return")
         .allowlist_function("rb_yjit_str_concat_codepoint")
-        .allowlist_type("robject_offsets")
         .allowlist_type("rstring_offsets")
 
         // From jit.c
         .allowlist_function("rb_assert_holding_vm_lock")
+        .allowlist_function("rb_jit_shape_too_complex_p")
+        .allowlist_type("robject_offsets")
 
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")

yjit/src/codegen.rs

@@ -3103,7 +3103,7 @@ fn gen_set_ivar(
 
         // If the VM ran out of shapes, or this class generated too many leaf,
         // it may be de-optimized into OBJ_TOO_COMPLEX_SHAPE (hash-table).
-        new_shape_too_complex = unsafe { rb_yjit_shape_too_complex_p(next_shape_id) };
+        new_shape_too_complex = unsafe { rb_jit_shape_too_complex_p(next_shape_id) };
         if new_shape_too_complex {
             Some((next_shape_id, None, 0_usize))
         } else {

yjit/src/cruby_bindings.inc.rs

@@ -337,12 +337,6 @@ rb_zjit_print_exception(void)
     rb_warn("Ruby error: %"PRIsVALUE"", rb_funcall(exception, rb_intern("full_message"), 0));
 }
 
-bool
-rb_zjit_shape_obj_too_complex_p(VALUE obj)
-{
-    return rb_shape_obj_too_complex_p(obj);
-}
-
 enum {
     RB_INVALID_SHAPE_ID = INVALID_SHAPE_ID,
 };

zjit.c

@@ -111,7 +111,6 @@ fn main() {
         .allowlist_function("rb_shape_id_offset")
         .allowlist_function("rb_shape_get_iv_index")
         .allowlist_function("rb_shape_transition_add_ivar_no_warnings")
-        .allowlist_function("rb_zjit_shape_obj_too_complex_p")
         .allowlist_var("SHAPE_ID_NUM_BITS")
 
         // From ruby/internal/intern/object.h
@@ -367,6 +366,8 @@ fn main() {
 
         // From jit.c
         .allowlist_function("rb_assert_holding_vm_lock")
+        .allowlist_function("rb_jit_shape_too_complex_p")
+        .allowlist_type("robject_offsets")
 
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")

zjit/bindgen/src/main.rs

@@ -405,6 +405,9 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::ToArray { val, state } => { gen_to_array(jit, asm, opnd!(val), &function.frame_state(state)) },
         &Insn::DefinedIvar { self_val, id, pushval, .. } => { gen_defined_ivar(asm, opnd!(self_val), id, pushval) },
         &Insn::ArrayExtend { left, right, state } => { no_output!(gen_array_extend(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state))) },
+        &Insn::GuardShape { val, shape, state } => gen_guard_shape(jit, asm, opnd!(val), shape, &function.frame_state(state)),
+        &Insn::LoadIvarEmbedded { self_val, id, index } => gen_load_ivar_embedded(asm, opnd!(self_val), id, index),
+        &Insn::LoadIvarExtended { self_val, id, index } => gen_load_ivar_extended(asm, opnd!(self_val), id, index),
         &Insn::ArrayMax { state, .. }
         | &Insn::FixnumDiv { state, .. }
         | &Insn::FixnumMod { state, .. }
@@ -716,6 +719,38 @@ fn gen_array_extend(jit: &mut JITState, asm: &mut Assembler, left: Opnd, right:
     asm_ccall!(asm, rb_ary_concat, left, right);
 }
 
+fn gen_guard_shape(jit: &mut JITState, asm: &mut Assembler, val: Opnd, shape: ShapeId, state: &FrameState) -> Opnd {
+    let shape_id_offset = unsafe { rb_shape_id_offset() };
+    let val = asm.load(val);
+    let shape_opnd = Opnd::mem(SHAPE_ID_NUM_BITS as u8, val, shape_id_offset);
+    asm.cmp(shape_opnd, Opnd::UImm(shape.0 as u64));
+    asm.jne(side_exit(jit, state, SideExitReason::GuardShape(shape)));
+    val
+}
+
+fn gen_load_ivar_embedded(asm: &mut Assembler, self_val: Opnd, id: ID, index: u16) -> Opnd {
+    // See ROBJECT_FIELDS() from include/ruby/internal/core/robject.h
+
+    asm_comment!(asm, "Load embedded ivar id={} index={}", id.contents_lossy(), index);
+    let offs = ROBJECT_OFFSET_AS_ARY as i32 + (SIZEOF_VALUE * index as usize) as i32;
+    let self_val = asm.load(self_val);
+    let ivar_opnd = Opnd::mem(64, self_val, offs);
+    asm.load(ivar_opnd)
+}
+
+fn gen_load_ivar_extended(asm: &mut Assembler, self_val: Opnd, id: ID, index: u16) -> Opnd {
+    asm_comment!(asm, "Load extended ivar id={} index={}", id.contents_lossy(), index);
+    // Compile time value is *not* embedded.
+
+    // Get a pointer to the extended table
+    let self_val = asm.load(self_val);
+    let tbl_opnd = asm.load(Opnd::mem(64, self_val, ROBJECT_OFFSET_AS_HEAP_FIELDS as i32));
+
+    // Read the ivar from the extended table
+    let ivar_opnd = Opnd::mem(64, tbl_opnd, (SIZEOF_VALUE * index as usize) as i32);
+    asm.load(ivar_opnd)
+}
+
 /// Compile an interpreter entry block to be inserted into an ISEQ
 fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
     asm_comment!(asm, "ZJIT entry point: {}", iseq_get_location(iseq, 0));
@@ -1270,6 +1305,12 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
 
         asm.cmp(klass, Opnd::Value(expected_class));
         asm.jne(side_exit);
+    } else if guard_type.bit_equal(types::HeapObject) {
+        let side_exit = side_exit(jit, state, GuardType(guard_type));
+        asm.cmp(val, Opnd::Value(Qfalse));
+        asm.je(side_exit.clone());
+        asm.test(val, (RUBY_IMMEDIATE_MASK as u64).into());
+        asm.jnz(side_exit);
     } else {
         unimplemented!("unsupported type: {guard_type}");
     }