Skip to content
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

YJIT: add support for calling bmethods #6489

Merged
merged 2 commits into from Oct 5, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
28 changes: 28 additions & 0 deletions bootstraptest/test_yjit.rb
Expand Up @@ -3309,3 +3309,31 @@ def respond_to_missing?(*) = true
end
foo(Foo.new)
}

# bmethod
assert_equal '[1, 2, 3]', %q{
one = 1
define_method(:foo) do
one
end

3.times.map { |i| foo + i }
}

# return inside bmethod
assert_equal 'ok', %q{
define_method(:foo) do
1.tap { return :ok }
end

foo
}

# bmethod optional and keywords
assert_equal '[[1, nil, 2]]', %q{
define_method(:opt_and_kwargs) do |a = {}, b: nil, c: nil|
[a, b, c]
end

5.times.map { opt_and_kwargs(1, c: 2) }.uniq
}
15 changes: 15 additions & 0 deletions yjit.c
Expand Up @@ -425,6 +425,14 @@ rb_RSTRING_PTR(VALUE str)
return RSTRING_PTR(str);
}

rb_proc_t *
rb_yjit_get_proc_ptr(VALUE procv)
{
rb_proc_t *proc;
GetProcPtr(procv, proc);
return proc;
}

// This is defined only as a named struct inside rb_iseq_constant_body.
// By giving it a separate typedef, we make it nameable by rust-bindgen.
// Bindgen's temp/anon name isn't guaranteed stable.
Expand Down Expand Up @@ -549,6 +557,13 @@ rb_get_def_iseq_ptr(rb_method_definition_t *def)
return def_iseq_ptr(def);
}

VALUE
rb_get_def_bmethod_proc(rb_method_definition_t *def)
{
RUBY_ASSERT(def->type == VM_METHOD_TYPE_BMETHOD);
return def->body.bmethod.proc;
}

const rb_iseq_t *
rb_get_iseq_body_local_iseq(const rb_iseq_t *iseq)
{
Expand Down
2 changes: 2 additions & 0 deletions yjit/bindgen/src/main.rs
Expand Up @@ -273,6 +273,7 @@ fn main() {
.allowlist_function("rb_RSTRING_PTR")
.allowlist_function("rb_RSTRING_LEN")
.allowlist_function("rb_ENCODING_GET")
.allowlist_function("rb_yjit_get_proc_ptr")
.allowlist_function("rb_yjit_exit_locations_dict")
.allowlist_function("rb_yjit_icache_invalidate")

Expand Down Expand Up @@ -332,6 +333,7 @@ fn main() {
.allowlist_function("rb_get_mct_argc")
.allowlist_function("rb_get_mct_func")
.allowlist_function("rb_get_def_iseq_ptr")
.allowlist_function("rb_get_def_bmethod_proc")
.allowlist_function("rb_iseq_encoded_size")
.allowlist_function("rb_get_iseq_body_local_iseq")
.allowlist_function("rb_get_iseq_body_iseq_encoded")
Expand Down
94 changes: 74 additions & 20 deletions yjit/src/codegen.rs
Expand Up @@ -4035,6 +4035,7 @@ struct ControlFrame {
pc: Option<u64>,
frame_type: u32,
block_handler: BlockHandler,
prev_ep: Option<*const VALUE>,
cme: *const rb_callable_method_entry_t,
local_size: i32
}
Expand All @@ -4051,14 +4052,14 @@ struct ControlFrame {
// * Provided sp should point to the new frame's sp, immediately following locals and the environment
// * At entry, CFP points to the caller (not callee) frame
// * At exit, ec->cfp is updated to the pushed CFP
// * CFP and SP registers are updated only if switch_in_jit is set
// * CFP and SP registers are updated only if set_sp_cfp is set
// * Stack overflow is not checked (should be done by the caller)
// * Interrupts are not checked (should be done by the caller)
fn gen_push_frame(
_jit: &mut JITState,
_ctx: &mut Context,
asm: &mut Assembler,
set_pc_cfp: bool, // if true CFP and SP will be switched to the callee
set_sp_cfp: bool, // if true CFP and SP will be switched to the callee
frame: ControlFrame,
) {
assert!(frame.local_size >= 0);
Expand All @@ -4076,26 +4077,32 @@ fn gen_push_frame(
}
}

asm.comment("push cme, block handler, frame type");
asm.comment("push cme, specval, frame type");

// Write method entry at sp[-3]
// sp[-3] = me;
// Use compile time cme. It's assumed to be valid because we are notified when
// any cme we depend on become outdated. See yjit_method_lookup_change().
asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -3), VALUE::from(frame.cme).into());

// Write block handler at sp[-2]
// sp[-2] = block_handler;
let block_handler: Opnd = match frame.block_handler {
BlockHandler::None => {
// Write special value at sp[-2]. It's either a block handler or a pointer to
// the outer environment depending on the frame type.
// sp[-2] = specval;
let specval: Opnd = match (frame.prev_ep, frame.block_handler) {
(None, BlockHandler::None) => {
VM_BLOCK_HANDLER_NONE.into()
},
BlockHandler::CurrentFrame => {
}
(None, BlockHandler::CurrentFrame) => {
let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
asm.or(cfp_self, Opnd::Imm(1))
},
}
(Some(prev_ep), BlockHandler::None) => {
let tagged_prev_ep = (prev_ep as usize) | 1;
VALUE(tagged_prev_ep).into()
}
(_, _) => panic!("specval can only be one of prev_ep or block_handler")
};
asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), block_handler);
asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);

// Write env flags at sp[-1]
// sp[-1] = frame_type;
Expand Down Expand Up @@ -4134,7 +4141,7 @@ fn gen_push_frame(
asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);
asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());

if set_pc_cfp {
if set_sp_cfp {
// Saving SP before calculating ep avoids a dependency on a register
// However this must be done after referencing frame.recv, which may be SP-relative
asm.mov(SP, sp);
Expand All @@ -4144,7 +4151,7 @@ fn gen_push_frame(

asm.comment("switch to new CFP");
let new_cfp = asm.lea(cfp_opnd(0));
if set_pc_cfp {
if set_sp_cfp {
asm.mov(CFP, new_cfp);
asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
} else {
Expand Down Expand Up @@ -4274,6 +4281,7 @@ fn gen_send_cfunc(
cme,
recv,
sp,
prev_ep: None,
pc: Some(0),
iseq: None,
local_size: 0,
Expand Down Expand Up @@ -4443,17 +4451,59 @@ fn push_splat_args(required_args: i32, ctx: &mut Context, asm: &mut Assembler, o
}
}

fn gen_send_bmethod(
jit: &mut JITState,
ctx: &mut Context,
asm: &mut Assembler,
ocb: &mut OutlinedCb,
ci: *const rb_callinfo,
cme: *const rb_callable_method_entry_t,
block: Option<IseqPtr>,
argc: i32,
) -> CodegenStatus {
let procv = unsafe { rb_get_def_bmethod_proc((*cme).def) };

let proc = unsafe { rb_yjit_get_proc_ptr(procv) };
let proc_block = unsafe { &(*proc).block };

if proc_block.type_ != block_type_iseq {
return CantCompile;
}

let capture = unsafe { proc_block.as_.captured.as_ref() };
let iseq = unsafe { *capture.code.iseq.as_ref() };

// Optimize for single ractor mode and avoid runtime check for
// "defined with an un-shareable Proc in a different Ractor"
if !assume_single_ractor_mode(jit, ocb) {
gen_counter_incr!(asm, send_bmethod_ractor);
return CantCompile;
}

// Passing a block to a block needs logic different from passing
// a block to a method and sometimes requires allocation. Bail for now.
if block.is_some() {
gen_counter_incr!(asm, send_bmethod_block_arg);
return CantCompile;
}

let frame_type = VM_FRAME_MAGIC_BLOCK | VM_FRAME_FLAG_BMETHOD | VM_FRAME_FLAG_LAMBDA;
gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, Some(capture.ep), cme, block, argc)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

😍

}

fn gen_send_iseq(
jit: &mut JITState,
ctx: &mut Context,
asm: &mut Assembler,
ocb: &mut OutlinedCb,
iseq: *const rb_iseq_t,
ci: *const rb_callinfo,
frame_type: u32,
prev_ep: Option<*const VALUE>,
cme: *const rb_callable_method_entry_t,
block: Option<IseqPtr>,
argc: i32,
) -> CodegenStatus {
let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
let mut argc = argc;

let flags = unsafe { vm_ci_flag(ci) };
Expand Down Expand Up @@ -4893,15 +4943,14 @@ fn gen_send_iseq(
BlockHandler::None
};

let frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;

// Setup the new frame
gen_push_frame(jit, ctx, asm, true, ControlFrame {
frame_type,
block_handler: frame_block_handler,
cme,
recv,
sp: callee_sp,
prev_ep,
iseq: Some(iseq),
pc: None, // We are calling into jitted code, which will set the PC as necessary
local_size: num_locals
Expand Down Expand Up @@ -5173,7 +5222,9 @@ fn gen_send_general(

match def_type {
VM_METHOD_TYPE_ISEQ => {
return gen_send_iseq(jit, ctx, asm, ocb, ci, cme, block, argc);
let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
let frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;
return gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, None, cme, block, argc);
}
VM_METHOD_TYPE_CFUNC => {
return gen_send_cfunc(
Expand Down Expand Up @@ -5243,8 +5294,7 @@ fn gen_send_general(
}
// Block method, e.g. define_method(:foo) { :my_block }
VM_METHOD_TYPE_BMETHOD => {
gen_counter_incr!(asm, send_bmethod);
return CantCompile;
return gen_send_bmethod(jit, ctx, asm, ocb, ci, cme, block, argc);
}
VM_METHOD_TYPE_ZSUPER => {
gen_counter_incr!(asm, send_zsuper_method);
Expand Down Expand Up @@ -5481,7 +5531,11 @@ fn gen_invokesuper(
ctx.clear_local_types();

match cme_def_type {
VM_METHOD_TYPE_ISEQ => gen_send_iseq(jit, ctx, asm, ocb, ci, cme, block, argc),
VM_METHOD_TYPE_ISEQ => {
let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
let frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;
gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, None, cme, block, argc)
}
VM_METHOD_TYPE_CFUNC => {
gen_send_cfunc(jit, ctx, asm, ocb, ci, cme, block, argc, ptr::null())
}
Expand Down