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

YJIT: Extract keyword handling to make gen_send_iseq() shorter

Browse Source 
This should make it easier to grok the order of operations.
This commit is contained in:
Alan Wu 2024-02-14 15:23:38 -05:00
parent ae8db4b65a
commit 72f8883a2f
1 changed files with 299 additions and 275 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

574
yjit/src/codegen.rs
View File

@ -6572,76 +6572,7 @@ perf_fn!(fn gen_send_iseq(
// Check that required keyword arguments are supplied and find any extras
// that should go into the keyword rest parameter (**kw_rest).
if doing_kw_call {
// This struct represents the metadata about the callee-specified
// keyword parameters.
let keyword = unsafe { get_iseq_body_param_keyword(iseq) };
let keyword_num: usize = unsafe { (*keyword).num }.try_into().unwrap();
let keyword_required_num: usize = unsafe { (*keyword).required_num }.try_into().unwrap();
let mut required_kwargs_filled = 0;
if keyword_num > 30 || kw_arg_num > 64 {
// We have so many keywords that (1 << num) encoded as a FIXNUM
// (which shifts it left one more) no longer fits inside a 32-bit
// immediate. Similarly, we use a u64 in case of keyword rest parameter.
gen_counter_incr(asm, Counter::send_iseq_too_many_kwargs);
return None;
}
// Check that the kwargs being passed are valid
if supplying_kws {
// This is the list of keyword arguments that the callee specified
// in its initial declaration.
// SAFETY: see compile.c for sizing of this slice.
let callee_kwargs = if keyword_num == 0 {
&[]
} else {
unsafe { slice::from_raw_parts((*keyword).table, keyword_num) }
};
// Here we're going to build up a list of the IDs that correspond to
// the caller-specified keyword arguments. If they're not in the
// same order as the order specified in the callee declaration, then
// we're going to need to generate some code to swap values around
// on the stack.
let kw_arg_keyword_len: usize =
unsafe { get_cikw_keyword_len(kw_arg) }.try_into().unwrap();
let mut caller_kwargs: Vec<ID> = vec![0; kw_arg_keyword_len];
for kwarg_idx in 0..kw_arg_keyword_len {
let sym = unsafe { get_cikw_keywords_idx(kw_arg, kwarg_idx.try_into().unwrap()) };
caller_kwargs[kwarg_idx] = unsafe { rb_sym2id(sym) };
}
// First, we're going to be sure that the names of every
// caller-specified keyword argument correspond to a name in the
// list of callee-specified keyword parameters.
for caller_kwarg in caller_kwargs {
let search_result = callee_kwargs
.iter()
.enumerate() // inject element index
.find(|(_, &kwarg)| kwarg == caller_kwarg);
match search_result {
None if !has_kwrest => {
// If the keyword was never found, then we know we have a
// mismatch in the names of the keyword arguments, so we need to
// bail.
gen_counter_incr(asm, Counter::send_iseq_kwargs_mismatch);
return None;
}
Some((callee_idx, _)) if callee_idx < keyword_required_num => {
// Keep a count to ensure all required kwargs are specified
required_kwargs_filled += 1;
}
_ => (),
}
}
}
assert!(required_kwargs_filled <= keyword_required_num);
if required_kwargs_filled != keyword_required_num {
gen_counter_incr(asm, Counter::send_iseq_kwargs_mismatch);
return None;
}
gen_iseq_kw_call_checks(asm, iseq, kw_arg, has_kwrest, kw_arg_num)?;
}
let splat_array_length = if flags & VM_CALL_ARGS_SPLAT != 0 {
@ -7047,211 +6978,7 @@ perf_fn!(fn gen_send_iseq(
// Keyword argument passing
if doing_kw_call {
// This struct represents the metadata about the caller-specified
// keyword arguments.
let ci_kwarg = unsafe { vm_ci_kwarg(ci) };
let caller_keyword_len_i32: i32 = if ci_kwarg.is_null() {
0
} else {
unsafe { get_cikw_keyword_len(ci_kwarg) }
};
let caller_keyword_len: usize = caller_keyword_len_i32.try_into().unwrap();
// This struct represents the metadata about the callee-specified
// keyword parameters.
let keyword = unsafe { get_iseq_body_param_keyword(iseq) };
asm_comment!(asm, "keyword args");
// This is the list of keyword arguments that the callee specified
// in its initial declaration.
let callee_kwargs = unsafe { (*keyword).table };
let callee_kw_count_i32: i32 = unsafe { (*keyword).num };
let callee_kw_count: usize = callee_kw_count_i32.try_into().unwrap();
let keyword_required_num: usize = unsafe { (*keyword).required_num }.try_into().unwrap();
// Here we're going to build up a list of the IDs that correspond to
// the caller-specified keyword arguments. If they're not in the
// same order as the order specified in the callee declaration, then
// we're going to need to generate some code to swap values around
// on the stack.
let mut kwargs_order: Vec<ID> = vec![0; cmp::max(caller_keyword_len, callee_kw_count)];
for kwarg_idx in 0..caller_keyword_len {
let sym = unsafe { get_cikw_keywords_idx(ci_kwarg, kwarg_idx.try_into().unwrap()) };
kwargs_order[kwarg_idx] = unsafe { rb_sym2id(sym) };
}
let mut unspecified_bits = 0;
// The stack_opnd() index to the 0th keyword argument.
let kwargs_stack_base = caller_keyword_len_i32 - 1;
// Build the keyword rest parameter hash before we make any changes to the order of
// the supplied keyword arguments
if has_kwrest {
c_callable! {
fn build_kw_rest(rest_mask: u64, stack_kwargs: *const VALUE, keywords: *const rb_callinfo_kwarg) -> VALUE {
if keywords.is_null() {
return unsafe { rb_hash_new() };
}
// Use the total number of supplied keywords as a size upper bound
let keyword_len = unsafe { (*keywords).keyword_len } as usize;
let hash = unsafe { rb_hash_new_with_size(keyword_len as u64) };
// Put pairs into the kwrest hash as the mask describes
for kwarg_idx in 0..keyword_len {
if (rest_mask & (1 << kwarg_idx)) != 0 {
unsafe {
let keyword_symbol = (*keywords).keywords.as_ptr().add(kwarg_idx).read();
let keyword_value = stack_kwargs.add(kwarg_idx).read();
rb_hash_aset(hash, keyword_symbol, keyword_value);
}
}
}
return hash;
}
}
asm_comment!(asm, "build kwrest hash");
// Make a bit mask describing which keywords should go into kwrest.
let mut rest_mask: u64 = 0;
// Index for one argument that will go into kwrest.
let mut rest_collected_idx = None;
for (supplied_kw_idx, &supplied_kw) in kwargs_order.iter().take(caller_keyword_len).enumerate() {
let mut found = false;
for callee_idx in 0..callee_kw_count {
let callee_kw = unsafe { callee_kwargs.add(callee_idx).read() };
if callee_kw == supplied_kw {
found = true;
break;
}
}
if !found {
rest_mask |= 1 << supplied_kw_idx;
if rest_collected_idx.is_none() {
rest_collected_idx = Some(supplied_kw_idx as i32);
}
}
}
// Save PC and SP before allocating
jit_save_pc(jit, asm);
gen_save_sp(asm);
// Build the kwrest hash. `struct rb_callinfo_kwarg` is malloc'd, so no GC concerns.
let kwargs_start = asm.lea(asm.ctx.sp_opnd(-(caller_keyword_len_i32 * SIZEOF_VALUE_I32) as isize));
let kwrest = asm.ccall(
build_kw_rest as _,
vec![rest_mask.into(), kwargs_start, Opnd::const_ptr(ci_kwarg.cast())]
);
// The kwrest parameter sits after `unspecified_bits` if the callee specifies any
// keywords.
let stack_kwrest_idx = kwargs_stack_base - callee_kw_count_i32 - i32::from(callee_kw_count > 0);
let stack_kwrest = asm.stack_opnd(stack_kwrest_idx);
// If `stack_kwrest` already has another argument there, we need to stow it elsewhere
// first before putting kwrest there. Use `rest_collected_idx` because that value went
// into kwrest so the slot is now free.
let kwrest_idx = callee_kw_count + usize::from(callee_kw_count > 0);
if let (Some(rest_collected_idx), true) = (rest_collected_idx, kwrest_idx < caller_keyword_len) {
let rest_collected = asm.stack_opnd(kwargs_stack_base - rest_collected_idx);
let mapping = asm.ctx.get_opnd_mapping(stack_kwrest.into());
asm.mov(rest_collected, stack_kwrest);
asm.ctx.set_opnd_mapping(rest_collected.into(), mapping);
// Update our bookkeeping to inform the reordering step later.
kwargs_order[rest_collected_idx as usize] = kwargs_order[kwrest_idx];
kwargs_order[kwrest_idx] = 0;
}
// Put kwrest straight into memory, since we might pop it later
asm.ctx.dealloc_temp_reg(stack_kwrest.stack_idx());
asm.mov(stack_kwrest, kwrest);
if stack_kwrest_idx >= 0 {
asm.ctx.set_opnd_mapping(stack_kwrest.into(), TempMapping::map_to_stack(Type::THash));
}
}
// Ensure the stack is large enough for the callee
for _ in caller_keyword_len..callee_kw_count {
argc += 1;
asm.stack_push(Type::Unknown);
}
// Now this is the stack_opnd() index to the 0th keyword argument.
let kwargs_stack_base = kwargs_order.len() as i32 - 1;
// Next, we're going to loop through every keyword that was
// specified by the caller and make sure that it's in the correct
// place. If it's not we're going to swap it around with another one.
for kwarg_idx in 0..callee_kw_count {
let callee_kwarg = unsafe { callee_kwargs.add(kwarg_idx).read() };
// If the argument is already in the right order, then we don't
// need to generate any code since the expected value is already
// in the right place on the stack.
if callee_kwarg == kwargs_order[kwarg_idx] {
continue;
}
// In this case the argument is not in the right place, so we
// need to find its position where it _should_ be and swap with
// that location.
for swap_idx in 0..kwargs_order.len() {
if callee_kwarg == kwargs_order[swap_idx] {
// First we're going to generate the code that is going
// to perform the actual swapping at runtime.
let swap_idx_i32: i32 = swap_idx.try_into().unwrap();
let kwarg_idx_i32: i32 = kwarg_idx.try_into().unwrap();
let offset0 = kwargs_stack_base - swap_idx_i32;
let offset1 = kwargs_stack_base - kwarg_idx_i32;
stack_swap(asm, offset0, offset1);
// Next we're going to do some bookkeeping on our end so
// that we know the order that the arguments are
// actually in now.
kwargs_order.swap(kwarg_idx, swap_idx);
break;
}
}
}
// Now that every caller specified kwarg is in the right place, filling
// in unspecified default paramters won't overwrite anything.
for kwarg_idx in keyword_required_num..callee_kw_count {
if kwargs_order[kwarg_idx] != unsafe { callee_kwargs.add(kwarg_idx).read() } {
let default_param_idx = kwarg_idx - keyword_required_num;
let mut default_value = unsafe { (*keyword).default_values.add(default_param_idx).read() };
if default_value == Qundef {
// Qundef means that this value is not constant and must be
// recalculated at runtime, so we record it in unspecified_bits
// (Qnil is then used as a placeholder instead of Qundef).
unspecified_bits |= 0x01 << default_param_idx;
default_value = Qnil;
}
let default_param = asm.stack_opnd(kwargs_stack_base - kwarg_idx as i32);
let param_type = Type::from(default_value);
asm.mov(default_param, default_value.into());
asm.ctx.set_opnd_mapping(default_param.into(), TempMapping::map_to_stack(param_type));
}
}
// Pop extra arguments that went into kwrest now that they're at stack top
if has_kwrest && caller_keyword_len > callee_kw_count {
let extra_kwarg_count = caller_keyword_len - callee_kw_count;
asm.stack_pop(extra_kwarg_count);
argc = argc - extra_kwarg_count as i32;
}
// Keyword arguments cause a special extra local variable to be
// pushed onto the stack that represents the parameters that weren't
// explicitly given a value and have a non-constant default.
if callee_kw_count > 0 {
let unspec_opnd = VALUE::fixnum_from_usize(unspecified_bits).as_u64();
asm.ctx.dealloc_temp_reg(asm.stack_opnd(-1).stack_idx()); // avoid using a register for unspecified_bits
asm.mov(asm.stack_opnd(-1), unspec_opnd.into());
}
argc = gen_iseq_kw_call(jit, asm, kw_arg, iseq, argc, has_kwrest);
}
// Same as vm_callee_setup_block_arg_arg0_check and vm_callee_setup_block_arg_arg0_splat
@ -7466,6 +7193,303 @@ perf_fn!(fn gen_send_iseq(
Some(EndBlock)
});
// Check if we can handle a keyword call
fn gen_iseq_kw_call_checks(
asm: &mut Assembler,
iseq: *const rb_iseq_t,
kw_arg: *const rb_callinfo_kwarg,
has_kwrest: bool,
caller_kw_num: i32
) -> Option<()> {
// This struct represents the metadata about the callee-specified
// keyword parameters.
let keyword = unsafe { get_iseq_body_param_keyword(iseq) };
let keyword_num: usize = unsafe { (*keyword).num }.try_into().unwrap();
let keyword_required_num: usize = unsafe { (*keyword).required_num }.try_into().unwrap();
let mut required_kwargs_filled = 0;
if keyword_num > 30 || caller_kw_num > 64 {
// We have so many keywords that (1 << num) encoded as a FIXNUM
// (which shifts it left one more) no longer fits inside a 32-bit
// immediate. Similarly, we use a u64 in case of keyword rest parameter.
gen_counter_incr(asm, Counter::send_iseq_too_many_kwargs);
return None;
}
// Check that the kwargs being passed are valid
if caller_kw_num > 0 {
// This is the list of keyword arguments that the callee specified
// in its initial declaration.
// SAFETY: see compile.c for sizing of this slice.
let callee_kwargs = if keyword_num == 0 {
&[]
} else {
unsafe { slice::from_raw_parts((*keyword).table, keyword_num) }
};
// Here we're going to build up a list of the IDs that correspond to
// the caller-specified keyword arguments. If they're not in the
// same order as the order specified in the callee declaration, then
// we're going to need to generate some code to swap values around
// on the stack.
let kw_arg_keyword_len = caller_kw_num as usize;
let mut caller_kwargs: Vec<ID> = vec![0; kw_arg_keyword_len];
for kwarg_idx in 0..kw_arg_keyword_len {
let sym = unsafe { get_cikw_keywords_idx(kw_arg, kwarg_idx.try_into().unwrap()) };
caller_kwargs[kwarg_idx] = unsafe { rb_sym2id(sym) };
}
// First, we're going to be sure that the names of every
// caller-specified keyword argument correspond to a name in the
// list of callee-specified keyword parameters.
for caller_kwarg in caller_kwargs {
let search_result = callee_kwargs
.iter()
.enumerate() // inject element index
.find(|(_, &kwarg)| kwarg == caller_kwarg);
match search_result {
None if !has_kwrest => {
// If the keyword was never found, then we know we have a
// mismatch in the names of the keyword arguments, so we need to
// bail.
gen_counter_incr(asm, Counter::send_iseq_kwargs_mismatch);
return None;
}
Some((callee_idx, _)) if callee_idx < keyword_required_num => {
// Keep a count to ensure all required kwargs are specified
required_kwargs_filled += 1;
}
_ => (),
}
}
}
assert!(required_kwargs_filled <= keyword_required_num);
if required_kwargs_filled != keyword_required_num {
gen_counter_incr(asm, Counter::send_iseq_kwargs_mismatch);
return None;
}
Some(())
}
// Codegen for keyword argument handling. Essentially private to gen_send_iseq() since
// there are a lot of preconditions to check before reaching this code.
fn gen_iseq_kw_call(
jit: &mut JITState,
asm: &mut Assembler,
ci_kwarg: *const rb_callinfo_kwarg,
iseq: *const rb_iseq_t,
mut argc: i32,
has_kwrest: bool,
) -> i32 {
let caller_keyword_len_i32: i32 = if ci_kwarg.is_null() {
0
} else {
unsafe { get_cikw_keyword_len(ci_kwarg) }
};
let caller_keyword_len: usize = caller_keyword_len_i32.try_into().unwrap();
// This struct represents the metadata about the callee-specified
// keyword parameters.
let keyword = unsafe { get_iseq_body_param_keyword(iseq) };
asm_comment!(asm, "keyword args");
// This is the list of keyword arguments that the callee specified
// in its initial declaration.
let callee_kwargs = unsafe { (*keyword).table };
let callee_kw_count_i32: i32 = unsafe { (*keyword).num };
let callee_kw_count: usize = callee_kw_count_i32.try_into().unwrap();
let keyword_required_num: usize = unsafe { (*keyword).required_num }.try_into().unwrap();
// Here we're going to build up a list of the IDs that correspond to
// the caller-specified keyword arguments. If they're not in the
// same order as the order specified in the callee declaration, then
// we're going to need to generate some code to swap values around
// on the stack.
let mut kwargs_order: Vec<ID> = vec![0; cmp::max(caller_keyword_len, callee_kw_count)];
for kwarg_idx in 0..caller_keyword_len {
let sym = unsafe { get_cikw_keywords_idx(ci_kwarg, kwarg_idx.try_into().unwrap()) };
kwargs_order[kwarg_idx] = unsafe { rb_sym2id(sym) };
}
let mut unspecified_bits = 0;
// The stack_opnd() index to the 0th keyword argument.
let kwargs_stack_base = caller_keyword_len_i32 - 1;
// Build the keyword rest parameter hash before we make any changes to the order of
// the supplied keyword arguments
if has_kwrest {
c_callable! {
fn build_kw_rest(rest_mask: u64, stack_kwargs: *const VALUE, keywords: *const rb_callinfo_kwarg) -> VALUE {
if keywords.is_null() {
return unsafe { rb_hash_new() };
}
// Use the total number of supplied keywords as a size upper bound
let keyword_len = unsafe { (*keywords).keyword_len } as usize;
let hash = unsafe { rb_hash_new_with_size(keyword_len as u64) };
// Put pairs into the kwrest hash as the mask describes
for kwarg_idx in 0..keyword_len {
if (rest_mask & (1 << kwarg_idx)) != 0 {
unsafe {
let keyword_symbol = (*keywords).keywords.as_ptr().add(kwarg_idx).read();
let keyword_value = stack_kwargs.add(kwarg_idx).read();
rb_hash_aset(hash, keyword_symbol, keyword_value);
}
}
}
return hash;
}
}
asm_comment!(asm, "build kwrest hash");
// Make a bit mask describing which keywords should go into kwrest.
let mut rest_mask: u64 = 0;
// Index for one argument that will go into kwrest.
let mut rest_collected_idx = None;
for (supplied_kw_idx, &supplied_kw) in kwargs_order.iter().take(caller_keyword_len).enumerate() {
let mut found = false;
for callee_idx in 0..callee_kw_count {
let callee_kw = unsafe { callee_kwargs.add(callee_idx).read() };
if callee_kw == supplied_kw {
found = true;
break;
}
}
if !found {
rest_mask |= 1 << supplied_kw_idx;
if rest_collected_idx.is_none() {
rest_collected_idx = Some(supplied_kw_idx as i32);
}
}
}
// Save PC and SP before allocating
jit_save_pc(jit, asm);
gen_save_sp(asm);
// Build the kwrest hash. `struct rb_callinfo_kwarg` is malloc'd, so no GC concerns.
let kwargs_start = asm.lea(asm.ctx.sp_opnd(-(caller_keyword_len_i32 * SIZEOF_VALUE_I32) as isize));
let kwrest = asm.ccall(
build_kw_rest as _,
vec![rest_mask.into(), kwargs_start, Opnd::const_ptr(ci_kwarg.cast())]
);
// The kwrest parameter sits after `unspecified_bits` if the callee specifies any
// keywords.
let stack_kwrest_idx = kwargs_stack_base - callee_kw_count_i32 - i32::from(callee_kw_count > 0);
let stack_kwrest = asm.stack_opnd(stack_kwrest_idx);
// If `stack_kwrest` already has another argument there, we need to stow it elsewhere
// first before putting kwrest there. Use `rest_collected_idx` because that value went
// into kwrest so the slot is now free.
let kwrest_idx = callee_kw_count + usize::from(callee_kw_count > 0);
if let (Some(rest_collected_idx), true) = (rest_collected_idx, kwrest_idx < caller_keyword_len) {
let rest_collected = asm.stack_opnd(kwargs_stack_base - rest_collected_idx);
let mapping = asm.ctx.get_opnd_mapping(stack_kwrest.into());
asm.mov(rest_collected, stack_kwrest);
asm.ctx.set_opnd_mapping(rest_collected.into(), mapping);
// Update our bookkeeping to inform the reordering step later.
kwargs_order[rest_collected_idx as usize] = kwargs_order[kwrest_idx];
kwargs_order[kwrest_idx] = 0;
}
// Put kwrest straight into memory, since we might pop it later
asm.ctx.dealloc_temp_reg(stack_kwrest.stack_idx());
asm.mov(stack_kwrest, kwrest);
if stack_kwrest_idx >= 0 {
asm.ctx.set_opnd_mapping(stack_kwrest.into(), TempMapping::map_to_stack(Type::THash));
}
}
// Ensure the stack is large enough for the callee
for _ in caller_keyword_len..callee_kw_count {
argc += 1;
asm.stack_push(Type::Unknown);
}
// Now this is the stack_opnd() index to the 0th keyword argument.
let kwargs_stack_base = kwargs_order.len() as i32 - 1;
// Next, we're going to loop through every keyword that was
// specified by the caller and make sure that it's in the correct
// place. If it's not we're going to swap it around with another one.
for kwarg_idx in 0..callee_kw_count {
let callee_kwarg = unsafe { callee_kwargs.add(kwarg_idx).read() };
// If the argument is already in the right order, then we don't
// need to generate any code since the expected value is already
// in the right place on the stack.
if callee_kwarg == kwargs_order[kwarg_idx] {
continue;
}
// In this case the argument is not in the right place, so we
// need to find its position where it _should_ be and swap with
// that location.
for swap_idx in 0..kwargs_order.len() {
if callee_kwarg == kwargs_order[swap_idx] {
// First we're going to generate the code that is going
// to perform the actual swapping at runtime.
let swap_idx_i32: i32 = swap_idx.try_into().unwrap();
let kwarg_idx_i32: i32 = kwarg_idx.try_into().unwrap();
let offset0 = kwargs_stack_base - swap_idx_i32;
let offset1 = kwargs_stack_base - kwarg_idx_i32;
stack_swap(asm, offset0, offset1);
// Next we're going to do some bookkeeping on our end so
// that we know the order that the arguments are
// actually in now.
kwargs_order.swap(kwarg_idx, swap_idx);
break;
}
}
}
// Now that every caller specified kwarg is in the right place, filling
// in unspecified default paramters won't overwrite anything.
for kwarg_idx in keyword_required_num..callee_kw_count {
if kwargs_order[kwarg_idx] != unsafe { callee_kwargs.add(kwarg_idx).read() } {
let default_param_idx = kwarg_idx - keyword_required_num;
let mut default_value = unsafe { (*keyword).default_values.add(default_param_idx).read() };
if default_value == Qundef {
// Qundef means that this value is not constant and must be
// recalculated at runtime, so we record it in unspecified_bits
// (Qnil is then used as a placeholder instead of Qundef).
unspecified_bits |= 0x01 << default_param_idx;
default_value = Qnil;
}
let default_param = asm.stack_opnd(kwargs_stack_base - kwarg_idx as i32);
let param_type = Type::from(default_value);
asm.mov(default_param, default_value.into());
asm.ctx.set_opnd_mapping(default_param.into(), TempMapping::map_to_stack(param_type));
}
}
// Pop extra arguments that went into kwrest now that they're at stack top
if has_kwrest && caller_keyword_len > callee_kw_count {
let extra_kwarg_count = caller_keyword_len - callee_kw_count;
asm.stack_pop(extra_kwarg_count);
argc = argc - extra_kwarg_count as i32;
}
// Keyword arguments cause a special extra local variable to be
// pushed onto the stack that represents the parameters that weren't
// explicitly given a value and have a non-constant default.
if callee_kw_count > 0 {
let unspec_opnd = VALUE::fixnum_from_usize(unspecified_bits).as_u64();
asm.ctx.dealloc_temp_reg(asm.stack_opnd(-1).stack_idx()); // avoid using a register for unspecified_bits
asm.mov(asm.stack_opnd(-1), unspec_opnd.into());
}
argc
}
/// This is a helper function to allow us to exit early
/// during code generation if a predicate is true.
/// We return Option<()> here because we will be able to