Implement Proper Tail Call [Discussion]

[abchatra]

This is a discussion thread on implementing proper tail calls (PTC) in Chakra. For simplicity, only full jit implementation is discussed.

Introduction

Wiki & in Javascript

x86 implementation details of a PTC

Chakra uses caller cleanup for the arguments in the stack. This means caller is required to clean up the stack space for the arguments it pushed into the stack (similar to C calling convention).

Let us take a example:

"use strict"
function foo()
{ 
  bar(1);
}
function bar()
{
  return foobar(1,2,3,4); //tail call
}

Without PTC, asm code for a tail call to foobar(1,2,3,4) looks as below

   return foobar(1,2,3,4)

    //#1 allocate the stack for arguments first
    (esp).i32       =  LEA            [(esp).i32-24].i32

    //#2 push the arguments
    [esp]           = MOV           0xXXXXXXXX (undefined)[Undefined].var
    [esp + 4]       = MOV           0x00000003.var
    [esp + 8]       = MOV           0x00000005.var
    [esp + 12]      = MOV           0x00000007.var
    [esp + 16]      = MOV           0x00000009.var

    // push two meta arguments, callinfo (which has a count of arguments) and the function object. 
                        PUSH           268435461 (0x10000005).i32
                        PUSH           s7(esi).var

    //#3 Call the function
                       CALL           s21(eax).u32

    //#4 Re-adjust the stack to 24+8(2*4 meta arguments)                   
    (esp).i32       =  LEA            [(esp).i32+32].i32

     //#5 Epilog and return
    (ebx).i32       =  POP
    (esi).i32       =  POP
    (edi).i32       =  POP
    (esp).i32       =  MOV            (ebp).i32
    (ebp).i32       =  POP
                       RET            0 (0x0).i32, (eax).i32    

On implementing PTC asm code:

    return foobar(1,2,3,4)

    //#2 push the arguments to the caller allocated stack space. Note ebp is a frame pointer.
    [ebp + 16]      = MOV           0xXXXXXXXX (undefined)[Undefined].var
    [ebp + 20]      = MOV           0x00000003.var
    [ebp + 24]      = MOV           0x00000005.var
    [ebp + 28]      = MOV           0x00000007.var
    [ebp + 32]      = MOV           0x00000009.var

    //push two meta arguments, callinfo (which has number of arguments) and function object. 
    [ebp + 12]      = MOV           268435461 (0x10000005).i32
    [ebp + 8]       = MOV           s7(esi).var


     //#4 Execute epilog before the jump, as we want to restore non-volatile registers
    (ebx).i32       =  [ebp - 16]
    (esi).i32       =  [ebp - 12]
    (edi).i32       =  [ebp - 8]
    (esp).i32       =  [ebp + 4] //this should point to return address
    (ebp).i32       =  [ebp]

    //#3 Jump to callee
                     JMP           s21(eax).u32

This looks cool as we avoided a new stack frame allocation. Though there is a caveat, bar does not know the number of arguments passed to it. So bar call frame has no clue for the step 2 whether it has enough space on the stack or not. In above example, foo passed only one argument and it doesn't have space for 3 more arguments to move to the stack. What do we do now? Check in the main path.

    CMP            [ebp + 12], 4 
    JLT            $helper

Branches are expensive, though a smart compiler can move that to a prolog of the function.

Now let us look at the helper. How would this allocate the extra stack space required for arguments?

    $helper:
    //#7 Now you are mucking with the stack be carefull. Let us burn registers here as it is slow path
     eax            = <extra args required>
     esp            =  LEA        [ebp + eax]

    //#6 Mov return address and frame pointer
    [esp + 4]     =  MOV [ebp + 4]
    [esp]           =  MOV [ebp]

    //#2 push the arguments to the caller-allocated stack space and additional stack space

    //#4 Execute epilog before the jump, as we want to restore non-volatile registers

    //#3 Jump to callee
                     JMP           s21(eax).u32

Are we done? Not yet, foo which called bar has no clue that bar has changed stack pointer (esp) while doing a PTC to foobar. If foo has any locals referenced by the stack pointer, which is altered now. So it has to check and restore the stack pointer, something like following:

bar(1);
   // Save and re-adjust the stack to where it was before the call from a non-volatile register or known stack location.                   
        edi         =  MOV     esp
                          CALL           s21(eax).u32
        esp        =  MOV     edi

This is now burning a register in non-PTC position on the main path. Every call we need to burn an extra non-volatile register as the callee can make a PTC. This is bad news for performance.

There are couple of alternate solution to this:

• In case there is not enough stack space, do a CALL instead of PTC which will cleanup the stack space when control is returned to it. This will work though there will be an extra fake frame on the stack which needs to be hidden.
• Change the calling convention in the Chakra runtime to do a callee cleanup. This is riddled with issues as all our built-ins such as Math.sin are C++ helpers which accept variable arguments and MSVC can't do callee cleanup. The solution is to build individual thunks for every built-in in the runtime 😟.

x64 implementation

Chakra uses C++ exception to propagate errors and relies on C++ exception unwinding to catch the errors and report to the users.

Windows x64 ABI mandates setting up the stack pointer (rsp) in the prolog and disallows changing it on the fly. This is to facilitate generating static unwind data, so OS can statically unwind the call frames.

Changing the stack pointer as in x86 is ruled out. We do have to create a whole new frame when the number of arguments to function at PTC position is higher than passed to the call frame. This is again a fake frame which needs to be hidden from stack walking etc.

[bterlson]

So just to summarize the options we have in front of us for x86:

1. Regress call performance.
2. Use CALL to allocate additional stack space when needed (thus violating the spec in certain cases)
3. Invest months of effort changing our calling conventions.

For x64, we only have option 2/3.

None of these options seem good :(

[abchatra]

@bterlson Yep that is my understanding.

[msaboff]

I have some clarifying questions about how Chakra works and what is presented.

1. Does the callinfo meta argument describe the actual number of parameters or the formal number of parameters? If the actual, the generated PTC code could perform the math needed to make sure SP has enough space. If formal, then presumably the JIT knows that value as a constant for both the caller and callee and can generate the appropriate fixup code in the case that the parameter counts don't match.
2. Why does bar() need to save and restore the SP when that is presumably provided by the callee saving SP in BP as part of a prologue?

[abchatra]

Thanks for taking a look @msaboff

Does the callinfo meta argument describe the actual number of parameters or the formal number of parameters? If the actual, the generated PTC code could perform the math needed to make sure SP has enough space. If formal, then presumably the JIT knows that value as a constant for both the caller and callee and can generate the appropriate fixup code in the case that the parameter counts don't match.

CallInfo has an actual number of parameters. We follow caller cleanup (similar to C calling convention). If esp is altered, PTC target (foobar) won't clean up the stack. The current frame (bar) which did PTC doesn't exist. The caller (foo) of the current frame needs to readjust the stack.

Why does bar() need to save and restore the SP when that is presumably provided by the callee saving SP in BP as part of a prologue?

Function bar doesn't have to save and restore the esp, but foo which did the call to bar(1) need to do so. As explained above, both bar and foobar won't clean up the stack from pushing extra arguments. If foo doesn't restore the esp, all locals referenced through esp points to garbage. One such example, nested calls x(y(z(1),2),3);

One of the alternatives I suggested was to change to callee cleanup, though as you might know it is super expensive, in terms of dev cost to change a calling convention in a runtime. Even then this won't work on x64. As x64 windows ABI mandates C calling convention which is caller cleanup and our built-ins are C++ functions and we rely on C++ exceptions for error handling.

Do let me know if you have any other questions. Appreciate your help.

[msaboff]

We had to change our calling convention to make PTC work. The biggest change was saving and restoring callee saves in prologues and epilogues for the 64 bit platforms we support. I spent about a month making callee saves work in all tiers including the various special cases we have.

We always use a frame pointer for every frame on all platforms, which makes restoring the caller's sp fairly straightforward. We use a little trick to restore the caller's sp to it's proper value. We know the sp that we allocated in the prologue of a function which is based on the sp at entry. After each call, we set the sp using the same calculation, but we can use bp since it has sp on entry.. This allows us to not require the sp to be a certain value after returning from a call. Basically you'll see an lea of bp - constant after every call instruction in our code. In most cases our sp is right and doesn't need the restoration, but with the lea, it is always right. In our experience, the lea is cheap to execute and allows us to make simplifying assumptions in our code.

Concerning making space for outgoing tail call arguments, why can't you always make sp right by calculating where it should be based on what you find at [ebp + 12] and the number of outgoing arguments? Basically always do what you have above after helper:. This is part of what our call frame shuffler does to make the frame right for the tail callee. Certainly you could add an optimization when doing self recursion. Note that you don't need to move the previous frame pointer. If you just load it from its original spot and don't restore it in the epilogue before tail calling (jump), you eliminate a store and a dependent load.

Couldn't you employ the same trick that we do to restore your sp using a lea of an offset to bp?

With these two changes, you eliminate the compare and branch as well as burning a register to save and restore sp in every caller.

[abchatra]

We had to change our calling convention to make PTC work. The biggest change was saving and restoring callee saves in prologues and epilogues for the 64 bit platforms we support. I spent about a month making callee saves work in all tiers including the various special cases we have

I don't understand this comment. We already save and restore callee save registers (non-volatile) in prolog and epilog of a callee. It is mandatory by Windows ABI (may be any ABI). What is the relation to callee cleanup and caller cleanup?

We always use a frame pointer for every frame on all platforms, which makes restoring the caller's sp fairly straightforward. We use a little trick to restore the caller's sp to it's proper value. We know the sp that we allocated in the prologue of a function which is based on the sp at entry. After each call, we set the sp using the same calculation, but we can use bp since it has sp on entry.. This allows us to not require the sp to be a certain value after returning from a call. Basically you'll see an lea of bp - constant after every call instruction in our code. In most cases our sp is right and doesn't need the restoration, but with the lea , it is always right. In our experience, the lea is cheap to execute and allows us to make simplifying assumptions in our code.

Couldn't you employ the same trick that we do to restore your sp using a lea of an offset to bp ?

Cool! Nice to JSC has a constant stack frame size for a call frame. We don't have a fixed frame for x86, we consume stack as necessary. For x64 and ARM we do have a fixed stack frame. But as I wrote earlier windows ABI mandates callee can't modify someone else (caller) stack, callee growing caller stack is ruled out. Next question would be can we not adhere to ABI? Yes if we didn't care about C++ exception handling to work. Our entire runtime error propagation is based on C++ exception handling. We throw C++ exception, expect destructors to be invoked, catch to report error or execute javascript catch blocks. Changing it to setjmp\longjmp or error code based model is again expensive.

Concerning making space for outgoing tail call arguments, why can't you always make sp right by calculating where it should be based on what you find at [ebp + 12] and the number of outgoing arguments? Basically, always do what you have above after helper: . This is part of what our call frame shuffler does to make the frame right for the tail callee. Certainly you could add an optimization when doing self recursion. Note that you don't need to move the previous frame pointer. If you just load it from its original spot and don't restore it in the epilogue before tail calling (jump), you eliminate a store and a dependent load.

Yes, this is doable. (Side note our [ebp + 12] contains more than callinfo, it has top 8 bits representing various metadata such as eval call). Even then don't you still need a CMP to see if you need to move the return address and ebp? Or you generate a move regardless?

Thank you for trying to help.

[msaboff]

I don't understand this comment. We already save and restore callee save registers (non-volatile) in prolog and epilog of a callee. It is mandatory by Windows ABI (may be any ABI). What is the relation to callee cleanup and caller cleanup?

I understand that you already handle callee saves. Just pointing out that we had to make these changes ourselves to implement PTC.

Our caller cleanup is minimal, just restoring sp.

But as I wrote earlier windows ABI mandates callee can't modify someone else (caller) stack, callee growing caller stack is ruled out.

In JSC, a callee cannot change the size of a caller's frame. The top most call frames can and do change their size. The lea sp, bp - offset restoration puts the stack pointer back in the default location. This default is the calculated maximum stack frame size. Accesses to arguments and locals are bp relative.

[abchatra]

Just pointing out that we had to make these changes ourselves to implement PTC.

Makes sense, thanks!

In JSC, a callee cannot change the size of a caller's frame. The top most call frames can and do change their size. The lea sp, bp - offset restoration puts the stack pointer back in the default location. This default is the calculated maximum stack frame size. Accesses to arguments and locals are bp relative.

Note as per x64 (& arm) windows ABI, even caller can't change its own stack except in prolog and epilog (strict frame size restrictions). This is to statically unwind the frame, without a frame pointer, which is not mandatory to setup in x64. In short, no modification to rsp (or r11) is allowed apart from prolog & epilog.