ZJIT: Optimize setivar with shape transition (ruby#15375)

Since we do a decent job of pre-sizing objects, don't handle the case where we would need to re-size an object. Also don't handle too-complex shapes.

lobsters stats before:

```
Top-20 calls to C functions from JIT code (79.4% of total 90,051,140):
                             rb_vm_opt_send_without_block: 19,762,433 (21.9%)
                                rb_vm_setinstancevariable:  7,698,314 ( 8.5%)
                                             rb_hash_aref:  6,767,461 ( 7.5%)
                                          rb_vm_env_write:  5,373,080 ( 6.0%)
                                               rb_vm_send:  5,049,229 ( 5.6%)
                                rb_vm_getinstancevariable:  4,535,259 ( 5.0%)
                                        rb_obj_is_kind_of:  3,746,306 ( 4.2%)
                           rb_ivar_get_at_no_ractor_check:  3,745,237 ( 4.2%)
                                        rb_vm_invokesuper:  3,037,467 ( 3.4%)
                                             rb_ary_entry:  2,351,983 ( 2.6%)
                               rb_vm_opt_getconstant_path:  1,344,740 ( 1.5%)
                                        rb_vm_invokeblock:  1,184,474 ( 1.3%)
                                                 Hash#[]=:  1,064,288 ( 1.2%)
                                       rb_gc_writebarrier:  1,006,972 ( 1.1%)
                                rb_ec_ary_new_from_values:    902,687 ( 1.0%)
                                                    fetch:    898,667 ( 1.0%)
                                        rb_str_buf_append:    833,787 ( 0.9%)
                               rb_class_allocate_instance:    822,024 ( 0.9%)
                                               Hash#fetch:    699,580 ( 0.8%)
                                                    _bi20:    682,068 ( 0.8%)
Top-4 setivar fallback reasons (100.0% of total 7,732,326):
  shape_transition: 6,032,109 (78.0%)
   not_monomorphic: 1,469,300 (19.0%)
      not_t_object:   172,636 ( 2.2%)
       too_complex:    58,281 ( 0.8%)
```

lobsters stats after:

```
Top-20 calls to C functions from JIT code (79.0% of total 88,322,656):
                             rb_vm_opt_send_without_block: 19,777,880 (22.4%)
                                             rb_hash_aref:  6,771,589 ( 7.7%)
                                          rb_vm_env_write:  5,372,789 ( 6.1%)
                                       rb_gc_writebarrier:  5,195,527 ( 5.9%)
                                               rb_vm_send:  5,049,145 ( 5.7%)
                                rb_vm_getinstancevariable:  4,538,485 ( 5.1%)
                                        rb_obj_is_kind_of:  3,746,241 ( 4.2%)
                           rb_ivar_get_at_no_ractor_check:  3,745,172 ( 4.2%)
                                        rb_vm_invokesuper:  3,037,157 ( 3.4%)
                                             rb_ary_entry:  2,351,968 ( 2.7%)
                                rb_vm_setinstancevariable:  1,703,337 ( 1.9%)
                               rb_vm_opt_getconstant_path:  1,344,730 ( 1.5%)
                                        rb_vm_invokeblock:  1,184,290 ( 1.3%)
                                                 Hash#[]=:  1,061,868 ( 1.2%)
                                rb_ec_ary_new_from_values:    902,666 ( 1.0%)
                                                    fetch:    898,666 ( 1.0%)
                                        rb_str_buf_append:    833,784 ( 0.9%)
                               rb_class_allocate_instance:    821,778 ( 0.9%)
                                               Hash#fetch:    755,913 ( 0.9%)
Top-4 setivar fallback reasons (100.0% of total 1,703,337):
            not_monomorphic: 1,472,405 (86.4%)
               not_t_object:   172,629 (10.1%)
                too_complex:    58,281 ( 3.4%)
  new_shape_needs_extension:        22 ( 0.0%)
```

I also noticed that primitive printing in HIR was broken so I fixed that.

Co-authored-by: Aaron Patterson <[email protected]>

Author: tekknolagi

jit.c

@@ -767,3 +767,9 @@ rb_yarv_str_eql_internal(VALUE str1, VALUE str2)
 }
 
 void rb_jit_str_concat_codepoint(VALUE str, VALUE codepoint);
+
+attr_index_t
+rb_jit_shape_capacity(shape_id_t shape_id)
+{
+    return RSHAPE_CAPACITY(shape_id);
+}

zjit/bindgen/src/main.rs

@@ -100,6 +100,7 @@ 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_jit_shape_capacity")
         .allowlist_var("rb_invalid_shape_id")
         .allowlist_type("shape_id_fl_type")
         .allowlist_var("VM_KW_SPECIFIED_BITS_MAX")

zjit/src/codegen.rs

@@ -350,6 +350,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Const { val: Const::CPtr(val) } => gen_const_cptr(val),
         &Insn::Const { val: Const::CInt64(val) } => gen_const_long(val),
         &Insn::Const { val: Const::CUInt16(val) } => gen_const_uint16(val),
+        &Insn::Const { val: Const::CUInt32(val) } => gen_const_uint32(val),
         Insn::Const { .. } => panic!("Unexpected Const in gen_insn: {insn}"),
         Insn::NewArray { elements, state } => gen_new_array(asm, opnds!(elements), &function.frame_state(*state)),
         Insn::NewHash { elements, state } => gen_new_hash(jit, asm, opnds!(elements), &function.frame_state(*state)),
@@ -475,7 +476,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::LoadEC => gen_load_ec(),
         Insn::LoadSelf => gen_load_self(),
         &Insn::LoadField { recv, id, offset, return_type: _ } => gen_load_field(asm, opnd!(recv), id, offset),
-        &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val))),
+        &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val), function.type_of(val))),
         &Insn::WriteBarrier { recv, val } => no_output!(gen_write_barrier(asm, opnd!(recv), opnd!(val), function.type_of(val))),
         &Insn::IsBlockGiven => gen_is_block_given(jit, asm),
         Insn::ArrayInclude { elements, target, state } => gen_array_include(jit, asm, opnds!(elements), opnd!(target), &function.frame_state(*state)),
@@ -1099,10 +1100,10 @@ fn gen_load_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32) -> Opnd
     asm.load(Opnd::mem(64, recv, offset))
 }
 
-fn gen_store_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32, val: Opnd) {
+fn gen_store_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32, val: Opnd, val_type: Type) {
     asm_comment!(asm, "Store field id={} offset={}", id.contents_lossy(), offset);
     let recv = asm.load(recv);
-    asm.store(Opnd::mem(64, recv, offset), val);
+    asm.store(Opnd::mem(val_type.num_bits(), recv, offset), val);
 }
 
 fn gen_write_barrier(asm: &mut Assembler, recv: Opnd, val: Opnd, val_type: Type) {
@@ -1159,6 +1160,10 @@ fn gen_const_uint16(val: u16) -> lir::Opnd {
     Opnd::UImm(val as u64)
 }
 
+fn gen_const_uint32(val: u32) -> lir::Opnd {
+    Opnd::UImm(val as u64)
+}
+
 /// Compile a basic block argument
 fn gen_param(asm: &mut Assembler, idx: usize) -> lir::Opnd {
     // Allocate a register or a stack slot

zjit/src/cruby.rs

@@ -999,8 +999,9 @@ mod manual_defs {
     use super::*;
 
     pub const SIZEOF_VALUE: usize = 8;
+    pub const BITS_PER_BYTE: usize = 8;
     pub const SIZEOF_VALUE_I32: i32 = SIZEOF_VALUE as i32;
-    pub const VALUE_BITS: u8 = 8 * SIZEOF_VALUE as u8;
+    pub const VALUE_BITS: u8 = BITS_PER_BYTE as u8 * SIZEOF_VALUE as u8;
 
     pub const RUBY_LONG_MIN: isize = std::os::raw::c_long::MIN as isize;
     pub const RUBY_LONG_MAX: isize = std::os::raw::c_long::MAX as isize;
@@ -1382,6 +1383,7 @@ pub(crate) mod ids {
         name: thread_ptr
         name: self_              content: b"self"
         name: rb_ivar_get_at_no_ractor_check
+        name: _shape_id
     }
 
     /// Get an CRuby `ID` to an interned string, e.g. a particular method name.

zjit/src/cruby_bindings.inc.rs

@@ -2949,10 +2949,34 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         }
                         let mut ivar_index: u16 = 0;
+                        let mut next_shape_id = recv_type.shape();
                         if !unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } {
-                            // TODO(max): Shape transition
-                            self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_shape_transition));
-                            self.push_insn_id(block, insn_id); continue;
+                            // Current shape does not contain this ivar; do a shape transition.
+                            let current_shape_id = recv_type.shape();
+                            let class = recv_type.class();
+                            // We're only looking at T_OBJECT so ignore all of the imemo stuff.
+                            assert!(recv_type.flags().is_t_object());
+                            next_shape_id = ShapeId(unsafe { rb_shape_transition_add_ivar_no_warnings(class, current_shape_id.0, id) });
+                            let ivar_result = unsafe { rb_shape_get_iv_index(next_shape_id.0, id, &mut ivar_index) };
+                            assert!(ivar_result, "New shape must have the ivar index");
+                            // 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).
+                            let new_shape_too_complex = unsafe { rb_jit_shape_too_complex_p(next_shape_id.0) };
+                            // TODO(max): Is it OK to bail out here after making a shape transition?
+                            if new_shape_too_complex {
+                                self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_too_complex));
+                                self.push_insn_id(block, insn_id); continue;
+                            }
+                            let current_capacity = unsafe { rb_jit_shape_capacity(current_shape_id.0) };
+                            let next_capacity = unsafe { rb_jit_shape_capacity(next_shape_id.0) };
+                            // If the new shape has a different capacity, or is TOO_COMPLEX, we'll have to
+                            // reallocate it.
+                            let needs_extension = next_capacity != current_capacity;
+                            if needs_extension {
+                                self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_needs_extension));
+                                self.push_insn_id(block, insn_id); continue;
+                            }
+                            // Fall through to emitting the ivar write
                         }
                         let self_val = self.push_insn(block, Insn::GuardType { val: self_val, guard_type: types::HeapBasicObject, state });
                         let self_val = self.push_insn(block, Insn::GuardShape { val: self_val, shape: recv_type.shape(), state });
@@ -2968,6 +2992,13 @@ impl Function {
                         };
                         self.push_insn(block, Insn::StoreField { recv: ivar_storage, id, offset, val });
                         self.push_insn(block, Insn::WriteBarrier { recv: self_val, val });
+                        if next_shape_id != recv_type.shape() {
+                            // Write the new shape ID
+                            assert_eq!(SHAPE_ID_NUM_BITS, 32);
+                            let shape_id = self.push_insn(block, Insn::Const { val: Const::CUInt32(next_shape_id.0) });
+                            let shape_id_offset = unsafe { rb_shape_id_offset() };
+                            self.push_insn(block, Insn::StoreField { recv: self_val, id: ID!(_shape_id), offset: shape_id_offset, val: shape_id });
+                        }
                     }
                     _ => { self.push_insn_id(block, insn_id); }
                 }

zjit/src/hir.rs

@@ -3484,6 +3484,39 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_dont_specialize_complex_shape_definedivar() {
+        eval(r#"
+            class C
+              def test = defined?(@a)
+            end
+            obj = C.new
+            (0..1000).each do |i|
+              obj.instance_variable_set(:"@v#{i}", i)
+            end
+            (0..1000).each do |i|
+              obj.remove_instance_variable(:"@v#{i}")
+            end
+            obj.test
+            TEST = C.instance_method(:test)
+        "#);
+        assert_snapshot!(hir_string_proc("TEST"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          IncrCounter definedivar_fallback_too_complex
+          v10:StringExact|NilClass = DefinedIvar v6, :@a
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
     #[test]
     fn test_specialize_monomorphic_setivar_already_in_shape() {
         eval("
@@ -3513,7 +3546,7 @@ mod hir_opt_tests {
     }
 
     #[test]
-    fn test_dont_specialize_monomorphic_setivar_with_shape_transition() {
+    fn test_specialize_monomorphic_setivar_with_shape_transition() {
         eval("
             def test = @foo = 5
             test
@@ -3530,13 +3563,57 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           v10:Fixnum[5] = Const Value(5)
           PatchPoint SingleRactorMode
-          IncrCounter setivar_fallback_shape_transition
-          SetIvar v6, :@foo, v10
+          v19:HeapBasicObject = GuardType v6, HeapBasicObject
+          v20:HeapBasicObject = GuardShape v19, 0x1000
+          StoreField v20, :@foo@0x1001, v10
+          WriteBarrier v20, v10
+          v23:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v20, :_shape_id@0x1002, v23
           CheckInterrupts
           Return v10
         ");
     }
 
+    #[test]
+    fn test_specialize_multiple_monomorphic_setivar_with_shape_transition() {
+        eval("
+            def test
+              @foo = 1
+              @bar = 2
+            end
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v10:Fixnum[1] = Const Value(1)
+          PatchPoint SingleRactorMode
+          v25:HeapBasicObject = GuardType v6, HeapBasicObject
+          v26:HeapBasicObject = GuardShape v25, 0x1000
+          StoreField v26, :@foo@0x1001, v10
+          WriteBarrier v26, v10
+          v29:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v26, :_shape_id@0x1002, v29
+          v16:Fixnum[2] = Const Value(2)
+          PatchPoint SingleRactorMode
+          v31:HeapBasicObject = GuardType v6, HeapBasicObject
+          v32:HeapBasicObject = GuardShape v31, 0x1003
+          StoreField v32, :@bar@0x1004, v16
+          WriteBarrier v32, v16
+          v35:CUInt32[4194312] = Const CUInt32(4194312)
+          StoreField v32, :_shape_id@0x1002, v35
+          CheckInterrupts
+          Return v16
+        ");
+    }
+
     #[test]
     fn test_dont_specialize_setivar_with_t_data() {
         eval("
@@ -3611,6 +3688,9 @@ mod hir_opt_tests {
             (0..1000).each do |i|
               obj.instance_variable_set(:"@v#{i}", i)
             end
+            (0..1000).each do |i|
+              obj.remove_instance_variable(:"@v#{i}")
+            end
             obj.test
             TEST = C.instance_method(:test)
         "#);
@@ -3626,7 +3706,7 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           v10:Fixnum[5] = Const Value(5)
           PatchPoint SingleRactorMode
-          IncrCounter setivar_fallback_shape_transition
+          IncrCounter setivar_fallback_too_complex
           SetIvar v6, :@a, v10
           CheckInterrupts
           Return v10
@@ -5417,6 +5497,43 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_dont_optimize_getivar_with_too_complex_shape() {
+        eval(r#"
+            class C
+              attr_accessor :foo
+            end
+            obj = C.new
+            (0..1000).each do |i|
+              obj.instance_variable_set(:"@v#{i}", i)
+            end
+            (0..1000).each do |i|
+              obj.remove_instance_variable(:"@v#{i}")
+            end
+            def test(o) = o.foo
+            test obj
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:12:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          PatchPoint MethodRedefined(C@0x1000, foo@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(C@0x1000)
+          v21:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
+          IncrCounter getivar_fallback_too_complex
+          v22:BasicObject = GetIvar v21, :@foo
+          CheckInterrupts
+          Return v22
+        ");
+    }
+
     #[test]
     fn test_optimize_send_with_block() {
         eval(r#"
@@ -5697,8 +5814,11 @@ mod hir_opt_tests {
           v16:Fixnum[5] = Const Value(5)
           PatchPoint MethodRedefined(C@0x1000, foo=@0x1008, cme:0x1010)
           v26:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
-          IncrCounter setivar_fallback_shape_transition
-          SetIvar v26, :@foo, v16
+          v29:HeapObject[class_exact:C] = GuardShape v26, 0x1038
+          StoreField v29, :@foo@0x1039, v16
+          WriteBarrier v29, v16
+          v32:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v29, :_shape_id@0x103a, v32
           CheckInterrupts
           Return v16
         ");
@@ -5729,8 +5849,11 @@ mod hir_opt_tests {
           v16:Fixnum[5] = Const Value(5)
           PatchPoint MethodRedefined(C@0x1000, foo=@0x1008, cme:0x1010)
           v26:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
-          IncrCounter setivar_fallback_shape_transition
-          SetIvar v26, :@foo, v16
+          v29:HeapObject[class_exact:C] = GuardShape v26, 0x1038
+          StoreField v29, :@foo@0x1039, v16
+          WriteBarrier v29, v16
+          v32:CUInt32[4194311] = Const CUInt32(4194311)
+          StoreField v29, :_shape_id@0x103a, v32
           CheckInterrupts
           Return v16
         ");
@@ -9109,18 +9232,22 @@ mod hir_opt_tests {
          SetLocal l0, EP@3, v27
          v39:BasicObject = GetLocal l0, EP@3
          PatchPoint SingleRactorMode
-         IncrCounter setivar_fallback_shape_transition
-         SetIvar v14, :@formatted, v39
-         v45:Class[VMFrozenCore] = Const Value(VALUE(0x1000))
-         PatchPoint MethodRedefined(Class@0x1008, lambda@0x1010, cme:0x1018)
-         PatchPoint NoSingletonClass(Class@0x1008)
-         v60:BasicObject = CCallWithFrame v45, :RubyVM::FrozenCore.lambda@0x1040, block=0x1048
+         v56:HeapBasicObject = GuardType v14, HeapBasicObject
+         v57:HeapBasicObject = GuardShape v56, 0x1000
+         StoreField v57, :@formatted@0x1001, v39
+         WriteBarrier v57, v39
+         v60:CUInt32[4194311] = Const CUInt32(4194311)
+         StoreField v57, :_shape_id@0x1002, v60
+         v45:Class[VMFrozenCore] = Const Value(VALUE(0x1008))
+         PatchPoint MethodRedefined(Class@0x1010, lambda@0x1018, cme:0x1020)
+         PatchPoint NoSingletonClass(Class@0x1010)
+         v65:BasicObject = CCallWithFrame v45, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050
          v48:BasicObject = GetLocal l0, EP@6
          v49:BasicObject = GetLocal l0, EP@5
          v50:BasicObject = GetLocal l0, EP@4
          v51:BasicObject = GetLocal l0, EP@3
          CheckInterrupts
-         Return v60
+         Return v65
        ");
     }
 }