Hot-cold splitting for JIT stencils

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

Nice idea!

@markshannon

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

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

@markshannon

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.

@waterlens

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

More textual assembly hacks (not preferred).
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:

Create a new jit function pointer type called jit_func_preserve_all
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.
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!

Thanks for the detailed explanation! I see it's not the layout problem I originally assumed. The goal is to use preserve_all for 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:

Besides _EXIT_TRACE and _DEOPT, should EXIT_IF, DEOPT_IF, and ERROR_IF also emit cold jumps?
To wire this up, I'd edit tier2_generator.py to emit a new JUMP_TO_JUMP_TARGET_COLD() for these cases. Is that the right place to make that change?

@waterlens

The other thing: I think preserve_all will 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 _PyUOpInstruction must 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!

@Fidget-Spinner

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

diff --git a/tmp/cpython-jit-stencils-base.7hqjz7//assembly/_BINARY_OP_ADD_INT_r23.s b/tmp/cpython-jit-stencil-diff.J10rNA/assembly-produced-all/_BINARY_OP_ADD_INT_r23.s
index dc0c28ce015..6b77be8ffe5 100644
--- a/tmp/cpython-jit-stencils-base.7hqjz7//assembly/_BINARY_OP_ADD_INT_r23.s
+++ b/tmp/cpython-jit-stencil-diff.J10rNA/assembly-produced-all/_BINARY_OP_ADD_INT_r23.s
@@ -416,14 +416,7 @@ Lloh3:
 ; %bb.12:
 	bl	__PyCompactLong_Add
 	cmp	x0, #1
-	b.ne	LBB0_14
-# JIT: COLD ####################################################################
-; %bb.13:
-	mov	x24, x25
-	mov	x25, x19
-	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
-	b	__JIT_JUMP_TARGET
-# JIT: HOT #####################################################################
+	b.eq	L_JIT_LABEL_0
 LBB0_14:
 	ldrb	w8, [x21, #74]
 	cmp	w8, #3
@@ -597,6 +590,14 @@ Lloh59:
 	mov	w2, #192                        ; =0xc0
 	bl	___assert_rtn
 L_JIT_CONTINUE:
+# JIT: COLD ####################################################################
+__JIT_COLD_START:
+L_JIT_LABEL_0:
+; %bb.13:
+	mov	x24, x25
+	mov	x25, x19
+	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
+	b	__JIT_JUMP_TARGET
 	.loh AdrpLdrGot	Lloh0, Lloh1
 	.loh AdrpLdrGot	Lloh2, Lloh3
 	.loh AdrpLdrGot	Lloh4, Lloh5

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?

That still looks quite impressive, have you tried out the bigger uops to see if they have more cold code?

Yeah I have tried every uops. For example, _BINARY_OP_ADD_FLOAT_r23:

diff --git a/tmp/cpython-jit-stencils-base.7hqjz7//assembly/_BINARY_OP_ADD_FLOAT_r23.s b/tmp/cpython-jit-stencil-diff.J10rNA/assembly-produced-all/_BINARY_OP_ADD_FLOAT_r23.s
index 520dbdb93a3..6100bc71ca0 100644
--- a/tmp/cpython-jit-stencils-base.7hqjz7//assembly/_BINARY_OP_ADD_FLOAT_r23.s
+++ b/tmp/cpython-jit-stencil-diff.J10rNA/assembly-produced-all/_BINARY_OP_ADD_FLOAT_r23.s
@@ -472,23 +472,6 @@ LBB0_20:
 	mov	x26, x19
 	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
 	b	L_JIT_CONTINUE
-# JIT: COLD ####################################################################
-LBB0_21:
-	stp	x25, x19, [x22]
-Lloh6:
-	adrp	x2, l_.str@PAGE
-Lloh7:
-	add	x2, x2, l_.str@PAGEOFF
-	add	x1, x22, #16
-	mov	x0, x21
-	mov	w3, #185                        ; =0xb9
-	bl	__Py_assert_within_stack_bounds
-	add	x22, x22, #16
-	mov	x24, x25
-	mov	x25, x19
-	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
-	b	__JIT_ERROR_TARGET
-# JIT: HOT #####################################################################
 LBB0_22:
 Lloh8:
 	adrp	x0, l___func__.PyStackRef_AsPyObjectBorrow@PAGE
@@ -625,6 +608,23 @@ Lloh61:
 	mov	w2, #502                        ; =0x1f6
 	bl	___assert_rtn
 L_JIT_CONTINUE:
+# JIT: COLD ####################################################################
+__JIT_COLD_START:
+LBB0_21:
+	stp	x25, x19, [x22]
+Lloh6:
+	adrp	x2, l_.str@PAGE
+Lloh7:
+	add	x2, x2, l_.str@PAGEOFF
+	add	x1, x22, #16
+	mov	x0, x21
+	mov	w3, #185                        ; =0xb9
+	bl	__Py_assert_within_stack_bounds
+	add	x22, x22, #16
+	mov	x24, x25
+	mov	x25, x19
+	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
+	b	__JIT_ERROR_TARGET
 	.loh AdrpLdrGot	Lloh0, Lloh1
 	.loh AdrpLdrGot	Lloh2, Lloh3
 	.loh AdrpLdrGot	Lloh4, Lloh5

This is how the uop be like without the hot cold splitting:

	.build_version macos, 15, 7	sdk_version 26, 2
	.section	__TEXT,__text,regular,pure_instructions
	.globl	__JIT_ENTRY                     ; -- Begin function _JIT_ENTRY
	.p2align	2
__JIT_ENTRY:                            ; @_JIT_ENTRY
	.cfi_startproc
; %bb.0:
	stp	x29, x30, [sp, #-16]!           ; 16-byte Folded Spill
	mov	x29, sp
	.cfi_def_cfa w29, 16
	.cfi_offset w30, -8
	.cfi_offset w29, -16
	mov	x19, x25
	mov	x25, x24
	ldrb	w8, [x21, #74]
	cmp	w8, #3
	b.eq	LBB0_6
; %bb.1:
	ldr	x8, [x21]
	cmp	x8, #1
	b.eq	LBB0_25
; %bb.2:
	mvn	w9, w8
	tst	x9, #0x3
	b.eq	LBB0_22
; %bb.3:
	and	x8, x8, #0xfffffffffffffffe
	ldr	x9, [x8, #8]
Lloh0:
	adrp	x10, _PyCode_Type@GOTPAGE
Lloh1:
	ldr	x10, [x10, _PyCode_Type@GOTPAGEOFF]
	cmp	x9, x10
	b.ne	LBB0_26
; %bb.4:
	ldrsw	x9, [x8, #72]
	add	x9, x21, x9, lsl #3
	sub	x9, x22, x9
	sub	x9, x9, #80
	tbnz	x9, #34, LBB0_27
; %bb.5:
	lsr	x9, x9, #3
	ldr	w8, [x8, #64]
	cmp	w8, w9
	b.lt	LBB0_27
LBB0_6:
	mvn	w8, w25
	tst	x8, #0x3
	b.eq	LBB0_22
; %bb.7:
	mvn	w8, w19
	tst	x8, #0x3
	b.eq	LBB0_22
; %bb.8:
	and	x8, x25, #0xfffffffffffffffe
	ldr	x9, [x8, #8]
Lloh2:
	adrp	x10, _PyFloat_Type@GOTPAGE
Lloh3:
	ldr	x10, [x10, _PyFloat_Type@GOTPAGEOFF]
	cmp	x9, x10
	b.ne	LBB0_23
; %bb.9:
	and	x9, x19, #0xfffffffffffffffe
	ldr	x11, [x9, #8]
	cmp	x11, x10
	b.ne	LBB0_24
; %bb.10:
	ldr	d0, [x8, #16]
	ldr	d1, [x9, #16]
	fadd	d0, d0, d1
	bl	_PyFloat_FromDouble
	cbz	x0, LBB0_21
; %bb.11:
	ldrh	w8, [x0, #6]
	and	x8, x8, #0x1
	orr	x24, x8, x0
	and	w8, w24, #0x3
	cmp	w8, #1
	b.eq	LBB0_14
; %bb.12:
	cbnz	w8, LBB0_30
; %bb.13:
	and	x8, x0, #0xfffffffffffffffe
	ldrh	w8, [x8, #6]
	tbnz	w8, #2, LBB0_29
LBB0_14:
	ldrb	w8, [x21, #74]
	cmp	w8, #3
	b.eq	LBB0_20
; %bb.15:
	ldr	x8, [x21]
	cmp	x8, #1
	b.eq	LBB0_25
; %bb.16:
	mvn	w9, w8
	tst	x9, #0x3
	b.eq	LBB0_22
; %bb.17:
	and	x8, x8, #0xfffffffffffffffe
	ldr	x9, [x8, #8]
Lloh4:
	adrp	x10, _PyCode_Type@GOTPAGE
Lloh5:
	ldr	x10, [x10, _PyCode_Type@GOTPAGEOFF]
	cmp	x9, x10
	b.ne	LBB0_26
; %bb.18:
	ldrsw	x9, [x8, #72]
	add	x9, x21, x9, lsl #3
	sub	x9, x22, x9
	sub	x9, x9, #80
	tbnz	x9, #34, LBB0_28
; %bb.19:
	lsr	x9, x9, #3
	ldr	w8, [x8, #64]
	cmp	w8, w9
	b.lt	LBB0_28
LBB0_20:
	mov	x26, x19
	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
	b	L_JIT_CONTINUE
# JIT: COLD ####################################################################
LBB0_21:
	stp	x25, x19, [x22]
Lloh6:
	adrp	x2, l_.str@PAGE
Lloh7:
	add	x2, x2, l_.str@PAGEOFF
	add	x1, x22, #16
	mov	x0, x21
	mov	w3, #185                        ; =0xb9
	bl	__Py_assert_within_stack_bounds
	add	x22, x22, #16
	mov	x24, x25
	mov	x25, x19
	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
	b	__JIT_ERROR_TARGET
# JIT: HOT #####################################################################
LBB0_22:
Lloh8:
	adrp	x0, l___func__.PyStackRef_AsPyObjectBorrow@PAGE
Lloh9:
	add	x0, x0, l___func__.PyStackRef_AsPyObjectBorrow@PAGEOFF
Lloh10:
	adrp	x1, l_.str.8@PAGE
Lloh11:
	add	x1, x1, l_.str.8@PAGEOFF
Lloh12:
	adrp	x3, l_.str.9@PAGE
Lloh13:
	add	x3, x3, l_.str.9@PAGEOFF
	mov	w2, #527                        ; =0x20f
	bl	___assert_rtn
LBB0_23:
Lloh14:
	adrp	x0, l___func__._JIT_ENTRY@PAGE
Lloh15:
	add	x0, x0, l___func__._JIT_ENTRY@PAGEOFF
Lloh16:
	adrp	x1, l_.str@PAGE
Lloh17:
	add	x1, x1, l_.str@PAGEOFF
Lloh18:
	adrp	x3, l_.str.2@PAGE
Lloh19:
	add	x3, x3, l_.str.2@PAGEOFF
	mov	w2, #174                        ; =0xae
	bl	___assert_rtn
LBB0_24:
Lloh20:
	adrp	x0, l___func__._JIT_ENTRY@PAGE
Lloh21:
	add	x0, x0, l___func__._JIT_ENTRY@PAGEOFF
Lloh22:
	adrp	x1, l_.str@PAGE
Lloh23:
	add	x1, x1, l_.str@PAGEOFF
Lloh24:
	adrp	x3, l_.str.3@PAGE
Lloh25:
	add	x3, x3, l_.str.3@PAGEOFF
	mov	w2, #175                        ; =0xaf
	bl	___assert_rtn
LBB0_25:
Lloh26:
	adrp	x0, l___func__._PyFrame_GetCode@PAGE
Lloh27:
	add	x0, x0, l___func__._PyFrame_GetCode@PAGEOFF
Lloh28:
	adrp	x1, l_.str.5@PAGE
Lloh29:
	add	x1, x1, l_.str.5@PAGEOFF
Lloh30:
	adrp	x3, l_.str.6@PAGE
Lloh31:
	add	x3, x3, l_.str.6@PAGEOFF
	mov	w2, #21                         ; =0x15
	bl	___assert_rtn
LBB0_26:
Lloh32:
	adrp	x0, l___func__._PyFrame_GetCode@PAGE
Lloh33:
	add	x0, x0, l___func__._PyFrame_GetCode@PAGEOFF
Lloh34:
	adrp	x1, l_.str.5@PAGE
Lloh35:
	add	x1, x1, l_.str.5@PAGEOFF
Lloh36:
	adrp	x3, l_.str.7@PAGE
Lloh37:
	add	x3, x3, l_.str.7@PAGEOFF
	mov	w2, #23                         ; =0x17
	bl	___assert_rtn
LBB0_27:
Lloh38:
	adrp	x0, l___func__._JIT_ENTRY@PAGE
Lloh39:
	add	x0, x0, l___func__._JIT_ENTRY@PAGEOFF
Lloh40:
	adrp	x1, l_.str@PAGE
Lloh41:
	add	x1, x1, l_.str@PAGEOFF
Lloh42:
	adrp	x3, l_.str.1@PAGE
Lloh43:
	add	x3, x3, l_.str.1@PAGEOFF
	mov	w2, #162                        ; =0xa2
	bl	___assert_rtn
LBB0_28:
Lloh44:
	adrp	x0, l___func__._JIT_ENTRY@PAGE
Lloh45:
	add	x0, x0, l___func__._JIT_ENTRY@PAGEOFF
Lloh46:
	adrp	x1, l_.str@PAGE
Lloh47:
	add	x1, x1, l_.str@PAGEOFF
Lloh48:
	adrp	x3, l_.str.1@PAGE
Lloh49:
	add	x3, x3, l_.str.1@PAGEOFF
	mov	w2, #196                        ; =0xc4
	bl	___assert_rtn
LBB0_29:
Lloh50:
	adrp	x0, l___func__.PyStackRef_CheckValid@PAGE
Lloh51:
	add	x0, x0, l___func__.PyStackRef_CheckValid@PAGEOFF
Lloh52:
	adrp	x1, l_.str.8@PAGE
Lloh53:
	add	x1, x1, l_.str.8@PAGEOFF
Lloh54:
	adrp	x3, l_.str.12@PAGE
Lloh55:
	add	x3, x3, l_.str.12@PAGEOFF
	mov	w2, #497                        ; =0x1f1
	bl	___assert_rtn
LBB0_30:
Lloh56:
	adrp	x0, l___func__.PyStackRef_CheckValid@PAGE
Lloh57:
	add	x0, x0, l___func__.PyStackRef_CheckValid@PAGEOFF
Lloh58:
	adrp	x1, l_.str.8@PAGE
Lloh59:
	add	x1, x1, l_.str.8@PAGEOFF
Lloh60:
	adrp	x3, l_.str.13@PAGE
Lloh61:
	add	x3, x3, l_.str.13@PAGEOFF
	mov	w2, #502                        ; =0x1f6
	bl	___assert_rtn
L_JIT_CONTINUE:
	.loh AdrpLdrGot	Lloh0, Lloh1
	.loh AdrpLdrGot	Lloh2, Lloh3
	.loh AdrpLdrGot	Lloh4, Lloh5
	.loh AdrpAdd	Lloh6, Lloh7
	.loh AdrpAdd	Lloh12, Lloh13
	.loh AdrpAdd	Lloh10, Lloh11
	.loh AdrpAdd	Lloh8, Lloh9
	.loh AdrpAdd	Lloh18, Lloh19
	.loh AdrpAdd	Lloh16, Lloh17
	.loh AdrpAdd	Lloh14, Lloh15
	.loh AdrpAdd	Lloh24, Lloh25
	.loh AdrpAdd	Lloh22, Lloh23
	.loh AdrpAdd	Lloh20, Lloh21
	.loh AdrpAdd	Lloh30, Lloh31
	.loh AdrpAdd	Lloh28, Lloh29
	.loh AdrpAdd	Lloh26, Lloh27
	.loh AdrpAdd	Lloh36, Lloh37
	.loh AdrpAdd	Lloh34, Lloh35
	.loh AdrpAdd	Lloh32, Lloh33
	.loh AdrpAdd	Lloh42, Lloh43
	.loh AdrpAdd	Lloh40, Lloh41
	.loh AdrpAdd	Lloh38, Lloh39
	.loh AdrpAdd	Lloh48, Lloh49
	.loh AdrpAdd	Lloh46, Lloh47
	.loh AdrpAdd	Lloh44, Lloh45
	.loh AdrpAdd	Lloh54, Lloh55
	.loh AdrpAdd	Lloh52, Lloh53
	.loh AdrpAdd	Lloh50, Lloh51
	.loh AdrpAdd	Lloh60, Lloh61
	.loh AdrpAdd	Lloh58, Lloh59
	.loh AdrpAdd	Lloh56, Lloh57
	.cfi_endproc
                                        ; -- End function
	.section	__TEXT,__cstring,cstring_literals
l___func__._JIT_ENTRY:                  ; @__func__._JIT_ENTRY
	.asciz	"_JIT_ENTRY"

l_.str:                                 ; @.str
	.asciz	"_BINARY_OP_ADD_FLOAT_r23.c"

l_.str.1:                               ; @.str.1
	.asciz	"WITHIN_STACK_BOUNDS_IGNORING_CACHE()"

l_.str.2:                               ; @.str.2
	.asciz	"PyFloat_CheckExact(left_o)"

l_.str.3:                               ; @.str.3
	.asciz	"PyFloat_CheckExact(right_o)"

l___func__._PyFrame_GetCode:            ; @__func__._PyFrame_GetCode
	.asciz	"_PyFrame_GetCode"

l_.str.5:                               ; @.str.5
	.asciz	"pycore_interpframe.h"

l_.str.6:                               ; @.str.6
	.asciz	"!PyStackRef_IsNull(f->f_executable)"

l_.str.7:                               ; @.str.7
	.asciz	"PyCode_Check(executable)"

l___func__.PyStackRef_AsPyObjectBorrow: ; @__func__.PyStackRef_AsPyObjectBorrow
	.asciz	"PyStackRef_AsPyObjectBorrow"

l_.str.8:                               ; @.str.8
	.asciz	"pycore_stackref.h"

l_.str.9:                               ; @.str.9
	.asciz	"!PyStackRef_IsTaggedInt(ref)"

l___func__.PyStackRef_CheckValid:       ; @__func__.PyStackRef_CheckValid
	.asciz	"PyStackRef_CheckValid"

l_.str.12:                              ; @.str.12
	.asciz	"!_Py_IsStaticImmortal(obj)"

l_.str.13:                              ; @.str.13
	.asciz	"0"

.subsections_via_symbols

It looks like you're compiling with a debug build, can you compile the jit stencils without debug and compare?

Thanks for reminding! It's much better now.

_BINARY_OP_ADD_INT_r23 before motion:

	.build_version macos, 15, 0	sdk_version 26, 2
	.section	__TEXT,__text,regular,pure_instructions
	.globl	__JIT_ENTRY                     ; -- Begin function _JIT_ENTRY
	.p2align	2
__JIT_ENTRY:                            ; @_JIT_ENTRY
	.cfi_startproc
; %bb.0:
	stp	x29, x30, [sp, #-16]!           ; 16-byte Folded Spill
	mov	x29, sp
	.cfi_def_cfa w29, 16
	.cfi_offset w30, -8
	.cfi_offset w29, -16
	mov	x19, x25
	mov	x25, x24
	and	x0, x24, #0xfffffffffffffffe
	and	x1, x19, #0xfffffffffffffffe
	bl	__PyCompactLong_Add
	cmp	x0, #1
	b.ne	LBB0_2
# JIT: COLD ####################################################################
; %bb.1:
	mov	x24, x25
	mov	x25, x19
	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
	b	__JIT_JUMP_TARGET
# JIT: HOT #####################################################################
LBB0_2:
	mov	x24, x0
	mov	x26, x19
	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
L_JIT_CONTINUE:
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

After motion:

	.build_version macos, 15, 0	sdk_version 26, 2
	.section	__TEXT,__text,regular,pure_instructions
	.globl	__JIT_ENTRY                     ; -- Begin function _JIT_ENTRY
	.p2align	2
__JIT_ENTRY:                            ; @_JIT_ENTRY
	.cfi_startproc
; %bb.0:
	stp	x29, x30, [sp, #-16]!           ; 16-byte Folded Spill
	mov	x29, sp
	.cfi_def_cfa w29, 16
	.cfi_offset w30, -8
	.cfi_offset w29, -16
	mov	x19, x25
	mov	x25, x24
	and	x0, x24, #0xfffffffffffffffe
	and	x1, x19, #0xfffffffffffffffe
	bl	__PyCompactLong_Add
	cmp	x0, #1
	b.eq	L_JIT_LABEL_0
LBB0_2:
	mov	x24, x0
	mov	x26, x19
	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
L_JIT_CONTINUE:
# JIT: COLD ####################################################################
__JIT_COLD_START:
L_JIT_LABEL_0:
; %bb.1:
	mov	x24, x25
	mov	x25, x19
	ldp	x29, x30, [sp], #16             ; 16-byte Folded Reload
	b	__JIT_JUMP_TARGET
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

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.

I don't think I can assign issues to non repo members, sorry!

Never mind. I'll create a new draft PR for the issue soon.

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:

        op (_GUARD_IS_NONE_POP, (val -- )) {
            int is_none = PyStackRef_IsNone(val);
            if (!is_none) {
                PyStackRef_CLOSE(val);
                AT_END_EXIT_IF(1);
            }
            DEAD(val);
        }

into this:

        op (_GUARD_IS_NONE_POP, (val -- )) {
            int is_none = PyStackRef_IsNone(val);
            if (!is_none) {
                goto _GUARD_IS_NONE_POP_cold_1; // tail call
            }
            DEAD(val);
        }

        op (_GUARD_IS_NONE_POP_cold_1, (val -- )) {
              PyStackRef_CLOSE(val);
              AT_END_EXIT_IF(1);
        }

Suggested steps:

Refactor EXIT_IF(1) into EXIT() and EXIT_IF(cond) into if (cond) EXIT() and change code generators to accept EXIT not EXIT_IF
Likewise for DEOPT_IF
In analyser.py mark basic blocks as "cold" which end in EXIT() or DEOPT(), and transitively those BB whose successors are "cold"
Split uops into hot/cold pairs by converting any cold basic blocks in the hot uop into a jump/tailcall to the cold uop, and generating the matching cold uop. Note there might be more than one cold BB that is an entry point to the cold code, in which case we will need multiple cold uops.

@markshannon

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

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_r23 it 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_INLINE which has a lot of code in the cold path.

Even if we do this, I'd like to do the splitting in the optimizer first.

@Fidget-Spinner

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.

@waterlens

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.

I've created #150712 to hot-cold splitting in the optimizer

Navigation Menu

Repository navigation

Hot-cold splitting for JIT stencils #143158

Feature or enhancement

Proposal:

Has this already been discussed elsewhere?

Links to previous discussion of this feature:

Linked PRs

1 remaining item

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Hot-cold splitting for JIT stencils #143158

Description

Feature or enhancement

Proposal:

Has this already been discussed elsewhere?

Links to previous discussion of this feature:

Linked PRs

Activity

corona10 commented on Dec 25, 2025

Fidget-Spinner commented on Dec 25, 2025

chris-eibl commented on Dec 26, 2025

waterlens commented on Apr 19, 2026

Fidget-Spinner commented on Apr 19, 2026

waterlens commented on Apr 19, 2026

Fidget-Spinner commented on Apr 19, 2026

1 remaining item

waterlens commented on Apr 25, 2026

Fidget-Spinner commented on Apr 25, 2026

waterlens commented on Apr 26, 2026

Fidget-Spinner commented on Apr 26, 2026

waterlens commented on Apr 26, 2026

Fidget-Spinner commented on Apr 27, 2026

waterlens commented on May 1, 2026

markshannon commented on May 26, 2026

markshannon commented on May 26, 2026

Fidget-Spinner commented on May 26, 2026

markshannon commented on May 27, 2026

markshannon commented on May 27, 2026

waterlens commented on May 27, 2026

markshannon commented on Jun 1, 2026

markshannon commented on Jun 1, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions