gh-139109: A new tracing JIT compiler frontend for CPython #140310

Fidget-Spinner · 2025-10-18T21:49:17Z

This PR changes the current JIT model from trace projection to trace recording. Benchmarking: 0.6% slower pyperformance versus TC+current JIT, However, 70% faster Richards on the most improved benchmark versus the current JIT. https://github.com/facebookexperimental/free-threading-benchmarking/blob/main/results/bm-20251020-3.15.0a1%2B-ec2971f-CLANG%2CJIT/bm-20251020-vultr-x86_64-Fidget%252dSpinner-tracing_jit-3.15.0a1%2B-ec2971f-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

Some tests for the optimizers are disabled for now not because the optimizer doesn't work, but that it's actually unrolling recursion, thus leading to different results in the trace!

The full test suite passes on my system.

Not close enough to a holiday, so no poem this time!

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. 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 gh-139922: Tail calling for MSVC (VS 2026) #139962
Build times for JIT is now a lot higher as we build two interpreters, not one. computed goto interpreter's build times is very long on my system. Tail call is roughly half the build times of computed goto as there's no blowup of CFG. I recommend building with tail call in the future if you want fast build times.

Design:

We (ab)use the cases generator DSL to duplicate the normal interpreter --- one specializing interpreter, one tracing version.
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.
This was used instead of not duplicating the interpreter as it offers more powerful customization and specialization (see below).
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. This is the power of having another interpreter.
Because we do not share interpreters, there is no measurable slowdown to the original specializing interpreter (on Tailcall only, the design needs to be changed for computed goto). 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 new instruction will cause a re-trace when hot enough starting from the target instruction.
As such, traces "grow" like a web/graph rather than a tree. It starts from a loop, then grows outwards naturally consuming more code around it as the hotspot gets hotter.
I believe this might be the first time trace recording is combined with an inline caching + quickening interpreter (in literature by Brunthaler at least it's called that) to achieve specialization. Let's hope this goes well.

Issue: Trace recording JIT #139109

.github/workflows/jit.yml

Fidget-Spinner · 2025-10-18T23:02:12Z

Closing because this is triggering some UB in CI and wasting resources. Will reopen once I fix it.

markshannon

This is exciting stuff.
I've done a preliminary review. More to come later.

@Fidget-Spinner do you have any ideas as to why the go benchmark is 20% slower?

markshannon · 2025-10-21T13:16:01Z

Include/internal/pycore_interp_structs.h


 typedef _Py_CODEUNIT *(*_PyJitEntryFuncPtr)(struct _PyExecutorObject *exec, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate);

+typedef struct _PyJitState {


All but one of these fields have a 'jit_tracer' prefix. Maybe remove the prefix, so we don't have too many redundant "jit"s in the code.
Maybe rename the struct to PyJitTracerState.

In general, drop "jit" wherever the context makes it clear, to help readability

Include/internal/pycore_optimizer.h

markshannon · 2025-10-21T13:45:45Z

Objects/codeobject.c

        PyMem_Free(co_extra);
    }
 #ifdef _Py_TIER2
+    _Py_JITTracer_InvalidateDependency(tstate, self);


This is correct, but it made me curious.
OOI, did you just spot this, or did the JIT manage to end up running code from a deallocated code object?

The main problem isn't running deallocated code. The problem was that the JIT actually traced then tried to compile code that happened to deallocate while tracing. Thus leading to segfaults while optimizing.

This came up in a few tests from the test suite.

Python/optimizer.c

markshannon · 2025-10-21T13:59:04Z

Python/optimizer_analysis.c

-    if (length == 0) {
-        return length;
-    }
+    optimize_uops(


We may still want to prune the trace in the optimizer. For example, if the optimizer's confidence drops too low.
So leave optimize_uops returning the length.

The current code is wrong. The optimizer never errors. So length can never be 0.

The optimizer currently just returns the length of the unoptimized trace if it doesn't optimize. So this entire code is dead code.

However, I'll leave it there in case you need it in the future.

markshannon · 2025-10-21T15:49:45Z

Python/ceval_macros.h

+    DISPATCH_TABLE_VAR = TRACING_DISPATCH_TABLE;
+#  define LEAVE_TRACING() \
+    DISPATCH_TABLE_VAR = DISPATCH_TABLE;
+#  define BAIL_TRACING_NO_DISPATCH() \


We've already got too many macros here, and it makes the code really hard to follow.
It also creates code bloat as we replicate all this code 256 times.

At the very least, factor out as much as you can into helper functions.
Almost everything in this macro after LEAVE_TRACING() could go in a helper function.

markshannon · 2025-10-21T15:51:24Z

Python/ceval.c

+}
+#endif

 /* _PyEval_EvalFrameDefault is too large to optimize for speed with PGO on MSVC.


Never more true 🙂

markshannon · 2025-10-21T15:51:46Z

Python/ceval_macros.h

 } while (0)

-#define GOTO_TIER_ONE(TARGET)                                         \
+#define GOTO_TIER_ONE(TARGET, SHOULD_CONTINUE_TRACING)                \


SHOULD_CONTINUE_TRACING doesn't seem to be used anywhere

It's used by _COLD_EXIT to tell the interpreter to resume tracing upon exiting to tier 1. This is what allows it to form the side exit trace.

Python/optimizer.c

markshannon · 2025-10-21T16:00:22Z

Python/optimizer.c

-            ADD_TO_TRACE(_SET_IP, 0, (uintptr_t)instr, target);
+    DPRINTF(2, "%p %d: %s(%d) %d\n", old_code, target, _PyOpcode_OpName[opcode], oparg, progress_needed);
+
+    bool needs_guard_ip = _PyOpcode_NeedsGuardIp[opcode] &&


Can we fix up this table, so we don't need all the special cases afterwards?

Fidget-Spinner added 30 commits September 18, 2025 17:00

Basic trace recording

795ef49

WIP generators

40bf6c1

refactor to translate on the go

13188a9

working python startup

e63de39

fix a bug with specializzation

fba9d2d

Fully working bm_generators

7192671

fix jit build

07542dd

Fix exception tracing

021fc44

Fix jump tracing

f886c43

Fix handling of ENTER_EXECUTOR

3066963

Fix ENTER_EXECUTOR bug

20b283b

Fix over-tracing bug

36554a5

fix JIT + debug builds

92bba64

Fix double-initialization

2e3ddc1

fix exception bug

aada168

Fix dispatch_inlined

7b5c655

Fix handling of EXTENDED_ARG

3e9f782

Fix chain depth bug

9a66605

remove printf

108ab7f

fix problem with jumping labels

fac8c74

Point to previous executor when side-exiting

fd3bb48

Fix progress needed and warmup

96b7bb2

Fix unsupported opcode bug, turn off optimizer again

396818b

fix branch tracing

57f417e

fix branch prediction for real

2c603cc

fix non-sstandard C

02f1fb4

Track from JUMP_BACKWARD rather than FOR_ITER

dc414a3

Fix bug where code/func get freed halfway

0ffc2dd

add back replaced, move jit tracing env var to

2032b9c

Handle recursive tracing and CALL_ALLOC_AND_ENTER_INIT

299a068

Fidget-Spinner added 2 commits October 18, 2025 23:21

fix linter and mypy?

87c0b72

more cleanup to fix CI

24cd7f9

efimov-mikhail reviewed Oct 18, 2025

View reviewed changes

.github/workflows/jit.yml Show resolved Hide resolved

Fidget-Spinner closed this Oct 18, 2025

Fidget-Spinner reopened this Oct 18, 2025

Fidget-Spinner closed this Oct 18, 2025

Fidget-Spinner mentioned this pull request Oct 18, 2025

gh-139269: Fix unaligned memory access in JIT code patching functions #139271

Merged

Merge remote-tracking branch 'upstream/main' into tracing_jit

8ae2e4c

Fidget-Spinner reopened this Oct 18, 2025

Fidget-Spinner closed this Oct 18, 2025

Fix lltrace on jit debug builds

f38ef69

Fidget-Spinner reopened this Oct 18, 2025

Fidget-Spinner added 16 commits October 19, 2025 14:53

Turn off tracing on dynamic exit

960d647

Fix _CHECK_PERIODIC insertion

1798ab1

Increase uop length to compensate

681485f

Handle EXTENDED_ARG

9bb03a8

Handle unstable branches

4b26cde

Don't JIT short traces except if they end in a loop

d820e22

revert last 2 changes

00c81fa

Support BINARY_OP_INPLACE_ADD_UNICODE

ba64a5b

Trace through BINARY_OP_SUBSCR_GETITEM

b00252e

Close loops

754b3b7

Specialize on deopt when tracing

6045a67

make mypy happy

ec2971f

remedies against trace explosion

d49e367

lint

55892a4

Fix a bug with where the executors get inserted during EXTENDED_ARG

dd0e16f

Revert remedies against trace explosion

d18c1a1

markshannon reviewed Oct 21, 2025

View reviewed changes


		typedef _Py_CODEUNIT (_PyJitEntryFuncPtr)(struct _PyExecutorObject exec, _PyInterpreterFrame frame, _PyStackRef stack_pointer, PyThreadState tstate);

		typedef struct _PyJitState {

Uh oh!

gh-139109: A new tracing JIT compiler frontend for CPython #140310

Are you sure you want to change the base?

gh-139109: A new tracing JIT compiler frontend for CPython #140310

Conversation

Fidget-Spinner commented Oct 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Fidget-Spinner commented Oct 18, 2025

Uh oh!

markshannon left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

markshannon Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Fidget-Spinner commented Oct 18, 2025 •

edited

Loading

markshannon Oct 21, 2025 •

edited

Loading