ZJIT: Implement SingleRactorMode invalidation (#14121)

* ZJIT: Implement SingleRactorMode invalidation

* ZJIT: Add macro for compiling jumps

* ZJIT: Fix typo in comment

* YJIT: Fix typo in comment

* ZJIT: Avoid using unexported types in zjit.h

`enum ruby_vminsn_type` is declared in `insns.inc` and is not exported.
Using it in `zjit.h` would cause build errors when the file including it
doesn't include `insns.inc`.
This commit is contained in:
Stan Lo 2025-08-06 21:51:41 +01:00 committed by GitHub
parent e378a21a32
commit 4a70f946a7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 85 additions and 35 deletions

View file

@ -4,7 +4,7 @@ use std::ffi::{c_int, c_void};
use crate::asm::Label;
use crate::backend::current::{Reg, ALLOC_REGS};
use crate::invariants::{track_bop_assumption, track_cme_assumption, track_stable_constant_names_assumption};
use crate::invariants::{track_bop_assumption, track_cme_assumption, track_single_ractor_assumption, track_stable_constant_names_assumption};
use crate::gc::{get_or_create_iseq_payload, append_gc_offsets};
use crate::state::ZJITState;
use crate::stats::{counter_ptr, Counter};
@ -542,9 +542,9 @@ fn gen_patch_point(jit: &mut JITState, asm: &mut Assembler, invariant: &Invarian
let side_exit_ptr = cb.resolve_label(label);
track_stable_constant_names_assumption(idlist, code_ptr, side_exit_ptr);
}
_ => {
debug!("ZJIT: gen_patch_point: unimplemented invariant {invariant:?}");
return;
Invariant::SingleRactorMode => {
let side_exit_ptr = cb.resolve_label(label);
track_single_ractor_assumption(code_ptr, side_exit_ptr);
}
}
});

View file

@ -1,7 +1,20 @@
use std::{collections::{HashMap, HashSet}};
use std::{collections::{HashMap, HashSet}, mem};
use crate::{backend::lir::{asm_comment, Assembler}, cruby::{rb_callable_method_entry_t, ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag, ID}, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
macro_rules! compile_jumps {
($cb:expr, $jumps:expr, $($comment_args:tt)*) => {
for jump in $jumps {
$cb.with_write_ptr(jump.from, |cb| {
let mut asm = Assembler::new();
asm_comment!(asm, $($comment_args)*);
asm.jmp(jump.to.into());
asm.compile(cb).expect("can write existing code");
});
}
};
}
#[derive(Debug, Eq, Hash, PartialEq)]
struct Jump {
from: CodePtr,
@ -26,6 +39,9 @@ pub struct Invariants {
/// Map from constant ID to patch points that assume the constant hasn't been redefined
constant_state_patch_points: HashMap<ID, HashSet<Jump>>,
/// Set of patch points that assume that the interpreter is running with only one ractor
single_ractor_patch_points: HashSet<Jump>,
}
/// Called when a basic operator is redefined. Note that all the blocks assuming
@ -46,14 +62,7 @@ pub extern "C" fn rb_zjit_bop_redefined(klass: RedefinitionFlag, bop: ruby_basic
debug!("BOP is redefined: {}", bop);
// Invalidate all patch points for this BOP
for jump in jumps {
cb.with_write_ptr(jump.from, |cb| {
let mut asm = Assembler::new();
asm_comment!(asm, "BOP is redefined: {}", bop);
asm.jmp(jump.to.into());
asm.compile(cb).expect("can write existing code");
});
}
compile_jumps!(cb, jumps, "BOP is redefined: {}", bop);
cb.mark_all_executable();
}
@ -159,14 +168,8 @@ pub extern "C" fn rb_zjit_cme_invalidate(cme: *const rb_callable_method_entry_t)
debug!("CME is invalidated: {:?}", cme);
// Invalidate all patch points for this CME
for jump in jumps {
cb.with_write_ptr(jump.from, |cb| {
let mut asm = Assembler::new();
asm_comment!(asm, "CME is invalidated: {:?}", cme);
asm.jmp(jump.to.into());
asm.compile(cb).expect("can write existing code");
});
}
compile_jumps!(cb, jumps, "CME is invalidated: {:?}", cme);
cb.mark_all_executable();
}
});
@ -187,16 +190,38 @@ pub extern "C" fn rb_zjit_constant_state_changed(id: ID) {
debug!("Constant state changed: {:?}", id);
// Invalidate all patch points for this constant ID
for jump in jumps {
cb.with_write_ptr(jump.from, |cb| {
let mut asm = Assembler::new();
asm_comment!(asm, "Constant state changed: {:?}", id);
asm.jmp(jump.to.into());
asm.compile(cb).expect("can write existing code");
});
}
compile_jumps!(cb, jumps, "Constant state changed: {:?}", id);
cb.mark_all_executable();
}
});
}
/// Track the JIT code that assumes that the interpreter is running with only one ractor
pub fn track_single_ractor_assumption(patch_point_ptr: CodePtr, side_exit_ptr: CodePtr) {
let invariants = ZJITState::get_invariants();
invariants.single_ractor_patch_points.insert(Jump {
from: patch_point_ptr,
to: side_exit_ptr,
});
}
/// Callback for when Ruby is about to spawn a ractor. In that case we need to
/// invalidate every block that is assuming single ractor mode.
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_before_ractor_spawn() {
// If ZJIT isn't enabled, do nothing
if !zjit_enabled_p() {
return;
}
with_vm_lock(src_loc!(), || {
let cb = ZJITState::get_code_block();
let jumps = mem::take(&mut ZJITState::get_invariants().single_ractor_patch_points);
// Invalidate all patch points for single ractor mode
compile_jumps!(cb, jumps, "Another ractor spawned, invalidating single ractor mode assumption");
cb.mark_all_executable();
});
}

View file

@ -39,10 +39,10 @@ impl Profiler {
/// API called from zjit_* instruction. opcode is the bare (non-zjit_*) instruction.
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
pub extern "C" fn rb_zjit_profile_insn(bare_opcode: u32, ec: EcPtr) {
with_vm_lock(src_loc!(), || {
let mut profiler = Profiler::new(ec);
profile_insn(&mut profiler, bare_opcode);
profile_insn(&mut profiler, bare_opcode as ruby_vminsn_type);
});
}