gh-141594: A free-threaded JIT

Include/internal/pycore_interp_structs.h

@@ -930,7 +930,7 @@ struct _is {
     struct types_state types;
     struct callable_cache callable_cache;
     PyObject *common_consts[NUM_COMMON_CONSTANTS];
-    bool jit;
+    uint8_t jit;
     bool compiling;
     struct _PyExecutorObject *executor_list_head;
     struct _PyExecutorObject *executor_deletion_list_head;

Include/internal/pycore_optimizer.h

@@ -372,8 +372,6 @@ _PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame,
 
 void _PyJit_FinalizeTracing(PyThreadState *tstate);
 
-void _PyJit_Tracer_InvalidateDependency(PyThreadState *old_tstate, void *obj);
-
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_capi/test_opt.py

@@ -9,9 +9,13 @@
 
 import _opcode
 
-from test.support import (script_helper, requires_specialization,
-                          import_helper, Py_GIL_DISABLED, requires_jit_enabled,
-                          reset_code)
+from test.support import (
+    script_helper,
+    import_helper,
+    Py_GIL_DISABLED,
+    requires_jit_enabled,
+    reset_code
+)
 
 _testinternalcapi = import_helper.import_module("_testinternalcapi")
 
@@ -61,8 +65,6 @@ def get_opnames(ex):
     return list(iter_opnames(ex))
 
 
-@requires_specialization
-@unittest.skipIf(Py_GIL_DISABLED, "optimizer not yet supported in free-threaded builds")
 @requires_jit_enabled
 class TestExecutorInvalidation(unittest.TestCase):
 
@@ -130,8 +132,6 @@ def f():
         self.assertIsNone(exe)
 
 
-@requires_specialization
-@unittest.skipIf(Py_GIL_DISABLED, "optimizer not yet supported in free-threaded builds")
 @requires_jit_enabled
 @unittest.skipIf(os.getenv("PYTHON_UOPS_OPTIMIZE") == "0", "Needs uop optimizer to run.")
 class TestUops(unittest.TestCase):
@@ -434,8 +434,6 @@ def testfunc(n, m):
         self.assertIn("_FOR_ITER_TIER_TWO", uops)
 
 
-@requires_specialization
-@unittest.skipIf(Py_GIL_DISABLED, "optimizer not yet supported in free-threaded builds")
 @requires_jit_enabled
 @unittest.skipIf(os.getenv("PYTHON_UOPS_OPTIMIZE") == "0", "Needs uop optimizer to run.")
 class TestUopsOptimization(unittest.TestCase):
@@ -2052,6 +2050,7 @@ def testfunc(n):
         self.assertNotIn("_GUARD_NOS_INT", uops)
         self.assertNotIn("_GUARD_TOS_INT", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build immortalizes constants")
     def test_call_len_known_length_small_int(self):
         # Make sure that len(t) is optimized for a tuple of length 5.
         # See https://github.com/python/cpython/issues/139393.
@@ -2076,6 +2075,7 @@ def testfunc(n):
         self.assertNotIn("_POP_CALL_LOAD_CONST_INLINE_BORROW", uops)
         self.assertNotIn("_POP_TOP_LOAD_CONST_INLINE_BORROW", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build immortalizes constants")
     def test_call_len_known_length(self):
         # Make sure that len(t) is not optimized for a tuple of length 2048.
         # See https://github.com/python/cpython/issues/139393.
@@ -2497,6 +2497,7 @@ def testfunc(n):
 
         self.assertIn("_POP_TOP_NOP", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT might immortalize this.")
     def test_pop_top_specialize_int(self):
         def testfunc(n):
             for _ in range(n):
@@ -2510,6 +2511,7 @@ def testfunc(n):
 
         self.assertIn("_POP_TOP_INT", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT might immortalize this.")
     def test_pop_top_specialize_float(self):
         def testfunc(n):
             for _ in range(n):

Misc/NEWS.d/next/Core_and_Builtins/2025-11-15-16-30-46.gh-issue-141594.PSsC5J.rst

@@ -0,0 +1,3 @@
+Add free-threading support to the JIT. The JIT is only enabled on
+single-threaded code in free-threading, and is disabled when multiple
+threads are spawned. Patch by Ken Jin.

Modules/_testinternalcapi.c

@@ -1240,8 +1240,7 @@ add_executor_dependency(PyObject *self, PyObject *args)
 static PyObject *
 invalidate_executors(PyObject *self, PyObject *obj)
 {
-    PyInterpreterState *interp = PyInterpreterState_Get();
-    _Py_Executors_InvalidateDependency(interp, obj, 1);
+    _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), obj, 1);
     Py_RETURN_NONE;
 }
 

Objects/codeobject.c

@@ -2432,7 +2432,7 @@ code_dealloc(PyObject *self)
         PyMem_Free(co_extra);
     }
 #ifdef _Py_TIER2
-    _PyJit_Tracer_InvalidateDependency(tstate, self);
+    _Py_Executors_InvalidateDependency(tstate->interp, self, 1);
     if (co->co_executors != NULL) {
         clear_executors(co);
     }
@@ -3363,8 +3363,12 @@ deopt_code_unit(PyCodeObject *code, int i)
         inst.op.code = _PyOpcode_Deopt[opcode];
         assert(inst.op.code < MIN_SPECIALIZED_OPCODE);
     }
-    // JIT should not be enabled with free-threading
-    assert(inst.op.code != ENTER_EXECUTOR);
+    if (inst.op.code == ENTER_EXECUTOR) {
+        _PyExecutorObject *exec = code->co_executors->executors[inst.op.arg];
+        assert(exec != NULL);
+        inst.op.code = exec->vm_data.opcode;
+        inst.op.arg = exec->vm_data.oparg;
+    }
     return inst;
 }
 

Objects/frameobject.c

@@ -263,7 +263,6 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
 
 #if _Py_TIER2
         _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), co, 1);
-        _PyJit_Tracer_InvalidateDependency(_PyThreadState_GET(), co);
 #endif
 
         _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);

Objects/funcobject.c

@@ -11,7 +11,7 @@
 #include "pycore_setobject.h"     // _PySet_NextEntry()
 #include "pycore_stats.h"
 #include "pycore_weakref.h"       // FT_CLEAR_WEAKREFS()
-#include "pycore_optimizer.h"     // _PyJit_Tracer_InvalidateDependency
+#include "pycore_optimizer.h"     // _Py_Executors_InvalidateDependency
 
 static const char *
 func_event_name(PyFunction_WatchEvent event) {
@@ -298,7 +298,7 @@ functions is running.
 
 */
 
-#ifndef Py_GIL_DISABLED
+#if _Py_TIER2
 static inline struct _func_version_cache_item *
 get_cache_item(PyInterpreterState *interp, uint32_t version)
 {
@@ -315,11 +315,13 @@ _PyFunction_SetVersion(PyFunctionObject *func, uint32_t version)
     // This should only be called from MAKE_FUNCTION. No code is specialized
     // based on the version, so we do not need to stop the world to set it.
     func->func_version = version;
-#ifndef Py_GIL_DISABLED
+#if _Py_TIER2
     PyInterpreterState *interp = _PyInterpreterState_GET();
+    FT_MUTEX_LOCK(&interp->func_state.mutex);
     struct _func_version_cache_item *slot = get_cache_item(interp, version);
     slot->func = func;
     slot->code = func->func_code;
+    FT_MUTEX_UNLOCK(&interp->func_state.mutex);
 #endif
 }
 
@@ -330,13 +332,15 @@ func_clear_version(PyInterpreterState *interp, PyFunctionObject *func)
         // Version was never set or has already been cleared.
         return;
     }
-#ifndef Py_GIL_DISABLED
+#if _Py_TIER2
+    FT_MUTEX_LOCK(&interp->func_state.mutex);
     struct _func_version_cache_item *slot =
         get_cache_item(interp, func->func_version);
     if (slot->func == func) {
         slot->func = NULL;
         // Leave slot->code alone, there may be use for it.
     }
+    FT_MUTEX_UNLOCK(&interp->func_state.mutex);
 #endif
     func->func_version = FUNC_VERSION_CLEARED;
 }
@@ -358,8 +362,9 @@ _PyFunction_ClearVersion(PyFunctionObject *func)
 void
 _PyFunction_ClearCodeByVersion(uint32_t version)
 {
-#ifndef Py_GIL_DISABLED
+#if _Py_TIER2
     PyInterpreterState *interp = _PyInterpreterState_GET();
+    FT_MUTEX_LOCK(&interp->func_state.mutex);
     struct _func_version_cache_item *slot = get_cache_item(interp, version);
     if (slot->code) {
         assert(PyCode_Check(slot->code));
@@ -369,15 +374,17 @@ _PyFunction_ClearCodeByVersion(uint32_t version)
             slot->func = NULL;
         }
     }
+    FT_MUTEX_UNLOCK(&interp->func_state.mutex);
 #endif
 }
 
 PyFunctionObject *
 _PyFunction_LookupByVersion(uint32_t version, PyObject **p_code)
 {
-#ifdef Py_GIL_DISABLED
-    return NULL;
-#else
+#if _Py_TIER2
+    // This function does not need locking/atomics as it can only be
+    // called from the optimizer, which is currently disabled
+    // when there are multiple threads.
     PyInterpreterState *interp = _PyInterpreterState_GET();
     struct _func_version_cache_item *slot = get_cache_item(interp, version);
     if (slot->code) {
@@ -395,12 +402,18 @@ _PyFunction_LookupByVersion(uint32_t version, PyObject **p_code)
         return slot->func;
     }
     return NULL;
+#else
+    return NULL;
 #endif
 }
 
 uint32_t
 _PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
 {
+    // This function does not need locking/atomics as it can only be
+    // called from the specializing interpreter or optimizer.
+    // The specializing interpreter holds a strong reference to the function.
+    // The optimizer is currently disabled when there are multiple threads.
     return func->func_version;
 }
 
@@ -1153,7 +1166,6 @@ func_dealloc(PyObject *self)
     }
 #if _Py_TIER2
     _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), self, 1);
-    _PyJit_Tracer_InvalidateDependency(_PyThreadState_GET(), self);
 #endif
     _PyObject_GC_UNTRACK(op);
     FT_CLEAR_WEAKREFS(self, op->func_weakreflist);

Objects/listobject.c

@@ -79,7 +79,9 @@ ensure_shared_on_resize(PyListObject *self)
     // We can't use _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED here because
     // the `CALL_LIST_APPEND` bytecode handler may lock the list without
     // a critical section.
-    assert(Py_REFCNT(self) == 1 || PyMutex_IsLocked(&_PyObject_CAST(self)->ob_mutex));
+    assert(Py_REFCNT(self) == 1 ||
+        (_Py_IsOwnedByCurrentThread((PyObject *)self) && !_PyObject_GC_IS_SHARED(self)) ||
+        PyMutex_IsLocked(&_PyObject_CAST(self)->ob_mutex));
 
     // Ensure that the list array is freed using QSBR if we are not the
     // owning thread.

Objects/typeobject.c

@@ -1148,15 +1148,18 @@ static void
 set_version_unlocked(PyTypeObject *tp, unsigned int version)
 {
     assert(version == 0 || (tp->tp_versions_used != _Py_ATTR_CACHE_UNUSED));
-#ifndef Py_GIL_DISABLED
+#if _Py_TIER2
     PyInterpreterState *interp = _PyInterpreterState_GET();
+    BEGIN_TYPE_LOCK();
     // lookup the old version and set to null
     if (tp->tp_version_tag != 0) {
         PyTypeObject **slot =
             interp->types.type_version_cache
             + (tp->tp_version_tag % TYPE_VERSION_CACHE_SIZE);
         *slot = NULL;
     }
+#endif
+#ifndef Py_GIL_DISABLED
     if (version) {
         tp->tp_versions_used++;
     }
@@ -1166,13 +1169,14 @@ set_version_unlocked(PyTypeObject *tp, unsigned int version)
     }
 #endif
     FT_ATOMIC_STORE_UINT_RELAXED(tp->tp_version_tag, version);
-#ifndef Py_GIL_DISABLED
+#if _Py_TIER2
     if (version != 0) {
         PyTypeObject **slot =
             interp->types.type_version_cache
             + (version % TYPE_VERSION_CACHE_SIZE);
         *slot = tp;
     }
+    END_TYPE_LOCK();
 #endif
 }
 
@@ -1357,9 +1361,12 @@ _PyType_SetVersion(PyTypeObject *tp, unsigned int version)
 PyTypeObject *
 _PyType_LookupByVersion(unsigned int version)
 {
-#ifdef Py_GIL_DISABLED
+#ifndef _Py_TIER2
     return NULL;
 #else
+    // This function does not need locking/atomics as it can only be
+    // called from the optimizer, which is currently disabled
+    // when there are multiple threads.
     PyInterpreterState *interp = _PyInterpreterState_GET();
     PyTypeObject **slot =
         interp->types.type_version_cache

Python/bytecodes.c

@@ -2939,9 +2939,9 @@ dummy_func(
         };
 
         specializing tier1 op(_SPECIALIZE_JUMP_BACKWARD, (--)) {
-        #if ENABLE_SPECIALIZATION
+        #if ENABLE_SPECIALIZATION_FT
             if (this_instr->op.code == JUMP_BACKWARD) {
-                uint8_t desired = tstate->interp->jit ? JUMP_BACKWARD_JIT : JUMP_BACKWARD_NO_JIT;
+                uint8_t desired = FT_ATOMIC_LOAD_UINT8(tstate->interp->jit) ? JUMP_BACKWARD_JIT : JUMP_BACKWARD_NO_JIT;
                 FT_ATOMIC_STORE_UINT8_RELAXED(this_instr->op.code, desired);
                 // Need to re-dispatch so the warmup counter isn't off by one:
                 next_instr = this_instr;
@@ -3284,11 +3284,9 @@ dummy_func(
 
         // Only used by Tier 2
         op(_GUARD_NOT_EXHAUSTED_LIST, (iter, null_or_index -- iter, null_or_index)) {
-#ifndef Py_GIL_DISABLED
             PyObject *list_o = PyStackRef_AsPyObjectBorrow(iter);
             assert(Py_TYPE(list_o) == &PyList_Type);
             EXIT_IF((size_t)PyStackRef_UntagInt(null_or_index) >= (size_t)PyList_GET_SIZE(list_o));
-#endif
         }
 
         replaced op(_ITER_NEXT_LIST, (iter, null_or_index -- iter, null_or_index, next)) {
@@ -5281,6 +5279,19 @@ dummy_func(
         }
 
         tier2 op(_CHECK_VALIDITY, (--)) {
+            // For FT:
+            // This doesn't need atomics (for now) as there is only a single time
+            // where a write from another thread is possible:
+            // when a new thread is spawned and it invalidates all current
+            // executors.
+            // The new thread can only be created by an executing uop prior to the
+            // _CHECK_VALIDITY check. New thread creation is synchronized by
+            // locking of the runtime, and the current thread is naturally
+            // paused/waiting for the new thread to be created. Thus,
+            // there is a strict happens-before relation between that
+            // uop's invalidation of validity and this check.
+            // So for now, while the JIT does not run on multiple threads,
+            // it is safe for this to be non-atomic.
             DEOPT_IF(!current_executor->vm_data.valid);
         }
 

Python/ceval.c

@@ -1300,7 +1300,7 @@ _PyTier2Interpreter(
     for (;;) {
         uopcode = next_uop->opcode;
 #ifdef Py_DEBUG
-        if (frame->lltrace >= 3) {
+        if (frame->lltrace >= 4) {
             dump_stack(frame, stack_pointer);
             if (next_uop->opcode == _START_EXECUTOR) {
                 printf("%4d uop: ", 0);

Python/ceval_gil.c

@@ -1395,11 +1395,13 @@ _Py_HandlePending(PyThreadState *tstate)
         _Py_RunGC(tstate);
     }
 
+#ifdef _Py_TIER2
     if ((breaker & _PY_EVAL_JIT_INVALIDATE_COLD_BIT) != 0) {
         _Py_unset_eval_breaker_bit(tstate, _PY_EVAL_JIT_INVALIDATE_COLD_BIT);
         _Py_Executors_InvalidateCold(tstate->interp);
         tstate->interp->executor_creation_counter = JIT_CLEANUP_THRESHOLD;
     }
+#endif
 
     /* GIL drop request */
     if ((breaker & _PY_GIL_DROP_REQUEST_BIT) != 0) {