Post merge fixups

Author: markshannon

Include/internal/pycore_optimizer.h

@@ -35,7 +35,8 @@ typedef struct {
     uint8_t linked:1;
     uint8_t chain_depth:6;  // Must be big enough for MAX_CHAIN_DEPTH - 1.
     bool warm;
-    int index;           // Index of ENTER_EXECUTOR (if code isn't NULL, below).
+    uint8_t tos_cache;
+    int16_t index;           // Index of ENTER_EXECUTOR (if code isn't NULL, below).
     _PyBloomFilter bloom;
     _PyExecutorLinkListNode links;
     PyCodeObject *code;  // Weak (NULL if no corresponding ENTER_EXECUTOR).
@@ -94,7 +95,7 @@ typedef struct _PyExecutorObject {
 // Export for '_opcode' shared extension (JIT compiler).
 PyAPI_FUNC(_PyExecutorObject*) _Py_GetExecutor(PyCodeObject *code, int offset);
 
-void _Py_ExecutorInit(_PyExecutorObject *, const _PyBloomFilter *);
+void _Py_ExecutorInit(_PyExecutorObject *, const _PyBloomFilter *, int tos_cache);
 void _Py_ExecutorDetach(_PyExecutorObject *);
 void _Py_BloomFilter_Init(_PyBloomFilter *);
 void _Py_BloomFilter_Add(_PyBloomFilter *bloom, void *obj);

Python/bytecodes.c

@@ -2972,10 +2972,7 @@ dummy_func(
                 else {
                     this_instr[1].counter = initial_jump_backoff_counter();
                     assert(tstate->current_executor == NULL);
-                    tstate->jit_exit = NULL;
-#if defined(_Py_TIER2) & defined(Py_DEBUG)
-                    current_cached_values = 0;
-#endif
+                    assert(executor->vm_data.tos_cache == 0);
                     TIER1_TO_TIER2(executor);
                 }
             }
@@ -3040,10 +3037,6 @@ dummy_func(
                 }
                 DISPATCH_GOTO();
             }
-#if defined(_Py_TIER2) & defined(Py_DEBUG)
-            current_cached_values = 0;
-#endif
-            tstate->jit_exit = NULL;
             TIER1_TO_TIER2(executor);
             #else
             Py_FatalError("ENTER_EXECUTOR is not supported in this build");
@@ -5432,13 +5425,19 @@ dummy_func(
                 GOTO_TIER_ONE(target);
                 Py_UNREACHABLE();
             }
-            _PyExecutorObject *executor;
+            _PyExecutorObject *executor = NULL;
             if (target->op.code == ENTER_EXECUTOR) {
                 PyCodeObject *code = _PyFrame_GetCode(frame);
                 executor = code->co_executors->executors[target->op.arg];
-                Py_INCREF(executor);
+                int tos_cache = uopcode - _COLD_EXIT_r00;
+                if (tos_cache == executor->vm_data.tos_cache) {
+                    Py_INCREF(executor);
+                }
+                else {
+                    executor = NULL;
+                }
             }
-            else {
+            if (executor == NULL) {
                 _PyExecutorObject *previous_executor = _PyExecutor_FromExit(exit);
                 assert(tstate->current_executor == (PyObject *)previous_executor);
                 int chain_depth = previous_executor->vm_data.chain_depth + 1;

Python/ceval.c

@@ -1132,10 +1132,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 #endif
     }
 
-#if defined(_Py_TIER2) & defined(Py_DEBUG)
-    int current_cached_values = -1;
-#endif
-
 #if Py_TAIL_CALL_INTERP
 #   if Py_STATS
         return _TAIL_CALL_start_frame(frame, NULL, tstate, NULL, 0, lastopcode);
@@ -1175,15 +1171,24 @@ _PyJitEntryFuncPtr _Py_jit_entry = _PyTier2Interpreter;
 _Py_CODEUNIT *
 _PyTier2Interpreter(
     _PyExecutorObject *current_executor, _PyInterpreterFrame *frame,
-    _PyStackRef *stack_pointer, PyThreadState *tstate,
-    _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2
+    _PyStackRef *stack_pointer, PyThreadState *tstate
 ) {
     const _PyUOpInstruction *next_uop;
     int oparg;
-tier2_start:
 
+    /* Set up "jit" state after entry from tier 1.
+     * This mimics what the jit trampoline function does. */
+    int current_cached_values = current_executor->vm_data.tos_cache;
+    _PyStackRef _tos_cache0 = current_cached_values ? stack_pointer[-current_cached_values] : PyStackRef_ZERO_BITS;
+    _PyStackRef _tos_cache1 = stack_pointer[-1-(current_cached_values&1)]; /* Correct value for 2 or 3, harmless junk for 0 or 1 */
+    _PyStackRef _tos_cache2 = stack_pointer[-1]; /* Correct value for 3, harmless junk otherwise */
+    stack_pointer -= current_cached_values;
+    tstate->jit_exit = NULL;
+
+tier2_start:
     next_uop = current_executor->trace;
-    assert(next_uop->opcode == _START_EXECUTOR || next_uop->opcode == _COLD_EXIT);
+    assert(next_uop->opcode == _START_EXECUTOR_r00 + current_cached_values ||
+           next_uop->opcode == _COLD_EXIT_r00 + current_cached_values);
 
 #undef LOAD_IP
 #define LOAD_IP(UNUSED) (void)0

Python/ceval_macros.h

@@ -358,6 +358,8 @@ _PyFrame_SetStackPointer(frame, stack_pointer)
 
 /* Tier-switching macros. */
 
+/* It is always safe to read stack_pointer[-1] and stack_pointer[-2]
+ * due to stack layout, even if the value is meaningless */
 #define TIER1_TO_TIER2(EXECUTOR)                        \
 do {                                                   \
     OPT_STAT_INC(traces_executed);                     \
@@ -375,9 +377,6 @@ do {                                                   \
 do {                                                   \
     OPT_STAT_INC(traces_executed);                     \
     current_executor = (EXECUTOR);                     \
-    goto tier2_start;                                  \
-    assert(next_uop->opcode == _START_EXECUTOR_r00 + current_cached_values || \
-    next_uop->opcode == _COLD_EXIT_r00 + current_cached_values); \
     goto tier2_start; \
 } while (0)
 

Python/executor_cases.c.h

@@ -535,7 +535,8 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     unsigned char *code = memory;
     state.trampolines.mem = memory + code_size;
     unsigned char *data = memory + code_size + state.trampolines.size + code_padding;
-    assert(trace[0].opcode == _START_EXECUTOR_r00 || trace[0].opcode == _COLD_EXIT_r00);
+    assert(trace[0].opcode == _START_EXECUTOR_r00 + executor->vm_data.tos_cache ||
+           trace[0].opcode == _COLD_EXIT_r00 + executor->vm_data.tos_cache);
     for (size_t i = 0; i < length; i++) {
         const _PyUOpInstruction *instruction = &trace[i];
         group = &stencil_groups[instruction->opcode];