ZJIT: Optimize Send

Author: st0012

insns.def

@@ -846,6 +846,7 @@ send
 (CALL_DATA cd, ISEQ blockiseq)
 (...)
 (VALUE val)
+// attr bool zjit_profile = true;
 // attr rb_snum_t sp_inc = sp_inc_of_sendish(cd->ci);
 // attr rb_snum_t comptime_sp_inc = sp_inc_of_sendish(ci);
 {

zjit/src/codegen.rs

@@ -20,7 +20,7 @@ use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::{compile_time_
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, asm_ccall, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, NATIVE_BASE_PTR, SP};
 use crate::hir::{iseq_to_hir, BlockId, BranchEdge, Invariant, RangeType, SideExitReason::{self, *}, SpecialBackrefSymbol, SpecialObjectType};
-use crate::hir::{Const, FrameState, Function, Insn, InsnId, SendFallbackReason};
+use crate::hir::{Const, FrameState, Function, Insn, InsnId, MethodType, SendFallbackReason};
 use crate::hir_type::{types, Type};
 use crate::options::get_option;
 use crate::cast::IntoUsize;
@@ -378,6 +378,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::SendWithoutBlockDirectTooManyArgs),
         Insn::SendWithoutBlockDirect { cme, iseq, recv, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), &function.frame_state(*state)),
+        // Give up SendDirect for 6+ args since asm.ccall() doesn't support it.
+        Insn::SendDirect { cd, blockiseq, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
+            gen_send(jit, asm, *cd, *blockiseq, &function.frame_state(*state), SendFallbackReason::SendDirectTooManyArgs),
+        Insn::SendDirect { cd, cme, blockiseq, recv, args, state, .. } => gen_send_direct(jit, asm, *cd, *cme, *blockiseq, opnd!(recv), opnds!(args), &function.frame_state(*state)),
         &Insn::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::InvokeBlock { cd, state, reason, .. } => gen_invokeblock(jit, asm, cd, &function.frame_state(state), reason),
         // Ensure we have enough room fit ec, self, and arguments
@@ -682,6 +686,8 @@ fn gen_ccall_with_frame(jit: &mut JITState, asm: &mut Assembler, cfunc: *const u
         iseq: None,
         cme,
         frame_type: VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL,
+        block_handler: None,
+        block_iseq: None,
     });
 
     asm_comment!(asm, "switch to new SP register");
@@ -737,6 +743,8 @@ fn gen_ccall_variadic(
         iseq: None,
         cme,
         frame_type: VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL,
+        block_handler: None,
+        block_iseq: None,
     });
 
     asm_comment!(asm, "switch to new SP register");
@@ -1129,6 +1137,8 @@ fn gen_send_without_block_direct(
         iseq: Some(iseq),
         cme,
         frame_type: VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL,
+        block_handler: None,
+        block_iseq: None,
     });
 
     asm_comment!(asm, "switch to new SP register");
@@ -1166,6 +1176,90 @@ fn gen_send_without_block_direct(
     ret
 }
 
+/// Compile a direct jump to an ISEQ call with block
+fn gen_send_direct(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cd: *const rb_call_data,
+    cme: *const rb_callable_method_entry_t,
+    blockiseq: IseqPtr,
+    recv: Opnd,
+    args: Vec<Opnd>,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_incr_counter(asm, Counter::iseq_optimized_send_count);
+
+    let cfunc = unsafe { get_cme_def_body_cfunc(cme) };
+    let argc = args.len() as i32;
+
+    let ci = unsafe { get_call_data_ci(cd) };
+    let flags = unsafe { vm_ci_flag(ci) };
+    // When seeing &block argument, fall back to dynamic dispatch for now
+    // TODO: Support block forwarding
+    if flags & VM_CALL_ARGS_BLOCKARG != 0 {
+        return gen_send(
+            jit,
+            asm,
+            cd,
+            blockiseq,
+            state,
+            SendFallbackReason::SendNotOptimizedMethodType(MethodType::Cfunc),
+        );
+    }
+
+    let stack_growth = state.stack_size();
+    gen_stack_overflow_check(jit, asm, state, stack_growth);
+
+    // Save cfp->pc and cfp->sp for the caller frame
+    gen_prepare_call_with_gc(asm, state, false);
+    // Special SP math. Can't use gen_prepare_non_leaf_call
+    gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver
+
+    gen_spill_locals(jit, asm, state);
+    gen_spill_stack(jit, asm, state);
+
+    let block_handler = if !blockiseq.is_null() {
+        asm.store(
+            Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE),
+            VALUE::from(blockiseq).into(),
+        );
+        let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
+        Some(asm.or(cfp_self, Opnd::Imm(1)))
+    } else {
+        None
+    };
+
+    let frame_type = VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL;
+    gen_push_frame(asm, argc as usize, state, ControlFrame {
+        recv,
+        iseq: None,
+        cme,
+        frame_type,
+        block_handler,
+        block_iseq: None,
+    });
+
+    asm_comment!(asm, "set ec->cfp");
+    let new_cfp = asm.lea(Opnd::mem(64, CFP, -(RUBY_SIZEOF_CONTROL_FRAME as i32)));
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+
+    let mut c_args = vec![recv];
+    c_args.extend(args);
+
+    asm_comment!(asm, "call C function");
+    let cfunc_ptr = unsafe { get_mct_func(cfunc) }.cast();
+    let ret = asm.ccall(cfunc_ptr, c_args);
+
+    asm_comment!(asm, "restore ec->cfp");
+    let ec_cfp_opnd = Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP);
+    let caller_cfp = asm.add(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
+    asm.mov(CFP, caller_cfp);
+    asm.store(ec_cfp_opnd, CFP);
+
+    ret
+}
+
 /// Compile for invokeblock
 fn gen_invokeblock(
     jit: &mut JITState,
@@ -1709,6 +1803,8 @@ struct ControlFrame {
     iseq: Option<IseqPtr>,
     cme: *const rb_callable_method_entry_t,
     frame_type: u32,
+    block_handler: Option<Opnd>,
+    block_iseq: Option<IseqPtr>,
 }
 
 /// Compile an interpreter frame
@@ -1725,9 +1821,24 @@ fn gen_push_frame(asm: &mut Assembler, argc: usize, state: &FrameState, frame: C
     };
     let ep_offset = state.stack().len() as i32 + local_size - argc as i32 + VM_ENV_DATA_SIZE as i32 - 1;
     asm.store(Opnd::mem(64, SP, (ep_offset - 2) * SIZEOF_VALUE_I32), VALUE::from(frame.cme).into());
-    // ep[-1]: block_handler or prev EP
-    // block_handler is not supported for now
-    asm.store(Opnd::mem(64, SP, (ep_offset - 1) * SIZEOF_VALUE_I32), VM_BLOCK_HANDLER_NONE.into());
+
+    let block_handler_opnd = if let Some(handler) = frame.block_handler {
+        handler
+    } else if let Some(block_iseq) = frame.block_iseq {
+        asm.store(
+            Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE),
+            VALUE::from(block_iseq).into(),
+        );
+        let cfp_self_addr = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
+        asm.or(cfp_self_addr, Opnd::Imm(1))
+    } else {
+        VM_BLOCK_HANDLER_NONE.into()
+    };
+
+    asm.store(
+        Opnd::mem(64, SP, (ep_offset - 1) * SIZEOF_VALUE_I32),
+        block_handler_opnd,
+    );
     // ep[0]: ENV_FLAGS
     asm.store(Opnd::mem(64, SP, ep_offset * SIZEOF_VALUE_I32), frame.frame_type.into());