add support for AFL_LLVM_THREADSAFE_INST to other LLVM passes

1 files changed, 19 insertions, 17 deletions

diff --git a/instrumentation/afl-llvm-lto-instrumentation.so.cc b/instrumentation/afl-llvm-lto-instrumentation.so.cc
index 5ed13ff0..10cfa579 100644
--- a/instrumentation/afl-llvm-lto-instrumentation.so.cc
+++ b/instrumentation/afl-llvm-lto-instrumentation.so.cc
@@ -93,7 +93,8 @@ class AFLLTOPass : public ModulePass {
   uint32_t function_minimum_size = 1;
   uint32_t inst_blocks = 0, inst_funcs = 0, total_instr = 0;
   unsigned long long int map_addr = 0x10000;
-  char *   skip_nozero = NULL;
+  const char *skip_nozero = NULL;
+  const char *use_threadsafe_counters = nullptr;
 
 };
 
@@ -131,6 +132,8 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
     be_quiet = 1;
 
+  use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
+
   if ((ptr = getenv("AFL_LLVM_DOCUMENT_IDS")) != NULL) {
 
     if ((documentFile = fopen(ptr, "a")) == NULL)
@@ -839,29 +842,28 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
           /* Update bitmap */
 
-#if 1 /* Atomic */
-          IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
-              llvm::AtomicOrdering::Monotonic);
+          if (use_threadsafe_counters) {
+            IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                                llvm::AtomicOrdering::Monotonic);
+          } else {
+            LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+            Counter->setMetadata(M.getMDKindID("nosanitize"),
+                                 MDNode::get(C, None));
 
-#else
-          LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-          Counter->setMetadata(M.getMDKindID("nosanitize"),
-                               MDNode::get(C, None));
+            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)
+                ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
           }
 
-          IRB.CreateStore(Incr, MapPtrIdx)
-              ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
-#endif
-
           // done :)
 
           inst_blocks++;