Implement branch stub

lib/ruby_vm/mjit/assembler.rb

@@ -169,6 +169,10 @@ module RubyVM::MJIT
       in Label => dst_label
         # 74 cb
         insn(opcode: 0x74, imm: dst_label)
+      # JZ rel32
+      in Integer => dst_addr
+        # 0F 84 cd
+        insn(opcode: [0x0f, 0x84], imm: rel32(dst_addr))
       else
         raise NotImplementedError, "jz: not-implemented operands: #{dst.inspect}"
       end
@@ -338,7 +342,7 @@ module RubyVM::MJIT
     def test(left, right)
       case [left, right]
       # TEST r/m8*, imm8 (Mod 01: [reg]+disp8)
-      in [[Symbol => left_reg, Integer => left_disp], Integer => right_imm] if imm8?(right_imm)
+      in [[Symbol => left_reg, Integer => left_disp], Integer => right_imm] if imm8?(right_imm) && right_imm >= 0
         # REX + F6 /0 ib
         # MI: Operand 1: ModRM:r/m (r), Operand 2: imm8/16/32
         insn(
@@ -347,6 +351,17 @@ module RubyVM::MJIT
           disp: left_disp,
           imm: imm8(right_imm),
         )
+      # TEST r/m64, imm32 (Mod 01: [reg]+disp8)
+      in [[Symbol => left_reg, Integer => left_disp], Integer => right_imm] if imm32?(right_imm)
+        # REX.W + F7 /0 id
+        # MI: Operand 1: ModRM:r/m (r), Operand 2: imm8/16/32
+        insn(
+          prefix: REX_W,
+          opcode: 0xf7,
+          mod_rm: ModRM[mod: Mod01, reg: 0, rm: left_reg],
+          disp: left_disp,
+          imm: imm32(right_imm),
+        )
       # TEST r/m32, r32 (Mod 11: reg)
       in [Symbol => left_reg, Symbol => right_reg] if r32?(left_reg) && r32?(right_reg)
         # 85 /r
@@ -574,7 +589,6 @@ module RubyVM::MJIT
         end
         @stub_ends.fetch(index, []).each do |stub|
           stub.end_addr = write_addr + index
-          stub.freeze
         end
       end
     end

lib/ruby_vm/mjit/block_stub.rb

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

lib/ruby_vm/mjit/branch_stub.rb

@@ -0,0 +1,14 @@
+class RubyVM::MJIT::BranchStub < Struct.new(
+  :iseq,               # @param [RubyVM::MJIT::CPointer::Struct_rb_iseq_struct] Branch target ISEQ
+  :ctx,                # @param [RubyVM::MJIT::Context] Branch target context
+  :branch_target_pc,   # @param [Integer] Branch target PC
+  :branch_target_addr, # @param [Integer] Branch target address
+  :branch_target_next, # @param [Proc] Compile branch target next
+  :fallthrough_pc,     # @param [Integer] Fallthrough PC
+  :fallthrough_addr,   # @param [Integer] Fallthrough address
+  :fallthrough_next,   # @param [Proc] Compile fallthrough next
+  :neither_next,       # @param [Proc] Compile neither branch target nor fallthrough next
+  :start_addr,         # @param [Integer] Stub source start address to be re-generated
+  :end_addr,           # @param [Integer] Stub source end address to be re-generated
+)
+end

lib/ruby_vm/mjit/code_block.rb

@@ -44,7 +44,7 @@ module RubyVM::MJIT
 
     def with_write_addr(addr)
       old_write_pos = @write_pos
-      set_addr(addr)
+      set_write_addr(addr)
       yield
     ensure
       @write_pos = old_write_pos
@@ -54,6 +54,10 @@ module RubyVM::MJIT
       @mem_block + @write_pos
     end
 
+    def include?(addr)
+      (@mem_block...(@mem_block + @mem_size)).include?(addr)
+    end
+
     private
 
     def dump_disasm(from, to)

lib/ruby_vm/mjit/compiler.rb

@@ -1,6 +1,7 @@
 require 'ruby_vm/mjit/assembler'
 require 'ruby_vm/mjit/block'
 require 'ruby_vm/mjit/block_stub'
+require 'ruby_vm/mjit/branch_stub'
 require 'ruby_vm/mjit/code_block'
 require 'ruby_vm/mjit/context'
 require 'ruby_vm/mjit/exit_compiler'
@@ -61,29 +62,29 @@ module RubyVM::MJIT
       $stderr.puts e.full_message # TODO: check verbose
     end
 
-    # Continue compilation from a stub.
+    # Continue compilation from a block stub.
-    # @param stub [RubyVM::MJIT::BlockStub]
+    # @param block_stub [RubyVM::MJIT::BlockStub]
     # @param cfp `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t`
-    # @return [Integer] The starting address of a compiled stub
+    # @return [Integer] The starting address of the compiled block stub
-    def stub_hit(stub, cfp)
+    def block_stub_hit(block_stub, cfp)
       # Update cfp->pc for `jit.at_current_insn?`
-      cfp.pc = stub.pc
+      cfp.pc = block_stub.pc
 
       # Prepare the jump target
       new_asm = Assembler.new.tap do |asm|
-        jit = JITState.new(iseq: stub.iseq, cfp:)
+        jit = JITState.new(iseq: block_stub.iseq, cfp:)
-        compile_block(asm, jit:, pc: stub.pc, ctx: stub.ctx)
+        compile_block(asm, jit:, pc: block_stub.pc, ctx: block_stub.ctx)
       end
 
-      # Rewrite the stub
+      # Rewrite the block stub
-      if @cb.write_addr == stub.end_addr
+      if @cb.write_addr == block_stub.end_addr
-        # If the stub jump is the last code, overwrite the jump with the new code.
+        # If the block stub's jump is the last code, overwrite the jump with the new code.
-        @cb.set_write_addr(stub.start_addr)
+        @cb.set_write_addr(block_stub.start_addr)
         @cb.write(new_asm)
       else
-        # If the stub jump is old code, change the jump target to the new code.
+        # If the block stub's jump is old code, change the jump target to the new code.
         new_addr = @cb.write(new_asm)
-        @cb.with_write_addr(stub.start_addr) do
+        @cb.with_write_addr(block_stub.start_addr) do
           asm = Assembler.new
           asm.comment('regenerate block stub')
           asm.jmp(new_addr)
@@ -92,6 +93,74 @@ module RubyVM::MJIT
       end
     end
 
+    # Compile a branch stub.
+    # @param branch_stub [RubyVM::MJIT::BranchStub]
+    # @param cfp `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t`
+    # @param branch_target_p [TrueClass,FalseClass]
+    # @return [Integer] The starting address of the compiled branch stub
+    def branch_stub_hit(branch_stub, cfp, branch_target_p)
+      # Update cfp->pc for `jit.at_current_insn?`
+      pc = branch_target_p ? branch_stub.branch_target_pc : branch_stub.fallthrough_pc
+      cfp.pc = pc
+
+      # Prepare the jump target
+      new_asm = Assembler.new.tap do |asm|
+        jit = JITState.new(iseq: branch_stub.iseq, cfp:)
+        compile_block(asm, jit:, pc:, ctx: branch_stub.ctx.dup)
+      end
+
+      # Rewrite the branch stub
+      if @cb.write_addr == branch_stub.end_addr
+        # If the branch stub's jump is the last code, overwrite the jump with the new code.
+        @cb.set_write_addr(branch_stub.start_addr)
+        Assembler.new.tap do |branch_asm|
+          if branch_target_p
+            branch_stub.branch_target_next.call(branch_asm)
+          else
+            branch_stub.fallthrough_next.call(branch_asm)
+          end
+          @cb.write(branch_asm)
+        end
+
+        # Compile a fallthrough over the jump
+        if branch_target_p
+          branch_stub.branch_target_addr = @cb.write(new_asm)
+        else
+          branch_stub.fallthrough_addr = @cb.write(new_asm)
+        end
+      else
+        # Otherwise, just prepare the new code somewhere
+        if branch_target_p
+          unless @cb.include?(branch_stub.branch_target_addr)
+            branch_stub.branch_target_addr = @cb.write(new_asm)
+          end
+        else
+          unless @cb.include?(branch_stub.fallthrough_addr)
+            branch_stub.fallthrough_addr = @cb.write(new_asm)
+          end
+        end
+
+        # Update jump destinations
+        branch_asm = Assembler.new
+        if branch_stub.end_addr == branch_stub.branch_target_addr # branch_target_next has been used
+          branch_stub.branch_target_next.call(branch_asm)
+        elsif branch_stub.end_addr == branch_stub.fallthrough_addr # fallthrough_next has been used
+          branch_stub.fallthrough_next.call(branch_asm)
+        else
+          branch_stub.neither_next.call(branch_asm)
+        end
+        @cb.with_write_addr(branch_stub.start_addr) do
+          @cb.write(branch_asm)
+        end
+      end
+
+      if branch_target_p
+        branch_stub.branch_target_addr
+      else
+        branch_stub.fallthrough_addr
+      end
+    end
+
     private
 
     # Callee-saved: rbx, rsp, rbp, r12, r13, r14, r15
@@ -127,15 +196,18 @@ module RubyVM::MJIT
         insn = self.class.decode_insn(iseq.body.iseq_encoded[index])
         jit.pc = (iseq.body.iseq_encoded + index).to_i
 
-        case @insn_compiler.compile(jit, ctx, asm, insn)
+        case status = @insn_compiler.compile(jit, ctx, asm, insn)
+        when KeepCompiling
+          index += insn.len
         when EndBlock
           # TODO: pad nops if entry exit exists
           break
         when CantCompile
           @exit_compiler.compile_side_exit(jit, ctx, asm)
           break
+        else
+          raise "compiling #{insn.name} returned unexpected status: #{status.inspect}"
         end
-        index += insn.len
       end
     end
   end

lib/ruby_vm/mjit/exit_compiler.rb

@@ -47,19 +47,30 @@ module RubyVM::MJIT
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
-    # @param stub [RubyVM::MJIT::BlockStub]
+    # @param block_stub [RubyVM::MJIT::BlockStub]
-    def compile_jump_stub(jit, ctx, asm, stub)
+    def compile_block_stub(jit, ctx, asm, block_stub)
-      case stub
+      # Call rb_mjit_block_stub_hit
-      when BlockStub
+      asm.comment("block stub hit: #{block_stub.iseq.body.location.label}@#{C.rb_iseq_path(block_stub.iseq)}:#{iseq_lineno(block_stub.iseq, block_stub.pc)}")
-        asm.comment("block stub hit: #{stub.iseq.body.location.label}@#{C.rb_iseq_path(stub.iseq)}:#{stub.iseq.body.location.first_lineno}")
+      asm.mov(:rdi, to_value(block_stub))
-      else
-        raise "unexpected stub object: #{stub.inspect}"
-      end
-
-      # Call rb_mjit_stub_hit
-      asm.mov(:rdi, to_value(stub))
       asm.mov(:esi, ctx.sp_offset)
-      asm.call(C.rb_mjit_stub_hit)
+      asm.call(C.rb_mjit_block_stub_hit)
+
+      # Jump to the address returned by rb_mjit_stub_hit
+      asm.jmp(:rax)
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    # @param branch_stub [RubyVM::MJIT::BranchStub]
+    # @param branch_target_p [TrueClass,FalseClass]
+    def compile_branch_stub(jit, ctx, asm, branch_stub, branch_target_p)
+      # Call rb_mjit_branch_stub_hit
+      asm.comment("branch stub hit: #{branch_stub.iseq.body.location.label}@#{C.rb_iseq_path(branch_stub.iseq)}:#{iseq_lineno(branch_stub.iseq, branch_target_p ? branch_stub.branch_target_pc : branch_stub.fallthrough_pc)}")
+      asm.mov(:rdi, to_value(branch_stub))
+      asm.mov(:esi, ctx.sp_offset)
+      asm.mov(:edx, branch_target_p ? 1 : 0)
+      asm.call(C.rb_mjit_branch_stub_hit)
 
       # Jump to the address returned by rb_mjit_stub_hit
       asm.jmp(:rax)
@@ -103,5 +114,9 @@ module RubyVM::MJIT
       @gc_refs << obj
       C.to_value(obj)
     end
+
+    def iseq_lineno(iseq, pc)
+      C.rb_iseq_line_no(iseq, (pc - iseq.body.iseq_encoded.to_i) / C.VALUE.size)
+    end
   end
 end

lib/ruby_vm/mjit/insn_compiler.rb

@@ -17,7 +17,7 @@ module RubyVM::MJIT
       asm.incr_counter(:mjit_insns_count)
       asm.comment("Insn: #{insn.name}")
 
-      # 5/101
+      # 6/101
       case insn.name
       # nop
       # getlocal
@@ -83,7 +83,7 @@ module RubyVM::MJIT
       # throw
       # jump
       # branchif
-      # branchunless
+      when :branchunless then branchunless(jit, ctx, asm)
       # branchnil
       # once
       # opt_case_dispatch
@@ -247,7 +247,61 @@ module RubyVM::MJIT
     # throw
     # jump
     # branchif
-    # branchunless
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def branchunless(jit, ctx, asm)
+      # TODO: check ints for backward branches
+      # TODO: skip check for known truthy
+
+      # This `test` sets ZF only for Qnil and Qfalse, which let jz jump.
+      asm.test([SP, C.VALUE.size * (ctx.stack_size - 1)], ~Qnil)
+      ctx.stack_pop(1)
+
+      # Set stubs
+      # TODO: reuse already-compiled blocks
+      branch_stub = BranchStub.new(
+        iseq: jit.iseq,
+        ctx:  ctx.dup,
+        branch_target_pc: jit.pc + (jit.insn.len + jit.operand(0)) * C.VALUE.size,
+        fallthrough_pc:   jit.pc + jit.insn.len * C.VALUE.size,
+      )
+      branch_stub.branch_target_addr = Assembler.new.then do |ocb_asm|
+        @exit_compiler.compile_branch_stub(jit, ctx, ocb_asm, branch_stub, true)
+        @ocb.write(ocb_asm)
+      end
+      branch_stub.fallthrough_addr = Assembler.new.then do |ocb_asm|
+        @exit_compiler.compile_branch_stub(jit, ctx, ocb_asm, branch_stub, false)
+        @ocb.write(ocb_asm)
+      end
+
+      # Prepare codegen for all cases
+      branch_stub.branch_target_next = proc do |branch_asm|
+        branch_asm.stub(branch_stub) do
+          branch_asm.comment('branch_target_next')
+          branch_asm.jnz(branch_stub.fallthrough_addr)
+        end
+      end
+      branch_stub.fallthrough_next = proc do |branch_asm|
+        branch_asm.stub(branch_stub) do
+          branch_asm.comment('fallthrough_next')
+          branch_asm.jz(branch_stub.branch_target_addr)
+        end
+      end
+      branch_stub.neither_next = proc do |branch_asm|
+        branch_asm.stub(branch_stub) do
+          branch_asm.comment('neither_next')
+          branch_asm.jz(branch_stub.branch_target_addr)
+          branch_asm.jmp(branch_stub.fallthrough_addr)
+        end
+      end
+
+      # Just jump to stubs
+      branch_stub.neither_next.call(asm)
+      EndBlock
+    end
+
     # branchnil
     # once
     # opt_case_dispatch
@@ -370,12 +424,12 @@ module RubyVM::MJIT
       # Make a stub to compile the current insn
       block_stub = BlockStub.new(
         iseq: jit.iseq,
-        pc:   jit.pc,
         ctx:  ctx.dup,
+        pc:   jit.pc,
       )
 
       stub_hit = Assembler.new.then do |ocb_asm|
-        @exit_compiler.compile_jump_stub(jit, ctx, ocb_asm, block_stub)
+        @exit_compiler.compile_block_stub(jit, ctx, ocb_asm, block_stub)
         @ocb.write(ocb_asm)
       end
 

lib/ruby_vm/mjit/jit_state.rb

@@ -8,6 +8,10 @@ module RubyVM::MJIT
   )
     def initialize(side_exits: {}, **) = super
 
+    def insn
+      Compiler.decode_insn(C.VALUE.new(pc).*)
+    end
+
     def operand(index)
       C.VALUE.new(pc)[index + 1]
     end

mjit.c

@@ -345,7 +345,7 @@ rb_mjit_compile(const rb_iseq_t *iseq)
 }
 
 void *
-rb_mjit_stub_hit(VALUE branch_stub, int sp_offset)
+rb_mjit_block_stub_hit(VALUE block_stub, int sp_offset)
 {
     VALUE result;
 
@@ -359,7 +359,33 @@ rb_mjit_stub_hit(VALUE branch_stub, int sp_offset)
     cfp->sp += sp_offset; // preserve stack values, also using the actual sp_offset to make jit.peek_at_stack work
 
     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);
+    result = rb_funcall(rb_MJITCompiler, rb_intern("block_stub_hit"), 2, block_stub, cfp_ptr);
+
+    cfp->sp -= sp_offset; // reset for consistency with the code without the stub
+
+    mjit_stats_p = mjit_opts.stats;
+    mjit_call_p = original_call_p;
+    RB_VM_LOCK_LEAVE();
+
+    return (void *)NUM2SIZET(result);
+}
+
+void *
+rb_mjit_branch_stub_hit(VALUE branch_stub, int sp_offset, int branch_target_p)
+{
+    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;
+    cfp->sp += sp_offset; // preserve stack values, also using the actual sp_offset to make jit.peek_at_stack work
+
+    VALUE cfp_ptr = rb_funcall(rb_cMJITCfpPtr, rb_intern("new"), 1, SIZET2NUM((size_t)cfp));
+    result = rb_funcall(rb_MJITCompiler, rb_intern("branch_stub_hit"), 3, branch_stub, cfp_ptr, RBOOL(branch_target_p));
 
     cfp->sp -= sp_offset; // reset for consistency with the code without the stub
 

mjit_c.rb

@@ -36,10 +36,17 @@ module RubyVM::MJIT # :nodoc: all
       CType::Immediate.parse("size_t").new(addr)
     end
 
-    def rb_mjit_stub_hit
+    def rb_mjit_block_stub_hit
       Primitive.cstmt! %{
-        extern void *rb_mjit_stub_hit(VALUE stub);
+        extern void *rb_mjit_block_stub_hit(VALUE block_stub, int sp_offset);
-        return SIZET2NUM((size_t)rb_mjit_stub_hit);
+        return SIZET2NUM((size_t)rb_mjit_block_stub_hit);
+      }
+    end
+
+    def rb_mjit_branch_stub_hit
+      Primitive.cstmt! %{
+        extern void *rb_mjit_branch_stub_hit(VALUE branch_stub, int sp_offset, int branch_target_p);
+        return SIZET2NUM((size_t)rb_mjit_branch_stub_hit);
       }
     end
 
@@ -69,6 +76,11 @@ module RubyVM::MJIT # :nodoc: all
       Primitive.cexpr! 'RBOOL(BASIC_OP_UNREDEFINED_P(NUM2INT(op), NUM2INT(klass)))'
     end
 
+    def rb_iseq_line_no(iseq, pos)
+      _iseq_addr = iseq.to_i
+      Primitive.cexpr! 'UINT2NUM(rb_iseq_line_no((const rb_iseq_t *)NUM2SIZET(_iseq_addr), NUM2SIZET(pos)))'
+    end
+
     #========================================================================================
     #
     # Old stuff