YJIT: Verify the assumption of leaf C calls (#10002)

This commit is contained in:
Takashi Kokubun 2024-02-20 13:42:29 -08:00 committed by GitHub
parent d4b4b53bc0
commit 9216a2ac43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 139 additions and 19 deletions

View File

@ -2046,6 +2046,16 @@ assert_equal '[97, :nil, 97, :nil, :raised]', %q{
[getbyte("a", 0), getbyte("a", 1), getbyte("a", -1), getbyte("a", -2), getbyte("a", "a")]
} unless rjit_enabled? # Not yet working on RJIT
# non-leaf String#byteslice
assert_equal 'TypeError', %q{
def ccall = "".byteslice(nil, nil)
begin
ccall
rescue => e
e.class
end
}
# Test << operator on string subclass
assert_equal 'abab', %q{
class MyString < String; end

View File

@ -337,6 +337,17 @@ vm_push_frame_debug_counter_inc(
#define vm_push_frame_debug_counter_inc(ec, cfp, t) /* void */
#endif
// Return a poison value to be set above the stack top to verify leafness.
VALUE
rb_vm_stack_canary(void)
{
#if VM_CHECK_MODE > 0
return vm_stack_canary;
#else
return 0;
#endif
}
STATIC_ASSERT(VM_ENV_DATA_INDEX_ME_CREF, VM_ENV_DATA_INDEX_ME_CREF == -2);
STATIC_ASSERT(VM_ENV_DATA_INDEX_SPECVAL, VM_ENV_DATA_INDEX_SPECVAL == -1);
STATIC_ASSERT(VM_ENV_DATA_INDEX_FLAGS, VM_ENV_DATA_INDEX_FLAGS == -0);

View File

@ -3,8 +3,8 @@ use std::fmt;
use std::convert::From;
use std::mem::take;
use crate::codegen::{gen_outlined_exit, gen_counted_exit};
use crate::cruby::{VALUE, SIZEOF_VALUE_I32};
use crate::virtualmem::{CodePtr};
use crate::cruby::{vm_stack_canary, SIZEOF_VALUE_I32, VALUE};
use crate::virtualmem::CodePtr;
use crate::asm::{CodeBlock, OutlinedCb};
use crate::core::{Context, RegTemps, MAX_REG_TEMPS};
use crate::options::*;
@ -229,11 +229,16 @@ impl Opnd
/// Calculate Opnd::Stack's index from the stack bottom.
pub fn stack_idx(&self) -> u8 {
self.get_stack_idx().unwrap()
}
/// Calculate Opnd::Stack's index from the stack bottom if it's Opnd::Stack.
pub fn get_stack_idx(&self) -> Option<u8> {
match self {
Opnd::Stack { idx, stack_size, .. } => {
(*stack_size as isize - *idx as isize - 1) as u8
Some((*stack_size as isize - *idx as isize - 1) as u8)
},
_ => unreachable!(),
_ => None
}
}
@ -1017,6 +1022,9 @@ pub struct Assembler {
/// Stack size for Target::SideExit
side_exit_stack_size: Option<u8>,
/// If true, the next ccall() should verify its leafness
leaf_ccall: bool,
}
impl Assembler
@ -1034,6 +1042,7 @@ impl Assembler
side_exits,
side_exit_pc: None,
side_exit_stack_size: None,
leaf_ccall: false,
}
}
@ -1561,6 +1570,21 @@ impl Assembler
pub fn into_draining_iter(self) -> AssemblerDrainingIterator {
AssemblerDrainingIterator::new(self)
}
/// Return true if the next ccall() is expected to be leaf.
pub fn get_leaf_ccall(&mut self) -> bool {
self.leaf_ccall
}
/// Assert that the next ccall() is going to be leaf.
pub fn expect_leaf_ccall(&mut self) {
self.leaf_ccall = true;
}
/// Undo expect_leaf_ccall() as an exception.
pub fn allow_non_leaf_ccall(&mut self) {
self.leaf_ccall = false;
}
}
/// A struct that allows iterating through an assembler's instructions and
@ -1661,6 +1685,9 @@ impl Assembler {
}
pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd {
// Let vm_check_canary() assert this ccall's leafness if leaf_ccall is set
let canary_opnd = self.set_stack_canary(&opnds);
let old_temps = self.ctx.get_reg_temps(); // with registers
// Spill stack temp registers since they are caller-saved registers.
// Note that this doesn't spill stack temps that are already popped
@ -1680,9 +1707,37 @@ impl Assembler {
// so rollback the manipulated RegTemps to a spilled version.
self.ctx.set_reg_temps(new_temps);
// Clear the canary after use
if let Some(canary_opnd) = canary_opnd {
self.mov(canary_opnd, 0.into());
}
out
}
/// Let vm_check_canary() assert the leafness of this ccall if leaf_ccall is set
fn set_stack_canary(&mut self, opnds: &Vec<Opnd>) -> Option<Opnd> {
// Use the slot right above the stack top for verifying leafness.
let canary_opnd = self.stack_opnd(-1);
// If the slot is already used, which is a valid optimization to avoid spills,
// give up the verification.
let canary_opnd = if cfg!(debug_assertions) && self.leaf_ccall && opnds.iter().all(|opnd|
opnd.get_stack_idx() != canary_opnd.get_stack_idx()
) {
asm_comment!(self, "set stack canary");
self.mov(canary_opnd, vm_stack_canary().into());
Some(canary_opnd)
} else {
None
};
// Avoid carrying the flag to the next instruction whether we verified it or not.
self.leaf_ccall = false;
canary_opnd
}
pub fn cmp(&mut self, left: Opnd, right: Opnd) {
self.push_insn(Insn::Cmp { left, right });
}
@ -1959,6 +2014,16 @@ impl Assembler {
out
}
/// Verify the leafness of the given block
pub fn with_leaf_ccall<F, R>(&mut self, mut block: F) -> R
where F: FnMut(&mut Self) -> R {
let old_leaf_ccall = self.leaf_ccall;
self.leaf_ccall = true;
let ret = block(self);
self.leaf_ccall = old_leaf_ccall;
ret
}
/// Add a label at the current position
pub fn write_label(&mut self, target: Target) {
assert!(target.unwrap_label_idx() < self.label_names.len());

View File

@ -411,8 +411,10 @@ fn jit_prepare_non_leaf_call(
jit: &mut JITState,
asm: &mut Assembler
) {
// Prepare for GC. This also sets PC, which prepares for showing a backtrace.
jit_prepare_call_with_gc(jit, asm);
// Prepare for GC. Setting PC also prepares for showing a backtrace.
jit.record_boundary_patch_point = true; // VM lock could trigger invalidation
jit_save_pc(jit, asm); // for allocation tracing
gen_save_sp(asm); // protect objects from GC
// In case the routine calls Ruby methods, it can set local variables
// through Kernel#binding, rb_debug_inspector API, and other means.
@ -429,6 +431,9 @@ fn jit_prepare_call_with_gc(
jit.record_boundary_patch_point = true; // VM lock could trigger invalidation
jit_save_pc(jit, asm); // for allocation tracing
gen_save_sp(asm); // protect objects from GC
// Expect a leaf ccall(). You should use jit_prepare_non_leaf_call() if otherwise.
asm.expect_leaf_ccall();
}
/// Record the current codeblock write position for rewriting into a jump into
@ -1083,6 +1088,9 @@ pub fn gen_single_block(
jit_perf_symbol_push!(jit, &mut asm, &insn_name(opcode), PerfMap::Codegen);
status = gen_fn(&mut jit, &mut asm, ocb);
jit_perf_symbol_pop!(jit, &mut asm, PerfMap::Codegen);
#[cfg(debug_assertions)]
assert!(!asm.get_leaf_ccall(), "ccall() wasn't used after leaf_ccall was set in {}", insn_name(opcode));
}
// If we can't compile this instruction
@ -5376,18 +5384,26 @@ fn jit_rb_str_byteslice(
argc: i32,
_known_recv_class: Option<VALUE>,
) -> bool {
asm_comment!(asm, "String#byteslice");
if argc != 2 {
return false
}
// Raises when non-integers are passed in
jit_prepare_non_leaf_call(jit, asm);
let len = asm.stack_opnd(0);
let beg = asm.stack_opnd(1);
let recv = asm.stack_opnd(2);
// rb_str_byte_substr should be leaf if indexes are fixnums
match (asm.ctx.get_opnd_type(beg.into()), asm.ctx.get_opnd_type(len.into())) {
(Type::Fixnum, Type::Fixnum) => {},
// Raises when non-integers are passed in, which requires the method frame
// to be pushed for the backtrace
_ => return false,
}
asm_comment!(asm, "String#byteslice");
// rb_str_byte_substr allocates a substring
jit_prepare_call_with_gc(jit, asm);
let ret_opnd = asm.ccall(rb_str_byte_substr as *const u8, vec![recv, beg, len]);
asm.stack_pop(3);
@ -5398,7 +5414,7 @@ fn jit_rb_str_byteslice(
}
fn jit_rb_str_getbyte(
jit: &mut JITState,
_jit: &mut JITState,
asm: &mut Assembler,
_ocb: &mut OutlinedCb,
_ci: *const rb_callinfo,
@ -5407,7 +5423,6 @@ fn jit_rb_str_getbyte(
_argc: i32,
_known_recv_class: Option<VALUE>,
) -> bool {
asm_comment!(asm, "String#getbyte");
extern "C" {
fn rb_str_getbyte(str: VALUE, index: VALUE) -> VALUE;
}
@ -5417,9 +5432,11 @@ fn jit_rb_str_getbyte(
// rb_str_getbyte should be leaf if the index is a fixnum
if asm.ctx.get_opnd_type(index.into()) != Type::Fixnum {
// Raises when non-integers are passed in
jit_prepare_non_leaf_call(jit, asm);
// Raises when non-integers are passed in, which requires the method frame
// to be pushed for the backtrace
return false;
}
asm_comment!(asm, "String#getbyte");
let ret_opnd = asm.ccall(rb_str_getbyte as *const u8, vec![recv, index]);
asm.stack_pop(2); // Keep them on stack during ccall for GC
@ -5507,10 +5524,11 @@ fn jit_rb_str_concat(
// Guard that the concat argument is a string
guard_object_is_string(asm, asm.stack_opnd(0), StackOpnd(0), Counter::guard_send_not_string);
// Guard buffers from GC since rb_str_buf_append may allocate. During the VM lock on GC,
// other Ractors may trigger global invalidation, so we need ctx.clear_local_types().
// PC is used on errors like Encoding::CompatibilityError raised by rb_str_buf_append.
// Guard buffers from GC since rb_str_buf_append may allocate.
jit_prepare_non_leaf_call(jit, asm);
// rb_str_buf_append may raise Encoding::CompatibilityError, but we accept compromised
// backtraces on this method since the interpreter does the same thing on opt_ltlt.
asm.allow_non_leaf_ccall();
asm.spill_temps(); // For ccall. Unconditionally spill them for RegTemps consistency.
let concat_arg = asm.stack_pop(1);
@ -6069,7 +6087,10 @@ fn gen_send_cfunc(
(cfunc_argc == -1 || argc == cfunc_argc) {
let expected_stack_after = asm.ctx.get_stack_size() as i32 - argc;
if let Some(known_cfunc_codegen) = lookup_cfunc_codegen(unsafe { (*cme).def }) {
if perf_call!("gen_send_cfunc: ", known_cfunc_codegen(jit, asm, ocb, ci, cme, block, argc, recv_known_class)) {
// We don't push a frame for specialized cfunc codegen, so the generated code must be leaf.
if asm.with_leaf_ccall(|asm|
perf_call!("gen_send_cfunc: ", known_cfunc_codegen(jit, asm, ocb, ci, cme, block, argc, recv_known_class))
) {
assert_eq!(expected_stack_after, asm.ctx.get_stack_size() as i32);
gen_counter_incr(asm, Counter::num_send_cfunc_inline);
// cfunc codegen generated code. Terminate the block so

View File

@ -141,6 +141,7 @@ extern "C" {
ic: ICVARC,
) -> VALUE;
pub fn rb_vm_ic_hit_p(ic: IC, reg_ep: *const VALUE) -> bool;
pub fn rb_vm_stack_canary() -> VALUE;
}
// Renames
@ -260,6 +261,18 @@ pub fn iseq_opcode_at_idx(iseq: IseqPtr, insn_idx: u32) -> u32 {
unsafe { rb_iseq_opcode_at_pc(iseq, pc) as u32 }
}
/// Return a poison value to be set above the stack top to verify leafness.
#[cfg(not(test))]
pub fn vm_stack_canary() -> u64 {
unsafe { rb_vm_stack_canary() }.as_u64()
}
/// Avoid linking the C function in `cargo test`
#[cfg(test)]
pub fn vm_stack_canary() -> u64 {
0
}
/// Opaque execution-context type from vm_core.h
#[repr(C)]
pub struct rb_execution_context_struct {