JIT compile a trampoline function to avoid shim code

markshannon · markshannon · commit 0dd2c4973060 · 2025-08-19T15:19:28.000+01:00
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -123,6 +123,15 @@ _PyEval_EvalFrame(PyThreadState *tstate, _PyInterpreterFrame *frame, int throwfl
     return tstate->interp->eval_frame(tstate, frame, throwflag);
 }
 
+#if defined(_Py_TIER2) && !defined(_Py_JIT)
+_Py_CODEUNIT *_PyTier2Interpreter(
+    struct _PyExecutorObject *current_executor, _PyInterpreterFrame *frame,
+    _PyStackRef *stack_pointer, PyThreadState *tstate
+);
+#endif
+
+extern _PyJitEntryFuncPtr _Py_jit_entry;
+
 extern PyObject*
 _PyEval_Vector(PyThreadState *tstate,
             PyFunctionObject *func, PyObject *locals,
diff --git a/Include/internal/pycore_interp_structs.h b/Include/internal/pycore_interp_structs.h
@@ -765,6 +765,7 @@ struct _Py_unique_id_pool {
 
 #endif
 
+typedef _Py_CODEUNIT *(*_PyJitEntryFuncPtr)(struct _PyExecutorObject *exec, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate);
 
 /* PyInterpreterState holds the global state for one of the runtime's
    interpreters.  Typically the initial (main) interpreter is the only one.
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -2971,7 +2971,7 @@ dummy_func(
                     assert(tstate->current_executor == NULL);
                     assert(executor != tstate->interp->cold_executor);
                     tstate->jit_exit = NULL;
-                    GOTO_TIER_TWO(executor);
+                    TIER1_TO_TIER2(executor);
                 }
             }
             else {
@@ -3037,7 +3037,7 @@ dummy_func(
             }
             assert(executor != tstate->interp->cold_executor);
             tstate->jit_exit = NULL;
-            GOTO_TIER_TWO(executor);
+            TIER1_TO_TIER2(executor);
             #else
             Py_FatalError("ENTER_EXECUTOR is not supported in this build");
             #endif /* _Py_TIER2 */
@@ -5257,7 +5257,7 @@ dummy_func(
             }
         #endif
             tstate->jit_exit = exit;
-            GOTO_TIER_TWO(exit->executor);
+            TIER2_TO_TIER2(exit->executor);
         }
 
         tier2 op(_CHECK_VALIDITY, (--)) {
@@ -5353,7 +5353,7 @@ dummy_func(
 
         tier2 op(_START_EXECUTOR, (executor/4 --)) {
 #ifndef _Py_JIT
-            current_executor = (_PyExecutorObject*)executor;
+            assert(current_executor == (_PyExecutorObject*)executor);
 #endif
             assert(tstate->jit_exit == NULL || tstate->jit_exit->executor == current_executor);
             tstate->current_executor = (PyObject *)executor;
@@ -5434,7 +5434,7 @@ dummy_func(
             }
             assert(tstate->jit_exit == exit);
             exit->executor = executor;
-            GOTO_TIER_TWO(exit->executor);
+            TIER2_TO_TIER2(exit->executor);
         }
 
         label(pop_2_error) {
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -275,7 +275,8 @@ maybe_lltrace_resume_frame(_PyInterpreterFrame *frame, PyObject *globals)
     }
     int r = PyDict_Contains(globals, &_Py_ID(__lltrace__));
     if (r < 0) {
-        return -1;
+        PyErr_Clear();
+        return 0;
     }
     int lltrace = r * 5;  // Levels 1-4 only trace uops
     if (!lltrace) {
@@ -1109,11 +1110,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 #endif
     }
 
-#if defined(_Py_TIER2) && !defined(_Py_JIT)
-    /* Tier 2 interpreter state */
-    _PyExecutorObject *current_executor = NULL;
-    const _PyUOpInstruction *next_uop = NULL;
-#endif
 #if Py_TAIL_CALL_INTERP
 #   if Py_STATS
         return _TAIL_CALL_start_frame(frame, NULL, tstate, NULL, 0, lastopcode);
@@ -1126,14 +1122,41 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 #endif
 
 
+early_exit:
+    assert(_PyErr_Occurred(tstate));
+    _Py_LeaveRecursiveCallPy(tstate);
+    assert(frame->owner != FRAME_OWNED_BY_INTERPRETER);
+    // GH-99729: We need to unlink the frame *before* clearing it:
+    _PyInterpreterFrame *dying = frame;
+    frame = tstate->current_frame = dying->previous;
+    _PyEval_FrameClearAndPop(tstate, dying);
+    frame->return_offset = 0;
+    assert(frame->owner == FRAME_OWNED_BY_INTERPRETER);
+    /* Restore previous frame and exit */
+    tstate->current_frame = frame->previous;
+    return NULL;
+}
 #ifdef _Py_TIER2
-
-// Tier 2 is also here!
-enter_tier_two:
-
 #ifdef _Py_JIT
-    assert(0);
+_PyJitEntryFuncPtr _Py_jit_entry = _Py_LazyJitTrampoline;
 #else
+_PyJitEntryFuncPtr _Py_jit_entry = _PyTier2Interpreter;
+#endif
+#endif
+
+#if defined(_Py_TIER2) && !defined(_Py_JIT)
+
+_Py_CODEUNIT *
+_PyTier2Interpreter(
+    _PyExecutorObject *current_executor, _PyInterpreterFrame *frame,
+    _PyStackRef *stack_pointer, PyThreadState *tstate
+) {
+    const _PyUOpInstruction *next_uop;
+    int oparg;
+tier2_start:
+
+    next_uop = current_executor->trace;
+    assert(next_uop->opcode == _START_EXECUTOR || next_uop->opcode == _COLD_EXIT);
 
 #undef LOAD_IP
 #define LOAD_IP(UNUSED) (void)0
@@ -1151,7 +1174,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 #undef ENABLE_SPECIALIZATION_FT
 #define ENABLE_SPECIALIZATION_FT 0
 
-    ; // dummy statement after a label, before a declaration
     uint16_t uopcode;
 #ifdef Py_STATS
     int lastuop = 0;
@@ -1225,24 +1247,9 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     next_uop = current_executor->trace + target;
     goto tier2_dispatch;
 
-#endif  // _Py_JIT
-
+}
 #endif // _Py_TIER2
 
-early_exit:
-    assert(_PyErr_Occurred(tstate));
-    _Py_LeaveRecursiveCallPy(tstate);
-    assert(frame->owner != FRAME_OWNED_BY_INTERPRETER);
-    // GH-99729: We need to unlink the frame *before* clearing it:
-    _PyInterpreterFrame *dying = frame;
-    frame = tstate->current_frame = dying->previous;
-    _PyEval_FrameClearAndPop(tstate, dying);
-    frame->return_offset = 0;
-    assert(frame->owner == FRAME_OWNED_BY_INTERPRETER);
-    /* Restore previous frame and exit */
-    tstate->current_frame = frame->previous;
-    return NULL;
-}
 
 #ifdef DO_NOT_OPTIMIZE_INTERP_LOOP
 #  pragma optimize("", on)
diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
@@ -133,9 +133,6 @@ do { \
     _PyFrame_SetStackPointer(frame, stack_pointer); \
     int lltrace = maybe_lltrace_resume_frame(frame, GLOBALS()); \
     stack_pointer = _PyFrame_GetStackPointer(frame); \
-    if (lltrace < 0) { \
-        JUMP_TO_LABEL(exit_unwind); \
-    } \
     frame->lltrace = lltrace; \
 } while (0)
 #else
@@ -354,16 +351,10 @@ _PyFrame_SetStackPointer(frame, stack_pointer)
 
 /* Tier-switching macros. */
 
-#ifdef _Py_JIT
-#define GOTO_TIER_TWO(EXECUTOR)                        \
+#define TIER1_TO_TIER2(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);                              \
-    next_instr = jitted(frame, stack_pointer, tstate); \
-    Py_DECREF(_executor);                              \
+    next_instr = _Py_jit_entry((EXECUTOR), frame, stack_pointer, tstate); \
     frame = tstate->current_frame;                     \
     stack_pointer = _PyFrame_GetStackPointer(frame);   \
     if (next_instr == NULL) {                          \
@@ -372,31 +363,22 @@ do {                                                   \
     }                                                  \
     DISPATCH();                                        \
 } while (0)
-#else
-#define GOTO_TIER_TWO(EXECUTOR) \
-do { \
-    OPT_STAT_INC(traces_executed); \
-    _PyExecutorObject *_executor = (EXECUTOR); \
-    next_uop = _executor->trace; \
-    assert(next_uop->opcode == _START_EXECUTOR || next_uop->opcode == _COLD_EXIT); \
-    goto enter_tier_two; \
+
+#define TIER2_TO_TIER2(EXECUTOR) \
+do {                                                   \
+    OPT_STAT_INC(traces_executed);                     \
+    current_executor = (EXECUTOR);                     \
+    next_uop = current_executor->trace;                \
+    goto tier2_start;                                  \
 } while (0)
-#endif
 
 #define GOTO_TIER_ONE(TARGET)                                         \
     do                                                                \
     {                                                                 \
         tstate->current_executor = NULL;                              \
-        next_instr = (TARGET);                                        \
         OPT_HIST(trace_uop_execution_counter, trace_run_length_hist); \
         _PyFrame_SetStackPointer(frame, stack_pointer);               \
-        stack_pointer = _PyFrame_GetStackPointer(frame);              \
-        if (next_instr == NULL)                                       \
-        {                                                             \
-            next_instr = frame->instr_ptr;                            \
-            goto error;                                               \
-        }                                                             \
-        DISPATCH();                                                   \
+        return TARGET;                                                \
     } while (0)
 
 #define CURRENT_OPARG()    (next_uop[-1].oparg)
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/jit.c b/Python/jit.c
@@ -571,6 +571,65 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     return 0;
 }
 
+// Compiles executor in-place. Don't forget to call _PyJIT_Free later!
+static _PyJitEntryFuncPtr
+compile_trampoline(void)
+{
+    _PyExecutorObject dummy;
+    const StencilGroup *group;
+    // Loop once to find the total compiled size:
+    size_t code_size = 0;
+    size_t data_size = 0;
+    jit_state state = {0};
+    group = &trampoline;
+    code_size += group->code_size;
+    data_size += group->data_size;
+    combine_symbol_mask(group->trampoline_mask, state.trampolines.mask);
+    // Round up to the nearest page:
+    size_t page_size = get_page_size();
+    assert((page_size & (page_size - 1)) == 0);
+    size_t code_padding = DATA_ALIGN - ((code_size + state.trampolines.size) & (DATA_ALIGN - 1));
+    size_t padding = page_size - ((code_size + state.trampolines.size + code_padding + data_size) & (page_size - 1));
+    size_t total_size = code_size + state.trampolines.size + code_padding + data_size + padding;
+    unsigned char *memory = jit_alloc(total_size);
+    if (memory == NULL) {
+        return NULL;
+    }
+    // Loop again to emit the code:
+    unsigned char *code = memory;
+    state.trampolines.mem = memory + code_size;
+    unsigned char *data = memory + code_size + state.trampolines.size + code_padding;
+    // Compile the shim, which handles converting between the native
+    // calling convention and the calling convention used by jitted code
+    // (which may be different for efficiency reasons).
+    group = &trampoline;
+    group->emit(code, data, &dummy, NULL, &state);
+    code += group->code_size;
+    data += group->data_size;
+    assert(code == memory + code_size);
+    assert(data == memory + code_size + state.trampolines.size + code_padding + data_size);
+    if (mark_executable(memory, total_size)) {
+        jit_free(memory, total_size);
+        return NULL;
+    }
+    return (_PyJitEntryFuncPtr)memory;
+}
+
+
+_Py_CODEUNIT *
+_Py_LazyJitTrampoline(
+    _PyExecutorObject *executor, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate
+) {
+    assert(_Py_jit_entry == _Py_LazyJitTrampoline);
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    _PyJitEntryFuncPtr trampoline = compile_trampoline();
+    if (trampoline == NULL) {
+        Py_FatalError("Cannot allocate core JIT code");
+    }
+    _Py_jit_entry = trampoline;
+    return trampoline(executor, frame, stack_pointer, tstate);
+}
+
 void
 _PyJIT_Free(_PyExecutorObject *executor)
 {
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -494,6 +494,11 @@ free_interpreter(PyInterpreterState *interp)
 static inline int check_interpreter_whence(long);
 #endif
 
+extern _Py_CODEUNIT *
+_Py_LazyJitTrampoline(
+    struct _PyExecutorObject *exec, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate
+);
+
 /* Get the interpreter state to a minimal consistent state.
    Further init happens in pylifecycle.c before it can be used.
    All fields not initialized here are expected to be zeroed out,
@@ -574,6 +579,8 @@ init_interpreter(PyInterpreterState *interp,
         /* Fix the self-referential, statically initialized fields. */
         interp->dtoa = (struct _dtoa_state)_dtoa_state_INIT(interp);
     }
+
+
 #if !defined(Py_GIL_DISABLED) && defined(Py_STACKREF_DEBUG)
     interp->next_stackref = INITIAL_STACKREF_INDEX;
     _Py_hashtable_allocator_t alloc = {
diff --git a/Tools/jit/_targets.py b/Tools/jit/_targets.py
@@ -193,6 +193,8 @@ async def _build_stencils(self) -> dict[str, _stencils.StencilGroup]:
             async with asyncio.TaskGroup() as group:
                 coro = self._compile("shim", TOOLS_JIT / "shim.c", work)
                 tasks.append(group.create_task(coro, name="shim"))
+                coro = self._compile("trampoline", TOOLS_JIT / "trampoline.c", work)
+                tasks.append(group.create_task(coro, name="trampoline"))
                 template = TOOLS_JIT_TEMPLATE_C.read_text()
                 for case, opname in cases_and_opnames:
                     # Write out a copy of the template with *only* this case
diff --git a/Tools/jit/_writer.py b/Tools/jit/_writer.py
@@ -23,10 +23,11 @@ def _dump_footer(
     yield "} StencilGroup;"
     yield ""
     yield f"static const StencilGroup shim = {groups['shim'].as_c('shim')};"
+    yield f"static const StencilGroup trampoline = {groups['trampoline'].as_c('trampoline')};"
     yield ""
     yield "static const StencilGroup stencil_groups[MAX_UOP_ID + 1] = {"
     for opname, group in sorted(groups.items()):
-        if opname == "shim":
+        if opname in ("shim", "trampoline"):
             continue
         yield f"    [{opname}] = {group.as_c(opname)},"
     yield "};"
diff --git a/Tools/jit/template.c b/Tools/jit/template.c

Original file line number	Diff line number	Diff line change
`@@ -2971,7 +2971,7 @@ dummy_func(`
`2971`	`2971`	`assert(tstate->current_executor == NULL);`
`2972`	`2972`	`assert(executor != tstate->interp->cold_executor);`
`2973`	`2973`	`tstate->jit_exit = NULL;`
`2974`		`- GOTO_TIER_TWO(executor);`
	`2974`	`+ TIER1_TO_TIER2(executor);`
`2975`	`2975`	`}`
`2976`	`2976`	`}`
`2977`	`2977`	`else {`
`@@ -3037,7 +3037,7 @@ dummy_func(`
`3037`	`3037`	`}`
`3038`	`3038`	`assert(executor != tstate->interp->cold_executor);`
`3039`	`3039`	`tstate->jit_exit = NULL;`
`3040`		`- GOTO_TIER_TWO(executor);`
	`3040`	`+ TIER1_TO_TIER2(executor);`
`3041`	`3041`	`#else`
`3042`	`3042`	`Py_FatalError("ENTER_EXECUTOR is not supported in this build");`
`3043`	`3043`	`#endif /* _Py_TIER2 */`
`@@ -5257,7 +5257,7 @@ dummy_func(`
`5257`	`5257`	`}`
`5258`	`5258`	`#endif`
`5259`	`5259`	`tstate->jit_exit = exit;`
`5260`		`- GOTO_TIER_TWO(exit->executor);`
	`5260`	`+ TIER2_TO_TIER2(exit->executor);`
`5261`	`5261`	`}`
`5262`	`5262`
`5263`	`5263`	`tier2 op(_CHECK_VALIDITY, (--)) {`
`@@ -5353,7 +5353,7 @@ dummy_func(`
`5353`	`5353`
`5354`	`5354`	`tier2 op(_START_EXECUTOR, (executor/4 --)) {`
`5355`	`5355`	`#ifndef _Py_JIT`
`5356`		`- current_executor = (_PyExecutorObject*)executor;`
	`5356`	`+ assert(current_executor == (_PyExecutorObject*)executor);`
`5357`	`5357`	`#endif`
`5358`	`5358`	`assert(tstate->jit_exit == NULL \|\| tstate->jit_exit->executor == current_executor);`
`5359`	`5359`	`tstate->current_executor = (PyObject *)executor;`
`@@ -5434,7 +5434,7 @@ dummy_func(`
`5434`	`5434`	`}`
`5435`	`5435`	`assert(tstate->jit_exit == exit);`
`5436`	`5436`	`exit->executor = executor;`
`5437`		`- GOTO_TIER_TWO(exit->executor);`
	`5437`	`+ TIER2_TO_TIER2(exit->executor);`
`5438`	`5438`	`}`
`5439`	`5439`
`5440`	`5440`	`label(pop_2_error) {`