Add ability for external executor

Author: DinoV

Include/cpython/pyframe.h

@@ -40,6 +40,7 @@ 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
-#define PyUnstable_EXECUTABLE_KINDS 5
+#define PyUnstable_EXECUTABLE_KIND_JIT 5
+#define PyUnstable_EXECUTABLE_KINDS 6
 
 PyAPI_DATA(const PyTypeObject *) const PyUnstable_ExecutableKinds[PyUnstable_EXECUTABLE_KINDS+1];

Include/internal/pycore_genobject.h

@@ -8,15 +8,15 @@ extern "C" {
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
-#include "pycore_interpframe_structs.h" // _PyGenObject
+#include "pycore_interpframe_structs.h" // _PyInterpreterFrame
 
 #include <stddef.h>               // offsetof()
 
 
 static inline
 PyGenObject *_PyGen_GetGeneratorFromFrame(_PyInterpreterFrame *frame)
 {
-    assert(frame->owner == FRAME_OWNED_BY_GENERATOR);
+    assert(frame->owner & FRAME_OWNED_BY_GENERATOR);
     size_t offset_in_gen = offsetof(PyGenObject, gi_iframe);
     return (PyGenObject *)(((char *)frame) - offset_in_gen);
 }

Include/internal/pycore_interpframe.h

@@ -17,9 +17,37 @@ extern "C" {
 #define _PyInterpreterFrame_LASTI(IF) \
     ((int)((IF)->instr_ptr - _PyFrame_GetBytecode((IF))))
 
+PyAPI_DATA(PyTypeObject) PyUnstable_ExternalExecutable_Type;
+
+#define PyUnstable_ExternalExecutable_Check(op) Py_IS_TYPE((op), &PyUnstable_ExternalExecutable_Type)
+
+// Initialize a potentially external frame and make it safe to access the
+// all of the members of the returned _PyInterpreterFrame. The returned
+// value will be the same address as the passed in pointer.
+PyAPI_FUNC(void) _PyFrame_InitializeExternalFrame(_PyInterpreterFrame *frame);
+
+PyAPI_FUNC(PyObject *) PyUnstable_MakeExternalExecutable(_PyFrame_Reifier reifier, PyCodeObject *code, PyObject *state);
+
+static bool _PyFrame_IsExternalFrame(_PyInterpreterFrame *frame)
+{
+    return frame->owner & FRAME_OWNED_EXTERNALLY;
+}
+
+static inline void
+_PyFrame_EnsureFrameFullyInitialized(_PyInterpreterFrame *frame)
+{
+    if (_PyFrame_IsExternalFrame(frame)) {
+        _PyFrame_InitializeExternalFrame(frame);
+    }
+}
+
 static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
     assert(!PyStackRef_IsNull(f->f_executable));
     PyObject *executable = PyStackRef_AsPyObjectBorrow(f->f_executable);
+    if (f->owner & FRAME_OWNED_EXTERNALLY) {
+        assert(PyUnstable_ExternalExecutable_Check(executable));
+        return ((PyUnstable_PyExternalExecutable *)executable)->ef_code;
+    }
     assert(PyCode_Check(executable));
     return (PyCodeObject *)executable;
 }
@@ -30,12 +58,6 @@ static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
 static inline PyCodeObject* _Py_NO_SANITIZE_THREAD
 _PyFrame_SafeGetCode(_PyInterpreterFrame *f)
 {
-    // globals and builtins may be NULL on a legit frame, but it's unlikely.
-    // It's more likely that it's a sign of an invalid frame.
-    if (f->f_globals == NULL || f->f_builtins == NULL) {
-        return NULL;
-    }
-
     if (PyStackRef_IsNull(f->f_executable)) {
         return NULL;
     }
@@ -48,6 +70,18 @@ _PyFrame_SafeGetCode(_PyInterpreterFrame *f)
     if (_PyObject_IsFreed(executable)) {
         return NULL;
     }
+    if (_PyFrame_IsExternalFrame(f)) {
+        executable = (PyObject *)((PyUnstable_PyExternalExecutable *)executable)->ef_code;
+        if (_PyObject_IsFreed(executable)) {
+            return NULL;
+        }
+    } else {
+        // globals and builtins may be NULL on a legit frame, but it's unlikely.
+        // It's more likely that it's a sign of an invalid frame.
+        if (f->f_globals == NULL || f->f_builtins == NULL) {
+            return NULL;
+        }
+    }
     if (!PyCode_Check(executable)) {
         return NULL;
     }
@@ -59,6 +93,7 @@ _PyFrame_GetBytecode(_PyInterpreterFrame *f)
 {
 #ifdef Py_GIL_DISABLED
     PyCodeObject *co = _PyFrame_GetCode(f);
+    _PyFrame_EnsureFrameFullyInitialized(f);
     _PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
     assert(f->tlbc_index >= 0 && f->tlbc_index < tlbc->size);
     return (_Py_CODEUNIT *)tlbc->entries[f->tlbc_index];
@@ -81,6 +116,7 @@ _PyFrame_SafeGetLasti(struct _PyInterpreterFrame *f)
     }
 
     _Py_CODEUNIT *bytecode;
+    _PyFrame_EnsureFrameFullyInitialized(f);
 #ifdef Py_GIL_DISABLED
     _PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
     assert(f->tlbc_index >= 0 && f->tlbc_index < tlbc->size);
@@ -256,10 +292,11 @@ _PyFrame_IsIncomplete(_PyInterpreterFrame *frame)
 {
     if (frame->owner >= FRAME_OWNED_BY_INTERPRETER) {
         return true;
+    } else if (frame->owner & (FRAME_OWNED_BY_GENERATOR|FRAME_OWNED_EXTERNALLY)) {
+        return false;
     }
-    return frame->owner != FRAME_OWNED_BY_GENERATOR &&
-           frame->instr_ptr < _PyFrame_GetBytecode(frame) +
-                                  _PyFrame_GetCode(frame)->_co_firsttraceable;
+    return frame->instr_ptr < _PyFrame_GetBytecode(frame) +
+           _PyFrame_GetCode(frame)->_co_firsttraceable;
 }
 
 static inline _PyInterpreterFrame *
@@ -271,12 +308,66 @@ _PyFrame_GetFirstComplete(_PyInterpreterFrame *frame)
     return frame;
 }
 
+#if Py_DEBUG
+
+static inline bool _Py_NO_SANITIZE_THREAD
+_PyFrame_IsIncompleteOrUninitialized(_PyInterpreterFrame *frame)
+{
+    if (frame->owner >= FRAME_OWNED_BY_INTERPRETER || _PyFrame_IsExternalFrame(frame)) {
+        return true;
+    }
+    return !(frame->owner & FRAME_OWNED_BY_GENERATOR) &&
+           frame->instr_ptr < _PyFrame_GetBytecode(frame) +
+                             _PyFrame_GetCode(frame)->_co_firsttraceable;
+}
+
+static inline _PyInterpreterFrame *
+_PyFrame_GetFirstCompleteInitialized(_PyInterpreterFrame *frame)
+{
+    while (frame && _PyFrame_IsIncompleteOrUninitialized(frame)) {
+        frame = frame->previous;
+    }
+    return frame;
+}
+
+#endif
+
+static inline bool
+_PyFrame_StackpointerSaved(void)
+{
+#if Py_DEBUG
+    PyThreadState *tstate = PyThreadState_GET();
+    return _PyFrame_GetFirstCompleteInitialized(tstate->current_frame) == NULL ||
+           _PyFrame_GetFirstCompleteInitialized(tstate->current_frame)->stackpointer != NULL;
+#else
+    return true;
+#endif
+}
+
+
+
 static inline _PyInterpreterFrame *
 _PyThreadState_GetFrame(PyThreadState *tstate)
 {
     return _PyFrame_GetFirstComplete(tstate->current_frame);
 }
 
+static inline PyObject *
+_PyFrame_GetGlobals(_PyInterpreterFrame *frame) {
+    if (frame->f_globals == NULL) {
+        frame->f_globals = _PyFrame_GetFunction(frame)->func_globals;
+    }
+    return frame->f_globals;
+}
+
+static inline PyObject *
+_PyFrame_GetBuiltins(_PyInterpreterFrame *frame) {
+    if (frame->f_builtins == NULL) {
+        frame->f_builtins = _PyFrame_GetFunction(frame)->func_builtins;
+    }
+    return frame->f_builtins;
+}
+
 /* For use by _PyFrame_GetFrameObject
   Do not call directly. */
 PyAPI_FUNC(PyFrameObject *)
@@ -288,9 +379,9 @@ _PyFrame_MakeAndSetFrameObject(_PyInterpreterFrame *frame);
 static inline PyFrameObject *
 _PyFrame_GetFrameObject(_PyInterpreterFrame *frame)
 {
-
+    _PyFrame_EnsureFrameFullyInitialized(frame);
     assert(!_PyFrame_IsIncomplete(frame));
-    PyFrameObject *res =  frame->frame_obj;
+    PyFrameObject *res = frame->frame_obj;
     if (res != NULL) {
         return res;
     }
@@ -309,7 +400,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
@@ -358,7 +449,7 @@ _PyFrame_PushUnchecked(PyThreadState *tstate, _PyStackRef func, int null_locals_
 /* Pushes a trampoline frame without checking for space.
  * Must be guarded by _PyThreadState_HasStackSpace() */
 static inline _PyInterpreterFrame *
-_PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int stackdepth, _PyInterpreterFrame * previous)
+_PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int stackdepth, _PyInterpreterFrame *previous)
 {
     CALL_STAT_INC(frames_pushed);
     _PyInterpreterFrame *frame = (_PyInterpreterFrame *)tstate->datastack_top;

Include/internal/pycore_interpframe_structs.h

@@ -20,10 +20,20 @@ extern "C" {
 #endif
 
 enum _frameowner {
-    FRAME_OWNED_BY_THREAD = 0,
-    FRAME_OWNED_BY_GENERATOR = 1,
-    FRAME_OWNED_BY_FRAME_OBJECT = 2,
-    FRAME_OWNED_BY_INTERPRETER = 3,
+    // The frame is allocated on per-thread memory that will be freed or transferred when
+    // the frame unwinds.
+    FRAME_OWNED_BY_THREAD = 0x00,
+    // The frame is allocated in a generator and may out-live the execution.
+    FRAME_OWNED_BY_GENERATOR = 0x01,
+    // A flag which indicates the frame is owned externally. May be combined with
+    // FRAME_OWNED_BY_THREAD or FRAME_OWNED_BY_GENERATOR. The frame may only have
+    // _PyInterpreterFrameFields. To access other fields and ensure they are up to
+    // date _PyFrame_EnsureFrameFullyInitialized must be called first.
+    FRAME_OWNED_EXTERNALLY = 0x02,
+    // The frame is owned by the frame object (indicating the frame has unwound).
+    FRAME_OWNED_BY_FRAME_OBJECT = 0x04,
+    // The frame is a sentinel frame for entry to the interpreter loop
+    FRAME_OWNED_BY_INTERPRETER = 0x08,
 };
 
 struct _PyInterpreterFrame {
@@ -85,6 +95,15 @@ struct _PyAsyncGenObject {
     _PyGenObject_HEAD(ag)
 };
 
+typedef void (*_PyFrame_Reifier)(struct _PyInterpreterFrame *, PyObject *reifier);
+
+typedef struct {
+    PyObject_HEAD
+    PyCodeObject *ef_code;
+    PyObject *ef_state;
+    _PyFrame_Reifier ef_reifier;
+} PyUnstable_PyExternalExecutable;
+
 #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
@@ -2869,6 +2870,61 @@ 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_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):