python · markshannon · May 3, 2024
diff --git a/Include/internal/pycore_uop_ids.h b/Include/internal/pycore_uop_ids.h
diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.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 //

    }

@@ -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



diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
@@ -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.)

@@ -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));


diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_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;
 }


diff --git a/Python/specialize.c b/Python/specialize.c
@@ -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);