Skip to content

Commit 3efd8c6

Browse files
tekknolagitekknolagi
authored
ZJIT: Inline Kernel#class (ruby#15397)
We generally know the receiver's class from profile info. I see 600k of these when running lobsters.
1 parent 19f0df0 commit 3efd8c6

File tree

5 files changed

+138
-16
lines changed

5 files changed

+138
-16
lines changed

‎zjit/bindgen/src/main.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ fn main() {
107107
.allowlist_function("rb_obj_is_kind_of")
108108
.allowlist_function("rb_obj_frozen_p")
109109
.allowlist_function("rb_class_inherited_p")
110+
.allowlist_function("rb_class_real")
110111
.allowlist_type("ruby_encoding_consts")
111112
.allowlist_function("rb_hash_new")
112113
.allowlist_function("rb_hash_new_with_size")

‎zjit/src/cruby_bindings.inc.rs‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎zjit/src/cruby_methods.rs‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,11 +186,18 @@ pub fn init() -> Annotations {
186186
($module:ident, $method_name:literal, $return_type:expr) => {
187187
annotate_builtin!($module, $method_name, $return_type, no_gc, leaf, elidable)
188188
};
189-
($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
189+
($module:ident, $method_name:literal, $return_type:expr $(, $properties:ident)*) => {
190190
let mut props = FnProperties::default();
191191
props.return_type = $return_type;
192192
$(props.$properties = true;)+
193193
annotate_builtin_method(builtin_funcs, unsafe { $module }, $method_name, props);
194+
};
195+
($module:ident, $method_name:literal, $inline:ident, $return_type:expr $(, $properties:ident)*) => {
196+
let mut props = FnProperties::default();
197+
props.return_type = $return_type;
198+
props.inline = $inline;
199+
$(props.$properties = true;)+
200+
annotate_builtin_method(builtin_funcs, unsafe { $module }, $method_name, props);
194201
}
195202
}
196203

@@ -256,7 +263,7 @@ pub fn init() -> Annotations {
256263
annotate_builtin!(rb_mKernel, "Integer", types::Integer);
257264
// TODO(max): Annotate rb_mKernel#class as returning types::Class. Right now there is a subtle
258265
// type system bug that causes an issue if we make it return types::Class.
259-
annotate_builtin!(rb_mKernel, "class", types::HeapObject, leaf);
266+
annotate_builtin!(rb_mKernel, "class", inline_kernel_class, types::HeapObject, leaf);
260267
annotate_builtin!(rb_mKernel, "frozen?", types::BoolExact);
261268
annotate_builtin!(rb_cSymbol, "name", types::StringExact);
262269
annotate_builtin!(rb_cSymbol, "to_s", types::StringExact);
@@ -785,3 +792,10 @@ fn inline_kernel_respond_to_p(
785792
}
786793
Some(fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(result) }))
787794
}
795+
796+
fn inline_kernel_class(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
797+
let &[recv] = args else { return None; };
798+
let recv_class = fun.type_of(recv).runtime_exact_ruby_class()?;
799+
let real_class = unsafe { rb_class_real(recv_class) };
800+
Some(fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(real_class) }))
801+
}

‎zjit/src/hir.rs‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -862,6 +862,7 @@ pub enum Insn {
862862
// Invoke a builtin function
863863
InvokeBuiltin {
864864
bf: rb_builtin_function,
865+
recv: InsnId,
865866
args: Vec<InsnId>,
866867
state: InsnId,
867868
leaf: bool,
@@ -1922,7 +1923,7 @@ impl Function {
19221923
state,
19231924
reason,
19241925
},
1925-
&InvokeBuiltin { bf, ref args, state, leaf, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, leaf, return_type },
1926+
&InvokeBuiltin { bf, recv, ref args, state, leaf, return_type } => InvokeBuiltin { bf, recv: find!(recv), args: find_vec!(args), state, leaf, return_type },
19261927
&ArrayDup { val, state } => ArrayDup { val: find!(val), state },
19271928
&HashDup { val, state } => HashDup { val: find!(val), state },
19281929
&HashAref { hash, key, state } => HashAref { hash: find!(hash), key: find!(key), state },
@@ -2811,6 +2812,7 @@ impl Function {
28112812
self.push_insn(block, Insn::IncrCounter(Counter::inline_iseq_optimized_send_count));
28122813
let replacement = self.push_insn(block, Insn::InvokeBuiltin {
28132814
bf,
2815+
recv,
28142816
args: vec![recv],
28152817
state,
28162818
leaf: true,
@@ -3388,6 +3390,25 @@ impl Function {
33883390
continue;
33893391
}
33903392
}
3393+
Insn::InvokeBuiltin { bf, recv, args, state, .. } => {
3394+
let props = ZJITState::get_method_annotations().get_builtin_properties(&bf).unwrap_or_default();
3395+
// Try inlining the cfunc into HIR
3396+
let tmp_block = self.new_block(u32::MAX);
3397+
if let Some(replacement) = (props.inline)(self, tmp_block, recv, &args, state) {
3398+
// Copy contents of tmp_block to block
3399+
assert_ne!(block, tmp_block);
3400+
let insns = std::mem::take(&mut self.blocks[tmp_block.0].insns);
3401+
self.blocks[block.0].insns.extend(insns);
3402+
self.push_insn(block, Insn::IncrCounter(Counter::inline_cfunc_optimized_send_count));
3403+
self.make_equal_to(insn_id, replacement);
3404+
if self.type_of(replacement).bit_equal(types::Any) {
3405+
// Not set yet; infer type
3406+
self.insn_types[replacement.0] = self.infer_type(replacement);
3407+
}
3408+
self.remove_block(tmp_block);
3409+
continue;
3410+
}
3411+
}
33913412
_ => {}
33923413
}
33933414
self.push_insn_id(block, insn_id);
@@ -3704,13 +3725,13 @@ impl Function {
37043725
| &Insn::CCallVariadic { recv, ref args, state, .. }
37053726
| &Insn::CCallWithFrame { recv, ref args, state, .. }
37063727
| &Insn::SendWithoutBlockDirect { recv, ref args, state, .. }
3728+
| &Insn::InvokeBuiltin { recv, ref args, state, .. }
37073729
| &Insn::InvokeSuper { recv, ref args, state, .. } => {
37083730
worklist.push_back(recv);
37093731
worklist.extend(args);
37103732
worklist.push_back(state);
37113733
}
3712-
&Insn::InvokeBuiltin { ref args, state, .. }
3713-
| &Insn::InvokeBlock { ref args, state, .. } => {
3734+
&Insn::InvokeBlock { ref args, state, .. } => {
37143735
worklist.extend(args);
37153736
worklist.push_back(state)
37163737
}
@@ -4323,6 +4344,7 @@ impl Function {
43234344
| Insn::InvokeSuper { recv, ref args, .. }
43244345
| Insn::CCallWithFrame { recv, ref args, .. }
43254346
| Insn::CCallVariadic { recv, ref args, .. }
4347+
| Insn::InvokeBuiltin { recv, ref args, .. }
43264348
| Insn::ArrayInclude { target: recv, elements: ref args, .. } => {
43274349
self.assert_subtype(insn_id, recv, types::BasicObject)?;
43284350
for &arg in args {
@@ -4331,8 +4353,7 @@ impl Function {
43314353
Ok(())
43324354
}
43334355
// Instructions with a Vec of Ruby objects
4334-
Insn::InvokeBuiltin { ref args, .. }
4335-
| Insn::InvokeBlock { ref args, .. }
4356+
Insn::InvokeBlock { ref args, .. }
43364357
| Insn::NewArray { elements: ref args, .. }
43374358
| Insn::ArrayHash { elements: ref args, .. }
43384359
| Insn::ArrayMax { elements: ref args, .. } => {
@@ -5785,6 +5806,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
57855806

57865807
let insn_id = fun.push_insn(block, Insn::InvokeBuiltin {
57875808
bf,
5809+
recv: self_param,
57885810
args,
57895811
state: exit_id,
57905812
leaf,
@@ -5813,6 +5835,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
58135835

58145836
let insn_id = fun.push_insn(block, Insn::InvokeBuiltin {
58155837
bf,
5838+
recv: self_param,
58165839
args,
58175840
state: exit_id,
58185841
leaf,

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 92 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2611,9 +2611,10 @@ mod hir_opt_tests {
26112611
PatchPoint MethodRedefined(Module@0x1010, class@0x1018, cme:0x1020)
26122612
PatchPoint NoSingletonClass(Module@0x1010)
26132613
IncrCounter inline_iseq_optimized_send_count
2614-
v25:HeapObject = InvokeBuiltin leaf _bi20, v20
2614+
v26:Class[Module@0x1010] = Const Value(VALUE(0x1010))
2615+
IncrCounter inline_cfunc_optimized_send_count
26152616
CheckInterrupts
2616-
Return v25
2617+
Return v26
26172618
");
26182619
}
26192620

@@ -8937,15 +8938,15 @@ mod hir_opt_tests {
89378938
PatchPoint NoSingletonClass(C@0x1000)
89388939
v40:HeapObject[class_exact:C] = GuardType v6, HeapObject[class_exact:C]
89398940
IncrCounter inline_iseq_optimized_send_count
8940-
v43:HeapObject = InvokeBuiltin leaf _bi20, v40
8941+
v44:Class[C@0x1000] = Const Value(VALUE(0x1000))
8942+
IncrCounter inline_cfunc_optimized_send_count
89418943
v13:StaticSymbol[:_lex_actions] = Const Value(VALUE(0x1038))
89428944
v15:TrueClass = Const Value(true)
89438945
PatchPoint MethodRedefined(Class@0x1040, respond_to?@0x1048, cme:0x1050)
89448946
PatchPoint NoSingletonClass(Class@0x1040)
8945-
v47:ModuleSubclass[class_exact*:Class@VALUE(0x1040)] = GuardType v43, ModuleSubclass[class_exact*:Class@VALUE(0x1040)]
89468947
PatchPoint MethodRedefined(Class@0x1040, _lex_actions@0x1078, cme:0x1080)
89478948
PatchPoint NoSingletonClass(Class@0x1040)
8948-
v51:TrueClass = Const Value(true)
8949+
v52:TrueClass = Const Value(true)
89498950
IncrCounter inline_cfunc_optimized_send_count
89508951
CheckInterrupts
89518952
v24:StaticSymbol[:CORRECT] = Const Value(VALUE(0x10a8))
@@ -8976,14 +8977,96 @@ mod hir_opt_tests {
89768977
PatchPoint NoSingletonClass(C@0x1000)
89778978
v23:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
89788979
IncrCounter inline_iseq_optimized_send_count
8979-
v26:HeapObject = InvokeBuiltin leaf _bi20, v23
8980+
v27:Class[C@0x1000] = Const Value(VALUE(0x1000))
8981+
IncrCounter inline_cfunc_optimized_send_count
89808982
PatchPoint MethodRedefined(Class@0x1038, name@0x1040, cme:0x1048)
89818983
PatchPoint NoSingletonClass(Class@0x1038)
8982-
v30:ModuleSubclass[class_exact*:Class@VALUE(0x1038)] = GuardType v26, ModuleSubclass[class_exact*:Class@VALUE(0x1038)]
89838984
IncrCounter inline_cfunc_optimized_send_count
8984-
v32:StringExact|NilClass = CCall v30, :Module#name@0x1070
8985+
v33:StringExact|NilClass = CCall v27, :Module#name@0x1070
89858986
CheckInterrupts
8986-
Return v32
8987+
Return v33
8988+
");
8989+
}
8990+
8991+
#[test]
8992+
fn test_fold_kernel_class() {
8993+
eval(r#"
8994+
class C; end
8995+
def test(o) = o.class
8996+
test(C.new)
8997+
"#);
8998+
assert_snapshot!(hir_string("test"), @r"
8999+
fn test@<compiled>:3:
9000+
bb0():
9001+
EntryPoint interpreter
9002+
v1:BasicObject = LoadSelf
9003+
v2:BasicObject = GetLocal l0, SP@4
9004+
Jump bb2(v1, v2)
9005+
bb1(v5:BasicObject, v6:BasicObject):
9006+
EntryPoint JIT(0)
9007+
Jump bb2(v5, v6)
9008+
bb2(v8:BasicObject, v9:BasicObject):
9009+
PatchPoint MethodRedefined(C@0x1000, class@0x1008, cme:0x1010)
9010+
PatchPoint NoSingletonClass(C@0x1000)
9011+
v21:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
9012+
IncrCounter inline_iseq_optimized_send_count
9013+
v25:Class[C@0x1000] = Const Value(VALUE(0x1000))
9014+
IncrCounter inline_cfunc_optimized_send_count
9015+
CheckInterrupts
9016+
Return v25
9017+
");
9018+
}
9019+
9020+
#[test]
9021+
fn test_fold_fixnum_class() {
9022+
eval(r#"
9023+
def test = 5.class
9024+
test
9025+
"#);
9026+
assert_snapshot!(hir_string("test"), @r"
9027+
fn test@<compiled>:2:
9028+
bb0():
9029+
EntryPoint interpreter
9030+
v1:BasicObject = LoadSelf
9031+
Jump bb2(v1)
9032+
bb1(v4:BasicObject):
9033+
EntryPoint JIT(0)
9034+
Jump bb2(v4)
9035+
bb2(v6:BasicObject):
9036+
v10:Fixnum[5] = Const Value(5)
9037+
PatchPoint MethodRedefined(Integer@0x1000, class@0x1008, cme:0x1010)
9038+
IncrCounter inline_iseq_optimized_send_count
9039+
v21:Class[Integer@0x1000] = Const Value(VALUE(0x1000))
9040+
IncrCounter inline_cfunc_optimized_send_count
9041+
CheckInterrupts
9042+
Return v21
9043+
");
9044+
}
9045+
9046+
#[test]
9047+
fn test_fold_singleton_class() {
9048+
eval(r#"
9049+
def test = self.class
9050+
test
9051+
"#);
9052+
assert_snapshot!(hir_string("test"), @r"
9053+
fn test@<compiled>:2:
9054+
bb0():
9055+
EntryPoint interpreter
9056+
v1:BasicObject = LoadSelf
9057+
Jump bb2(v1)
9058+
bb1(v4:BasicObject):
9059+
EntryPoint JIT(0)
9060+
Jump bb2(v4)
9061+
bb2(v6:BasicObject):
9062+
PatchPoint MethodRedefined(Object@0x1000, class@0x1008, cme:0x1010)
9063+
PatchPoint NoSingletonClass(Object@0x1000)
9064+
v18:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
9065+
IncrCounter inline_iseq_optimized_send_count
9066+
v22:Class[Object@0x1038] = Const Value(VALUE(0x1038))
9067+
IncrCounter inline_cfunc_optimized_send_count
9068+
CheckInterrupts
9069+
Return v22
89879070
");
89889071
}
89899072

0 commit comments

Comments
 (0)