[GlobalISel] When a tail call is emitted in a block, stop translating it

[paquette]

Summary

This fixes a crash in tail call translation caused by assume and lifetime_end intrinsics.

It's possible to have instructions other than a return after a tail call which will still have Analysis::isInTailCallPosition return true. (Namely, lifetime_end and assume intrinsics.)

If we emit a tail call instruction, we should stop translating instructions in the block. Otherwise, we can end up emitting an extra return, or dead instructions in general. This makes the verifier unhappy, and is generally unfortunate for codegen.

This also removes the code from AArch64CallLowering that checks if we have a tail call when lowering a return. This is covered by the new code now.

Also update call-translator-tail-call.ll to show that we now properly tail call in the presence of lifetime_end and assume.

[aemerson]

LGTM with nit.

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

1574

Assert here that HasTailCall == false?

llvm/include/llvm/CodeGen/GlobalISel/IRTranslator.h

private:
   std::unique_ptr<OptimizationRemarkEmitter> ORE;
 
   FunctionLoweringInfo FuncInfo;
 
   // True when either the Target Machine specifies no optimizations or the
   // function has the optnone attribute.
   bool EnableOpts = false;
 
+  /// True when the block contains a tail call. This allows the IRTranslator to
+  /// stop translating such blocks early.
+  bool HasTailCall = false;
+
   /// Switch analysis and optimization.
   class GISelSwitchLowering : public SwitchCG::SwitchLowering {
   public:
     GISelSwitchLowering(IRTranslator *irt, FunctionLoweringInfo &funcinfo)
         : SwitchLowering(funcinfo), IRT(irt) {
       assert(irt && "irt is null!");
     }
 

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
for (auto &Arg : CS.args()) {
     Args.push_back(getOrCreateVRegs(*Arg));
   }
 
   MF->getFrameInfo().setHasCalls(true);
   bool Success =
       CLI->lowerCall(MIRBuilder, CS, Res, Args, SwiftErrorVReg,
                      [&]() { return getOrCreateVReg(*CS.getCalledValue()); });
 
+  // Check if we just inserted a tail call.
+  if (Success) {
+    const TargetInstrInfo *TII = MF->getSubtarget().getInstrInfo();
+    HasTailCall = TII->isTailCall(*std::prev(MIRBuilder.getInsertPt()));
+  }
+
   return Success;
 }
 
 bool IRTranslator::translateCall(const User &U, MachineIRBuilder &MIRBuilder) {
   const CallInst &CI = cast<CallInst>(U);
   auto TII = MF->getTarget().getIntrinsicInfo();
   const Function *F = CI.getCalledFunction();
 
#ifndef NDEBUG
     WrapperObserver.addObserver(&Verifier);
 #endif // ifndef NDEBUG
     RAIIDelegateInstaller DelInstall(*MF, &WrapperObserver);
     for (const BasicBlock *BB : RPOT) {
       MachineBasicBlock &MBB = getMBB(*BB);
       // Set the insertion point of all the following translations to
       // the end of this basic block.
       CurBuilder->setMBB(MBB);
-
+      HasTailCall = false;
       for (const Instruction &Inst : *BB) {
+        // If we translated a tail call in the last step, then we know
+        // everything after the call is either a return, or something that is
+        // handled by the call itself. (E.g. a lifetime marker or assume
+        // intrinsic.) In this case, we should stop translating the block and
+        // move on.
+        if (HasTailCall)
+          break;
 #ifndef NDEBUG
         Verifier.setCurrentInst(&Inst);
 #endif // ifndef NDEBUG
         if (translate(Inst))
           continue;
 
         OptimizationRemarkMissed R("gisel-irtranslator", "GISelFailure",
                                    Inst.getDebugLoc(), BB);

llvm/lib/Target/AArch64/AArch64CallLowering.cpp

void AArch64CallLowering::splitToValueTypes(
 
   SplitArgs.back().Flags[0].setInConsecutiveRegsLast();
 }
 
 bool AArch64CallLowering::lowerReturn(MachineIRBuilder &MIRBuilder,
                                       const Value *Val,
                                       ArrayRef<Register> VRegs,
                                       Register SwiftErrorVReg) const {
-
-  // Check if a tail call was lowered in this block. If so, we already handled
-  // the terminator.
-  MachineFunction &MF = MIRBuilder.getMF();
-  if (MF.getFrameInfo().hasTailCall()) {
-    MachineBasicBlock &MBB = MIRBuilder.getMBB();
-    auto FirstTerm = MBB.getFirstTerminator();
-    if (FirstTerm != MBB.end() && FirstTerm->isCall())
-      return true;
-  }
-
   auto MIB = MIRBuilder.buildInstrNoInsert(AArch64::RET_ReallyLR);
   assert(((Val && !VRegs.empty()) || (!Val && VRegs.empty())) &&
          "Return value without a vreg");
 
   bool Success = true;
   if (!VRegs.empty()) {
     MachineFunction &MF = MIRBuilder.getMF();
     const Function &F = MF.getFunction();

llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll

define void @test_mismatched_caller() {
   ; COMMON: bb.1 (%ir-block.0):
   ; COMMON:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
   ; COMMON:   BL @fast_fn, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
   ; COMMON:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
   ; COMMON:   RET_ReallyLR
   tail call fastcc void @fast_fn()
   ret void
 }
+
+; Verify that lifetime markers and llvm.assume don't impact tail calling.
+declare void @llvm.assume(i1)
+define void @test_assume() local_unnamed_addr {
+  ; COMMON-LABEL: name: test_assume
+  ; COMMON: bb.1.entry:
+  ; COMMON:   TCRETURNdi @nonvoid_ret, 0, csr_aarch64_aapcs, implicit $sp
+entry:
+  %x = tail call i32 @nonvoid_ret()
+  %y = icmp ne i32 %x, 0
+  tail call void @llvm.assume(i1 %y)
+  ret void
+}
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
+define void @test_lifetime() local_unnamed_addr {
+  ; COMMON-LABEL: name: test_lifetime
+  ; COMMON: bb.1.entry:
+  ; COMMON:   [[FRAME_INDEX:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.0.t
+  ; COMMON:   LIFETIME_START %stack.0.t
+  ; COMMON:   TCRETURNdi @nonvoid_ret, 0, csr_aarch64_aapcs, implicit $sp
+entry:
+  %t = alloca i8, align 1
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %t)
+  %x = tail call i32 @nonvoid_ret()
+  %y = icmp ne i32 %x, 0
+  tail call void @llvm.lifetime.end.p0i8(i64 1, i8* %t)
+  ret void
+}