Skip to content

Commit 3c84858

Browse files
rwstaunerk0kubun
authored andcommitted
ZJIT: Optimize send with a nil block to SendWithoutBlockDirect
1 parent 9a5c5bf commit 3c84858

3 files changed

Lines changed: 58 additions & 24 deletions

File tree

zjit/src/hir.rs

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -531,6 +531,7 @@ pub enum SideExitReason {
531531
UnhandledYARVInsn(u32),
532532
UnhandledCallType(CallType),
533533
UnhandledBlockArg,
534+
BlockArgNotNil,
534535
TooManyKeywordParameters,
535536
TooManyArgsForLir,
536537
FixnumAddOverflow,
@@ -694,6 +695,8 @@ pub enum SendFallbackReason {
694695
SendCfuncArrayVariadic,
695696
SendNotOptimizedMethodType(MethodType),
696697
SendNotOptimizedNeedPermission,
698+
/// The block argument is not nil, so we can't optimize to SendWithoutBlockDirect
699+
SendBlockArgNotNil,
697700
CCallWithFrameTooManyArgs,
698701
ObjToStringNotString,
699702
TooManyArgsForLir,
@@ -768,6 +771,7 @@ impl Display for SendFallbackReason {
768771
SendCfuncVariadic => write!(f, "Send: C function is variadic"),
769772
SendCfuncArrayVariadic => write!(f, "Send: C function expects array variadic"),
770773
SendNotOptimizedMethodType(method_type) => write!(f, "Send: unsupported method type {:?}", method_type),
774+
SendBlockArgNotNil => write!(f, "Send: block argument is not nil"),
771775
CCallWithFrameTooManyArgs => write!(f, "CCallWithFrame: too many arguments"),
772776
ObjToStringNotString => write!(f, "ObjToString: result is not a string"),
773777
TooManyArgsForLir => write!(f, "Too many arguments for LIR"),
@@ -2479,7 +2483,7 @@ pub enum ValidationError {
24792483
}
24802484

24812485
/// Check if we can do a direct send to the given iseq with the given args.
2482-
fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq_t, ci: *const rb_callinfo, send_insn: InsnId, args: &[InsnId], has_block: bool) -> bool {
2486+
fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq_t, ci: *const rb_callinfo, send_insn: InsnId, args: &[InsnId], has_block: bool, caller_passes_block_arg: bool) -> bool {
24832487
let mut can_send = true;
24842488
let mut count_failure = |counter| {
24852489
can_send = false;
@@ -2488,8 +2492,6 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
24882492
let params = unsafe { iseq.params() };
24892493

24902494
let callee_has_block_param = 0 != params.flags.has_block();
2491-
let caller_passes_block_arg = (unsafe { rb_vm_ci_flag(ci) } & VM_CALL_ARGS_BLOCKARG) != 0;
2492-
24932495
use Counter::*;
24942496
if 0 != params.flags.has_rest() { count_failure(complex_arg_pass_param_rest) }
24952497
if 0 != params.flags.forwardable() { count_failure(complex_arg_pass_param_forwardable) }
@@ -3703,14 +3705,17 @@ impl Function {
37033705
/// Try trivially inlining the method. If we can't, emit a SendDirect instruction instead and
37043706
/// leave it to the general-purpose inliner to handle.
37053707
fn try_inline_send_direct(&mut self, block: BlockId, insn: Insn) -> InsnId {
3706-
let Insn::SendDirect { recv, iseq, cd, ref args, state, .. } = insn else {
3708+
let Insn::SendDirect { recv, iseq, cd, ref args, state, block: send_block, .. } = insn else {
37073709
panic!("try_inline_send_direct called with non-SendDirect instruction");
37083710
};
37093711
// The trivial inliner runs first to handle simple cases (constant returns,
37103712
// parameter returns, etc.) without frame push/pop overhead. The general
37113713
// inliner then handles more complex methods that require full inlining.
37123714
let call_info = unsafe { (*cd).ci };
3713-
let ci_flags = unsafe { vm_ci_flag(call_info) };
3715+
let mut ci_flags = unsafe { vm_ci_flag(call_info) };
3716+
if send_block.is_none() {
3717+
ci_flags &= !VM_CALL_ARGS_BLOCKARG;
3718+
}
37143719
// .send call is not currently supported for builtins
37153720
if ci_flags & VM_CALL_OPT_SEND != 0 {
37163721
return self.push_insn(block, insn);
@@ -3822,9 +3827,32 @@ impl Function {
38223827
def_type = unsafe { get_cme_def_type(cme) };
38233828
}
38243829

3830+
// Check if we can optimize `foo(&block)` where block is nil to a send without block
3831+
let mut send_block = send_block;
3832+
let mut args = args;
3833+
let mut stripped_nil_block = false;
3834+
if send_block == Some(BlockHandler::BlockArg) && def_type == VM_METHOD_TYPE_ISEQ {
3835+
// The block arg is the last element in args
3836+
if let Some(&block_arg) = args.last() {
3837+
let block_arg_type = self.type_of(block_arg);
3838+
if block_arg_type.is_subtype(types::NilClass) {
3839+
// Block arg is known to be nil - strip it and treat as no block
3840+
args = args[..args.len() - 1].to_vec();
3841+
send_block = None;
3842+
stripped_nil_block = true;
3843+
} else {
3844+
// Can't prove block arg is nil
3845+
self.set_dynamic_send_reason(insn_id, SendBlockArgNotNil);
3846+
self.push_insn_id(block, insn_id); continue;
3847+
}
3848+
}
3849+
}
3850+
38253851
// If the call site info indicates that the `Function` has overly complex arguments, then do not optimize into a `SendDirect`.
38263852
// Optimized methods(`VM_METHOD_TYPE_OPTIMIZED`) handle their own argument constraints (e.g., kw_splat for Proc call).
3827-
if def_type != VM_METHOD_TYPE_OPTIMIZED && unspecializable_call_type(flags) {
3853+
// Mask out ARGS_BLOCKARG only if we've already handled the nil block arg case above.
3854+
let flags_for_check = if stripped_nil_block { flags & !VM_CALL_ARGS_BLOCKARG } else { flags };
3855+
if def_type != VM_METHOD_TYPE_OPTIMIZED && unspecializable_call_type(flags_for_check) {
38283856
self.count_complex_call_features(block, flags);
38293857
self.set_dynamic_send_reason(insn_id, ComplexArgPass);
38303858
self.push_insn_id(block, insn_id); continue;
@@ -3835,7 +3863,7 @@ impl Function {
38353863
// Only specialize positional-positional calls
38363864
// TODO(max): Handle other kinds of parameter passing
38373865
let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
3838-
if !can_direct_send(self, block, iseq, ci, insn_id, args.as_slice(), has_block) {
3866+
if !can_direct_send(self, block, iseq, ci, insn_id, args.as_slice(), send_block.is_some(), send_block == Some(BlockHandler::BlockArg)) {
38393867
self.push_insn_id(block, insn_id); continue;
38403868
}
38413869

@@ -3874,7 +3902,7 @@ impl Function {
38743902
let capture = unsafe { proc_block.as_.captured.as_ref() };
38753903
let iseq = unsafe { *capture.code.iseq.as_ref() };
38763904

3877-
if !can_direct_send(self, block, iseq, ci, insn_id, args.as_slice(), has_block) {
3905+
if !can_direct_send(self, block, iseq, ci, insn_id, args.as_slice(), has_block, false) {
38783906
self.push_insn_id(block, insn_id); continue;
38793907
}
38803908

@@ -4216,7 +4244,7 @@ impl Function {
42164244
// If not, we can't do direct dispatch.
42174245
let super_iseq = unsafe { get_def_iseq_ptr((*super_cme).def) };
42184246
// TODO: pass Option<blockiseq> to can_direct_send when we start specializing `super { ... }`.
4219-
if !can_direct_send(self, block, super_iseq, ci, insn_id, args.as_slice(), false) {
4247+
if !can_direct_send(self, block, super_iseq, ci, insn_id, args.as_slice(), false, false) {
42204248
self.push_insn_id(block, insn_id);
42214249
self.set_dynamic_send_reason(insn_id, SuperTargetComplexArgsPass);
42224250
continue;

zjit/src/hir/opt_tests.rs

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5238,7 +5238,7 @@ mod hir_opt_tests {
52385238
v26:TrueClass = GuardBitEquals v25, Value(true) recompile
52395239
Jump bb6(v24, v10)
52405240
bb6(v16:BasicObject, v17:BasicObject):
5241-
v29:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Complex argument passing
5241+
v29:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Send: block argument is not nil
52425242
CheckInterrupts
52435243
Return v29
52445244
");
@@ -5294,7 +5294,7 @@ mod hir_opt_tests {
52945294
v37:NilClass = Const Value(nil)
52955295
Jump bb8(v37, v13)
52965296
bb8(v27:BasicObject, v28:BasicObject):
5297-
v40:BasicObject = Send v25, &block, :then, v27 # SendFallbackReason: Complex argument passing
5297+
v40:BasicObject = Send v25, &block, :then, v27 # SendFallbackReason: Send: block argument is not nil
52985298
CheckInterrupts
52995299
Return v40
53005300
bb4(v45:BasicObject, v46:Falsy, v47:BasicObject):
@@ -5459,7 +5459,7 @@ mod hir_opt_tests {
54595459
v34:NilClass = Const Value(nil)
54605460
Jump bb6(v34, v10)
54615461
bb6(v16:BasicObject, v17:BasicObject):
5462-
v38:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Complex argument passing
5462+
v38:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Send: block argument is not nil
54635463
CheckInterrupts
54645464
Return v38
54655465
bb10():
@@ -5527,7 +5527,7 @@ mod hir_opt_tests {
55275527
v41:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1010))
55285528
Jump bb6(v41, v10)
55295529
bb6(v16:BasicObject, v17:BasicObject):
5530-
v45:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Complex argument passing
5530+
v45:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Send: block argument is not nil
55315531
CheckInterrupts
55325532
Return v45
55335533
bb13():
@@ -9550,7 +9550,7 @@ mod hir_opt_tests {
95509550
v26:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
95519551
Jump bb6(v26, v10)
95529552
bb6(v16:BasicObject, v17:BasicObject):
9553-
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
9553+
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
95549554
CheckInterrupts
95559555
Return v29
95569556
");
@@ -9590,7 +9590,7 @@ mod hir_opt_tests {
95909590
v26:NilClass = Const Value(nil)
95919591
Jump bb6(v26, v10)
95929592
bb6(v16:BasicObject, v17:BasicObject):
9593-
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
9593+
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
95949594
CheckInterrupts
95959595
Return v29
95969596
");
@@ -9631,7 +9631,7 @@ mod hir_opt_tests {
96319631
v21:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
96329632
Jump bb6(v21)
96339633
bb6(v12:BasicObject):
9634-
v24:BasicObject = Send v10, &block, :map, v12 # SendFallbackReason: Complex argument passing
9634+
v24:BasicObject = Send v10, &block, :map, v12 # SendFallbackReason: Send: block argument is not nil
96359635
CheckInterrupts
96369636
Return v24
96379637
");
@@ -9799,7 +9799,7 @@ mod hir_opt_tests {
97999799
Jump bb3(v5, v6)
98009800
bb3(v8:BasicObject, v9:NilClass):
98019801
v13:StaticSymbol[:to_s] = Const Value(VALUE(0x1000))
9802-
v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Complex argument passing
9802+
v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Send: block argument is not nil
98039803
CheckInterrupts
98049804
Return v19
98059805
");
@@ -9830,9 +9830,11 @@ mod hir_opt_tests {
98309830
Jump bb3(v5, v6)
98319831
bb3(v8:BasicObject, v9:NilClass):
98329832
v13:NilClass = Const Value(nil)
9833-
v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Complex argument passing
9833+
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
9834+
v26:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v8, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile
9835+
v27:Fixnum[42] = Const Value(42)
98349836
CheckInterrupts
9835-
Return v19
9837+
Return v27
98369838
");
98379839
}
98389840

@@ -13193,7 +13195,7 @@ mod hir_opt_tests {
1319313195
Jump bb3(v4)
1319413196
bb3(v6:BasicObject):
1319513197
v11:StaticSymbol[:the_block] = Const Value(VALUE(0x1000))
13196-
v13:BasicObject = Send v6, &block, :callee, v11 # SendFallbackReason: Complex argument passing
13198+
v13:BasicObject = Send v6, &block, :callee, v11 # SendFallbackReason: Send: block argument is not nil
1319713199
CheckInterrupts
1319813200
Return v13
1319913201
");
@@ -13238,7 +13240,7 @@ mod hir_opt_tests {
1323813240
v26:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
1323913241
Jump bb6(v26, v10)
1324013242
bb6(v16:BasicObject, v17:BasicObject):
13241-
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
13243+
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
1324213244
CheckInterrupts
1324313245
Return v29
1324413246
");
@@ -17085,7 +17087,7 @@ mod hir_opt_tests {
1708517087
test(true, block)
1708617088
");
1708717089

17088-
assert_snapshot!(hir_string("test"), @r"
17090+
assert_snapshot!(hir_string("test"), @"
1708917091
fn test@<compiled>:7:
1709017092
bb1():
1709117093
EntryPoint interpreter
@@ -17108,7 +17110,7 @@ mod hir_opt_tests {
1710817110
bb5():
1710917111
v22:Truthy = RefineType v12, Truthy
1711017112
v26:Fixnum[42] = Const Value(42)
17111-
v29:BasicObject = Send v11, &block, :passthrough_recompile_blockarg, v26, v13 # SendFallbackReason: Complex argument passing
17113+
v29:BasicObject = Send v11, &block, :passthrough_recompile_blockarg, v26, v13 # SendFallbackReason: Send: block argument is not nil
1711217114
CheckInterrupts
1711317115
Return v29
1711417116
bb4(v34:BasicObject, v35:Falsy, v36:BasicObject):
@@ -19316,7 +19318,7 @@ mod hir_opt_tests {
1931619318
v47:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1050))
1931719319
Jump bb8(v47, v55)
1931819320
bb8(v37:BasicObject, v38:BasicObject):
19319-
v50:BasicObject = Send v27, &block, :inner, v10, v37 # SendFallbackReason: Complex argument passing
19321+
v50:BasicObject = Send v27, &block, :inner, v10, v37 # SendFallbackReason: Send: block argument is not nil
1932019322
CheckInterrupts
1932119323
PopInlineFrame
1932219324
PatchPoint NoEPEscape(test)

zjit/src/stats.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,7 @@ make_counters! {
195195
exit_unhandled_splat,
196196
exit_unhandled_kwarg,
197197
exit_unhandled_block_arg,
198+
exit_block_arg_not_nil,
198199
exit_unknown_special_variable,
199200
exit_unhandled_hir_insn,
200201
exit_unhandled_yarv_insn,
@@ -267,6 +268,7 @@ make_counters! {
267268
send_fallback_send_no_profiles,
268269
send_fallback_send_not_optimized_method_type,
269270
send_fallback_send_not_optimized_need_permission,
271+
send_fallback_send_block_arg_not_nil,
270272
send_fallback_ccall_with_frame_too_many_args,
271273
send_fallback_argc_param_mismatch,
272274
// The call has at least one feature on the caller or callee side
@@ -611,6 +613,7 @@ pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter {
611613
UnhandledHIRUnknown(_) => exit_unhandled_hir_insn,
612614
UnhandledYARVInsn(_) => exit_unhandled_yarv_insn,
613615
UnhandledBlockArg => exit_unhandled_block_arg,
616+
BlockArgNotNil => exit_block_arg_not_nil,
614617
FixnumAddOverflow => exit_fixnum_add_overflow,
615618
FixnumSubOverflow => exit_fixnum_sub_overflow,
616619
FixnumMultOverflow => exit_fixnum_mult_overflow,
@@ -702,6 +705,7 @@ pub fn send_fallback_counter(reason: crate::hir::SendFallbackReason) -> Counter
702705
BmethodNonIseqProc => send_fallback_bmethod_non_iseq_proc,
703706
SendNotOptimizedMethodType(_) => send_fallback_send_not_optimized_method_type,
704707
SendNotOptimizedNeedPermission => send_fallback_send_not_optimized_need_permission,
708+
SendBlockArgNotNil => send_fallback_send_block_arg_not_nil,
705709
CCallWithFrameTooManyArgs => send_fallback_ccall_with_frame_too_many_args,
706710
ObjToStringNotString => send_fallback_obj_to_string_not_string,
707711
SuperCallWithBlock => send_fallback_super_call_with_block,

0 commit comments

Comments
 (0)