ZJIT: Re-profile whole ISEQ in non-final versions

Reset each existing instruction profile and reinstall zjit_* instructions after compiling a non-final version. Do the same for inlined ISEQs so materialized frames can collect fresh profiles.

Let interpreter profiling instructions collect the samples and make recompile exits only observe whether the exiting instruction has finished profiling.

Author: k0kubun

zjit/bindgen/src/main.rs

@@ -110,6 +110,7 @@ fn main() {
         .allowlist_function("rb_funcallv")
         .allowlist_function("rb_protect")
         .allowlist_function("rb_zjit_profile_disable")
+        .allowlist_function("rb_zjit_profile_enable")
         .allowlist_function("rb_zjit_insn_to_bare_insn")
         .allowlist_function("rb_zjit_hash_new_size")
         .allowlist_function("rb_zjit_class_allocate_instance_fastpath")

zjit/src/codegen.rs

@@ -17,6 +17,7 @@ use crate::invariants::{
 };
 use crate::gc::append_gc_offsets;
 use crate::payload::{IseqCodePtrs, IseqStatus, IseqVersion, IseqVersionRef, JITFrame, get_or_create_iseq_payload};
+use crate::profile::reset_profiles_remaining;
 use crate::state::ZJITState;
 use crate::stats::{CompileError, exit_counter_for_compile_error, exit_counter_for_unhandled_hir_insn, incr_counter, incr_counter_by, send_fallback_counter, send_fallback_counter_for_method_type, send_fallback_counter_for_super_method_type, send_fallback_counter_ptr_for_opcode, send_fallback_counter_for_optimized_method_type};
 use crate::stats::{counter_ptr, with_time_stat, trace_compile_phase, Counter, Counter::{compile_time_ns, exit_compile_error}};
@@ -3361,6 +3362,10 @@ fn compile_iseq(iseq: IseqPtr) -> Result<Function, CompileError> {
         trace_compile_phase("optimize", || function.optimize());
     }
     function.dump_hir();
+    let non_final_version = get_or_create_iseq_payload(iseq).versions.len() + 1 < max_iseq_versions();
+    if non_final_version {
+        reset_profiles_remaining(iseq);
+    }
     Ok(function)
 }
 

zjit/src/codegen_tests.rs

@@ -122,6 +122,28 @@ fn test_putobject() {
     "), @"1");
 }
 
+#[test]
+fn test_recompile_exit_waits_for_interpreter_profiles() {
+    set_call_threshold(2);
+    eval("
+        def recompile_profile_window(a, b) = a + b
+        recompile_profile_window(1, 2)
+        recompile_profile_window(1, 2)
+    ");
+
+    let iseq = get_method_iseq("self", "recompile_profile_window");
+    let num_profiles = get_option!(num_profiles);
+    for _ in 0..num_profiles {
+        eval("recompile_profile_window(1.5, 2.5)");
+    }
+    let payload = get_or_create_iseq_payload(iseq);
+    assert!(!unsafe { payload.versions.last().unwrap().as_ref() }.is_invalidated());
+
+    eval("recompile_profile_window(1.5, 2.5)");
+    let payload = get_or_create_iseq_payload(iseq);
+    assert!(unsafe { payload.versions.last().unwrap().as_ref() }.is_invalidated());
+}
+
 #[test]
 fn test_dupstring() {
     eval(r##"

zjit/src/cruby_bindings.inc.rs

@@ -6,9 +6,7 @@
 #![allow(clippy::if_same_then_else)]
 #![allow(clippy::match_like_matches_macro)]
 use crate::{
-    backend::lir::C_ARG_OPNDS,
-    cast::IntoUsize, codegen::{max_iseq_versions}, cruby::*, invariants::{self, iseq_seen_ep_escape}, payload::get_or_create_iseq_payload, options::{debug, get_option, DumpHIR, InlineDepth}, state::ZJITState, json::Json,
-    state,
+    backend::lir::C_ARG_OPNDS, cast::IntoUsize, codegen::max_iseq_versions, cruby::*, invariants::{self, iseq_seen_ep_escape}, json::Json, options::{DumpHIR, InlineDepth, debug, get_option}, payload::get_or_create_iseq_payload, profile::reset_profiles_remaining, state::{self, ZJITState},
 };
 use std::{
     cell::RefCell, collections::{HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, c_int, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter,
@@ -9751,6 +9749,10 @@ fn add_iseq_to_hir(
             None => {},
         }
     }
+    if matches!(mode, AddIseqMode::Inlined { .. }) {
+        // Materialized inlined frames also need fresh interpreter profiles.
+        reset_profiles_remaining(iseq);
+    }
 
     Ok(AddIseqResult { body_entry_block, profiles })
 }

zjit/src/hir.rs

@@ -5602,11 +5602,15 @@ mod hir_opt_tests {
           v36:NilClass = Const Value(nil)
           Jump bb8(v36, v13)
         bb8(v26:BasicObject, v27:BasicObject):
-          v39:BasicObject = Send v24, &block, :then, v26 # SendFallbackReason: Send: block argument is not nil
+          v56:NilClass = GuardBitEquals v26, Value(nil) recompile
+          PatchPoint MethodRedefined(Integer@0x1008, then@0x1010, cme:0x1018)
+          PushInlineFrame v24 (0x1040)
+          v76:BasicObject = InvokeBuiltin <inline_expr>, v24
           CheckInterrupts
-          Return v39
+          PopInlineFrame
+          Return v76
         bb4(v44:BasicObject, v45:Falsy, v46:BasicObject):
-          v50:StaticSymbol[:skip] = Const Value(VALUE(0x1008))
+          v50:StaticSymbol[:skip] = Const Value(VALUE(0x1048))
           CheckInterrupts
           Return v50
         ");
@@ -9175,12 +9179,17 @@ mod hir_opt_tests {
           Jump bb3(v6, v7)
         bb3(v9:HeapBasicObject, v10:BasicObject):
           v17:Fixnum[5] = Const Value(5)
+          v21:CBool = HasType v10, ObjectSubclass[class_exact:C]
+          CondBranch v21, bb5(), bb6()
+        bb5():
+          v24:ObjectSubclass[class_exact:C] = RefineType v10, ObjectSubclass[class_exact:C]
           PatchPoint MethodRedefined(C@0x1008, foo=@0x1010, cme:0x1018)
-          v28:ObjectSubclass[class_exact:C] = GuardType v10, ObjectSubclass[class_exact:C] recompile
-          v30:CShape = LoadField v28, :shape_id@0x1040
-          v31:CShape[0x1041] = GuardBitEquals v30, CShape(0x1041)
-          StoreField v28, :@foo@0x1042, v17
-          WriteBarrier v28, v17
+          SetIvar v24, :@foo, v17
+          Jump bb4(v17)
+        bb6():
+          v27:BasicObject = Send v10, :foo=, v17 # SendFallbackReason: Send: polymorphic fallback
+          Jump bb4(v27)
+        bb4(v20:BasicObject):
           CheckInterrupts
           Return v17
         ");
@@ -18809,27 +18818,25 @@ mod hir_opt_tests {
         bb4():
           PatchPoint NoEPEscape(f)
           v44:Fixnum[1] = Const Value(1)
-          v48:CBool = HasType v12, Flonum
+          v48:CBool = HasType v12, Fixnum
           CondBranch v48, bb10(), bb11()
         bb10():
-          v51:Flonum = RefineType v12, Flonum
-          PatchPoint MethodRedefined(Float@0x1008, +@0x1010, cme:0x1018)
-          v86:Float = FloatAdd v51, v44
+          v51:Fixnum = RefineType v12, Fixnum
+          PatchPoint MethodRedefined(Integer@0x1008, +@0x1010, cme:0x1018)
+          v86:Fixnum = FixnumAdd v51, v44
           Jump bb9(v86)
         bb11():
-          v54:CBool = HasType v12, Fixnum
+          v54:CBool = HasType v12, Flonum
           CondBranch v54, bb12(), bb13()
         bb12():
-          v57:Fixnum = RefineType v12, Fixnum
-          PatchPoint MethodRedefined(Integer@0x1040, +@0x1010, cme:0x1048)
-          v89:Fixnum = FixnumAdd v57, v44
+          v57:Flonum = RefineType v12, Flonum
+          PatchPoint MethodRedefined(Float@0x1040, +@0x1010, cme:0x1048)
+          v89:Float = FloatAdd v57, v44
           Jump bb9(v89)
         bb13():
-          PatchPoint MethodRedefined(Float@0x1008, +@0x1010, cme:0x1018)
-          v92:Flonum = GuardType v12, Flonum recompile
-          v93:Float = FloatAdd v92, v44
-          Jump bb9(v93)
-        bb9(v47:Float|Fixnum):
+          v60:BasicObject = Send v12, :+, v44 # SendFallbackReason: Send: polymorphic fallback
+          Jump bb9(v60)
+        bb9(v47:BasicObject):
           PatchPoint SingleRactorMode
           v69:CShape = LoadField v11, :shape_id@0x1001
           v70:CShape[0x1002] = Const CShape(0x1002)

zjit/src/hir/opt_tests.rs

@@ -122,36 +122,21 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
     }
 }
 
-/// Profile the instruction at the current CFP for a recompile side exit.
+/// Return whether the interpreter finished profiling the current instruction.
 pub fn profile_recompile_insn(ec: EcPtr) -> bool {
-    let profiler = &mut Profiler::new(ec);
-    let pc = unsafe { get_cfp_pc(profiler.cfp) };
-    let bare_opcode = unsafe {
-        rb_zjit_insn_to_bare_insn(rb_iseq_opcode_at_pc(profiler.iseq, pc))
-    } as ruby_vminsn_type;
-    let profile = &mut get_or_create_iseq_payload(profiler.iseq).profile;
-
-    let is_send = matches!(bare_opcode, YARVINSN_send | YARVINSN_opt_send_without_block);
-    // For now, send recompile exits only fill in missing profiles. Once the send site
-    // has finished profiling, don't recompile it on later exits.
-    if is_send && profile.done_profiling_at(profiler.insn_idx) {
-        return false;
-    }
-    // For now, non-send recompile exits reset the profiling counter before requesting recompilation
-    // so that we can collect enough samples.
-    if !is_send && profile.done_profiling_at(profiler.insn_idx) {
-        profile.entry_mut(profiler.insn_idx)
-            .set_profiles_remaining(get_option!(num_profiles));
-    }
+    let profiler = &Profiler::new(ec);
+    get_or_create_iseq_payload(profiler.iseq).profile.done_profiling_at(profiler.insn_idx)
+}
 
-    // If this opcode can't be sampled here, this exit has no profile data to collect.
-    if !profile_insn_sample(bare_opcode, profiler, profile) {
-        return false;
+/// Reset existing profile counters and install profiling instructions throughout an ISEQ.
+/// Newly reached instructions initialize their counters from the same option.
+pub(crate) fn reset_profiles_remaining(iseq: IseqPtr) {
+    let profile = &mut get_or_create_iseq_payload(iseq).profile;
+    let num_profiles = get_option!(num_profiles);
+    for entry in &mut profile.entries {
+        entry.profiles_remaining = num_profiles;
     }
-
-    let entry = profile.entry_mut(profiler.insn_idx);
-    entry.profiles_remaining = entry.profiles_remaining.saturating_sub(1);
-    entry.profiles_remaining == 0
+    unsafe { rb_zjit_profile_enable(iseq) };
 }
 
 /// Return the argc as stated in the calldata plus: