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

[Xun Li via Phabricator via llvm-commits]

lxfind updated this revision to Diff 291653.
lxfind added a comment.

fix test failures


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D87620/new/

https://reviews.llvm.org/D87620

Files:
  llvm/lib/Transforms/Utils/EscapeEnumerator.cpp
  llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll


Index: llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll
===================================================================
--- /dev/null
+++ llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll
@@ -0,0 +1,30 @@
+; 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
Index: llvm/lib/Transforms/Utils/EscapeEnumerator.cpp
===================================================================
--- llvm/lib/Transforms/Utils/EscapeEnumerator.cpp
+++ llvm/lib/Transforms/Utils/EscapeEnumerator.cpp
@@ -41,7 +41,27 @@
     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;
   }
 
@@ -54,11 +74,12 @@
     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())


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D87620.291653.patch
Type: text/x-patch
Size: 3117 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200914/e5451b47/attachment.bin>