gh-115999: Implement thread-local bytecode and enable specialization for BINARY_OP

Include/cpython/code.h

@@ -72,6 +72,24 @@ 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`. The
+ * first entry in the array always points to the "main" copy of the bytecode
+ * that is stored at the end of the code object.
+ */
+typedef struct {
+    Py_ssize_t size;
+    char *entries[1];
+} _PyCodeArray;
+
+#define _PyCode_DEF_THREAD_LOCAL_BYTECODE() \
+    _PyCodeArray *co_tlbc;
+#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) {                                                    \
@@ -138,6 +156,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)];                                             \
 }
 

Include/cpython/initconfig.h

@@ -183,6 +183,7 @@ typedef struct PyConfig {
     int cpu_count;
 #ifdef Py_GIL_DISABLED
     int enable_gil;
+    int tlbc_enabled;
 #endif
 
     /* --- Path configuration inputs ------------ */

Include/internal/pycore_ceval.h

@@ -174,6 +174,18 @@ _PyEval_IsGILEnabled(PyThreadState *tstate)
 extern int _PyEval_EnableGILTransient(PyThreadState *tstate);
 extern int _PyEval_EnableGILPermanent(PyThreadState *tstate);
 extern int _PyEval_DisableGIL(PyThreadState *state);
+
+
+static inline _Py_CODEUNIT *
+_PyEval_GetExecutableCode(PyThreadState *tstate, PyCodeObject *co)
+{
+    _Py_CODEUNIT *bc = _PyCode_GetTLBCFast(tstate, co);
+    if (bc != NULL) {
+        return bc;
+    }
+    return _PyCode_GetTLBC(co);
+}
+
 #endif
 
 extern void _PyEval_DeactivateOpCache(void);

Include/internal/pycore_code.h

@@ -11,6 +11,7 @@ extern "C" {
 #include "pycore_stackref.h"    // _PyStackRef
 #include "pycore_lock.h"        // PyMutex
 #include "pycore_backoff.h"     // _Py_BackoffCounter
+#include "pycore_tstate.h"      // _PyThreadStateImpl
 
 
 /* Each instruction in a code object is a fixed-width value,
@@ -313,11 +314,17 @@ extern int _PyLineTable_PreviousAddressRange(PyCodeAddressRange *range);
 /** API for executors */
 extern void _PyCode_Clear_Executors(PyCodeObject *code);
 
+
 #ifdef Py_GIL_DISABLED
 // gh-115999 tracks progress on addressing this.
 #define ENABLE_SPECIALIZATION 0
+// Use this to enable specialization families once they are thread-safe. All
+// uses will be replaced with ENABLE_SPECIALIZATION once all families are
+// thread-safe.
+#define ENABLE_SPECIALIZATION_FT 1
 #else
 #define ENABLE_SPECIALIZATION 1
+#define ENABLE_SPECIALIZATION_FT ENABLE_SPECIALIZATION
 #endif
 
 /* Specialization functions */
@@ -600,6 +607,40 @@ struct _PyCode8 _PyCode_DEF(8);
 
 PyAPI_DATA(const struct _PyCode8) _Py_InitCleanup;
 
+#ifdef Py_GIL_DISABLED
+
+// Return a pointer to the thread-local bytecode for the current thread, if it
+// exists.
+static inline _Py_CODEUNIT *
+_PyCode_GetTLBCFast(PyThreadState *tstate, PyCodeObject *co)
+{
+    _PyCodeArray *code = _Py_atomic_load_ptr_acquire(&co->co_tlbc);
+    int32_t idx = ((_PyThreadStateImpl*) tstate)->tlbc_index;
+    if (idx < code->size && code->entries[idx] != NULL) {
+        return (_Py_CODEUNIT *) code->entries[idx];
+    }
+    return NULL;
+}
+
+// Return a pointer to the thread-local bytecode for the current thread,
+// creating it if necessary.
+extern _Py_CODEUNIT *_PyCode_GetTLBC(PyCodeObject *co);
+
+// Reserve an index for the current thread into thread-local bytecode
+// arrays
+//
+// Returns the reserved index or -1 on error.
+extern int32_t _Py_ReserveTLBCIndex(PyInterpreterState *interp);
+
+// Release the current thread's index into thread-local bytecode arrays
+extern void _Py_ClearTLBCIndex(_PyThreadStateImpl *tstate);
+
+// Free all TLBC copies not associated with live threads.
+//
+// Returns 0 on success or -1 on error.
+extern int _Py_ClearUnusedTLBC(PyInterpreterState *interp);
+#endif
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_frame.h

@@ -68,6 +68,10 @@ 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
+    /* Index of thread-local bytecode containing instr_ptr. */
+    int32_t tlbc_index;
+#endif
     _PyStackRef *stackpointer;
     uint16_t return_offset;  /* Only relevant during a function call */
     char owner;
@@ -76,14 +80,27 @@ typedef struct _PyInterpreterFrame {
 } _PyInterpreterFrame;
 
 #define _PyInterpreterFrame_LASTI(IF) \
-    ((int)((IF)->instr_ptr - _PyCode_CODE(_PyFrame_GetCode(IF))))
+    ((int)((IF)->instr_ptr - _PyFrame_GetBytecode((IF))))
 
 static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
     PyObject *executable = PyStackRef_AsPyObjectBorrow(f->f_executable);
     assert(PyCode_Check(executable));
     return (PyCodeObject *)executable;
 }
 
+static inline _Py_CODEUNIT *
+_PyFrame_GetBytecode(_PyInterpreterFrame *f)
+{
+#ifdef Py_GIL_DISABLED
+    PyCodeObject *co = _PyFrame_GetCode(f);
+    _PyCodeArray *tlbc = _Py_atomic_load_ptr_acquire(&co->co_tlbc);
+    assert(f->tlbc_index >= 0 && f->tlbc_index < tlbc->size);
+    return (_Py_CODEUNIT *)tlbc->entries[f->tlbc_index];
+#else
+    return _PyCode_CODE(_PyFrame_GetCode(f));
+#endif
+}
+
 static inline PyFunctionObject *_PyFrame_GetFunction(_PyInterpreterFrame *f) {
     PyObject *func = PyStackRef_AsPyObjectBorrow(f->f_funcobj);
     assert(PyFunction_Check(func));
@@ -144,13 +161,33 @@ static inline void _PyFrame_Copy(_PyInterpreterFrame *src, _PyInterpreterFrame *
 #endif
 }
 
+#ifdef Py_GIL_DISABLED
+static inline void
+_PyFrame_InitializeTLBC(PyThreadState *tstate, _PyInterpreterFrame *frame,
+                        PyCodeObject *code)
+{
+    _Py_CODEUNIT *tlbc = _PyCode_GetTLBCFast(tstate, code);
+    if (tlbc == NULL) {
+        // No thread-local bytecode exists for this thread yet; use the main
+        // thread's copy, deferring thread-local bytecode creation to the
+        // execution of RESUME.
+        frame->instr_ptr = _PyCode_CODE(code);
+        frame->tlbc_index = 0;
+    }
+    else {
+        frame->instr_ptr = tlbc;
+        frame->tlbc_index = ((_PyThreadStateImpl *)tstate)->tlbc_index;
+    }
+}
+#endif
+
 /* Consumes reference to func and locals.
    Does not initialize frame->previous, which happens
    when frame is linked into the frame stack.
  */
 static inline void
 _PyFrame_Initialize(
-    _PyInterpreterFrame *frame, _PyStackRef func,
+    PyThreadState *tstate, _PyInterpreterFrame *frame, _PyStackRef func,
     PyObject *locals, PyCodeObject *code, int null_locals_from, _PyInterpreterFrame *previous)
 {
     frame->previous = previous;
@@ -162,7 +199,12 @@ _PyFrame_Initialize(
     frame->f_locals = locals;
     frame->stackpointer = frame->localsplus + code->co_nlocalsplus;
     frame->frame_obj = NULL;
+#ifdef Py_GIL_DISABLED
+    _PyFrame_InitializeTLBC(tstate, frame, code);
+#else
+    (void)tstate;
     frame->instr_ptr = _PyCode_CODE(code);
+#endif
     frame->return_offset = 0;
     frame->owner = FRAME_OWNED_BY_THREAD;
 
@@ -224,7 +266,8 @@ _PyFrame_IsIncomplete(_PyInterpreterFrame *frame)
         return true;
     }
     return frame->owner != FRAME_OWNED_BY_GENERATOR &&
-        frame->instr_ptr < _PyCode_CODE(_PyFrame_GetCode(frame)) + _PyFrame_GetCode(frame)->_co_firsttraceable;
+           frame->instr_ptr < _PyFrame_GetBytecode(frame) +
+                                  _PyFrame_GetCode(frame)->_co_firsttraceable;
 }
 
 static inline _PyInterpreterFrame *
@@ -315,7 +358,8 @@ _PyFrame_PushUnchecked(PyThreadState *tstate, _PyStackRef func, int null_locals_
     _PyInterpreterFrame *new_frame = (_PyInterpreterFrame *)tstate->datastack_top;
     tstate->datastack_top += code->co_framesize;
     assert(tstate->datastack_top < tstate->datastack_limit);
-    _PyFrame_Initialize(new_frame, func, NULL, code, null_locals_from, previous);
+    _PyFrame_Initialize(tstate, new_frame, func, NULL, code, null_locals_from,
+                        previous);
     return new_frame;
 }
 
@@ -339,7 +383,11 @@ _PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int
     assert(stackdepth <= code->co_stacksize);
     frame->stackpointer = frame->localsplus + code->co_nlocalsplus + stackdepth;
     frame->frame_obj = NULL;
+#ifdef Py_GIL_DISABLED
+    _PyFrame_InitializeTLBC(tstate, frame, code);
+#else
     frame->instr_ptr = _PyCode_CODE(code);
+#endif
     frame->owner = FRAME_OWNED_BY_THREAD;
     frame->return_offset = 0;
 

Include/internal/pycore_gc.h

@@ -389,6 +389,10 @@ extern int _PyGC_VisitStackRef(union _PyStackRef *ref, visitproc visit, void *ar
         }                                                               \
     } while (0)
 
+#ifdef Py_GIL_DISABLED
+extern void _PyGC_VisitObjectsWorldStopped(PyInterpreterState *interp,
+                                           gcvisitobjects_t callback, void *arg);
+#endif
 
 #ifdef __cplusplus
 }

Include/internal/pycore_index_pool.h

@@ -0,0 +1,56 @@
+#ifndef Py_INTERNAL_INDEX_POOL_H
+#define Py_INTERNAL_INDEX_POOL_H
+
+#include "Python.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#ifdef Py_GIL_DISABLED
+
+// This contains code for allocating unique indices in an array. It is used by
+// the free-threaded build to assign each thread a globally unique index into
+// each code object's thread-local bytecode array.
+
+// A min-heap of indices
+typedef struct _PyIndexHeap {
+    int32_t *values;
+
+    // Number of items stored in values
+    Py_ssize_t size;
+
+    // Maximum number of items that can be stored in values
+    Py_ssize_t capacity;
+} _PyIndexHeap;
+
+// An unbounded pool of indices. Indices are allocated starting from 0. They
+// may be released back to the pool once they are no longer in use.
+typedef struct _PyIndexPool {
+    PyMutex mutex;
+
+    // Min heap of indices available for allocation
+    _PyIndexHeap free_indices;
+
+    // Next index to allocate if no free indices are available
+    int32_t next_index;
+} _PyIndexPool;
+
+// Allocate the smallest available index. Returns -1 on error.
+extern int32_t _PyIndexPool_AllocIndex(_PyIndexPool *indices);
+
+// Release `index` back to the pool
+extern void _PyIndexPool_FreeIndex(_PyIndexPool *indices, int32_t index);
+
+extern void _PyIndexPool_Fini(_PyIndexPool *indices);
+
+#endif // Py_GIL_DISABLED
+
+#ifdef __cplusplus
+}
+#endif
+#endif // !Py_INTERNAL_INDEX_POOL_H

Include/internal/pycore_interp.h

@@ -26,6 +26,7 @@ extern "C" {
 #include "pycore_genobject.h"     // _PyGen_FetchStopIterationValue
 #include "pycore_global_objects.h"// struct _Py_interp_cached_objects
 #include "pycore_import.h"        // struct _import_state
+#include "pycore_index_pool.h"     // _PyIndexPool
 #include "pycore_instruments.h"   // _PY_MONITORING_EVENTS
 #include "pycore_list.h"          // struct _Py_list_state
 #include "pycore_mimalloc.h"      // struct _mimalloc_interp_state
@@ -222,6 +223,7 @@ struct _is {
     struct _brc_state brc;  // biased reference counting state
     struct _Py_unique_id_pool unique_ids;  // object ids for per-thread refcounts
     PyMutex weakref_locks[NUM_WEAKREF_LIST_LOCKS];
+    _PyIndexPool tlbc_indices;
 #endif
 
     // Per-interpreter state for the obmalloc allocator.  For the main

Include/internal/pycore_tstate.h

@@ -42,6 +42,9 @@ typedef struct _PyThreadStateImpl {
         int is_finalized;
     } refcounts;
 
+    // Index to use to retrieve thread-local bytecode for this thread
+    int32_t tlbc_index;
+
     // When >1, code objects do not immortalize their non-string constants.
     int suppress_co_const_immortalization;
 #endif
@@ -52,7 +55,6 @@ typedef struct _PyThreadStateImpl {
 
 } _PyThreadStateImpl;
 
-
 #ifdef __cplusplus
 }
 #endif