diff --git a/llvm/lib/Transforms/Utils/EscapeEnumerator.cpp b/llvm/lib/Transforms/Utils/EscapeEnumerator.cpp --- a/llvm/lib/Transforms/Utils/EscapeEnumerator.cpp +++ b/llvm/lib/Transforms/Utils/EscapeEnumerator.cpp @@ -41,7 +41,27 @@ if (!isa(TI) && !isa(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(I); + if (!RI || !RI->getPrevNode()) + return I; + auto *CI = dyn_cast(RI->getPrevNode()); + if (CI && CI->isMustTailCall()) + return CI; + auto *BI = dyn_cast(RI->getPrevNode()); + if (!BI || !BI->getPrevNode()) + return I; + CI = dyn_cast(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 Calls; for (BasicBlock &BB : F) for (Instruction &II : BB) if (CallInst *CI = dyn_cast(&II)) - if (!CI->doesNotThrow()) + if (!CI->doesNotThrow() && !CI->isMustTailCall()) Calls.push_back(CI); if (Calls.empty()) diff --git a/llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll b/llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Instrumentation/ThreadSanitizer/tsan_musttail.ll @@ -0,0 +1,28 @@ +; 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: 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: 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