Skip to content

Commit cb9d0e2

Browse files
authored
Merge pull request #2 from Shopify/yjit-landing-pad
Use a context-free landing pad to optimize `leave`
2 parents 51733e8 + 5ed796d commit cb9d0e2

File tree

1 file changed

+90
-72
lines changed

1 file changed

+90
-72
lines changed

yjit_codegen.c

Lines changed: 90 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ codeblock_t* cb = NULL;
2626
static codeblock_t outline_block;
2727
codeblock_t* ocb = NULL;
2828

29+
// Code for exiting back to the interpreter from the leave insn
30+
static void *leave_exit_code;
31+
2932
// Print the current source location for debugging purposes
3033
RBIMPL_ATTR_MAYBE_UNUSED()
3134
static void
@@ -112,6 +115,61 @@ jit_peek_at_self(jitstate_t *jit, ctx_t *ctx)
112115

113116
static bool jit_guard_known_klass(jitstate_t *jit, ctx_t* ctx, VALUE known_klass, insn_opnd_t insn_opnd, const int max_chain_depth, uint8_t *side_exit);
114117

118+
#if RUBY_DEBUG
119+
120+
// Increment a profiling counter with counter_name
121+
#define GEN_COUNTER_INC(cb, counter_name) _gen_counter_inc(cb, &(yjit_runtime_counters . counter_name))
122+
static void
123+
_gen_counter_inc(codeblock_t *cb, int64_t *counter)
124+
{
125+
if (!rb_yjit_opts.gen_stats) return;
126+
mov(cb, REG0, const_ptr_opnd(counter));
127+
cb_write_lock_prefix(cb); // for ractors.
128+
add(cb, mem_opnd(64, REG0, 0), imm_opnd(1));
129+
}
130+
131+
// Increment a counter then take an existing side exit.
132+
#define COUNTED_EXIT(side_exit, counter_name) _counted_side_exit(side_exit, &(yjit_runtime_counters . counter_name))
133+
static uint8_t *
134+
_counted_side_exit(uint8_t *existing_side_exit, int64_t *counter)
135+
{
136+
if (!rb_yjit_opts.gen_stats) return existing_side_exit;
137+
138+
uint8_t *start = cb_get_ptr(ocb, ocb->write_pos);
139+
_gen_counter_inc(ocb, counter);
140+
jmp_ptr(ocb, existing_side_exit);
141+
return start;
142+
}
143+
144+
// Add a comment at the current position in the code block
145+
static void
146+
_add_comment(codeblock_t* cb, const char* comment_str)
147+
{
148+
// Avoid adding duplicate comment strings (can happen due to deferred codegen)
149+
size_t num_comments = rb_darray_size(yjit_code_comments);
150+
if (num_comments > 0) {
151+
struct yjit_comment last_comment = rb_darray_get(yjit_code_comments, num_comments - 1);
152+
if (last_comment.offset == cb->write_pos && strcmp(last_comment.comment, comment_str) == 0) {
153+
return;
154+
}
155+
}
156+
157+
struct yjit_comment new_comment = (struct yjit_comment){ cb->write_pos, comment_str };
158+
rb_darray_append(&yjit_code_comments, new_comment);
159+
}
160+
161+
// Comments for generated machine code
162+
#define ADD_COMMENT(cb, comment) _add_comment((cb), (comment))
163+
yjit_comment_array_t yjit_code_comments;
164+
165+
#else
166+
167+
#define GEN_COUNTER_INC(cb, counter_name) ((void)0)
168+
#define COUNTED_EXIT(side_exit, counter_name) side_exit
169+
#define ADD_COMMENT(cb, comment) ((void)0)
170+
171+
#endif // if RUBY_DEBUG
172+
115173
// Save YJIT registers prior to a C call
116174
static void
117175
yjit_save_regs(codeblock_t* cb)
@@ -136,7 +194,7 @@ yjit_load_regs(codeblock_t* cb)
136194
static uint8_t *
137195
yjit_gen_exit(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
138196
{
139-
uint8_t *code_ptr = cb_get_ptr(ocb, ocb->write_pos);
197+
uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
140198

141199
VALUE *exit_pc = jit->pc;
142200

@@ -183,74 +241,38 @@ yjit_gen_exit(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
183241
return code_ptr;
184242
}
185243

186-
187-
// A shorthand for generating an exit in the outline block
244+
// Generate a continuation for gen_leave() that exits to the interpreter at REG_CFP->pc.
188245
static uint8_t *
189-
yjit_side_exit(jitstate_t *jit, ctx_t *ctx)
246+
yjit_gen_leave_exit(codeblock_t *cb)
190247
{
191-
return yjit_gen_exit(jit, ctx, ocb);
192-
}
248+
uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
193249

194-
#if RUBY_DEBUG
250+
// Note, gen_leave() fully reconstructs interpreter state before
251+
// coming here.
195252

196-
// Increment a profiling counter with counter_name
197-
#define GEN_COUNTER_INC(cb, counter_name) _gen_counter_inc(cb, &(yjit_runtime_counters . counter_name))
198-
static void
199-
_gen_counter_inc(codeblock_t *cb, int64_t *counter)
200-
{
201-
if (!rb_yjit_opts.gen_stats) return;
202-
mov(cb, REG0, const_ptr_opnd(counter));
203-
cb_write_lock_prefix(cb); // for ractors.
204-
add(cb, mem_opnd(64, REG0, 0), imm_opnd(1));
205-
}
253+
// Every exit to the interpreter should be counted
254+
GEN_COUNTER_INC(cb, leave_interp_return);
206255

207-
// Increment a counter then take an existing side exit.
208-
#define COUNTED_EXIT(side_exit, counter_name) _counted_side_exit(side_exit, &(yjit_runtime_counters . counter_name))
209-
static uint8_t *
210-
_counted_side_exit(uint8_t *existing_side_exit, int64_t *counter)
211-
{
212-
if (!rb_yjit_opts.gen_stats) return existing_side_exit;
256+
// Put PC into the return register, which the post call bytes dispatches to
257+
mov(cb, RAX, member_opnd(REG_CFP, rb_control_frame_t, pc));
213258

214-
uint8_t *start = cb_get_ptr(ocb, ocb->write_pos);
215-
_gen_counter_inc(ocb, counter);
216-
jmp_ptr(ocb, existing_side_exit);
217-
return start;
259+
cb_write_post_call_bytes(cb);
260+
261+
return code_ptr;
218262
}
219263

220-
// Add a comment at the current position in the code block
221-
static void
222-
_add_comment(codeblock_t* cb, const char* comment_str)
264+
// A shorthand for generating an exit in the outline block
265+
static uint8_t *
266+
yjit_side_exit(jitstate_t *jit, ctx_t *ctx)
223267
{
224-
// Avoid adding duplicate comment strings (can happen due to deferred codegen)
225-
size_t num_comments = rb_darray_size(yjit_code_comments);
226-
if (num_comments > 0) {
227-
struct yjit_comment last_comment = rb_darray_get(yjit_code_comments, num_comments - 1);
228-
if (last_comment.offset == cb->write_pos && strcmp(last_comment.comment, comment_str) == 0) {
229-
return;
230-
}
231-
}
232-
233-
struct yjit_comment new_comment = (struct yjit_comment){ cb->write_pos, comment_str };
234-
rb_darray_append(&yjit_code_comments, new_comment);
268+
return yjit_gen_exit(jit, ctx, ocb);
235269
}
236270

237-
// Comments for generated machine code
238-
#define ADD_COMMENT(cb, comment) _add_comment((cb), (comment))
239-
yjit_comment_array_t yjit_code_comments;
240-
241-
#else
242-
243-
#define GEN_COUNTER_INC(cb, counter_name) ((void)0)
244-
#define COUNTED_EXIT(side_exit, counter_name) side_exit
245-
#define ADD_COMMENT(cb, comment) ((void)0)
246-
247-
#endif // if RUBY_DEBUG
248-
249271
/*
250272
Compile an interpreter entry block to be inserted into an iseq
251273
Returns `NULL` if compilation fails.
252274
*/
253-
uint8_t*
275+
uint8_t *
254276
yjit_entry_prologue(void)
255277
{
256278
RUBY_ASSERT(cb != NULL);
@@ -270,6 +292,11 @@ yjit_entry_prologue(void)
270292
// Load the current SP from the CFP into REG_SP
271293
mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
272294

295+
// Setup cfp->jit_return
296+
// TODO: this could use an IP relative LEA instead of an 8 byte immediate
297+
mov(cb, REG0, const_ptr_opnd(leave_exit_code));
298+
mov(cb, member_opnd(REG_CFP, rb_control_frame_t, jit_return), REG0);
299+
273300
return code_ptr;
274301
}
275302

@@ -2062,9 +2089,6 @@ gen_leave(jitstate_t* jit, ctx_t* ctx)
20622089
// Load the return value
20632090
mov(cb, REG0, ctx_stack_pop(ctx, 1));
20642091

2065-
// Load the JIT return address
2066-
mov(cb, REG1, member_opnd(REG_CFP, rb_control_frame_t, jit_return));
2067-
20682092
// Pop the current frame (ec->cfp++)
20692093
// Note: the return PC is already in the previous CFP
20702094
add(cb, REG_CFP, imm_opnd(sizeof(rb_control_frame_t)));
@@ -2076,20 +2100,9 @@ gen_leave(jitstate_t* jit, ctx_t* ctx)
20762100
mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
20772101
mov(cb, mem_opnd(64, REG_SP, -SIZEOF_VALUE), REG0);
20782102

2079-
// If the return address is NULL, fall back to the interpreter
2080-
ADD_COMMENT(cb, "check for jit return");
2081-
int FALLBACK_LABEL = cb_new_label(cb, "FALLBACK");
2082-
test(cb, REG1, REG1);
2083-
jz_label(cb, FALLBACK_LABEL);
2084-
2085-
// Jump to the JIT return address
2086-
jmp_rm(cb, REG1);
2087-
2088-
// Fall back to the interpreter
2089-
cb_write_label(cb, FALLBACK_LABEL);
2090-
cb_link_labels(cb);
2091-
GEN_COUNTER_INC(cb, leave_interp_return);
2092-
cb_write_post_call_bytes(cb);
2103+
// Jump to the JIT return address on the frame that was just popped
2104+
const int32_t offset_to_jit_return = -((int32_t)sizeof(rb_control_frame_t)) + (int32_t)offsetof(rb_control_frame_t, jit_return);
2105+
jmp_rm(cb, mem_opnd(64, REG_CFP, offset_to_jit_return));
20932106

20942107
return YJIT_END_BLOCK;
20952108
}
@@ -2151,12 +2164,17 @@ yjit_init_codegen(void)
21512164
{
21522165
// Initialize the code blocks
21532166
uint32_t mem_size = 128 * 1024 * 1024;
2154-
uint8_t* mem_block = alloc_exec_mem(mem_size);
2167+
uint8_t *mem_block = alloc_exec_mem(mem_size);
2168+
21552169
cb = █
21562170
cb_init(cb, mem_block, mem_size/2);
2171+
21572172
ocb = &outline_block;
21582173
cb_init(ocb, mem_block + mem_size/2, mem_size/2);
21592174

2175+
// Generate the interpreter exit code for leave
2176+
leave_exit_code = yjit_gen_leave_exit(cb);
2177+
21602178
// Map YARV opcodes to the corresponding codegen functions
21612179
yjit_reg_op(BIN(dup), gen_dup);
21622180
yjit_reg_op(BIN(nop), gen_nop);

0 commit comments

Comments
 (0)