🌐 AI搜索 & 代理 主页
Skip to content

Commit c9178d6

Browse files
DinoVDinoV
committed
Add new type of executable to support external JITs
1 parent fa44845 commit c9178d6

File tree

12 files changed

+369
-8
lines changed

12 files changed

+369
-8
lines changed

‎Include/cpython/pyframe.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ PyAPI_FUNC(int) PyUnstable_InterpreterFrame_GetLine(struct _PyInterpreterFrame *
4040
#define PyUnstable_EXECUTABLE_KIND_PY_FUNCTION 1
4141
#define PyUnstable_EXECUTABLE_KIND_BUILTIN_FUNCTION 3
4242
#define PyUnstable_EXECUTABLE_KIND_METHOD_DESCRIPTOR 4
43-
#define PyUnstable_EXECUTABLE_KINDS 5
43+
#define PyUnstable_EXECUTABLE_KIND_JIT 5
44+
#define PyUnstable_EXECUTABLE_KINDS 6
4445

4546
PyAPI_DATA(const PyTypeObject *) const PyUnstable_ExecutableKinds[PyUnstable_EXECUTABLE_KINDS+1];

‎Include/internal/pycore_interpframe.h‎

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,29 @@ extern "C" {
1717
#define _PyInterpreterFrame_LASTI(IF) \
1818
((int)((IF)->instr_ptr - _PyFrame_GetBytecode((IF))))
1919

20+
PyAPI_DATA(PyTypeObject) PyUnstable_JITExecutable_Type;
21+
22+
#define PyUnstable_JITExecutable_Check(op) Py_IS_TYPE((op), &PyUnstable_JITExecutable_Type)
23+
24+
PyAPI_FUNC(PyObject *) PyUnstable_MakeJITExecutable(_PyFrame_Reifier reifier, PyCodeObject *code, PyObject *state);
25+
26+
PyAPI_FUNC(int) _PyFrame_InitializeExternalFrame(_PyInterpreterFrame *frame);
27+
28+
static inline int
29+
_PyFrame_EnsureFrameFullyInitialized(_PyInterpreterFrame *frame)
30+
{
31+
if (PyUnstable_JITExecutable_Check(PyStackRef_AsPyObjectBorrow(frame->f_executable))) {
32+
return _PyFrame_InitializeExternalFrame(frame);
33+
}
34+
return 0;
35+
}
36+
2037
static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
2138
assert(!PyStackRef_IsNull(f->f_executable));
2239
PyObject *executable = PyStackRef_AsPyObjectBorrow(f->f_executable);
40+
if (PyUnstable_JITExecutable_Check(executable)) {
41+
return ((PyUnstable_PyJitExecutable *)executable)->je_code;
42+
}
2343
assert(PyCode_Check(executable));
2444
return (PyCodeObject *)executable;
2545
}
@@ -48,6 +68,12 @@ _PyFrame_SafeGetCode(_PyInterpreterFrame *f)
4868
if (_PyObject_IsFreed(executable)) {
4969
return NULL;
5070
}
71+
if (PyUnstable_JITExecutable_Check(executable)) {
72+
executable = (PyObject *)((PyUnstable_PyJitExecutable *)executable)->je_code;
73+
if (_PyObject_IsFreed(executable)) {
74+
return NULL;
75+
}
76+
}
5177
if (!PyCode_Check(executable)) {
5278
return NULL;
5379
}
@@ -58,6 +84,10 @@ static inline _Py_CODEUNIT *
5884
_PyFrame_GetBytecode(_PyInterpreterFrame *f)
5985
{
6086
#ifdef Py_GIL_DISABLED
87+
if (_PyFrame_EnsureFrameFullyInitialized(f) < 0) {
88+
return NULL;
89+
}
90+
6191
PyCodeObject *co = _PyFrame_GetCode(f);
6292
_PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
6393
assert(f->tlbc_index >= 0 && f->tlbc_index < tlbc->size);
@@ -80,6 +110,10 @@ _PyFrame_SafeGetLasti(struct _PyInterpreterFrame *f)
80110
return -1;
81111
}
82112

113+
if (_PyFrame_EnsureFrameFullyInitialized(f) < 0) {
114+
return -1;
115+
}
116+
83117
_Py_CODEUNIT *bytecode;
84118
#ifdef Py_GIL_DISABLED
85119
_PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
@@ -277,6 +311,22 @@ _PyThreadState_GetFrame(PyThreadState *tstate)
277311
return _PyFrame_GetFirstComplete(tstate->current_frame);
278312
}
279313

314+
static inline PyObject *
315+
_PyFrame_GetGlobals(_PyInterpreterFrame *frame) {
316+
if (_PyFrame_EnsureFrameFullyInitialized(frame) < 0) {
317+
return NULL;
318+
}
319+
return frame->f_globals;
320+
}
321+
322+
static inline PyObject *
323+
_PyFrame_GetBuiltins(_PyInterpreterFrame *frame) {
324+
if (_PyFrame_EnsureFrameFullyInitialized(frame) < 0) {
325+
return NULL;
326+
}
327+
return frame->f_builtins;
328+
}
329+
280330
/* For use by _PyFrame_GetFrameObject
281331
Do not call directly. */
282332
PyFrameObject *
@@ -288,6 +338,9 @@ _PyFrame_MakeAndSetFrameObject(_PyInterpreterFrame *frame);
288338
static inline PyFrameObject *
289339
_PyFrame_GetFrameObject(_PyInterpreterFrame *frame)
290340
{
341+
if (PyUnstable_JITExecutable_Check(PyStackRef_AsPyObjectBorrow(frame->f_executable))) {
342+
return _PyFrame_MakeAndSetFrameObject(frame);
343+
}
291344

292345
assert(!_PyFrame_IsIncomplete(frame));
293346
PyFrameObject *res = frame->frame_obj;
@@ -309,7 +362,7 @@ _PyFrame_ClearLocals(_PyInterpreterFrame *frame);
309362
* take should be set to 1 for heap allocated
310363
* frames like the ones in generators and coroutines.
311364
*/
312-
void
365+
PyAPI_FUNC(void)
313366
_PyFrame_ClearExceptCode(_PyInterpreterFrame * frame);
314367

315368
int

‎Include/internal/pycore_interpframe_structs.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,15 @@ struct _PyAsyncGenObject {
8585
_PyGenObject_HEAD(ag)
8686
};
8787

88+
typedef int (*_PyFrame_Reifier)(struct _PyInterpreterFrame *, PyObject *reifier);
89+
90+
typedef struct {
91+
PyObject_HEAD
92+
PyCodeObject *je_code;
93+
PyObject *je_state;
94+
_PyFrame_Reifier je_reifier;
95+
} PyUnstable_PyJitExecutable;
96+
8897
#undef _PyGenObject_HEAD
8998

9099

‎Lib/test/test_capi/test_misc.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import _thread
55
from collections import deque
66
import contextlib
7+
import dis
78
import importlib.machinery
89
import importlib.util
910
import json
@@ -2857,6 +2858,85 @@ def func():
28572858
names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
28582859
self.do_test(func, names)
28592860

2861+
def test_jit_frame(self):
2862+
def fakefunc():
2863+
pass
2864+
2865+
def f():
2866+
return sys._getframe(1)
2867+
2868+
res = _testinternalcapi.call_with_jit_frame(fakefunc, f, ())
2869+
2870+
def test_jit_frame_globals(self):
2871+
"""jit executable can fill in globals when accessed"""
2872+
def fakefunc():
2873+
pass
2874+
2875+
fake_globals = {"abc":42}
2876+
def callback():
2877+
return {"globals": fake_globals}
2878+
2879+
res = _testinternalcapi.call_with_jit_frame(fakefunc, globals, (), callback)
2880+
self.assertEqual(res, fake_globals)
2881+
2882+
def test_jit_frame_builtins(self):
2883+
"""jit executable can fill in builtins when accessed"""
2884+
def fakefunc():
2885+
pass
2886+
2887+
fake_builtins = {"abc":42}
2888+
def callback():
2889+
return {"builtins": fake_builtins}
2890+
2891+
res = _testinternalcapi.call_with_jit_frame(fakefunc, _testlimitedcapi.eval_getbuiltins, (), callback)
2892+
self.assertEqual(res, fake_builtins)
2893+
2894+
def test_jit_frame_instr_ptr(self):
2895+
"""jit executable can fill in the instr ptr each time the frame is queried"""
2896+
def fakefunc():
2897+
pass
2898+
pass
2899+
pass
2900+
pass
2901+
2902+
offset = 0
2903+
linenos = []
2904+
def test():
2905+
for op in dis.get_instructions(fakefunc):
2906+
if op.opname in ("RESUME", "NOP", "RETURN_VALUE"):
2907+
nonlocal offset
2908+
offset = op.offset//2
2909+
linenos.append(sys._getframe(1).f_lineno)
2910+
2911+
def callback():
2912+
return {"instr_ptr": offset}
2913+
2914+
_testinternalcapi.call_with_jit_frame(fakefunc, test, (), callback)
2915+
base = fakefunc.__code__.co_firstlineno
2916+
self.assertEqual(linenos, [base, base + 1, base + 2, base + 3, base + 4])
2917+
2918+
def test_jit_frame_code(self):
2919+
"""internal C api checks the for a code executor"""
2920+
def fakefunc():
2921+
pass
2922+
2923+
def callback():
2924+
return _testinternalcapi.iframe_getcode(sys._getframe(1))
2925+
2926+
res = _testinternalcapi.call_with_jit_frame(fakefunc, callback, ())
2927+
self.assertEqual(res, fakefunc.__code__)
2928+
2929+
def test_jit_frame_line(self):
2930+
"""internal C api checks the for a code executor"""
2931+
def fakefunc():
2932+
pass
2933+
2934+
def callback():
2935+
return _testinternalcapi.iframe_getline(sys._getframe(1))
2936+
2937+
res = _testinternalcapi.call_with_jit_frame(fakefunc, callback, ())
2938+
self.assertEqual(res, fakefunc.__code__.co_firstlineno)
2939+
28602940

28612941
@unittest.skipUnless(support.Py_GIL_DISABLED, 'need Py_GIL_DISABLED')
28622942
class TestPyThreadId(unittest.TestCase):

‎Modules/_testinternalcapi.c‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
#include "pycore_initconfig.h" // _Py_GetConfigsAsDict()
2929
#include "pycore_instruction_sequence.h" // _PyInstructionSequence_New()
3030
#include "pycore_interpframe.h" // _PyFrame_GetFunction()
31+
#include "pycore_interpframe_structs.h" // _PyInterpreterFrame
3132
#include "pycore_object.h" // _PyObject_IsFreed()
3233
#include "pycore_optimizer.h" // _Py_Executor_DependsOn
3334
#include "pycore_pathconfig.h" // _PyPathConfig_ClearGlobal()
@@ -693,6 +694,113 @@ set_eval_frame_record(PyObject *self, PyObject *list)
693694
Py_RETURN_NONE;
694695
}
695696

697+
typedef struct {
698+
bool initialized;
699+
_PyInterpreterFrame frame;
700+
} JitFrame;
701+
702+
int
703+
reifier(_PyInterpreterFrame *frame, PyObject *executable)
704+
{
705+
JitFrame *jitframe = (JitFrame*)((char *)frame - offsetof(JitFrame, frame));
706+
PyFunctionObject *func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
707+
if (!jitframe->initialized) {
708+
frame->f_locals = NULL;
709+
frame->f_globals = func->func_globals; // borrowed
710+
frame->f_builtins = func->func_builtins; // borrowed
711+
frame->frame_obj = NULL;
712+
#ifdef Py_GIL_DISABLED
713+
frame->tlbc_index = 0;
714+
#endif
715+
jitframe->initialized = true;
716+
}
717+
PyUnstable_PyJitExecutable *jit_exec = (PyUnstable_PyJitExecutable*)executable;
718+
if (jit_exec->je_state == NULL) {
719+
return 0;
720+
}
721+
722+
PyObject *res = PyObject_CallNoArgs(jit_exec->je_state);
723+
if (res == NULL) {
724+
return -1;
725+
}
726+
727+
// let the test-state function fill in details on the frame
728+
if (PyDict_Check(res)) {
729+
PyObject *globals = PyDict_GetItemString(res, "globals");
730+
if (globals != NULL) {
731+
frame->f_globals = globals;
732+
}
733+
PyObject *builtins = PyDict_GetItemString(res, "builtins");
734+
if (builtins != NULL) {
735+
frame->f_builtins = builtins;
736+
}
737+
PyObject *instr_ptr = PyDict_GetItemString(res, "instr_ptr");
738+
if (instr_ptr != NULL) {
739+
frame->instr_ptr = _PyCode_CODE((PyCodeObject *)func->func_code) +
740+
PyLong_AsLong(instr_ptr);
741+
}
742+
}
743+
Py_DECREF(res);
744+
return 0;
745+
}
746+
747+
static PyObject *
748+
call_with_jit_frame(PyObject *self, PyObject *args)
749+
{
750+
PyObject *fakefunc; // used for f_funcobj as-if we were that JITed function
751+
PyObject *call; // the thing to call for testing purposes
752+
PyObject *callargs; // the arguments to provide for the test call
753+
PyObject *state = NULL; // a state object provided to the reifier, for tests we
754+
// callback on it to populate fields.
755+
if (!PyArg_ParseTuple(args, "OOO|O", &fakefunc, &call, &callargs, &state)) {
756+
return NULL;
757+
}
758+
if (!PyTuple_Check(callargs)) {
759+
PyErr_SetString(PyExc_TypeError, "callargs must be a tuple");
760+
return NULL;
761+
}
762+
763+
PyThreadState *tstate = PyThreadState_Get();
764+
PyCodeObject *code = (PyCodeObject *)((PyFunctionObject *)fakefunc)->func_code;
765+
PyObject *executable = PyUnstable_MakeJITExecutable(reifier, code, state);
766+
if (executable == NULL) {
767+
return NULL;
768+
}
769+
770+
// Create JIT frame and push onto the _PyInterprerFrame stack.
771+
JitFrame frame;
772+
frame.initialized = false;
773+
// Initialize minimal set of fields
774+
frame.frame.previous = tstate->current_frame;
775+
frame.frame.f_executable = PyStackRef_FromPyObjectSteal(executable);
776+
frame.frame.f_funcobj = PyStackRef_FromPyObjectNew(fakefunc);
777+
frame.frame.instr_ptr = _PyCode_CODE(code) + code->_co_firsttraceable;
778+
frame.frame.stackpointer = &frame.frame.localsplus[0];
779+
frame.frame.owner = FRAME_OWNED_BY_THREAD;
780+
tstate->current_frame = &frame.frame;
781+
782+
// call the test function
783+
PyObject *res = PyObject_Call(call, callargs, NULL);
784+
785+
tstate->current_frame = frame.frame.previous;
786+
// the test function may have caused the frame to get reified.
787+
if (frame.initialized && frame.frame.frame_obj != NULL) {
788+
// remove our reifier
789+
PyStackRef_CLOSE(frame.frame.f_executable);
790+
frame.frame.f_executable = PyStackRef_FromPyObjectNew(code);
791+
792+
// Transfer ownership to the reified frame object
793+
_PyFrame_ClearExceptCode(&frame.frame);
794+
PyStackRef_CLOSE(frame.frame.f_executable);
795+
}
796+
else {
797+
// Pop frame from the stack
798+
PyStackRef_CLOSE(frame.frame.f_executable);
799+
PyStackRef_CLOSE(frame.frame.f_funcobj);
800+
}
801+
return res;
802+
}
803+
696804
/*[clinic input]
697805
698806
_testinternalcapi.compiler_cleandoc -> object
@@ -2517,6 +2625,7 @@ static PyMethodDef module_functions[] = {
25172625
{"DecodeLocaleEx", decode_locale_ex, METH_VARARGS},
25182626
{"set_eval_frame_default", set_eval_frame_default, METH_NOARGS, NULL},
25192627
{"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
2628+
{"call_with_jit_frame", call_with_jit_frame, METH_VARARGS, NULL},
25202629
_TESTINTERNALCAPI_COMPILER_CLEANDOC_METHODDEF
25212630
_TESTINTERNALCAPI_NEW_INSTRUCTION_SEQUENCE_METHODDEF
25222631
_TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF

‎Objects/frameobject.c‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2286,6 +2286,10 @@ _PyFrame_HasHiddenLocals(_PyInterpreterFrame *frame)
22862286
PyObject *
22872287
_PyFrame_GetLocals(_PyInterpreterFrame *frame)
22882288
{
2289+
if (_PyFrame_EnsureFrameFullyInitialized(frame) < 0) {
2290+
return NULL;
2291+
}
2292+
22892293
// We should try to avoid creating the FrameObject if possible.
22902294
// So we check if the frame is a module or class level scope
22912295
PyCodeObject *co = _PyFrame_GetCode(frame);

‎Objects/object.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2514,6 +2514,7 @@ static PyTypeObject* static_types[] = {
25142514
&PyTuple_Type,
25152515
&PyUnicodeIter_Type,
25162516
&PyUnicode_Type,
2517+
&PyUnstable_JITExecutable_Type,
25172518
&PyWrapperDescr_Type,
25182519
&PyZip_Type,
25192520
&Py_GenericAliasType,

0 commit comments

Comments
 (0)