pd-vm: JIT native avoid jump back to interpreter as possible

Author: fffonion

pd-vm/src/vm/jit/runtime.rs

@@ -577,16 +577,24 @@ impl Vm {
                     {
                         continue;
                     }
-                    // Chain directly into another already-compiled trace at the current ip.
-                    if !has_yielding_call
-                        && let Some(next_trace_id) = self.jit.compiled_trace_for_ip(self.ip)
-                        && next_trace_id != current_trace_id
-                    {
-                        current_trace_id = next_trace_id;
-                        self.ensure_native_trace(current_trace_id)?;
-                        (entry, root_ip, terminal, has_yielding_call) =
-                            self.native_trace_state(current_trace_id)?;
-                        continue;
+                    if !has_yielding_call {
+                        let ip = self.ip;
+                        let mut next_trace_id = self.jit.compiled_trace_for_ip(ip);
+                        if next_trace_id.is_none() {
+                            next_trace_id = {
+                                let program = &self.program;
+                                self.jit.observe_exit_ip(ip, program)
+                            };
+                        }
+                        if let Some(next_trace_id) = next_trace_id
+                            && next_trace_id != current_trace_id
+                        {
+                            current_trace_id = next_trace_id;
+                            self.ensure_native_trace(current_trace_id)?;
+                            (entry, root_ip, terminal, has_yielding_call) =
+                                self.native_trace_state(current_trace_id)?;
+                            continue;
+                        }
                     }
                     return Ok(ExecOutcome::Continue);
                 }

pd-vm/src/vm/jit/trace.rs

@@ -252,9 +252,9 @@ impl TraceJitEngine {
         if self.blocked_roots.contains(&ip) {
             return None;
         }
-        if !self.is_loop_header(program, ip) {
-            return None;
-        }
+        if !self.is_loop_header(program, ip) {
+            return None;
+        }
 
         let count = self.hot_counts.entry(ip).or_insert(0);
         *count = count.saturating_add(1);
@@ -296,9 +296,57 @@ impl TraceJitEngine {
         }
     }
 
-    pub fn trace_clone(&self, trace_id: usize) -> Option<JitTrace> {
-        self.traces.get(trace_id).cloned()
-    }
+    pub fn trace_clone(&self, trace_id: usize) -> Option<JitTrace> {
+        self.traces.get(trace_id).cloned()
+    }
+
+    pub fn observe_exit_ip(&mut self, ip: usize, program: &Program) -> Option<usize> {
+        if !self.config.enabled || !native_jit_supported() {
+            return None;
+        }
+        // Keep default behavior unchanged: only aggressively chain-compile exit roots when the
+        // user requested the most aggressive hotness policy.
+        if self.config.hot_loop_threshold > 1 {
+            return None;
+        }
+        if let Some(&trace_id) = self.compiled_by_root.get(&ip) {
+            return Some(trace_id);
+        }
+        if self.blocked_roots.contains(&ip) {
+            return None;
+        }
+
+        let line = program
+            .debug
+            .as_ref()
+            .and_then(|debug| debug.line_for_offset(ip));
+        let result = if self.config.hot_loop_threshold == 0 {
+            Err(JitNyiReason::HotLoopThresholdZero)
+        } else {
+            self.compile_trace(program, ip)
+        };
+
+        match result {
+            Ok(trace_id) => {
+                self.attempts.push(JitAttempt {
+                    root_ip: ip,
+                    line,
+                    result: Ok(trace_id),
+                });
+                self.compiled_by_root.insert(ip, trace_id);
+                Some(trace_id)
+            }
+            Err(reason) => {
+                self.attempts.push(JitAttempt {
+                    root_ip: ip,
+                    line,
+                    result: Err(reason),
+                });
+                self.blocked_roots.insert(ip);
+                None
+            }
+        }
+    }
 
     pub fn trace_has_call(&self, trace_id: usize) -> bool {
         self.traces