Skip to content

Commit d81a11d

Browse files
rwstaunertekknolagi
rwstauner
authored and
tekknolagi
committed
ZJIT: Snapshot FrameState with reordered args before direct send
You can see the reordered args in the new Snapshot right before the DirectSend insn: v14:Any = Snapshot FrameState { pc: 0x00, stack: [v6, v11, v13], locals: [] } PatchPoint MethodRedefined(Object@0x00, a@0x00, cme:0x00) PatchPoint NoSingletonClass(Object@0x00) v22:HeapObject[class_exact*:Object@VALUE(0x00)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x00)] - v23:BasicObject = SendWithoutBlockDirect v22, :a (0x00), v13, v11 - v16:Any = Snapshot FrameState { pc: 0x00, stack: [v23], locals: [] } + v23:Any = Snapshot FrameState { pc: 0x00, stack: [v6, v13, v11], locals: [] } + v24:BasicObject = SendWithoutBlockDirect v22, :a (0x00), v13, v11 + v16:Any = Snapshot FrameState { pc: 0x00, stack: [v24], locals: [] }
1 parent 351616a commit d81a11d

4 files changed

Lines changed: 95 additions & 24 deletions

File tree

‎test/ruby/test_zjit.rb‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,43 @@ def f(a)
383383
}, call_threshold: 2
384384
end
385385

386+
def test_kwargs_with_exit_and_local_invalidation
387+
assert_compiles ':ok', %q{
388+
def a(b:, c:)
389+
if c == :b
390+
return -> {}
391+
end
392+
Class # invalidate locals
393+
394+
raise "c is :b!" if c == :b
395+
end
396+
397+
def test
398+
# note opposite order of kwargs
399+
a(c: :c, b: :b)
400+
end
401+
402+
4.times { test }
403+
:ok
404+
}, call_threshold: 2
405+
end
406+
407+
def test_kwargs_with_max_direct_send_arg_count
408+
# Ensure that we only reorder the args when we _can_ use direct send (< 6 args).
409+
assert_compiles '[[1, 2, 3, 4, 5, 6, 7, 8]]', %q{
410+
def kwargs(five, six, a:, b:, c:, d:, e:, f:)
411+
[a, b, c, d, five, six, e, f]
412+
end
413+
414+
5.times.flat_map do
415+
[
416+
kwargs(5, 6, d: 4, c: 3, a: 1, b: 2, e: 7, f: 8),
417+
kwargs(5, 6, d: 4, c: 3, b: 2, a: 1, e: 7, f: 8)
418+
]
419+
end.uniq
420+
}, call_threshold: 2
421+
end
422+
386423
def test_setlocal_on_eval
387424
assert_compiles '1', %q{
388425
@b = binding

‎zjit/src/hir.rs‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#![allow(clippy::if_same_then_else)]
77
#![allow(clippy::match_like_matches_macro)]
88
use crate::{
9+
backend::lir::C_ARG_OPNDS,
910
cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json
1011
};
1112
use std::{
@@ -2699,19 +2700,30 @@ impl Function {
26992700
}
27002701

27012702
let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
2702-
let processed_args = if !kwarg.is_null() {
2703+
let (send_state, processed_args) = if !kwarg.is_null() {
27032704
match self.reorder_keyword_arguments(&args, kwarg, iseq) {
2704-
Ok(reordered) => reordered,
2705+
Ok(reordered) => {
2706+
// Only use reordered state if args fit in C registers.
2707+
// Fallback to interpreter needs original order for kwarg handling.
2708+
// NOTE: This needs to match with the condition in codegen.rs.
2709+
if reordered.len() + 1 <= C_ARG_OPNDS.len() {
2710+
let new_state = self.frame_state(state).with_reordered_args(&reordered);
2711+
let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
2712+
(snapshot, reordered)
2713+
} else {
2714+
(state, reordered)
2715+
}
2716+
}
27052717
Err(reason) => {
27062718
self.set_dynamic_send_reason(insn_id, reason);
27072719
self.push_insn_id(block, insn_id); continue;
27082720
}
27092721
}
27102722
} else {
2711-
args.clone()
2723+
(state, args.clone())
27122724
};
27132725

2714-
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
2726+
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state: send_state });
27152727
self.make_equal_to(insn_id, send_direct);
27162728
} else if def_type == VM_METHOD_TYPE_BMETHOD {
27172729
let procv = unsafe { rb_get_def_bmethod_proc((*cme).def) };
@@ -2748,19 +2760,30 @@ impl Function {
27482760
}
27492761

27502762
let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
2751-
let processed_args = if !kwarg.is_null() {
2763+
let (send_state, processed_args) = if !kwarg.is_null() {
27522764
match self.reorder_keyword_arguments(&args, kwarg, iseq) {
2753-
Ok(reordered) => reordered,
2765+
Ok(reordered) => {
2766+
// Only use reordered state if args fit in C registers.
2767+
// Fallback to interpreter needs original order for kwarg handling.
2768+
// NOTE: This needs to match with the condition in codegen.rs.
2769+
if reordered.len() + 1 <= C_ARG_OPNDS.len() {
2770+
let new_state = self.frame_state(state).with_reordered_args(&reordered);
2771+
let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
2772+
(snapshot, reordered)
2773+
} else {
2774+
(state, reordered)
2775+
}
2776+
}
27542777
Err(reason) => {
27552778
self.set_dynamic_send_reason(insn_id, reason);
27562779
self.push_insn_id(block, insn_id); continue;
27572780
}
27582781
}
27592782
} else {
2760-
args.clone()
2783+
(state, args.clone())
27612784
};
27622785

2763-
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
2786+
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state: send_state });
27642787
self.make_equal_to(insn_id, send_direct);
27652788
} else if def_type == VM_METHOD_TYPE_IVAR && args.is_empty() {
27662789
// Check if we're accessing ivars of a Class or Module object as they require single-ractor mode.
@@ -5045,6 +5068,15 @@ impl FrameState {
50455068
state.stack.clear();
50465069
state
50475070
}
5071+
5072+
/// Return itself with send args reordered. Used when kwargs are reordered for callee.
5073+
fn with_reordered_args(&self, reordered_args: &[InsnId]) -> Self {
5074+
let mut state = self.clone();
5075+
let args_start = state.stack.len() - reordered_args.len();
5076+
state.stack.truncate(args_start);
5077+
state.stack.extend_from_slice(reordered_args);
5078+
state
5079+
}
50485080
}
50495081

50505082
/// Print adaptor for [`FrameState`]. See [`PtrPrintMap`].

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2816,9 +2816,9 @@ mod hir_opt_tests {
28162816
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
28172817
PatchPoint NoSingletonClass(Object@0x1000)
28182818
v22:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2819-
v23:BasicObject = SendWithoutBlockDirect v22, :foo (0x1038), v11, v13
2819+
v24:BasicObject = SendWithoutBlockDirect v22, :foo (0x1038), v11, v13
28202820
CheckInterrupts
2821-
Return v23
2821+
Return v24
28222822
");
28232823
}
28242824

@@ -2846,9 +2846,9 @@ mod hir_opt_tests {
28462846
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
28472847
PatchPoint NoSingletonClass(Object@0x1000)
28482848
v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2849-
v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v13, v15, v11
2849+
v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v13, v15, v11
28502850
CheckInterrupts
2851-
Return v25
2851+
Return v26
28522852
");
28532853
}
28542854

@@ -2876,9 +2876,9 @@ mod hir_opt_tests {
28762876
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
28772877
PatchPoint NoSingletonClass(Object@0x1000)
28782878
v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2879-
v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v11, v15, v13
2879+
v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v11, v15, v13
28802880
CheckInterrupts
2881-
Return v25
2881+
Return v26
28822882
");
28832883
}
28842884

@@ -2933,16 +2933,16 @@ mod hir_opt_tests {
29332933
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
29342934
PatchPoint NoSingletonClass(Object@0x1000)
29352935
v37:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2936-
v38:BasicObject = SendWithoutBlockDirect v37, :foo (0x1038), v11, v13, v15
2936+
v39:BasicObject = SendWithoutBlockDirect v37, :foo (0x1038), v11, v13, v15
29372937
v20:Fixnum[1] = Const Value(1)
29382938
v22:Fixnum[2] = Const Value(2)
29392939
v24:Fixnum[4] = Const Value(4)
29402940
v26:Fixnum[3] = Const Value(3)
29412941
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
29422942
PatchPoint NoSingletonClass(Object@0x1000)
2943-
v41:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2944-
v42:BasicObject = SendWithoutBlockDirect v41, :foo (0x1038), v20, v22, v26, v24
2945-
v30:ArrayExact = NewArray v38, v42
2943+
v42:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2944+
v44:BasicObject = SendWithoutBlockDirect v42, :foo (0x1038), v20, v22, v26, v24
2945+
v30:ArrayExact = NewArray v39, v44
29462946
CheckInterrupts
29472947
Return v30
29482948
");

‎zjit/src/hir/tests.rs‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,11 +79,12 @@ mod snapshot_tests {
7979
PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
8080
PatchPoint NoSingletonClass(Object@0x1010)
8181
v24:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
82-
v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1048), v13, v15, v11
83-
v18:Any = Snapshot FrameState { pc: 0x1050, stack: [v25], locals: [] }
82+
v25:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v13, v15, v11], locals: [] }
83+
v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1048), v13, v15, v11
84+
v18:Any = Snapshot FrameState { pc: 0x1050, stack: [v26], locals: [] }
8485
PatchPoint NoTracePoint
8586
CheckInterrupts
86-
Return v25
87+
Return v26
8788
");
8889
}
8990

@@ -113,11 +114,12 @@ mod snapshot_tests {
113114
PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
114115
PatchPoint NoSingletonClass(Object@0x1010)
115116
v22:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
116-
v23:BasicObject = SendWithoutBlockDirect v22, :foo (0x1048), v11, v13
117-
v16:Any = Snapshot FrameState { pc: 0x1050, stack: [v23], locals: [] }
117+
v23:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v11, v13], locals: [] }
118+
v24:BasicObject = SendWithoutBlockDirect v22, :foo (0x1048), v11, v13
119+
v16:Any = Snapshot FrameState { pc: 0x1050, stack: [v24], locals: [] }
118120
PatchPoint NoTracePoint
119121
CheckInterrupts
120-
Return v23
122+
Return v24
121123
");
122124
}
123125

0 commit comments

Comments
 (0)