[Instrumentation] Share InstrumentationIRBuilder between TSan and SanCov

[melver]

Summary

Factor our InstrumentationIRBuilder and share it between ThreadSanitizer
and SanitizerCoverage. Simplify its usage at the same time (use function
of passed Instruction or BasicBlock).

This class may be used in other instrumentation passes in future.

NFCI.

[nickdesaulniers]

llvm/include/llvm/Transforms/Instrumentation.h

18–24

If we're going to update the includes, please consider also including Function.h and Instruction.h.

203–212

Do we construct any IRBuilders with more than 1 arg?

Do we still need explicit if there's more than 1 arg?

[melver]

Add includes.

llvm/include/llvm/Transforms/Instrumentation.h

203–212

Yeah, I wanted to try this on AddressSanitizer.cpp, where there are some. It's not used right now, but for completeness I felt I should add it. If you think we should only add the support when it's actually used, I can remove.

WDYT?

Yes, the explicit should stay for vararg constructors where there can only be 1 arg.

[nickdesaulniers]

llvm/include/llvm/Transforms/Instrumentation.h

203–212

Unless you're actively planning to do that port for asan soon (why not do it in this patch, too), I feel like YAGNI, but I also don't care enough either way, hence the approval. I bring it up only because I was wondering what was being used from <utility> (which is std::forward).

[melver]

Remove unneeded constructors (and <utility>) for now. Can add back later when
we need it.

llvm/include/llvm/Transforms/Instrumentation.h

203–212

Fair point.

I tried ASan just now, but requires a bit more care. I will try to get to it eventually. (FYI, in the kernel we have LTO depends on !KASAN, which is why nobody noticed - but if I remove that restriction, things break as expected.)

llvm/include/llvm/Transforms/Instrumentation.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_INSTRUMENTATION_H
 #define LLVM_TRANSFORMS_INSTRUMENTATION_H
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
 #include <cassert>
 #include <cstdint>
 #include <limits>
 #include <string>
 #include <vector>
 
 namespace llvm {
 
 class Triple;
 class FunctionPass;
 class ModulePass;

 ///
 /// Scale a 64-bit weight down to 32-bits using \c Scale.
 ///
 static inline uint32_t scaleBranchCount(uint64_t Count, uint64_t Scale) {
   uint64_t Scaled = Count / Scale;
   assert(Scaled <= std::numeric_limits<uint32_t>::max() && "overflow 32-bits");
   return Scaled;
 }
+
+// Use to ensure the inserted instrumentation has a DebugLocation; if none is
+// attached to the source instruction, try to use a DILocation with offset 0
+// scoped to surrounding function (if it has a DebugLocation).
+//
+// Some non-call instructions may be missing debug info, but when inserting
+// instrumentation calls, some builds (e.g. LTO) want calls to have debug info
+// if the enclosing function does.
+struct InstrumentationIRBuilder : IRBuilder<> {
+  static void ensureDebugInfo(IRBuilder<> &IRB, const Function &F) {
+    if (IRB.getCurrentDebugLocation())
+      return;
+    if (DISubprogram *SP = F.getSubprogram())
+      IRB.SetCurrentDebugLocation(DILocation::get(SP->getContext(), 0, 0, SP));
+  }
+
+  explicit InstrumentationIRBuilder(Instruction *IP) : IRBuilder<>(IP) {
+    ensureDebugInfo(*this, *IP->getFunction());
+  }
+};
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_INSTRUMENTATION_H

llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp

void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
   DebugLoc EntryLoc;
   if (IsEntryBB) {
     if (auto SP = F.getSubprogram())
       EntryLoc = DILocation::get(SP->getContext(), SP->getScopeLine(), 0, SP);
     // Keep static allocas and llvm.localescape calls in the entry block.  Even
     // if we aren't splitting the block, it's nice for allocas to be before
     // calls.
     IP = PrepareToSplitEntryBlock(BB, IP);
-  } else {
-    EntryLoc = IP->getDebugLoc();
-    if (!EntryLoc)
-      if (auto *SP = F.getSubprogram())
-        EntryLoc = DILocation::get(SP->getContext(), 0, 0, SP);
   }
 
-  IRBuilder<> IRB(&*IP);
+  InstrumentationIRBuilder IRB(&*IP);
+  if (EntryLoc)
     IRB.SetCurrentDebugLocation(EntryLoc);
   if (Options.TracePC) {
     IRB.CreateCall(SanCovTracePC)
         ->setCannotMerge(); // gets the PC using GET_CALLER_PC.
   }
   if (Options.TracePCGuard) {
     auto GuardPtr = IRB.CreateIntToPtr(
         IRB.CreateAdd(IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
                       ConstantInt::get(IntptrTy, Idx * 4)),

llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp

 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/EscapeEnumerator.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "tsan"
struct InstructionInfo {
 
     explicit InstructionInfo(Instruction *Inst) : Inst(Inst) {}
 
     Instruction *Inst;
     unsigned Flags = 0;
   };
 
   void initialize(Module &M);
-  bool instrumentLoadOrStore(const InstructionInfo &II, const DataLayout &DL,
-                             const Function &F);
-  bool instrumentAtomic(Instruction *I, const DataLayout &DL,
-                        const Function &F);
+  bool instrumentLoadOrStore(const InstructionInfo &II, const DataLayout &DL);
+  bool instrumentAtomic(Instruction *I, const DataLayout &DL);
   bool instrumentMemIntrinsic(Instruction *I);
   void chooseInstructionsToInstrument(SmallVectorImpl<Instruction *> &Local,
                                       SmallVectorImpl<InstructionInfo> &All,
                                       const DataLayout &DL);
   bool addrPointsToConstantData(Value *Addr);
   int getMemoryAccessFuncIndex(Type *OrigTy, Value *Addr, const DataLayout &DL);
   void InsertRuntimeIgnores(Function &F);
 

 void insertModuleCtor(Module &M) {
   getOrCreateSanitizerCtorAndInitFunctions(
       M, kTsanModuleCtorName, kTsanInitName, /*InitArgTypes=*/{},
       /*InitArgs=*/{},
       // This callback is invoked when the functions are created the first
       // time. Hook them into the global ctors list in that case:
       [&](Function *Ctor, FunctionCallee) { appendToGlobalCtors(M, Ctor, 0); });
 }
-
-// Use to ensure the inserted instrumentation has a DebugLocation; if none is
-// attached to the source instruction, try to use a DILocation with offset 0
-// scoped to surrounding function (if it has a DebugLocation).
-//
-// Some non-call instructions may be missing debug info, but when inserting
-// instrumentation calls, some builds (e.g. LTO) want calls to have debug info
-// if the enclosing function does.
-struct InstrumentationIRBuilder : IRBuilder<> {
-  static void ensureDebugInfo(IRBuilder<> &IRB, const Function &F) {
-    if (IRB.getCurrentDebugLocation())
-      return;
-    if (DISubprogram *SP = F.getSubprogram())
-      IRB.SetCurrentDebugLocation(DILocation::get(SP->getContext(), 0, 0, SP));
-  }
-
-  explicit InstrumentationIRBuilder(Instruction *I, const Function &F)
-      : IRBuilder<>(I) {
-    ensureDebugInfo(*this, F);
-  }
-};
-
 }  // namespace
 
 PreservedAnalyses ThreadSanitizerPass::run(Function &F,
                                            FunctionAnalysisManager &FAM) {
   ThreadSanitizer TSan;
   if (TSan.sanitizeFunction(F, FAM.getResult<TargetLibraryAnalysis>(F)))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
static bool isTsanAtomic(const Instruction *I) {
   if (!SSID.hasValue())
     return false;
   if (isa<LoadInst>(I) || isa<StoreInst>(I))
     return SSID.getValue() != SyncScope::SingleThread;
   return true;
 }
 
 void ThreadSanitizer::InsertRuntimeIgnores(Function &F) {
-  InstrumentationIRBuilder IRB(F.getEntryBlock().getFirstNonPHI(), F);
+  InstrumentationIRBuilder IRB(F.getEntryBlock().getFirstNonPHI());
   IRB.CreateCall(TsanIgnoreBegin);
   EscapeEnumerator EE(F, "tsan_ignore_cleanup", ClHandleCxxExceptions);
   while (IRBuilder<> *AtExit = EE.Next()) {
     InstrumentationIRBuilder::ensureDebugInfo(*AtExit, F);
     AtExit->CreateCall(TsanIgnoreEnd);
   }
 }
 
bool ThreadSanitizer::sanitizeFunction(Function &F,
 
   // 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).
 
   // Instrument memory accesses only if we want to report bugs in the function.
   if (ClInstrumentMemoryAccesses && SanitizeFunction)
     for (const auto &II : AllLoadsAndStores) {
-      Res |= instrumentLoadOrStore(II, DL, F);
+      Res |= instrumentLoadOrStore(II, DL);
     }
 
   // Instrument atomic memory accesses in any case (they can be used to
   // implement synchronization).
   if (ClInstrumentAtomics)
     for (auto Inst : AtomicAccesses) {
-      Res |= instrumentAtomic(Inst, DL, F);
+      Res |= instrumentAtomic(Inst, DL);
     }
 
   if (ClInstrumentMemIntrinsics && SanitizeFunction)
     for (auto Inst : MemIntrinCalls) {
       Res |= instrumentMemIntrinsic(Inst);
     }
 
   if (F.hasFnAttribute("sanitize_thread_no_checking_at_run_time")) {
     assert(!F.hasFnAttribute(Attribute::SanitizeThread));
     if (HasCalls)
       InsertRuntimeIgnores(F);
   }
 
   // Instrument function entry/exit points if there were instrumented accesses.
   if ((Res || HasCalls) && ClInstrumentFuncEntryExit) {
-    InstrumentationIRBuilder IRB(F.getEntryBlock().getFirstNonPHI(), F);
+    InstrumentationIRBuilder IRB(F.getEntryBlock().getFirstNonPHI());
     Value *ReturnAddress = IRB.CreateCall(
         Intrinsic::getDeclaration(F.getParent(), 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, {});
     }
     Res = true;
   }
   return Res;
 }
 
 bool ThreadSanitizer::instrumentLoadOrStore(const InstructionInfo &II,
-                                            const DataLayout &DL,
-                                            const Function &F) {
-  InstrumentationIRBuilder IRB(II.Inst, F);
+                                            const DataLayout &DL) {
+  InstrumentationIRBuilder IRB(II.Inst);
   const bool IsWrite = isa<StoreInst>(*II.Inst);
   Value *Addr = IsWrite ? cast<StoreInst>(II.Inst)->getPointerOperand()
                         : cast<LoadInst>(II.Inst)->getPointerOperand();
   Type *OrigTy = getLoadStoreType(II.Inst);
 
   // swifterror memory addresses are mem2reg promoted by instruction selection.
   // As such they cannot have regular uses like an instrumentation function and
   // it makes no sense to track them as memory.

 // Both llvm and ThreadSanitizer atomic operations are based on C++11/C1x
 // standards.  For background see C++11 standard.  A slightly older, publicly
 // available draft of the standard (not entirely up-to-date, but close enough
 // for casual browsing) is available here:
 // http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2011/n3242.pdf
 // The following page contains more background information:
 // http://www.hpl.hp.com/personal/Hans_Boehm/c++mm/
 
-bool ThreadSanitizer::instrumentAtomic(Instruction *I, const DataLayout &DL,
-                                       const Function &F) {
-  InstrumentationIRBuilder IRB(I, F);
+bool ThreadSanitizer::instrumentAtomic(Instruction *I, const DataLayout &DL) {
+  InstrumentationIRBuilder IRB(I);
   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
     Value *Addr = LI->getPointerOperand();
     Type *OrigTy = LI->getType();
     int Idx = getMemoryAccessFuncIndex(OrigTy, Addr, DL);
     if (Idx < 0)
       return false;
     const unsigned ByteSize = 1U << Idx;
     const unsigned BitSize = ByteSize * 8;