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) {                                                    \
@@ -137,6 +155,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

@@ -177,6 +177,21 @@ _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(PyCodeObject *co)
+{
+    _Py_CODEUNIT *bc = _PyCode_GetTLBCFast(co);
+    if (bc != NULL) {
+        return bc;
+    }
+    if (!_PyInterpreterState_GET()->config.tlbc_enabled) {
+        return _PyCode_CODE(co);
+    }
+    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,41 @@ extern int _PyLineTable_PreviousAddressRange(PyCodeAddressRange *range);
 /** API for executors */
 extern void _PyCode_Clear_Executors(PyCodeObject *code);
 
+#define ENABLE_SPECIALIZATION 1
+
 #ifdef Py_GIL_DISABLED
 // gh-115999 tracks progress on addressing this.
-#define ENABLE_SPECIALIZATION 0
+#define ENABLE_SPECIALIZED_BINARY_OP ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_BINARY_SUBSCR 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_CALL 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_CALL_KW 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_COMPARE_OP 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_CONTAINS_OP 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_FOR_ITER 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_ATTR 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_GLOBAL 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_SUPER_ATTR 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_SEND 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_STORE_ATTR 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_STORE_SUBSCR 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_TO_BOOL 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_UNPACK_SEQUENCE 0 && ENABLE_SPECIALIZATION
 #else
-#define ENABLE_SPECIALIZATION 1
+#define ENABLE_SPECIALIZED_BINARY_OP ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_BINARY_SUBSCR ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_CALL ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_CALL_KW ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_COMPARE_OP ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_CONTAINS_OP ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_FOR_ITER ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_ATTR ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_GLOBAL ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_SUPER_ATTR ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_SEND ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_STORE_ATTR ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_STORE_SUBSCR ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_TO_BOOL ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_UNPACK_SEQUENCE ENABLE_SPECIALIZATION
 #endif
 
 /* Specialization functions */
@@ -600,6 +631,36 @@ 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(PyCodeObject *co)
+{
+    _PyCodeArray *code = _Py_atomic_load_ptr_acquire(&co->co_tlbc);
+    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *) PyThreadState_GET();
+    Py_ssize_t idx = 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 Py_ssize_t _Py_ReserveTLBCIndex(PyInterpreterState *interp);
+
+// Release the current thread's index into thread-local bytecode arrays
+extern void _Py_ClearTLBCIndex(_PyThreadStateImpl *tstate);
+#endif
+
 #ifdef __cplusplus
 }
 #endif

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;
@@ -76,14 +79,24 @@ 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
+    return f->bytecode;
+#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));
@@ -171,6 +184,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.
@@ -224,7 +238,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 *
@@ -344,6 +359,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;

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 {
+    Py_ssize_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
+    Py_ssize_t next_index;
+} _PyIndexPool;
+
+// Allocate the smallest available index. Returns -1 on error.
+extern Py_ssize_t _PyIndexPool_AllocIndex(_PyIndexPool *indices);
+
+// Release `index` back to the pool
+extern void _PyIndexPool_FreeIndex(_PyIndexPool *indices, Py_ssize_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
@@ -223,6 +224,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

@@ -41,6 +41,9 @@ typedef struct _PyThreadStateImpl {
         // If set, don't use per-thread refcounts
         int is_finalized;
     } refcounts;
+
+    // Index to use to retrieve thread-local bytecode for this thread
+    Py_ssize_t tlbc_index;
 #endif
 
 #if defined(Py_REF_DEBUG) && defined(Py_GIL_DISABLED)