[TSAN] Handle musttail call properly in EscapeEnumerator (and TSAN)

[lxfind]

Summary

Call instructions with musttail tag must be optimized as a tailcall, otherwise could lead to incorrect program behavior.
When TSAN is instrumenting functions, it broke the contract by adding a call to the tsan exit function inbetween the musttail call and return instruction, and also inserted exception handling code.
This happend throguh EscapeEnumerator, which adds exception handling code and returns ret instructions as the place to insert instrumentation calls.
This becomes especially problematic for coroutines, because coroutines rely on tail calls to do symmetric transfers properly.
To fix this, this patch moves the location to insert instrumentation calls prior to the musttail call for ret instructions that are following musttail calls, and also does not handle exception for musttail calls.

[wenlei]

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

111

Does TSAN rely on exit cleanup instrumentation for correctness? How does it work for existing Invoke?

This patch would be fine if missing some exit instrumentation is ok in general..

[lxfind]

fix test failures

[lxfind]

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

111

The change would just make the call to tsan exit function happen earlier, before the tail call. So we are not missing any exit instrumentation after this change.

[wenlei]

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

111

The change would just make the call to tsan exit function happen earlier, before the tail call. So we are not missing any exit instrumentation after this change.

That's the part from AdjustMustTailCall.

IIUC, the other part on line 82 would skip changing throwing tail call to invoke, which would also skip the clean up in landing pad. That's what I was referring to,

[lxfind]

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

111

Ah I see. So the reason EscapeEnumerator tries to turn calls into invoke with an exception return path is so that when exceptions are thrown, it still has a chance to call the tsan exit function before propagating the exception back.
Now that this changes make the call to tsan exit function before the musttail call, we no longer need to try to catch the exception after the musttail call and call exit function again if it ever throws.

[wenlei]

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

111

Ok, yeah. Tailcall is different in that it's always going to be the block terminator, so guaranteed to be covered by the loop over blocks, unlike regular calls. So we're not missing out exit cleanup here. (I'm still curious how pre-existing invoke is handled, because we don't add cleanup to existing landing pad. But that is not directly related to this patch)

I think there's a still slight difference though for the order of cleanup. Now between cleanup for the caller of tail call and the actual return from the caller, there's the entire tail callee. I would think that is ok, but it really depends on the contract of TSAN, so would be nice if someone familiar with TSAN can confirm.

Otherwise LGTM, accept to unblock.

llvm/lib/Transforms/Utils/EscapeEnumerator.cpp

while (StateBB != StateE) {
     BasicBlock *CurBB = &*StateBB++;
 
     // Branches and invokes do not escape, only unwind, resume, and return
     // do.
     Instruction *TI = CurBB->getTerminator();
     if (!isa<ReturnInst>(TI) && !isa<ResumeInst>(TI))
       continue;
 
-    Builder.SetInsertPoint(TI);
+    // If the ret instruction is followed by a musttaill call,
+    // or a bitcast instruction and then a musttail call, we should return
+    // the musttail call as the insertion point to not break the musttail
+    // contract.
+    auto AdjustMustTailCall = [&](Instruction *I) -> Instruction * {
+      auto *RI = dyn_cast<ReturnInst>(I);
+      if (!RI || !RI->getPrevNode())
+        return I;
+      auto *CI = dyn_cast<CallInst>(RI->getPrevNode());
+      if (CI && CI->isMustTailCall())
+        return CI;
+      auto *BI = dyn_cast<BitCastInst>(RI->getPrevNode());
+      if (!BI || !BI->getPrevNode())
+        return I;
+      CI = dyn_cast<CallInst>(BI->getPrevNode());
+      if (CI && CI->isMustTailCall())
+        return CI;
+      return I;
+    };
+
+    Builder.SetInsertPoint(AdjustMustTailCall(TI));
     return &Builder;
   }
 
   Done = true;
 
   if (!HandleExceptions)
     return nullptr;
 
   if (F.doesNotThrow())
     return nullptr;
 
   // Find all 'call' instructions that may throw.
+  // We cannot tranform calls with musttail tag.
   SmallVector<Instruction *, 16> Calls;
   for (BasicBlock &BB : F)
     for (Instruction &II : BB)
       if (CallInst *CI = dyn_cast<CallInst>(&II))
-        if (!CI->doesNotThrow())
+        if (!CI->doesNotThrow() && !CI->isMustTailCall())
           Calls.push_back(CI);
 
   if (Calls.empty())
     return nullptr;
 
   // Create a cleanup block.
   LLVMContext &C = F.getContext();
   BasicBlock *CleanupBB = BasicBlock::Create(C, CleanupBBName, &F);
IRBuilder<> *EscapeEnumerator::Next() {
   LPad->setCleanup(true);
   ResumeInst *RI = ResumeInst::Create(LPad, CleanupBB);
 
   // Transform the 'call' instructions into 'invoke's branching to the
   // cleanup block. Go in reverse order to make prettier BB names.
   SmallVector<Value *, 16> Args;
   for (unsigned I = Calls.size(); I != 0;) {
     CallInst *CI = cast<CallInst>(Calls[--I]);
     changeToInvokeAndSplitBasicBlock(CI, CleanupBB);
   }
 
   Builder.SetInsertPoint(RI);
   return &Builder;
 }

llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll

+; To test that __tsan_func_exit always happen before musttaill call and no exception handling code.
+; RUN: opt < %s -tsan -S | FileCheck %s
+
+define internal i32 @preallocated_musttail(i32* preallocated(i32) %p) sanitize_thread {
+  %rv = load i32, i32* %p
+  ret i32 %rv
+}
+
+define i32 @call_preallocated_musttail(i32* preallocated(i32) %a) sanitize_thread {
+  %r = musttail call i32 @preallocated_musttail(i32* preallocated(i32) %a)
+  ret i32 %r
+}
+
+; CHECK-LABEL:  define i32 @call_preallocated_musttail(i32* preallocated(i32) %a) 
+; CHECK:          call void @__tsan_func_exit()
+; CHECK-NEXT:     %r = musttail call i32 @preallocated_musttail(i32* preallocated(i32) %a)
+; CHECK-NEXT:     ret i32 %r
+
+
+define i32 @call_preallocated_musttail_cast(i32* preallocated(i32) %a) sanitize_thread {
+  %r = musttail call i32 @preallocated_musttail(i32* preallocated(i32) %a)
+  %t = bitcast i32 %r to i32
+  ret i32 %t
+}
+
+; CHECK-LABEL:  define i32 @call_preallocated_musttail_cast(i32* preallocated(i32) %a)
+; CHECK:          call void @__tsan_func_exit()
+; CHECK-NEXT:     %r = musttail call i32 @preallocated_musttail(i32* preallocated(i32) %a)
+; CHECK-NEXT:     %t = bitcast i32 %r to i32
+; CHECK-NEXT:     ret i32 %t