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

Some tests for the optimizers are disabled for now because their invariants have changed. The main one is the stack space check, as it's no longer valid to deal with underflow.

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. 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 gh-139922: Tail calling for MSVC (VS 2026) #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.

Future ideas:

_DYNAMIC_EXIT 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.

Issue: Trace recording JIT #139109

Fidget-Spinner · 2025-11-07T23:05:46Z

I'm deliberately not merging main in yet, as if we do I can't benchmark this on the meta runners, as those only have clang 19 to build the JIT with. Main has clang 20 (soon 21), which the meta runners don't have.

savannahostrowski

As I said in the group chat...KJ 4 PREZ ⭐ . This is very, very cool and I'm really excited about it. Just a couple comments to clarify my understanding and ask you to inline document some things 😄

savannahostrowski · 2025-11-08T19:06:55Z

.github/workflows/jit.yml

-          - i686-pc-windows-msvc/msvc
-          - x86_64-pc-windows-msvc/msvc
-          - aarch64-pc-windows-msvc/msvc
+# To re-enable later when we support these.


Mostly because I don't know the time horizon for getting this support, can we add a link in the comments to the issue where we're tracking this? I think it's #139922, right?

Mark says he will add normal switch-case support back to this PR by building on top of it I think. So we should get it nearly immediately after this lands.

savannahostrowski · 2025-11-08T19:44:04Z

Python/ceval_macros.h

+#  define IS_JIT_TRACING_MAKING_PROGRESS() (IS_JIT_TRACING() && ((_PyThreadStateImpl *)tstate)->jit_state.prev_state.specialize_counter < MAX_SPECIALIZATION_TRIES)
+#  define ENTER_TRACING() \
+    DISPATCH_TABLE_VAR = TRACING_DISPATCH_TABLE;
+#  define LEAVE_TRACING() \


I do think it might be worth inline documenting some of how the dual dispatch architecture works with regards to ENTER_TRACING/LEAVE_TRACING, since it's pretty novel if IIUC. Correct me if I'm wrong here but I think we're using the same interpreter loop for both normal and tracing.

You're right. I should document this in the internaldocs probably. I'll do it after this PR so we can discuss how to document it properly :)

savannahostrowski · 2025-11-08T19:49:08Z

Include/internal/pycore_uop.h

-// This is the length of the trace we project initially.
-#define UOP_MAX_TRACE_LENGTH 1200
+// This is the length of the trace we translate initially.
+#define UOP_MAX_TRACE_LENGTH 3000


Based on what you said in the PR comment about seeing a lot of "trace too long" aborts even with the higher limit, I'm wondering if you've benchmarked even higher limits like 5k or 10k?

Oops, the trace too long aborts are with the old limit (1200). I have not benchmarked the stats for the newer limit. We should definitely gather stats for this in the future and fine-tune it though

savannahostrowski · 2025-11-08T19:56:33Z

Include/internal/pycore_backoff.h

+// 2. Re-specialized from deopt, in which case the counter will be 1.
+// 3. Deopt -> Specialize -> Deopt -> Specialize, in which case the counter will be 2.
+// We do not want the 3rd case.
+#define MAX_SPECIALIZATION_TRIES 2


I wonder if it's at all worth instrumenting stats for this? Would it be worth understanding how often we're running up against this?

Python/bytecodes.c

savannahostrowski · 2025-11-08T21:23:34Z

Python/optimizer_bytecodes.c

+                REPLACE_OP((this_instr + 1), _NOP, 0, 0);
+            }
+        }
+        if (frame_pop(ctx, returning_code, returning_stacklevel)) {


Should we add stats here to understand how often frame_pop failures (primarily underflow) are causing trace optimization to fail and the trace to be discarded?

There's both overflow and underflow, but yes we should add stats. I'm mainly too lazy to add the pystats because it's currently a little complicated to add a stat to pystats (it's not documented anywhere and also I forgor how to do it)

markshannon

Three things remaining:

I think we should handle instructions that stop tracing, INTERPRETER_EXIT, RAISE_VARARGS, etc. when dispatching. Add a stop_tracing label to the record_previous_inst label in the dispatch table
There are still a few places that check for INTERPRETER_EXIT, that shouldn't need to.
It indicates a bug (at least it does to me) if we need to check for INTERPRETER_EXIT Fixing 1. above should fix this.
There is no need for specialize_counter in the JIT state:

Assuming there are no infinite loops in the current interpreter, we need make no changes to the base interpreter to ensure that there are no infinite loops in it with tracing.
Respecialization can be triggered by adding this to record_previous_inst:

   if _PyOpcode_Caches[opcode] set_backoff_counter(&this_instr[1], 0);

markshannon · 2025-11-10T14:37:42Z

Python/optimizer_analysis.c

+
 #define WITHIN_STACK_BOUNDS() \
-    (STACK_LEVEL() >= 0 && STACK_LEVEL() <= STACK_SIZE())
+    (CURRENT_FRAME_IS_INIT_SHIM() || (STACK_LEVEL() >= 0 && STACK_LEVEL() <= STACK_SIZE()))


The __init__ shim should be a normal frame, so this shouldn't need changing.
What happens if you leave this unchanged?

The problem is not the code object or frame. The problem is that the stack pointer does not correspond to the stack base in between the time the init frame is pushed, to the actual _PUSH_FRAME. A reminder:

_CREATE_INIT_FRAME _PUSH_FRAME

Between _CREATE_INIT_FRAME and _PUSH_FRAME, the stack pointer is the old frame's stack pointer but the stack base is of the init shim frame, thus leading to a misleading segfault.

markshannon · 2025-11-10T14:45:54Z

Python/bytecodes.c


        tier2 op(_DEOPT, (--)) {
-            GOTO_TIER_ONE(_PyFrame_GetBytecode(frame) + CURRENT_TARGET());
+            GOTO_TIER_ONE((frame->owner == FRAME_OWNED_BY_INTERPRETER)


I don't think that _DEOPT should ever see INTERPRETER_EXIT, it is only after a dynamic exit that the next instruction can be INTERPRETER_EXIT

You asked me to end the trace with a _DEOPT if I hit an INTERPRETER_EXIT in a previous review. Which I did. Thus it now _DEOPT exits to INTERPRETER_EXIT.

OK. I think I meant "return to tier 1" more generally, rather than _DEOPT specifically, but this works.
We can tweak it in another PR.

Fidget-Spinner · 2025-11-10T17:02:12Z

Assuming there are no infinite loops in the current interpreter, we need make no changes to the base interpreter to ensure that there are no infinite loops in it with tracing. Respecialization can be triggered by adding this to record_previous_inst:
   if _PyOpcode_Caches[opcode] set_backoff_counter(&this_instr[1], 0);

That isn't true for the FT interpreter. It is possible to enter an infinite specialization loop if we are unlucky enough that LOCK_OBJECT fails every time.

Fidget-Spinner · 2025-11-10T18:23:22Z

   if _PyOpcode_Caches[opcode] set_backoff_counter(&this_instr[1], 0);

So I tried this and we are now hitting bottom in the abstract interpreter, indicating the trace itself is inconsistent and theres a bug somewhere (a trace recording interpreter should never hit bottom). I've narrowed it down to LOAD_ATTR_INSTANCE_VALUE. Probably the specialization and the guard don't agree now. I already found one more here #141367 (comment)

I would really like to make the JIT more fault tolerant and not so dependent on the specializing interpreter. We should definitely fix those bugs, but we should also make the JIT optimizer as resilient as possible

Fidget-Spinner · 2025-11-10T23:27:41Z

Nevermind, I decided to do if _PyOpcode_Caches[opcode] set_backoff_counter(&this_instr[1], 0);. Thanks. We might lose a slight bit of performance from the bad specialization bugs. However, we can fix that in the future.

markshannon · 2025-11-11T13:33:53Z

Assuming there are no infinite loops in the current interpreter, we need make no changes to the base interpreter to ensure that there are no infinite loops in it with tracing. Respecialization can be triggered by adding this to record_previous_inst:
   if _PyOpcode_Caches[opcode] set_backoff_counter(&this_instr[1], 0);
That isn't true for the FT interpreter. It is possible to enter an infinite specialization loop if we are unlucky enough that LOCK_OBJECT fails every time.

If it fails every time, that's an infinite loop regardless. Setting the backoff counter to zero, just sets it to a state that it could be in anyway without tracing.

markshannon

This is starting to look good. I've a few more minor suggestions, but nothing deep.
Hopefully we can get this merged very soon.

Include/internal/pycore_interp_structs.h

markshannon · 2025-11-11T12:01:35Z

Lib/test/test_capi/test_opt.py

        self.assertIn("_JUMP_TO_TOP", uops)
        self.assertIn("_LOAD_FAST_BORROW_0", uops)

+    @unittest.skip("gh-139109 WIP")


How do these skipped tests fail, and how do we get them back again?
I'm OK leaving it to another PR to restore them, but only if there is a plan.

Woops this test works. I accidentally left them them out.

All the stack space optimization is broken. I think a possible fix is to check if the returning code is not matching, then just bail the optimization. That's a conservative and safe way. We'd need to handle YIELD_VALUE too.

Python/optimizer.c

markshannon · 2025-11-11T12:18:02Z

Python/optimizer.c

+    _tstate->jit_tracer_state.initial_state.exit = exit;
+    _tstate->jit_tracer_state.initial_state.stack_depth = curr_stackdepth;
+    _tstate->jit_tracer_state.initial_state.chain_depth = chain_depth;
+    _tstate->jit_tracer_state.prev_state.instr_frame = frame;


Why are we setting prev_state here? There is no previous instruction.
Could we either set instr == NULL or point it to some NOP so that record_previous_inst doesn't record anything.

I now see why you are using the no-tracing form of dispatch after starting a trace.
It would be simpler to set instr == NULL here and use normal dispatch. It would save some code duplication as well.

If we did that it would complicate the side exit tracing case. Right now that is very clean. I don't think we should change that.

When you say "side exit", you mean from _COLD_EXIT?
How would it complicate it?

Say the target of the _COLD_EXIT is at instruction target. We record target at the intiialization. When DISPATCH() at INSTR(target), it records target as the instruction, which is right.

If we set it to NULL, the first DISPATCH() at INSTR(target) is invalid, as there is no previous instruction for the target.

We cannot normal DISPATCH() when returning from tier 1 too, as we have not yet actually executed the instruction. So it will double count it.

If we don't record target at initialization, and DISPATCH() normally, everything works fine I think.
target is recorded at the start of the next instruction as normal.

It tracing we do the following:

Add the prev instruction to the trace

Set the prev to the instruction about to be executed, for the next dispatch to record

Execute the instruction as normal

If we set prev = NOP at a side exit, this should work fine.

Likewise, instead of doing this when we start tracing:

Set the prev to the instruction about to be executed, for the next dispatch to record

Do non-tracing dispatch

we can do this:

Set the prev toNOP/Null

Do tracing dispatch

This might be a tiny bit slower, but it saves a significant amount of code duplication.

Having said that, I think it would be fine to make these changes after this PR is merged.
I'll add a task to the issue.

markshannon · 2025-11-11T13:31:18Z

Python/bytecodes.c


        tier2 op(_DEOPT, (--)) {
-            GOTO_TIER_ONE(_PyFrame_GetBytecode(frame) + CURRENT_TARGET());
+            GOTO_TIER_ONE((frame->owner == FRAME_OWNED_BY_INTERPRETER)


OK. I think I meant "return to tier 1" more generally, rather than _DEOPT specifically, but this works.
We can tweak it in another PR.

Fidget-Spinner · 2025-11-12T14:59:45Z

The results after trusting the interpreter with specialization is slower. 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

This likely indicates bugs in the specializing interpreter. For example, deltablue is suffering from the CALL_LIST_APPEND bug.

I say we just merge the "slower" version though. We should fix the bugs in the interpreter separately.

markshannon · 2025-11-12T15:39:03Z

I'm not worried about the deltablue slowdown. It subclasses a list for no reason, which I don't think we care about.

markshannon · 2025-11-12T15:39:51Z

Is the iOS failure anything to do with this PR or is it failing elsewhere?

markshannon · 2025-11-12T16:57:59Z

Python/bytecodes.c

+            GOTO_TIER_ONE(target);
+        }
+
+        tier2 op(_GUARD_IP__PUSH_FRAME, (ip/4 --)) {


I thought the plan was to generate these exits, otherwise they get out of sync with the source instruction.

What is OFFSET_OF(x)? The offset of what?

If generating the whole uop is too cumbersome because of uop ids and such, then generate the body, leaving the declaration looking like this:

tier2 op(_GUARD_IP__PUSH_FRAME, (ip/4 --)) { IP_GUARD(_PUSH_FRAME); }

It is automatically generating it from the offset of the IP.

I've changed the name to IP_OFFSET_OF to make it clearer.

markshannon · 2025-11-12T16:59:29Z

Python/bytecodes.c

            DISPATCH();
        }

+        label(record_previous_inst) {


NOTE: To support the switch/case interpreter this will need to be an instruction.
No need to do so in this PR.

Python/ceval.c

mhsmith · 2025-11-12T21:05:16Z

Is the iOS failure anything to do with this PR or is it failing elsewhere?

It's been failing intermittently for a few weeks: #141358 (comment).

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 4 commits November 7, 2025 21:08

fix windows

0f978c4

make mypy happy

3c80446

Move to thread state

aaf6873

Update optimizer.c

547f587

Fidget-Spinner added 4 commits November 7, 2025 23:05

Update generated_cases.c.h

a55d766

Support __init__ in the optimizer

278bbe6

Fix a few perf regressions due to tracing thru optimizer

f547880

Some fixups

7e2bc1d

savannahostrowski reviewed Nov 8, 2025

View reviewed changes

markshannon requested changes Nov 10, 2025

View reviewed changes

Fidget-Spinner added 2 commits November 10, 2025 17:02

Clean up labels

fa3e285

fix TC

251e19e

Remove specialize_counter

08ec600

rename jit_state to jit_tracer_state

f7c26d4

markshannon reviewed Nov 11, 2025

View reviewed changes

Fidget-Spinner and others added 3 commits November 11, 2025 15:10

Restore a test, address review

ae1d6fe

remove CALL_LIST_APPEND fix

0658f1b

fix JIT builds

f8a764a

markshannon requested changes Nov 12, 2025

View reviewed changes

Fidget-Spinner added 2 commits November 12, 2025 18:42

Address review

b72ab8b

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

8993e6c

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!

Fidget-Spinner commented Nov 7, 2025

Uh oh!

savannahostrowski left a comment

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!

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!

markshannon left a comment

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!

Fidget-Spinner commented Nov 10, 2025

Uh oh!

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

Uh oh!

Uh oh!

Fidget-Spinner commented Nov 10, 2025

Uh oh!

markshannon commented Nov 11, 2025

Uh oh!

markshannon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Fidget-Spinner Nov 11, 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!

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!

Fidget-Spinner Nov 11, 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!

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

Fidget-Spinner commented Oct 18, 2025 •

edited

Loading

Fidget-Spinner commented Nov 10, 2025 •

edited

Loading

Fidget-Spinner Nov 11, 2025 •

edited

Loading

Fidget-Spinner Nov 11, 2025 •

edited

Loading

Fidget-Spinner commented Nov 12, 2025 •

edited

Loading