Skip to content

Commit

Permalink
YJIT: Stop using the starting_context pattern (#7610)
Browse files Browse the repository at this point in the history
  • Loading branch information
k0kubun committed Mar 28, 2023
1 parent 2488b4d commit 2f8a598
Showing 1 changed file with 64 additions and 72 deletions.
136 changes: 64 additions & 72 deletions yjit/src/codegen.rs
Expand Up @@ -806,9 +806,8 @@ pub fn gen_single_block(
// For each instruction to compile
// NOTE: could rewrite this loop with a std::iter::Iterator
while insn_idx < iseq_size {
// Set the starting_ctx so we can use it for side_exiting
// if status == CantCompile
let starting_ctx = ctx.clone();
// Set the initial stack size to check if it's safe to exit on CantCompile
let starting_stack_size = ctx.get_stack_size();

// Get the current pc and opcode
let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx.into()) };
Expand Down Expand Up @@ -872,11 +871,10 @@ pub fn gen_single_block(
println!("can't compile {}", insn_name(opcode));
}

// TODO: if the codegen function makes changes to ctx and then return YJIT_CANT_COMPILE,
// the exit this generates would be wrong. There are some changes that are safe to make
// like popping from the stack (but not writing). We are assuming here that only safe
// changes were made and using the starting_ctx.
gen_exit(jit.pc, &starting_ctx, &mut asm);
// Using the latest ctx instead of a starting ctx to allow spilling stack temps in the future.
// When you return CantCompile, you should not modify stack_size.
assert_eq!(ctx.get_stack_size(), starting_stack_size, "stack_size was modified despite CantCompile");
gen_exit(jit.pc, &ctx, &mut asm);

// If this is the first instruction in the block, then
// the entry address is the address for block_entry_exit
Expand Down Expand Up @@ -1997,7 +1995,6 @@ fn gen_get_ivar(
side_exit: Target,
) -> CodegenStatus {
let comptime_val_klass = comptime_receiver.class_of();
let starting_context = ctx.clone(); // make a copy for use with jit_chain_guard

// If recv isn't already a register, load it.
let recv = match recv {
Expand Down Expand Up @@ -2065,11 +2062,6 @@ fn gen_get_ivar(
// Guard heap object (recv_opnd must be used before stack_pop)
guard_object_is_heap(ctx, asm, recv, recv_opnd, side_exit);

// Pop receiver if it's on the temp stack
if recv_opnd != SelfOpnd {
ctx.stack_pop(1);
}

// Compile time self is embedded and the ivar index lands within the object
let embed_test_result = unsafe { FL_TEST_RAW(comptime_receiver, VALUE(ROBJECT_EMBED.as_usize())) != VALUE(0) };

Expand All @@ -2083,13 +2075,18 @@ fn gen_get_ivar(
jit_chain_guard(
JCC_JNE,
jit,
&starting_context,
&ctx,
asm,
ocb,
max_chain_depth,
megamorphic_side_exit,
);

// Pop receiver if it's on the temp stack
if recv_opnd != SelfOpnd {
ctx.stack_pop(1);
}

match ivar_index {
// If there is no IVAR index, then the ivar was undefined
// when we entered the compiler. That means we can just return
Expand Down Expand Up @@ -2207,8 +2204,6 @@ fn gen_setinstancevariable(
asm: &mut Assembler,
ocb: &mut OutlinedCb,
) -> CodegenStatus {
let starting_context = ctx.clone(); // make a copy for use with jit_chain_guard

// Defer compilation so we can specialize on a runtime `self`
if !jit.at_current_insn() {
defer_compilation(jit, ctx, asm, ocb);
Expand Down Expand Up @@ -2301,7 +2296,7 @@ fn gen_setinstancevariable(
jit_chain_guard(
JCC_JNE,
jit,
&starting_context,
&ctx,
asm,
ocb,
SET_IVAR_MAX_DEPTH,
Expand Down Expand Up @@ -3567,14 +3562,13 @@ fn gen_opt_case_dispatch(
defer_compilation(jit, ctx, asm, ocb);
return EndBlock;
}
let starting_context = ctx.clone();

let case_hash = jit.get_arg(0);
let else_offset = jit.get_arg(1).as_u32();

// Try to reorder case/else branches so that ones that are actually used come first.
// Supporting only Fixnum for now so that the implementation can be an equality check.
let key_opnd = ctx.stack_pop(1);
let key_opnd = ctx.stack_opnd(0);
let comptime_key = jit.peek_at_stack(ctx, 0);

// Check that all cases are fixnums to avoid having to register BOP assumptions on
Expand Down Expand Up @@ -3603,16 +3597,17 @@ fn gen_opt_case_dispatch(

// Check if the key is the same value
asm.cmp(key_opnd, comptime_key.into());
let side_exit = get_side_exit(jit, ocb, &starting_context);
let side_exit = get_side_exit(jit, ocb, &ctx);
jit_chain_guard(
JCC_JNE,
jit,
&starting_context,
&ctx,
asm,
ocb,
CASE_WHEN_MAX_DEPTH,
side_exit,
);
ctx.stack_pop(1); // Pop key_opnd

// Get the offset for the compile-time key
let mut offset = 0;
Expand All @@ -3630,6 +3625,7 @@ fn gen_opt_case_dispatch(
gen_direct_jump(jit, &ctx, jump_block, asm);
EndBlock
} else {
ctx.stack_pop(1); // Pop key_opnd
KeepCompiling // continue with === branches
}
}
Expand Down Expand Up @@ -5665,8 +5661,43 @@ fn gen_send_iseq(
}
}

// Number of locals that are not parameters
let num_locals = unsafe { get_iseq_body_local_table_size(iseq) as i32 } - (num_params as i32);
// Check if we need the arg0 splat handling of vm_callee_setup_block_arg
let arg_setup_block = captured_opnd.is_some(); // arg_setup_type: arg_setup_block (invokeblock)
let block_arg0_splat = arg_setup_block && argc == 1 && unsafe {
get_iseq_flags_has_lead(iseq) && !get_iseq_flags_ambiguous_param0(iseq)
};
if block_arg0_splat {
// If block_arg0_splat, we still need side exits after splat, but
// doing push_splat_args here disallows it. So bail out.
if flags & VM_CALL_ARGS_SPLAT != 0 && !iseq_has_rest {
gen_counter_incr!(asm, invokeblock_iseq_arg0_args_splat);
return CantCompile;
}
// The block_arg0_splat implementation is for the rb_simple_iseq_p case,
// but doing_kw_call means it's not a simple ISEQ.
if doing_kw_call {
gen_counter_incr!(asm, invokeblock_iseq_arg0_has_kw);
return CantCompile;
}
}
let splat_array_length = if flags & VM_CALL_ARGS_SPLAT != 0 && !iseq_has_rest {
let array = jit.peek_at_stack(ctx, if block_arg { 1 } else { 0 }) ;
let array_length = if array == Qnil {
0
} else {
unsafe { rb_yjit_array_len(array) as u32}
};

if opt_num == 0 && required_num != array_length as i32 + argc - 1 {
gen_counter_incr!(asm, send_iseq_splat_arity_error);
return CantCompile;
}
Some(array_length)
} else {
None
};

// We will not have CantCompile from here. You can use stack_pop / stack_pop.

match block_arg_type {
Some(Type::Nil) => {
Expand All @@ -5688,13 +5719,8 @@ fn gen_send_iseq(
let builtin_attrs = unsafe { rb_yjit_iseq_builtin_attrs(iseq) };
let builtin_func_raw = unsafe { rb_yjit_builtin_function(iseq) };
let builtin_func = if builtin_func_raw.is_null() { None } else { Some(builtin_func_raw) };
if let (None, Some(builtin_info), true) = (block, builtin_func, builtin_attrs & BUILTIN_ATTR_LEAF != 0) {
// this is a .send call not currently supported for builtins
if flags & VM_CALL_OPT_SEND != 0 {
gen_counter_incr!(asm, send_send_builtin);
return CantCompile;
}

let opt_send_call = flags & VM_CALL_OPT_SEND != 0; // .send call is not currently supported for builtins
if let (None, Some(builtin_info), true, false) = (block, builtin_func, builtin_attrs & BUILTIN_ATTR_LEAF != 0, opt_send_call) {
let builtin_argc = unsafe { (*builtin_info).argc };
if builtin_argc + 1 < (C_ARG_OPNDS.len() as i32) {
asm.comment("inlined leaf builtin");
Expand Down Expand Up @@ -5729,6 +5755,9 @@ fn gen_send_iseq(
}
}

// Number of locals that are not parameters
let num_locals = unsafe { get_iseq_body_local_table_size(iseq) as i32 } - (num_params as i32);

// Stack overflow check
// Note that vm_push_frame checks it against a decremented cfp, hence the multiply by 2.
// #define CHECK_VM_STACK_OVERFLOW0(cfp, sp, margin)
Expand All @@ -5740,37 +5769,11 @@ fn gen_send_iseq(
asm.cmp(CFP, stack_limit);
asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow));

// Check if we need the arg0 splat handling of vm_callee_setup_block_arg
let arg_setup_block = captured_opnd.is_some(); // arg_setup_type: arg_setup_block (invokeblock)
let block_arg0_splat = arg_setup_block && argc == 1 && unsafe {
get_iseq_flags_has_lead(iseq) && !get_iseq_flags_ambiguous_param0(iseq)
};

// push_splat_args does stack manipulation so we can no longer side exit
if flags & VM_CALL_ARGS_SPLAT != 0 && !iseq_has_rest {
// If block_arg0_splat, we still need side exits after this, but
// doing push_splat_args here disallows it. So bail out.
if block_arg0_splat {
gen_counter_incr!(asm, invokeblock_iseq_arg0_args_splat);
return CantCompile;
}

let array = jit.peek_at_stack(ctx, if block_arg { 1 } else { 0 }) ;
let array_length = if array == Qnil {
0
} else {
unsafe { rb_yjit_array_len(array) as u32}
};

if opt_num == 0 && required_num != array_length as i32 + argc - 1 {
gen_counter_incr!(asm, send_iseq_splat_arity_error);
return CantCompile;
}

if let Some(array_length) = splat_array_length {
let remaining_opt = (opt_num as u32 + required_num as u32).saturating_sub(array_length + (argc as u32 - 1));

if opt_num > 0 {

// We are going to jump to the correct offset based on how many optional
// params are remaining.
unsafe {
Expand Down Expand Up @@ -5801,13 +5804,6 @@ fn gen_send_iseq(
// Here we're calling a method with keyword arguments and specifying
// keyword arguments at this call site.

// The block_arg0_splat implementation is for the rb_simple_iseq_p case,
// but doing_kw_call means it's not a simple ISEQ.
if block_arg0_splat {
gen_counter_incr!(asm, invokeblock_iseq_arg0_has_kw);
return CantCompile;
}

// Number of positional arguments the callee expects before the first
// keyword argument
let args_before_kw = required_num + opt_num;
Expand Down Expand Up @@ -6522,8 +6518,6 @@ fn gen_send_general(
// instead we look up the method and call it,
// doing some stack shifting based on the VM_CALL_OPT_SEND flag

let starting_context = ctx.clone();

// Reject nested cases such as `send(:send, :alias_for_send, :foo))`.
// We would need to do some stack manipulation here or keep track of how
// many levels deep we need to stack manipulate. Because of how exits
Expand Down Expand Up @@ -6605,7 +6599,7 @@ fn gen_send_general(
jit_chain_guard(
JCC_JNE,
jit,
&starting_context,
&ctx,
asm,
ocb,
SEND_MAX_CHAIN_DEPTH,
Expand Down Expand Up @@ -7541,8 +7535,6 @@ fn gen_getblockparamproxy(
return EndBlock;
}

let starting_context = ctx.clone(); // make a copy for use with jit_chain_guard

// A mirror of the interpreter code. Checking for the case
// where it's pushing rb_block_param_proxy.
let side_exit = get_side_exit(jit, ocb, ctx);
Expand Down Expand Up @@ -7585,7 +7577,7 @@ fn gen_getblockparamproxy(
jit_chain_guard(
JCC_JNZ,
jit,
&starting_context,
&ctx,
asm,
ocb,
SEND_MAX_DEPTH,
Expand All @@ -7603,7 +7595,7 @@ fn gen_getblockparamproxy(
jit_chain_guard(
JCC_JNZ,
jit,
&starting_context,
&ctx,
asm,
ocb,
SEND_MAX_DEPTH,
Expand Down

0 comments on commit 2f8a598

Please sign in to comment.