[Backend][X86] Improved tail call optimization for functions marked as musttail

[huangjd]

Summary

Now check if a function call marked as musttail also qualifies as a sibcall,
which means on-stack arguments are already in place before handling the
tail call function.

This also suppress a bug when a musttail call has struct arguments
passed on the stack, the output code results in return address
overwriting. See https://github.com/llvm/llvm-project/issues/56891

[davidxl]

llvm/test/CodeGen/X86/musttail-struct.ll

41

is this the correct comment, or it is intended for a negative test (missing)?

[huangjd]

llvm/test/CodeGen/X86/musttail-struct.ll

41

This is just describing the purpose of the test. in x64 the first 6 are passed by register and the rest 2 are passed by stack

[davidxl]

llvm/test/CodeGen/X86/musttail-struct.ll

41

Ok. Can you expand the comment a little more.

Can you also add a negative test?

[huangjd]

Updated comments

[huangjd]

llvm/test/CodeGen/X86/musttail-struct.ll

41

I am not sure what should a negative test be in this case, since the front end rejects the code if "musttail" is applied on function call that can't be tail called. Also if tail call is not applied, the output assembly can be quite unpredictable (if someone else change anything else in codegen) since instructions can be shuffled. I have yet seen a negative test with tail call optimization

[davidxl]

llvm/lib/Target/X86/X86ISelLowering.cpp

4267

Can this be handled by line 4274?

4314

what is this change for?

llvm/test/CodeGen/X86/musttail-struct.ll

41

An IR test can be written manually. This is to test line 4262: IsEligibleForTailCallOptimization() logic for mustTail case.

[huangjd]

Restructured logic to determine isSibcall

[huangjd]

Added comments

llvm/lib/Target/X86/X86ISelLowering.cpp

4314

Sibcall should not adjust stack because all arguments are already in place. This is previously causing the bug I posted on github

[efriedma]

Example of copying arguments in memory for musttail follows (which is still broken with this patch). Was leaving this broken intentional?

struct A {
    int a[5];
};
extern int G(A a, A b);
int F(A a, A b) {
    [[clang::musttail]] return G(b, a);
}

[huangjd]

@efriedma This is a limitation of the current patch that it could only handle the case with arguments in matching order. After some investigation I found that the current function call lowering pass cannot handle in-place stack argument shuffling at all, as it doesn't take read/write dependency into account. It would require an extensive refactoring to completely fix it.

In addition, shuffling large stack arguments to force a tail call can get quite complicated, for instance, if we have F(x0, x1, x2, ... xn) calling G(x1, x2, ... xn, x0) where every argument is the same type and is passed by stack, in order to do an in-place argument swapping without further stack spill, the optimal code should be something like

for i in 0..k:
  tmpreg = x0.word_i
  x0.word_i = x1.word_i
  ...
  xn.word_i = tmpreg

where x0.word_i, ..., xn.word_i are cache-line sized partitions of the arguments x0, ..., xn.

The current patch is able to handle two common real world scenario of tail calls (a) a function that simply delegates to another function without doing anything else, and (b) recursive tail call where only the induction variable passed by register changes.

llvm/lib/Target/X86/X86ISelLowering.cpp

if (Subtarget.isPICStyleGOT() && !IsGuaranteeTCO && !IsMustTail) {
     // that require lazy function symbol resolution. Using musttail or
     // GuaranteedTailCallOpt will override this.
     GlobalAddressSDNode *G = dyn_cast<GlobalAddressSDNode>(Callee);
     if (!G || (!G->getGlobal()->hasLocalLinkage() &&
                G->getGlobal()->hasDefaultVisibility()))
       isTailCall = false;
   }
 
-  if (isTailCall && !IsMustTail) {
+  if (isTailCall) {
     // Check if it's really possible to do a tail call.
-    isTailCall = IsEligibleForTailCallOptimization(
+    bool CanTailCall = IsEligibleForTailCallOptimization(
         Callee, CallConv, IsCalleePopSRet, isVarArg, CLI.RetTy, Outs, OutVals,
         Ins, DAG);
 
-    // Sibcalls are automatically detected tailcalls which do not require
-    // ABI changes.
+    if (IsMustTail)
+      IsSibcall = CanTailCall;
+    else {
+      isTailCall = CanTailCall;
+
+      // Sibcalls are automatically detected tailcalls which do not require ABI
+      // changes.
       if (!IsGuaranteeTCO && isTailCall)
         IsSibcall = true;
+    }
 
     if (isTailCall)
       ++NumTailCalls;
   }
 
   if (IsMustTail && !isTailCall)
     report_fatal_error("failed to perform tail call elimination on a call "
                        "site marked musttail");
X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   if (IsSibcall)
     // This is a sibcall. The memory operands are available in caller's
     // own caller's stack.
     NumBytes = 0;
   else if (IsGuaranteeTCO && canGuaranteeTCO(CallConv))
     NumBytes = GetAlignedArgumentStackSize(NumBytes, DAG);
 
   int FPDiff = 0;
-  if (isTailCall &&
+  if (isTailCall && !IsSibcall &&
       shouldGuaranteeTCO(CallConv,
                          MF.getTarget().Options.GuaranteedTailCallOpt)) {
     // Lower arguments at fp - stackoffset + fpdiff.
     unsigned NumBytesCallerPushed = X86Info->getBytesToPopOnReturn();
 
     FPDiff = NumBytesCallerPushed - NumBytes;
 
     // Set the delta of movement of the returnaddr stackslot.

llvm/test/CodeGen/X86/musttail-struct.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s
+; RUN: llc < %s -mtriple=i686-unknown-unknown | FileCheck %s
+
+; Test correct handling of musttail call with struct.
+
+%struct.1xi32 = type { [1 x i32] }
+%struct.3xi32 = type { [3 x i32] }
+%struct.5xi32 = type { [5 x i32] }
+
+declare dso_local i32 @Func1(ptr byval(%struct.1xi32) %0)
+declare dso_local i32 @Func3(ptr byval(%struct.3xi32) %0)
+declare dso_local i32 @Func5(ptr byval(%struct.5xi32) %0)
+declare dso_local i32 @FuncManyArgs(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, i8 %6, ptr byval(%struct.5xi32) %7)
+
+define dso_local i32 @test1(ptr byval(%struct.1xi32) %0) {
+; CHECK-LABEL: test1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    jmp Func1 # TAILCALL
+  %r = musttail call i32 @Func1(ptr byval(%struct.1xi32) %0)
+  ret i32 %r
+}
+
+define dso_local i32 @test3(ptr byval(%struct.3xi32) %0) {
+; CHECK-LABEL: test3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    jmp Func3 # TAILCALL
+  %r = musttail call i32 @Func3(ptr byval(%struct.3xi32) %0)
+  ret i32 %r
+}
+
+; sizeof(%struct.5xi32) > 16, in x64 this is passed on stack.
+define dso_local i32 @test5(ptr byval(%struct.5xi32) %0) {
+; CHECK-LABEL: test5:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    jmp Func5 # TAILCALL
+  %r = musttail call i32 @Func5(ptr byval(%struct.5xi32) %0)
+  ret i32 %r
+}
+
+; Arguments of different sizes passed on stack.
+define dso_local i32 @testManyArgs(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, i8 %6, ptr byval(%struct.5xi32) %7) {
+; CHECK-LABEL: testManyArgs:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    jmp FuncManyArgs # TAILCALL
+  %r = musttail call i32 @FuncManyArgs(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, i8 %6, ptr byval(%struct.5xi32) %7)
+  ret i32 %r
+}
+
+define dso_local i32 @testRecursion(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, i8 %6, ptr byval(%struct.5xi32) %7) {
+; CHECK-LABEL: testRecursion:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    jmp testRecursion # TAILCALL
+  %r = musttail call i32 @testRecursion(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, i8 %6, ptr byval(%struct.5xi32) %7)
+  ret i32 %r
+}

llvm/test/CodeGen/X86/musttail-tailcc.ll

 
 define dso_local tailcc void @void_test(i32, i32, i32, i32) {
 ; X64-LABEL: void_test:
 ; X64:       # %bb.0: # %entry
 ; X64-NEXT:    jmp void_test # TAILCALL
 ;
 ; X32-LABEL: void_test:
 ; X32:       # %bb.0: # %entry
-; X32-NEXT:    pushl %esi
-; X32-NEXT:    .cfi_def_cfa_offset 8
-; X32-NEXT:    .cfi_offset %esi, -8
-; X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X32-NEXT:    movl {{[0-9]+}}(%esp), %esi
-; X32-NEXT:    movl %esi, {{[0-9]+}}(%esp)
-; X32-NEXT:    movl %eax, {{[0-9]+}}(%esp)
-; X32-NEXT:    popl %esi
-; X32-NEXT:    .cfi_def_cfa_offset 4
 ; X32-NEXT:    jmp void_test # TAILCALL
   entry:
    musttail call tailcc void @void_test( i32 %0, i32 %1, i32 %2, i32 %3)
    ret void
 }
 
 define dso_local tailcc i1 @i1test(i32, i32, i32, i32) {
 ; X64-LABEL: i1test:
 ; X64:       # %bb.0: # %entry
 ; X64-NEXT:    jmp i1test # TAILCALL
 ;
 ; X32-LABEL: i1test:
 ; X32:       # %bb.0: # %entry
-; X32-NEXT:    pushl %esi
-; X32-NEXT:    .cfi_def_cfa_offset 8
-; X32-NEXT:    .cfi_offset %esi, -8
-; X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X32-NEXT:    movl {{[0-9]+}}(%esp), %esi
-; X32-NEXT:    movl %esi, {{[0-9]+}}(%esp)
-; X32-NEXT:    movl %eax, {{[0-9]+}}(%esp)
-; X32-NEXT:    popl %esi
-; X32-NEXT:    .cfi_def_cfa_offset 4
 ; X32-NEXT:    jmp i1test # TAILCALL
   entry:
   %4 = musttail call tailcc i1 @i1test( i32 %0, i32 %1, i32 %2, i32 %3)
   ret i1 %4
 }