Skip to content

Commit f9a849a

Browse files
rwstaunerrwstauner
committed
ZJIT: Optimize send with a nil block to SendWithoutBlockDirect
1 parent c197707 commit f9a849a

File tree

4 files changed

+51
-9
lines changed

4 files changed

+51
-9
lines changed

‎zjit/src/hir.rs‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -506,6 +506,7 @@ pub enum SideExitReason {
506506
UnhandledYARVInsn(u32),
507507
UnhandledCallType(CallType),
508508
UnhandledBlockArg,
509+
BlockArgNotNil,
509510
TooManyKeywordParameters,
510511
FixnumAddOverflow,
511512
FixnumSubOverflow,
@@ -664,6 +665,8 @@ pub enum SendFallbackReason {
664665
SendCfuncArrayVariadic,
665666
SendNotOptimizedMethodType(MethodType),
666667
SendNotOptimizedNeedPermission,
668+
/// The block argument is not nil, so we can't optimize to SendWithoutBlockDirect
669+
SendBlockArgNotNil,
667670
CCallWithFrameTooManyArgs,
668671
ObjToStringNotString,
669672
TooManyArgsForLir,
@@ -738,6 +741,7 @@ impl Display for SendFallbackReason {
738741
SendCfuncVariadic => write!(f, "Send: C function is variadic"),
739742
SendCfuncArrayVariadic => write!(f, "Send: C function expects array variadic"),
740743
SendNotOptimizedMethodType(method_type) => write!(f, "Send: unsupported method type {:?}", method_type),
744+
SendBlockArgNotNil => write!(f, "Send: block argument is not nil"),
741745
CCallWithFrameTooManyArgs => write!(f, "CCallWithFrame: too many arguments"),
742746
ObjToStringNotString => write!(f, "ObjToString: result is not a string"),
743747
TooManyArgsForLir => write!(f, "Too many arguments for LIR"),
@@ -3704,9 +3708,38 @@ impl Function {
37043708
def_type = unsafe { get_cme_def_type(cme) };
37053709
}
37063710

3711+
// Check if we can optimize `foo(&block)` where block is nil to a send without block
3712+
let mut send_block = send_block;
3713+
let mut args = args;
3714+
let mut stripped_nil_block = false;
3715+
if send_block == Some(BlockHandler::BlockArg) && def_type == VM_METHOD_TYPE_ISEQ {
3716+
// The block arg is the last element in args
3717+
if let Some(&block_arg) = args.last() {
3718+
let block_arg_type = self.type_of(block_arg);
3719+
if block_arg_type.is_subtype(types::NilClass) {
3720+
// Block arg is known to be nil - strip it and treat as no block
3721+
self.push_insn(block, Insn::GuardBitEquals {
3722+
val: block_arg,
3723+
expected: Const::Value(Qnil),
3724+
reason: SideExitReason::BlockArgNotNil,
3725+
state,
3726+
});
3727+
args = args[..args.len() - 1].to_vec();
3728+
send_block = None;
3729+
stripped_nil_block = true;
3730+
} else {
3731+
// Can't prove block arg is nil
3732+
self.set_dynamic_send_reason(insn_id, SendBlockArgNotNil);
3733+
self.push_insn_id(block, insn_id); continue;
3734+
}
3735+
}
3736+
}
3737+
37073738
// If the call site info indicates that the `Function` has overly complex arguments, then do not optimize into a `SendDirect`.
37083739
// Optimized methods(`VM_METHOD_TYPE_OPTIMIZED`) handle their own argument constraints (e.g., kw_splat for Proc call).
3709-
if def_type != VM_METHOD_TYPE_OPTIMIZED && unspecializable_call_type(flags) {
3740+
// Mask out ARGS_BLOCKARG only if we've already handled the nil block arg case above.
3741+
let flags_for_check = if stripped_nil_block { flags & !VM_CALL_ARGS_BLOCKARG } else { flags };
3742+
if def_type != VM_METHOD_TYPE_OPTIMIZED && unspecializable_call_type(flags_for_check) {
37103743
self.count_complex_call_features(block, flags);
37113744
self.set_dynamic_send_reason(insn_id, ComplexArgPass);
37123745
self.push_insn_id(block, insn_id); continue;

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8019,7 +8019,7 @@ mod hir_opt_tests {
80198019
v22:BasicObject = LoadField v18, :block@0x1010
80208020
Jump bb6(v22, v22)
80218021
bb6(v16:BasicObject, v17:BasicObject):
8022-
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
8022+
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
80238023
CheckInterrupts
80248024
Return v29
80258025
");
@@ -8057,7 +8057,7 @@ mod hir_opt_tests {
80578057
v22:BasicObject = LoadField v18, :block@0x1002
80588058
Jump bb6(v22, v22)
80598059
bb6(v16:BasicObject, v17:BasicObject):
8060-
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
8060+
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
80618061
CheckInterrupts
80628062
Return v29
80638063
");
@@ -8096,7 +8096,7 @@ mod hir_opt_tests {
80968096
v17:BasicObject = LoadField v13, :block@0x1010
80978097
Jump bb6(v17)
80988098
bb6(v12:BasicObject):
8099-
v24:BasicObject = Send v10, &block, :map, v12 # SendFallbackReason: Complex argument passing
8099+
v24:BasicObject = Send v10, &block, :map, v12 # SendFallbackReason: Send: block argument is not nil
81008100
CheckInterrupts
81018101
Return v24
81028102
");
@@ -8218,7 +8218,7 @@ mod hir_opt_tests {
82188218
Jump bb3(v5, v6)
82198219
bb3(v8:BasicObject, v9:NilClass):
82208220
v13:StaticSymbol[:to_s] = Const Value(VALUE(0x1000))
8221-
v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Complex argument passing
8221+
v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Send: block argument is not nil
82228222
CheckInterrupts
82238223
Return v19
82248224
");
@@ -8249,9 +8249,12 @@ mod hir_opt_tests {
82498249
Jump bb3(v5, v6)
82508250
bb3(v8:BasicObject, v9:NilClass):
82518251
v13:NilClass = Const Value(nil)
8252-
v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Complex argument passing
8252+
v25:NilClass = GuardBitEquals v13, Value(nil)
8253+
PatchPoint MethodRedefined(NilClass@0x1000, foo@0x1008, cme:0x1010)
8254+
v27:NilClass = GuardType v8, NilClass
8255+
v29:Fixnum[42] = Const Value(42)
82538256
CheckInterrupts
8254-
Return v19
8257+
Return v29
82558258
");
82568259
}
82578260

@@ -11602,7 +11605,7 @@ mod hir_opt_tests {
1160211605
Jump bb3(v4)
1160311606
bb3(v6:BasicObject):
1160411607
v11:StaticSymbol[:the_block] = Const Value(VALUE(0x1000))
11605-
v13:BasicObject = Send v6, &block, :callee, v11 # SendFallbackReason: Complex argument passing
11608+
v13:BasicObject = Send v6, &block, :callee, v11 # SendFallbackReason: Send: block argument is not nil
1160611609
CheckInterrupts
1160711610
Return v13
1160811611
");
@@ -11645,7 +11648,7 @@ mod hir_opt_tests {
1164511648
v22:BasicObject = LoadField v18, :block@0x1010
1164611649
Jump bb6(v22, v22)
1164711650
bb6(v16:BasicObject, v17:BasicObject):
11648-
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
11651+
v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
1164911652
CheckInterrupts
1165011653
Return v29
1165111654
");

‎zjit/src/profile.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
9797
let cd: *const rb_call_data = profiler.insn_opnd(0).as_ptr();
9898
let argc = unsafe { vm_ci_argc((*cd).ci) };
9999
// Profile all the arguments and self (+1).
100+
// Note: For sends with ARGS_BLOCKARG, the block arg is NOT profiled here.
101+
// The block arg type is determined from HIR static type inference.
100102
profile_operands(profiler, profile, (argc + 1) as usize);
101103
}
102104
YARVINSN_splatkw => profile_operands(profiler, profile, 2),

‎zjit/src/stats.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,7 @@ make_counters! {
193193
exit_unhandled_splat,
194194
exit_unhandled_kwarg,
195195
exit_unhandled_block_arg,
196+
exit_block_arg_not_nil,
196197
exit_unknown_special_variable,
197198
exit_unhandled_hir_insn,
198199
exit_unhandled_yarv_insn,
@@ -263,6 +264,7 @@ make_counters! {
263264
send_fallback_send_no_profiles,
264265
send_fallback_send_not_optimized_method_type,
265266
send_fallback_send_not_optimized_need_permission,
267+
send_fallback_send_block_arg_not_nil,
266268
send_fallback_ccall_with_frame_too_many_args,
267269
send_fallback_argc_param_mismatch,
268270
// The call has at least one feature on the caller or callee side
@@ -587,6 +589,7 @@ pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter {
587589
UnhandledHIRUnknown(_) => exit_unhandled_hir_insn,
588590
UnhandledYARVInsn(_) => exit_unhandled_yarv_insn,
589591
UnhandledBlockArg => exit_unhandled_block_arg,
592+
BlockArgNotNil => exit_block_arg_not_nil,
590593
FixnumAddOverflow => exit_fixnum_add_overflow,
591594
FixnumSubOverflow => exit_fixnum_sub_overflow,
592595
FixnumMultOverflow => exit_fixnum_mult_overflow,
@@ -676,6 +679,7 @@ pub fn send_fallback_counter(reason: crate::hir::SendFallbackReason) -> Counter
676679
BmethodNonIseqProc => send_fallback_bmethod_non_iseq_proc,
677680
SendNotOptimizedMethodType(_) => send_fallback_send_not_optimized_method_type,
678681
SendNotOptimizedNeedPermission => send_fallback_send_not_optimized_need_permission,
682+
SendBlockArgNotNil => send_fallback_send_block_arg_not_nil,
679683
CCallWithFrameTooManyArgs => send_fallback_ccall_with_frame_too_many_args,
680684
ObjToStringNotString => send_fallback_obj_to_string_not_string,
681685
SuperCallWithBlock => send_fallback_super_call_with_block,

0 commit comments

Comments
 (0)