ZJIT: Avoid redundant SP save in codegen

Author: st0012

zjit/src/codegen.rs

@@ -789,8 +789,9 @@ fn gen_ccall_with_frame(
     let caller_stack_size = state.stack().len() - args_with_recv_len;
 
     // Can't use gen_prepare_non_leaf_call() because we need to adjust the SP
-    // to account for the receiver and arguments (and block arguments if any)
-    gen_prepare_call_with_gc(asm, state, false);
+    // to account for the receiver and arguments (and block arguments if any).
+    // Skip SP save in gen_prepare_call_with_gc since we call gen_save_sp right after.
+    gen_prepare_call_with_gc(asm, state, false, true);
     gen_save_sp(asm, caller_stack_size);
     gen_spill_stack(jit, asm, state);
     gen_spill_locals(jit, asm, state);
@@ -874,8 +875,9 @@ fn gen_ccall_variadic(
     let caller_stack_size = state.stack_size() - args_with_recv_len;
 
     // Can't use gen_prepare_non_leaf_call() because we need to adjust the SP
-    // to account for the receiver and arguments (like gen_ccall_with_frame does)
-    gen_prepare_call_with_gc(asm, state, false);
+    // to account for the receiver and arguments (like gen_ccall_with_frame does).
+    // Skip SP save in gen_prepare_call_with_gc since we call gen_save_sp right after.
+    gen_prepare_call_with_gc(asm, state, false, true);
     gen_save_sp(asm, caller_stack_size);
     gen_spill_stack(jit, asm, state);
     gen_spill_locals(jit, asm, state);
@@ -1303,8 +1305,9 @@ fn gen_send_without_block_direct(
     let stack_growth = state.stack_size() + local_size + unsafe { get_iseq_body_stack_max(iseq) }.to_usize();
     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);
+    // Save cfp->pc and cfp->sp for the caller frame.
+    // Skip SP save in gen_prepare_call_with_gc since we call gen_save_sp right after.
+    gen_prepare_call_with_gc(asm, state, false, true);
     // Special SP math. Can't use gen_prepare_non_leaf_call
     gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver
 
@@ -2016,15 +2019,21 @@ fn gen_incr_send_fallback_counter(asm: &mut Assembler, reason: SendFallbackReaso
 /// because the backend spills all live registers onto the C stack on CCall.
 /// However, to avoid marking uninitialized stack slots, this also updates SP,
 /// which may have cfp->sp for a past frame or a past non-leaf call.
-fn gen_prepare_call_with_gc(asm: &mut Assembler, state: &FrameState, leaf: bool) {
+///
+/// Set `skip_sp_save` to true if the caller will call gen_save_sp() immediately
+/// after with a different stack size (e.g., gen_ccall_with_frame adjusts SP to
+/// exclude receiver and arguments).
+fn gen_prepare_call_with_gc(asm: &mut Assembler, state: &FrameState, leaf: bool, skip_sp_save: bool) {
     let opcode: usize = state.get_opcode().try_into().unwrap();
     let next_pc: *const VALUE = unsafe { state.pc.offset(insn_len(opcode) as isize) };
 
     gen_incr_counter(asm, Counter::vm_write_pc_count);
     asm_comment!(asm, "save PC to CFP");
     asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC), Opnd::const_ptr(next_pc));
 
-    gen_save_sp(asm, state.stack_size());
+    if !skip_sp_save {
+        gen_save_sp(asm, state.stack_size());
+    }
     if leaf {
         asm.expect_leaf_ccall(state.stack_size());
     }
@@ -2044,7 +2053,7 @@ fn gen_prepare_leaf_call_with_gc(asm: &mut Assembler, state: &FrameState) {
     // We use state.without_stack() to pass stack_size=0 to gen_save_sp() because we don't write
     // VM stack slots on leaf calls, which leaves those stack slots uninitialized. ZJIT keeps
     // live objects on the C stack, so they are protected from GC properly.
-    gen_prepare_call_with_gc(asm, &state.without_stack(), true);
+    gen_prepare_call_with_gc(asm, &state.without_stack(), true, false);
 }
 
 /// Save the current SP on the CFP
@@ -2087,7 +2096,7 @@ fn gen_prepare_non_leaf_call(jit: &JITState, asm: &mut Assembler, state: &FrameS
     // TODO: Lazily materialize caller frames when needed
     // Save PC for backtraces and allocation tracing
     // and SP to avoid marking uninitialized stack slots
-    gen_prepare_call_with_gc(asm, state, false);
+    gen_prepare_call_with_gc(asm, state, false, false);
 
     // Spill the virtual stack in case it raises an exception
     // and the interpreter uses the stack for handling the exception