Implement defer_compilation

lib/ruby_vm/mjit/assembler.rb

@@ -14,9 +14,7 @@ module RubyVM::MJIT
     Rel32Pad = Object.new
 
     # A set of ModR/M values encoded on #insn
-    class ModRM < Data.define(:mod, :reg, :rm)
+    class ModRM < Data.define(:mod, :reg, :rm); end
-      def initialize(mod:, reg: nil, rm: nil) = super
-    end
     Mod00 = 0b00 # Mod 00: [reg]
     Mod01 = 0b01 # Mod 01: [reg]+disp8
     Mod10 = 0b10 # Mod 10: [reg]+disp16
@@ -33,9 +31,11 @@ module RubyVM::MJIT
       @labels = {}
       @label_id = 0
       @comments = Hash.new { |h, k| h[k] = [] }
+      @stubs = Hash.new { |h, k| h[k] = [] }
     end
 
     def assemble(addr)
+      set_stub_addrs(addr)
       resolve_rel32(addr)
       resolve_labels
 
@@ -67,7 +67,7 @@ module RubyVM::MJIT
         insn(
           prefix: REX_W,
           opcode: 0x83,
-          mod_rm: ModRM[mod: Mod11, rm: dst_reg],
+          mod_rm: ModRM[mod: Mod11, reg: 0, rm: dst_reg],
           imm: imm8(src_imm),
         )
       # ADD r/m64, imm8 (Mod 00: [reg])
@@ -77,7 +77,7 @@ module RubyVM::MJIT
         insn(
           prefix: REX_W,
           opcode: 0x83,
-          mod_rm: ModRM[mod: Mod00, rm: dst_reg],
+          mod_rm: ModRM[mod: Mod00, reg: 0, rm: dst_reg],
           imm: imm8(src_imm),
         )
       else
@@ -85,12 +85,34 @@ module RubyVM::MJIT
       end
     end
 
+    # @param addr [Integer]
+    def call(addr)
+      # CALL rel32
+      # E8 cd
+      insn(opcode: 0xe8, imm: rel32(addr))
+    end
+
+    def jmp(dst)
+      case dst
+      # JMP rel32
+      in Integer => dst_addr
+        # E9 cd
+        insn(opcode: 0xe9, imm: rel32(dst_addr))
+      # JMP r/m64 (Mod 11: reg)
+      in Symbol => dst_reg
+        # FF /4
+        insn(opcode: 0xff, mod_rm: ModRM[mod: Mod11, reg: 4, rm: dst_reg])
+      else
+        raise NotImplementedError, "jmp: not-implemented operands: #{dst.inspect}"
+      end
+    end
+
     def jnz(dst)
       case dst
       # JNZ rel32
-      in Integer => addr
+      in Integer => dst_addr
         # 0F 85 cd
-        insn(opcode: [0x0f, 0x85], imm: rel32(addr))
+        insn(opcode: [0x0f, 0x85], imm: rel32(dst_addr))
       else
         raise NotImplementedError, "jnz: not-implemented operands: #{dst.inspect}"
       end
@@ -99,9 +121,9 @@ module RubyVM::MJIT
     def jz(dst)
       case dst
       # JZ rel8
-      in Label => label
+      in Label => dst_label
         # 74 cb
-        insn(opcode: 0x74, imm: label)
+        insn(opcode: 0x74, imm: dst_label)
       else
         raise NotImplementedError, "jz: not-implemented operands: #{dst.inspect}"
       end
@@ -155,7 +177,7 @@ module RubyVM::MJIT
           insn(
             prefix: REX_W,
             opcode: 0xc7,
-            mod_rm: ModRM[mod: Mod11, rm: dst_reg],
+            mod_rm: ModRM[mod: Mod11, reg: 0, rm: dst_reg],
             imm: imm32(src_imm),
           )
         # MOV r64, imm64
@@ -180,7 +202,7 @@ module RubyVM::MJIT
           insn(
             prefix: REX_W,
             opcode: 0xc7,
-            mod_rm: ModRM[mod: Mod00, rm: dst_reg],
+            mod_rm: ModRM[mod: Mod00, reg: 0, rm: dst_reg],
             imm: imm32(src_imm),
           )
         # MOV r/m64, r64 (Mod 00: [reg])
@@ -207,7 +229,7 @@ module RubyVM::MJIT
           insn(
             prefix: REX_W,
             opcode: 0xc7,
-            mod_rm: ModRM[mod: Mod01, rm: dst_reg],
+            mod_rm: ModRM[mod: Mod01, reg: 0, rm: dst_reg],
             disp: dst_disp,
             imm: imm32(src_imm),
           )
@@ -284,6 +306,10 @@ module RubyVM::MJIT
       @comments[@bytes.size] << message
     end
 
+    def stub(stub)
+      @stubs[@bytes.size] << stub
+    end
+
     def new_label(name)
       Label.new(id: @label_id += 1, name:)
     end
@@ -314,8 +340,8 @@ module RubyVM::MJIT
         opcode += reg_code(rd)
       end
       if mod_rm
-        prefix |= REX_R if mod_rm.reg && extended_reg?(mod_rm.reg)
+        prefix |= REX_R if mod_rm.reg.is_a?(Symbol) && extended_reg?(mod_rm.reg)
-        prefix |= REX_B if mod_rm.rm && extended_reg?(mod_rm.rm)
+        prefix |= REX_B if mod_rm.rm.is_a?(Symbol) && extended_reg?(mod_rm.rm)
       end
 
       # Encode insn
@@ -326,8 +352,8 @@ module RubyVM::MJIT
       if mod_rm
         mod_rm_byte = encode_mod_rm(
           mod: mod_rm.mod,
-          reg: mod_rm.reg ? reg_code(mod_rm.reg) : 0,
+          reg: mod_rm.reg.is_a?(Symbol) ? reg_code(mod_rm.reg) : mod_rm.reg,
-          rm: mod_rm.rm ? reg_code(mod_rm.rm) : 0,
+          rm: mod_rm.rm.is_a?(Symbol) ? reg_code(mod_rm.rm) : mod_rm.rm,
         )
         @bytes.push(mod_rm_byte)
       end
@@ -413,7 +439,7 @@ module RubyVM::MJIT
       unless imm32?(imm)
         raise ArgumentError, "unexpected imm32: #{imm}"
       end
-      imm_bytes(imm, 4)
+      [imm].pack('l').unpack('c*')
     end
 
     # io: 8 bytes
@@ -465,6 +491,15 @@ module RubyVM::MJIT
       [Rel32.new(addr), Rel32Pad, Rel32Pad, Rel32Pad]
     end
 
+    def set_stub_addrs(write_addr)
+      @bytes.each_with_index do |byte, index|
+        @stubs.fetch(index, []).each do |stub|
+          stub.addr = write_addr + index
+          stub.freeze
+        end
+      end
+    end
+
     def resolve_rel32(write_addr)
       @bytes.each_with_index do |byte, index|
         if byte.is_a?(Rel32)

lib/ruby_vm/mjit/block_stub.rb

@@ -0,0 +1,7 @@
+class RubyVM::MJIT::BlockStub < Struct.new(
+  :iseq, # @param [RubyVM::MJIT::CPointer::Struct_rb_iseq_struct] Jump target ISEQ
+  :pc,   # @param [Integer] Jump target pc
+  :ctx,  # @param [RubyVM::MJIT::Context] Jump target context
+  :addr, # @param [Integer] Jump source address to be re-generated
+)
+end

lib/ruby_vm/mjit/code_block.rb

@@ -37,6 +37,15 @@ module RubyVM::MJIT
       start_addr
     end
 
+    def with_addr(addr)
+      old_write_pos = @write_pos
+      @write_pos = addr - @mem_block
+      @comments.delete(addr) # TODO: clean up old comments for all overwritten insns?
+      yield
+    ensure
+      @write_pos = old_write_pos
+    end
+
     private
 
     def write_addr

lib/ruby_vm/mjit/compiler.rb

@@ -1,4 +1,5 @@
 require 'ruby_vm/mjit/assembler'
+require 'ruby_vm/mjit/block_stub'
 require 'ruby_vm/mjit/code_block'
 require 'ruby_vm/mjit/context'
 require 'ruby_vm/mjit/exit_compiler'
@@ -38,31 +39,56 @@ module RubyVM::MJIT
       @insn_compiler = InsnCompiler.new(@ocb)
     end
 
-    # @param iseq [RubyVM::MJIT::CPointer::Struct]
+    # Compile an ISEQ from its entry point.
-    def compile(iseq)
+    # @param iseq `RubyVM::MJIT::CPointer::Struct_rb_iseq_t`
+    # @param cfp `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t`
+    def compile(iseq, cfp)
       # TODO: Support has_opt
       return if iseq.body.param.flags.has_opt
 
       asm = Assembler.new
-      asm.comment("Block: #{iseq.body.location.label}@#{pathobj_path(iseq.body.location.pathobj)}:#{iseq.body.location.first_lineno}")
+      asm.comment("Block: #{iseq.body.location.label}@#{C.rb_iseq_path(iseq)}:#{iseq.body.location.first_lineno}")
       compile_prologue(asm)
-      compile_block(asm, iseq)
+      compile_block(asm, jit: JITState.new(iseq:, cfp:))
       iseq.body.jit_func = @cb.write(asm)
     rescue Exception => e
       $stderr.puts e.full_message # TODO: check verbose
     end
 
+    # Continue compilation from a stub.
+    # @param stub [RubyVM::MJIT::BlockStub]
+    # @param cfp `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t`
+    # @return [Integer] The starting address of a compiled stub
+    def stub_hit(stub, cfp)
+      # Update cfp->pc for `jit.at_current_insn?`
+      cfp.pc = stub.pc
+
+      # Compile the jump target
+      new_addr = Assembler.new.then do |asm|
+        jit = JITState.new(iseq: stub.iseq, cfp:)
+        index = (stub.pc - stub.iseq.body.iseq_encoded.to_i) / C.VALUE.size
+        compile_block(asm, jit:, index:, ctx: stub.ctx)
+        @cb.write(asm)
+      end
+
+      # Re-generate the jump source
+      @cb.with_addr(stub.addr) do
+        asm = Assembler.new
+        asm.comment('regenerate block stub')
+        asm.jmp(new_addr)
+        @cb.write(asm)
+      end
+      new_addr
+    end
+
     private
 
-    #  ec: rdi
-    # cfp: rsi
-    #
     # Callee-saved: rbx, rsp, rbp, r12, r13, r14, r15
     # Caller-saved: rax, rdi, rsi, rdx, rcx, r8, r9, r10, r11
     #
     # @param asm [RubyVM::MJIT::Assembler]
     def compile_prologue(asm)
-      asm.comment("MJIT entry")
+      asm.comment('MJIT entry')
 
       # Save callee-saved registers used by JITed code
       asm.push(CFP)
@@ -78,11 +104,8 @@ module RubyVM::MJIT
     end
 
     # @param asm [RubyVM::MJIT::Assembler]
-    def compile_block(asm, iseq)
+    def compile_block(asm, jit:, index: 0, ctx: Context.new)
-      jit = JITState.new
+      iseq = jit.iseq
-      ctx = Context.new
-
-      index = 0
       while index < iseq.body.iseq_size
         insn = self.class.decode_insn(iseq.body.iseq_encoded[index])
         jit.pc = (iseq.body.iseq_encoded + index).to_i
@@ -181,7 +204,7 @@ module RubyVM::MJIT
       # opt_mod
       # opt_eq
       # opt_neq
-      # opt_lt
+      when :opt_lt then @insn_compiler.opt_lt(jit, ctx, asm)
       # opt_le
       # opt_gt
       # opt_ge

lib/ruby_vm/mjit/context.rb

@@ -1,8 +1,15 @@
 class RubyVM::MJIT::Context < Struct.new(
-  :stack_size, # @param [Integer]
+  :stack_size, # @param [Integer] The number of values on the stack
+  :sp_offset,  # @param [Integer] JIT sp offset relative to the interpreter's sp
 )
   def initialize(*)
     super
     self.stack_size ||= 0
+    self.sp_offset ||= 0
+  end
+
+  def stack_push(size)
+    self.stack_size += size
+    self.sp_offset += size
   end
 end

lib/ruby_vm/mjit/exit_compiler.rb

@@ -1,6 +1,9 @@
 module RubyVM::MJIT
   class ExitCompiler
-    def initialize = freeze
+    def initialize
+      # TODO: Use GC offsets
+      @gc_refs = []
+    end
 
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
@@ -12,19 +15,12 @@ module RubyVM::MJIT
         asm.mov(:rax, (C.mjit_insn_exits + insn.bin).to_i)
         asm.add([:rax], 1) # TODO: lock
       end
-      asm.comment("exit to interpreter")
 
-      # Update pc
+      # Fix pc/sp offsets for the interpreter
-      asm.mov(:rax, jit.pc) # rax = jit.pc
+      save_pc_and_sp(jit, ctx, asm)
-      asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax) # cfp->pc = rax
-
-      # Update sp
-      if ctx.stack_size > 0
-        asm.add(SP, C.VALUE.size * ctx.stack_size) # rbx += stack_size
-        asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP) # cfp->sp = rbx
-      end
 
       # Restore callee-saved registers
+      asm.comment('exit to interpreter')
       asm.pop(SP)
       asm.pop(EC)
       asm.pop(CFP)
@@ -32,5 +28,48 @@ module RubyVM::MJIT
       asm.mov(:rax, Qundef)
       asm.ret
     end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param asm [RubyVM::MJIT::Assembler]
+    # @param stub [RubyVM::MJIT::BlockStub]
+    def compile_jump_stub(jit, asm, stub)
+      case stub
+      when BlockStub
+        asm.comment("block stub hit: #{stub.iseq.body.location.label}@#{C.rb_iseq_path(stub.iseq)}:#{stub.iseq.body.location.first_lineno}")
+      else
+        raise "unexpected stub object: #{stub.inspect}"
+      end
+
+      # Call rb_mjit_stub_hit
+      asm.mov(:rdi, to_value(stub))
+      asm.call(C.rb_mjit_stub_hit)
+
+      # Jump to the address returned by rb_mjit_stub_hit
+      asm.jmp(:rax)
+    end
+
+    private
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def save_pc_and_sp(jit, ctx, asm)
+      # Update pc
+      asm.comment("save pc #{'and sp' if ctx.sp_offset != 0}")
+      asm.mov(:rax, jit.pc) # rax = jit.pc
+      asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax) # cfp->pc = rax
+
+      # Update sp
+      if ctx.sp_offset != 0
+        asm.add(SP, C.VALUE.size * ctx.sp_offset) # sp += stack_size
+        asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP) # cfp->sp = sp
+        ctx.sp_offset = 0
+      end
+    end
+
+    def to_value(obj)
+      @gc_refs << obj
+      C.to_value(obj)
+    end
   end
 end

lib/ruby_vm/mjit/insn_compiler.rb

@@ -1,7 +1,7 @@
 module RubyVM::MJIT
   # scratch regs: rax
   #
-  # 4/101
+  # 5/101
   class InsnCompiler
     # @param ocb [CodeBlock]
     def initialize(ocb)
@@ -33,7 +33,7 @@ module RubyVM::MJIT
     # @param asm [RubyVM::MJIT::Assembler]
     def putnil(jit, ctx, asm)
       asm.mov([SP, C.VALUE.size * ctx.stack_size], Qnil)
-      ctx.stack_size += 1
+      ctx.stack_push(1)
       KeepCompiling
     end
 
@@ -55,7 +55,7 @@ module RubyVM::MJIT
         asm.mov([SP, C.VALUE.size * ctx.stack_size], :rax)
       end
 
-      ctx.stack_size += 1
+      ctx.stack_push(1)
       KeepCompiling
     end
 
@@ -141,7 +141,18 @@ module RubyVM::MJIT
     # opt_mod
     # opt_eq
     # opt_neq
-    # opt_lt
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def opt_lt(jit, ctx, asm)
+      unless jit.at_current_insn?
+        defer_compilation(jit, ctx, asm)
+        return EndBlock
+      end
+      CantCompile
+    end
+
     # opt_le
     # opt_gt
     # opt_ge
@@ -178,7 +189,7 @@ module RubyVM::MJIT
 
       # Push it to the stack
       asm.mov([SP, C.VALUE.size * ctx.stack_size], :rax)
-      ctx.stack_size += 1
+      ctx.stack_push(1)
       KeepCompiling
     end
 
@@ -196,6 +207,27 @@ module RubyVM::MJIT
       end
     end
 
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def defer_compilation(jit, ctx, asm)
+      # Make a stub to compile the current insn
+      block_stub = BlockStub.new(
+        iseq: jit.iseq,
+        pc:   jit.pc,
+        ctx:  ctx.dup,
+      )
+
+      stub_hit = Assembler.new.then do |ocb_asm|
+        @exit_compiler.compile_jump_stub(jit, ocb_asm, block_stub)
+        @ocb.write(ocb_asm)
+      end
+
+      asm.comment('defer_compilation: block stub')
+      asm.stub(block_stub)
+      asm.jmp(stub_hit)
+    end
+
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     def compile_side_exit(jit, ctx)

lib/ruby_vm/mjit/jit_state.rb

@@ -1,9 +1,15 @@
 module RubyVM::MJIT
   class JITState < Struct.new(
-    :pc, # @param [Integer]
+    :iseq,
+    :pc,  # @param [Integer] The JIT target PC
+    :cfp, # @param `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t` The JIT source CFP (before MJIT is called)
   )
     def operand(index)
       C.VALUE.new(pc)[index + 1]
     end
+
+    def at_current_insn?
+      pc == cfp.pc.to_i
+    end
   end
 end

mjit.c

@@ -130,6 +130,8 @@ static VALUE rb_mMJITC = 0;
 static VALUE rb_MJITCompiler = 0;
 // RubyVM::MJIT::CPointer::Struct_rb_iseq_t
 static VALUE rb_cMJITIseqPtr = 0;
+// RubyVM::MJIT::CPointer::Struct_rb_control_frame_t
+static VALUE rb_cMJITCfpPtr = 0;
 
 void
 rb_mjit_add_iseq_to_process(const rb_iseq_t *iseq)
@@ -334,13 +336,43 @@ rb_mjit_compile(const rb_iseq_t *iseq)
     mjit_stats_p = false; // Avoid impacting JIT stats by itself
 
     VALUE iseq_ptr = rb_funcall(rb_cMJITIseqPtr, rb_intern("new"), 1, SIZET2NUM((size_t)iseq));
-    rb_funcall(rb_MJITCompiler, rb_intern("compile"), 1, iseq_ptr);
+    VALUE cfp_ptr = rb_funcall(rb_cMJITCfpPtr, rb_intern("new"), 1, SIZET2NUM((size_t)GET_EC()->cfp));
+    rb_funcall(rb_MJITCompiler, rb_intern("compile"), 2, iseq_ptr, cfp_ptr);
 
     mjit_stats_p = mjit_opts.stats;
     mjit_call_p = original_call_p;
     RB_VM_LOCK_LEAVE();
 }
 
+void *
+rb_mjit_stub_hit(VALUE branch_stub)
+{
+    VALUE result;
+
+    RB_VM_LOCK_ENTER();
+    rb_vm_barrier();
+    bool original_call_p = mjit_call_p;
+    mjit_call_p = false; // Avoid impacting JIT metrics by itself
+    mjit_stats_p = false; // Avoid impacting JIT stats by itself
+
+    rb_control_frame_t *cfp = GET_EC()->cfp;
+    // Given JIT's SP offset, temporarily update SP to preserve stack values.
+    // It's reset afterwards for consistency with the code without this stub.
+    unsigned int stack_max = cfp->iseq->body->stack_max;
+    cfp->sp += stack_max;
+
+    VALUE cfp_ptr = rb_funcall(rb_cMJITCfpPtr, rb_intern("new"), 1, SIZET2NUM((size_t)cfp));
+    result = rb_funcall(rb_MJITCompiler, rb_intern("stub_hit"), 2, branch_stub, cfp_ptr);
+
+    cfp->sp -= stack_max;
+
+    mjit_stats_p = mjit_opts.stats;
+    mjit_call_p = original_call_p;
+    RB_VM_LOCK_LEAVE();
+
+    return (void *)NUM2SIZET(result);
+}
+
 // Called by rb_vm_mark()
 void
 mjit_mark(void)
@@ -352,6 +384,7 @@ mjit_mark(void)
     // Mark objects used by the MJIT compiler
     rb_gc_mark(rb_MJITCompiler);
     rb_gc_mark(rb_cMJITIseqPtr);
+    rb_gc_mark(rb_cMJITCfpPtr);
 
     RUBY_MARK_LEAVE("mjit");
 }
@@ -377,6 +410,7 @@ mjit_init(const struct mjit_options *opts)
     rb_MJITCompiler = rb_funcall(rb_cMJITCompiler, rb_intern("new"), 2,
                                  SIZET2NUM((size_t)rb_mjit_mem_block), UINT2NUM(MJIT_CODE_SIZE));
     rb_cMJITIseqPtr = rb_funcall(rb_mMJITC, rb_intern("rb_iseq_t"), 0);
+    rb_cMJITCfpPtr = rb_funcall(rb_mMJITC, rb_intern("rb_control_frame_t"), 0);
 
     mjit_call_p = true;
     mjit_stats_p = mjit_opts.stats;

mjit_c.rb

@@ -36,6 +36,13 @@ module RubyVM::MJIT # :nodoc: all
       CType::Immediate.parse("size_t").new(addr)
     end
 
+    def rb_mjit_stub_hit
+      Primitive.cstmt! %{
+        extern void *rb_mjit_stub_hit(VALUE stub);
+        return SIZET2NUM((size_t)rb_mjit_stub_hit);
+      }
+    end
+
     def rb_mjit_counters
       addr = Primitive.cstmt! %{
         #if MJIT_STATS
@@ -53,6 +60,11 @@ module RubyVM::MJIT # :nodoc: all
       Primitive.dump_disasm(from, to)
     end
 
+    # Convert a Ruby object to a VALUE in Integer
+    def to_value(obj)
+      Primitive.cexpr! 'SIZET2NUM((size_t)obj)'
+    end
+
     #========================================================================================
     #
     # Old stuff