Skip to content

Commit 2a4fc46

Browse files
rwstaunerrwstauner
committed
[Backport #22426] 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: [] } 4.0 Backport of d81a11d
1 parent 229531a commit 2a4fc46

4 files changed

Lines changed: 208 additions & 18 deletions

File tree

‎test/ruby/test_zjit.rb‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,43 @@ def f(a)
406406
}, call_threshold: 2
407407
end
408408

409+
def test_kwargs_with_exit_and_local_invalidation
410+
assert_compiles ':ok', %q{
411+
def a(b:, c:)
412+
if c == :b
413+
return -> {}
414+
end
415+
Class # invalidate locals
416+
417+
raise "c is :b!" if c == :b
418+
end
419+
420+
def test
421+
# note opposite order of kwargs
422+
a(c: :c, b: :b)
423+
end
424+
425+
4.times { test }
426+
:ok
427+
}, call_threshold: 2
428+
end
429+
430+
def test_kwargs_with_max_direct_send_arg_count
431+
# Ensure that we only reorder the args when we _can_ use direct send (< 6 args).
432+
assert_compiles '[[1, 2, 3, 4, 5, 6, 7, 8]]', %q{
433+
def kwargs(five, six, a:, b:, c:, d:, e:, f:)
434+
[a, b, c, d, five, six, e, f]
435+
end
436+
437+
5.times.flat_map do
438+
[
439+
kwargs(5, 6, d: 4, c: 3, a: 1, b: 2, e: 7, f: 8),
440+
kwargs(5, 6, d: 4, c: 3, b: 2, a: 1, e: 7, f: 8)
441+
]
442+
end.uniq
443+
}, call_threshold: 2
444+
end
445+
409446
def test_setlocal_on_eval
410447
assert_compiles '1', %q{
411448
@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::{
@@ -2687,19 +2688,30 @@ impl Function {
26872688
}
26882689

26892690
let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
2690-
let processed_args = if !kwarg.is_null() {
2691+
let (send_state, processed_args) = if !kwarg.is_null() {
26912692
match self.reorder_keyword_arguments(&args, kwarg, iseq) {
2692-
Ok(reordered) => reordered,
2693+
Ok(reordered) => {
2694+
// Only use reordered state if args fit in C registers.
2695+
// Fallback to interpreter needs original order for kwarg handling.
2696+
// NOTE: This needs to match with the condition in codegen.rs.
2697+
if reordered.len() + 1 <= C_ARG_OPNDS.len() {
2698+
let new_state = self.frame_state(state).with_reordered_args(&reordered);
2699+
let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
2700+
(snapshot, reordered)
2701+
} else {
2702+
(state, reordered)
2703+
}
2704+
}
26932705
Err(reason) => {
26942706
self.set_dynamic_send_reason(insn_id, reason);
26952707
self.push_insn_id(block, insn_id); continue;
26962708
}
26972709
}
26982710
} else {
2699-
args.clone()
2711+
(state, args.clone())
27002712
};
27012713

2702-
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
2714+
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state: send_state });
27032715
self.make_equal_to(insn_id, send_direct);
27042716
} else if def_type == VM_METHOD_TYPE_BMETHOD {
27052717
let procv = unsafe { rb_get_def_bmethod_proc((*cme).def) };
@@ -2736,19 +2748,30 @@ impl Function {
27362748
}
27372749

27382750
let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
2739-
let processed_args = if !kwarg.is_null() {
2751+
let (send_state, processed_args) = if !kwarg.is_null() {
27402752
match self.reorder_keyword_arguments(&args, kwarg, iseq) {
2741-
Ok(reordered) => reordered,
2753+
Ok(reordered) => {
2754+
// Only use reordered state if args fit in C registers.
2755+
// Fallback to interpreter needs original order for kwarg handling.
2756+
// NOTE: This needs to match with the condition in codegen.rs.
2757+
if reordered.len() + 1 <= C_ARG_OPNDS.len() {
2758+
let new_state = self.frame_state(state).with_reordered_args(&reordered);
2759+
let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
2760+
(snapshot, reordered)
2761+
} else {
2762+
(state, reordered)
2763+
}
2764+
}
27422765
Err(reason) => {
27432766
self.set_dynamic_send_reason(insn_id, reason);
27442767
self.push_insn_id(block, insn_id); continue;
27452768
}
27462769
}
27472770
} else {
2748-
args.clone()
2771+
(state, args.clone())
27492772
};
27502773

2751-
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
2774+
let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state: send_state });
27522775
self.make_equal_to(insn_id, send_direct);
27532776
} else if def_type == VM_METHOD_TYPE_IVAR && args.is_empty() {
27542777
// Check if we're accessing ivars of a Class or Module object as they require single-ractor mode.
@@ -5002,6 +5025,15 @@ impl FrameState {
50025025
state.stack.clear();
50035026
state
50045027
}
5028+
5029+
/// Return itself with send args reordered. Used when kwargs are reordered for callee.
5030+
fn with_reordered_args(&self, reordered_args: &[InsnId]) -> Self {
5031+
let mut state = self.clone();
5032+
let args_start = state.stack.len() - reordered_args.len();
5033+
state.stack.truncate(args_start);
5034+
state.stack.extend_from_slice(reordered_args);
5035+
state
5036+
}
50055037
}
50065038

50075039
/// 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
@@ -2822,9 +2822,9 @@ mod hir_opt_tests {
28222822
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
28232823
PatchPoint NoSingletonClass(Object@0x1000)
28242824
v22:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2825-
v23:BasicObject = SendWithoutBlockDirect v22, :foo (0x1038), v11, v13
2825+
v24:BasicObject = SendWithoutBlockDirect v22, :foo (0x1038), v11, v13
28262826
CheckInterrupts
2827-
Return v23
2827+
Return v24
28282828
");
28292829
}
28302830

@@ -2852,9 +2852,9 @@ mod hir_opt_tests {
28522852
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
28532853
PatchPoint NoSingletonClass(Object@0x1000)
28542854
v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2855-
v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v13, v15, v11
2855+
v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v13, v15, v11
28562856
CheckInterrupts
2857-
Return v25
2857+
Return v26
28582858
");
28592859
}
28602860

@@ -2882,9 +2882,9 @@ mod hir_opt_tests {
28822882
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
28832883
PatchPoint NoSingletonClass(Object@0x1000)
28842884
v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2885-
v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v11, v15, v13
2885+
v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v11, v15, v13
28862886
CheckInterrupts
2887-
Return v25
2887+
Return v26
28882888
");
28892889
}
28902890

@@ -2939,16 +2939,16 @@ mod hir_opt_tests {
29392939
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
29402940
PatchPoint NoSingletonClass(Object@0x1000)
29412941
v37:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2942-
v38:BasicObject = SendWithoutBlockDirect v37, :foo (0x1038), v11, v13, v15
2942+
v39:BasicObject = SendWithoutBlockDirect v37, :foo (0x1038), v11, v13, v15
29432943
v20:Fixnum[1] = Const Value(1)
29442944
v22:Fixnum[2] = Const Value(2)
29452945
v24:Fixnum[4] = Const Value(4)
29462946
v26:Fixnum[3] = Const Value(3)
29472947
PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
29482948
PatchPoint NoSingletonClass(Object@0x1000)
2949-
v41:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2950-
v42:BasicObject = SendWithoutBlockDirect v41, :foo (0x1038), v20, v22, v26, v24
2951-
v30:ArrayExact = NewArray v38, v42
2949+
v42:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
2950+
v44:BasicObject = SendWithoutBlockDirect v42, :foo (0x1038), v20, v22, v26, v24
2951+
v30:ArrayExact = NewArray v39, v44
29522952
CheckInterrupts
29532953
Return v30
29542954
");

‎zjit/src/hir/tests.rs‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,16 @@ mod snapshot_tests {
1414
format!("{}", FunctionPrinter::with_snapshot(&function))
1515
}
1616

17+
#[track_caller]
18+
fn optimized_hir_string(method: &str) -> String {
19+
let iseq = crate::cruby::with_rubyvm(|| get_proc_iseq(&format!("{}.method(:{})", "self", method)));
20+
unsafe { crate::cruby::rb_zjit_profile_disable(iseq) };
21+
let mut function = iseq_to_hir(iseq).unwrap();
22+
function.optimize();
23+
function.validate().unwrap();
24+
format!("{}", FunctionPrinter::with_snapshot(&function))
25+
}
26+
1727
#[test]
1828
fn test_new_array_with_elements() {
1929
eval("def test(a, b) = [a, b]");
@@ -41,6 +51,117 @@ mod snapshot_tests {
4151
Return v18
4252
");
4353
}
54+
55+
#[test]
56+
fn test_send_direct_with_reordered_kwargs_has_snapshot() {
57+
eval("
58+
def foo(a:, b:, c:) = [a, b, c]
59+
def test = foo(c: 3, a: 1, b: 2)
60+
test
61+
test
62+
");
63+
assert_snapshot!(optimized_hir_string("test"), @r"
64+
fn test@<compiled>:3:
65+
bb0():
66+
EntryPoint interpreter
67+
v1:BasicObject = LoadSelf
68+
Jump bb2(v1)
69+
bb1(v4:BasicObject):
70+
EntryPoint JIT(0)
71+
Jump bb2(v4)
72+
bb2(v6:BasicObject):
73+
v8:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [] }
74+
PatchPoint NoTracePoint
75+
v11:Fixnum[3] = Const Value(3)
76+
v13:Fixnum[1] = Const Value(1)
77+
v15:Fixnum[2] = Const Value(2)
78+
v16:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v11, v13, v15], locals: [] }
79+
PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
80+
PatchPoint NoSingletonClass(Object@0x1010)
81+
v24:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
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: [] }
85+
PatchPoint NoTracePoint
86+
CheckInterrupts
87+
Return v26
88+
");
89+
}
90+
91+
#[test]
92+
fn test_send_direct_with_kwargs_in_order_has_snapshot() {
93+
eval("
94+
def foo(a:, b:) = [a, b]
95+
def test = foo(a: 1, b: 2)
96+
test
97+
test
98+
");
99+
assert_snapshot!(optimized_hir_string("test"), @r"
100+
fn test@<compiled>:3:
101+
bb0():
102+
EntryPoint interpreter
103+
v1:BasicObject = LoadSelf
104+
Jump bb2(v1)
105+
bb1(v4:BasicObject):
106+
EntryPoint JIT(0)
107+
Jump bb2(v4)
108+
bb2(v6:BasicObject):
109+
v8:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [] }
110+
PatchPoint NoTracePoint
111+
v11:Fixnum[1] = Const Value(1)
112+
v13:Fixnum[2] = Const Value(2)
113+
v14:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v11, v13], locals: [] }
114+
PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
115+
PatchPoint NoSingletonClass(Object@0x1010)
116+
v22:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
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: [] }
120+
PatchPoint NoTracePoint
121+
CheckInterrupts
122+
Return v24
123+
");
124+
}
125+
126+
#[test]
127+
fn test_send_direct_with_many_kwargs_no_reorder_snapshot() {
128+
eval("
129+
def foo(five, six, a:, b:, c:, d:, e:, f:) = [a, b, c, d, five, six, e, f]
130+
def test = foo(5, 6, d: 4, c: 3, a: 1, b: 2, e: 7, f: 8)
131+
test
132+
test
133+
");
134+
assert_snapshot!(optimized_hir_string("test"), @r"
135+
fn test@<compiled>:3:
136+
bb0():
137+
EntryPoint interpreter
138+
v1:BasicObject = LoadSelf
139+
Jump bb2(v1)
140+
bb1(v4:BasicObject):
141+
EntryPoint JIT(0)
142+
Jump bb2(v4)
143+
bb2(v6:BasicObject):
144+
v8:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [] }
145+
PatchPoint NoTracePoint
146+
v11:Fixnum[5] = Const Value(5)
147+
v13:Fixnum[6] = Const Value(6)
148+
v15:Fixnum[4] = Const Value(4)
149+
v17:Fixnum[3] = Const Value(3)
150+
v19:Fixnum[1] = Const Value(1)
151+
v21:Fixnum[2] = Const Value(2)
152+
v23:Fixnum[7] = Const Value(7)
153+
v25:Fixnum[8] = Const Value(8)
154+
v26:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v11, v13, v15, v17, v19, v21, v23, v25], locals: [] }
155+
PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
156+
PatchPoint NoSingletonClass(Object@0x1010)
157+
v34:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
158+
v35:BasicObject = SendWithoutBlockDirect v34, :foo (0x1048), v11, v13, v19, v21, v17, v15, v23, v25
159+
v28:Any = Snapshot FrameState { pc: 0x1050, stack: [v35], locals: [] }
160+
PatchPoint NoTracePoint
161+
CheckInterrupts
162+
Return v35
163+
");
164+
}
44165
}
45166

46167
#[cfg(test)]

0 commit comments

Comments
 (0)