Add patch '18-lightweight-frame-support' to third-party/python/3.14

Summary:
This diff adds a patch to the `third-party/python/3.14` Meta-internal fork.

Add lightweight frame support to CPython

Reviewed By: alexmalyshev

Differential Revision: D87386838

fbshipit-source-id: a1eb39c79d31d7db23f26b3ca91ba1c0259e316d

Author: DinoV

Include/cpython/pyframe.h

@@ -40,6 +40,12 @@ PyAPI_FUNC(int) PyUnstable_InterpreterFrame_GetLine(struct _PyInterpreterFrame *
 #define PyUnstable_EXECUTABLE_KIND_PY_FUNCTION 1
 #define PyUnstable_EXECUTABLE_KIND_BUILTIN_FUNCTION 3
 #define PyUnstable_EXECUTABLE_KIND_METHOD_DESCRIPTOR 4
+
+#ifdef META_PYTHON
+#define PyUnstable_EXECUTABLE_KIND_JIT 5
+#define PyUnstable_EXECUTABLE_KINDS 6
+#else
 #define PyUnstable_EXECUTABLE_KINDS 5
+#endif
 
 PyAPI_DATA(const PyTypeObject *) const PyUnstable_ExecutableKinds[PyUnstable_EXECUTABLE_KINDS+1];

Include/internal/pycore_interpframe.h

@@ -17,9 +17,35 @@ extern "C" {
 #define _PyInterpreterFrame_LASTI(IF) \
     ((int)((IF)->instr_ptr - _PyFrame_GetBytecode((IF))))
 
+#ifdef META_PYTHON
+
+__attribute__((weak)) PyAPI_DATA(PyTypeObject) PyUnstable_JITExecutable_Type;
+
+#define PyUnstable_JITExecutable_Check(op) Py_IS_TYPE((op), &PyUnstable_JITExecutable_Type)
+
+PyAPI_FUNC(PyObject *) PyUnstable_MakeJITExecutable(_PyFrame_Reifier reifier, PyCodeObject *code, PyObject *state);
+
+PyAPI_FUNC(int) _PyFrame_InitializeExternalFrame(_PyInterpreterFrame *frame);
+
+static inline int
+_PyFrame_EnsureFrameFullyInitialized(_PyInterpreterFrame *frame)
+{
+    if (PyUnstable_JITExecutable_Check(PyStackRef_AsPyObjectBorrow(frame->f_executable))) {
+        return _PyFrame_InitializeExternalFrame(frame);
+    }
+    return 0;
+}
+
+#endif
+
 static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
     assert(!PyStackRef_IsNull(f->f_executable));
     PyObject *executable = PyStackRef_AsPyObjectBorrow(f->f_executable);
+#ifdef META_PYTHON
+    if (PyUnstable_JITExecutable_Check(executable)) {
+        return ((PyUnstable_PyJitExecutable *)executable)->je_code;
+    }
+#endif
     assert(PyCode_Check(executable));
     return (PyCodeObject *)executable;
 }
@@ -48,6 +74,14 @@ _PyFrame_SafeGetCode(_PyInterpreterFrame *f)
     if (_PyObject_IsFreed(executable)) {
         return NULL;
     }
+#ifdef META_PYTHON
+    if (PyUnstable_JITExecutable_Check(executable)) {
+        executable = (PyObject *)((PyUnstable_PyJitExecutable *)executable)->je_code;
+        if (_PyObject_IsFreed(executable)) {
+            return NULL;
+        }
+    }
+#endif
     if (!PyCode_Check(executable)) {
         return NULL;
     }
@@ -58,6 +92,12 @@ static inline _Py_CODEUNIT *
 _PyFrame_GetBytecode(_PyInterpreterFrame *f)
 {
 #ifdef Py_GIL_DISABLED
+#ifdef META_PYTHON
+    if (_PyFrame_EnsureFrameFullyInitialized(f) < 0) {
+        return NULL;
+    }
+#endif
+
     PyCodeObject *co = _PyFrame_GetCode(f);
     _PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
     assert(f->tlbc_index >= 0 && f->tlbc_index < tlbc->size);
@@ -80,6 +120,12 @@ _PyFrame_SafeGetLasti(struct _PyInterpreterFrame *f)
         return -1;
     }
 
+#ifdef META_PYTHON
+    if (_PyFrame_EnsureFrameFullyInitialized(f) < 0) {
+        return -1;
+    }
+#endif
+
     _Py_CODEUNIT *bytecode;
 #ifdef Py_GIL_DISABLED
     _PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
@@ -277,6 +323,26 @@ _PyThreadState_GetFrame(PyThreadState *tstate)
     return _PyFrame_GetFirstComplete(tstate->current_frame);
 }
 
+static inline PyObject *
+_PyFrame_GetGlobals(_PyInterpreterFrame *frame) {
+#ifdef META_PYTHON
+    if (_PyFrame_EnsureFrameFullyInitialized(frame) < 0) {
+        return NULL;
+    }
+#endif
+    return frame->f_globals;
+}
+
+static inline PyObject *
+_PyFrame_GetBuiltins(_PyInterpreterFrame *frame) {
+#ifdef META_PYTHON
+    if (_PyFrame_EnsureFrameFullyInitialized(frame) < 0) {
+        return NULL;
+    }
+#endif
+    return frame->f_builtins;
+}
+
 /* For use by _PyFrame_GetFrameObject
   Do not call directly. */
 PyFrameObject *
@@ -288,6 +354,11 @@ _PyFrame_MakeAndSetFrameObject(_PyInterpreterFrame *frame);
 static inline PyFrameObject *
 _PyFrame_GetFrameObject(_PyInterpreterFrame *frame)
 {
+#ifdef META_PYTHON
+    if (PyUnstable_JITExecutable_Check(PyStackRef_AsPyObjectBorrow(frame->f_executable))) {
+        return _PyFrame_MakeAndSetFrameObject(frame);
+    }
+#endif
 
     assert(!_PyFrame_IsIncomplete(frame));
     PyFrameObject *res =  frame->frame_obj;
@@ -309,7 +380,7 @@ _PyFrame_ClearLocals(_PyInterpreterFrame *frame);
  * take should  be set to 1 for heap allocated
  * frames like the ones in generators and coroutines.
  */
-void
+PyAPI_FUNC(void)
 _PyFrame_ClearExceptCode(_PyInterpreterFrame * frame);
 
 int

Include/internal/pycore_interpframe_structs.h

@@ -86,6 +86,19 @@ struct _PyAsyncGenObject {
     _PyGenObject_HEAD(ag)
 };
 
+#ifdef META_PYTHON
+
+typedef int (*_PyFrame_Reifier)(struct _PyInterpreterFrame *, PyObject *reifier);
+
+typedef struct {
+    PyObject_HEAD
+    PyCodeObject *je_code;
+    PyObject *je_state;
+    _PyFrame_Reifier je_reifier;
+} PyUnstable_PyJitExecutable;
+
+#endif
+
 #undef _PyGenObject_HEAD
 
 

Lib/test/test_capi/test_misc.py

@@ -4,6 +4,7 @@
 import _thread
 from collections import deque
 import contextlib
+import dis
 import importlib.machinery
 import importlib.util
 import json
@@ -2794,6 +2795,85 @@ def func():
         names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
         self.do_test(func, names)
 
+    def test_jit_frame(self):
+        def fakefunc():
+            pass
+
+        def f():
+            return sys._getframe(1)
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, f, ())
+
+    def test_jit_frame_globals(self):
+        """jit executable can fill in globals when accessed"""
+        def fakefunc():
+            pass
+
+        fake_globals = {"abc":42}
+        def callback():
+            return {"globals": fake_globals}
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, globals, (), callback)
+        self.assertEqual(res, fake_globals)
+
+    def test_jit_frame_builtins(self):
+        """jit executable can fill in builtins when accessed"""
+        def fakefunc():
+            pass
+
+        fake_builtins = {"abc":42}
+        def callback():
+            return {"builtins": fake_builtins}
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, _testlimitedcapi.eval_getbuiltins, (), callback)
+        self.assertEqual(res, fake_builtins)
+
+    def test_jit_frame_instr_ptr(self):
+        """jit executable can fill in the instr ptr each time the frame is queried"""
+        def fakefunc():
+            pass
+            pass
+            pass
+            pass
+
+        offset = 0
+        linenos = []
+        def test():
+            for op in dis.get_instructions(fakefunc):
+                if op.opname in ("RESUME", "NOP", "RETURN_VALUE"):
+                    nonlocal offset
+                    offset = op.offset//2
+                    linenos.append(sys._getframe(1).f_lineno)
+
+        def callback():
+            return {"instr_ptr": offset}
+
+        _testinternalcapi.call_with_jit_frame(fakefunc, test, (), callback)
+        base = fakefunc.__code__.co_firstlineno
+        self.assertEqual(linenos, [base, base + 1, base + 2, base + 3, base + 4])
+
+    def test_jit_frame_code(self):
+        """internal C api checks the for a code executor"""
+        def fakefunc():
+            pass
+
+        def callback():
+            return _testinternalcapi.iframe_getcode(sys._getframe(1))
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, callback, ())
+        self.assertEqual(res, fakefunc.__code__)
+
+    def test_jit_frame_line(self):
+        """internal C api checks the for a code executor"""
+        def fakefunc():
+            pass
+
+        def callback():
+            return _testinternalcapi.iframe_getline(sys._getframe(1))
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, callback, ())
+        self.assertEqual(res, fakefunc.__code__.co_firstlineno)
+
 
 @unittest.skipUnless(support.Py_GIL_DISABLED, 'need Py_GIL_DISABLED')
 class TestPyThreadId(unittest.TestCase):

Modules/_testinternalcapi.c

@@ -28,6 +28,7 @@
 #include "pycore_initconfig.h"    // _Py_GetConfigsAsDict()
 #include "pycore_instruction_sequence.h"  // _PyInstructionSequence_New()
 #include "pycore_interpframe.h"   // _PyFrame_GetFunction()
+#include "pycore_interpframe_structs.h" // _PyInterpreterFrame
 #include "pycore_object.h"        // _PyObject_IsFreed()
 #include "pycore_optimizer.h"     // _Py_Executor_DependsOn
 #include "pycore_pathconfig.h"    // _PyPathConfig_ClearGlobal()
@@ -692,6 +693,117 @@ set_eval_frame_record(PyObject *self, PyObject *list)
     Py_RETURN_NONE;
 }
 
+#ifdef META_PYTHON
+
+typedef struct {
+    bool initialized;
+    _PyInterpreterFrame frame;
+} JitFrame;
+
+int
+reifier(_PyInterpreterFrame *frame, PyObject *executable)
+{
+    JitFrame *jitframe = (JitFrame*)((char *)frame - offsetof(JitFrame, frame));
+    PyFunctionObject *func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
+    if (!jitframe->initialized) {
+        jitframe->initialized = true;
+        frame->f_locals = NULL;
+        frame->f_globals = func->func_globals; // borrowed
+        frame->f_builtins = func->func_builtins; // borrowed
+        frame->frame_obj = NULL;
+#ifdef Py_GIL_DISABLED
+        frame->tlbc_index = 0;
+#endif
+    }
+    PyUnstable_PyJitExecutable *jit_exec = (PyUnstable_PyJitExecutable*)executable;
+    if (jit_exec->je_state == NULL) {
+        return 0;
+    }
+
+    PyObject *res = PyObject_CallNoArgs(jit_exec->je_state);
+    if (res == NULL) {
+        return -1;
+    }
+
+    // let the test-state function fill in details on the frame
+    if (PyDict_Check(res)) {
+        PyObject *globals = PyDict_GetItemString(res, "globals");
+        if (globals != NULL) {
+            frame->f_globals = globals;
+        }
+        PyObject *builtins = PyDict_GetItemString(res, "builtins");
+        if (builtins != NULL) {
+            frame->f_builtins = builtins;
+        }
+        PyObject *instr_ptr = PyDict_GetItemString(res, "instr_ptr");
+        if (instr_ptr != NULL) {
+            frame->instr_ptr = _PyCode_CODE((PyCodeObject *)func->func_code) +
+                                PyLong_AsLong(instr_ptr);
+        }
+    }
+    Py_DECREF(res);
+    return 0;
+}
+
+static PyObject *
+call_with_jit_frame(PyObject *self, PyObject *args)
+{
+    PyObject *fakefunc; // used for f_funcobj as-if we were that JITed function
+    PyObject *call;     // the thing to call for testing purposes
+    PyObject *callargs; // the arguments to provide for the test call
+    PyObject *state = NULL; // a state object provided to the reifier, for tests we
+                            // callback on it to populate fields.
+    if (!PyArg_ParseTuple(args, "OOO|O", &fakefunc, &call, &callargs, &state)) {
+        return NULL;
+    }
+    if (!PyTuple_Check(callargs)) {
+        PyErr_SetString(PyExc_TypeError, "callargs must be a tuple");
+        return NULL;
+    }
+
+    PyThreadState *tstate = PyThreadState_Get();
+    PyCodeObject *code = (PyCodeObject *)((PyFunctionObject *)fakefunc)->func_code;
+    PyObject *executable = PyUnstable_MakeJITExecutable(reifier, code, state);
+    if (executable == NULL) {
+        return NULL;
+    }
+
+    // Create JIT frame and push onto the _PyInterprerFrame stack.
+    JitFrame frame;
+    frame.initialized = false;
+    // Initialize minimal set of fields
+    frame.frame.previous = tstate->current_frame;
+    frame.frame.f_executable = PyStackRef_FromPyObjectSteal(executable);
+    frame.frame.f_funcobj = PyStackRef_FromPyObjectNew(fakefunc);
+    frame.frame.instr_ptr = _PyCode_CODE(code) + code->_co_firsttraceable;
+    frame.frame.stackpointer = &frame.frame.localsplus[0];
+    frame.frame.owner = FRAME_OWNED_BY_THREAD;
+    tstate->current_frame = &frame.frame;
+
+    // call the test function
+    PyObject *res = PyObject_Call(call, callargs, NULL);
+
+    tstate->current_frame = frame.frame.previous;
+    // the test function may have caused the frame to get reified.
+    if (frame.initialized && frame.frame.frame_obj != NULL) {
+        // remove our reifier
+        PyStackRef_CLOSE(frame.frame.f_executable);
+        frame.frame.f_executable = PyStackRef_FromPyObjectNew(code);
+
+        // Transfer ownership to the reified frame object
+        _PyFrame_ClearExceptCode(&frame.frame);
+        PyStackRef_CLOSE(frame.frame.f_executable);
+    }
+    else {
+        // Pop frame from the stack
+        PyStackRef_CLOSE(frame.frame.f_executable);
+        PyStackRef_CLOSE(frame.frame.f_funcobj);
+    }
+    return res;
+}
+
+#endif
+
 /*[clinic input]
 
 _testinternalcapi.compiler_cleandoc -> object
@@ -2479,6 +2591,9 @@ static PyMethodDef module_functions[] = {
     {"DecodeLocaleEx", decode_locale_ex, METH_VARARGS},
     {"set_eval_frame_default", set_eval_frame_default, METH_NOARGS, NULL},
     {"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
+#ifdef META_PYTHON
+    {"call_with_jit_frame", call_with_jit_frame, METH_VARARGS, NULL},
+#endif
     _TESTINTERNALCAPI_COMPILER_CLEANDOC_METHODDEF
     _TESTINTERNALCAPI_NEW_INSTRUCTION_SEQUENCE_METHODDEF
     _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF