ZJIT: Skip redundant cfp->sp stores

gen_save_sp() stored cfp->sp before every C call and JIT-to-JIT call,
even when an earlier call in the same basic block had already written
the same value and nothing could have changed it in between: callees
write only their own cfp->sp, side exits and exception unwinding never
return to the same code, and rb_zjit_materialize_frames() only reads
cfp->sp.

Track the stack size of the last cfp->sp store emitted in the current
basic block in the Assembler and skip stores that would write the same
value again. The cache is cleared at block boundaries (but preserved
across split_block_jump() fall-throughs, which have a single
predecessor), when PushInlineFrame/PopInlineFrame change which frame
the SP/CFP registers point at, and after gen_prepare_fallback_call(),
whose VM helpers execute an instruction on the current frame and move
cfp->sp (e.g. vm_sendish() popping the receiver and arguments), which
the test_regression_cfp_sp_set_correctly_before_leaf_gc_call test
caught.

Author: k0kubun

zjit/src/backend/lir.rs

@@ -2020,6 +2020,13 @@ pub struct Assembler {
     /// allocator spill slot each one is saved to right after FrameSetup. Every exit path
     /// restores them from these slots. Filled in by [`Self::plan_callee_saved_saves`].
     callee_saved_saves: Vec<(Reg, usize)>,
+
+    /// The stack size whose corresponding value is known to be already stored
+    /// in the current frame's cfp->sp by an earlier store in the current basic
+    /// block. Used by codegen to skip redundant cfp->sp stores. Cleared at
+    /// block boundaries and whenever codegen changes which frame the SP/CFP
+    /// registers point at.
+    saved_sp_stack_size: Option<usize>,
 }
 
 impl Assembler
@@ -2038,6 +2045,7 @@ impl Assembler
             stack_map: None,
             allow_callee_saved: false,
             callee_saved_saves: Vec::default(),
+            saved_sp_stack_size: None,
         }
     }
 
@@ -2120,6 +2128,37 @@ impl Assembler
 
     pub fn set_current_block(&mut self, block_id: BlockId) {
         self.current_block_id = block_id;
+        // Conservatively assume the new block can be entered with a different
+        // frame memory state, e.g. join blocks with multiple predecessors.
+        self.clear_frame_write_cache();
+    }
+
+    /// Like set_current_block(), but preserves the frame write cache. Only
+    /// safe when the new block is a fresh fall-through block whose single
+    /// predecessor is the current block, like the ones split_block_jump()
+    /// creates: control falls through linearly, so frame writes recorded so
+    /// far still describe the memory state.
+    pub fn set_current_block_fall_through(&mut self, block_id: BlockId) {
+        let saved_sp_stack_size = self.saved_sp_stack_size;
+        self.set_current_block(block_id);
+        self.saved_sp_stack_size = saved_sp_stack_size;
+    }
+
+    /// Record that the current frame's cfp->sp is being stored for `stack_size`.
+    /// Returns true if an earlier store in the current block already wrote the
+    /// same value, in which case the store can be skipped.
+    pub fn note_sp_save(&mut self, stack_size: usize) -> bool {
+        if self.saved_sp_stack_size == Some(stack_size) {
+            return true;
+        }
+        self.saved_sp_stack_size = Some(stack_size);
+        false
+    }
+
+    /// Forget all frame writes recorded by note_sp_save().
+    /// Call this when the SP/CFP registers change which frame they point at.
+    pub fn clear_frame_write_cache(&mut self) {
+        self.saved_sp_stack_size = None;
     }
 
     pub fn current_block(&mut self) -> &mut BasicBlock {

zjit/src/codegen.rs

@@ -141,7 +141,7 @@ impl Assembler {
         emit(self, target);
         self.jmp(Target::Block(Box::new(fall_through_edge)));
 
-        self.set_current_block(fall_through_target);
+        self.set_current_block_fall_through(fall_through_target);
 
         let label = jit.get_label(self, fall_through_target, hir_block_id);
         self.write_label(label);
@@ -3013,6 +3013,10 @@ fn gen_push_inline_frame(
     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 as i32), CFP);
+
+    // The SP/CFP registers now point at the callee frame, so frame writes
+    // recorded for the caller frame no longer describe the same locations.
+    asm.clear_frame_write_cache();
 }
 
 /// Pop the interpreter frame for an inlined callee, restoring the caller's SP and CFP.
@@ -3034,6 +3038,11 @@ fn gen_pop_inline_frame(
     asm_comment!(asm, "restore caller CFP after inline");
     asm.add_into(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
     asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP as i32), CFP);
+
+    // The SP/CFP registers now point at the caller frame again, so frame
+    // writes recorded for the callee frame no longer describe the same
+    // locations.
+    asm.clear_frame_write_cache();
 }
 
 /// Compile a direct call to an ISEQ method.
@@ -5366,6 +5375,14 @@ fn gen_prepare_leaf_call_with_gc(asm: &mut Assembler, state: &FrameState) {
 
 /// Save the current SP on the CFP
 fn gen_save_sp(asm: &mut Assembler, stack_size: usize) {
+    // Skip the store if an earlier store in this block already wrote the same
+    // value. Nothing modifies a frame's cfp->sp while it's not the innermost
+    // frame: callees write only their own cfp->sp, side exits and exception
+    // unwinding never return to this code, and rb_zjit_materialize_frames()
+    // only reads cfp->sp.
+    if asm.note_sp_save(stack_size) {
+        return;
+    }
     // Update cfp->sp which will be read by the interpreter. We also have the SP register in JIT
     // code, and ZJIT's codegen currently assumes the SP register doesn't move, e.g. gen_param().
     // So we don't update the SP register here. We could update the SP register to avoid using
@@ -5420,6 +5437,12 @@ fn gen_prepare_fallback_call(jit: &JITState, asm: &mut Assembler, function: &Fun
     gen_save_sp(asm, state.stack_size());
     gen_spill_locals(jit, asm, state);
     gen_spill_stack(jit, asm, function, state);
+
+    // VM fallback helpers execute an instruction on the current frame, which
+    // moves cfp->sp (e.g. vm_sendish() pops the receiver and arguments), so
+    // frame writes recorded above no longer describe the frame's memory once
+    // the helper returns. Forget them so that later calls store fresh values.
+    asm.clear_frame_write_cache();
 }
 
 /// Build entries for Ruby stack values that need materialization. The actual