_DYNAMIC_EXIT on underflow

brandtbucher · Jul 31, 2024 · b523cd2 · b523cd2
1 parent 7797182
commit b523cd2
Show file tree

Hide file tree

Showing 9 changed files with 128 additions and 43 deletions.
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -941,6 +941,9 @@ dummy_func(
             _PyEval_FrameClearAndPop(tstate, dying);
             LOAD_SP();
             LOAD_IP(frame->return_offset);
+        #if TIER_TWO
+            frame->instr_ptr += frame->return_offset;
+        #endif
             res = retval;
             LLTRACE_RESUME_FRAME();
         }
@@ -1191,6 +1194,9 @@ dummy_func(
                    _PyOpcode_Deopt[frame->instr_ptr->op.code] == ENTER_EXECUTOR);
             #endif
             LOAD_IP(1 + INLINE_CACHE_ENTRIES_SEND);
+        #if TIER_TWO
+            frame->instr_ptr += 1 + INLINE_CACHE_ENTRIES_SEND;
+        #endif
             LOAD_SP();
             value = retval;
             LLTRACE_RESUME_FRAME();
@@ -2829,7 +2835,8 @@ dummy_func(
                 }
                 /* iterator ended normally */
                 assert(next_instr[oparg].op.code == END_FOR ||
-                       next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
+                       next_instr[oparg].op.code == INSTRUMENTED_END_FOR ||
+                       next_instr[oparg].op.code == ENTER_EXECUTOR);
                 PyStackRef_CLOSE(iter);
                 STACK_SHRINK(1);
                 /* Jump forward oparg, then skip following END_FOR and POP_TOP instruction */
@@ -2881,7 +2888,8 @@ dummy_func(
                 }
                 /* iterator ended normally */
                 assert(next_instr[oparg].op.code == END_FOR ||
-                       next_instr[oparg].op.code == INSTRUMENTED_END_FOR);
+                       next_instr[oparg].op.code == INSTRUMENTED_END_FOR ||
+                       next_instr[oparg].op.code == ENTER_EXECUTOR);
                 STACK_SHRINK(1);
                 PyStackRef_CLOSE(iter_stackref);
                 /* Skip END_FOR and POP_TOP */
@@ -3566,6 +3574,9 @@ dummy_func(
             tstate->py_recursion_remaining--;
             LOAD_SP();
             LOAD_IP(0);
+        #if TIER_TWO
+            frame->instr_ptr += 0;
+        #endif
             LLTRACE_RESUME_FRAME();
         }
 
@@ -4349,6 +4360,9 @@ dummy_func(
             _PyThreadState_PopFrame(tstate, frame);
             frame = tstate->current_frame = prev;
             LOAD_IP(frame->return_offset);
+        #if TIER_TWO
+            frame->instr_ptr += frame->return_offset;
+        #endif
             LOAD_SP();
             LLTRACE_RESUME_FRAME();
         }

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
diff --git a/Python/optimizer.c b/Python/optimizer.c
@@ -178,6 +178,11 @@ _PyOptimizer_Optimize(
     _PyInterpreterFrame *frame, _Py_CODEUNIT *start,
     _PyStackRef *stack_pointer, _PyExecutorObject **executor_ptr)
 {
+    if (!PyCode_Check(frame->f_executable) ||
+        !PyFunction_Check(frame->f_funcobj))
+    {
+        return 0;
+    }
     PyCodeObject *code = _PyFrame_GetCode(frame);
     assert(PyCode_Check(code));
     PyInterpreterState *interp = _PyInterpreterState_GET();
@@ -506,19 +511,25 @@ add_to_trace(
         return 0; \
     } \
     assert(func == NULL || func->func_code == (PyObject *)code); \
+    assert(code != NULL); \
     trace_stack[trace_stack_depth].func = func; \
     trace_stack[trace_stack_depth].code = code; \
     trace_stack[trace_stack_depth].instr = instr; \
     trace_stack_depth++;
 #define TRACE_STACK_POP() \
     if (trace_stack_depth <= 0) { \
-        Py_FatalError("Trace stack underflow\n"); \
+        func = NULL; \
+        code = NULL; \
+        instr = NULL; \
     } \
-    trace_stack_depth--; \
-    func = trace_stack[trace_stack_depth].func; \
-    code = trace_stack[trace_stack_depth].code; \
-    assert(func == NULL || func->func_code == (PyObject *)code); \
-    instr = trace_stack[trace_stack_depth].instr;
+    else { \
+        trace_stack_depth--; \
+        func = trace_stack[trace_stack_depth].func; \
+        code = trace_stack[trace_stack_depth].code; \
+        assert(func == NULL || func->func_code == (PyObject *)code); \
+        assert(code != NULL); \
+        instr = trace_stack[trace_stack_depth].instr; \
+    }
 
 /* Returns the length of the trace on success,
  * 0 if it failed to produce a worthwhile trace,
@@ -704,17 +715,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 == _RETURN_VALUE || 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 + _RETURN_VALUE.
-                        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;
@@ -779,7 +779,9 @@ translate_bytecode_to_trace(
                                 operand = (uintptr_t)code | 1;
                             }
                             else {
-                                operand = 0;
+                                ADD_TO_TRACE(uop, oparg, 0, target);
+                                ADD_TO_TRACE(_DYNAMIC_EXIT, 0, 0, 0);
+                                goto done;
                             }
                             ADD_TO_TRACE(uop, oparg, operand, target);
                             DPRINTF(2,
@@ -847,11 +849,9 @@ translate_bytecode_to_trace(
                                 if (new_func != NULL) {
                                     operand = (uintptr_t)new_func;
                                 }
-                                else if (new_code != NULL) {
-                                    operand = (uintptr_t)new_code | 1;
-                                }
                                 else {
-                                    operand = 0;
+                                    assert(new_code != NULL);
+                                    operand = (uintptr_t)new_code | 1;
                                 }
                                 ADD_TO_TRACE(uop, oparg, operand, target);
                                 code = new_code;
@@ -909,7 +909,8 @@ translate_bytecode_to_trace(
     while (trace_stack_depth > 0) {
         TRACE_STACK_POP();
     }
-    assert(code == initial_code);
+    assert(code == NULL || code == initial_code);
+    code = initial_code;
     // Skip short traces where we can't even translate a single instruction:
     if (progress_needed) {
         OPT_STAT_INC(trace_too_short);

diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
@@ -448,10 +448,12 @@ optimize_uops(
                 DPRINTF(1, "\nUnknown opcode in abstract interpreter\n");
                 Py_UNREACHABLE();
         }
-        assert(ctx->frame != NULL);
-        DPRINTF(3, " stack_level %d\n", STACK_LEVEL());
-        ctx->frame->stack_pointer = stack_pointer;
-        assert(STACK_LEVEL() >= 0);
+        // assert(ctx->frame != NULL);
+        if (ctx->frame) {
+            DPRINTF(3, " stack_level %d\n", STACK_LEVEL());
+            ctx->frame->stack_pointer = stack_pointer;
+            assert(STACK_LEVEL() >= 0);
+        }
     }
     if (ctx->out_of_space) {
         DPRINTF(3, "\n");