frame reconstruction

gvanrossum · Feb 20, 2024 · Feb 20, 2024 · Feb 20, 2024 · Feb 21, 2024 · Feb 21, 2024
commit 7b69a6274598a1a2d6b706f796c33517353769b9
diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
@@ -289,13 +289,14 @@ static inline int
 _PyFrame_ConvertToTier2(PyThreadState *tstate, _PyInterpreterFrame *frame,
                         int localsplus_grow)
 {
-    if (frame->owner != FRAME_OWNED_BY_THREAD) {
-        return 1;
-    }
+    assert(localsplus_grow > 0);
     // Already grown previously
     if (frame->tier2_extra_size >= localsplus_grow) {
         return 0;
     }
+    if (frame->owner != FRAME_OWNED_BY_THREAD) {
+        return 1;
+    }
     if (!_PyThreadState_HasStackSpace(tstate, localsplus_grow)) {
         return 1;
     }

diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
@@ -11,7 +11,7 @@ extern "C" {
 #include "pycore_uop_ids.h"
 
 // This is the length of the trace we project initially.
-#define UOP_MAX_TRACE_LENGTH 512
+#define UOP_MAX_TRACE_LENGTH 1024
 
 #define TRACE_STACK_SIZE 5
 

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
diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
@@ -887,5 +887,24 @@ def testfunc(n):
         self.assertLessEqual(len(guard_both_float_count), 1)
         self.assertIn("_COMPARE_OP_STR", uops)
 
+    def test_function_inlining(self):
+        def testfunc(n):
+            a = 1
+            for _ in range(n):
+                x = foo(a, a)
+            return x
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertTrue(res)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertLessEqual(len(guard_both_float_count), 1)
+        self.assertIn("_COMPARE_OP_STR", uops)
+
+
+def foo(x, y):
+    return x + y
+
+
 if __name__ == "__main__":
     unittest.main()
@@ -4212,15 +4212,15 @@ dummy_func(
         // Inlining postlude
         op(_POST_INLINE, (reconstructer/4 -- retval)) {
             // clear the locals
-            PyObject **end = frame->localsplus + oparg;
             PyObject *ret = PEEK(1);
             stack_pointer--;
+            PyObject **end = stack_pointer - oparg;
             while (stack_pointer > end) {
                 Py_CLEAR(stack_pointer[-1]);
                 stack_pointer--;
             }
             retval = ret;
-            frame->frame_reconstruction_inst = ((int64_t)reconstructer == -1
+            frame->frame_reconstruction_inst = ((int64_t)reconstructer == 0
                 ? NULL
                 : current_executor->trace + (int64_t)reconstructer);
             CHECK_EVAL_BREAKER();
@@ -4230,6 +4230,10 @@ dummy_func(
             DEOPT_IF(_PyFrame_ConvertToTier2(tstate, frame, oparg));
         }
 
+        // Dummy instruction to indicate this is frame reconstruction data.
+        op(_RECONSTRUCT_FRAME_INFO, (--)) {
+        }
+
         // Sentinel for true end of trace.
         op(_TRUE_END, (--)) {}
 // END BYTECODES //

@@ -251,6 +251,8 @@ _PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
 static  _PyInterpreterFrame *
 _PyEvalFramePushAndInit_Ex(PyThreadState *tstate, PyFunctionObject *func,
     PyObject *locals, Py_ssize_t nargs, PyObject *callargs, PyObject *kwargs);
+static _PyInterpreterFrame *
+_PyEvalFrame_ReconstructTier2Frame(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject ***stackptr_ptr);
 
 #ifdef HAVE_ERRNO_H
 #include <errno.h>
@@ -1071,6 +1073,10 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     }
 #endif
     OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
+    frame = _PyEvalFrame_ReconstructTier2Frame(tstate, frame, &stack_pointer);
+    if (frame == NULL) {
+        goto resume_with_error;
+    }
     frame->return_offset = 0;  // Don't leave this random
     _PyFrame_SetStackPointer(frame, stack_pointer);
     Py_DECREF(current_executor);
@@ -1079,6 +1085,11 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 
 // Jump here from DEOPT_IF()
 deoptimize:
+    frame = _PyEvalFrame_ReconstructTier2Frame(tstate, frame, &stack_pointer);
+    // Unrecoverable memory error.
+    if (frame == NULL) {
+        goto error_tier_two;
+    }
     next_instr = next_uop[-1].target + _PyCode_CODE(_PyFrame_GetCode(frame));
 #ifdef Py_DEBUG
     if (lltrace >= 2) {
@@ -1671,8 +1682,9 @@ clear_thread_frame(PyThreadState *tstate, _PyInterpreterFrame * frame)
     assert(frame->owner == FRAME_OWNED_BY_THREAD);
     // Make sure that this is, indeed, the top frame. We can't check this in
     // _PyThreadState_PopFrame, since f_code is already cleared at that point:
-    assert((PyObject **)frame + _PyFrame_GetCode(frame)->co_framesize ==
-        tstate->datastack_top);
+    // This doesn't apply to tier 2 frames.
+    assert(frame->tier2_extra_size ==  0 ? (PyObject **)frame + _PyFrame_GetCode(frame)->co_framesize ==
+              tstate->datastack_top : 1);
     tstate->c_recursion_remaining--;
     assert(frame->frame_obj == NULL || frame->frame_obj->f_frame == frame);
     _PyFrame_ClearExceptCode(frame);
@@ -1786,6 +1798,118 @@ _PyEvalFramePushAndInit_Ex(PyThreadState *tstate, PyFunctionObject *func,
     return NULL;
 }
 
+// Tells the current frame how to reconstruct truly inlined function frames.
+// See optimizer_analysis.c for what each field represents.
+static _PyInterpreterFrame *
+_PyEvalFrame_ReconstructTier2Frame(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject ***stackptr_ptr)
+{
+    // Does not need reconstruction.
+    if (frame->frame_reconstruction_inst == NULL) {
+        return frame;
+    }
+#ifdef LLTRACE
+    printf("pre-reconstruction stack: \n");
+    dump_stack(frame, *stackptr_ptr);
+#endif
+    _PyInterpreterFrame *prev_frame = frame;
+    _PyInterpreterFrame *recentmost_frame = frame;
+    _PyUOpInstruction *curr = frame->frame_reconstruction_inst;
+    int opcode = curr->opcode;
+    while (opcode == _RECONSTRUCT_FRAME_INFO) {
+        // Hit the root frame.
+        if ((curr+1)->opcode != _RECONSTRUCT_FRAME_INFO) {
+            break;
+        }
+#ifdef LLTRACE
+        printf("reconstructing frame... \n");
+#endif
+        PyCodeObject* code = (PyCodeObject *)(uintptr_t)curr->operand;
+        assert(PyCode_Check(code));
+        assert((curr+1)->opcode == _RECONSTRUCT_FRAME_INFO);
+        assert(PyFunction_Check((PyObject*)(uintptr_t)(curr+1)->operand));
+        assert((curr+2)->opcode == _SAVE_RETURN_OFFSET);
+
+        // We must retrieve a cached function and code object because the user might have
+        // modified them since execution. Thus, to remain consistent and give the appearance
+        // that the frame has existed since before modification, we use a manual code object
+        // rather than obtaining the function's.
+        PyFunctionObject *callable = (PyFunctionObject *)(uintptr_t)((curr+1)->operand);
+        int code_flags = ((PyCodeObject*)code)->co_flags;
+        PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : Py_NewRef(PyFunction_GET_GLOBALS(callable));
+
+        _PyInterpreterFrame *new_frame = _PyThreadState_PushFrame(tstate, code->co_framesize);
+        if (new_frame == NULL) {
+            goto fail;
+        }
+
+        // TODO CONSUME callable from the stack to deal with refleak.
+        _PyFrame_Initialize(new_frame, (PyFunctionObject*)Py_NewRef(callable),
+                            locals, (PyCodeObject *)code,
+                            ((PyCodeObject *)code)->co_nlocalsplus);
+        new_frame->previous = prev_frame;
+        new_frame->return_offset = (curr+2)->oparg;
+        new_frame->instr_ptr = _PyCode_CODE(code) + (int)(curr+2)->operand;
+        prev_frame = new_frame;
+        // Copy over locals, stack and friends.
+#ifdef LLTRACE
+        printf("copying over stack with offset %d: , locals count: %d, stacksize: %d\n", curr->oparg, code->co_nlocalsplus, code->co_stacksize);
+        dump_stack(frame, frame->localsplus + curr->oparg);
+#endif
+        int total_len = (code->co_nlocalsplus + code->co_stacksize);
+        memcpy(new_frame->localsplus, frame->localsplus + curr->oparg,
+               sizeof(PyObject *) * total_len);
+
+#ifdef LLTRACE
+        printf("setting stacktop: %d + co_nlocalsplus\n", (curr+1)->oparg);
+#endif
+        // Finally, set the stack pointer
+        new_frame->stacktop = _PyFrame_GetCode(new_frame)->co_nlocalsplus + (curr+1)->oparg;
+        assert(new_frame->stacktop >= 0 || (int)(curr+1)->oparg < 0);
+
+//#ifdef LLTRACE
+//        if (!(((int16_t)(curr+2)->oparg) < 0)) {
+//            printf("the new frame %p has stack entries %d: \n", new_frame, (curr+2)->oparg);
+//            dump_stack(new_frame, &(new_frame->localsplus[new_frame->stacktop]));
+//        }
+//#endif
+        recentmost_frame = new_frame;
+        curr+=3;
+    }
+    PyObject **curr_stacklevel = *stackptr_ptr;
+    // Recentmost frame stack pointer is set by the current level.
+    int recentmost_stackentries = (int)(curr_stacklevel - (frame->localsplus + curr->oparg + (_PyFrame_GetCode(recentmost_frame)->co_nlocalsplus)));
+    *stackptr_ptr = recentmost_frame->localsplus + (_PyFrame_GetCode(recentmost_frame)->co_nlocalsplus) + recentmost_stackentries;
+#ifdef LLTRACE
+    printf("restoring offset %d\n", (int)(frame->instr_ptr - (_PyCode_CODE(_PyFrame_GetCode(frame)))));
+#endif
+    recentmost_frame->instr_ptr =  (_PyCode_CODE(_PyFrame_GetCode(recentmost_frame))) + (frame->instr_ptr - (_PyCode_CODE(_PyFrame_GetCode(frame))));
+    recentmost_frame->return_offset = -1;
+    recentmost_frame->stacktop = (*stackptr_ptr - recentmost_frame->localsplus);
+    // Set root frame stack pointer.
+    assert(curr->opcode == _RECONSTRUCT_FRAME_INFO);
+    assert((curr+1)->opcode == _SAVE_RETURN_OFFSET);
+
+    assert(curr->oparg >= 0);
+    frame->stacktop = _PyFrame_GetCode(frame)->co_nlocalsplus + curr->oparg;
+    frame->return_offset = (curr+1)->oparg;
+    frame->instr_ptr = _PyCode_CODE(_PyFrame_GetCode(frame)) + (int)(curr+1)->operand;
+    frame->f_names = Py_NewRef(_PyFrame_GetCode(frame)->co_names);
+    tstate->current_frame = recentmost_frame;
+    frame->frame_reconstruction_inst = NULL;
+
+#ifdef LLTRACE
+    printf("after reconstruction root stack, with n_stackentries %d: \n", curr->oparg);
+    dump_stack(frame, &(frame->localsplus[frame->stacktop]));
+    printf("after reconstruction topmost stack, with n_stackentries %d: \n", recentmost_stackentries);
+    dump_stack(recentmost_frame, *stackptr_ptr);
+#endif
+    return recentmost_frame;
+fail:
+    PyErr_NoMemory();
+    return NULL;
+}
+
+
 PyObject *
 _PyEval_Vector(PyThreadState *tstate, PyFunctionObject *func,
                PyObject *locals,

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h