ZJIT: Materialize side exits from JITFrame metadata

Author: k0kubun

vm.c

@@ -2853,10 +2853,9 @@ rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t
                 cfp->block_code = NULL;
             }
 
-            // Materialize Ruby stack slots kept off the VM stack. On side exit,
-            // the exiting frame is already written by compile_exit_save_state()
-            // and skipped here after materialize_exit_trampoline clears its
-            // jit_return, so this restores older ZJIT frames from stack maps.
+            // Materialize Ruby stack slots kept off the VM stack. Regular
+            // JITFrames restore one frame at a time. A side-exit JITFrame starts
+            // at the current frame and may span older inlined frames.
             int32_t stack_size = (int32_t)jit_frame->stack_size;
             if (stack_size > 0) {
                 VALUE *stack = cfp->sp;
@@ -2872,6 +2871,13 @@ rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t
                     else if (ZJIT_STACK_MAP_SKIP_P(entry)) {
                         stack -= ZJIT_STACK_MAP_SKIP_SIZE(entry);
                     }
+                    else if (ZJIT_STACK_MAP_CONST_P(entry)) {
+                        size_t index = ZJIT_STACK_MAP_CONST_INDEX(entry);
+                        VM_ASSERT(jit_frame->side_exit != NULL);
+                        VM_ASSERT(index < jit_frame->side_exit->constants_size);
+                        stack--;
+                        *stack = jit_frame->side_exit->constants[index];
+                    }
                     else {
                         stack--;
                         *stack = entry;
@@ -2884,6 +2890,31 @@ rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t
         cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
     }
 }
+
+void rb_zjit_record_exit_stack(const char *reason);
+void rb_zjit_exit_recompile(rb_execution_context_t *ec, VALUE compiled_iseq);
+
+// Materialize a side exit while its native frame and saved registers are still
+// available, then run optional cold-path instrumentation and recompilation.
+void
+rb_zjit_materialize_side_exit(rb_execution_context_t *ec, rb_control_frame_t *cfp, VALUE *sp)
+{
+    const zjit_jit_frame_t *jit_frame = CFP_ZJIT_FRAME(cfp);
+    zjit_side_exit_t *side_exit = jit_frame->side_exit;
+    VM_ASSERT(side_exit != NULL);
+
+    // EP may point to an escaped heap environment for block frames, so restore
+    // SP from the live ZJIT SP register rather than deriving it from cfp->ep.
+    cfp->sp = sp + jit_frame->sp_size;
+    rb_zjit_materialize_frames(ec, cfp);
+
+    if (side_exit->reason != NULL) {
+        rb_zjit_record_exit_stack(side_exit->reason);
+    }
+    if (side_exit->compiled_iseq != NULL) {
+        rb_zjit_exit_recompile(ec, (VALUE)side_exit->compiled_iseq);
+    }
+}
 #endif
 
 static inline VALUE

zjit.c

@@ -42,6 +42,8 @@ const zjit_jit_frame_t rb_zjit_c_frame = (zjit_jit_frame_t) {
     .pc = 0,
     .iseq = 0,
     .materialize_block_code = false,
+    .sp_size = 0,
+    .side_exit = NULL,
 };
 
 void rb_zjit_profile_disable(const rb_iseq_t *iseq);

zjit.h

@@ -17,6 +17,7 @@
 // zero, so RB_SPECIAL_CONST_P is false.
 #define ZJIT_STACK_MAP_VREG_TAG 0x08
 #define ZJIT_STACK_MAP_SKIP_TAG 0x10
+#define ZJIT_STACK_MAP_CONST_TAG 0x18
 #define ZJIT_STACK_MAP_TAG_MASK 0xff
 #define ZJIT_STACK_MAP_SHIFT 8
 
@@ -44,6 +45,28 @@ ZJIT_STACK_MAP_SKIP_SIZE(VALUE entry)
     return entry >> ZJIT_STACK_MAP_SHIFT;
 }
 
+static inline bool
+ZJIT_STACK_MAP_CONST_P(VALUE entry)
+{
+    return (entry & ZJIT_STACK_MAP_TAG_MASK) == ZJIT_STACK_MAP_CONST_TAG;
+}
+
+static inline size_t
+ZJIT_STACK_MAP_CONST_INDEX(VALUE entry)
+{
+    return entry >> ZJIT_STACK_MAP_SHIFT;
+}
+
+typedef struct zjit_side_exit {
+    // The compiled unit to recompile after taking this exit, or NULL.
+    const rb_iseq_t *compiled_iseq;
+    // Heap VALUE constants referenced by tagged stack-map entries.
+    VALUE *constants;
+    // Perfetto side-exit reason, or NULL when this exit is not traced.
+    const char *reason;
+    uint32_t constants_size;
+} zjit_side_exit_t;
+
 // JITFrame is defined here as the single source of truth and imported into
 // Rust via bindgen. C code reads fields directly; Rust uses an impl block.
 typedef struct zjit_jit_frame {
@@ -59,11 +82,19 @@ typedef struct zjit_jit_frame {
     // Always false for C frames.
     bool materialize_block_code;
 
+    // Number of operand-stack slots in the top frame. The side-exit
+    // materializer adds this to the live ZJIT SP register.
+    uint32_t sp_size;
+
+    // Side-exit-only metadata. NULL for entry and C-call JITFrames.
+    zjit_side_exit_t *side_exit;
+
     // Number of stack map entries in stack[].
     uint32_t stack_size;
     // Flexible array of stack map entries. Each entry is either an immediate
     // VALUE, a tagged native-stack index from cfp->jit_return for a value
-    // kept by the JIT, or a tagged count of VM stack slots to skip.
+    // kept by the JIT, a tagged side-exit constant-table index, or a tagged
+    // count of VM stack slots to skip.
     VALUE stack[];
 } zjit_jit_frame_t;
 
@@ -92,6 +123,7 @@ void rb_zjit_invalidate_no_singleton_class(VALUE klass);
 void rb_zjit_invalidate_root_box(void);
 void rb_zjit_jit_frame_update_references(zjit_jit_frame_t *jit_frame);
 void rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp);
+void rb_zjit_materialize_side_exit(rb_execution_context_t *ec, rb_control_frame_t *cfp, VALUE *sp);
 size_t rb_zjit_hash_new_size(void);
 bool rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, shape_id_t *shape_id_out);
 

zjit/bindgen/src/main.rs

@@ -336,6 +336,7 @@ fn main() {
         .allowlist_var("ZJIT_STACK_MAP_SHIFT")
         .allowlist_var("ZJIT_STACK_MAP_VREG_TAG")
         .allowlist_var("ZJIT_STACK_MAP_SKIP_TAG")
+        .allowlist_var("ZJIT_STACK_MAP_CONST_TAG")
         .allowlist_var("ZJIT_JIT_RETURN_C_FRAME")
         .allowlist_function("rb_assert_holding_vm_lock")
         .allowlist_function("rb_jit_shape_complex_p")

zjit/src/backend/arm64/mod.rs

@@ -1644,7 +1644,6 @@ impl Assembler {
             let (assignments, num_stack_slots) = trace_compile_phase("linear_scan", || asm.linear_scan(intervals.clone(), regs.len(), &preferred_registers));
 
             asm.stack_state.num_spill_slots = num_stack_slots;
-            asm.stack_state.num_side_exit_stack_map_slots = asm.side_exit_stack_map_slots(&assignments);
             let stack_slot_count = asm.stack_state.stack_slot_count();
             if stack_slot_count > Self::MAX_FRAME_STACK_SLOTS {
                 return Err(CompileError::NativeStackTooLarge);
@@ -1782,7 +1781,7 @@ mod tests {
 
         let val64 = asm.add(CFP, Opnd::UImm(64));
         asm.store(Opnd::mem(64, SP, 0x10), val64);
-        let side_exit = Target::SideExit(Box::new(SideExitTarget { reason: SideExitReason::Interrupt, exit: SideExit { pc: 0.into(), iseq: std::ptr::null(), stack: vec![], locals: vec![], stack_map: None, recompile: None } }));
+        let side_exit = Target::SideExit(Box::new(SideExitTarget { reason: SideExitReason::Interrupt, exit: SideExit { pc: 0.into(), iseq: std::ptr::null(), stack: vec![], locals: vec![], stack_map: None, recompile: None, trace_reason: None } }));
         asm.push_insn(Insn::Joz(val64, side_exit));
         asm.mov(C_ARG_OPNDS[0], C_RET_OPND.with_num_bits(32));
         asm.mov(C_ARG_OPNDS[1], Opnd::mem(64, SP, -8));