StackProtector: insert check before `musttail` call.

[t.p.northover]

Summary

musttail calls have to directly precede the return instruction (of their value), so any stack checks must be inserted before the call instead of in between it and the return. Fortunately this same restriction means we don't have to look too hard to find out if we're about to break a musttail call.

[t.p.northover]

We've just discovered there are problems on the SDAG path too for guaranteed tail calls. I'll update the patch once I've got a fix ready.

[t.p.northover]

Fixed the issue in SelectionDAG too. Not quite identical there, but still fundamentally hinged on exactly where the check should be placed.

[aemerson]

LGTM with comment nit.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

1713

An example of what code you're looking for in this comment would be helpful.

[t.p.northover]

Thanks Amara. Committed with a beefed up comment as 5e3d9fcc3a88.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

 /// blocks. Luckily, selectiondag always moves physical registers into vregs
 /// when they enter functions and moves them through a sequence of copies back
 /// into the physical registers right before the terminator creating a
 /// ``Terminator Sequence''. This function is searching for the beginning of the
 /// terminator sequence so that we can ensure that we splice off not just the
 /// terminator, but additionally the copies that move the vregs into the
 /// physical registers.
 static MachineBasicBlock::iterator
-FindSplitPointForStackProtector(MachineBasicBlock *BB) {
+FindSplitPointForStackProtector(MachineBasicBlock *BB,
+                                const TargetInstrInfo &TII) {
   MachineBasicBlock::iterator SplitPoint = BB->getFirstTerminator();
-  //
   if (SplitPoint == BB->begin())
     return SplitPoint;
 
   MachineBasicBlock::iterator Start = BB->begin();
   MachineBasicBlock::iterator Previous = SplitPoint;
   --Previous;
 
+  if (TII.isTailCall(*SplitPoint) &&
+      Previous->getOpcode() == TII.getCallFrameDestroyOpcode()) {
+    // Call frames cannot be nested, so if this frame is describing the tail
+    // call itself then we must insert before the sequence even starts. On the
+    // other hand, it could be an unrelated call in which case this tail call
+    // has no register moves of its own and should be the split point.
+    do {
+      --Previous;
+      if (Previous->isCall())
+        return SplitPoint;
+    } while(Previous->getOpcode() != TII.getCallFrameSetupOpcode());
+
+    return Previous;
+  }
+
   while (MIIsInTerminatorSequence(*Previous)) {
     SplitPoint = Previous;
     if (Previous == Start)
       break;
     --Previous;
   }
 
   return SplitPoint;
SelectionDAGISel::FinishBasicBlock() {
   // Handle stack protector.
   if (SDB->SPDescriptor.shouldEmitFunctionBasedCheckStackProtector()) {
     // The target provides a guard check function. There is no need to
     // generate error handling code or to split current basic block.
     MachineBasicBlock *ParentMBB = SDB->SPDescriptor.getParentMBB();
 
     // Add load and check to the basicblock.
     FuncInfo->MBB = ParentMBB;
     FuncInfo->InsertPt =
-        FindSplitPointForStackProtector(ParentMBB);
+        FindSplitPointForStackProtector(ParentMBB, *TII);
     SDB->visitSPDescriptorParent(SDB->SPDescriptor, ParentMBB);
     CurDAG->setRoot(SDB->getRoot());
     SDB->clear();
     CodeGenAndEmitDAG();
 
     // Clear the Per-BB State.
     SDB->SPDescriptor.resetPerBBState();
   } else if (SDB->SPDescriptor.shouldEmitStackProtector()) {
     MachineBasicBlock *ParentMBB = SDB->SPDescriptor.getParentMBB();
     MachineBasicBlock *SuccessMBB = SDB->SPDescriptor.getSuccessMBB();
 
     // Find the split point to split the parent mbb. At the same time copy all
     // physical registers used in the tail of parent mbb into virtual registers
     // before the split point and back into physical registers after the split
     // point. This prevents us needing to deal with Live-ins and many other
     // register allocation issues caused by us splitting the parent mbb. The
     // register allocator will clean up said virtual copies later on.
     MachineBasicBlock::iterator SplitPoint =
-        FindSplitPointForStackProtector(ParentMBB);
+        FindSplitPointForStackProtector(ParentMBB, *TII);
 
     // Splice the terminator of ParentMBB into SuccessMBB.
     SuccessMBB->splice(SuccessMBB->end(), ParentMBB,
                        SplitPoint,
                        ParentMBB->end());
 
     // Add compare/jump on neq/jump to the parent BB.
     FuncInfo->MBB = ParentMBB;

llvm/lib/CodeGen/StackProtector.cpp

if (!AI) {
       AI = cast<AllocaInst>(SPCall->getArgOperand(1));
     }
 
     // Set HasIRCheck to true, so that SelectionDAG will not generate its own
     // version. SelectionDAG called 'shouldEmitSDCheck' to check whether
     // instrumentation has already been generated.
     HasIRCheck = true;
 
+    // If we're instrumenting a block with a musttail call, the check has to be
+    // inserted before the call rather than between it and the return. The
+    // verifier guarantees that a musttail call is either directly before the
+    // return or with a single correct bitcast of the return value in between so
+    // we don't need to worry about many situations here.
+    Instruction *CheckLoc = RI;
+    Instruction *Prev = RI->getPrevNonDebugInstruction();
+    if (Prev && isa<CallInst>(Prev) && cast<CallInst>(Prev)->isMustTailCall())
+      CheckLoc = Prev;
+    else if (Prev) {
+      Prev = Prev->getPrevNonDebugInstruction();
+      if (Prev && isa<CallInst>(Prev) && cast<CallInst>(Prev)->isMustTailCall())
+        CheckLoc = Prev;
+    }
+
     // Generate epilogue instrumentation. The epilogue intrumentation can be
     // function-based or inlined depending on which mechanism the target is
     // providing.
     if (Function *GuardCheck = TLI->getSSPStackGuardCheck(*M)) {
       // Generate the function-based epilogue instrumentation.
       // The target provides a guard check function, generate a call to it.
-      IRBuilder<> B(RI);
+      IRBuilder<> B(CheckLoc);
       LoadInst *Guard = B.CreateLoad(B.getInt8PtrTy(), AI, true, "Guard");
       CallInst *Call = B.CreateCall(GuardCheck, {Guard});
       Call->setAttributes(GuardCheck->getAttributes());
       Call->setCallingConv(GuardCheck->getCallingConv());
     } else {
       // Generate the epilogue with inline instrumentation.
-      // If we do not support SelectionDAG based tail calls, generate IR level
-      // tail calls.
+      // If we do not support SelectionDAG based calls, generate IR level
+      // calls.
       //
       // For each block with a return instruction, convert this:
       //
       //   return:
       //     ...
       //     ret ...
       //
       // into this:
if (Function *GuardCheck = TLI->getSSPStackGuardCheck(*M)) {
       //     unreachable
 
       // Create the FailBB. We duplicate the BB every time since the MI tail
       // merge pass will merge together all of the various BB into one including
       // fail BB generated by the stack protector pseudo instruction.
       BasicBlock *FailBB = CreateFailBB();
 
       // Split the basic block before the return instruction.
-      BasicBlock *NewBB = BB->splitBasicBlock(RI->getIterator(), "SP_return");
+      BasicBlock *NewBB =
+          BB->splitBasicBlock(CheckLoc->getIterator(), "SP_return");
 
       // Update the dominator tree if we need to.
       if (DT && DT->isReachableFromEntry(BB)) {
         DT->addNewBlock(NewBB, BB);
         DT->addNewBlock(FailBB, BB);
       }
 
       // Remove default branch instruction to the new BB.

llvm/test/CodeGen/AArch64/stack-protector-musttail.ll

+; RUN: llc -mtriple=arm64-apple-macosx -fast-isel %s -o - -start-before=stack-protector -stop-after=stack-protector  | FileCheck %s
+
+@var = global [2 x i64]* null
+
+declare void @callee()
+
+define void @caller1() ssp {
+; CHECK-LABEL: define void @caller1()
+; Prologue:
+; CHECK: @llvm.stackguard
+
+; CHECK: [[GUARD:%.*]] = call i8* @llvm.stackguard()
+; CHECK: [[TOKEN:%.*]] = load volatile i8*, i8** {{%.*}}
+; CHECK: [[TST:%.*]] = icmp eq i8* [[GUARD]], [[TOKEN]]
+; CHECK: br i1 [[TST]]
+
+; CHECK: musttail call void @callee()
+; CHECK-NEXT: ret void
+  %var = alloca [2 x i64]
+  store [2 x i64]* %var, [2 x i64]** @var
+  musttail call void @callee()
+  ret void
+}
+
+define void @justret() ssp {
+; CHECK-LABEL: define void @justret()
+; Prologue:
+; CHECK: @llvm.stackguard
+
+; CHECK: [[GUARD:%.*]] = call i8* @llvm.stackguard()
+; CHECK: [[TOKEN:%.*]] = load volatile i8*, i8** {{%.*}}
+; CHECK: [[TST:%.*]] = icmp eq i8* [[GUARD]], [[TOKEN]]
+; CHECK: br i1 [[TST]]
+
+; CHECK: ret void
+  %var = alloca [2 x i64]
+  store [2 x i64]* %var, [2 x i64]** @var
+  br label %retblock
+
+retblock:
+  ret void
+}
+
+
+declare i64* @callee2()
+
+define i8* @caller2() ssp {
+; CHECK-LABEL: define i8* @caller2()
+; Prologue:
+; CHECK: @llvm.stackguard
+
+; CHECK: [[GUARD:%.*]] = call i8* @llvm.stackguard()
+; CHECK: [[TOKEN:%.*]] = load volatile i8*, i8** {{%.*}}
+; CHECK: [[TST:%.*]] = icmp eq i8* [[GUARD]], [[TOKEN]]
+; CHECK: br i1 [[TST]]
+
+; CHECK: [[TMP:%.*]] = musttail call i64* @callee2()
+; CHECK-NEXT: [[RES:%.*]] = bitcast i64* [[TMP]] to i8*
+; CHECK-NEXT: ret i8* [[RES]]
+
+  %var = alloca [2 x i64]
+  store [2 x i64]* %var, [2 x i64]** @var
+  %tmp = musttail call i64* @callee2()
+  %res = bitcast i64* %tmp to i8*
+  ret i8* %res
+}

llvm/test/CodeGen/ARM/Windows/stack-protector-musttail.ll

+; RUN: llc -mtriple=thumbv7-windows-msvc -fast-isel %s -o - -start-before=stack-protector -stop-after=stack-protector  | FileCheck %s
+
+@var = global [2 x i64]* null
+
+declare void @callee()
+
+define void @caller1() sspreq {
+; CHECK-LABEL: define void @caller1()
+; Prologue:
+
+; CHECK: call void @__security_check_cookie
+
+; CHECK: musttail call void @callee()
+; CHECK-NEXT: ret void
+  %var = alloca [2 x i64]
+  store [2 x i64]* %var, [2 x i64]** @var
+  musttail call void @callee()
+  ret void
+}
+
+define void @justret() sspreq {
+; CHECK-LABEL: define void @justret()
+; Prologue:
+; CHECK: @llvm.stackguard
+
+; CHECK: call void @__security_check_cookie
+
+; CHECK: ret void
+  %var = alloca [2 x i64]
+  store [2 x i64]* %var, [2 x i64]** @var
+  br label %retblock
+
+retblock:
+  ret void
+}
+
+
+declare i64* @callee2()
+
+define i8* @caller2() sspreq {
+; CHECK-LABEL: define i8* @caller2()
+; Prologue:
+; CHECK: @llvm.stackguard
+
+; CHECK: call void @__security_check_cookie
+
+; CHECK: [[TMP:%.*]] = musttail call i64* @callee2()
+; CHECK-NEXT: [[RES:%.*]] = bitcast i64* [[TMP]] to i8*
+; CHECK-NEXT: ret i8* [[RES]]
+
+  %var = alloca [2 x i64]
+  store [2 x i64]* %var, [2 x i64]** @var
+  %tmp = musttail call i64* @callee2()
+  %res = bitcast i64* %tmp to i8*
+  ret i8* %res
+}

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

+; RUN: llc -mtriple=x86_64-windows-msvc %s -o - -verify-machineinstrs | FileCheck %s
+
+declare void @h(i8*, i64, i8*)
+
+define tailcc void @tailcall_frame(i8* %0, i64 %1) sspreq {
+; CHECK-LABEL: tailcall_frame:
+; CHECK: callq __security_check_cookie
+; CHECK: xorl %ecx, %ecx
+; CHECK: jmp h
+
+   tail call tailcc void @h(i8* null, i64 0, i8* null)
+   ret void
+}
+
+declare void @bar()
+define void @tailcall_unrelated_frame() sspreq {
+; CHECK-LABEL: tailcall_unrelated_frame:
+; CHECK: subq [[STACK:\$.*]], %rsp
+; CHECK: callq bar
+; CHECK: callq __security_check_cookie
+; CHECK: addq [[STACK]], %rsp
+; CHECK: jmp bar
+  call void @bar()
+  tail call void @bar()
+  ret void
+}