[X86] Fold load with tail call more aggressively

[mejedi]

Summary

Enable folding for tail calls with multiple arguments.

Fixes bugzilla issue 50042.

[mejedi]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1209

This part I'm least happy about. It was copied verbatim from X86InstrCompiler.td (X86tcret_6regs).

TBH, I don't fully understand the implications of moving LOAD close to TC_RETURN. It looks safe. However, if this check is omitted, we can end up moving the LOAD and not folding it. This is causing malformed DAG during scheduling in test/CodeGen/X86/musttail-varargs.ll, can't figure why.

[pengfei]

I don't find any problem in the tests. So I think it looks good. But I'm not familiar with call lowering, I'd like others to sign off.

llvm/test/CodeGen/X86/sibcall-4.ll

7

Nit: This change is confusing. Maybe better to check all assemblys, or to use update_llc_test_checks.py to generate it.

[mejedi]

ping

[mejedi]

Ping

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

[waterlens]

This patch looks nice. Why wasn't it merged?

[pengfei]

LGTM. Let me know if you need help for committing the patch.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1209

LOAD folding requires 2 scratch registers in some cases. If the 6 parameter registers are occupied, there's only one scratch register available.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

 }
 
 /// Replace the original chain operand of the call with
 /// load's chain operand and move load below the call's chain operand.
 static void moveBelowOrigChain(SelectionDAG *CurDAG, SDValue Load,
                                SDValue Call, SDValue OrigChain) {
   SmallVector<SDValue, 8> Ops;
   SDValue Chain = OrigChain.getOperand(0);
-  if (Chain.getNode() == Load.getNode())
+  if (Chain.getNode() == Load.getNode()) {
+    // Multiple chains might branch out of Load, fix them.
     Ops.push_back(Load.getOperand(0));
-  else {
+    for (auto UI = Load.getNode()->use_begin(), UE = Load.getNode()->use_end();
+         UI != UE;) {
+      auto VT = UI.getUse().getValueType();
+      auto *U = *UI++;
+      if (VT == MVT::Other) {
+        Ops.append(U->op_begin() + 1, U->op_end());
+        CurDAG->UpdateNodeOperands(U, Ops);
+        Ops.resize(1);
+      }
+    }
+  } else {
     assert(Chain.getOpcode() == ISD::TokenFactor &&
            "Unexpected chain operand");
     for (unsigned i = 0, e = Chain.getNumOperands(); i != e; ++i)
       if (Chain.getOperand(i).getNode() == Load.getNode())
         Ops.push_back(Load.getOperand(0));
       else
         Ops.push_back(Chain.getOperand(i));
     SDValue NewChain =
       CurDAG->getNode(ISD::TokenFactor, SDLoc(Load), MVT::Other, Ops);
     Ops.clear();
     Ops.push_back(NewChain);
-  }
     Ops.append(OrigChain->op_begin() + 1, OrigChain->op_end());
     CurDAG->UpdateNodeOperands(OrigChain.getNode(), Ops);
+  }
   CurDAG->UpdateNodeOperands(Load.getNode(), Call.getOperand(0),
                              Load.getOperand(1), Load.getOperand(2));
 
   Ops.clear();
   Ops.push_back(SDValue(Load.getNode(), 1));
   Ops.append(Call->op_begin() + 1, Call->op_end());
   CurDAG->UpdateNodeOperands(Call.getNode(), Ops);
 }
if (Callee.getNode() == Chain.getNode() || !Callee.hasOneUse())
     return false;
   LoadSDNode *LD = dyn_cast<LoadSDNode>(Callee.getNode());
   if (!LD ||
       !LD->isSimple() ||
       LD->getAddressingMode() != ISD::UNINDEXED ||
       LD->getExtensionType() != ISD::NON_EXTLOAD)
     return false;
 
+  if (HasCallSeq) {
     // Now let's find the callseq_start.
-  while (HasCallSeq && Chain.getOpcode() != ISD::CALLSEQ_START) {
+    while (Chain.getOpcode() != ISD::CALLSEQ_START) {
       if (!Chain.hasOneUse())
         return false;
       Chain = Chain.getOperand(0);
     }
+  } else if (Chain.getOpcode() == ISD::CopyToReg) {
+    // Locate first CopyToReg in the sequence of CopyToReg-s.
+    while (Chain.getOperand(0).getOpcode() == ISD::CopyToReg) {
+      Chain = Chain.getOperand(0);
+    }
+  }
 
   if (!Chain.getNumOperands())
     return false;
   // Since we are not checking for AA here, conservatively abort if the chain
   // writes to memory. It's not safe to move the callee (a load) across a store.
   if (isa<MemSDNode>(Chain.getNode()) &&
       cast<MemSDNode>(Chain.getNode())->writeMem())
     return false;
if (OptLevel != CodeGenOpt::None &&
       ///      \        /
       ///       \      /
       ///       [CALL]
       bool HasCallSeq = N->getOpcode() == X86ISD::CALL;
       SDValue Chain = N->getOperand(0);
       SDValue Load  = N->getOperand(1);
       if (!isCalleeLoad(Load, Chain, HasCallSeq))
         continue;
+      if (N->getOpcode() == X86ISD::TC_RETURN) {
+        // The TC_RETURN instructions execute after the epilog, so they
+        // can never use callee-saved registers. Ensure that there are
+        // registers available for the load address and the index.
+        unsigned NumRegs = 0;
+        for (unsigned i = 3, e = N->getNumOperands(); i != e; ++i)
+          if (isa<RegisterSDNode>(N->getOperand(i)))
+            ++NumRegs;
+        if (NumRegs > 6)
+          continue;
+      }
       moveBelowOrigChain(CurDAG, Load, SDValue(N, 0), Chain);
       ++NumLoadMoved;
       MadeChange = true;
       continue;
     }
 
     // Lower fpround and fpextend nodes that target the FP stack to be store and
     // load to the stack.  This is a gross hack.  We would like to simply mark

llvm/test/CodeGen/X86/cfguard-checks.ll

entry:
   ; X32:       movl %esi, %eax
   ; X32:       jmpl       *%ecx                  # TAILCALL
   ; X32-NOT:   calll
 
   ; Use NEXT here because we previously had an extra instruction in this sequence.
   ; X64-LABEL: vmptr_thunk:
   ; X64:            movq (%rcx), %rax
   ; X64-NEXT:       movq 8(%rax), %rax
-  ; X64-NEXT:       movq __guard_dispatch_icall_fptr(%rip), %rdx
-  ; X64-NEXT:       rex64 jmpq *%rdx            # TAILCALL
+  ; X64-NEXT:       rex64 jmpq *__guard_dispatch_icall_fptr(%rip) # TAILCALL
   ; X64-NOT:   callq
 }
 
 ; Test that longjmp targets have public labels and are included in the .gljmp section.
 %struct._SETJMP_FLOAT128 = type { [2 x i64] }
 @buf1 = internal global [16 x %struct._SETJMP_FLOAT128] zeroinitializer, align 16
 
 define i32 @func_cf_setjmp() {

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

 ; LINUX-X32-NEXT:    movl %edi, %edi
 ; LINUX-X32-NEXT:    jmpq *%r11 # TAILCALL
 ;
 ; WINDOWS-LABEL: h_thunk:
 ; WINDOWS:       # %bb.0:
 ; WINDOWS-NEXT:    cmpb $1, (%rcx)
 ; WINDOWS-NEXT:    jne .LBB2_2
 ; WINDOWS-NEXT:  # %bb.1: # %then
-; WINDOWS-NEXT:    movq 8(%rcx), %rax
-; WINDOWS-NEXT:    rex64 jmpq *%rax # TAILCALL
+; WINDOWS-NEXT:    rex64 jmpq *8(%rcx) # TAILCALL
 ; WINDOWS-NEXT:  .LBB2_2: # %else
 ; WINDOWS-NEXT:    movq 16(%rcx), %rax
 ; WINDOWS-NEXT:    movl $42, {{.*}}(%rip)
 ; WINDOWS-NEXT:    rex64 jmpq *%rax # TAILCALL
 ;
 ; X86-LABEL: h_thunk:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax

llvm/test/CodeGen/X86/sibcall-4.ll

 ; RUN: llc < %s -mtriple=i386-pc-linux-gnu | FileCheck %s
 ; pr7610
 
 define ghccc void @t(i32* %Base_Arg, i32* %Sp_Arg, i32* %Hp_Arg, i32 %R1_Arg) nounwind {
 cm1:
 ; CHECK-LABEL: t:
-; CHECK: jmpl *%eax
+; CHECK: jmpl *(%eax)
   %nm3 = getelementptr i32, i32* %Sp_Arg, i32 1
   %nm9 = load i32, i32* %Sp_Arg
   %nma = inttoptr i32 %nm9 to void (i32*, i32*, i32*, i32)*
   tail call ghccc void %nma(i32* %Base_Arg, i32* %nm3, i32* %Hp_Arg, i32 %R1_Arg) nounwind
   ret void
 }

llvm/test/CodeGen/X86/tailcall-64.ll

 ; CHECK-NEXT:    jmpq *16(%rcx,%rax,8) ## TAILCALL
 entry:
   %dsplen = getelementptr inbounds [0 x %struct.funcs], [0 x %struct.funcs]* @func_table, i64 0, i64 %idxprom, i32 2
   %x1 = load i32 (i8*)*, i32 (i8*)** %dsplen, align 8
   %call = tail call i32 %x1(i8* %mbstr) nounwind
   ret void
 }
 
+; Check that we can fold an indexed load into a tail call instruction, multiple args.
+define void @fold_indexed_load_3(i8* %mbstr, i32* %fw1, i32 %fw2, i64 %idxprom) nounwind uwtable ssp {
+; CHECK-LABEL: fold_indexed_load_3:
+; CHECK:       ## %bb.0: ## %entry
+; CHECK-NEXT:    leaq (%rcx,%rcx,4), %rax
+; CHECK-NEXT:    movq _func_table@{{.*}}(%rip), %rcx
+; CHECK-NEXT:    jmpq *(%rcx,%rax,8) ## TAILCALL
+entry:
+  %dsplen = getelementptr inbounds [0 x %struct.funcs], [0 x %struct.funcs]* @func_table, i64 0, i64 %idxprom, i32 0
+  %x1 = load i32 (i8*, i32*, i32)*, i32 (i8*, i32*, i32)** %dsplen, align 8
+  %call = tail call i32 %x1(i8* %mbstr, i32* %fw1, i32 %fw2) nounwind
+  ret void
+}
+
 @funcs = external constant [0 x i32 (i8*, ...)*]
 
 ; <rdar://problem/12282281> Fold an indexed load into the tail call instruction.
 ; Calling a varargs function with 6 arguments requires 7 registers (%al is the
 ; vector count for varargs functions). This leaves %r11 as the only available
 ; scratch register.
 ;
 ; It is not possible to fold an indexed load into TCRETURNmi64 in that case.