Thanks to visit codestin.com
Credit goes to github.com

Skip to content

ZJIT: Split Send into Lookup+Call #13400

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions jit.c
Original file line number Diff line number Diff line change
Expand Up @@ -421,3 +421,19 @@ rb_assert_cme_handle(VALUE handle)
RUBY_ASSERT_ALWAYS(!rb_objspace_garbage_object_p(handle));
RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_ment));
}

extern
VALUE
rb_vm_call0(rb_execution_context_t *ec, VALUE recv, ID id, int argc, const VALUE *argv, const rb_callable_method_entry_t *cme, int kw_splat);

VALUE
rb_zjit_vm_call0_no_splat(rb_execution_context_t *ec, VALUE recv, ID id, int argc, const VALUE *argv, const rb_callable_method_entry_t *cme) {
const char *cme_name = rb_id2name(cme->def->original_id);
fprintf(stderr, "call0 ec %p recv %p id %p argc %d, argv %p, cme %p name %s\n",
ec, (void*)recv, (void*)id, argc, argv, cme, cme_name);
return rb_vm_call0(ec, recv, id, argc, argv, cme, /*kw_splat=*/0);
}

extern
VALUE
rb_vm_send(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd, ISEQ blockiseq);
2 changes: 2 additions & 0 deletions zjit/bindgen/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,8 @@ fn main() {
.allowlist_function("rb_optimized_call")
.allowlist_function("rb_zjit_icache_invalidate")
.allowlist_function("rb_zjit_print_exception")
.allowlist_function("rb_zjit_vm_call0_no_splat")
.allowlist_function("rb_vm_send")
.allowlist_type("robject_offsets")
.allowlist_type("rstring_offsets")

Expand Down
5 changes: 5 additions & 0 deletions zjit/src/backend/arm64/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ pub const _C_ARG_OPNDS: [Opnd; 6] = [
// C return value register on this platform
pub const C_RET_REG: Reg = X0_REG;
pub const _C_RET_OPND: Opnd = Opnd::Reg(X0_REG);
pub const _NATIVE_STACK_PTR: Opnd = Opnd::Reg(XZR_REG);

// These constants define the way we work with Arm64's stack pointer. The stack
// pointer always needs to be aligned to a 16-byte boundary.
Expand Down Expand Up @@ -679,6 +680,10 @@ impl Assembler
*opnd = opnd0;
asm.push_insn(insn);
},
Insn::CPush(opnd) => {
let opnd = split_load_operand(asm, *opnd);
asm.push_insn(Insn::CPush(opnd));
}
Insn::Store { dest, src } => {
// The value being stored must be in a register, so if it's
// not already one we'll load it first.
Expand Down
3 changes: 2 additions & 1 deletion zjit/src/backend/lir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ pub const SP: Opnd = _SP;

pub const C_ARG_OPNDS: [Opnd; 6] = _C_ARG_OPNDS;
pub const C_RET_OPND: Opnd = _C_RET_OPND;
pub const NATIVE_STACK_PTR: Opnd = _NATIVE_STACK_PTR;
pub use crate::backend::current::{Reg, C_RET_REG};

// Memory operand base
Expand Down Expand Up @@ -2225,7 +2226,7 @@ impl Assembler {
/// when not dumping disassembly.
macro_rules! asm_comment {
($asm:expr, $($fmt:tt)*) => {
if $crate::options::get_option!(dump_disasm) {
if $crate::options::get_option!(dump_disasm) || $crate::options::get_option!(dump_lir) {
$asm.push_insn(crate::backend::lir::Insn::Comment(format!($($fmt)*)));
}
};
Expand Down
1 change: 1 addition & 0 deletions zjit/src/backend/x86_64/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ pub const _C_ARG_OPNDS: [Opnd; 6] = [
// C return value register on this platform
pub const C_RET_REG: Reg = RAX_REG;
pub const _C_RET_OPND: Opnd = Opnd::Reg(RAX_REG);
pub const _NATIVE_STACK_PTR: Opnd = Opnd::Reg(RSP_REG);

impl CodeBlock {
// The number of bytes that are generated by jmp_ptr
Expand Down
138 changes: 116 additions & 22 deletions zjit/src/codegen.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@ use crate::profile::get_or_create_iseq_payload;
use crate::state::ZJITState;
use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX};
use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP, NATIVE_STACK_PTR};
use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, RangeType, SELF_PARAM_IDX};
use crate::hir::{Const, FrameState, Function, Insn, InsnId};
use crate::hir_type::{types::Fixnum, Type};
use crate::options::get_option;
Expand Down Expand Up @@ -257,8 +257,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
Insn::Jump(branch) => return gen_jump(jit, asm, branch),
Insn::IfTrue { val, target } => return gen_if_true(jit, asm, opnd!(val), target),
Insn::IfFalse { val, target } => return gen_if_false(jit, asm, opnd!(val), target),
Insn::SendWithoutBlock { call_info, cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state), self_val, args)?,
Insn::SendWithoutBlockDirect { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
Insn::LookupMethod { self_val, method_id, state } => gen_lookup_method(jit, asm, opnd!(self_val), *method_id, &function.frame_state(*state))?,
Insn::CallMethod { callable, cd, self_val, args, state } => gen_call_method(jit, asm, opnd!(callable), *cd, *self_val, args, &function.frame_state(*state))?,
Insn::CallCFunc { cfunc, self_val, args, .. } => gen_call_cfunc(jit, asm, *cfunc, opnd!(self_val), args)?,
Insn::CallIseq { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Expand Down Expand Up @@ -455,38 +457,126 @@ fn gen_if_false(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, branch:
}

/// Compile a dynamic dispatch without block
fn gen_send_without_block(
fn gen_lookup_method(
jit: &mut JITState,
asm: &mut Assembler,
call_info: &CallInfo,
cd: *const rb_call_data,
recv: Opnd,
method_id: ID,
state: &FrameState,
self_val: &InsnId,
args: &Vec<InsnId>,
) -> Option<lir::Opnd> {
// Spill the receiver and the arguments onto the stack. They need to be marked by GC and may be caller-saved registers.
asm_comment!(asm, "get the class of the receiver with rb_obj_class");
let class = asm.ccall(rb_obj_class as *const u8, vec![recv]);
// TODO(max): Figure out if we need to do anything here to save state to CFP
let method_opnd = Opnd::UImm(method_id.0.into());
asm_comment!(asm, "call rb_callable_method_entry");
let result = asm.ccall(
rb_callable_method_entry as *const u8,
vec![class, method_opnd],
);
// rb_callable_method_entry may be NULL, and we don't want to handle method_missing shenanigans
// in JIT code.
asm.test(result, result);
asm.jz(side_exit(jit, state)?);
Some(result)
}

fn gen_call_method(
jit: &mut JITState,
asm: &mut Assembler,
callable: Opnd,
cd: CallDataPtr,
recv: InsnId,
args: &Vec<InsnId>,
state: &FrameState,
) -> Option<Opnd> {
// Spill the receiver and the arguments onto the stack.
// They need to be on the interpreter stack to let the interpreter access them.
// TODO: Avoid spilling operands that have been spilled before.
for (idx, &insn_id) in [*self_val].iter().chain(args.iter()).enumerate() {
asm_comment!(asm, "spill receiver and arguments");
for (idx, &insn_id) in [recv].iter().chain(args.iter()).enumerate() {
// Currently, we don't move the SP register. So it's equal to the base pointer.
let stack_opnd = Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32);
asm.mov(stack_opnd, jit.get_opnd(insn_id)?);
}

// Save PC and SP
// Don't push recv; it is passed in separately.
asm_comment!(asm, "make stack-allocated array of {} args", args.len());
let sp_adjust = if args.len() % 2 == 0 { args.len() } else { args.len() + 1 };
let new_sp = asm.sub(NATIVE_STACK_PTR, (sp_adjust*SIZEOF_VALUE).into());
asm.mov(NATIVE_STACK_PTR, new_sp);
for (idx, &arg) in args.iter().rev().enumerate() {
asm.mov(Opnd::mem(VALUE_BITS, NATIVE_STACK_PTR, ((idx)*SIZEOF_VALUE).try_into().unwrap()), jit.get_opnd(arg)?);
}
// Save PC for GC
gen_save_pc(asm, state);
gen_save_sp(asm, 1 + args.len()); // +1 for receiver
// Call rb_zjit_vm_call0_no_splat, which will push a frame
// TODO(max): Figure out if we need to manually handle stack alignment and how to do it
let call_info = unsafe { rb_get_call_data_ci(cd) };
let method_id = unsafe { rb_vm_ci_mid(call_info) };
asm_comment!(asm, "get stack pointer");
let sp = asm.lea(Opnd::mem(VALUE_BITS, NATIVE_STACK_PTR, VALUE_BITS.into()));
asm_comment!(asm, "call rb_zjit_vm_call0_no_splat");
let recv = jit.get_opnd(recv)?;
let result = asm.ccall(
rb_zjit_vm_call0_no_splat as *const u8,
vec![EC, recv, Opnd::UImm(method_id.0), Opnd::UImm(args.len().try_into().unwrap()), sp, callable],
);
// Pop all the args off the stack
asm_comment!(asm, "clear stack-allocated array of {} args", args.len());
let new_sp = asm.add(NATIVE_STACK_PTR, (sp_adjust*SIZEOF_VALUE).into());
asm.mov(NATIVE_STACK_PTR, new_sp);
Some(result)
}

/// Compile an interpreter frame
fn gen_push_frame(asm: &mut Assembler, recv: Opnd) {
// Write to a callee CFP
fn cfp_opnd(offset: i32) -> Opnd {
Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
}

asm_comment!(asm, "call #{} with dynamic dispatch", call_info.method_name);
unsafe extern "C" {
fn rb_vm_opt_send_without_block(ec: EcPtr, cfp: CfpPtr, cd: VALUE) -> VALUE;
asm_comment!(asm, "push callee control frame");
asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), recv);
// TODO: Write more fields as needed
}

fn gen_call_cfunc(
jit: &mut JITState,
asm: &mut Assembler,
cfunc: CFuncPtr,
recv: Opnd,
args: &Vec<InsnId>,
) -> Option<lir::Opnd> {
let cfunc_argc = unsafe { get_mct_argc(cfunc) };
// NB: The presence of self is assumed (no need for +1).
if args.len() != cfunc_argc as usize {
// TODO(max): We should check this at compile-time. If we have an arity mismatch at this
// point, we should side-exit (we're definitely going to raise) and if we don't, we should
// not check anything.
todo!("Arity mismatch");
}
let ret = asm.ccall(
rb_vm_opt_send_without_block as *const u8,
vec![EC, CFP, (cd as usize).into()],
);

// Set up the new frame
gen_push_frame(asm, recv);

asm_comment!(asm, "switch to new CFP");
let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
asm.mov(CFP, new_cfp);
asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);

// Set up arguments
let mut c_args: Vec<Opnd> = vec![recv];
for &arg in args.iter() {
c_args.push(jit.get_opnd(arg)?);
}

// Make a method call. The target address will be rewritten once compiled.
let cfun = unsafe { get_mct_func(cfunc) }.cast();
// TODO(max): Add a PatchPoint here that can side-exit the function if the callee messed with
// the frame's locals

let ret = asm.ccall(cfun, c_args);
gen_pop_frame(asm);
Some(ret)
}

Expand Down Expand Up @@ -590,14 +680,18 @@ fn gen_new_range(
new_range
}

/// Compile code that exits from JIT code with a return value
fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
fn gen_pop_frame(asm: &mut Assembler) {
// Pop the current frame (ec->cfp++)
// Note: the return PC is already in the previous CFP
asm_comment!(asm, "pop stack frame");
let incr_cfp = asm.add(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
asm.mov(CFP, incr_cfp);
asm.mov(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
}

/// Compile code that exits from JIT code with a return value
fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
gen_pop_frame(asm);

asm.frame_teardown();

Expand Down
24 changes: 15 additions & 9 deletions zjit/src/cruby.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,9 @@ pub use autogened::*;
#[cfg_attr(test, allow(unused))] // We don't link against C code when testing
unsafe extern "C" {
pub fn rb_check_overloaded_cme(
me: *const rb_callable_method_entry_t,
me: CmePtr,
ci: *const rb_callinfo,
) -> *const rb_callable_method_entry_t;
) -> CmePtr;

// Floats within range will be encoded without creating objects in the heap.
// (Range is 0x3000000000000001 to 0x4fffffffffffffff (1.7272337110188893E-77 to 2.3158417847463237E+77).
Expand All @@ -144,8 +144,8 @@ unsafe extern "C" {
pub fn rb_vm_set_ivar_id(obj: VALUE, idx: u32, val: VALUE) -> VALUE;
pub fn rb_vm_setinstancevariable(iseq: IseqPtr, obj: VALUE, id: ID, val: VALUE, ic: IVC);
pub fn rb_aliased_callable_method_entry(
me: *const rb_callable_method_entry_t,
) -> *const rb_callable_method_entry_t;
me: CmePtr,
) -> CmePtr;
pub fn rb_vm_getclassvariable(iseq: IseqPtr, cfp: CfpPtr, id: ID, ic: ICVARC) -> VALUE;
pub fn rb_vm_setclassvariable(
iseq: IseqPtr,
Expand All @@ -156,7 +156,7 @@ unsafe extern "C" {
) -> VALUE;
pub fn rb_vm_ic_hit_p(ic: IC, reg_ep: *const VALUE) -> bool;
pub fn rb_vm_stack_canary() -> VALUE;
pub fn rb_vm_push_cfunc_frame(cme: *const rb_callable_method_entry_t, recv_idx: c_int);
pub fn rb_vm_push_cfunc_frame(cme: CmePtr, recv_idx: c_int);
}

// Renames
Expand Down Expand Up @@ -259,6 +259,12 @@ pub struct ID(pub ::std::os::raw::c_ulong);
/// Pointer to an ISEQ
pub type IseqPtr = *const rb_iseq_t;

pub type CallDataPtr = *const rb_call_data;

pub type CmePtr = *const rb_callable_method_entry_t;

pub type CFuncPtr = *const rb_method_cfunc_t;

// Given an ISEQ pointer, convert PC to insn_idx
pub fn iseq_pc_to_insn_idx(iseq: IseqPtr, pc: *mut VALUE) -> Option<u16> {
let pc_zero = unsafe { rb_iseq_pc_at_idx(iseq, 0) };
Expand Down Expand Up @@ -579,8 +585,8 @@ impl VALUE {
}

/// Assert that `self` is a method entry in debug builds
pub fn as_cme(self) -> *const rb_callable_method_entry_t {
let ptr: *const rb_callable_method_entry_t = self.as_ptr();
pub fn as_cme(self) -> CmePtr {
let ptr: CmePtr = self.as_ptr();

#[cfg(debug_assertions)]
if !ptr.is_null() {
Expand Down Expand Up @@ -611,9 +617,9 @@ impl From<IseqPtr> for VALUE {
}
}

impl From<*const rb_callable_method_entry_t> for VALUE {
impl From<CmePtr> for VALUE {
/// For `.into()` convenience
fn from(cme: *const rb_callable_method_entry_t) -> Self {
fn from(cme: CmePtr) -> Self {
VALUE(cme as usize)
}
}
Expand Down
16 changes: 16 additions & 0 deletions zjit/src/cruby_bindings.inc.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 5 additions & 3 deletions zjit/src/cruby_methods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ impl Annotations {
if VM_METHOD_TYPE_CFUNC != get_cme_def_type(method) {
return None;
}
get_mct_func(get_cme_def_body_cfunc(method.cast()))
rb_get_mct_func(rb_get_cme_def_body_cfunc(method.cast()))
};
self.cfuncs.get(&fn_ptr).copied()
}
Expand Down Expand Up @@ -65,11 +65,12 @@ pub fn init() -> Annotations {
let cfuncs = &mut HashMap::new();

macro_rules! annotate {
($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),*) => {
#[allow(unused_mut)]
let mut props = FnProperties { no_gc: false, leaf: false, elidable: false, return_type: $return_type };
$(
props.$properties = true;
)+
)*
annotate_c_method(cfuncs, unsafe { $module }, $method_name, props);
}
}
Expand All @@ -80,6 +81,7 @@ pub fn init() -> Annotations {
annotate!(rb_cModule, "===", types::BoolExact, no_gc, leaf);
annotate!(rb_cArray, "length", types::Fixnum, no_gc, leaf, elidable);
annotate!(rb_cArray, "size", types::Fixnum, no_gc, leaf, elidable);
annotate!(rb_cInteger, "+", types::IntegerExact,);

Annotations {
cfuncs: std::mem::take(cfuncs)
Expand Down
Loading