Optimize on side exits

Include/cpython/optimizer.h

@@ -36,10 +36,17 @@ typedef struct {
     uint64_t operand;  // A cache entry
 } _PyUOpInstruction;
 
+typedef struct _exit_data {
+    uint32_t target;
+    int16_t hotness;
+    struct _PyExecutorObject *executor;
+} _PyExitData;
+
 typedef struct _PyExecutorObject {
     PyObject_VAR_HEAD
     _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
-    _PyUOpInstruction trace[1];
+    _PyUOpInstruction *trace;
+    _PyExitData exits[1];
 } _PyExecutorObject;
 
 typedef struct _PyOptimizerObject _PyOptimizerObject;

Include/internal/pycore_optimizer.h

@@ -17,12 +17,6 @@ extern PyTypeObject _PyDefaultOptimizer_Type;
 extern PyTypeObject _PyUOpExecutor_Type;
 extern PyTypeObject _PyUOpOptimizer_Type;
 
-typedef struct _exit_data {
-    uint32_t target;
-    int16_t hotness;
-    _PyExecutorObject *exit;
-} _PyExitData;
-
 #ifdef __cplusplus
 }
 #endif

Python/bytecodes.c

@@ -4013,21 +4013,21 @@ dummy_func(
         ///////// Tier-2 only opcodes /////////
 
         op (_GUARD_IS_TRUE_POP, (flag -- )) {
-            DEOPT_IF(Py_IsFalse(flag));
+            EXIT_IF(Py_IsFalse(flag));
             assert(Py_IsTrue(flag));
         }
 
         op (_GUARD_IS_FALSE_POP, (flag -- )) {
-            DEOPT_IF(Py_IsTrue(flag));
+            EXIT_IF(Py_IsTrue(flag));
             assert(Py_IsFalse(flag));
         }
 
         op (_GUARD_IS_NONE_POP, (val -- )) {
-            DEOPT_IF(!Py_IsNone(val));
+            EXIT_IF(!Py_IsNone(val));
         }
 
         op (_GUARD_IS_NOT_NONE_POP, (val -- )) {
-            DEOPT_IF(Py_IsNone(val));
+            EXIT_IF(Py_IsNone(val));
             Py_DECREF(val);
         }
 
@@ -4078,24 +4078,29 @@ dummy_func(
 
         op(_COLD_EXIT, (--)) {
             TIER_TWO_ONLY
-            _PyExitData *exit = PyOptimizer_GetExit(current_executor, oparg);
+            _PyExitData *exit = &current_executor->exits[oparg];
             Py_DECREF(current_executor);
             exit->hotness++;
             DEOPT_IF(exit->hotness < 0);
             PyCodeObject *code = _PyFrame_GetCode(frame);
             _Py_CODEUNIT *target = _PyCode_CODE(code) + exit->target;
-            _PyOptimizerObject *opt = interp->optimizer;
             _PyExecutorObject *executor = NULL;
-            int optimized = opt->optimize(opt, code, target, &executor, (int)(stack_pointer - _PyFrame_Stackbase(frame)));
-            ERROR_IF(optimized < 0, error);
-            if (optimized) {
-                exit->exit = executor;
+            if (target->op.code == ENTER_EXECUTOR) {
+                _PyExecutorObject *executor = code->co_executors->executors[oparg & 255];
                 Py_INCREF(executor);
-                GOTO_TIER_TWO();
-            }
-            else {
-                exit->hotness = -10000; /* Choose a better number */
+            } else {
+                _PyOptimizerObject *opt = tstate->interp->optimizer;
+                int optimized = opt->optimize(opt, code, target, &executor, (int)(stack_pointer - _PyFrame_Stackbase(frame)));
+                if (optimized <= 0) {
+                    next_instr = target;
+                    exit->hotness = -10000; /* Choose a better number */
+                    ERROR_IF(optimized < 0, error);
+                    DISPATCH();
+                }
             }
+            exit->executor = current_executor = executor;
+            Py_INCREF(current_executor);
+            GOTO_TIER_TWO();
         }
 
 

Python/ceval.c

@@ -1073,7 +1073,22 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     Py_DECREF(current_executor);
     DISPATCH();
 
+// Jump here from EXIT_IF()
+side_exit:
+    OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
+    UOP_STAT_INC(uopcode, miss);
+    _PyExitData *exit = &current_executor->exits[next_uop[-1].target];
+    DPRINTF(2, "SIDE EXIT: [UOp %d (%s), oparg %d, operand %" PRIu64 ", target %d @ %d -> %s]\n",
+            uopcode, _PyUOpName(uopcode), next_uop[-1].oparg, next_uop[-1].operand, exit->target,
+            (int)(next_uop - current_executor->trace - 1),
+            _PyOpcode_OpName[frame->instr_ptr->op.code]);
+    Py_INCREF(exit->executor);
+    Py_DECREF(current_executor);
+    current_executor = exit->executor;
+    GOTO_TIER_TWO();
+
 }
+
 #if defined(__GNUC__)
 #  pragma GCC diagnostic pop
 #elif defined(_MSC_VER) /* MS_WINDOWS */

Python/optimizer.c

@@ -281,8 +281,8 @@ PySequenceMethods uop_as_sequence = {
 PyTypeObject _PyUOpExecutor_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     .tp_name = "uop_executor",
-    .tp_basicsize = offsetof(_PyExecutorObject, trace),
-    .tp_itemsize = sizeof(_PyUOpInstruction),
+    .tp_basicsize = offsetof(_PyExecutorObject, exits),
+    .tp_itemsize = 1,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
     .tp_dealloc = (destructor)uop_dealloc,
     .tp_as_sequence = &uop_as_sequence,
@@ -700,9 +700,10 @@ translate_bytecode_to_trace(
  * NOPs are excluded from the count.
 */
 static int
-compute_used(_PyUOpInstruction *buffer, uint32_t *used)
+compute_used(_PyUOpInstruction *buffer, uint32_t *used, int *exit_count_ptr)
 {
     int count = 0;
+    int exit_count = 0;
     SET_BIT(used, 0);
     for (int i = 0; i < UOP_MAX_TRACE_LENGTH; i++) {
         if (!BIT_IS_SET(used, i)) {
@@ -720,14 +721,62 @@ compute_used(_PyUOpInstruction *buffer, uint32_t *used)
             /* Mark target as reachable */
             SET_BIT(used, buffer[i].oparg);
         }
+        if (_PyUop_Flags[opcode] & HAS_EXIT_FLAG) {
+            exit_count++;
+        }
         if (opcode == NOP) {
             count--;
             UNSET_BIT(used, i);
         }
     }
+    *exit_count_ptr = exit_count;
     return count;
 }
 
+/* Executor side exits */
+
+typedef struct _cold_exit {
+    PyObject_VAR_HEAD
+    _PyVMData vm_data;
+    union {
+        void *jit_code; // Will be a function pointer
+        _PyUOpInstruction *trace;
+    };
+    _PyUOpInstruction cold;
+} _PyColdExitObject;
+
+static _PyColdExitObject COLD_EXITS[UOP_MAX_TRACE_LENGTH];
+
+static _PyExecutorObject *
+allocate_executor(int exit_count, int length)
+{
+    int size = exit_count*sizeof(_PyExitData) + length*sizeof(_PyUOpInstruction);
+    _PyExecutorObject *res = PyObject_NewVar(_PyExecutorObject, &_PyUOpExecutor_Type, size);
+    if (res != NULL) {
+        res->trace = (_PyUOpInstruction *)(res->exits + exit_count);
+    }
+    return res;
+}
+
+static int cold_exits_initialized = 0;
+
+static void
+initialize_cold_exits(void) {
+    if (cold_exits_initialized) {
+        return;
+    }
+    cold_exits_initialized = 1;
+    for (int i = 0; i < UOP_MAX_TRACE_LENGTH; i++) {
+        COLD_EXITS[i].ob_base.ob_base.ob_refcnt = _Py_IMMORTAL_REFCNT;
+        COLD_EXITS[i].ob_base.ob_base.ob_type = &_PyUOpExecutor_Type;
+        COLD_EXITS[i].ob_base.ob_size = 0;
+        COLD_EXITS[i].vm_data = (_PyVMData){ 0 };
+        COLD_EXITS[i].cold.opcode = _COLD_EXIT;
+        COLD_EXITS[i].cold.oparg = i;
+        COLD_EXITS[i].trace = &COLD_EXITS[i].cold;
+    }
+}
+
 /* Makes an executor from a buffer of uops.
  * Account for the buffer having gaps and NOPs by computing a "used"
  * bit vector and only copying the used uops. Here "used" means reachable
@@ -736,12 +785,22 @@ compute_used(_PyUOpInstruction *buffer, uint32_t *used)
 static _PyExecutorObject *
 make_executor_from_uops(_PyUOpInstruction *buffer, _PyBloomFilter *dependencies)
 {
+    if (!cold_exits_initialized) {
+        initialize_cold_exits();
+    }
     uint32_t used[(UOP_MAX_TRACE_LENGTH + 31)/32] = { 0 };
-    int length = compute_used(buffer, used);
-    _PyExecutorObject *executor = PyObject_NewVar(_PyExecutorObject, &_PyUOpExecutor_Type, length);
+    int exit_count;
+    int length = compute_used(buffer, used, &exit_count);
+    _PyExecutorObject *executor = allocate_executor(exit_count, length);
     if (executor == NULL) {
         return NULL;
     }
+    /* Initialize exits */
+    for (int i = 0; i < exit_count; i++) {
+        executor->exits[i].executor = (_PyExecutorObject *)&COLD_EXITS[i];
+        executor->exits[i].hotness = -273; /* Physics reference */
+    }
+    int next_exit = 0;
     int dest = length - 1;
     /* Scan backwards, so that we see the destinations of jumps before the jumps themselves. */
     for (int i = UOP_MAX_TRACE_LENGTH-1; i >= 0; i--) {
@@ -757,11 +816,17 @@ make_executor_from_uops(_PyUOpInstruction *buffer, _PyBloomFilter *dependencies)
             int oparg = executor->trace[dest].oparg;
             executor->trace[dest].oparg = buffer[oparg].oparg;
         }
+        if (_PyUop_Flags[opcode] & HAS_EXIT_FLAG) {
+            executor->exits[next_exit].target = buffer[i].target;
+            executor->trace[dest].target = next_exit;
+            next_exit++;
+        }
         /* Set the oparg to be the destination offset,
          * so that we can set the oparg of earlier jumps correctly. */
         buffer[i].oparg = dest;
         dest--;
     }
+    assert(next_exit == exit_count);
     assert(dest == -1);
     _Py_ExecutorInit(executor, dependencies);
 #ifdef Py_DEBUG
@@ -1138,36 +1203,5 @@ _Py_Executors_InvalidateAll(PyInterpreterState *interp)
     interp->executor_list_head = NULL;
 }
 
-/* Executor side exits */
-
-
-typedef struct _cold_exit {
-    PyObject_VAR_HEAD
-    _PyVMData vm_data;
-    union {
-        void *jit_code; // Will be a function pointer
-        _PyUOpInstruction *trace;
-    };
-    _PyUOpInstruction cold;
-} _PyColdExitObject;
-
-static int cold_exits_initialized = 0;
-_PyColdExitObject COLD_EXITS[_Py_UOP_MAX_TRACE_LENGTH];
 
-static void
-initialize_cold_exits(void) {
-    if (cold_exits_initialized) {
-        return;
-    }
-    cold_exits_initialized = 1;
-    for (int i = 0; i < _Py_UOP_MAX_TRACE_LENGTH; i++) {
-        COLD_EXITS[i].ob_base.ob_base.ob_refcnt = _Py_IMMORTAL_REFCNT;
-        COLD_EXITS[i].ob_base.ob_base.ob_type = &_PyUOpExecutor_Type;
-        COLD_EXITS[i].ob_base.ob_size = 0;
-        COLD_EXITS[i].vm_data = (_PyVMData){ 0 };
-        COLD_EXITS[i].cold.opcode = _COLD_EXIT;
-        COLD_EXITS[i].cold.oparg = i;
-        COLD_EXITS[i].trace = &COLD_EXITS[i].cold;
-    }
-}