1 files changed, 21 insertions, 14 deletions
diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc
index dd2e1459..e1e922be 100644
--- a/instrumentation/SanitizerCoveragePCGUARD.so.cc
+++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -96,7 +96,8 @@ static const char *const SanCovPCsSectionName = "sancov_pcs";
 
 static const char *const SanCovLowestStackName = "__sancov_lowest_stack";
 
-static char *skip_nozero;
+static const char *skip_nozero;
+static const char *use_threadsafe_counters;
 
 namespace {
 
@@ -396,6 +397,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
     be_quiet = 1;
 
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
+  use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
 
   initInstrumentList();
   scanForDangerousFunctions(&M);
@@ -1081,27 +1083,32 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
 
     Value *   MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
 
-#if 1 /* Atomic */
-    IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
-        llvm::AtomicOrdering::Monotonic);
+    if (use_threadsafe_counters) {
 
-#else
-    LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                          llvm::AtomicOrdering::Monotonic);
+
+    }
+    else
+    {
 
-    /* Update bitmap */
+      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      /* Update bitmap */
 
-    Value *Incr = IRB.CreateAdd(Counter, One);
+      Value *Incr = IRB.CreateAdd(Counter, One);
 
-    if (skip_nozero == NULL) {
+      if (skip_nozero == NULL) {
 
-      auto cf = IRB.CreateICmpEQ(Incr, Zero);
-      auto carry = IRB.CreateZExt(cf, Int8Ty);
-      Incr = IRB.CreateAdd(Incr, carry);
+        auto cf = IRB.CreateICmpEQ(Incr, Zero);
+        auto carry = IRB.CreateZExt(cf, Int8Ty);
+        Incr = IRB.CreateAdd(Incr, carry);
+
+      }
+
+      IRB.CreateStore(Incr, MapPtrIdx);
 
     }
 
-    IRB.CreateStore(Incr, MapPtrIdx);
-#endif
     // done :)
 
     //    IRB.CreateCall(SanCovTracePCGuard, Offset)->setCannotMerge();