From 703eee7745935c54e7b80b22b9b695e99f53fc5e Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Mon, 22 Jan 2024 11:55:44 -0500 Subject: [PATCH] YJIT: Drop extra arguments passed by yield (#9596) Support dropping extra arguments passed by `yield` in blocks. For example `10.times { work }` drops the count argument. This is common enough that it's about 3% of fallback reasons in `lobsters`. Only support simple cases where the surplus arguments are at the top of the stack, that way they just need to be popped, which takes no work. --- bootstraptest/test_yjit.rb | 39 +++++++++++++++ yjit/src/codegen.rs | 97 ++++++++++++++++++++++++-------------- yjit/src/stats.rs | 2 +- 3 files changed, 101 insertions(+), 37 deletions(-) diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb index 3bf715f888..926eaaf8f6 100644 --- a/bootstraptest/test_yjit.rb +++ b/bootstraptest/test_yjit.rb @@ -1,3 +1,42 @@ +# test discarding extra yield arguments +assert_equal "2210150001501015", %q{ + def splat_kw(ary) = yield *ary, a: 1 + + def splat(ary) = yield *ary + + def kw = yield 1, 2, a: 0 + + def simple = yield 0, 1 + + def calls + [ + splat([1, 1, 2]) { |x, y| x + y }, + splat([1, 1, 2]) { |y, opt = raise| opt + y}, + splat_kw([0, 1]) { |a:| a }, + kw { |a:| a }, + kw { |a| a }, + simple { 5.itself }, + simple { |a| a }, + simple { |opt = raise| opt }, + simple { |*rest| rest }, + simple { |opt_kw: 5| opt_kw }, + # autosplat ineractions + [0, 1, 2].yield_self { |a, b| [a, b] }, + [0, 1, 2].yield_self { |a, opt = raise| [a, opt] }, + [1].yield_self { |a, opt = 4| a + opt }, + ] + end + + calls.join +} + +# test autosplat with empty splat +assert_equal "ok", %q{ + def m(pos, splat) = yield pos, *splat + + m([:ok], []) {|v0,| v0 } +} + # regression test for send stack shifting assert_normal_exit %q{ def foo(a, b) diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index b50cd10163..fcb4c47608 100644 --- a/yjit/src/codegen.rs +++ b/yjit/src/codegen.rs @@ -6117,6 +6117,7 @@ fn gen_send_iseq( let supplying_kws = unsafe { vm_ci_flag(ci) & VM_CALL_KWARG } != 0; let iseq_has_rest = unsafe { get_iseq_flags_has_rest(iseq) }; let iseq_has_block_param = unsafe { get_iseq_flags_has_block(iseq) }; + let arg_setup_block = captured_opnd.is_some(); // arg_setup_type: arg_setup_block (invokeblock) // For computing offsets to callee locals let num_params = unsafe { get_iseq_body_param_size(iseq) }; @@ -6137,10 +6138,10 @@ fn gen_send_iseq( // Arity handling and optional parameter setup let mut opts_filled = argc - required_num - kw_arg_num; let opt_num = unsafe { get_iseq_body_param_opt_num(iseq) }; - // We have a rest parameter so there could be more args - // than are required + optional. Those will go in rest. + // With a rest parameter or a yield to a block, + // callers can pass more than required + optional. // So we cap ops_filled at opt_num. - if iseq_has_rest { + if iseq_has_rest || arg_setup_block { opts_filled = min(opts_filled, opt_num); } let mut opts_missing: i32 = opt_num - opts_filled; @@ -6159,11 +6160,17 @@ fn gen_send_iseq( exit_if_supplying_kws_and_accept_no_kwargs(asm, supplying_kws, iseq)?; exit_if_splat_and_zsuper(asm, flags)?; exit_if_doing_kw_and_splat(asm, doing_kw_call, flags)?; - exit_if_wrong_number_arguments(asm, opts_filled, flags, opt_num, iseq_has_rest)?; + exit_if_wrong_number_arguments(asm, arg_setup_block, opts_filled, flags, opt_num, iseq_has_rest)?; exit_if_doing_kw_and_opts_missing(asm, doing_kw_call, opts_missing)?; exit_if_has_rest_and_optional_and_block(asm, iseq_has_rest, opt_num, iseq, block_arg)?; let block_arg_type = exit_if_unsupported_block_arg_type(jit, asm, block_arg)?; + // Bail if we can't drop extra arguments for a yield by just popping them + if supplying_kws && arg_setup_block && argc > (kw_arg_num + required_num + opt_num) { + gen_counter_incr(asm, Counter::send_iseq_complex_discard_extras); + return None; + } + // Block parameter handling. This mirrors setup_parameters_complex(). if iseq_has_block_param { if unsafe { get_iseq_body_local_iseq(iseq) == iseq } { @@ -6249,35 +6256,6 @@ fn gen_send_iseq( } } - // Check if we need the arg0 splat handling of vm_callee_setup_block_arg() - // Also known as "autosplat" inside setup_parameters_complex() - 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) || opt_num > 1) - && !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, Counter::invokeblock_iseq_arg0_args_splat); - return None; - } - // 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, Counter::invokeblock_iseq_arg0_has_kw); - return None; - } - // The block_arg0_splat implementation cannot deal with optional parameters. - // This is a setup_parameters_complex() situation and interacts with the - // starting position of the callee. - if opt_num > 1 { - gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_optional); - return None; - } - } - let splat_array_length = if flags & VM_CALL_ARGS_SPLAT != 0 { let array = jit.peek_at_stack(&asm.ctx, if block_arg { 1 } else { 0 }) ; let array_length = if array == Qnil { @@ -6318,6 +6296,33 @@ fn gen_send_iseq( None }; + // Check if we need the arg0 splat handling of vm_callee_setup_block_arg() + // Also known as "autosplat" inside setup_parameters_complex(). + // Autosplat checks argc == 1 after splat and kwsplat processing, so make + // sure to amend this if we start support kw_splat. + let block_arg0_splat = arg_setup_block + && (argc == 1 || (argc == 2 && splat_array_length == Some(0))) + && !supplying_kws && !doing_kw_call + && unsafe { + (get_iseq_flags_has_lead(iseq) || opt_num > 1) + && !get_iseq_flags_ambiguous_param0(iseq) + }; + if block_arg0_splat { + // If block_arg0_splat, we still need side exits after splat, but + // the splat modifies the stack which breaks side exits. So bail out. + if flags & VM_CALL_ARGS_SPLAT != 0 { + gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_args_splat); + return None; + } + // The block_arg0_splat implementation cannot deal with optional parameters. + // This is a setup_parameters_complex() situation and interacts with the + // starting position of the callee. + if opt_num > 1 { + gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_optional); + return None; + } + } + // Adjust `opts_filled` and `opts_missing` taking // into account the size of the splat expansion. if let Some(len) = splat_array_length { @@ -6605,6 +6610,19 @@ fn gen_send_iseq( asm.store(rest_param, rest_param_array); } + // Pop surplus positional arguments when yielding + if arg_setup_block { + let extras = argc - required_num - opt_num; + if extras > 0 { + // Checked earlier. If there are keyword args, then + // the positional arguments are not at the stack top. + assert_eq!(0, kw_arg_num); + + asm.stack_pop(extras as usize); + argc = required_num + opt_num; + } + } + if doing_kw_call { // Here we're calling a method with keyword arguments and specifying // keyword arguments at this call site. @@ -7034,11 +7052,18 @@ fn exit_if_doing_kw_and_splat(asm: &mut Assembler, doing_kw_call: bool, flags: u } #[must_use] -fn exit_if_wrong_number_arguments(asm: &mut Assembler, opts_filled: i32, flags: u32, opt_num: i32, iseq_has_rest: bool) -> Option<()> { +fn exit_if_wrong_number_arguments( + asm: &mut Assembler, + args_setup_block: bool, + opts_filled: i32, + flags: u32, + opt_num: i32, + iseq_has_rest: bool, +) -> Option<()> { // Too few arguments and no splat to make up for it let too_few = opts_filled < 0 && flags & VM_CALL_ARGS_SPLAT == 0; - // Too many arguments and no place to put them (i.e. rest arg) - let too_many = opts_filled > opt_num && !iseq_has_rest; + // Too many arguments and no sink that take them + let too_many = opts_filled > opt_num && !(iseq_has_rest || args_setup_block); exit_if(asm, too_few || too_many, Counter::send_iseq_arity_error) } diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs index 1dd5f57b2b..7df01448a4 100644 --- a/yjit/src/stats.rs +++ b/yjit/src/stats.rs @@ -331,6 +331,7 @@ make_counters! { send_iseq_arity_error, send_iseq_block_arg_type, send_iseq_clobbering_block_arg, + send_iseq_complex_discard_extras, send_iseq_leaf_builtin_block_arg_block_param, send_iseq_only_keywords, send_iseq_kw_splat, @@ -391,7 +392,6 @@ make_counters! { invokeblock_megamorphic, invokeblock_none, invokeblock_iseq_arg0_optional, - invokeblock_iseq_arg0_has_kw, invokeblock_iseq_arg0_args_splat, invokeblock_iseq_arg0_not_array, invokeblock_iseq_arg0_wrong_len,