Fix profile count

Author: jacob-shops

zjit/src/backend/lir.rs

@@ -4,12 +4,12 @@ use std::mem::take;
 use std::rc::Rc;
 use crate::bitset::BitSet;
 use crate::codegen::{local_size_and_idx_to_ep_offset, perf_symbol_range_start, perf_symbol_range_end};
-use crate::cruby::{IseqPtr, RUBY_OFFSET_CFP_ISEQ, RUBY_OFFSET_CFP_JIT_RETURN, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, vm_stack_canary};
+use crate::cruby::{vm_stack_canary, IseqPtr, YarvInsnIdx, RUBY_OFFSET_CFP_ISEQ, RUBY_OFFSET_CFP_JIT_RETURN, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32};
 use crate::hir::{Invariant, SideExitReason};
 use crate::hir;
 use crate::options::{TraceExits, PerfMap, get_option};
 use crate::cruby::VALUE;
-use crate::payload::IseqVersionRef;
+use crate::payload::{IseqVersionRef, get_or_create_iseq_payload};
 use crate::stats::{exit_counter_ptr, exit_counter_ptr_for_opcode, side_exit_counter, CompileError};
 use crate::virtualmem::CodePtr;
 use crate::asm::{CodeBlock, Label};
@@ -2402,7 +2402,8 @@ impl Assembler
 
         fn compile_exit_recompile(asm: &mut Assembler, exit: &SideExit) {
             if let Some(recompile) = &exit.recompile {
-
+                let payload = get_or_create_iseq_payload(exit.iseq);
+                payload.reset_profiles_remaining(recompile.insn_idx as YarvInsnIdx);
                 use crate::codegen::exit_recompile;
                 asm_comment!(asm, "profile and maybe recompile");
                 asm_ccall!(asm, exit_recompile,

zjit/src/payload.rs

@@ -3,6 +3,7 @@ use std::ptr::NonNull;
 use crate::codegen::IseqCallRef;
 use crate::stats::CompileError;
 use crate::{cruby::*, profile::IseqProfile, virtualmem::CodePtr};
+use crate::options::get_option;
 
 pub use crate::jit_frame::JITFrame;
 
@@ -26,6 +27,11 @@ impl IseqPayload {
             was_invalidated_for_singleton_class_creation: false,
         }
     }
+
+    pub fn reset_profiles_remaining(&mut self, insn_idx: YarvInsnIdx) {
+        let num_profiles = get_option!(num_profiles);
+        self.profile.entry_mut(insn_idx).set_profiles_remaining(num_profiles);
+    }
 }
 
 /// JIT code version. When the same ISEQ is compiled with a different assumption, a new version is created.

zjit/src/profile.rs

@@ -417,7 +417,7 @@ impl ProfiledType {
 
 /// Per-instruction profile entry, stored sparsely in a sorted Vec.
 #[derive(Debug)]
-struct ProfileEntry {
+pub struct ProfileEntry {
     /// YARV instruction index
     insn_idx: u32,
     /// Type information of YARV instruction operands
@@ -426,6 +426,12 @@ struct ProfileEntry {
     profiles_remaining: NumProfiles,
 }
 
+impl ProfileEntry {
+    pub fn set_profiles_remaining(&mut self, num_profiles: NumProfiles) {
+        self.profiles_remaining = num_profiles;
+    }
+}
+
 #[derive(Debug)]
 pub struct IseqProfile {
     /// Sparse storage of per-instruction profile data, sorted by instruction index.
@@ -445,7 +451,7 @@ impl IseqProfile {
     }
 
     /// Get or create a mutable profile entry for the given instruction index.
-    fn entry_mut(&mut self, insn_idx: YarvInsnIdx) -> &mut ProfileEntry {
+    pub fn entry_mut(&mut self, insn_idx: YarvInsnIdx) -> &mut ProfileEntry {
         let idx = insn_idx as u32;
         match self.entries.binary_search_by_key(&idx, |e| e.insn_idx) {
             Ok(i) => &mut self.entries[i],
@@ -479,6 +485,11 @@ impl IseqProfile {
     pub fn profile_send_at(&mut self, iseq: IseqPtr, insn_idx: YarvInsnIdx, sp: *const VALUE, argc: usize) -> bool {
         let n = argc + 1; // args + receiver
         let entry = self.entry_mut(insn_idx);
+        // Reset profiling if the previous round already finished (stale YARV profiles).
+        // This ensures we collect num_profiles samples of the new shapes before recompiling.
+        if entry.profiles_remaining == 0 {
+            entry.profiles_remaining = get_option!(num_profiles);
+        }
         if entry.opnd_types.is_empty() {
             entry.opnd_types.resize(n, TypeDistribution::new());
         }