Stowage/cpython
2
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-12-08 06:10:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

gh-139109: A new tracing JIT compiler frontend for CPython (GH-140310)

Browse source 
This PR changes the current JIT model from trace projection to trace recording. Benchmarking: better pyperformance (about 1.7% overall) geomean versus current https://raw.githubusercontent.com/facebookexperimental/free-threading-benchmarking/refs/heads/main/results/bm-20251108-3.15.0a1%2B-7e2bc1d-JIT/bm-20251108-vultr-x86_64-Fidget%252dSpinner-tracing_jit-3.15.0a1%2B-7e2bc1d-vs-base.svg, 100% faster Richards on the most improved benchmark versus the current JIT. Slowdown of about 10-15% on the worst benchmark versus the current JIT. **Note: the fastest version isn't the one merged, as it relies on fixing bugs in the specializing interpreter, which is left to another PR**. The speedup in the merged version is about 1.1%. https://raw.githubusercontent.com/facebookexperimental/free-threading-benchmarking/refs/heads/main/results/bm-20251112-3.15.0a1%2B-f8a764a-JIT/bm-20251112-vultr-x86_64-Fidget%252dSpinner-tracing_jit-3.15.0a1%2B-f8a764a-vs-base.svg

Stats: 50% more uops executed, 30% more traces entered the last time we ran them. It also suggests our trace lengths for a real trace recording JIT are too short, as a lot of trace too long aborts https://github.com/facebookexperimental/free-threading-benchmarking/blob/main/results/bm-20251023-3.15.0a1%2B-eb73378-CLANG%2CJIT/bm-20251023-vultr-x86_64-Fidget%252dSpinner-tracing_jit-3.15.0a1%2B-eb73378-pystats-vs-base.md .

This new JIT frontend is already able to record/execute significantly more instructions than the previous JIT frontend. In this PR, we are now able to record through custom dunders, simple object creation, generators, etc. None of these were done by the old JIT frontend. Some custom dunders uops were discovered to be broken as part of this work gh-140277

The optimizer stack space check is disabled, as it's no longer valid to deal with underflow.

Pros:
* Ignoring the generated tracer code as it's automatically created, this is only additional 1k lines of code. The maintenance burden is handled by the DSL and code generator.
* `optimizer.c` is now significantly simpler, as we don't have to do strange things to recover the bytecode from a trace.
* The new JIT frontend is able to handle a lot more control-flow than the old one.
* Tracing is very low overhead. We use the tail calling interpreter/computed goto interpreter to switch between tracing mode and non-tracing mode. I call this mechanism dual dispatch, as we have two dispatch tables dispatching to each other. Specialization is still enabled while tracing.
* Better handling of polymorphism. We leverage the specializing interpreter for this.

Cons:
* (For now) requires tail calling interpreter or computed gotos. This means no Windows JIT for now :(. Not to fret, tail calling is coming soon to Windows though https://github.com/python/cpython/pull/139962

Design:
* After each instruction, the `record_previous_inst` function/label is executed. This does as the name suggests.
* The tracing interpreter lowers bytecode to uops directly so that it can obtain "fresh" values at the point of lowering.
* The tracing version behaves nearly identical to the normal interpreter, in fact it even has specialization! This allows it to run without much of a slowdown when tracing. The actual cost of tracing is only a function call and writes to memory.
* The tracing interpreter uses the specializing interpreter's deopt to naturally form the side exit chains. This allows it to side exit chain effectively, without repeating much code. We force a re-specializing when tracing a deopt.
* The tracing interpreter can even handle goto errors/exceptions, but I chose to disable them for now as it's not tested.
* Because we do not share interpreter dispatch, there is should be no significant slowdown to the original specializing interpreter on tailcall and computed got with JIT disabled. With JIT enabled, there might be a slowdown in the form of the JIT trying to trace.
* Things that could have dynamic instruction pointer effects are guarded on. The guard deopts to a new instruction --- `_DYNAMIC_EXIT`.
This commit is contained in:
Ken Jin 2025-11-14 02:08:32 +08:00 • committed by GitHub
parent 196f1519cd
commit 4fa80ce74c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 2409 additions and 1065 deletions
Download patch file Download diff file Expand all files Collapse all files

194
Python/bytecodes.c
View file

@ -2938,8 +2938,8 @@ dummy_func(
JUMP_BACKWARD_JIT,
};
tier1 op(_SPECIALIZE_JUMP_BACKWARD, (--)) {
#if ENABLE_SPECIALIZATION_FT
specializing tier1 op(_SPECIALIZE_JUMP_BACKWARD, (--)) {
#if ENABLE_SPECIALIZATION
if (this_instr->op.code == JUMP_BACKWARD) {
uint8_t desired = tstate->interp->jit ? JUMP_BACKWARD_JIT : JUMP_BACKWARD_NO_JIT;
FT_ATOMIC_STORE_UINT8_RELAXED(this_instr->op.code, desired);
@ -2953,25 +2953,21 @@ dummy_func(
tier1 op(_JIT, (--)) {
#ifdef _Py_TIER2
_Py_BackoffCounter counter = this_instr[1].counter;
if (backoff_counter_triggers(counter) && this_instr->op.code == JUMP_BACKWARD_JIT) {
_Py_CODEUNIT *start = this_instr;
/* Back up over EXTENDED_ARGs so optimizer sees the whole instruction */
if (!IS_JIT_TRACING() && backoff_counter_triggers(counter) &&
this_instr->op.code == JUMP_BACKWARD_JIT &&
next_instr->op.code != ENTER_EXECUTOR) {
/* Back up over EXTENDED_ARGs so executor is inserted at the correct place */
_Py_CODEUNIT *insert_exec_at = this_instr;
while (oparg > 255) {
oparg >>= 8;
start--;
insert_exec_at--;
}
_PyExecutorObject *executor;
int optimized = _PyOptimizer_Optimize(frame, start, &executor, 0);
if (optimized <= 0) {
this_instr[1].counter = restart_backoff_counter(counter);
ERROR_IF(optimized < 0);
int succ = _PyJit_TryInitializeTracing(tstate, frame, this_instr, insert_exec_at, next_instr, STACK_LEVEL(), 0, NULL, oparg);
if (succ) {
ENTER_TRACING();
}
else {
this_instr[1].counter = initial_jump_backoff_counter();
assert(tstate->current_executor == NULL);
assert(executor != tstate->interp->cold_executor);
tstate->jit_exit = NULL;
TIER1_TO_TIER2(executor);
this_instr[1].counter = restart_backoff_counter(counter);
}
}
else {
@ -3017,6 +3013,10 @@ dummy_func(
tier1 inst(ENTER_EXECUTOR, (--)) {
#ifdef _Py_TIER2
if (IS_JIT_TRACING()) {
next_instr = this_instr;
goto stop_tracing;
}
PyCodeObject *code = _PyFrame_GetCode(frame);
_PyExecutorObject *executor = code->co_executors->executors[oparg & 255];
assert(executor->vm_data.index == INSTR_OFFSET() - 1);
@ -3078,7 +3078,7 @@ dummy_func(
macro(POP_JUMP_IF_NOT_NONE) = unused/1 + _IS_NONE + _POP_JUMP_IF_FALSE;
tier1 inst(JUMP_BACKWARD_NO_INTERRUPT, (--)) {
replaced inst(JUMP_BACKWARD_NO_INTERRUPT, (--)) {
/* This bytecode is used in the `yield from` or `await` loop.
* If there is an interrupt, we want it handled in the innermost
* generator or coroutine, so we deliberately do not check it here.
@ -5245,19 +5245,40 @@ dummy_func(
tier2 op(_EXIT_TRACE, (exit_p/4 --)) {
_PyExitData *exit = (_PyExitData *)exit_p;
#if defined(Py_DEBUG) && !defined(_Py_JIT)
_Py_CODEUNIT *target = _PyFrame_GetBytecode(frame) + exit->target;
const _Py_CODEUNIT *target = ((frame->owner == FRAME_OWNED_BY_INTERPRETER)
? _Py_INTERPRETER_TRAMPOLINE_INSTRUCTIONS_PTR : _PyFrame_GetBytecode(frame))
+ exit->target;
OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
if (frame->lltrace >= 2) {
if (frame->lltrace >= 3) {
printf("SIDE EXIT: [UOp ");
_PyUOpPrint(&next_uop[-1]);
printf(", exit %tu, temp %d, target %d -> %s, is_control_flow %d]\n",
exit - current_executor->exits, exit->temperature.value_and_backoff,
(int)(target - _PyFrame_GetBytecode(frame)),
_PyOpcode_OpName[target->op.code], exit->is_control_flow);
}
#endif
tstate->jit_exit = exit;
TIER2_TO_TIER2(exit->executor);
}
tier2 op(_DYNAMIC_EXIT, (exit_p/4 --)) {
#if defined(Py_DEBUG) && !defined(_Py_JIT)
_PyExitData *exit = (_PyExitData *)exit_p;
_Py_CODEUNIT *target = frame->instr_ptr;
OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
if (frame->lltrace >= 3) {
printf("DYNAMIC EXIT: [UOp ");
_PyUOpPrint(&next_uop[-1]);
printf(", exit %tu, temp %d, target %d -> %s]\n",
exit - current_executor->exits, exit->temperature.value_and_backoff,
(int)(target - _PyFrame_GetBytecode(frame)),
_PyOpcode_OpName[target->op.code]);
}
#endif
tstate->jit_exit = exit;
TIER2_TO_TIER2(exit->executor);
#endif
// Disabled for now (gh-139109) as it slows down dynamic code tremendously.
// Compile and jump to the cold dynamic executors in the future.
GOTO_TIER_ONE(frame->instr_ptr);
}
tier2 op(_CHECK_VALIDITY, (--)) {
@ -5369,7 +5390,8 @@ dummy_func(
}
tier2 op(_DEOPT, (--)) {
GOTO_TIER_ONE(_PyFrame_GetBytecode(frame) + CURRENT_TARGET());
GOTO_TIER_ONE((frame->owner == FRAME_OWNED_BY_INTERPRETER)
? _Py_INTERPRETER_TRAMPOLINE_INSTRUCTIONS_PTR : _PyFrame_GetBytecode(frame) + CURRENT_TARGET());
}
tier2 op(_HANDLE_PENDING_AND_DEOPT, (--)) {
@ -5399,32 +5421,76 @@ dummy_func(
tier2 op(_COLD_EXIT, ( -- )) {
_PyExitData *exit = tstate->jit_exit;
assert(exit != NULL);
assert(frame->owner < FRAME_OWNED_BY_INTERPRETER);
_Py_CODEUNIT *target = _PyFrame_GetBytecode(frame) + exit->target;
_Py_BackoffCounter temperature = exit->temperature;
if (!backoff_counter_triggers(temperature)) {
exit->temperature = advance_backoff_counter(temperature);
GOTO_TIER_ONE(target);
}
_PyExecutorObject *executor;
if (target->op.code == ENTER_EXECUTOR) {
PyCodeObject *code = _PyFrame_GetCode(frame);
executor = code->co_executors->executors[target->op.arg];
Py_INCREF(executor);
assert(tstate->jit_exit == exit);
exit->executor = executor;
TIER2_TO_TIER2(exit->executor);
}
else {
if (!backoff_counter_triggers(temperature)) {
exit->temperature = advance_backoff_counter(temperature);
GOTO_TIER_ONE(target);
}
_PyExecutorObject *previous_executor = _PyExecutor_FromExit(exit);
assert(tstate->current_executor == (PyObject *)previous_executor);
int chain_depth = previous_executor->vm_data.chain_depth + 1;
int optimized = _PyOptimizer_Optimize(frame, target, &executor, chain_depth);
if (optimized <= 0) {
exit->temperature = restart_backoff_counter(temperature);
GOTO_TIER_ONE(optimized < 0 ? NULL : target);
// For control-flow guards, we don't want to increase the chain depth, as those don't actually
// represent deopts but rather just normal programs!
int chain_depth = previous_executor->vm_data.chain_depth + !exit->is_control_flow;
// Note: it's safe to use target->op.arg here instead of the oparg given by EXTENDED_ARG.
// The invariant in the optimizer is the deopt target always points back to the first EXTENDED_ARG.
// So setting it to anything else is wrong.
int succ = _PyJit_TryInitializeTracing(tstate, frame, target, target, target, STACK_LEVEL(), chain_depth, exit, target->op.arg);
exit->temperature = restart_backoff_counter(exit->temperature);
if (succ) {
GOTO_TIER_ONE_CONTINUE_TRACING(target);
}
exit->temperature = initial_temperature_backoff_counter();
GOTO_TIER_ONE(target);
}
}
tier2 op(_COLD_DYNAMIC_EXIT, ( -- )) {
// TODO (gh-139109): This should be similar to _COLD_EXIT in the future.
_Py_CODEUNIT *target = frame->instr_ptr;
GOTO_TIER_ONE(target);
}
tier2 op(_GUARD_IP__PUSH_FRAME, (ip/4 --)) {
_Py_CODEUNIT *target = frame->instr_ptr + IP_OFFSET_OF(_PUSH_FRAME);
if (target != (_Py_CODEUNIT *)ip) {
frame->instr_ptr += IP_OFFSET_OF(_PUSH_FRAME);
EXIT_IF(true);
}
}
tier2 op(_GUARD_IP_YIELD_VALUE, (ip/4 --)) {
_Py_CODEUNIT *target = frame->instr_ptr + IP_OFFSET_OF(YIELD_VALUE);
if (target != (_Py_CODEUNIT *)ip) {
frame->instr_ptr += IP_OFFSET_OF(YIELD_VALUE);
EXIT_IF(true);
}
}
tier2 op(_GUARD_IP_RETURN_VALUE, (ip/4 --)) {
_Py_CODEUNIT *target = frame->instr_ptr + IP_OFFSET_OF(RETURN_VALUE);
if (target != (_Py_CODEUNIT *)ip) {
frame->instr_ptr += IP_OFFSET_OF(RETURN_VALUE);
EXIT_IF(true);
}
}
tier2 op(_GUARD_IP_RETURN_GENERATOR, (ip/4 --)) {
_Py_CODEUNIT *target = frame->instr_ptr + IP_OFFSET_OF(RETURN_GENERATOR);
if (target != (_Py_CODEUNIT *)ip) {
frame->instr_ptr += IP_OFFSET_OF(RETURN_GENERATOR);
EXIT_IF(true);
}
assert(tstate->jit_exit == exit);
exit->executor = executor;
TIER2_TO_TIER2(exit->executor);
}
label(pop_2_error) {
@ -5571,6 +5637,62 @@ dummy_func(
DISPATCH();
}
label(record_previous_inst) {
#if _Py_TIER2
assert(IS_JIT_TRACING());
int opcode = next_instr->op.code;
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);
if (full) {
LEAVE_TRACING();
int err = stop_tracing_and_jit(tstate, frame);
ERROR_IF(err < 0);
DISPATCH_GOTO_NON_TRACING();
}
// Super instructions. Instruction deopted. There's a mismatch in what the stack expects
// in the optimizer. So we have to reflect in the trace correctly.
_PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
if ((_tstate->jit_tracer_state.prev_state.instr->op.code == CALL_LIST_APPEND &&
opcode == POP_TOP) ||
(_tstate->jit_tracer_state.prev_state.instr->op.code == BINARY_OP_INPLACE_ADD_UNICODE &&
opcode == STORE_FAST)) {
_tstate->jit_tracer_state.prev_state.instr_is_super = true;
}
else {
_tstate->jit_tracer_state.prev_state.instr = next_instr;
}
PyObject *prev_code = PyStackRef_AsPyObjectBorrow(frame->f_executable);
if (_tstate->jit_tracer_state.prev_state.instr_code != (PyCodeObject *)prev_code) {
Py_SETREF(_tstate->jit_tracer_state.prev_state.instr_code, (PyCodeObject*)Py_NewRef((prev_code)));
}
_tstate->jit_tracer_state.prev_state.instr_frame = frame;
_tstate->jit_tracer_state.prev_state.instr_oparg = oparg;
_tstate->jit_tracer_state.prev_state.instr_stacklevel = PyStackRef_IsNone(frame->f_executable) ? 2 : STACK_LEVEL();
if (_PyOpcode_Caches[_PyOpcode_Deopt[opcode]]) {
(&next_instr[1])->counter = trigger_backoff_counter();
}
DISPATCH_GOTO_NON_TRACING();
#else
Py_FatalError("JIT label executed in non-jit build.");
#endif
}
label(stop_tracing) {
#if _Py_TIER2
assert(IS_JIT_TRACING());
int opcode = next_instr->op.code;
_PyJit_translate_single_bytecode_to_trace(tstate, frame, NULL, true);
LEAVE_TRACING();
int err = stop_tracing_and_jit(tstate, frame);
ERROR_IF(err < 0);
DISPATCH_GOTO_NON_TRACING();
#else
Py_FatalError("JIT label executed in non-jit build.");
#endif
}
// END BYTECODES //