[DevTSAN] Sync thread clock to global state before kernel exit

clang/lib/Driver/SanitizerArgs.cpp

@@ -1281,10 +1281,6 @@ void SanitizerArgs::addArgs(const ToolChain &TC, const llvm::opt::ArgList &Args,
       CmdArgs.push_back("-msan-poison-stack-with-call=1");
     } else if (Sanitizers.has(SanitizerKind::Thread)) {
       CmdArgs.push_back("-fsanitize=thread");
-      // The tsan function entry/exit builtins are used to record stack
-      // position, we don't need them in device offloading.
-      CmdArgs.push_back("-mllvm");
-      CmdArgs.push_back("-tsan-instrument-func-entry-exit=0");
       // In device offloading, user can't call memory instrinsics explicitly, so
       // we can safely skip them.
       CmdArgs.push_back("-mllvm");

libdevice/sanitizer/tsan_rtl.cpp

@@ -464,6 +464,13 @@ DEVICE_EXTERN_C_NOINLINE void __tsan_cleanup_private(uptr addr, uint32_t size) {
 #endif
 }
 
+DEVICE_EXTERN_C_NOINLINE void __tsan_func_exit() {
+  Sid sid = GetCurrentSid();
+  // sync current thread clock to global state
+  TsanLaunchInfo->Clock[kThreadSlotCount].clk_[sid] =
+      TsanLaunchInfo->Clock[sid].clk_[sid];
+}
+
 DEVICE_EXTERN_C_INLINE void __tsan_device_barrier() {
   Sid sid = GetCurrentSid();
   __spirv_ControlBarrier(__spv::Scope::Device, __spv::Scope::Device,

llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp

@@ -361,8 +361,7 @@ void ThreadSanitizerOnSpirv::appendDebugInfoToArgs(
   auto &Loc = I->getDebugLoc();
 
   // SPIR constant address space
-  PointerType *ConstASPtrTy =
-      PointerType::get(Type::getInt8Ty(C), kSpirOffloadConstantAS);
+  PointerType *ConstASPtrTy = PointerType::get(C, kSpirOffloadConstantAS);
 
   // File & Line
   if (Loc) {
@@ -959,19 +958,25 @@ bool ThreadSanitizer::sanitizeFunction(Function &F,
     Res |= Spirv->instrumentAllocInst(&F, Allocas);
 
   // Instrument function entry/exit points if there were instrumented accesses.
-  if ((Res || HasCalls) && ClInstrumentFuncEntryExit) {
-    InstrumentationIRBuilder IRB(&F.getEntryBlock(),
-                                 F.getEntryBlock().getFirstNonPHIIt());
-    Value *ReturnAddress =
-        IRB.CreateIntrinsic(Intrinsic::returnaddress, IRB.getInt32(0));
-    IRB.CreateCall(TsanFuncEntry, ReturnAddress);
-
-    EscapeEnumerator EE(F, "tsan_cleanup", ClHandleCxxExceptions);
-    while (IRBuilder<> *AtExit = EE.Next()) {
-      InstrumentationIRBuilder::ensureDebugInfo(*AtExit, F);
-      AtExit->CreateCall(TsanFuncExit, {});
+  // Alway instrument function exit points for spir kernel.
+  if ((Res || HasCalls || Spirv) && ClInstrumentFuncEntryExit) {
+    if (!Spirv) {
+      InstrumentationIRBuilder IRB(&F.getEntryBlock(),
+                                   F.getEntryBlock().getFirstNonPHIIt());
+      Value *ReturnAddress =
+          IRB.CreateIntrinsic(Intrinsic::returnaddress, IRB.getInt32(0));
+      IRB.CreateCall(TsanFuncEntry, ReturnAddress);
+      Res |= true;
+    }
+
+    if (!Spirv || F.getCallingConv() == CallingConv::SPIR_KERNEL) {
+      EscapeEnumerator EE(F, "tsan_cleanup", ClHandleCxxExceptions);
+      while (IRBuilder<> *AtExit = EE.Next()) {
+        InstrumentationIRBuilder::ensureDebugInfo(*AtExit, F);
+        AtExit->CreateCall(TsanFuncExit, {});
+      }
+      Res |= true;
     }
-    Res = true;
   }
   return Res;
 }

llvm/test/Instrumentation/ThreadSanitizer/SPIRV/basic.ll

@@ -1,7 +1,15 @@
-; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-func-entry-exit=0 -tsan-instrument-memintrinsics=0 -S | FileCheck %s
+; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-memintrinsics=0 -S | FileCheck %s
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
 target triple = "spir64-unknown-unknown"
 
+; Function Attrs: sanitize_thread
+define spir_kernel void @test(i32 %val) #0 {
+; CHECK-LABEL: void @test
+entry:
+  ; CHECK: call void @__tsan_func_exit()
+  ret void
+}
+
 ; Function Attrs: sanitize_thread
 define linkonce_odr dso_local spir_func void @write_1_byte(ptr addrspace(4) %a) #0 {
 ; CHECK-LABEL: void @write_1_byte
@@ -11,6 +19,7 @@ entry:
   ; CHECK: ptrtoint ptr addrspace(4) %a to i64
   ; CHECK: call void @__tsan_write1
   store i8 %inc, ptr addrspace(4) %a, align 1
+  ; CHECK-NOT: call void @__tsan_func_exit()
   ret void
 }
 

llvm/test/Instrumentation/ThreadSanitizer/SPIRV/cleanup_private_shadow.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-func-entry-exit=0 -tsan-instrument-memintrinsics=0 -S | FileCheck %s
+; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-memintrinsics=0 -S | FileCheck %s
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
 target triple = "spir64-unknown-unknown"
 

llvm/test/Instrumentation/ThreadSanitizer/SPIRV/instrument_barrier.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-func-entry-exit=0 -tsan-instrument-memintrinsics=0 -S | FileCheck %s
+; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-memintrinsics=0 -S | FileCheck %s
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
 target triple = "spir64-unknown-unknown"
 

llvm/test/Instrumentation/ThreadSanitizer/SPIRV/instrument_device_global.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-func-entry-exit=0 -tsan-instrument-memintrinsics=0 -S | FileCheck %s
+; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-memintrinsics=0 -S | FileCheck %s
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
 target triple = "spir64-unknown-unknown"
 

llvm/test/Instrumentation/ThreadSanitizer/SPIRV/kernel_metadata.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-func-entry-exit=0 -tsan-instrument-memintrinsics=0 -S | FileCheck %s
+; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-memintrinsics=0 -S | FileCheck %s
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
 target triple = "spir64-unknown-unknown"