Specialize on deopt when tracing

Author: Fidget-Spinner

Include/internal/pycore_interp_structs.h

@@ -765,6 +765,7 @@ typedef struct _PyJitState {
     _PyUOpInstruction *jit_tracer_code_buffer;
     _Py_CODEUNIT *jit_tracer_insert_exec_instr;
     _Py_CODEUNIT *jit_tracer_close_loop_instr;
+    _Py_CODEUNIT *last_specialized_instr;
     int jit_tracer_initial_stack_depth;
     int jit_tracer_initial_chain_depth;
     PyCodeObject *jit_tracer_initial_code; // Strong

Python/ceval_macros.h

@@ -137,6 +137,9 @@
 
 #if _Py_TAIL_CALL_INTERP || USE_COMPUTED_GOTOS
 #  define IS_JIT_TRACING() (DISPATCH_TABLE_VAR == TRACING_DISPATCH_TABLE)
+// tstate->interp->jit_state.last_specialized_instr != this_instr is required to not get stuck in infinite
+// specialization loops due to specialization failure.
+#  define IS_JIT_TRACING_MAKING_PROGRESS() (IS_JIT_TRACING() && tstate->interp->jit_state.last_specialized_instr != this_instr)
 #  define ENTER_TRACING() \
     DISPATCH_TABLE_VAR = TRACING_DISPATCH_TABLE;
 #  define LEAVE_TRACING() \
@@ -214,6 +217,14 @@ do { \
         DISPATCH_GOTO(); \
     }
 
+#define TRACING_SPECIALIZE_DISPATCH_SAME_OPARG() \
+{ \
+    tstate->interp->jit_state.last_specialized_instr = this_instr; \
+    opcode = next_instr->op.code; \
+    PRE_DISPATCH_GOTO(); \
+    DISPATCH_GOTO(); \
+}
+
 #define DISPATCH_INLINED(NEW_FRAME)                     \
     do {                                                \
         assert(tstate->interp->eval_frame == NULL);     \
@@ -225,11 +236,13 @@ do { \
     } while (0)
 
 #define TRACING_DISPATCH_INLINED(NEW_FRAME) \
+    tstate->interp->jit_state.last_specialized_instr = this_instr; \
     RECORD_TRACE_NO_DISPATCH(); \
     DISPATCH_INLINED(NEW_FRAME);
 
 #define TRACING_DISPATCH() \
     { \
+        tstate->interp->jit_state.last_specialized_instr = this_instr; \
         assert(frame->stackpointer == NULL); \
         RECORD_TRACE_NO_DISPATCH(); \
         NEXTOPARG(); \
@@ -337,9 +350,14 @@ GETITEM(PyObject *v, Py_ssize_t i) {
 /* This takes a uint16_t instead of a _Py_BackoffCounter,
  * because it is used directly on the cache entry in generated code,
  * which is always an integral type. */
+#if _Py_TIER2
+// Force re-specialization when tracing a side exit to get good side exits.
+#define ADAPTIVE_COUNTER_TRIGGERS(COUNTER) \
+    backoff_counter_triggers(forge_backoff_counter((COUNTER))) || IS_JIT_TRACING_MAKING_PROGRESS()
+#else
 #define ADAPTIVE_COUNTER_TRIGGERS(COUNTER) \
     backoff_counter_triggers(forge_backoff_counter((COUNTER)))
-
+#endif
 #define ADVANCE_ADAPTIVE_COUNTER(COUNTER) \
     do { \
         (COUNTER) = advance_backoff_counter((COUNTER)); \

Python/generated_tracer_cases.c.h

@@ -904,9 +904,10 @@ _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_
     tstate->interp->jit_state.jit_tracer_previous_exit = exit;
     _Py_BloomFilter_Init(&tstate->interp->jit_state.jit_tracer_dependencies);
     tstate->interp->jit_state.jit_tracer_initial_stack_depth = curr_stackdepth;
-    tstate->interp->jit_state.jit_tracer_initial_chain_depth = chain_depth;
+    tstate->interp->jit_state.jit_tracer_initial_chain_depth = chain_depth % MAX_CHAIN_DEPTH;
     tstate->interp->jit_state.jit_tracer_current_frame = frame;
     tstate->interp->jit_state.jit_tracer_dependencies_still_valid = true;
+    tstate->interp->jit_state.last_specialized_instr = NULL;
 }
 
 void

Python/optimizer.c

@@ -44,6 +44,7 @@ def __init__(self, out: CWriter, labels: dict[str, Label], cannot_escape: bool =
             **self._replacers,
             "DISPATCH": self.dispatch,
             "DISPATCH_INLINED": self.dispatch_inlined,
+            "DISPATCH_SAME_OPARG": self.dispatch_same_oparg,
         }
 
     def dispatch(
@@ -75,6 +76,26 @@ def dispatch_inlined(
         self.out.start_line()
         self.emit("TRACING_DISPATCH_INLINED")
         return False
+
+    def dispatch_same_oparg(
+        self,
+        tkn: Token,
+        tkn_iter: TokenIterator,
+        uop: CodeSection,
+        storage: Storage,
+        inst: Instruction | None,
+    ) -> bool:
+        if storage.spilled:
+            raise analysis_error("stack_pointer needs reloading before dispatch", tkn)
+        storage.stack.flush(self.out)
+        self.out.start_line()
+        if "specializing" in uop.annotations:
+            self.emit("TRACING_SPECIALIZE_DISPATCH_SAME_OPARG")
+        else:
+            self.emit(tkn)
+        emit_to(self.out, tkn_iter, "SEMI")
+        return False
+
     def record_jump_taken(
         self,
         tkn: Token,