Hot-cold splitting for JIT stencils

[Fidget-Spinner]

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 %rbp

With 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-0x4

Running 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-0x4

We 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

Linked PRs

• gh-143158: Hot cold code splitting for JIT compiler #149292

[corona10]

Nice idea!

[Fidget-Spinner]

On second thought, I'm going to assign this to @markshannon instead, as I can instead complete #142982 quicker

[chris-eibl]

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.

There is also data like

 // 0: "cannot access local variable '%s' where it is not associated with a value\x00"

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]

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_blocks in _optimizers.py. However, _PyJIT_Compile emits 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 way JUMP_TARGET and ERROR_TARGET are 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]

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

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

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 the preserve_none calling convention.

To get this working, I think there are two possible options:

1. More textual assembly hacks (not preferred).
2. Experiment with the preserve_all calling convention. (preferred)

To elaborate more on the second point, preserve_all tells LLVM that the cold code will save all registers by themselves, leaving the caller not to require saving registers.

So the idea:

1. Create a new jit function pointer type called jit_func_preserve_all
2. For all JIT exits (uops starting with _EXIT_TRACE or _DEOPT), we compile/generate two versions. a. the current preserve_none ones and the preserve_all version.
3. Mark all _JIT_JUMP_TARGETS as a new PATCH_JUMP_COLD, and create a new DECLARE_TARGET_COLD.

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!