Skip to content

Commit 6902bc7

Browse files
st0012st0012
committed
ZJIT: Refactor receiver type resolution
1 parent a731080 commit 6902bc7

File tree

1 file changed

+98
-57
lines changed

1 file changed

+98
-57
lines changed

‎zjit/src/hir.rs‎

Lines changed: 98 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,22 @@ impl std::fmt::Display for SideExitReason {
563563
}
564564
}
565565

566+
/// Result of resolving the receiver type for method dispatch optimization.
567+
/// Represents whether we know the receiver's class statically at compile-time,
568+
/// have profiled type information, or know nothing about it.
569+
pub enum ReceiverTypeResolution {
570+
/// The receiver's class is statically known at JIT compile-time (no guard needed)
571+
StaticallyKnown { class: VALUE },
572+
/// The receiver has a monomorphic profile (single type observed, guard needed)
573+
Monomorphic { class: VALUE, profiled_type: ProfiledType },
574+
/// The receiver has a skewed polymorphic profile (dominant type with some other types, guard needed)
575+
SkewedPolymorphic { class: VALUE, profiled_type: ProfiledType },
576+
/// The receiver is polymorphic (multiple types, none dominant)
577+
Polymorphic,
578+
/// No profile information available for the receiver
579+
NoProfile,
580+
}
581+
566582
/// Reason why a send-ish instruction cannot be optimized from a fallback instruction
567583
#[derive(Debug, Clone, Copy)]
568584
pub enum SendFallbackReason {
@@ -2133,26 +2149,49 @@ impl Function {
21332149
None
21342150
}
21352151

2136-
/// Return whether a given HIR instruction as profiled by the interpreter is polymorphic or
2137-
/// whether it lacks a profile entirely.
2152+
/// Resolve the receiver type for method dispatch optimization.
21382153
///
2139-
/// * `Some(true)` if polymorphic
2140-
/// * `Some(false)` if monomorphic
2141-
/// * `None` if no profiled information so far
2142-
fn is_polymorphic_at(&self, insn: InsnId, iseq_insn_idx: usize) -> Option<bool> {
2143-
let profiles = self.profiles.as_ref()?;
2144-
let entries = profiles.types.get(&iseq_insn_idx)?;
2145-
let insn = self.chase_insn(insn);
2154+
/// Takes the receiver's Type, receiver HIR instruction, and ISEQ instruction index.
2155+
/// Performs a single iteration through profile data to determine the receiver type.
2156+
///
2157+
/// Returns:
2158+
/// - `StaticallyKnown` if the receiver's exact class is known at compile-time
2159+
/// - `Monomorphic`/`SkewedPolymorphic` if we have usable profile data
2160+
/// - `Polymorphic` if the receiver has multiple types
2161+
/// - `NoProfile` if we have no type information
2162+
fn resolve_receiver_type(&self, recv: InsnId, recv_type: Type, insn_idx: usize) -> ReceiverTypeResolution {
2163+
if let Some(class) = recv_type.runtime_exact_ruby_class() {
2164+
return ReceiverTypeResolution::StaticallyKnown { class };
2165+
}
2166+
let Some(profiles) = self.profiles.as_ref() else {
2167+
return ReceiverTypeResolution::NoProfile;
2168+
};
2169+
let Some(entries) = profiles.types.get(&insn_idx) else {
2170+
return ReceiverTypeResolution::NoProfile;
2171+
};
2172+
let recv = self.chase_insn(recv);
2173+
21462174
for (entry_insn, entry_type_summary) in entries {
2147-
if self.union_find.borrow().find_const(*entry_insn) == insn {
2148-
if !entry_type_summary.is_monomorphic() && !entry_type_summary.is_skewed_polymorphic() {
2149-
return Some(true);
2150-
} else {
2151-
return Some(false);
2175+
if self.union_find.borrow().find_const(*entry_insn) == recv {
2176+
if entry_type_summary.is_monomorphic() {
2177+
let profiled_type = entry_type_summary.bucket(0);
2178+
return ReceiverTypeResolution::Monomorphic {
2179+
class: profiled_type.class(),
2180+
profiled_type,
2181+
};
2182+
} else if entry_type_summary.is_skewed_polymorphic() {
2183+
let profiled_type = entry_type_summary.bucket(0);
2184+
return ReceiverTypeResolution::SkewedPolymorphic {
2185+
class: profiled_type.class(),
2186+
profiled_type,
2187+
};
2188+
} else if entry_type_summary.is_polymorphic() {
2189+
return ReceiverTypeResolution::Polymorphic;
21522190
}
21532191
}
21542192
}
2155-
None
2193+
2194+
ReceiverTypeResolution::NoProfile
21562195
}
21572196

21582197
pub fn assume_expected_cfunc(&mut self, block: BlockId, class: VALUE, method_id: ID, cfunc: *mut c_void, state: InsnId) -> bool {
@@ -2299,24 +2338,24 @@ impl Function {
22992338
self.try_rewrite_uminus(block, insn_id, recv, state),
23002339
Insn::SendWithoutBlock { mut recv, cd, args, state, .. } => {
23012340
let frame_state = self.frame_state(state);
2302-
let (klass, profiled_type) = if let Some(klass) = self.type_of(recv).runtime_exact_ruby_class() {
2303-
// If we know the class statically, use it to fold the lookup at compile-time.
2304-
(klass, None)
2305-
} else {
2306-
// If we know that self is reasonably monomorphic from profile information, guard and use it to fold the lookup at compile-time.
2307-
// TODO(max): Figure out how to handle top self?
2308-
let Some(recv_type) = self.profiled_type_of_at(recv, frame_state.insn_idx) else {
2341+
let (klass, profiled_type) = match self.resolve_receiver_type(recv, self.type_of(recv), frame_state.insn_idx) {
2342+
ReceiverTypeResolution::StaticallyKnown { class } => (class, None),
2343+
ReceiverTypeResolution::Monomorphic { class, profiled_type }
2344+
| ReceiverTypeResolution::SkewedPolymorphic { class, profiled_type } => (class, Some(profiled_type)),
2345+
ReceiverTypeResolution::Polymorphic => {
23092346
if get_option!(stats) {
2310-
match self.is_polymorphic_at(recv, frame_state.insn_idx) {
2311-
Some(true) => self.set_dynamic_send_reason(insn_id, SendWithoutBlockPolymorphic),
2312-
// If the class isn't known statically, then it should not also be monomorphic
2313-
Some(false) => panic!("Should not have monomorphic profile at this point in this branch"),
2314-
None => self.set_dynamic_send_reason(insn_id, SendWithoutBlockNoProfiles),
2315-
}
2347+
self.set_dynamic_send_reason(insn_id, SendWithoutBlockPolymorphic);
23162348
}
2317-
self.push_insn_id(block, insn_id); continue;
2318-
};
2319-
(recv_type.class(), Some(recv_type))
2349+
self.push_insn_id(block, insn_id);
2350+
continue;
2351+
}
2352+
ReceiverTypeResolution::NoProfile => {
2353+
if get_option!(stats) {
2354+
self.set_dynamic_send_reason(insn_id, SendWithoutBlockNoProfiles);
2355+
}
2356+
self.push_insn_id(block, insn_id);
2357+
continue;
2358+
}
23202359
};
23212360
let ci = unsafe { get_call_data_ci(cd) }; // info about the call site
23222361

@@ -2492,22 +2531,24 @@ impl Function {
24922531
// The actual optimization is done in reduce_send_to_ccall.
24932532
Insn::Send { recv, cd, state, .. } => {
24942533
let frame_state = self.frame_state(state);
2495-
let klass = if let Some(klass) = self.type_of(recv).runtime_exact_ruby_class() {
2496-
// If we know the class statically, use it to fold the lookup at compile-time.
2497-
klass
2498-
} else {
2499-
let Some(recv_type) = self.profiled_type_of_at(recv, frame_state.insn_idx) else {
2534+
let klass = match self.resolve_receiver_type(recv, self.type_of(recv), frame_state.insn_idx) {
2535+
ReceiverTypeResolution::StaticallyKnown { class } => class,
2536+
ReceiverTypeResolution::Monomorphic { class, .. }
2537+
| ReceiverTypeResolution::SkewedPolymorphic { class, .. } => class,
2538+
ReceiverTypeResolution::Polymorphic => {
25002539
if get_option!(stats) {
2501-
match self.is_polymorphic_at(recv, frame_state.insn_idx) {
2502-
Some(true) => self.set_dynamic_send_reason(insn_id, SendPolymorphic),
2503-
// If the class isn't known statically, then it should not also be monomorphic
2504-
Some(false) => panic!("Should not have monomorphic profile at this point in this branch"),
2505-
None => self.set_dynamic_send_reason(insn_id, SendNoProfiles),
2506-
}
2540+
self.set_dynamic_send_reason(insn_id, SendPolymorphic);
25072541
}
2508-
self.push_insn_id(block, insn_id); continue;
2509-
};
2510-
recv_type.class()
2542+
self.push_insn_id(block, insn_id);
2543+
continue;
2544+
}
2545+
ReceiverTypeResolution::NoProfile => {
2546+
if get_option!(stats) {
2547+
self.set_dynamic_send_reason(insn_id, SendNoProfiles);
2548+
}
2549+
self.push_insn_id(block, insn_id);
2550+
continue;
2551+
}
25112552
};
25122553
let ci = unsafe { get_call_data_ci(cd) }; // info about the call site
25132554
let mid = unsafe { vm_ci_mid(ci) };
@@ -2769,12 +2810,12 @@ impl Function {
27692810
let method_id = unsafe { rb_vm_ci_mid(call_info) };
27702811

27712812
// If we have info about the class of the receiver
2772-
let (recv_class, profiled_type) = if let Some(class) = self_type.runtime_exact_ruby_class() {
2773-
(class, None)
2774-
} else {
2775-
let iseq_insn_idx = fun.frame_state(state).insn_idx;
2776-
let Some(recv_type) = fun.profiled_type_of_at(recv, iseq_insn_idx) else { return Err(()) };
2777-
(recv_type.class(), Some(recv_type))
2813+
let iseq_insn_idx = fun.frame_state(state).insn_idx;
2814+
let (recv_class, profiled_type) = match fun.resolve_receiver_type(recv, self_type, iseq_insn_idx) {
2815+
ReceiverTypeResolution::StaticallyKnown { class } => (class, None),
2816+
ReceiverTypeResolution::Monomorphic { class, profiled_type }
2817+
| ReceiverTypeResolution::SkewedPolymorphic { class, profiled_type } => (class, Some(profiled_type)),
2818+
ReceiverTypeResolution::Polymorphic | ReceiverTypeResolution::NoProfile => return Err(()),
27782819
};
27792820

27802821
// Do method lookup
@@ -2874,12 +2915,12 @@ impl Function {
28742915
let method_id = unsafe { rb_vm_ci_mid(call_info) };
28752916

28762917
// If we have info about the class of the receiver
2877-
let (recv_class, profiled_type) = if let Some(class) = self_type.runtime_exact_ruby_class() {
2878-
(class, None)
2879-
} else {
2880-
let iseq_insn_idx = fun.frame_state(state).insn_idx;
2881-
let Some(recv_type) = fun.profiled_type_of_at(recv, iseq_insn_idx) else { return Err(()) };
2882-
(recv_type.class(), Some(recv_type))
2918+
let iseq_insn_idx = fun.frame_state(state).insn_idx;
2919+
let (recv_class, profiled_type) = match fun.resolve_receiver_type(recv, self_type, iseq_insn_idx) {
2920+
ReceiverTypeResolution::StaticallyKnown { class } => (class, None),
2921+
ReceiverTypeResolution::Monomorphic { class, profiled_type }
2922+
| ReceiverTypeResolution::SkewedPolymorphic { class, profiled_type } => (class, Some(profiled_type)),
2923+
ReceiverTypeResolution::Polymorphic | ReceiverTypeResolution::NoProfile => return Err(()),
28832924
};
28842925

28852926
// Do method lookup

0 commit comments

Comments
 (0)