Use _DYNAMIC_EXIT to stitch traces ending in unbalanced retuens or yields.

Author: markshannon

Include/internal/pycore_uop_ids.h

@@ -2589,8 +2589,7 @@ dummy_func(
                     _PyErr_Clear(tstate);
                 }
                 /* iterator ended normally */
-                assert(next_instr[oparg].op.code == END_FOR ||
-                       next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
+                assert(base_opcode(_PyFrame_GetCode(frame), INSTR_OFFSET() + oparg) == END_FOR);
                 Py_DECREF(iter);
                 STACK_SHRINK(1);
                 /* Jump forward oparg, then skip following END_FOR and POP_TOP instruction */
@@ -4351,6 +4350,14 @@ dummy_func(
             assert(tstate->tracing || eval_breaker == FT_ATOMIC_LOAD_UINTPTR_ACQUIRE(_PyFrame_GetCode(frame)->_co_instrumentation_version));
         }
 
+        tier2 op(_RETURN_OFFSET, (--)) {
+            frame->instr_ptr += frame->return_offset;
+        }
+
+        tier2 op(_YIELD_OFFSET, (--)) {
+            frame->instr_ptr += 1 + INLINE_CACHE_ENTRIES_SEND;
+        }
+
 // END BYTECODES //
 
     }

Include/internal/pycore_uop_metadata.h

@@ -655,6 +655,18 @@ extern const struct _PyCode_DEF(8) _Py_InitCleanup;
 
 #ifdef Py_DEBUG
 extern void _PyUOpPrint(const _PyUOpInstruction *uop);
+
+static int
+base_opcode(PyCodeObject *code, int offset)
+{
+    int opcode = _Py_GetBaseOpcode(code, offset);
+    if (opcode == ENTER_EXECUTOR) {
+        int oparg = _PyCode_CODE(code)[offset].op.arg;
+        _PyExecutorObject *ex = code->co_executors->executors[oparg];
+        return ex->vm_data.opcode;
+    }
+    return opcode;
+}
 #endif
 
 

Python/bytecodes.c

@@ -211,6 +211,9 @@ _PyOptimizer_Optimize(
     _PyInterpreterFrame *frame, _Py_CODEUNIT *start,
     PyObject **stack_pointer, _PyExecutorObject **executor_ptr)
 {
+    if (start->op.code == INTERPRETER_EXIT || start->op.code == EXIT_INIT_CHECK) {
+        return 0;
+    }
     PyCodeObject *code = _PyFrame_GetCode(frame);
     assert(PyCode_Check(code));
     PyInterpreterState *interp = _PyInterpreterState_GET();
@@ -747,17 +750,6 @@ translate_bytecode_to_trace(
                     // Reserve space for nuops (+ _SET_IP + _EXIT_TRACE)
                     int nuops = expansion->nuops;
                     RESERVE(nuops + 1); /* One extra for exit */
-                    int16_t last_op = expansion->uops[nuops-1].uop;
-                    if (last_op == _POP_FRAME || last_op == _RETURN_GENERATOR || last_op == _YIELD_VALUE) {
-                        // Check for trace stack underflow now:
-                        // We can't bail e.g. in the middle of
-                        // LOAD_CONST + _POP_FRAME.
-                        if (trace_stack_depth == 0) {
-                            DPRINTF(2, "Trace stack underflow\n");
-                            OPT_STAT_INC(trace_stack_underflow);
-                            goto done;
-                        }
-                    }
                     uint32_t orig_oparg = oparg;  // For OPARG_TOP/BOTTOM
                     for (int i = 0; i < nuops; i++) {
                         oparg = orig_oparg;
@@ -811,6 +803,19 @@ translate_bytecode_to_trace(
                         }
 
                         if (uop == _POP_FRAME || uop == _RETURN_GENERATOR || uop == _YIELD_VALUE) {
+                            if (trace_stack_depth == 0) {
+                                DPRINTF(2, "Trace stack underflow\n");
+                                OPT_STAT_INC(trace_stack_underflow);
+                                ADD_TO_TRACE(uop, oparg, 0, target);
+                                if (uop == _POP_FRAME || uop == _RETURN_GENERATOR) {
+                                    ADD_TO_TRACE(_RETURN_OFFSET, 0, 0, 0);
+                                }
+                                else {
+                                    ADD_TO_TRACE(_YIELD_OFFSET, 0, 0, 0);
+                                }
+                                ADD_TO_TRACE(_DYNAMIC_EXIT, 0, 0, 0);
+                                goto done;
+                            }
                             TRACE_STACK_POP();
                             /* Set the operand to the function or code object returned to,
                              * to assist optimization passes. (See _PUSH_FRAME below.)

Python/ceval.c

@@ -641,7 +641,9 @@ dummy_func(void) {
     op(_POP_FRAME, (retval -- res)) {
         SYNC_SP();
         ctx->frame->stack_pointer = stack_pointer;
-        frame_pop(ctx);
+        if (frame_pop(ctx)) {
+            goto done;
+        }
         stack_pointer = ctx->frame->stack_pointer;
         res = retval;
 
@@ -663,7 +665,9 @@ dummy_func(void) {
     op(_RETURN_GENERATOR, ( -- res)) {
         SYNC_SP();
         ctx->frame->stack_pointer = stack_pointer;
-        frame_pop(ctx);
+        if (frame_pop(ctx)) {
+            goto done;
+        }
         stack_pointer = ctx->frame->stack_pointer;
         OUT_OF_SPACE_IF_NULL(res = sym_new_unknown(ctx));
 

Python/executor_cases.c.h

@@ -384,9 +384,11 @@ _Py_uop_frame_pop(_Py_UOpsContext *ctx)
     _Py_UOpsAbstractFrame *frame = ctx->frame;
     ctx->n_consumed = frame->locals;
     ctx->curr_frame_depth--;
-    assert(ctx->curr_frame_depth >= 1);
+    if (ctx->curr_frame_depth == 0) {
+        ctx->frame = NULL;
+        return -1;
+    }
     ctx->frame = &ctx->frames[ctx->curr_frame_depth - 1];
-
     return 0;
 }
 

Python/generated_cases.c.h

@@ -2357,7 +2357,8 @@ _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr, int oparg)
     }
     else if (tp == &PyGen_Type && oparg <= SHRT_MAX) {
         assert(instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == END_FOR  ||
-            instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == INSTRUMENTED_END_FOR
+            instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == INSTRUMENTED_END_FOR ||
+            instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == ENTER_EXECUTOR
         );
         if (_PyInterpreterState_GET()->eval_frame) {
             SPECIALIZATION_FAIL(FOR_ITER, SPEC_FAIL_OTHER);