Re-enable the JIT

Author: Fidget-Spinner

Include/internal/pycore_opcode_metadata.h

@@ -78,7 +78,9 @@ 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_InvalidateAllLockHeld(PyInterpreterState *interp, int is_invalidation);
 PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyThreadState *tstate);
 PyAPI_FUNC(void) _Py_Executors_ClearExecutorList(PyObject *invalidate, int is_invalidation);
 #else

Include/internal/pycore_optimizer.h

@@ -54,8 +54,7 @@ typedef struct _PyJitTracerState {
 } _PyJitTracerState;
 
 typedef struct _PyJitExecutorState {
-    bool jit;
-    bool compiling;
+    char jit;
     struct _PyExecutorObject *executor_list_head;
     struct _PyExecutorObject *executor_deletion_list_head;
     struct _PyExecutorObject *cold_executor;

Include/internal/pycore_tstate.h

@@ -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):

Lib/test/test_capi/test_opt.py

@@ -2441,7 +2441,9 @@ code_dealloc(PyObject *self)
 #ifdef _Py_TIER2
     _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
 
@@ -3370,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/codeobject.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 = ((_PyThreadStateImpl*)tstate)->jit_executor_state.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 != ((_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");

Python/bytecodes.c

@@ -294,9 +294,13 @@ GETITEM(PyObject *v, Py_ssize_t i) {
 #ifdef Py_GIL_DISABLED
 #  define LOCK_OBJECT(op) PyMutex_LockFast(&(_PyObject_CAST(op))->ob_mutex)
 #  define UNLOCK_OBJECT(op) PyMutex_Unlock(&(_PyObject_CAST(op))->ob_mutex)
+#  define LOCK_OBJECT_SLOW(op) PyMutex_LockFast(&(_PyObject_CAST(op))->ob_mutex)
+#  define UNLOCK_OBJECT_SLOW(op) PyMutex_Unlock(&(_PyObject_CAST(op))->ob_mutex)
 #else
 #  define LOCK_OBJECT(op) (1)
 #  define UNLOCK_OBJECT(op) ((void)0)
+#  define LOCK_OBJECT_SLOW(op) ((void)0)
+#  define UNLOCK_OBJECT_SLOW(op) ((void)0)
 #endif
 
 #define GLOBALS() frame->f_globals

Python/ceval_macros.h

@@ -1785,7 +1785,7 @@ force_instrument_lock_held(PyCodeObject *code, PyInterpreterState *interp)
     if (code->co_executors != NULL) {
         _PyCode_Clear_Executors(code);
     }
-    _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), code, 1);
+    _Py_Executors_InvalidateDependencyLockHeld(_PyInterpreterState_GET(), code, 1);
 #endif
     int code_len = (int)Py_SIZE(code);
     /* Exit early to avoid creating instrumentation
@@ -2028,7 +2028,7 @@ _PyMonitoring_SetEvents(int tool_id, _PyMonitoringEventSet events)
     }
     set_global_version(tstate, new_version);
 #ifdef _Py_TIER2
-    _Py_Executors_InvalidateAll(interp, 1);
+    _Py_Executors_InvalidateAllLockHeld(interp, 1);
 #endif
     return instrument_all_executing_code_objects(interp);
 }