Rework sys._jit.is_active to support the tail-calling interpreter

brandtbucher · brandtbucher · commit fc4824b38686 · 2025-05-01T13:27:31.000-07:00
diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h
@@ -206,6 +206,10 @@ struct _ts {
     */
     PyObject *threading_local_sentinel;
     _PyRemoteDebuggerSupport remote_debugger_support;
+
+    // Any frame between the topmost frame where the JIT is currently active,
+    // and the interpreter entry frame below it (used by sys._jit.is_active):
+    struct _PyInterpreterFrame *jit_entry;
 };
 
 /* other API */
diff --git a/Include/internal/pycore_interpframe.h b/Include/internal/pycore_interpframe.h
@@ -146,7 +146,6 @@ _PyFrame_Initialize(
     frame->return_offset = 0;
     frame->owner = FRAME_OWNED_BY_THREAD;
     frame->visited = 0;
-    frame->jit_active = false;
 #ifdef Py_DEBUG
     frame->lltrace = 0;
 #endif
@@ -328,7 +327,6 @@ _PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int
 #endif
     frame->owner = FRAME_OWNED_BY_THREAD;
     frame->visited = 0;
-    frame->jit_active = false;
 #ifdef Py_DEBUG
     frame->lltrace = 0;
 #endif
diff --git a/Include/internal/pycore_interpframe_structs.h b/Include/internal/pycore_interpframe_structs.h
@@ -43,11 +43,11 @@ struct _PyInterpreterFrame {
 #endif
     uint16_t return_offset;  /* Only relevant during a function call */
     char owner;
-    uint8_t visited:1;
-    // Set on interpreter entry frames when the JIT is active:
-    uint8_t jit_active:1;
 #ifdef Py_DEBUG
-    uint8_t lltrace:6;
+    uint8_t visited:1;
+    uint8_t lltrace:7;
+#else
+    uint8_t visited;
 #endif
     /* Locals and stack */
     _PyStackRef localsplus[1];
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -1191,6 +1191,9 @@ dummy_func(
             LOAD_IP(frame->return_offset);
             res = temp;
             LLTRACE_RESUME_FRAME();
+        #if TIER_TWO
+            tstate->jit_entry = frame;
+        #endif
         }
 
         tier1 op(_RETURN_VALUE_EVENT, (val -- val)) {
@@ -1377,6 +1380,9 @@ dummy_func(
             LOAD_IP(1 + INLINE_CACHE_ENTRIES_SEND);
             value = PyStackRef_MakeHeapSafe(temp);
             LLTRACE_RESUME_FRAME();
+        #if TIER_TWO
+            tstate->jit_entry = frame;
+        #endif
         }
 
         tier1 op(_YIELD_VALUE_EVENT, (val -- val)) {
@@ -4962,6 +4968,9 @@ dummy_func(
             RELOAD_STACK();
             res = PyStackRef_FromPyObjectStealMortal((PyObject *)gen);
             LLTRACE_RESUME_FRAME();
+        #if TIER_TWO
+            tstate->jit_entry = frame;
+        #endif
         }
 
         inst(BUILD_SLICE, (args[oparg] -- slice)) {
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -1035,10 +1035,9 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     entry_frame.f_executable = PyStackRef_None;
     entry_frame.instr_ptr = (_Py_CODEUNIT *)_Py_INTERPRETER_TRAMPOLINE_INSTRUCTIONS + 1;
     entry_frame.stackpointer = entry_frame.localsplus;
-    entry_frame.return_offset = 0;
     entry_frame.owner = FRAME_OWNED_BY_INTERPRETER;
     entry_frame.visited = 0;
-    entry_frame.jit_active = false;
+    entry_frame.return_offset = 0;
 #ifdef Py_DEBUG
     entry_frame.lltrace = 0;
 #endif
diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
@@ -355,45 +355,43 @@ _PyFrame_SetStackPointer(frame, stack_pointer)
 /* Tier-switching macros. */
 
 #ifdef _Py_JIT
-#define GOTO_TIER_TWO(EXECUTOR)                        \
-do {                                                   \
-    OPT_STAT_INC(traces_executed);                     \
-    _PyExecutorObject *_executor = (EXECUTOR);         \
-    jit_func jitted = _executor->jit_code;             \
-    /* Keep the shim frame alive via the executor: */  \
-    Py_INCREF(_executor);                              \
-    assert(!entry_frame.jit_active);                   \
-    entry_frame.jit_active = true;                     \
-    next_instr = jitted(frame, stack_pointer, tstate); \
-    assert(entry_frame.jit_active);                    \
-    entry_frame.jit_active = false;                    \
-    Py_DECREF(_executor);                              \
-    Py_CLEAR(tstate->previous_executor);               \
-    frame = tstate->current_frame;                     \
-    stack_pointer = _PyFrame_GetStackPointer(frame);   \
-    if (next_instr == NULL) {                          \
-        next_instr = frame->instr_ptr;                 \
-        JUMP_TO_LABEL(error);                          \
-    }                                                  \
-    DISPATCH();                                        \
+#define GOTO_TIER_TWO(EXECUTOR)                         \
+do {                                                    \
+    OPT_STAT_INC(traces_executed);                      \
+    _PyExecutorObject *_executor = (EXECUTOR);          \
+    jit_func jitted = _executor->jit_code;              \
+    /* Keep the shim frame alive via the executor: */   \
+    Py_INCREF(_executor);                               \
+    _PyInterpreterFrame *jit_entry = tstate->jit_entry; \
+    tstate->jit_entry = frame;                          \
+    next_instr = jitted(frame, stack_pointer, tstate);  \
+    tstate->jit_entry = jit_entry;                      \
+    Py_DECREF(_executor);                               \
+    Py_CLEAR(tstate->previous_executor);                \
+    frame = tstate->current_frame;                      \
+    stack_pointer = _PyFrame_GetStackPointer(frame);    \
+    if (next_instr == NULL) {                           \
+        next_instr = frame->instr_ptr;                  \
+        JUMP_TO_LABEL(error);                           \
+    }                                                   \
+    DISPATCH();                                         \
 } while (0)
 #else
 #define GOTO_TIER_TWO(EXECUTOR) \
 do { \
     OPT_STAT_INC(traces_executed); \
     next_uop = (EXECUTOR)->trace; \
     assert(next_uop->opcode == _START_EXECUTOR); \
-    assert(!entry_frame.jit_active); \
-    entry_frame.jit_active = true; \
+    jit_entry = tstate->jit_entry; \
+    tstate->jit_entry = frame; \
     goto enter_tier_two; \
 } while (0)
 #endif
 
 #define GOTO_TIER_ONE(TARGET)                                         \
     do                                                                \
     {                                                                 \
-        assert(entry_frame.jit_active);                               \
-        entry_frame.jit_active = false;                               \
+        tstate->jit_entry = jit_entry;                                \
         next_instr = (TARGET);                                        \
         OPT_HIST(trace_uop_execution_counter, trace_run_length_hist); \
         _PyFrame_SetStackPointer(frame, stack_pointer);               \
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -1582,6 +1582,8 @@ init_threadstate(_PyThreadStateImpl *_tstate,
 
     tstate->delete_later = NULL;
 
+    tstate->jit_entry = NULL;
+
     llist_init(&_tstate->mem_free_queue);
     llist_init(&_tstate->asyncio_tasks_head);
     if (interp->stoptheworld.requested || _PyRuntime.stoptheworld.requested) {
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
@@ -4033,13 +4033,21 @@ _jit_is_active_impl(PyObject *module)
 /*[clinic end generated code: output=7facca06b10064d4 input=a7e31db659d40a0b]*/
 {
     (void)module;
-    _PyInterpreterFrame *frame = _PyThreadState_GET()->current_frame;
+    PyThreadState *tstate = _PyThreadState_GET();
+    _PyInterpreterFrame *frame = tstate->current_frame;
+    _PyInterpreterFrame *jit_entry_frame = tstate->jit_entry;
     while (true) {
-        frame = frame->previous;
+        // If we hit the JIT's "entry" frame first, we're in JIT code:
+        if (frame == jit_entry_frame) {
+            return true;
+        }
+        // If we hit the interpreter's "entry" frame first, we're not:
         if (frame->owner == FRAME_OWNED_BY_INTERPRETER) {
-            return frame->jit_active;
+            return false;
         }
+        frame = frame->previous;
     }
+    Py_UNREACHABLE();
 }
 
 static PyMethodDef _jit_methods[] = {

Original file line number	Diff line number	Diff line change
`@@ -4033,13 +4033,21 @@ _jit_is_active_impl(PyObject *module)`
`4033`	`4033`	`/[clinic end generated code: output=7facca06b10064d4 input=a7e31db659d40a0b]/`
`4034`	`4034`	`{`
`4035`	`4035`	`(void)module;`
`4036`		`- _PyInterpreterFrame *frame = _PyThreadState_GET()->current_frame;`
	`4036`	`+ PyThreadState *tstate = _PyThreadState_GET();`
	`4037`	`+ _PyInterpreterFrame *frame = tstate->current_frame;`
	`4038`	`+ _PyInterpreterFrame *jit_entry_frame = tstate->jit_entry;`
`4037`	`4039`	`while (true) {`
`4038`		`- frame = frame->previous;`
	`4040`	`+ // If we hit the JIT's "entry" frame first, we're in JIT code:`
	`4041`	`+ if (frame == jit_entry_frame) {`
	`4042`	`+ return true;`
	`4043`	`+ }`
	`4044`	`+ // If we hit the interpreter's "entry" frame first, we're not:`
`4039`	`4045`	`if (frame->owner == FRAME_OWNED_BY_INTERPRETER) {`
`4040`		`- return frame->jit_active;`
	`4046`	`+ return false;`
`4041`	`4047`	`}`
	`4048`	`+ frame = frame->previous;`
`4042`	`4049`	`}`
	`4050`	`+ Py_UNREACHABLE();`
`4043`	`4051`	`}`
`4044`	`4052`
`4045`	`4053`	`static PyMethodDef _jit_methods[] = {`