gh-141594: A free-threaded JIT

Include/cpython/code.h

@@ -16,6 +16,7 @@ typedef struct {
 } _PyCoCached;
 
 typedef struct {
+    uint8_t is_finalizing;
     int size;
     int capacity;
     struct _PyExecutorObject *executors[1];

Include/internal/pycore_interp_structs.h

@@ -930,14 +930,6 @@ struct _is {
     struct types_state types;
     struct callable_cache callable_cache;
     PyObject *common_consts[NUM_COMMON_CONSTANTS];
-    bool jit;
-    bool compiling;
-    struct _PyExecutorObject *executor_list_head;
-    struct _PyExecutorObject *executor_deletion_list_head;
-    struct _PyExecutorObject *cold_executor;
-    struct _PyExecutorObject *cold_dynamic_executor;
-    int executor_deletion_list_remaining_capacity;
-    size_t executor_creation_counter;
     _rare_events rare_events;
     PyDict_WatchCallback builtins_dict_watcher;
 

Include/internal/pycore_optimizer.h

@@ -24,7 +24,11 @@ typedef struct _PyExecutorLinkListNode {
 typedef struct {
     uint8_t opcode;
     uint8_t oparg;
+#ifdef Py_GIL_DISABLED
+    uint8_t valid;
+#else
     uint8_t valid:1;
+#endif
     uint8_t linked:1;
     uint8_t chain_depth:6;  // Must be big enough for MAX_CHAIN_DEPTH - 1.
     bool warm;
@@ -45,6 +49,7 @@ typedef struct _PyExitData {
 
 typedef struct _PyExecutorObject {
     PyObject_VAR_HEAD
+    PyThreadState *tstate;
     const _PyUOpInstruction *trace;
     _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
     uint32_t exit_count;
@@ -73,14 +78,16 @@ PyAPI_FUNC(void) _Py_Executor_DependsOn(_PyExecutorObject *executor, void *obj);
 
 #ifdef _Py_TIER2
 PyAPI_FUNC(void) _Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is_invalidation);
+PyAPI_FUNC(void) _Py_Executors_InvalidateDependencyLockHeld(PyInterpreterState *interp, void *obj, int is_invalidation);
 PyAPI_FUNC(void) _Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation);
-PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyInterpreterState *interp);
-
+PyAPI_FUNC(void) _Py_Executors_InvalidateAllLockHeld(PyInterpreterState *interp, int is_invalidation);
+PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyThreadState *tstate);
+PyAPI_FUNC(void) _Py_Executors_InvalidateColdGC(PyInterpreterState *interp);
 #else
 #  define _Py_Executors_InvalidateDependency(A, B, C) ((void)0)
 #  define _Py_Executors_InvalidateAll(A, B) ((void)0)
 #  define _Py_Executors_InvalidateCold(A) ((void)0)
-
+#  define _Py_Executors_InvalidateColdGC(A) ((void)0)
 #endif
 
 // Used as the threshold to trigger executor invalidation when
@@ -359,7 +366,7 @@ extern void _PyExecutor_Free(_PyExecutorObject *self);
 
 PyAPI_FUNC(int) _PyDumpExecutors(FILE *out);
 #ifdef _Py_TIER2
-extern void _Py_ClearExecutorDeletionList(PyInterpreterState *interp);
+extern void _Py_ClearExecutorDeletionList(PyThreadState *tstate);
 #endif
 
 int _PyJit_translate_single_bytecode_to_trace(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr, int stop_tracing_opcode);
@@ -372,8 +379,6 @@ _PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame,
 
 void _PyJit_FinalizeTracing(PyThreadState *tstate);
 
-void _PyJit_Tracer_InvalidateDependency(PyThreadState *old_tstate, void *obj);
-
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_tstate.h

@@ -35,7 +35,7 @@ typedef struct _PyJitTracerInitialState {
 } _PyJitTracerInitialState;
 
 typedef struct _PyJitTracerPreviousState {
-    bool dependencies_still_valid;
+    uint8_t dependencies_still_valid;
     bool instr_is_super;
     int code_max_size;
     int code_curr_size;
@@ -48,11 +48,25 @@ typedef struct _PyJitTracerPreviousState {
 } _PyJitTracerPreviousState;
 
 typedef struct _PyJitTracerState {
+    PyMutex lock;
     _PyUOpInstruction *code_buffer;
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
 } _PyJitTracerState;
+
+#endif
+
+typedef struct _PyJitExecutorState {
+    char jit;
+#if _Py_TIER2
+    struct _PyExecutorObject *executor_list_head;
+    struct _PyExecutorObject *executor_deletion_list_head;
+    struct _PyExecutorObject *cold_executor;
+    struct _PyExecutorObject *cold_dynamic_executor;
+    int executor_deletion_list_remaining_capacity;
+    size_t executor_creation_counter;
 #endif
+} _PyJitExecutorState;
 
 // Every PyThreadState is actually allocated as a _PyThreadStateImpl. The
 // PyThreadState fields are exposed as part of the C API, although most fields
@@ -121,6 +135,7 @@ typedef struct _PyThreadStateImpl {
 #if _Py_TIER2
     _PyJitTracerState jit_tracer_state;
 #endif
+    _PyJitExecutorState jit_executor_state;
 } _PyThreadStateImpl;
 
 #ifdef __cplusplus

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")
 
@@ -60,9 +64,8 @@ def iter_opnames(ex):
 def get_opnames(ex):
     return list(iter_opnames(ex))
 
+FT_TEST_DISABLED_REVERSE_TYPE_CACHE = "FT build reverse type cache is disabled for now"
 
-@requires_specialization
-@unittest.skipIf(Py_GIL_DISABLED, "optimizer not yet supported in free-threaded builds")
 @requires_jit_enabled
 class TestExecutorInvalidation(unittest.TestCase):
 
@@ -126,12 +129,8 @@ def f():
         self.assertTrue(exe.is_valid())
         sys._clear_internal_caches()
         self.assertFalse(exe.is_valid())
-        exe = get_first_executor(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 +433,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):
@@ -1186,6 +1183,7 @@ def test_modified_local_is_seen_by_optimized_code(self):
         self.assertIs(type(s), float)
         self.assertEqual(s, 1024.0)
 
+    @unittest.skipIf(Py_GIL_DISABLED, FT_TEST_DISABLED_REVERSE_TYPE_CACHE)
     def test_guard_type_version_removed(self):
         def thing(a):
             x = 0
@@ -1204,6 +1202,7 @@ class Foo:
         guard_type_version_count = opnames.count("_GUARD_TYPE_VERSION")
         self.assertEqual(guard_type_version_count, 1)
 
+    @unittest.skipIf(Py_GIL_DISABLED, FT_TEST_DISABLED_REVERSE_TYPE_CACHE)
     def test_guard_type_version_removed_inlined(self):
         """
         Verify that the guard type version if we have an inlined function
@@ -1230,6 +1229,7 @@ class Foo:
         guard_type_version_count = opnames.count("_GUARD_TYPE_VERSION")
         self.assertEqual(guard_type_version_count, 1)
 
+    @unittest.skipIf(Py_GIL_DISABLED, FT_TEST_DISABLED_REVERSE_TYPE_CACHE)
     def test_guard_type_version_removed_invalidation(self):
 
         def thing(a):
@@ -1260,6 +1260,7 @@ class Bar:
         self.assertEqual(opnames[:load_attr_top].count("_GUARD_TYPE_VERSION"), 1)
         self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 2)
 
+    @unittest.skipIf(Py_GIL_DISABLED, FT_TEST_DISABLED_REVERSE_TYPE_CACHE)
     def test_guard_type_version_removed_escaping(self):
 
         def thing(a):
@@ -1283,6 +1284,7 @@ class Foo:
         self.assertEqual(opnames[:load_attr_top].count("_GUARD_TYPE_VERSION"), 1)
         self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 2)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build only specializes on deferred methods/functions/classes")
     def test_guard_type_version_executor_invalidated(self):
         """
         Verify that the executor is invalided on a type change.
@@ -2052,6 +2054,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 +2079,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.
@@ -2327,6 +2331,7 @@ def testfunc(n):
         self.assertNotIn("_TO_BOOL_BOOL", uops)
         self.assertIn("_GUARD_IS_TRUE_POP", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build only specializes on deferred methods/functions/classes")
     def test_set_type_version_sets_type(self):
         class C:
             A = 1
@@ -2359,6 +2364,7 @@ def testfunc(n):
         self.assertNotIn("_LOAD_SMALL_INT", uops)
         self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, FT_TEST_DISABLED_REVERSE_TYPE_CACHE)
     def test_cached_attributes(self):
         class C:
             A = 1
@@ -2497,6 +2503,7 @@ def testfunc(n):
 
         self.assertIn("_POP_TOP_NOP", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build immortalizes constants")
     def test_pop_top_specialize_int(self):
         def testfunc(n):
             for _ in range(n):
@@ -2510,6 +2517,7 @@ def testfunc(n):
 
         self.assertIn("_POP_TOP_INT", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build immortalizes constants")
     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 @@
+Add free-threading support to the JIT. 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

@@ -2185,6 +2185,13 @@ _PyCode_ReturnsOnlyNone(PyCodeObject *co)
 static void
 clear_executors(PyCodeObject *co)
 {
+#ifdef Py_GIL_DISABLED
+    uint8_t expected = 0;
+    if (!_Py_atomic_compare_exchange_uint8(&co->co_executors->is_finalizing, &expected, 1)) {
+        // Another thread is already finalizing this.
+        return;
+    }
+#endif
     assert(co->co_executors);
     for (int i = 0; i < co->co_executors->size; i++) {
         if (co->co_executors->executors[i]) {
@@ -2432,9 +2439,11 @@ 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) {
+        Py_BEGIN_CRITICAL_SECTION(co);
         clear_executors(co);
+        Py_END_CRITICAL_SECTION();
     }
 #endif
 
@@ -3363,8 +3372,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) {
@@ -1153,7 +1153,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.