Repository navigation
Hot-cold splitting for JIT stencils #143158
Description
Feature or enhancement
Proposal:
We have a textual assembly parser for the stencils. It already knows what blocks are cold and what blocks are hot. With that, it's now not too hard to teach it to section-up blocks.
Currently this is _BINARY_OP_ADD_INT:
// _BINARY_OP_ADD_INT_r23.o: file format elf64-x86-64
//
// Disassembly of section .text:
//
// 0000000000000000 <_JIT_ENTRY>:
// 0: 55 pushq %rbp
// 1: 48 83 ec 10 subq $0x10, %rsp
// 5: 48 89 74 24 08 movq %rsi, 0x8(%rsp)
// a: 48 89 fb movq %rdi, %rbx
// d: 4c 89 fd movq %r15, %rbp
// 10: 4c 89 ff movq %r15, %rdi
// 13: 48 83 e7 fe andq $-0x2, %rdi
// 17: 48 89 de movq %rbx, %rsi
// 1a: 48 83 e6 fe andq $-0x2, %rsi
// 1e: ff 15 00 00 00 00 callq *(%rip) # 0x24 <_JIT_ENTRY+0x24>
// 0000000000000020: R_X86_64_GOTPCRELX _PyCompactLong_Add-0x4
// 24: 48 83 f8 01 cmpq $0x1, %rax
// 28: 75 15 jne 0x3f <_JIT_ENTRY+0x3f>
// 2a: 49 89 ef movq %rbp, %r15
// 2d: 48 89 df movq %rbx, %rdi
// 30: 48 8b 74 24 08 movq 0x8(%rsp), %rsi
// 35: 48 83 c4 10 addq $0x10, %rsp
// 39: 5d popq %rbp
// 3a: e9 00 00 00 00 jmp 0x3f <_JIT_ENTRY+0x3f>
// 000000000000003b: R_X86_64_PLT32 _JIT_JUMP_TARGET-0x4
// 3f: 49 89 c7 movq %rax, %r15
// 42: 48 89 ef movq %rbp, %rdi
// 45: 48 89 de movq %rbx, %rsi
// 48: 48 83 c4 10 addq $0x10, %rsp
// 4c: 5d popq %rbpWith hot-cold splitting, it will be split into:
_BINARY_OP_ADD_INT_r23.HOT:
// 0000000000000000 <_JIT_ENTRY>:
// 0: 55 pushq %rbp
// 1: 48 83 ec 10 subq $0x10, %rsp
// 5: 48 89 74 24 08 movq %rsi, 0x8(%rsp)
// a: 48 89 fb movq %rdi, %rbx
// d: 4c 89 fd movq %r15, %rbp
// 10: 4c 89 ff movq %r15, %rdi
// 13: 48 83 e7 fe andq $-0x2, %rdi
// 17: 48 89 de movq %rbx, %rsi
// 1a: 48 83 e6 fe andq $-0x2, %rsi
// 1e: ff 15 00 00 00 00 callq *(%rip) # 0x24 <_JIT_ENTRY+0x24>
// 0000000000000020: R_X86_64_GOTPCRELX _PyCompactLong_Add-0x4
// 24: 48 83 f8 01 cmpq $0x1, %rax
// 28: 75 15 jne 0x3f <_JIT_ENTRY+0x3f>
// 3f: 49 89 c7 movq %rax, %r15
// 42: 48 89 ef movq %rbp, %rdi
// 45: 48 89 de movq %rbx, %rsi
// 48: 48 83 c4 10 addq $0x10, %rsp
// 4c: 5d popq %rbp
_BINARY_OP_ADD_INT_r23.COLD:
// 2a: 49 89 ef movq %rbp, %r15
// 2d: 48 89 df movq %rbx, %rdi
// 30: 48 8b 74 24 08 movq 0x8(%rsp), %rsi
// 35: 48 83 c4 10 addq $0x10, %rsp
// 39: 5d popq %rbp
// 3a: e9 00 00 00 00 jmp 0x3f <_JIT_ENTRY+0x3f>
// 000000000000003b: R_X86_64_PLT32 _JIT_JUMP_TARGET-0x4Running the current jump inversion and zero length jump removal then gives us:
_BINARY_OP_ADD_INT_r23.HOT:
// 0000000000000000 <_JIT_ENTRY>:
// 0: 55 pushq %rbp
// 1: 48 83 ec 10 subq $0x10, %rsp
// 5: 48 89 74 24 08 movq %rsi, 0x8(%rsp)
// a: 48 89 fb movq %rdi, %rbx
// d: 4c 89 fd movq %r15, %rbp
// 10: 4c 89 ff movq %r15, %rdi
// 13: 48 83 e7 fe andq $-0x2, %rdi
// 17: 48 89 de movq %rbx, %rsi
// 1a: 48 83 e6 fe andq $-0x2, %rsi
// 1e: ff 15 00 00 00 00 callq *(%rip) # 0x24 <_JIT_ENTRY+0x24>
// 0000000000000020: R_X86_64_GOTPCRELX _PyCompactLong_Add-0x4
// 24: 48 83 f8 01 cmpq $0x1, %rax
// 28: 75 15 je _BINARY_OP_ADD_INT_r23.COLD
// 3f: 49 89 c7 movq %rax, %r15
// 42: 48 89 ef movq %rbp, %rdi
// 45: 48 89 de movq %rbx, %rsi
// 48: 48 83 c4 10 addq $0x10, %rsp
// 4c: 5d popq %rbp
_BINARY_OP_ADD_INT_r23.COLD:
// 2a: 49 89 ef movq %rbp, %r15
// 2d: 48 89 df movq %rbx, %rdi
// 30: 48 8b 74 24 08 movq 0x8(%rsp), %rsi
// 35: 48 83 c4 10 addq $0x10, %rsp
// 39: 5d popq %rbp
// 3a: e9 00 00 00 00 jmp 0x3f <_JIT_ENTRY+0x3f>
// 000000000000003b: R_X86_64_PLT32 _JIT_JUMP_TARGET-0x4We then lay out the traces using only the HOT sections and leave the COLD sections at the end. I think this is as good as it gets for machine code flow/layout unless we start writing things by hand.
This builds on #142228.
In the future, to reduce the jitted memory even further, we can de-duplicate common cold stencil fragments. E.g. if we see multiple _BINARY_OP_ADD_INT_r23 in a trace, we can all jump to the common _BINARY_OP_ADD_INT_r23.COLD instead of having one copy for each stencil. That should be a separate PR from this however.
I will work on this.
Has this already been discussed elsewhere?
No response given
Links to previous discussion of this feature:
No response
Activity
corona10 commentedon Dec 25, 2025
Nice idea!
Fidget-Spinner commentedon Dec 25, 2025
On second thought, I'm going to assign this to @markshannon instead, as I can instead complete #142982 quicker
chris-eibl commentedon Dec 26, 2025
There is also data like
in
_LOAD_FAST_CHECK_r*that could be de-duplicated, and data like this is almost always in the cold part.FWIW, this string is also part of
DELETE_FAST_r00, i.e. an "unrelated" opcode, but this is quite rare ...waterlens commentedon Apr 19, 2026
Hi, I'm interested in picking this up if no one is actively working on it. @markshannon, are you still planning to take this on? My background is in compiler engineering and this would be my first CPython contribution.
Here are my thoughts on this issue. The hot/cold block classification is already handled by
_mark_hot_blocksin_optimizers.py. However,_PyJIT_Compileemits each stencil as a single unit into the code section, interleaving hot and cold code across uops. The difficulty of separating cold code into an independent region is that the exact jump offset from a hot block to its cold companion cannot be determined until the hot code region is finalized. My idea is to introduce a new HoleValue for hot-to-cold branch relocations, so these jumps become patchable holes filled in during the emission pass, which is the same wayJUMP_TARGETandERROR_TARGETare resolved.Does this direction make sense? Are there constraints or ongoing work I should be aware of?
I'm open to suggestions if this is already being worked on or there's a different area where help would be more useful. I'm looking for a good entry point into the JIT part of this codebase.
Fidget-Spinner commentedon Apr 19, 2026
@waterlens you're welcome to contribute! Thanks for picking this up.
EDIT: I'M TOTALLY WRONG HERE. SORRY, IT'S BEEN TOO LONG SINCE I'VE SEEN THIS ISSUE. Check my comment below.
I think introducing new relocations might not work. The current "cold" paths are already the external
_JIT_JUMP_TARGET. For example, this is a cold block:The problem here is that we need all the register shuffling and stack adjustments before we jump (tail call) into the
_JIT_JUMP_TARGET. This is mainly because we need to respect thepreserve_nonecalling convention.To get this working, I think there are two possible options:
To elaborate more on the second point,
preserve_alltells LLVM that the cold code will save all registers by themselves, leaving the caller not to require saving registers.So the idea:
preserve_noneones and thepreserve_allversion.This issue is slightly challenging, but very likely to pay off in terms of performance. Let me know if you would like to work on it. If you find it too hard, we can always find another issue for you to work on!
waterlens commentedon Apr 19, 2026
Thanks for the detailed explanation! I see it's not the layout problem I originally assumed. The goal is to use
preserve_allfor cold code so the registers and frame adjustment can sink from the caller to the callee.I'd like to work on this. Just for confirmation:
_EXIT_TRACEand_DEOPT, shouldEXIT_IF,DEOPT_IF, andERROR_IFalso emit cold jumps?tier2_generator.pyto emit a newJUMP_TO_JUMP_TARGET_COLD()for these cases. Is that the right place to make that change?Fidget-Spinner commentedon Apr 19, 2026
The other thing: I think
preserve_allwill not work.@waterlens I think you are on the right track with the initial comment. We need the textual assembly pass to do this. However, I'm not too clear on the implementation. My current suggestion (which may be wrong), is to split each stencil's textual assembly into 1 hot function, and n or more cold functions. This also means
_PyUOpInstructionmust have more than one jump relocation (to the cold function) like you said.If you want to take this up, feel free. This is a tough one!
1 remaining item
waterlens commentedon Apr 25, 2026
@Fidget-Spinner FYI, I have started my work on https://github.com/waterlens/cpython/tree/hot-cold-splitting. Currently I am implementing cold block motion that put code blocks after
L_JIT_CONTINUE, to make later processing easier. Here is an example for_BINARY_OP_ADD_INT_r23:Although I observed that code in cold blocks are quite rare, the performance gain might not be significant. Shall we have more discussion about what we want it to look like in the end?
Fidget-Spinner commentedon Apr 25, 2026
That still looks quite impressive, have you tried out the bigger uops to see if they have more cold code?
waterlens commentedon Apr 26, 2026
Yeah I have tried every uops. For example,
_BINARY_OP_ADD_FLOAT_r23:This is how the uop be like without the hot cold splitting:
Fidget-Spinner commentedon Apr 26, 2026
It looks like you're compiling with a debug build, can you compile the jit stencils without debug and compare?
waterlens commentedon Apr 26, 2026
Thanks for reminding! It's much better now.
_BINARY_OP_ADD_INT_r23before motion:After motion:
I think I can continue working on this. I wonder if you could assign this issue to me. I would have a PR associated to track the progress.
Fidget-Spinner commentedon Apr 27, 2026
I don't think I can assign issues to non repo members, sorry!
waterlens commentedon May 1, 2026
Never mind. I'll create a new draft PR for the issue soon.
markshannon commentedon May 26, 2026
I think we need to do hot/cold splitting in the code generator.
Any future allocation removal/partial escaping analysis/PE pass will need to move compensation code onto exits paths, so it will be too late to handle hot/cold splitting in the machine code generator.
We will want to add compensation code in the optimizer before register allocation, so hot/cold splitting will need to done before generating code for different caching depths.
For example, we want to automatically transform this:
into this:
markshannon commentedon May 26, 2026
Suggested steps:
EXIT_IF(1)intoEXIT()andEXIT_IF(cond)intoif (cond) EXIT()and change code generators to acceptEXITnotEXIT_IFDEOPT_IFEXIT()orDEOPT(), and transitively those BB whose successors are "cold"Fidget-Spinner commentedon May 26, 2026
@markshannon I think you're misunderstanding the hot-cold splitting I'm suggesting in this issue. What I'm suggesting is at a lower-level: that is even after the code generator does hot-cold splitting, there are still certain "cold" fragments left in the machine code due to inefficiencies in the copy and patch generator. This just makes the hot-cold split properly work in the machine code generator. If you look at my original example, the hot-cold split already happens with the code generator's help, it's just that there are register shuffles/inefficient jumps left in the machine code even after processing with the code generator. This is inevitable due to the tail calling convention required by the JIT.
markshannon commentedon May 27, 2026
I'm worried that doing this in the JIT is going to be overly complex, and of limited value.
I don't object to doing it, as long as it doesn't negatively impact anything else.
In the case of
_BINARY_OP_ADD_INT_r23it reduces the size of the hot stencil by 25%, but it doesn't reduce either the overall stencil size, nor the number of instructions executed on the hot path.I'd be curious to see how the two approaches to splitting compare for
_ITER_NEXT_INLINEwhich has a lot of code in the cold path.markshannon commentedon May 27, 2026
Even if we do this, I'd like to do the splitting in the optimizer first.
waterlens commentedon May 27, 2026
Hi everyone, I can create both implementations for comparison. It just takes longer time : )
By the way, could you share how to benchmark or profile the JIT implementation? I found it vague which benchmark set to choose and how to collect data with the benchmarks.
cc @Fidget-Spinner @markshannon.
markshannon commentedon Jun 1, 2026
We will want hot/cold splitting in the optimizer.
We might want hot/cold splitting in the stencils as well.
So we don't need to compare the implementations, just each implementation to the status quo.
@waterlens Don't feel that you have to do hot/cold splitting in the optimizer.
If you just want to do it for the JIT, that's fine.
I'm not sure the speedup will be worth the complexity, but if you want to try it out, feel free to do so.
markshannon commentedon Jun 1, 2026
I've created #150712 to hot-cold splitting in the optimizer