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

Include/internal/pycore_interp_structs.h

@@ -943,7 +943,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

@@ -48,6 +48,8 @@ typedef struct _PyExitData {
 typedef struct _PyExecutorObject {
  PyObject_VAR_HEAD
  const _PyUOpInstruction *trace;
+ // The interpreter this executor belongs to.
+ PyInterpreterState *interp;
  _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
  uint32_t exit_count;
  uint32_t code_size;

Include/internal/pycore_pylifecycle.h

@@ -129,6 +129,7 @@ PyAPI_FUNC(int) _PyInterpreterConfig_UpdateFromDict(
  PyInterpreterConfig *,
  PyObject *);
 
+extern void _PyInterpreter_SetJitWithEnvVar(const PyConfig *config, PyInterpreterState *interp);
 
 #ifdef __cplusplus
 }

Lib/test/test_capi/test_opt.py

@@ -9,9 +9,14 @@
 
 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,
+ threading_helper,
+ reset_code
+)
 
 _testinternalcapi = import_helper.import_module("_testinternalcapi")
 
@@ -71,8 +76,6 @@ def count_ops(ex, name):
  return len([opname for opname in iter_opnames(ex) if opname == name])
 
 
-@requires_specialization
-@unittest.skipIf(Py_GIL_DISABLED, "optimizer not yet supported in free-threaded builds")
 @requires_jit_enabled
 class TestExecutorInvalidation(unittest.TestCase):
 
@@ -140,8 +143,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):
@@ -462,8 +463,35 @@ 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
+@threading_helper.requires_working_threading()
+@unittest.skipIf(not Py_GIL_DISABLED, "Requires FT and JIT")
+class TestJitFreeThreading(unittest.TestCase):
+ def tests_reenabled_with_multiple_threads(self):
+ import threading
+ def testfunc(x, expected_value):
+ for i in range(x):
+ pass
+
+ ex = get_first_executor(testfunc)
+ self.assertIsNone(ex)
+ # JIT
+ testfunc(TIER2_THRESHOLD+1, True)
+ ex = get_first_executor(testfunc)
+ self.assertIsNotNone(ex)
+ # Spawn threads (turn off the JIT).
+ t = threading.Thread(target=lambda:None, args=())
+ t.start()
+ t.join()
+ # JIT is invalidated after spawning threads.
+ ex = get_first_executor(testfunc)
+ self.assertIsNone(ex)
+ # JIT.
+ testfunc(TIER2_THRESHOLD+1, True)
+ ex = get_first_executor(testfunc)
+ self.assertIsNotNone(ex)
+
+
 @requires_jit_enabled
 @unittest.skipIf(os.getenv("PYTHON_UOPS_OPTIMIZE") == "0", "Needs uop optimizer to run.")
 class TestUopsOptimization(unittest.TestCase):
@@ -2185,6 +2213,7 @@ def testfunc(n):
  self.assertNotIn("_GUARD_TOS_INT", uops)
  self.assertIn("_POP_TOP_NOP", 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.
@@ -2209,6 +2238,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.
@@ -3007,6 +3037,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):
@@ -3020,6 +3051,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.

Objects/codeobject.c

@@ -2433,6 +2433,7 @@ code_dealloc(PyObject *self)
  }
 #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 +3364,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/funcobject.c

@@ -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,19 @@ _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
+#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();
  struct _func_version_cache_item *slot = get_cache_item(interp, version);
  if (slot->code) {
@@ -401,6 +410,10 @@ _PyFunction_LookupByVersion(uint32_t version, PyObject **p_code)
 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;
 }
 

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

@@ -1149,7 +1149,8 @@ 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
+ ASSERT_TYPE_LOCK_HELD();
  PyInterpreterState *interp = _PyInterpreterState_GET();
  // lookup the old version and set to null
  if (tp->tp_version_tag != 0) {
@@ -1158,6 +1159,8 @@ set_version_unlocked(PyTypeObject *tp, unsigned int version)
  + (tp->tp_version_tag % TYPE_VERSION_CACHE_SIZE);
  *slot = NULL;
  }
+#endif
+#ifndef Py_GIL_DISABLED
  if (version) {
  tp->tp_versions_used++;
  }
@@ -1167,7 +1170,7 @@ 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
@@ -1358,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

@@ -2967,9 +2967,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;
@@ -3312,11 +3312,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)) {
@@ -5326,6 +5324,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);
  }