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

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

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.

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

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.

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.

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?

I don't suppose you have stats for this?

No I do not.

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?

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.

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

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.

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

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.

@markshannon test fixed and verified on my machine.

Android build failure looks unrelated

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-app Bot added the awaiting core review label Jan 1, 2026

bedevere-app Bot 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-it Bot and 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 optimizer Jan 2, 2026

markshannon requested changes Jan 5, 2026

View reviewed changes

Comment thread Include/internal/pycore_tstate.h Outdated

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

bedevere-app Bot 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-app Bot added awaiting merge and removed awaiting changes labels Jan 6, 2026

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

bedevere-app Bot 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)

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

Repository navigation

gh-131798: Support generator frames in the JIT optimizer - #143340

gh-131798: Support generator frames in the JIT optimizer#143340Fidget-Spinner merged 8 commits intopython:mainpython/cpython:mainfrom Fidget-Spinner:cover_more_framesFidget-Spinner/cpython:cover_more_framesCopy head branch name to clipboard

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-app Bot 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

gh-131798: Support generator frames in the JIT optimizer#143340
Fidget-Spinner merged 8 commits into
python:mainfrom
Fidget-Spinner:cover_more_frames

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