Invalidate blocks on constant IC updates

lib/ruby_vm/mjit/block.rb

@@ -3,5 +3,7 @@ class RubyVM::MJIT::Block < Struct.new(
   :ctx,        # @param [RubyVM::MJIT::Context] **Starting** Context (TODO: freeze?)
   :start_addr, # @param [Integer] Starting address of this block's JIT code
   :entry_exit, # @param [Integer] Address of entry exit (optional)
+  :incoming,   # @param [Array<RubyVM::MJIT::BranchStub>] Incoming branches
 )
+  def initialize(incoming: [], **) = super
 end

lib/ruby_vm/mjit/compiler.rb

@@ -32,7 +32,11 @@ module RubyVM::MJIT
   class Compiler
     attr_accessor :write_pos
 
-    IseqBlocks = Hash.new { |h, k| h[k] = {} }
+    IseqBlocks = Hash.new do |iseq_pc_ctx, iseq|
+      iseq_pc_ctx[iseq] = Hash.new do |pc_ctx, pc|
+        pc_ctx[pc] = {}
+      end
+    end
     DeadBlocks = [] # invalidated IseqBlocks, but kept for safety
 
     def self.reset_blocks
@@ -51,7 +55,7 @@ module RubyVM::MJIT
       @ocb = CodeBlock.new(mem_block: mem_block + mem_size / 2, mem_size: mem_size / 2, outlined: true)
       @exit_compiler = ExitCompiler.new
       @insn_compiler = InsnCompiler.new(@cb, @ocb, @exit_compiler)
-      Invariants.initialize(@cb, @ocb, @exit_compiler)
+      Invariants.initialize(@cb, @ocb, self, @exit_compiler)
 
       @leave_exit = Assembler.new.then do |asm|
         @exit_compiler.compile_leave_exit(asm)
@@ -113,8 +117,9 @@ module RubyVM::MJIT
           compile_block(asm, jit:, pc: target.pc, ctx: target.ctx.dup)
           @cb.write(asm)
         end
-        set_block(branch_stub.iseq, target.pc, target.ctx, jit.block)
+        block = jit.block
       end
+      block.incoming << branch_stub # prepare for invalidate_block
 
       # Re-generate the branch code for non-fallthrough cases
       unless fallthrough
@@ -131,6 +136,14 @@ module RubyVM::MJIT
       exit 1
     end
 
+    # @param iseq `RubyVM::MJIT::CPointer::Struct_rb_iseq_t`
+    # @param pc [Integer]
+    def invalidate_blocks(iseq, pc)
+      list_blocks(iseq, pc).each do |block|
+        invalidate_block(iseq, block)
+      end
+    end
+
     private
 
     # Callee-saved: rbx, rsp, rbp, r12, r13, r14, r15
@@ -160,8 +173,9 @@ module RubyVM::MJIT
     # @param asm [RubyVM::MJIT::Assembler]
     def compile_block(asm, jit:, pc: jit.iseq.body.iseq_encoded.to_i, ctx: Context.new)
       # Mark the block start address and prepare an exit code storage
-      jit.block = Block.new(pc:, ctx: ctx.dup)
-      asm.block(jit.block)
+      block = Block.new(pc:, ctx: ctx.dup)
+      jit.block = block
+      asm.block(block)
 
       # Compile each insn
       iseq = jit.iseq
@@ -185,7 +199,7 @@ module RubyVM::MJIT
         when KeepCompiling
           index += insn.len
         when EndBlock
-          # TODO: pad nops if entry exit exists
+          # TODO: pad nops if entry exit exists (not needed for x86_64?)
           break
         when CantCompile
           @exit_compiler.compile_side_exit(jit.pc, ctx, asm)
@@ -196,6 +210,7 @@ module RubyVM::MJIT
       end
 
       incr_counter(:compiled_block_count)
+      set_block(iseq, block)
     end
 
     def incr_counter(name)
@@ -204,18 +219,59 @@ module RubyVM::MJIT
       end
     end
 
-    # @param [Integer] pc
-    # @param [RubyVM::MJIT::Context] ctx
-    # @return [RubyVM::MJIT::Block,NilClass]
-    def find_block(iseq, pc, ctx)
-      IseqBlocks[iseq.to_i][[pc, ctx]]
+    def invalidate_block(iseq, block)
+      # Remove this block from the version array
+      remove_block(iseq, block)
+
+      # Invalidate the block with entry exit
+      @cb.with_write_addr(block.start_addr) do
+        asm = Assembler.new
+        asm.comment('invalidate_block')
+        asm.jmp(block.entry_exit)
+        @cb.write(asm)
+      end
+
+      # Re-stub incoming branches
+      block.incoming.each do |branch_stub|
+        target = [branch_stub.target0, branch_stub.target1].compact.find do |target|
+          target.pc == block.pc && target.ctx == block.ctx
+        end
+        next if target.nil?
+        # TODO: Could target.address be a stub address? Is invalidation not needed in that case?
+
+        target.address = Assembler.new.then do |ocb_asm|
+          @exit_compiler.compile_branch_stub(block.ctx, ocb_asm, branch_stub, target == branch_stub.target0)
+          @ocb.write(ocb_asm)
+        end
+        @cb.with_write_addr(branch_stub.start_addr) do
+          branch_asm = Assembler.new
+          branch_stub.shape = Default # cancel fallthrough. TODO: It seems fine for defer_compilation, but is this always safe?
+          branch_stub.compile.call(branch_asm)
+          @cb.write(branch_asm)
+        end
+      end
+      # TODO: Reset jit_func and total_calls if it's the first block after prelude
+    end
+
+    def list_blocks(iseq, pc)
+      IseqBlocks[iseq.to_i][pc].values
     end
 
     # @param [Integer] pc
     # @param [RubyVM::MJIT::Context] ctx
+    # @return [RubyVM::MJIT::Block,NilClass]
+    def find_block(iseq, pc, ctx)
+      IseqBlocks[iseq.to_i][pc][ctx]
+    end
+
     # @param [RubyVM::MJIT::Block] block
-    def set_block(iseq, pc, ctx, block)
-      IseqBlocks[iseq.to_i][[pc, ctx]] = block
+    def set_block(iseq, block)
+      IseqBlocks[iseq.to_i][block.pc][block.ctx] = block
+    end
+
+    # @param [RubyVM::MJIT::Block] block
+    def remove_block(iseq, block)
+      IseqBlocks[iseq.to_i][block.pc].delete(block.ctx)
     end
   end
 end

lib/ruby_vm/mjit/hooks.rb

@@ -5,6 +5,7 @@ module RubyVM::MJIT
     end
 
     def self.on_cme_invalidate(cme)
+      cme = C.rb_callable_method_entry_struct.new(cme)
       Invariants.on_cme_invalidate(cme)
     end
 
@@ -16,8 +17,10 @@ module RubyVM::MJIT
       # to be used later
     end
 
-    def self.on_constant_ic_update(_iseq, _ic, _insn_idx)
-      # to be used later
+    def self.on_constant_ic_update(iseq, ic, insn_idx)
+      iseq = C.rb_iseq_t.new(iseq)
+      ic = C.IC.new(ic)
+      Invariants.on_constant_ic_update(iseq, ic, insn_idx)
     end
 
     def self.on_tracing_invalidate_all(_new_iseq_events)

lib/ruby_vm/mjit/insn_compiler.rb

@@ -214,9 +214,20 @@ module RubyVM::MJIT
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def opt_getconstant_path(jit, ctx, asm)
+      # Cut the block for invalidation
+      unless jit.at_current_insn?
+        defer_compilation(jit, ctx, asm)
+        return EndBlock
+      end
+
       ic = C.iseq_inline_constant_cache.new(jit.operand(0))
       idlist = ic.segments
 
+      # Make sure there is an exit for this block as the interpreter might want
+      # to invalidate this block from rb_mjit_constant_ic_update().
+      # For now, we always take an entry exit even if it was a side exit.
+      Invariants.ensure_block_entry_exit(jit, cause: 'opt_getconstant_path')
+
       # See vm_ic_hit_p(). The same conditions are checked in yjit_constant_ic_update().
       ice = ic.entry
       if ice.nil?
@@ -226,10 +237,6 @@ module RubyVM::MJIT
         return CantCompile
       end
 
-      # Make sure there is an exit for this block as the interpreter might want
-      # to invalidate this block from yjit_constant_ic_update().
-      Invariants.ensure_block_entry_exit(jit, cause: 'opt_getconstant_path')
-
       if ice.ic_cref # with cref
         # Not supported yet
         asm.incr_counter(:optgetconst_cref)

lib/ruby_vm/mjit/invariants.rb

@@ -6,10 +6,12 @@ module RubyVM::MJIT
       # Called by RubyVM::MJIT::Compiler to lazily initialize this
       # @param cb [CodeBlock]
       # @param ocb [CodeBlock]
+      # @param compiler [RubyVM::MJIT::Compiler]
       # @param exit_compiler [RubyVM::MJIT::ExitCompiler]
-      def initialize(cb, ocb, exit_compiler)
+      def initialize(cb, ocb, compiler, exit_compiler)
         @cb = cb
         @ocb = ocb
+        @compiler = compiler
         @exit_compiler = exit_compiler
         @bop_blocks = Set.new # TODO: actually invalidate this
         @cme_blocks = Hash.new { |h, k| h[k] = Set.new }
@@ -57,6 +59,24 @@ module RubyVM::MJIT
         end
       end
 
+      def on_constant_ic_update(iseq, ic, insn_idx)
+        # TODO: check multi ractor as well
+        if ic.entry.ic_cref
+          # No need to recompile the slowpath
+          return
+        end
+
+        pc = iseq.body.iseq_encoded + insn_idx
+        insn_name = Compiler.decode_insn(pc.*).name
+        if insn_name != :opt_getconstant_path && insn_name != :trace_opt_getconstant_path
+          raise 'insn_idx was not at opt_getconstant_path'
+        end
+        if ic.to_i != pc[1]
+          raise 'insn_idx + 1 was not at the updated IC'
+        end
+        @compiler.invalidate_blocks(iseq, pc.to_i)
+      end
+
       def on_tracing_invalidate_all
         # On-Stack Replacement
         @patches.each do |address, target|

mjit.c

@@ -343,11 +343,8 @@ static void
 mjit_cme_invalidate(void *data)
 {
     if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
-    rb_callable_method_entry_t *cme = (rb_callable_method_entry_t *)data;
     WITH_MJIT_ISOLATED({
-        VALUE cme_klass = rb_funcall(rb_mMJITC, rb_intern("rb_callable_method_entry_struct"), 0);
-        VALUE cme_ptr = rb_funcall(cme_klass, rb_intern("new"), 1, SIZET2NUM((size_t)cme));
-        rb_funcall(rb_mMJITHooks, rb_intern("on_cme_invalidate"), 1, cme_ptr);
+        rb_funcall(rb_mMJITHooks, rb_intern("on_cme_invalidate"), 1, SIZET2NUM((size_t)data));
     });
 }
 
@@ -367,10 +364,26 @@ rb_mjit_before_ractor_spawn(void)
     mjit_call_p = false;
 }
 
+void
+rb_mjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+
+    RB_VM_LOCK_ENTER();
+    rb_vm_barrier();
+
+    WITH_MJIT_ISOLATED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_constant_ic_update"), 3,
+                   SIZET2NUM((size_t)iseq), SIZET2NUM((size_t)ic), UINT2NUM(insn_idx));
+    });
+
+    RB_VM_LOCK_LEAVE();
+}
+
 void
 rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events)
 {
-    if (!mjit_call_p) return;
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
     WITH_MJIT_ISOLATED({
         rb_funcall(rb_mMJITHooks, rb_intern("on_tracing_invalidate_all"), 1, UINT2NUM(new_iseq_events));
     });

vm_insnhelper.c

@@ -5530,6 +5530,7 @@ vm_ic_update(const rb_iseq_t *iseq, IC ic, VALUE val, const VALUE *reg_ep, const
     RUBY_ASSERT(pc >= ISEQ_BODY(iseq)->iseq_encoded);
     unsigned pos = (unsigned)(pc - ISEQ_BODY(iseq)->iseq_encoded);
     rb_yjit_constant_ic_update(iseq, ic, pos);
+    rb_mjit_constant_ic_update(iseq, ic, pos);
 }
 
 static VALUE