gh-131798: Support generator frames in the JIT optimizer by Fidget-Spinner · Pull Request #143340 · python/cpython

Fidget-Spinner · 2026-01-01T23:40:40Z

This PR covers the iteration of generators in the JIT optimizer.

Benchmark results: 1.5% faster on AArch64 macOS https://github.com/facebookexperimental/free-threading-benchmarking/blob/main/results/bm-20260102-3.15.0a3%2B-8f7b4f4-JIT/bm-20260102-macm4pro-arm64-Fidget%252dSpinner-cover_more_frames-3.15.0a3%2B-8f7b4f4-vs-base.md

0% faster on x86-64 Linux. https://github.com/facebookexperimental/free-threading-benchmarking/blob/main/results/bm-20260102-3.15.0a3%2B-8f7b4f4-JIT/bm-20260102-vultr-x86_64-Fidget%252dSpinner-cover_more_frames-3.15.0a3%2B-8f7b4f4-vs-base.md

I think the macOS results are a little suspicious, maybe the baseline was when the machine was too hot or something? No clue.

Issue: Better uop coverage in the JIT optimizer #131798

markshannon

The changes to optimizer_bytecodes.c look good.

This PR also does loop peeling.
We should probably do that, but in another PR where we can assess its merits and when to peel, and when not.

bedevere-app · 2026-01-05T13:56:08Z

When you're done making the requested changes, leave the comment: I have made the requested changes; please review again.

Fidget-Spinner · 2026-01-05T14:22:09Z

This PR also does loop peeling. We should probably do that, but in another PR where we can assess its merits and when to peel, and when not.

No this doesn't peel. The problem is that generators have a JUMP_BACKWARD in them. If we don't do this, the generators are not traced through and just end early. In fact, none of the generator tests form a loop or pass if I don't add this change. So to test generators, I need this in this PR.

Fidget-Spinner · 2026-01-05T14:26:35Z

Note that we still close a loop when we see it in one iteration, not a peeled iteration. This only affects tracing inner loops, such as when tracing through a generator, or a double nested for loop. In those cases, the inner loop is now traced one more iteration, which is what we want for generators.

markshannon · 2026-01-05T16:27:13Z

Going through a backward jump is equivalent to loop peeling. Even if the loops are strangely shaped.
If we stop tracing on the backwards jump then we should, in theory, still be able to jit generators, just spread across more than one trace. (The first trace stops at the back edge, and the second one extends that back into the caller)
Unfortunately we don't record state on exits, so we would lose a lot of information in the second trace if we did that.

I don't suppose you have stats for this?

Fidget-Spinner · 2026-01-05T16:56:20Z

I don't suppose you have stats for this?

No I do not.

markshannon · 2026-01-05T17:02:45Z

How can you tell the impact of allowing the extra backward edge in the trace?
I suspect that it is beneficial, but some data would be good. Any toy examples with a dump of traces?

Fidget-Spinner · 2026-01-05T17:11:01Z

For this example:

def f(n):
for i in range(n):
yield i + i
def testfunc(n):
for _ in f(n):
pass
import sys
testfunc(4096)
sys._dump_tracelets("hello.gvz")

Before this PR dumps:

digraph ideal {
rankdir = "LR"
executor_0x590ef51207a0 [
shape = none
label = <<table border="0" cellspacing="0">
<tr><td port="start" border="1" ><b>Executor</b></td></tr>
<tr><td border="1" >testfunc: 6</td></tr>
<tr><td port="i0" border="1" >_START_EXECUTOR op0=97920776013728</td></tr>
<tr><td port="i1" border="1" >_MAKE_WARM op0=0</td></tr>
<tr><td port="i2" border="1" >_SET_IP op0=129295294395342</td></tr>
<tr><td port="i3" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i4" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i5" border="1" >_SET_IP op0=129295294395336</td></tr>
<tr><td port="i6" border="1" >_FOR_ITER_GEN_FRAME op0=0</td></tr>
<tr><td port="i7" border="1" >_SPILL_OR_RELOAD op0=129295304702230</td></tr>
<tr><td port="i8" border="1" >_PUSH_FRAME op0=129295294845520</td></tr>
<tr><td port="i9" border="1" >_GUARD_IP__PUSH_FRAME op0=129295304702226</td></tr>
<tr><td port="i10" border="1" >_TIER2_RESUME_CHECK op0=0</td></tr>
<tr><td port="i11" border="1" >_SET_IP op0=129295304702228</td></tr>
<tr><td port="i12" border="1" >_SPILL_OR_RELOAD op0=0</td></tr>
<tr><td port="i13" border="1" >_POP_TOP op0=0</td></tr>
<tr><td port="i14" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i15" border="1" >_SET_IP op0=129295304702230</td></tr>
<tr><td port="i16" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i17" border="1" >_EXIT_TRACE op0=97920776013856</td></tr>
</table>>
]
}

After this PR:

digraph ideal {
rankdir = "LR"
executor_0x63db39341ae0 [
shape = none
label = <<table border="0" cellspacing="0">
<tr><td port="start" border="1" ><b>Executor</b></td></tr>
<tr><td border="1" >testfunc: 6</td></tr>
<tr><td port="i0" border="1" >_START_EXECUTOR op0=109793208703712</td></tr>
<tr><td port="i1" border="1" >_MAKE_WARM op0=0</td></tr>
<tr><td port="i2" border="1" >_SET_IP op0=128718879585230</td></tr>
<tr><td port="i3" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i4" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i5" border="1" >_SET_IP op0=128718879585224</td></tr>
<tr><td port="i6" border="1" >_FOR_ITER_GEN_FRAME op0=0</td></tr>
<tr><td port="i7" border="1" >_SPILL_OR_RELOAD op0=128718889892118</td></tr>
<tr><td port="i8" border="1" >_PUSH_FRAME op0=128718880035408</td></tr>
<tr><td port="i9" border="1" >_GUARD_IP__PUSH_FRAME op0=128718889892114</td></tr>
<tr><td port="i10" border="1" >_TIER2_RESUME_CHECK op0=0</td></tr>
<tr><td port="i11" border="1" >_POP_TOP_NOP op0=0</td></tr>
<tr><td port="i12" border="1" >_SET_IP op0=128718889892118</td></tr>
<tr><td port="i13" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i14" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i15" border="1" >_ITER_CHECK_RANGE op0=0</td></tr>
<tr><td port="i16" border="1" >_GUARD_NOT_EXHAUSTED_RANGE op0=0</td></tr>
<tr><td port="i17" border="1" >_ITER_NEXT_RANGE op0=0</td></tr>
<tr><td port="i18" border="1" >_SET_IP op0=128718889892096</td></tr>
<tr><td port="i19" border="1" >_SPILL_OR_RELOAD op0=0</td></tr>
<tr><td port="i20" border="1" >_STORE_FAST_1 op0=0</td></tr>
<tr><td port="i21" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i22" border="1" >_LOAD_FAST_BORROW_1 op0=0</td></tr>
<tr><td port="i23" border="1" >_LOAD_FAST_BORROW_1 op0=0</td></tr>
<tr><td port="i24" border="1" >_GUARD_TOS_OVERFLOWED op0=0</td></tr>
<tr><td port="i25" border="1" >_BINARY_OP_ADD_INT op0=0</td></tr>
<tr><td port="i26" border="1" >_POP_TOP_NOP op0=0</td></tr>
<tr><td port="i27" border="1" >_POP_TOP_NOP op0=0</td></tr>
<tr><td port="i28" border="1" >_SET_IP op0=128718889892112</td></tr>
<tr><td port="i29" border="1" >_YIELD_VALUE op0=128718880031568</td></tr>
<tr><td port="i30" border="1" >_GUARD_IP_YIELD_VALUE op0=128718879585228</td></tr>
<tr><td port="i31" border="1" >_SET_IP op0=128718879585228</td></tr>
<tr><td port="i32" border="1" >_STORE_FAST_1 op0=0</td></tr>
<tr><td port="i33" border="1" >_JUMP_TO_TOP op0=0</td></tr>
</table>>
]
}

The trace is significantly improved.

Fidget-Spinner · 2026-01-05T17:11:55Z

In the above example, the trace is wrongly cut short because of the JUMP_BACKWARD caused by the inner generator.

markshannon · 2026-01-05T18:01:49Z

I patched main with just the changes to optimizer_bytecodes.c, but not the change to JUMP_BACKWARD handling, and I get this pair of traces (running for a bit longer to warm the side exit):

digraph ideal {
rankdir = "LR"
executor_0xaca35b8e89d0 [
shape = none
label = <<table border="0" cellspacing="0">
<tr><td port="start" border="1" ><b>Executor</b></td></tr>
<tr><td border="1" >f: 3</td></tr>
<tr><td port="i0" border="1" >_START_EXECUTOR op0=189817615714768</td></tr>
<tr><td port="i1" border="1" >_MAKE_WARM op0=0</td></tr>
<tr><td port="i2" border="1" >_SET_IP op0=248217676888342</td></tr>
<tr><td port="i3" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i4" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i5" border="1" >_ITER_CHECK_RANGE op0=0</td></tr>
<tr><td port="i6" border="1" >_GUARD_NOT_EXHAUSTED_RANGE op0=0</td></tr>
<tr><td port="i7" border="1" >_ITER_NEXT_RANGE op0=0</td></tr>
<tr><td port="i8" border="1" >_SET_IP op0=248217676888320</td></tr>
<tr><td port="i9" border="1" >_SWAP_FAST_1 op0=0</td></tr>
<tr><td port="i10" border="1" >_SPILL_OR_RELOAD op0=248217676888336</td></tr>
<tr><td port="i11" border="1" >_POP_TOP op0=0</td></tr>
<tr><td port="i12" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i13" border="1" >_LOAD_FAST_BORROW_1 op0=0</td></tr>
<tr><td port="i14" border="1" >_LOAD_FAST_BORROW_1 op0=0</td></tr>
<tr><td port="i15" border="1" >_GUARD_TOS_OVERFLOWED op0=0</td></tr>
<tr><td port="i16" border="1" >_BINARY_OP_ADD_INT op0=0</td></tr>
<tr><td port="i17" border="1" >_POP_TOP_NOP op0=0</td></tr>
<tr><td port="i18" border="1" >_POP_TOP_NOP op0=0</td></tr>
<tr><td port="i19" border="1" >_SET_IP op0=248217676888336</td></tr>
<tr><td port="i20" border="1" >_YIELD_VALUE op0=248217676203216</td></tr>
<tr><td port="i21" border="1" >_GUARD_IP_YIELD_VALUE op0=248217675777164</td></tr>
<tr><td port="i22" border="1" >_SET_IP op0=248217675777164</td></tr>
<tr><td port="i23" border="1" >_SWAP_FAST_1 op0=0</td></tr>
<tr><td port="i24" border="1" >_POP_TOP op0=0</td></tr>
<tr><td port="i25" border="1" >_EXIT_TRACE op0=189817615714896</td></tr>
</table>>
]
executor_0xaca35b8e89d0:i25 -> executor_0xaca35b8f5fc0:start
executor_0xaca35b8f5fc0 [
shape = none
label = <<table border="0" cellspacing="0">
<tr><td port="start" border="1" ><b>Executor</b></td></tr>
<tr><td border="1" >testfunc: 6</td></tr>
<tr><td port="i0" border="1" >_START_EXECUTOR op0=189817615769536</td></tr>
<tr><td port="i1" border="1" >_MAKE_WARM op0=0</td></tr>
<tr><td port="i2" border="1" >_SET_IP op0=248217675777166</td></tr>
<tr><td port="i3" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i4" border="1" >_CHECK_VALIDITY op0=0</td></tr>
<tr><td port="i5" border="1" >_SET_IP op0=248217675777160</td></tr>
<tr><td port="i6" border="1" >_FOR_ITER_GEN_FRAME op0=0</td></tr>
<tr><td port="i7" border="1" >_SPILL_OR_RELOAD op0=248217676888342</td></tr>
<tr><td port="i8" border="1" >_PUSH_FRAME op0=248217675654416</td></tr>
<tr><td port="i9" border="1" >_GUARD_IP__PUSH_FRAME op0=248217676888338</td></tr>
<tr><td port="i10" border="1" >_TIER2_RESUME_CHECK op0=0</td></tr>
<tr><td port="i11" border="1" >_POP_TOP_NOP op0=0</td></tr>
<tr><td port="i12" border="1" >_SET_IP op0=248217676888342</td></tr>
<tr><td port="i13" border="1" >_CHECK_PERIODIC op0=0</td></tr>
<tr><td port="i14" border="1" >_EXIT_TRACE op0=189817615769664</td></tr>
</table>>
]
executor_0xaca35b8f5fc0:i14 -> executor_0xaca35b8e89d0:start
}

Which covers the same code as the single trace in your example.

So, we don't need the changes to JUMP_BACKWARD to support generator frames, but it does give a better trace in this case.
Given that, could we make the changes to backward edge handling in another PR, as I expect they are more general than just for generators and I would like to keep the changes separate.

Fidget-Spinner · 2026-01-05T18:09:12Z

@markshannon removed the changes and limited it only to the optimizer.

markshannon · 2026-01-06T09:45:11Z

You'll need to update the test, which is a bit awkward as there are two executors to check if you're looking for the _YIELD_VALUE.

We really need a proper test framework for the optimizer.

Fidget-Spinner · 2026-01-06T09:53:33Z

@markshannon test fixed and verified on my machine.

markshannon · 2026-01-06T16:38:05Z

Android build failure looks unrelated

markshannon

Looks good. Thanks.

…nGH-143340)

Cover generator frames in the JIT optimizer
5aee627

Fidget-Spinner requested review from ericsnowcurrently, markshannon and tomasr8 as code owners January 1, 2026 23:40

bedevere-appBot added the awaiting core review label Jan 1, 2026

bedevere-appBot mentioned this pull request Jan 1, 2026
Better uop coverage in the JIT optimizer #131798
Open

Fidget-Spinner added skip news and removed skip news labels Jan 1, 2026

blurb-itBotand others added 3 commits January 1, 2026 23:41

📜🤖 Added by blurb_it.
b1e4adf

end trace earlier when full
c4a6bee

Reduce limit
8f7b4f4

Fidget-Spinner changed the title ~~gh-131798: Optimize generator frames in the JIT optimizer~~gh-131798: Support generator frames in the JIT optimizerJan 2, 2026

markshannon requested changes Jan 5, 2026
View reviewed changes

Comment threadInclude/internal/pycore_tstate.h Outdated

bedevere-appBot removed the awaiting core review label Jan 5, 2026

bedevere-appBot added the awaiting changes label Jan 5, 2026

Move state to elsewhere
13ad45d

Remove changes to JUMP_BACKWARD
5f03d23

Fix test
ce482c5

Merge branch 'main' into cover_more_frames
a2e47b3

markshannon self-requested a review January 6, 2026 16:38

markshannon approved these changes Jan 6, 2026
View reviewed changes

bedevere-appBot added awaiting merge and removed awaiting changes labels Jan 6, 2026

Fidget-Spinner merged commit 90c44bc into python:mainJan 6, 2026
66 of 68 checks passed

bedevere-appBot removed the awaiting merge label Jan 6, 2026

Fidget-Spinner deleted the cover_more_frames branch January 6, 2026 16:40

thunder-coding pushed a commit to thunder-coding/cpython that referenced this pull request Feb 15, 2026
pythongh-131798: Support generator frames in the JIT optimizer (pytho…
10fb802
…nGH-143340)

Uh oh!

Uh oh!

Conversation

Fidget-Spinner commented Jan 1, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

markshannon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bedevere-appBot commented Jan 5, 2026

Uh oh!

Fidget-Spinner commented Jan 5, 2026

Uh oh!

Fidget-Spinner commented Jan 5, 2026

Uh oh!

markshannon commented Jan 5, 2026

Uh oh!

Fidget-Spinner commented Jan 5, 2026

Uh oh!

markshannon commented Jan 5, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fidget-Spinner commented Jan 5, 2026

Uh oh!

Fidget-Spinner commented Jan 5, 2026

Uh oh!

markshannon commented Jan 5, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fidget-Spinner commented Jan 5, 2026

Uh oh!

markshannon commented Jan 6, 2026

Uh oh!

Fidget-Spinner commented Jan 6, 2026

Uh oh!

markshannon commented Jan 6, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

markshannon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Fidget-Spinner commented Jan 1, 2026•
edited
Loading

markshannon commented Jan 5, 2026•
edited
Loading

markshannon commented Jan 5, 2026•
edited
Loading

markshannon commented Jan 6, 2026•
edited
Loading