Add __tsan_enable/__tsan_disalbe intrinsics

Author: apaznikov

compiler-rt/include/sanitizer/tsan_interface.h

@@ -309,6 +309,12 @@ int SANITIZER_CDECL __tsan_get_report_unique_tid(void *report,
 /// \returns An opaque pointer to the current report. Otherwise returns NULL.
 void *SANITIZER_CDECL __tsan_get_current_report();
 
+/////////////////////////////////////////////////////////////////////////////
+// This code is for disabling/enabling instrumentation
+// for specific code section
+void SANITIZER_CDECL __tsan_enable();
+void SANITIZER_CDECL __tsan_disable();
+
 #ifdef __cplusplus
 } // extern "C"
 #endif

compiler-rt/lib/tsan/rtl/tsan_interface.h

@@ -437,6 +437,10 @@ SANITIZER_INTERFACE_ATTRIBUTE
 void __tsan_go_atomic64_compare_exchange(ThreadState *thr, uptr cpc, uptr pc,
                                          u8 *a);
 
+// Functions for disabling/enabling instrumentation
+// for specific code section
+SANITIZER_INTERFACE_ATTRIBUTE void __tsan_disable();
+SANITIZER_INTERFACE_ATTRIBUTE void __tsan_enable();
 }  // extern "C"
 
 }  // namespace __tsan

compiler-rt/lib/tsan/rtl/tsan_interface.inc

@@ -188,3 +188,9 @@ void __tsan_read_range_pc(void *addr, uptr size, void *pc) {
 void __tsan_write_range_pc(void *addr, uptr size, void *pc) {
   MemoryAccessRange(cur_thread(), STRIP_PAC_PC(pc), (uptr)addr, size, true);
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// This is for disabling/enabling TSan instrumentation
+int GV;
+SANITIZER_INTERFACE_ATTRIBUTE void __tsan_disable() { GV++; };
+SANITIZER_INTERFACE_ATTRIBUTE void __tsan_enable() { GV++; };

llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp

@@ -211,6 +211,10 @@ struct ThreadSanitizer {
   FunctionCallee TsanVptrUpdate;
   FunctionCallee TsanVptrLoad;
   FunctionCallee MemmoveFn, MemcpyFn, MemsetFn;
+
+  // Instrinsics for disabling/enabling instrumentation
+  // for specific code section
+  FunctionCallee TsanDisableFn, TsanEnableFn;
 };
 
 void insertModuleCtor(Module &M) {
@@ -414,6 +418,12 @@ void ThreadSanitizer::initialize(Module &M, const TargetLibraryInfo &TLI) {
       "__tsan_memset",
       TLI.getAttrList(&Ctx, {1}, /*Signed=*/true, /*Ret=*/false, Attr),
       IRB.getPtrTy(), IRB.getPtrTy(), IRB.getInt32Ty(), IntptrTy);
+
+  /////////////////////////////////////////////////////////////////////////////
+  // This code is for disabling/enabling instrumentation
+  // for specific code section
+  TsanEnableFn = M.getOrInsertFunction("__tsan_enable", Attr, IRB.getVoidTy());
+  TsanDisableFn = M.getOrInsertFunction("__tsan_disable", Attr, IRB.getVoidTy());
 }
 
 static bool isVtableAccess(Instruction *I) {
@@ -649,11 +659,42 @@ bool ThreadSanitizer::sanitizeFunction(
   bool SanitizeFunction = F.hasFnAttribute(Attribute::SanitizeThread);
   const DataLayout &DL = F.getParent()->getDataLayout();
 
+  ///////////////////////////////////////////////////////////////////////////////
+  // List of instructions (function calls) to delete
+  SmallVector<CallInst*, 8> EnableDisableFuncToDelete;
+
+  // Counter for considering nesting __tsan_disable/__tsan_enable
+  int disableEnableNesting = 0;
+
   // Traverse all instructions, collect loads/stores/returns, check for calls.
   for (auto &BB : F) {
     LLVM_DEBUG(dbgs() << "\nInstrumenting BB: " << BB.getName() << "\n");
     for (auto &Inst : BB) {
       LLVM_DEBUG(dbgs() << "Instrumenting I: " << Inst << "\n");
+            ///////////////////////////////////////////////////////////////////////////////
+      // This code is for disabling/enabling instrumentation
+      // for specific code section
+      //
+      // Check whether we encountered TSan disable/enable intrinsics
+      if (auto *CI = dyn_cast<CallInst>(&Inst)) {
+        Function *CalledFunc = CI->getCalledFunction();
+        if (CalledFunc) {
+          // errs() << "CalledFunc: " << CalledFunc->getName() << "\n";
+          if (CalledFunc->getName() == "__tsan_disable") {
+            disableEnableNesting++;
+            EnableDisableFuncToDelete.push_back(CI);
+            continue;
+          } else if (CalledFunc->getName() == "__tsan_enable") {
+            disableEnableNesting--;
+            EnableDisableFuncToDelete.push_back(CI);
+            continue;
+          }
+        }
+      }
+
+      if (disableEnableNesting > 0)
+        continue;
+      ///////////////////////////////////////////////////////////////////////////////
 
       // Skip instructions inserted by another instrumentation.
       if (Inst.hasMetadata(LLVMContext::MD_nosanitize))
@@ -688,6 +729,22 @@ bool ThreadSanitizer::sanitizeFunction(
                                    EAI, EAIGlobal);
   }
 
+  //////////////////////////////////////////////////////////////////////////////
+  // This is for disabling/enabling TSan instrumentation
+  // Disable for now check of matching
+  /*
+  if (disableEnableNesting != 0) {
+    // Error: __tsan_enable without __tsan_disable
+    report_fatal_error("Unmatched __tsan_enable__ call in function " +
+                       F.getName());
+  }
+  */
+
+  // Erase all __tsan_disable/enable functions
+  for (auto *CI: EnableDisableFuncToDelete)
+    CI->eraseFromParent();
+  //////////////////////////////////////////////////////////////////////////////
+
   // We have collected all loads and stores.
   // FIXME: many of these accesses do not need to be checked for races
   // (e.g. variables that do not escape, etc).