1berry/ruby
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Implement branch stub

Browse Source
This commit is contained in:
Takashi Kokubun 2023-01-07 13:21:14 -08:00
parent eddec7bc20
commit 62d36dd127
10 changed files with 256 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
lib/ruby_vm/mjit/assembler.rb
View File

@ -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

2
lib/ruby_vm/mjit/block_stub.rb
View File

@ -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
)

14
lib/ruby_vm/mjit/branch_stub.rb Normal file
View File

@ -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

6
lib/ruby_vm/mjit/code_block.rb
View File

@ -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)

102
lib/ruby_vm/mjit/compiler.rb
View File

@ -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.
# @param stub [RubyVM::MJIT::BlockStub]
# Continue compilation from a block stub.
# @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
def stub_hit(stub, cfp)
# @return [Integer] The starting address of the compiled block stub
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:)
compile_block(asm, jit:, pc: stub.pc, ctx: stub.ctx)
jit = JITState.new(iseq: block_stub.iseq, cfp:)
compile_block(asm, jit:, pc: block_stub.pc, ctx: block_stub.ctx)
end
# Rewrite the stub
if @cb.write_addr == stub.end_addr
# If the stub jump is the last code, overwrite the jump with the new code.
@cb.set_write_addr(stub.start_addr)
# Rewrite the block stub
if @cb.write_addr == block_stub.end_addr
# If the block stub's jump is the last code, overwrite the jump with the new code.
@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

39
lib/ruby_vm/mjit/exit_compiler.rb
View File

@ -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]
def compile_jump_stub(jit, ctx, 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))
# @param block_stub [RubyVM::MJIT::BlockStub]
def compile_block_stub(jit, ctx, asm, block_stub)
# Call rb_mjit_block_stub_hit
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.mov(:rdi, to_value(block_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

64
lib/ruby_vm/mjit/insn_compiler.rb
View File

@ -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

4
lib/ruby_vm/mjit/jit_state.rb
View File

@ -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

30
mjit.c
View File

@ -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

18
mjit_c.rb
View File

@ -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);
return SIZET2NUM((size_t)rb_mjit_stub_hit);
extern void *rb_mjit_block_stub_hit(VALUE block_stub, int sp_offset);
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