From 21a63450236af6eaae954b367385c071231dc03f Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Mon, 5 Apr 2021 17:59:25 -0400 Subject: [PATCH] YJIT: adjust branch shape properly when target already exists The old code decides branch->shape based on the write position of the native code block, which is unsound in case the block already exists and no new code is written to the write position. Make this decision with the start address of the target block instead. Also handle when the branch becomes smaller after patching. --- yjit_core.c | 31 ++++++++++++++++--------------- yjit_core.h | 2 +- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/yjit_core.c b/yjit_core.c index ec793a053a..70c7da3291 100644 --- a/yjit_core.c +++ b/yjit_core.c @@ -418,7 +418,7 @@ uint8_t* gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx, rb_execution_ // Called by the generated code when a branch stub is executed // Triggers compilation of branches and code patching static uint8_t * -branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_context_t* ec) +branch_stub_hit(const uint32_t branch_idx, const uint32_t target_idx, rb_execution_context_t* ec) { uint8_t* dst_addr; @@ -449,18 +449,6 @@ branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_context_t // may be out of sync in JITted code ec->cfp->pc = iseq_pc_at_idx(target.iseq, target.idx); - // If either of the target blocks will be placed next - if (cb->write_pos == branch->end_pos) - { - //fprintf(stderr, "target idx %d will be placed next\n", target_idx); - branch->shape = (uint8_t)target_idx; - - // Rewrite the branch with the new, potentially more compact shape - cb_set_pos(cb, branch->start_pos); - branch->gen_fn(cb, branch->dst_addrs[0], branch->dst_addrs[1], branch->shape); - RUBY_ASSERT(cb->write_pos <= branch->end_pos); - } - // Try to find an existing compiled version of this block block_t* p_block = find_block_version(target, target_ctx); @@ -487,13 +475,26 @@ branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_context_t dst_addr = cb_get_ptr(cb, p_block->start_pos); branch->dst_addrs[target_idx] = dst_addr; + // Adjust brach shape base on block placement relative to the branch + if (branch->end_pos == p_block->start_pos) { + branch->shape = (branch_shape_t)target_idx; + } + // Rewrite the branch with the new jump target address RUBY_ASSERT(branch->dst_addrs[0] != NULL); uint32_t cur_pos = cb->write_pos; cb_set_pos(cb, branch->start_pos); branch->gen_fn(cb, branch->dst_addrs[0], branch->dst_addrs[1], branch->shape); - RUBY_ASSERT(cb->write_pos <= branch->end_pos); - branch->end_pos = cb->write_pos; + RUBY_ASSERT(cb->write_pos <= branch->end_pos && "can't enlarge a branch"); + + // If the branch got smaller + if (cb->write_pos < branch->end_pos) { + // fill the difference with nops + uint32_t shrinkage = branch->end_pos - cb->write_pos; + nop(cb, shrinkage); + } + + // Done patching the branch. Restore write position. cb_set_pos(cb, cur_pos); // Restore interpreter sp, since the code hitting the stub expects the original. diff --git a/yjit_core.h b/yjit_core.h index 1a55408a7b..658a0cc691 100644 --- a/yjit_core.h +++ b/yjit_core.h @@ -156,7 +156,7 @@ typedef void (*branchgen_fn)(codeblock_t* cb, uint8_t* target0, uint8_t* target1 Store info about an outgoing branch in a code segment Note: care must be taken to minimize the size of branch_t objects */ -typedef struct BranchEntry +typedef struct yjit_branch_entry { // Positions where the generated code starts and ends uint32_t start_pos;