Get bytecode copying working

mpage · Nov 4, 2024 · Aug 15, 2024 · Aug 27, 2024 · Aug 20, 2024 · Aug 30, 2024
commit 344d7adb8a8f9eafad15c8864d39c1cb825beae3
diff --git a/Include/cpython/code.h b/Include/cpython/code.h
@@ -68,6 +68,20 @@ typedef struct {
     uint8_t *per_instruction_tools;
 } _PyCoMonitoringData;
 
+#ifdef Py_GIL_DISABLED
+/* Each thread specializes a thread-local copy of the bytecode in free-threaded
+ * builds. These copies are stored on the code object in a `_PyCodeArray`.
+ */
+typedef struct {
+    Py_ssize_t size;
+    char *entries[];
+} _PyCodeArray;
+
+#define _PyCode_DEF_THREAD_LOCAL_BYTECODE() _PyCodeArray *co_specialized_code;
+#else
+#define _PyCode_DEF_THREAD_LOCAL_BYTECODE()
+#endif
+
 // To avoid repeating ourselves in deepfreeze.py, all PyCodeObject members are
 // defined in this macro:
 #define _PyCode_DEF(SIZE) {                                                    \
@@ -133,6 +147,7 @@ typedef struct {
        Type is a void* to keep the format private in codeobject.c to force     \
        people to go through the proper APIs. */                                \
     void *co_extra;                                                            \
+    _PyCode_DEF_THREAD_LOCAL_BYTECODE()                                        \
     char co_code_adaptive[(SIZE)];                                             \
 }
 

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
@@ -602,6 +602,26 @@ struct _PyCode8 _PyCode_DEF(8);
 PyAPI_DATA(const struct _PyCode8) _Py_InitCleanup;
 
 #ifdef Py_GIL_DISABLED
+
+extern _Py_CODEUNIT *_PyCode_CreateSpecializableCode(PyCodeObject *co);
+/* Return bytecode that should be executed.
+ * Will not return NULL, but may disable specialization, in which case the
+ * returned bytecode should not be specialized.
+ *
+ * XXX - This is a confusing contract.
+ */
+static inline _Py_CODEUNIT *
+_PyCode_GetSpecializableCode(PyCodeObject *co)
+{
+    _PyCodeArray *code = _Py_atomic_load_ptr_acquire(&co->co_specialized_code);
+    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *) PyThreadState_GET();
+    Py_ssize_t idx = tstate->specialized_code_index;
+    if (idx < code->size && code->entries[idx] != NULL) {
+        // XXX - Do we need to worry about alignment here?
+        return (_Py_CODEUNIT *) code->entries[idx];
+    }
+    return _PyCode_CreateSpecializableCode(co);
+}
 extern int _Py_ReserveSpecializedCodeIndex(PyInterpreterState *interp);
 extern void _Py_ClearSpecializedCodeIndex(_PyThreadStateImpl *tstate);
 #endif

diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
@@ -68,6 +68,9 @@ typedef struct _PyInterpreterFrame {
     PyObject *f_locals; /* Strong reference, may be NULL. Only valid if not on C stack */
     PyFrameObject *frame_obj; /* Strong reference, may be NULL. Only valid if not on C stack */
     _Py_CODEUNIT *instr_ptr; /* Instruction currently executing (or about to begin) */
+#ifdef Py_GIL_DISABLED
+    _Py_CODEUNIT *bytecode;
+#endif
     _PyStackRef *stackpointer;
     uint16_t return_offset;  /* Only relevant during a function call */
     char owner;
@@ -84,7 +87,11 @@ static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
 }
 
 static inline _Py_CODEUNIT *_PyFrame_GetBytecode(_PyInterpreterFrame *f) {
+#ifdef Py_GIL_DISABLED
+    return f->bytecode;
+#else
     return _PyCode_CODE(_PyFrame_GetCode(f));
+#endif
 }
 
 static inline _PyStackRef *_PyFrame_Stackbase(_PyInterpreterFrame *f) {
@@ -167,6 +174,7 @@ _PyFrame_Initialize(
     }
 
 #ifdef Py_GIL_DISABLED
+    frame->bytecode = frame->instr_ptr;
     // On GIL disabled, we walk the entire stack in GC. Since stacktop
     // is not always in sync with the real stack pointer, we have
     // no choice but to traverse the entire stack.
@@ -339,6 +347,7 @@ _PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int
     frame->return_offset = 0;
 
 #ifdef Py_GIL_DISABLED
+    frame->bytecode = frame->instr_ptr;
     assert(code->co_nlocalsplus == 0);
     for (int i = 0; i < code->co_stacksize; i++) {
         frame->localsplus[i] = PyStackRef_NULL;

diff --git a/Include/internal/pycore_uop_ids.h b/Include/internal/pycore_uop_ids.h
@@ -12,11 +12,14 @@
 #include "pycore_object.h"        // _PyObject_SetDeferredRefcount
 #include "pycore_opcode_metadata.h" // _PyOpcode_Deopt, _PyOpcode_Caches
 #include "pycore_opcode_utils.h"  // RESUME_AT_FUNC_START
+#include "pycore_pymem.h"         // _PyMem_FreeDelayed
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
 #include "pycore_setobject.h"     // _PySet_NextEntry()
 #include "pycore_tuple.h"         // _PyTuple_ITEMS()
 #include "clinic/codeobject.c.h"
 
+#define INITIAL_SPECIALIZED_CODE_SIZE 16
+
 static const char *
 code_event_name(PyCodeEvent event) {
     switch (event) {
@@ -450,7 +453,7 @@ _PyCode_Validate(struct _PyCodeConstructor *con)
 
 extern void _PyCode_Quicken(_Py_CODEUNIT *instructions, Py_ssize_t size);
 
-static void
+static int
 init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
 {
     int nlocalsplus = (int)PyTuple_GET_SIZE(con->localsplusnames);
@@ -513,6 +516,12 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
 
     memcpy(_PyCode_CODE(co), PyBytes_AS_STRING(con->code),
            PyBytes_GET_SIZE(con->code));
+#ifdef Py_GIL_DISABLED
+    // XXX - initialize code array
+    co->co_specialized_code = PyMem_Calloc(1, sizeof(_PyCodeArray) + sizeof(void*) * INITIAL_SPECIALIZED_CODE_SIZE);
+    co->co_specialized_code->size = INITIAL_SPECIALIZED_CODE_SIZE;
+    co->co_specialized_code->entries[0] = (char *) _PyCode_CODE(co);
+#endif
     int entry_point = 0;
     while (entry_point < Py_SIZE(co) &&
         _PyCode_CODE(co)[entry_point].op.code != RESUME) {
@@ -521,6 +530,7 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     co->_co_firsttraceable = entry_point;
     _PyCode_Quicken(_PyCode_CODE(co), Py_SIZE(co));
     notify_code_watchers(PY_CODE_EVENT_CREATE, co);
+    return 0;
 }
 
 static int
@@ -675,7 +685,12 @@ _PyCode_New(struct _PyCodeConstructor *con)
         PyErr_NoMemory();
         return NULL;
     }
-    init_code(co, con);
+
+    if (init_code(co, con) < 0) {
+        Py_DECREF(co);
+        return NULL;
+    }
+
 #ifdef Py_GIL_DISABLED
     _PyObject_SetDeferredRefcount((PyObject *)co);
     _PyObject_GC_TRACK(co);
@@ -1871,6 +1886,15 @@ code_dealloc(PyCodeObject *co)
         PyObject_ClearWeakRefs((PyObject*)co);
     }
     free_monitoring_data(co->_co_monitoring);
+#ifdef Py_GIL_DISABLED
+    // The first element always points to the bytecode that follows the fixed
+    // part of the code object, which will be freed when the code object is
+    // freed.
+    for (Py_ssize_t i = 1; i < co->co_specialized_code->size; i++) {
+        PyMem_Free(co->co_specialized_code->entries[i]);
+    }
+    PyMem_Free(co->co_specialized_code);
+#endif
     PyObject_Free(co);
 }
 
@@ -2637,3 +2661,56 @@ _PyCode_Fini(PyInterpreterState *interp)
     _PyIndexPool_Fini(&interp->specialized_code_indices);
 #endif
 }
+
+#ifdef Py_GIL_DISABLED
+
+static void
+copy_code(_Py_CODEUNIT *dst, _Py_CODEUNIT *src, Py_ssize_t nbytes)
+{
+    int code_len = Py_SIZE(co);
+    _Py_CODEUNIT *dst_bytecode = (_Py_CODEUNIT *) dst->bytecode;
+    for (int i = 0; i < code_len; i += _PyInstruction_GetLength(co, i)) {
+        dst_bytecode[i] = _Py_GetBaseCodeUnit(co, i);
+    }
+    _PyCode_Quicken(dst_bytecode, code_len);
+}
+
+static _Py_CODEUNIT *
+create_specializable_code_lock_held(PyCodeObject *co)
+{
+    _PyCodeArray *spec_code = co->co_specialized_code;
+    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *) PyThreadState_GET();
+    Py_ssize_t idx = tstate->specialized_code_index;
+    if (idx >= spec_code->size) {
+        Py_ssize_t new_size = spec_code->size * 2;
+        _PyCodeArray *new_spec_code = PyMem_Calloc(sizeof(_PyCodeArray) + sizeof(char*) * new_size, 1);
+        if (new_spec_code == NULL) {
+            PyErr_NoMemory();
+            return NULL;
+        }
+        new_spec_code->size = new_size;
+        memcpy(new_spec_code->entries, spec_code->entries, spec_code->size * sizeof(char*));
+        _Py_atomic_store_ptr_release(&co->co_specialized_code, new_spec_code);
+        _PyMem_FreeDelayed(spec_code);
+        spec_code = new_spec_code;
+    }
+    spec_code->entries[idx] = PyMem_Malloc(_PyCode_NBYTES(co));
+    if (spec_code->entries[idx] == NULL) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+    copy_code((_Py_CODEUNIT *) spec_code->entries[idx], _PyCode_CODE(co), _PyCode_NBYTES(co));
+    return (_Py_CODEUNIT *) spec_code->entries[idx];
+}
+
+_Py_CODEUNIT *
+_PyCode_CreateSpecializableCode(PyCodeObject *co)
+{
+    _Py_CODEUNIT *result;
+    Py_BEGIN_CRITICAL_SECTION(co);
+    result = create_specializable_code_lock_held(co);
+    Py_END_CRITICAL_SECTION();
+    return result;
+}
+
+#endif