Merge pull request ruby#96 from Shopify/yjit-thread-bug

maximecb · web-flow · commit ff12b7a0e1e1 · 2021-04-06T10:16:45.000-04:00
YJIT: adjust branch shape properly when target already exists
diff --git 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 based 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
@@ -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;

Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,7 @@ typedef void (branchgen_fn)(codeblock_t cb, uint8_t* target0, uint8_t* target1`
`156`	`156`	`Store info about an outgoing branch in a code segment`
`157`	`157`	`Note: care must be taken to minimize the size of branch_t objects`
`158`	`158`	`*/`
`159`		`-typedef struct BranchEntry`
	`159`	`+typedef struct yjit_branch_entry`
`160`	`160`	`{`
`161`	`161`	`// Positions where the generated code starts and ends`
`162`	`162`	`uint32_t start_pos;`