properly restore exponential backoffs

Author: Fidget-Spinner

Include/internal/pycore_interp_structs.h

@@ -775,6 +775,8 @@ typedef struct _PyJitTracerState {
     _Py_CODEUNIT *prev_instr;
     PyCodeObject *prev_instr_code; // Strong
     struct _PyExitData *prev_exit;
+    struct _PyExecutorObject *prev_executor; // Strong
+    _Py_CODEUNIT *jump_backward_instr;
     _PyInterpreterFrame *prev_instr_frame;
     _PyBloomFilter dependencies;
 } _PyJitTracerState;

Include/internal/pycore_optimizer.h

@@ -365,7 +365,8 @@ int _PyJit_translate_single_bytecode_to_trace(PyThreadState *tstate, _PyInterpre
 int
 _PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame,
     _Py_CODEUNIT *curr_instr, _Py_CODEUNIT *insert_exec_instr,
-    _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit, int oparg);
+    _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit,
+    _PyExecutorObject *prev_exec, int oparg);
 
 void _PyJit_FinalizeTracing(PyThreadState *tstate);
 

Python/bytecodes.c

@@ -2956,7 +2956,6 @@ dummy_func(
             if (!IS_JIT_TRACING() && backoff_counter_triggers(counter) &&
                 this_instr->op.code == JUMP_BACKWARD_JIT &&
                 next_instr->op.code != ENTER_EXECUTOR) {
-                this_instr[1].counter = restart_backoff_counter(counter);
                 if (tstate->interp->jit_state.code_buffer == NULL) {
                     tstate->interp->jit_state.code_buffer = (_PyUOpInstruction *)_PyObject_VirtualAlloc(UOP_BUFFER_SIZE);
                     if (tstate->interp->jit_state.code_buffer == NULL) {
@@ -2972,7 +2971,7 @@ dummy_func(
                         oparg >>= 8;
                         insert_exec_at--;
                     }
-                    int succ = _PyJit_TryInitializeTracing(tstate, frame, this_instr, insert_exec_at, next_instr, STACK_LEVEL(), 0, NULL, oparg);
+                    int succ = _PyJit_TryInitializeTracing(tstate, frame, this_instr, insert_exec_at, next_instr, STACK_LEVEL(), 0, NULL, NULL, oparg);
                     if (succ) {
                         ENTER_TRACING();
                     }
@@ -5432,8 +5431,7 @@ dummy_func(
                 // Note: it's safe to use target->op.arg here instead of the oparg given by EXTENDED_ARG.
                 // The invariant in the optimizer is the deopt target always points back to the first EXTENDED_ARG.
                 // So setting it to anything else is wrong.
-                int succ = _PyJit_TryInitializeTracing(tstate, frame, target, target, target, STACK_LEVEL(), chain_depth, exit, target->op.arg);
-                exit->temperature = initial_temperature_backoff_counter();
+                int succ = _PyJit_TryInitializeTracing(tstate, frame, target, target, target, STACK_LEVEL(), chain_depth, exit, previous_executor, target->op.arg);
                 if (succ) {
                     GOTO_TIER_ONE_CONTINUE_TRACING(target);
                 }

Python/ceval.c

@@ -998,6 +998,29 @@ bail_tracing_and_jit(PyThreadState *tstate, _PyInterpreterFrame *frame)
     if (!_PyErr_Occurred(tstate) && !_is_sys_tracing) {
         err = _PyOptimizer_Optimize(frame, tstate);
     }
+    // Deal with backoffs
+    _PyExitData *exit = tstate->interp->jit_state.prev_exit;
+    if (exit == NULL) {
+        // We hold a strong reference to the code object, so the instruction won't be freed.
+        if (err <= 0) {
+            assert(tstate->interp->jit_state.jump_backward_instr->op.code == JUMP_BACKWARD_JIT);
+            _Py_BackoffCounter counter = tstate->interp->jit_state.jump_backward_instr[1].counter;
+            tstate->interp->jit_state.jump_backward_instr[1].counter = restart_backoff_counter(counter);
+        }
+        else {
+            tstate->interp->jit_state.jump_backward_instr[1].counter = initial_jump_backoff_counter();
+        }
+    }
+    else {
+        // Likewise, we hold a strong reference to the executor containing this exit, so the exit is guaranteed
+        // to be valid to access.
+        if (err <= 0) {
+            exit->temperature = restart_backoff_counter(exit->temperature);
+        }
+        else {
+            exit->temperature = initial_temperature_backoff_counter();
+        }
+    }
     _PyJit_FinalizeTracing(tstate);
     return err;
 }

Python/executor_cases.c.h

@@ -949,7 +949,7 @@ _PyJit_translate_single_bytecode_to_trace(
 
 // Returns 0 for do not enter tracing, 1 on enter tracing.
 int
-_PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *curr_instr, _Py_CODEUNIT *insert_exec_instr, _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit, int oparg)
+_PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *curr_instr, _Py_CODEUNIT *insert_exec_instr, _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit, _PyExecutorObject *prev_exec, int oparg)
 {
     // A recursive trace.
     // Don't trace into the inner call because it will stomp on the previous trace, causing endless retraces.
@@ -991,6 +991,9 @@ _PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _
     tstate->interp->jit_state.prev_instr_oparg = oparg;
     tstate->interp->jit_state.prev_instr_stacklevel = curr_stackdepth;
     tstate->interp->jit_state.prev_instr_is_super = false;
+    assert(curr_instr->op.code == JUMP_BACKWARD_JIT || (prev_exec != NULL && exit != NULL));
+    tstate->interp->jit_state.jump_backward_instr = curr_instr;
+    tstate->interp->jit_state.prev_executor = (_PyExecutorObject *)Py_XNewRef(prev_exec);
     _Py_BloomFilter_Init(&tstate->interp->jit_state.dependencies);
     return 1;
 }
@@ -1001,6 +1004,7 @@ _PyJit_FinalizeTracing(PyThreadState *tstate)
     Py_CLEAR(tstate->interp->jit_state.initial_code);
     Py_CLEAR(tstate->interp->jit_state.initial_func);
     Py_CLEAR(tstate->interp->jit_state.prev_instr_code);
+    Py_CLEAR(tstate->interp->jit_state.prev_executor);
     tstate->interp->jit_state.code_curr_size = 2;
     tstate->interp->jit_state.code_max_size = UOP_MAX_TRACE_LENGTH - 1;
 }