jit: gh#394 warmup follow-ons — cache FrameRoot shadow-stack cell + correct walk-end handback vsd

majit/majit-gc/src/shadow_stack.rs

@@ -212,7 +212,7 @@ thread_local! {
     static MAX_SHADOW_STACK_DEPTH: Cell<usize> = const { Cell::new(DEFAULT_SHADOW_STACK_DEPTH) };
 
     /// Thread-local shadow stack for individual GcRef roots.
-    static SHADOW_STACK: RefCell<ShadowStack> = RefCell::new(ShadowStack::new());
+    static SHADOW_STACK: RefCell<ShadowStack> = const { RefCell::new(ShadowStack::new()) };
 
     /// Thread-local flat jitframe root stack. Rust tests run multiple
     /// JIT/GC tests in parallel, so using one process-global root stack lets
@@ -372,9 +372,9 @@ struct ShadowStack {
 }
 
 impl ShadowStack {
-    fn new() -> Self {
+    const fn new() -> Self {
         ShadowStack {
-            entries: Vec::with_capacity(64),
+            entries: Vec::new(),
         }
     }
 }
@@ -430,6 +430,37 @@ pub fn get(index: usize) -> GcRef {
     })
 }
 
+/// An opaque handle to this thread's `SHADOW_STACK` cell.
+///
+/// The pointer is stable for the life of the owning thread (the same
+/// invariant `MutatorEntry` relies on for STW walks): the thread-local cell is
+/// never reallocated once created. Resolving it once and reusing it lets a hot
+/// caller re-read roots without paying the thread-local resolution
+/// (`_tlv_get_addr` on macOS) on every access — the root-stack base is held in
+/// a fixed location and re-read cheaply, matching the C backend's
+/// once-per-function `gc_enter_roots_frame` resolution and the x86 JIT's
+/// register-cached root-stack top.
+#[derive(Clone, Copy)]
+pub struct ShadowStackSlot(*const RefCell<ShadowStack>);
+
+/// Resolve this thread's `SHADOW_STACK` cell once, for reuse by [`slot_get`].
+pub fn shadow_stack_slot() -> ShadowStackSlot {
+    ShadowStackSlot(SHADOW_STACK.with(|ss| ss as *const _))
+}
+
+/// Get a GcRef at `index` through a previously resolved [`ShadowStackSlot`].
+///
+/// # Safety
+///
+/// `slot` must have been produced by [`shadow_stack_slot`] on the current
+/// thread and the thread must still be alive (its `SHADOW_STACK` not yet torn
+/// down). No `&mut` borrow of the cell may be held across this call.
+pub unsafe fn slot_get(slot: ShadowStackSlot, index: usize) -> GcRef {
+    // SAFETY: per the contract, `slot.0` points at the live owning-thread cell.
+    let ss = unsafe { &*slot.0 }.borrow();
+    ss.entries[index]
+}
+
 /// Walk all entries on the GcRef shadow stack.
 pub fn walk_roots(mut visitor: impl FnMut(&mut GcRef)) {
     SHADOW_STACK.with(|ss| {

pyre/pyre-jit/src/eval.rs

@@ -172,18 +172,26 @@ struct FrameLocalsRoot {
 /// collection.
 struct FrameRoot {
     depth: usize,
+    slot: majit_gc::shadow_stack::ShadowStackSlot,
 }
 
 impl FrameRoot {
     #[majit_macros::dont_look_inside]
     fn new(frame: &mut PyFrame) -> Self {
         let depth = majit_gc::shadow_stack::push(majit_ir::GcRef(frame as *mut PyFrame as usize));
-        Self { depth }
+        // Resolve the thread-local shadow-stack cell once; `frame()` re-reads
+        // the root through this cached slot instead of paying the thread-local
+        // resolution on every access.
+        let slot = majit_gc::shadow_stack::shadow_stack_slot();
+        Self { depth, slot }
     }
 
     #[majit_macros::dont_look_inside]
     fn frame(&mut self) -> &mut PyFrame {
-        let frame = majit_gc::shadow_stack::get(self.depth).0 as *mut PyFrame;
+        // SAFETY: `slot` was resolved on this thread in `new` and the thread is
+        // still running; no `&mut` borrow of the cell is held here.
+        let frame =
+            unsafe { majit_gc::shadow_stack::slot_get(self.slot, self.depth) }.0 as *mut PyFrame;
         unsafe { &mut *frame }
     }
 
@@ -5745,33 +5753,9 @@ fn eval_loop_jit(frame: &mut PyFrame) -> LoopResult {
         }
         // The ec block above may have run bytecode_trace / perform_actions
         // (collection points) on a fall-through path; re-seed before the
-        // stack-effect reads and the opcode dispatch.
+        // opcode dispatch.
         let f: *mut PyFrame = frame_root.frame() as *mut PyFrame;
         let mut next_instr = unsafe { &*f }.next_instr();
-        let raw_arg: u32 = op_arg.into();
-        let delta = instruction.stack_effect(raw_arg);
-        if delta > 0 {
-            // `frame` is a shared reborrow used only in this block's reads and
-            // the `if` condition below; its last use ends before the `&mut *f`
-            // reborrow in the taken branch, so the two never alias. Keep any
-            // future `frame` use above that write — the raw pointer means the
-            // borrow checker will not catch an overlap introduced here.
-            let frame = unsafe { &*f };
-            let pushed_top = frame.valuestackdepth.saturating_add(delta as usize);
-            let next_pc = opcode_pc + 1;
-            // A JIT handoff can arrive with the stack depth for the point just
-            // after a super-instruction while `last_instr` still names the
-            // super-instruction itself. If metadata proves the current depth
-            // belongs to the next opcode, advance the pc instead of re-running
-            // pushes that are already reflected in the frame stack.
-            if pushed_top > frame.locals_w().len()
-                && pyre_jit_trace::state::depth_based_vsd_for_wcode(frame.pycode as usize, next_pc)
-                    == Some(frame.valuestackdepth)
-            {
-                unsafe { &mut *f }.set_last_instr_from_next_instr(next_pc);
-                continue;
-            }
-        }
         let step_result =
             execute_opcode_step(unsafe { &mut *f }, code, instruction, op_arg, next_instr);
         match step_result {
@@ -6756,9 +6740,15 @@ fn compile_and_run_once(
                     .frame()
                     .restore_resume_state_from(&executed_frame);
             } else if let Some(restart_pc) = walk_end_restart_pc {
-                frame_root
-                    .frame()
-                    .set_last_instr_from_next_instr(restart_pc);
+                // A marker inside a super-instruction closes the loop at
+                // `loop_header_pc + 1`, and the walk already advanced
+                // `valuestackdepth` through the super-instruction. Set both the
+                // resume pc and its operand depth so the handed-back frame is
+                // self-consistent, mirroring the flush leg above and the
+                // blackhole legs (`apply_blackhole_crn_handoff`).
+                let frame = frame_root.frame();
+                frame.set_last_instr_from_next_instr(restart_pc);
+                correct_resume_vsd(frame, restart_pc);
             }
             propagated_exception = pyre_jit_trace::trace::take_walk_end_propagated_exception();
             action