ZJIT: Count exits coming from jit_exception (ruby#14428)

Author: k0kubun

vm.c

@@ -445,7 +445,7 @@ jit_compile(rb_execution_context_t *ec)
 
         // At call-threshold, compile the ISEQ with ZJIT.
         if (body->jit_entry_calls == rb_zjit_call_threshold) {
-            rb_zjit_compile_iseq(iseq, ec, false);
+            rb_zjit_compile_iseq(iseq, false);
         }
     }
 #endif
@@ -493,8 +493,25 @@ jit_compile_exception(rb_execution_context_t *ec)
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
 
-    // Increment the ISEQ's call counter and trigger JIT compilation if not compiled
+#if USE_ZJIT
+    if (body->jit_exception == NULL && rb_zjit_enabled_p) {
+        body->jit_exception_calls++;
+
+        // At profile-threshold, rewrite some of the YARV instructions
+        // to zjit_* instructions to profile these instructions.
+        if (body->jit_exception_calls == rb_zjit_profile_threshold) {
+            rb_zjit_profile_enable(iseq);
+        }
+
+        // At call-threshold, compile the ISEQ with ZJIT.
+        if (body->jit_exception_calls == rb_zjit_call_threshold) {
+            rb_zjit_compile_iseq(iseq, true);
+        }
+    }
+#endif
+
 #if USE_YJIT
+    // Increment the ISEQ's call counter and trigger JIT compilation if not compiled
     if (body->jit_exception == NULL && rb_yjit_enabled_p) {
         body->jit_exception_calls++;
         if (body->jit_exception_calls == rb_yjit_call_threshold) {

yjit.c

@@ -655,7 +655,7 @@ rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit
         else {
             iseq->body->jit_entry = (rb_jit_func_t)code_ptr;
         }
-}
+    }
 }
 
 // GC root for interacting with the GC

zjit.c

@@ -159,18 +159,22 @@ rb_zjit_reserve_addr_space(uint32_t mem_size)
 void rb_zjit_profile_disable(const rb_iseq_t *iseq);
 
 void
-rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception)
+rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception)
 {
     RB_VM_LOCKING() {
         rb_vm_barrier();
 
         // Compile a block version starting at the current instruction
-        uint8_t *rb_zjit_iseq_gen_entry_point(const rb_iseq_t *iseq, rb_execution_context_t *ec); // defined in Rust
-        uintptr_t code_ptr = (uintptr_t)rb_zjit_iseq_gen_entry_point(iseq, ec);
+        uint8_t *rb_zjit_iseq_gen_entry_point(const rb_iseq_t *iseq, bool jit_exception); // defined in Rust
+        uintptr_t code_ptr = (uintptr_t)rb_zjit_iseq_gen_entry_point(iseq, jit_exception);
 
-        // TODO: support jit_exception
-        iseq->body->jit_entry = (rb_jit_func_t)code_ptr;
-}
+        if (jit_exception) {
+            iseq->body->jit_exception = (rb_jit_func_t)code_ptr;
+        }
+        else {
+            iseq->body->jit_entry = (rb_jit_func_t)code_ptr;
+        }
+    }
 }
 
 extern VALUE *rb_vm_base_ptr(struct rb_control_frame_struct *cfp);

zjit.h

@@ -13,7 +13,7 @@
 extern bool rb_zjit_enabled_p;
 extern uint64_t rb_zjit_call_threshold;
 extern uint64_t rb_zjit_profile_threshold;
-void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception);
+void rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception);
 void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec);
 void rb_zjit_profile_enable(const rb_iseq_t *iseq);
 void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
@@ -26,7 +26,7 @@ void rb_zjit_before_ractor_spawn(void);
 void rb_zjit_tracing_invalidate_all(void);
 #else
 #define rb_zjit_enabled_p false
-static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
+static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception) {}
 static inline void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec) {}
 static inline void rb_zjit_profile_enable(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}

zjit/src/codegen.rs

@@ -65,9 +65,11 @@ impl JITState {
     }
 }
 
-/// CRuby API to compile a given ISEQ
+/// CRuby API to compile a given ISEQ.
+/// If jit_exception is true, compile JIT code for handling exceptions.
+/// See jit_compile_exception() for details.
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *const u8 {
+pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, jit_exception: bool) -> *const u8 {
     // Do not test the JIT code in HIR tests
     if cfg!(test) {
         return std::ptr::null();
@@ -77,11 +79,12 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *co
     // with_vm_lock() does nothing if the program doesn't use Ractors.
     with_vm_lock(src_loc!(), || {
         let cb = ZJITState::get_code_block();
-        let mut code_ptr = with_time_stat(compile_time_ns, || gen_iseq_entry_point(cb, iseq));
+        let mut code_ptr = with_time_stat(compile_time_ns, || gen_iseq_entry_point(cb, iseq, jit_exception));
 
         if let Err(err) = &code_ptr {
-            // Assert that the ISEQ compiles if RubyVM::ZJIT.assert_compiles is enabled
-            if ZJITState::assert_compiles_enabled() {
+            // Assert that the ISEQ compiles if RubyVM::ZJIT.assert_compiles is enabled.
+            // We assert only `jit_exception: false` cases until we support exception handlers.
+            if ZJITState::assert_compiles_enabled() && !jit_exception {
                 let iseq_location = iseq_get_location(iseq, 0);
                 panic!("Failed to compile: {iseq_location}");
             }
@@ -102,7 +105,12 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *co
 }
 
 /// Compile an entry point for a given ISEQ
-fn gen_iseq_entry_point(cb: &mut CodeBlock, iseq: IseqPtr) -> Result<CodePtr, CompileError> {
+fn gen_iseq_entry_point(cb: &mut CodeBlock, iseq: IseqPtr, jit_exception: bool) -> Result<CodePtr, CompileError> {
+    // We don't support exception handlers yet
+    if jit_exception {
+        return Err(CompileError::ExceptionHandler);
+    }
+
     // Compile ISEQ into High-level IR
     let function = compile_iseq(iseq).inspect_err(|_| {
         incr_counter!(failed_iseq_count);

zjit/src/stats.rs

@@ -116,6 +116,7 @@ make_counters! {
 
     // compile_error_: Compile error reasons
     compile_error_iseq_stack_too_large,
+    compile_error_exception_handler,
     compile_error_out_of_memory,
     compile_error_register_spill_on_ccall,
     compile_error_register_spill_on_alloc,
@@ -180,6 +181,7 @@ pub fn exit_counter_ptr_for_call_type(call_type: crate::hir::CallType) -> *mut u
 #[derive(Clone, Debug, PartialEq)]
 pub enum CompileError {
     IseqStackTooLarge,
+    ExceptionHandler,
     OutOfMemory,
     RegisterSpillOnAlloc,
     RegisterSpillOnCCall,
@@ -194,6 +196,7 @@ pub fn exit_counter_for_compile_error(compile_error: &CompileError) -> Counter {
     use crate::stats::Counter::*;
     match compile_error {
         IseqStackTooLarge     => compile_error_iseq_stack_too_large,
+        ExceptionHandler      => compile_error_exception_handler,
         OutOfMemory           => compile_error_out_of_memory,
         RegisterSpillOnAlloc  => compile_error_register_spill_on_alloc,
         RegisterSpillOnCCall  => compile_error_register_spill_on_ccall,