GH-109369: Exit tier 2 if executor is invalid (GH-111657)

6 files changed, 34 insertions, 5 deletions

diff --git a/Python/abstract_interp_cases.c.h b/Python/abstract_interp_cases.c.h
index 384a112..8892d26 100644
--- a/Python/abstract_interp_cases.c.h
+++ b/Python/abstract_interp_cases.c.h
@@ -944,3 +944,7 @@
             PARTITIONNODE_OVERWRITE((_Py_PARTITIONNODE_t *)PARTITIONNODE_NULLROOT, PEEK(-(-1 - oparg)), true);
             break;
         }
+
+        case _CHECK_VALIDITY: {
+            break;
+        }
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index f879ea5..7010042 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -4021,6 +4021,11 @@ dummy_func(
             memmove(&stack_pointer[-1 - oparg], &stack_pointer[-oparg], oparg * sizeof(stack_pointer[0]));
         }
 
+        op(_CHECK_VALIDITY, (--)) {
+            TIER_TWO_ONLY
+            DEOPT_IF(!current_executor->base.vm_data.valid);
+        }
+
 
 // END BYTECODES //
 
diff --git a/Python/ceval.c b/Python/ceval.c
index 670f312..68bed93 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1057,7 +1057,7 @@ error_tier_two:
 deoptimize:
     // On DEOPT_IF we just repeat the last instruction.
     // This presumes nothing was popped from the stack (nor pushed).
-    DPRINTF(2, "DEOPT: [Opcode %d, operand %" PRIu64 "]\n", opcode, operand);
+    DPRINTF(2, "DEOPT: [Opcode %d, operand %" PRIu64 " @ %d]\n", opcode, operand, (int)(next_uop-current_executor->trace-1));
     OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
     UOP_STAT_INC(opcode, miss);
     frame->return_offset = 0;  // Dispatch to frame->instr_ptr
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index d94a7cc..2f4d55b 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -3273,4 +3273,10 @@
             break;
         }
 
+        case _CHECK_VALIDITY: {
+            TIER_TWO_ONLY
+            DEOPT_IF(!current_executor->base.vm_data.valid, _CHECK_VALIDITY);
+            break;
+        }
+
 #undef TIER_TWO
diff --git a/Python/optimizer.c b/Python/optimizer.c
index 065e127..42279be 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -472,8 +472,8 @@ translate_bytecode_to_trace(
     } \
     reserved = (n);  // Keep ADD_TO_TRACE / ADD_TO_STUB honest
 
-// Reserve space for main+stub uops, plus 2 for _SET_IP and _EXIT_TRACE
-#define RESERVE(main, stub) RESERVE_RAW((main) + (stub) + 2, uop_name(opcode))
+// Reserve space for main+stub uops, plus 3 for _SET_IP, _CHECK_VALIDITY and _EXIT_TRACE
+#define RESERVE(main, stub) RESERVE_RAW((main) + (stub) + 3, uop_name(opcode))
 
 // Trace stack operations (used by _PUSH_FRAME, _POP_FRAME)
 #define TRACE_STACK_PUSH() \
@@ -503,8 +503,9 @@ translate_bytecode_to_trace(
 
 top:  // Jump here after _PUSH_FRAME or likely branches
     for (;;) {
-        RESERVE_RAW(2, "epilogue");  // Always need space for _SET_IP and _EXIT_TRACE
+        RESERVE_RAW(3, "epilogue");  // Always need space for _SET_IP, _CHECK_VALIDITY and _EXIT_TRACE
         ADD_TO_TRACE(_SET_IP, INSTR_IP(instr, code), 0);
+        ADD_TO_TRACE(_CHECK_VALIDITY, 0, 0);
 
         uint32_t opcode = instr->op.code;
         uint32_t oparg = instr->op.arg;
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 61bda80..3c85964 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -12,13 +12,13 @@
 #include <stddef.h>
 #include "pycore_optimizer.h"
 
-
 static void
 remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
 {
     // Note that we don't enter stubs, those SET_IPs are needed.
     int last_set_ip = -1;
     bool need_ip = true;
+    bool maybe_invalid = false;
     for (int pc = 0; pc < buffer_size; pc++) {
         int opcode = buffer[pc].opcode;
         if (opcode == _SET_IP) {
@@ -28,6 +28,16 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
             need_ip = false;
             last_set_ip = pc;
         }
+        else if (opcode == _CHECK_VALIDITY) {
+            if (maybe_invalid) {
+                /* Exiting the trace requires that IP is correct */
+                need_ip = true;
+                maybe_invalid = false;
+            }
+            else {
+                buffer[pc].opcode = NOP;
+            }
+        }
         else if (opcode == _JUMP_TO_TOP || opcode == _EXIT_TRACE) {
             break;
         }
@@ -36,6 +46,9 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
             if (_PyOpcode_opcode_metadata[opcode].flags & (HAS_ERROR_FLAG | HAS_DEOPT_FLAG) || opcode == _PUSH_FRAME) {
                 need_ip = true;
             }
+            if (_PyOpcode_opcode_metadata[opcode].flags & HAS_ESCAPES_FLAG) {
+                maybe_invalid = true;
+            }
         }
     }
 }