[AddressSanitizer] Split out memory intrinsic handling

Summary:
In both AddressSanitizer and HWAddressSanitizer, we first collect
instructions whose operands should be instrumented and memory intrinsics,
then instrument them. Both during collection and when inserting
instrumentation, they are handled separately.

Collect them separately and instrument them separately. This is a bit
more straightforward, and prepares for collecting operands instead of
instructions in a future patch.

This is patch 2/4 of a patch series:
https://reviews.llvm.org/D77616 [PATCH 1/4] [AddressSanitizer] Refactor ClDebug{Min,Max} handling
https://reviews.llvm.org/D77617 [PATCH 2/4] [AddressSanitizer] Split out memory intrinsic handling
https://reviews.llvm.org/D77618 [PATCH 3/4] [AddressSanitizer] Refactor: Permit >1 interesting operands per instruction
https://reviews.llvm.org/D77619 [PATCH 4/4] [AddressSanitizer] Instrument byval call arguments

Reviewers: kcc, glider

Reviewed By: glider

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D77617

Author: thejh

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -2652,6 +2652,7 @@ bool AddressSanitizer::instrumentFunction(Function &F,
   // are calls between uses).
   SmallPtrSet<Value *, 16> TempsToInstrument;
   SmallVector<Instruction *, 16> ToInstrument;
+  SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
   SmallVector<Instruction *, 8> NoReturnCalls;
   SmallVector<BasicBlock *, 16> AllBlocks;
   SmallVector<Instruction *, 16> PointerComparisonsOrSubtracts;
@@ -2688,8 +2689,11 @@ bool AddressSanitizer::instrumentFunction(Function &F,
                   isInterestingPointerSubtraction(&Inst))) {
         PointerComparisonsOrSubtracts.push_back(&Inst);
         continue;
-      } else if (isa<MemIntrinsic>(Inst)) {
+      } else if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(&Inst)) {
         // ok, take it.
+        IntrinToInstrument.push_back(MI);
+        NumInsnsPerBB++;
+        continue;
       } else {
         if (isa<AllocaInst>(Inst)) NumAllocas++;
         if (auto *CB = dyn_cast<CallBase>(&Inst)) {
@@ -2708,9 +2712,9 @@ bool AddressSanitizer::instrumentFunction(Function &F,
     }
   }
 
-  bool UseCalls =
-      (ClInstrumentationWithCallsThreshold >= 0 &&
-       ToInstrument.size() > (unsigned)ClInstrumentationWithCallsThreshold);
+  bool UseCalls = (ClInstrumentationWithCallsThreshold >= 0 &&
+                   ToInstrument.size() + IntrinToInstrument.size() >
+                       (unsigned)ClInstrumentationWithCallsThreshold);
   const DataLayout &DL = F.getParent()->getDataLayout();
   ObjectSizeOpts ObjSizeOpts;
   ObjSizeOpts.RoundToAlign = true;
@@ -2723,9 +2727,11 @@ bool AddressSanitizer::instrumentFunction(Function &F,
       if (isInterestingMemoryAccess(Inst, &IsWrite, &TypeSize, &Alignment))
         instrumentMop(ObjSizeVis, Inst, UseCalls,
                       F.getParent()->getDataLayout());
-      else
-        instrumentMemIntrinsic(cast<MemIntrinsic>(Inst));
     }
+  }
+  for (auto Inst : IntrinToInstrument) {
+    if (!suppressInstrumentationSiteForDebug(NumInstrumented))
+      instrumentMemIntrinsic(Inst);
     FunctionModified = true;
   }
 

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -720,11 +720,6 @@ bool HWAddressSanitizer::instrumentMemAccess(Instruction *I) {
   uint64_t TypeSize = 0;
   Value *MaybeMask = nullptr;
 
-  if (ClInstrumentMemIntrinsics && isa<MemIntrinsic>(I)) {
-    instrumentMemIntrinsic(cast<MemIntrinsic>(I));
-    return true;
-  }
-
   Value *Addr =
       isInterestingMemoryAccess(I, &IsWrite, &TypeSize, &Alignment, &MaybeMask);
 
@@ -1090,6 +1085,7 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F) {
   LLVM_DEBUG(dbgs() << "Function: " << F.getName() << "\n");
 
   SmallVector<Instruction*, 16> ToInstrument;
+  SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
   SmallVector<AllocaInst*, 8> AllocasToInstrument;
   SmallVector<Instruction*, 8> RetVec;
   SmallVector<Instruction*, 8> LandingPadVec;
@@ -1121,8 +1117,11 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F) {
       uint64_t TypeSize;
       Value *Addr = isInterestingMemoryAccess(&Inst, &IsWrite, &TypeSize,
                                               &Alignment, &MaybeMask);
-      if (Addr || isa<MemIntrinsic>(Inst))
+      if (Addr)
         ToInstrument.push_back(&Inst);
+
+      if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(&Inst))
+        IntrinToInstrument.push_back(MI);
     }
   }
 
@@ -1138,7 +1137,8 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F) {
     F.setPersonalityFn(nullptr);
   }
 
-  if (AllocasToInstrument.empty() && ToInstrument.empty())
+  if (AllocasToInstrument.empty() && ToInstrument.empty() &&
+      IntrinToInstrument.empty())
     return false;
 
   assert(!LocalDynamicShadow);
@@ -1219,6 +1219,12 @@ bool HWAddressSanitizer::sanitizeFunction(Function &F) {
   for (auto Inst : ToInstrument)
     Changed |= instrumentMemAccess(Inst);
 
+  if (ClInstrumentMemIntrinsics && !IntrinToInstrument.empty()) {
+    for (auto Inst : IntrinToInstrument)
+      instrumentMemIntrinsic(cast<MemIntrinsic>(Inst));
+    Changed = true;
+  }
+
   LocalDynamicShadow = nullptr;
   StackBaseTag = nullptr;