Load target address earlier for tail call interpreter

[aconz2]

My working branch is here main...aconz2:cpython:aconz2/early-tail-call-load

I saw the recent merge of the tail call interpreter (#128718), very nice! I have played with this style of interpreter before and one thing that comes up is when to calculate the target address. As it is, the current interpreter does it in DISPATCH() by doing

DEF_TARGET(foo) {
    // ...
    TAIL return INSTRUCTION_TABLE[opcode](ARGS);
}

this results in assembly like:

0000000000289580 <_TAIL_CALL_GET_LEN>:
  289580: 50                           	push	rax
  289581: 89 fb                        	mov	ebx, edi
  289583: 4d 89 7c 24 38               	mov	qword ptr [r12 + 0x38], r15
  289588: 49 83 c7 02                  	add	r15, 0x2
  28958c: 49 8b 7d f8                  	mov	rdi, qword ptr [r13 - 0x8]
  289590: 4d 89 6c 24 40               	mov	qword ptr [r12 + 0x40], r13
  289595: e8 f6 fc ea ff               	call	0x139290 <PyObject_Size>
  28959a: 4d 8b 6c 24 40               	mov	r13, qword ptr [r12 + 0x40]
  28959f: 49 c7 44 24 40 00 00 00 00   	mov	qword ptr [r12 + 0x40], 0x0
  2895a8: 48 85 c0                     	test	rax, rax
  2895ab: 78 2b                        	js	0x2895d8 <_TAIL_CALL_GET_LEN+0x58>
  2895ad: 48 89 c7                     	mov	rdi, rax
  2895b0: e8 eb 54 ec ff               	call	0x14eaa0 <PyLong_FromSsize_t>
  2895b5: 48 85 c0                     	test	rax, rax
  2895b8: 74 1e                        	je	0x2895d8 <_TAIL_CALL_GET_LEN+0x58>
  2895ba: 49 89 45 00                  	mov	qword ptr [r13], rax
  2895be: 49 83 c5 08                  	add	r13, 0x8
  2895c2: 41 0f b7 3f                  	movzx	edi, word ptr [r15]  #<-- Load next_instr
  2895c6: 40 0f b6 c7                  	movzx	eax, dil             #<-- grab opcode
  2895ca: c1 ef 08                     	shr	edi, 0x8
  2895cd: 48 8d 0d 7c 50 1f 00         	lea	rcx, [rip + 0x1f507c]   # 0x47e650 <INSTRUCTION_TABLE>
  2895d4: 5a                           	pop	rdx
  2895d5: ff 24 c1                     	jmp	qword ptr [rcx + 8*rax] #<-- jmp with addr calculation
  2895d8: 89 df                        	mov	edi, ebx
  2895da: 58                           	pop	rax
  2895db: e9 30 dc ff ff               	jmp	0x287210 <_TAIL_CALL_error>

where we jmp to a computed adress which is dependent on the lea and a memory load a few instructions prior.

Another method looks like

DEF_TARGET(foo) {
  // ...
  tail_funcptr next_f = INSTRUCTION_TABLE[next_opcode];
  // ...
  TAIL return next_f(ARGS);
}

where we try to get the compiler to compute the target earlier and then have a jmp reg. We have to pay special attention to places where next_instr is modified and reload the pointer (though hopefully the optimizer will just wait to do the calculation until the latest place).

In this early branch, I was able to get this working enough to see what asm it would generate. For _TAIL_CALL_GET_LEN, the sequence now looks like

00000000002896b0 <_TAIL_CALL_GET_LEN>:
  2896b0: 55                           	push	rbp
  2896b1: 89 fb                        	mov	ebx, edi
  2896b3: 4d 89 7c 24 38               	mov	qword ptr [r12 + 0x38], r15
  2896b8: 41 0f b6 47 02               	movzx	eax, byte ptr [r15 + 0x2]  #<-- Load next instr opcode
  2896bd: 49 83 c7 02                  	add	r15, 0x2
  2896c1: 48 8d 0d 88 5f 1f 00         	lea	rcx, [rip + 0x1f5f88]   # 0x47f650 <INSTRUCTION_TABLE>
  2896c8: 48 8b 2c c1                  	mov	rbp, qword ptr [rcx + 8*rax]  #<-- load next target addr
  2896cc: 49 8b 7d f8                  	mov	rdi, qword ptr [r13 - 0x8]
  2896d0: 4d 89 6c 24 40               	mov	qword ptr [r12 + 0x40], r13
  2896d5: e8 b6 fb ea ff               	call	0x139290 <PyObject_Size>
  2896da: 4d 8b 6c 24 40               	mov	r13, qword ptr [r12 + 0x40]
  2896df: 49 c7 44 24 40 00 00 00 00   	mov	qword ptr [r12 + 0x40], 0x0
  2896e8: 48 85 c0                     	test	rax, rax
  2896eb: 78 20                        	js	0x28970d <_TAIL_CALL_GET_LEN+0x5d>
  2896ed: 48 89 c7                     	mov	rdi, rax
  2896f0: e8 ab 53 ec ff               	call	0x14eaa0 <PyLong_FromSsize_t>
  2896f5: 48 85 c0                     	test	rax, rax
  2896f8: 74 13                        	je	0x28970d <_TAIL_CALL_GET_LEN+0x5d>
  2896fa: 49 89 45 00                  	mov	qword ptr [r13], rax
  2896fe: 49 83 c5 08                  	add	r13, 0x8
  289702: 41 0f b6 7f 01               	movzx	edi, byte ptr [r15 + 0x1]
  289707: 48 89 e8                     	mov	rax, rbp                  #<-- register rename
  28970a: 5d                           	pop	rbp
  28970b: ff e0                        	jmp	rax                       #<-- jmp to target addr
  28970d: 89 df                        	mov	edi, ebx
  28970f: 5d                           	pop	rbp
  289710: e9 fb da ff ff               	jmp	0x287210 <_TAIL_CALL_error>
  289715: 66 66 2e 0f 1f 84 00 00 00 00 00     	nop	word ptr cs:[rax + rax]
  2896c1: 48 8d 0d 88 5f 1f 00         	lea	rcx, [rip + 0x1f5f88]   # 0x47f650 <INSTRUCTION_TABLE>
  2896c8: 48 8b 2c c1                  	mov	rbp, qword ptr [rcx + 8*rax]

Specifically in this case, both PyObject_Size and PyLong_FromSsize_t don't touch rbp so there isn't any additional register pressure. But I haven't looked extensively so may not be universally true.

My theory is that this could be better for the CPU because in this example once it gets back from PyLong_FromSsize_t, the jump target is already in a register and could maybe prefetch better.

Have not benchmarked anything yet.

Looking at another example _TAIL_CALL_BINARY_OP_SUBSCR_GETITEM, this does a LOAD_IP() towards the end so we have to reload our target address. It does seem like the optimizer is smart enough to avoid double loading, but this just ends up with an almost identical ending:

# main
  28eba5: 41 0f b7 3f                   movzx   edi, word ptr [r15]
  28eba9: 40 0f b6 cf                   movzx   ecx, dil
  28ebad: c1 ef 08                      shr     edi, 0x8
  28ebb0: 48 8d 15 99 fa 1e 00          lea     rdx, [rip + 0x1efa99]   # 0x47e650 <INSTRUCTION_TABLE>
  28ebb7: 49 89 c4                      mov     r12, rax
  28ebba: ff 24 ca                      jmp     qword ptr [rdx + 8*rcx]

# this branch
  28f185: 41 0f b6 0f                   movzx   ecx, byte ptr [r15]
  28f189: 48 8d 15 c0 04 1f 00          lea     rdx, [rip + 0x1f04c0]   # 0x47f650 <INSTRUCTION_TABLE>
  28f190: 41 0f b6 7f 01                movzx   edi, byte ptr [r15 + 0x1]
  28f195: 49 89 c4                      mov     r12, rax
  28f198: ff 24 ca                      jmp     qword ptr [rdx + 8*rcx]

I did this a bit half-hazardly through a combination of modifying macros and manual changes to anything that assigns to next_instr and a few special cases like exit_unwind that didn't fit. Could clean up with some direction.

One super naive metric is

# should be a tab after jmp
llvm-objdump --x86-asm-syntax=intel -D python | grep 'jmp   r' | wc -l

which is 916 for this modification and 731 originally, so 185 more places where we jmp to a register instead of a computed address.

[Fidget-Spinner]

Thanks for your work on this! Let me know if it passes the PGO test suite (./python -m test --pgo), after that we can bench it!

[aconz2]

Okay it is passing ./python -m test and ./python -m test --pgo which took a bit of trial and error to figure out which ops can call into things that modify bytecode where we have to reload next_op_f later than the entrypoint. These are unceremoniously scattered right now in bytecodes.c sometimes as early as I think is possible and some I just stuck at the end because I couldn't tell what was going on.

An example of one of these is STORE_ATTR, where in for example ./python -m test test_descr the next op is LOAD_GLOBAL initially but after calling PyObject_SetAttr, it reaches specialize and turns the op into LOAD_GLOBAL_BUILTIN. Or at least I think that is what is going on. Some others that I think can do this are PyObject_SetItem, PyObject_GetItem, PyStackRef_CLOSE, and PyStackRef_XCLOSE.

I'm not sure if passing all the tests is conclusive that this is correct; is there a more principled way of knowing who/what could change the value of next_instr between the start of the op and the end?

After fixing bugs the jmp to reg count went up to 933 from 917.

One thing I noticed is that sometimes the target address is loaded early and then stashed on the interpreter stack, like:

00000000001a8920 <_TAIL_CALL_BINARY_OP>:
  1a8920: 55                            push    rbp
  1a8921: 48 83 ec 10                   sub     rsp, 0x10
  1a8925: 41 89 f9                      mov     r9d, edi
  1a8928: 4d 89 7c 24 38                mov     qword ptr [r12 + 0x38], r15
  1a892d: 48 8d 15 1c cd 2d 00          lea     rdx, [rip + 0x2dcd1c]   # 0x485650 <INSTRUCTION_TABLE>
  1a8934: 49 8b 6d f0                   mov     rbp, qword ptr [r13 - 0x10]
  1a8938: 49 8b 5d f8                   mov     rbx, qword ptr [r13 - 0x8]
  1a893c: 41 0f b7 47 02                movzx   eax, word ptr [r15 + 0x2]
  1a8941: 66 83 f8 0e                   cmp     ax, 0xe
  1a8945: 0f 86 97 00 00 00             jbe     0x1a89e2 <_TAIL_CALL_BINARY_OP+0xc2>
  1a894b: 4c 89 34 24                   mov     qword ptr [rsp], r14
  1a894f: 41 0f b6 4f 0c                movzx   ecx, byte ptr [r15 + 0xc]
  1a8954: 48 8b 0c ca                   mov     rcx, qword ptr [rdx + 8*rcx]  #<-- compute target address
  1a8958: 48 89 4c 24 08                mov     qword ptr [rsp + 0x8], rcx    #<-- store to stack
  ...
  ...
  1a89d6: 48 8b 44 24 08                mov     rax, qword ptr [rsp + 0x8]    #<-- load from stack
  1a89db: 48 83 c4 10                   add     rsp, 0x10
  1a89df: 5d                            pop     rbp
  1a89e0: ff e0                         jmp     rax                           #<-- jmp
  1a89e2: 4d 89 6c 24 40                mov     qword ptr [r12 + 0x40], r13
  1a89e7: 4d 8d 44 24 50                lea     r8, [r12 + 0x50]
  1a89ec: 48 89 ef                      mov     rdi, rbp
  1a89ef: 48 89 de                      mov     rsi, rbx
  1a89f2: 48 89 d5                      mov     rbp, rdx                      #<-- rdx still has INSTRUCTION_TABLE base, "save" to rbp
  1a89f5: 4c 89 fa                      mov     rdx, r15
  1a89f8: 44 89 c9                      mov     ecx, r9d
  1a89fb: 44 89 cb                      mov     ebx, r9d
  1a89fe: e8 5d ea 11 00                call    0x2c7460 <_Py_Specialize_BinaryOp>
  1a8a03: 4d 8b 6c 24 40                mov     r13, qword ptr [r12 + 0x40]
  1a8a08: 49 c7 44 24 40 00 00 00 00    mov     qword ptr [r12 + 0x40], 0x0
  1a8a11: 41 0f b6 07                   movzx   eax, byte ptr [r15]
  1a8a15: 89 df                         mov     edi, ebx
  1a8a17: 48 89 e9                      mov     rcx, rbp                    #<-- rename to get INSTRUCTION_TABLE
  1a8a1a: 48 83 c4 10                   add     rsp, 0x10
  1a8a1e: 5d                            pop     rbp
  1a8a1f: ff 24 c1                      jmp     qword ptr [rcx + 8*rax]     #<-- jump to computed address
  ...

and then there is another variant seen above where it has loaded the INSTRUCTION_TABLE address early and then uses it later with a computed address. Not sure if these are still a win.

And this is all without the jit in mind right now.

[Fidget-Spinner]

I've requested a benchmarking run for this. Will report back soon.

[Fidget-Spinner]

Ok we have received the benchmark results (thanks to Mike):

https://github.com/faster-cpython/benchmarking-public/tree/main/results/bm-20250212-3.14.0a4+-e9c43a0-CLANG

Roughly 1% faster on AMD64 Xeon W-2255

No speedup on macOS M1.

This makes sense. The Xeon W-2255 has presumably weaker out of order execution capabilities than the M1. So there's benefit in doing this for older processors, or those with weaker OOO execution.

[Fidget-Spinner]

Note that a 1% win is pretty significant. So great job! If we were to upstream this though, it would need to be less likely to trip up the cpython contributors, so let me think how to approach this in a less error-prone way using the code generator.

[aconz2]

One thought I've had to eliminate the special casing is something like:

DEF_TARGET() {
    u8 next_opcode = load_next_opcode();
    tail_funcptr next_f = INSTRUCTION_TABLE[next_opcode];
    // ...
    u8 next_opcode2 = load_next_opcode();
    if (next_opcode == next_opcode2) TAIL return next_f(ARGS);
    next_f = INSTRUCTION_TABLE[next_opcode];
    TAIL return next_f(ARGS);
}

But not clear without testing whether this would still be a win. I can try this out and post a branch when I get some time

[aconz2]

I tried it out and is passing -m test pgo. Branch is main...aconz2:cpython:aconz2/early-tail-call-load-2

The diff is much smaller, only exception_unwind and start_frame needed changes beyond the generator.

Here is one op for comparison

0000000000296870 <_TAIL_CALL_GET_LEN>:
  296870: 55                            push    rbp
  296871: 4d 89 7c 24 38                mov     qword ptr [r12 + 0x38], r15
  296876: 41 0f b6 5f 02                movzx   ebx, byte ptr [r15 + 0x2]
  29687b: 49 83 c7 02                   add     r15, 0x2
  29687f: 48 8d 05 ca fd 1e 00          lea     rax, [rip + 0x1efdca]   # 0x486650 <INSTRUCTION_TABLE>
  296886: 48 8b 2c d8                   mov     rbp, qword ptr [rax + 8*rbx]
  29688a: 49 8b 7d f8                   mov     rdi, qword ptr [r13 - 0x8]
  29688e: 4d 89 6c 24 40                mov     qword ptr [r12 + 0x40], r13
  296893: e8 98 27 ea ff                call    0x139030 <PyObject_Size>
  296898: 4d 8b 6c 24 40                mov     r13, qword ptr [r12 + 0x40]
  29689d: 49 c7 44 24 40 00 00 00 00    mov     qword ptr [r12 + 0x40], 0x0
  2968a6: 48 85 c0                      test    rax, rax
  2968a9: 78 28                         js      0x2968d3 <_TAIL_CALL_GET_LEN+0x63>
  2968ab: 48 89 c7                      mov     rdi, rax
  2968ae: e8 cd 7e eb ff                call    0x14e780 <PyLong_FromSsize_t>
  2968b3: 48 85 c0                      test    rax, rax
  2968b6: 74 1b                         je      0x2968d3 <_TAIL_CALL_GET_LEN+0x63>
  2968b8: 49 89 45 00                   mov     qword ptr [r13], rax
  2968bc: 49 83 c5 08                   add     r13, 0x8
  2968c0: 41 0f b7 07                   movzx   eax, word ptr [r15]
  2968c4: 89 c7                         mov     edi, eax
  2968c6: c1 ef 08                      shr     edi, 0x8
  2968c9: 38 c3                         cmp     bl, al
  2968cb: 75 0c                         jne     0x2968d9 <_TAIL_CALL_GET_LEN+0x69>
  2968cd: 48 89 e8                      mov     rax, rbp
  2968d0: 5d                            pop     rbp
  2968d1: ff e0                         jmp     rax
  2968d3: 5d                            pop     rbp
  2968d4: e9 c7 ce f0 ff                jmp     0x1a37a0 <_TAIL_CALL_error>
  2968d9: 0f b6 c0                      movzx   eax, al
  2968dc: 48 8d 0d 6d fd 1e 00          lea     rcx, [rip + 0x1efd6d]   # 0x486650 <INSTRUCTION_TABLE>
  2968e3: 5d                            pop     rbp
  2968e4: ff 24 c1                      jmp     qword ptr [rcx + 8*rax]
  2968e7: 66 0f 1f 84 00 00 00 00 00    nop     word ptr [rax + rax]

and bigger ops are stashing the original next opcode on the stack for later comparison.