msan: Handle musttail calls

[rnk]

Summary

First, avoid calling setTailCall(false) on musttail calls. The funciton
prototypes should be "congruent", so the shadow layout should be exactly
the same.

Second, avoid inserting instrumentation after a musttail call to
propagate the return value shadow. We don't need to propagate the
result of a tail call, it should already be in the right place.

[eugenis]

LGTM w/ nits

include/llvm/IR/Instructions.h

2418

Are you sure it should be part of ReturnInst interface?
Makes more sense as a utility function to me.
Anyway, this is an unrelated refactoring, please commit separately.

lib/Transforms/Instrumentation/MemorySanitizer.cpp

2276

I no longer understand what is this code trying to achieve. The comment suggests that there is a problem when caller returns void, and that's exactly the case you are enabling for musttail calls.

But this whole thing does not make sense, and I can't figure out what is it that tail call optimization can break. Perhaps it had something to do with the hybrid tool and had been made obsolete by us resetting retval shadow to 0 _before_ every call.

Anyway, the tests still pass if this chunk is removed. I suggest you just do that. If I'm missing something, we will know it soon enough.

[rnk]

include/llvm/IR/Instructions.h

2418

I agree, the IR classes are generally very lightweight and don't have these kinds of analysis helpers.

+chandlerc, WDYT? Is there a better place to put this?

lib/Transforms/Instrumentation/MemorySanitizer.cpp

2276

OK. :) I'll commit that as a separate change without any of the actual musttail handling stuff.

[rnk]

rebase

[rnk]

Closed by commit rL215415 (authored by @rnk).

include/llvm/IR/Instructions.h

public:
   /// Provide fast operand accessors
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Value);
 
   /// Convenience accessor. Returns null if there is no return value.
   Value *getReturnValue() const {
     return getNumOperands() != 0 ? getOperand(0) : nullptr;
   }
 
+  /// Returns a musttail call instruction if one immediately precedes the given
+  /// return instruction with an optional bitcast instruction between them.
+  CallInst *getPrecedingMustTailCallIfPresent();
+
   unsigned getNumSuccessors() const { return 0; }
 
   // Methods for support type inquiry through isa, cast, and dyn_cast:
   static inline bool classof(const Instruction *I) {
     return (I->getOpcode() == Instruction::Ret);
   }
   static inline bool classof(const Value *V) {
     return isa<Instruction>(V) && classof(cast<Instruction>(V));

lib/IR/Instructions.cpp

 
 BasicBlock *ReturnInst::getSuccessorV(unsigned idx) const {
   llvm_unreachable("ReturnInst has no successors!");
 }
 
 ReturnInst::~ReturnInst() {
 }
 
+CallInst *ReturnInst::getPrecedingMustTailCallIfPresent() {
+  if (this == getParent()->begin())
+    return nullptr;
+
+  Instruction *Prev = getPrevNode();
+  if (!Prev)
+    return nullptr;
+
+  if (Value *RV = getReturnValue()) {
+    if (RV != Prev)
+      return nullptr;
+
+    // Look through the optional bitcast.
+    if (auto *BI = dyn_cast<BitCastInst>(Prev)) {
+      RV = BI->getOperand(0);
+      Prev = BI->getPrevNode();
+      if (!Prev || RV != Prev)
+        return nullptr;
+    }
+  }
+
+  if (auto *CI = dyn_cast<CallInst>(Prev)) {
+    if (CI->isMustTailCall())
+      return CI;
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 //                        ResumeInst Implementation
 //===----------------------------------------------------------------------===//
 
 ResumeInst::ResumeInst(const ResumeInst &RI)
   : TerminatorInst(Type::getVoidTy(RI.getContext()), Instruction::Resume,
                    OperandTraits<ResumeInst>::op_begin(this), 1) {
   Op<0>() = RI.Op<0>();

lib/Transforms/Instrumentation/MemorySanitizer.cpp

if (CS.isCall()) {
       // outputs as clean. Note that any side effects of the inline asm that are
       // not immediately visible in its constraints are not handled.
       if (Call->isInlineAsm()) {
         visitInstruction(I);
         return;
       }
 
       assert(!isa<IntrinsicInst>(&I) && "intrinsics are handled elsewhere");
 
       // We are going to insert code that relies on the fact that the callee
       // will become a non-readonly function after it is instrumented by us. To
       // prevent this code from being optimized out, mark that function
       // non-readonly in advance.
       if (Function *Func = Call->getCalledFunction()) {
         // Clear out readonly/readnone attributes.
         AttrBuilder B;
         B.addAttribute(Attribute::ReadOnly)
void visitCallSite(CallSite CS) {
     DEBUG(dbgs() << "  done with call args\n");
 
     FunctionType *FT =
       cast<FunctionType>(CS.getCalledValue()->getType()->getContainedType(0));
     if (FT->isVarArg()) {
       VAHelper->visitCallSite(CS, IRB);
     }
 
+    // If this is a musttail call site, we can't insert propagation code here.
+    // The return type of the caller must match the callee, so the shadow should
+    // already be set up for an immediate return.
+    if (CS.isMustTailCall())
+      return;
+
     // Now, get the shadow for the RetVal.
     if (!I.getType()->isSized()) return;
     IRBuilder<> IRBBefore(&I);
     // Until we have full dynamic coverage, make sure the retval shadow is 0.
     Value *Base = getShadowPtrForRetval(&I, IRBBefore);
     IRBBefore.CreateAlignedStore(getCleanShadow(&I), Base, kShadowTLSAlignment);
     Instruction *NextInsn = nullptr;
     if (CS.isCall()) {
Value *RetvalShadow =
       IRBAfter.CreateAlignedLoad(getShadowPtrForRetval(&I, IRBAfter),
                                  kShadowTLSAlignment, "_msret");
     setShadow(&I, RetvalShadow);
     if (MS.TrackOrigins)
       setOrigin(&I, IRBAfter.CreateLoad(getOriginPtrForRetval(IRBAfter)));
   }
 
   void visitReturnInst(ReturnInst &I) {
+    // Don't propagate shadow between musttail calls and the return.
+    if (I.getPrecedingMustTailCallIfPresent())
+      return;
+
     IRBuilder<> IRB(&I);
     Value *RetVal = I.getReturnValue();
     if (!RetVal) return;
     Value *ShadowPtr = getShadowPtrForRetval(RetVal, IRB);
     if (CheckReturnValue) {
       insertShadowCheck(RetVal, &I);
       Value *Shadow = getCleanShadow(RetVal);
       IRB.CreateAlignedStore(Shadow, ShadowPtr, kShadowTLSAlignment);

lib/Transforms/Utils/InlineFunction.cpp

for (BasicBlock::iterator BI = FI->begin(), BE = FI->end();
           DVI->setOperand(2, createInlinedVariable(DVI->getVariable(), 
                                                    InlinedAt, Ctx));
         }
       }
     }
   }
 }
 
-/// Returns a musttail call instruction if one immediately precedes the given
-/// return instruction with an optional bitcast instruction between them.
-static CallInst *getPrecedingMustTailCall(ReturnInst *RI) {
-  Instruction *Prev = RI->getPrevNode();
-  if (!Prev)
-    return nullptr;
-
-  if (Value *RV = RI->getReturnValue()) {
-    if (RV != Prev)
-      return nullptr;
-
-    // Look through the optional bitcast.
-    if (auto *BI = dyn_cast<BitCastInst>(Prev)) {
-      RV = BI->getOperand(0);
-      Prev = BI->getPrevNode();
-      if (!Prev || RV != Prev)
-        return nullptr;
-    }
-  }
-
-  if (auto *CI = dyn_cast<CallInst>(Prev)) {
-    if (CI->isMustTailCall())
-      return CI;
-  }
-  return nullptr;
-}
-
 /// InlineFunction - This function inlines the called function into the basic
 /// block of the caller.  This returns false if it is not possible to inline
 /// this call.  The program is still in a well defined state if this occurs
 /// though.
 ///
 /// Note that this only does one level of inlining.  For example, if the
 /// instruction 'call B' is inlined, and 'B' calls 'C', then the call to 'C' now
 /// exists in the instruction stream.  Similarly this will inline a recursive
for (unsigned ai = 0, ae = IFI.StaticAllocas.size(); ai != ae; ++ai) {
           }
         }
       }
 
       builder.CreateLifetimeStart(AI, AllocaSize);
       for (ReturnInst *RI : Returns) {
         // Don't insert llvm.lifetime.end calls between a musttail call and a
         // return.  The return kills all local allocas.
-        if (InlinedMustTailCalls && getPrecedingMustTailCall(RI))
+        if (InlinedMustTailCalls && RI->getPrecedingMustTailCallIfPresent())
           continue;
         IRBuilder<>(RI).CreateLifetimeEnd(AI, AllocaSize);
       }
     }
   }
 
   // If the inlined code contained dynamic alloca instructions, wrap the inlined
   // code with llvm.stacksave/llvm.stackrestore intrinsics.
   if (InlinedFunctionInfo.ContainsDynamicAllocas) {
     Module *M = Caller->getParent();
     // Get the two intrinsics we care about.
     Function *StackSave = Intrinsic::getDeclaration(M, Intrinsic::stacksave);
     Function *StackRestore=Intrinsic::getDeclaration(M,Intrinsic::stackrestore);
 
     // Insert the llvm.stacksave.
     CallInst *SavedPtr = IRBuilder<>(FirstNewBlock, FirstNewBlock->begin())
       .CreateCall(StackSave, "savedstack");
 
     // Insert a call to llvm.stackrestore before any return instructions in the
     // inlined function.
     for (ReturnInst *RI : Returns) {
       // Don't insert llvm.stackrestore calls between a musttail call and a
       // return.  The return will restore the stack pointer.
-      if (InlinedMustTailCalls && getPrecedingMustTailCall(RI))
+      if (InlinedMustTailCalls && RI->getPrecedingMustTailCallIfPresent())
         continue;
       IRBuilder<>(RI).CreateCall(StackRestore, SavedPtr);
     }
   }
 
   // If we are inlining for an invoke instruction, we must make sure to rewrite
   // any call instructions into invoke instructions.
   if (InvokeInst *II = dyn_cast<InvokeInst>(TheCall))
     HandleInlinedInvoke(II, FirstNewBlock, InlinedFunctionInfo);
 
   // Handle any inlined musttail call sites.  In order for a new call site to be
   // musttail, the source of the clone and the inlined call site must have been
   // musttail.  Therefore it's safe to return without merging control into the
   // phi below.
   if (InlinedMustTailCalls) {
     // Check if we need to bitcast the result of any musttail calls.
     Type *NewRetTy = Caller->getReturnType();
     bool NeedBitCast = !TheCall->use_empty() && TheCall->getType() != NewRetTy;
 
     // Handle the returns preceded by musttail calls separately.
     SmallVector<ReturnInst *, 8> NormalReturns;
     for (ReturnInst *RI : Returns) {
-      CallInst *ReturnedMustTail = getPrecedingMustTailCall(RI);
+      CallInst *ReturnedMustTail = RI->getPrecedingMustTailCallIfPresent();
       if (!ReturnedMustTail) {
         NormalReturns.push_back(RI);
         continue;
       }
       if (!NeedBitCast)
         continue;
 
       // Delete the old return and any preceding bitcast.

test/Instrumentation/MemorySanitizer/msan_basic.ll

 }
 
 ; We used to strip off the 'tail' modifier, but now that we unpoison return slot
 ; shadow before the call, we don't need to anymore.
 
 ; CHECK-LABEL: define void @MismatchedReturnTypeTailCall
 ; CHECK: tail call i32 @InnerTailCall
 ; CHECK: ret void
+
+declare i32 @InnerMustTailCall(i32 %a)
+
+define i32 @MustTailCall(i32 %a) {
+  %b = musttail call i32 @InnerMustTailCall(i32 %a)
+  ret i32 %b
+}
+
+; Test that 'musttail' is preserved.  The ABI should make this work.
+
+; CHECK-LABEL: define i32 @MustTailCall
+; CHECK: musttail call i32 @InnerMustTailCall
+; CHECK-NEXT: ret i32