gh-141594: A free-threaded JIT (Part 1)

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,9 +78,10 @@ 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);
 #else
 #  define _Py_Executors_InvalidateDependency(A, B, C) ((void)0)
 #  define _Py_Executors_InvalidateAll(A, B) ((void)0)
@@ -359,10 +365,10 @@ 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, bool stop_tracing);
+int _PyJit_translate_single_bytecode_to_trace(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr, int stop_tracing_exit);
 
 int
 _PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame,

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,10 +48,22 @@ typedef struct _PyJitTracerPreviousState {
 } _PyJitTracerPreviousState;
 
 typedef struct _PyJitTracerState {
+    PyMutex lock;
     _PyUOpInstruction *code_buffer;
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
 } _PyJitTracerState;
+
+typedef struct _PyJitExecutorState {
+    char jit;
+    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;
+} _PyJitExecutorState;
+
 #endif
 
 // Every PyThreadState is actually allocated as a _PyThreadStateImpl. The
@@ -120,6 +132,7 @@ typedef struct _PyThreadStateImpl {
 #endif
 #if _Py_TIER2
     _PyJitTracerState jit_tracer_state;
+    _PyJitExecutorState jit_executor_state;
 #endif
 } _PyThreadStateImpl;
 

Lib/test/test_capi/test_opt.py

@@ -50,8 +50,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):
 
@@ -115,12 +113,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):
@@ -423,8 +417,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):
@@ -1175,6 +1167,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 build only specializes on deferred methods/functions/classes")
     def test_guard_type_version_removed(self):
         def thing(a):
             x = 0
@@ -1193,6 +1186,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 build only specializes on deferred methods/functions/classes")
     def test_guard_type_version_removed_inlined(self):
         """
         Verify that the guard type version if we have an inlined function
@@ -1219,6 +1213,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 build only specializes on deferred methods/functions")
     def test_guard_type_version_removed_invalidation(self):
 
         def thing(a):
@@ -1249,6 +1244,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 build only specializes on deferred methods/functions")
     def test_guard_type_version_removed_escaping(self):
 
         def thing(a):
@@ -1272,6 +1268,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.
@@ -2041,6 +2038,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.
@@ -2065,6 +2063,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.
@@ -2316,6 +2315,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
@@ -2348,6 +2348,7 @@ def testfunc(n):
         self.assertNotIn("_LOAD_SMALL_INT", uops)
         self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
 
+    @unittest.skipIf(Py_GIL_DISABLED, "FT build only specializes on deferred methods/functions")
     def test_cached_attributes(self):
         class C:
             A = 1

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

@@ -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,8 @@ 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_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.

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_CHAR_RELAXED(((_PyThreadStateImpl*)tstate)->jit_executor_state.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;
@@ -3018,11 +3018,25 @@ dummy_func(
                 goto stop_tracing;
             }
             PyCodeObject *code = _PyFrame_GetCode(frame);
+#ifdef Py_GIL_DISABLED
+            LOCK_OBJECT_SLOW(code);
+            _PyExecutorObject *executor = code->co_executors->executors[oparg & 255];
+            // On FT, we are responsible for cleaning up after ourselves.
+            if (!_Py_atomic_load_uint8_relaxed(&executor->vm_data.valid)) {
+                opcode = executor->vm_data.opcode;
+                oparg = (oparg & ~255) | executor->vm_data.oparg;
+                next_instr = this_instr;
+                _Py_ExecutorDetach(executor);
+                UNLOCK_OBJECT_SLOW(code);
+                DISPATCH_GOTO();
+            }
+#else
             _PyExecutorObject *executor = code->co_executors->executors[oparg & 255];
-            assert(executor->vm_data.index == INSTR_OFFSET() - 1);
-            assert(executor->vm_data.code == code);
             assert(executor->vm_data.valid);
+#endif
             assert(tstate->current_executor == NULL);
+            assert(executor->vm_data.index == INSTR_OFFSET() - 1);
+            assert(executor->vm_data.code == code);
             /* If the eval breaker is set then stay in tier 1.
              * This avoids any potentially infinite loops
              * involving _RESUME_CHECK */
@@ -3033,10 +3047,15 @@ dummy_func(
                 if (_PyOpcode_Caches[_PyOpcode_Deopt[opcode]]) {
                     PAUSE_ADAPTIVE_COUNTER(this_instr[1].counter);
                 }
+#ifdef Py_GIL_DISABLED
+                UNLOCK_OBJECT_SLOW(code);
+#endif
                 DISPATCH_GOTO();
             }
-            assert(executor != tstate->interp->cold_executor);
+            assert(executor != ((_PyThreadStateImpl *)tstate)->jit_executor_state.cold_executor);
             tstate->jit_exit = NULL;
+            Py_INCREF(executor);
+            UNLOCK_OBJECT_SLOW(code);
             TIER1_TO_TIER2(executor);
             #else
             Py_FatalError("ENTER_EXECUTOR is not supported in this build");
@@ -3284,11 +3303,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)) {
@@ -5282,7 +5299,7 @@ dummy_func(
         }
 
         tier2 op(_CHECK_VALIDITY, (--)) {
-            DEOPT_IF(!current_executor->vm_data.valid);
+            DEOPT_IF(!FT_ATOMIC_LOAD_UINT8(current_executor->vm_data.valid));
         }
 
         tier2 pure op(_LOAD_CONST_INLINE, (ptr/4 -- value)) {
@@ -5644,7 +5661,7 @@ dummy_func(
             bool stop_tracing = (opcode == WITH_EXCEPT_START ||
                 opcode == RERAISE || opcode == CLEANUP_THROW ||
                 opcode == PUSH_EXC_INFO || opcode == INTERPRETER_EXIT);
-            int full = !_PyJit_translate_single_bytecode_to_trace(tstate, frame, next_instr, stop_tracing);
+            int full = !_PyJit_translate_single_bytecode_to_trace(tstate, frame, next_instr, stop_tracing ? _DEOPT : 0);
             if (full) {
                 LEAVE_TRACING();
                 int err = stop_tracing_and_jit(tstate, frame);
@@ -5684,7 +5701,7 @@ dummy_func(
 #if _Py_TIER2
             assert(IS_JIT_TRACING());
             int opcode = next_instr->op.code;
-            _PyJit_translate_single_bytecode_to_trace(tstate, frame, NULL, true);
+            _PyJit_translate_single_bytecode_to_trace(tstate, frame, NULL, _EXIT_TRACE);
             LEAVE_TRACING();
             int err = stop_tracing_and_jit(tstate, frame);
             ERROR_IF(err < 0);