[llvm] [Instrumentor] Add Instrumentor pass (PR #138958)

[Matt Arsenault via llvm-commits]

================
@@ -0,0 +1,774 @@
+//===-- Instrumentor.cpp - Highly configurable instrumentation pass -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/IPO/Instrumentor.h"
+#include "llvm/Transforms/IPO/InstrumentorConfigFile.h"
+
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/iterator.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/Verifier.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/Regex.h"
+
+#include <cassert>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <memory>
+#include <string>
+#include <system_error>
+#include <type_traits>
+
+using namespace llvm;
+using namespace llvm::instrumentor;
+
+#define DEBUG_TYPE "instrumentor"
+
+namespace {
+
+/// The user option to specify an output JSON file to write the configuration.
+static cl::opt<std::string> WriteConfigFile(
+    "instrumentor-write-config-file",
+    cl::desc(
+        "Write the instrumentor configuration into the specified JSON file"),
+    cl::init(""));
+
+/// The user option to specify an input JSON file to read the configuration.
+static cl::opt<std::string> ReadConfigFile(
+    "instrumentor-read-config-file",
+    cl::desc(
+        "Read the instrumentor configuration from the specified JSON file"),
+    cl::init(""));
+
+/// Set the debug location, if not set, after changing the insertion point of
+/// the IR builder \p IRB.
+template <typename IRBuilderTy> void ensureDbgLoc(IRBuilderTy &IRB) {
+  if (IRB.getCurrentDebugLocation())
+    return;
+  auto *BB = IRB.GetInsertBlock();
+  if (auto *SP = BB->getParent()->getSubprogram())
+    IRB.SetCurrentDebugLocation(DILocation::get(BB->getContext(), 0, 0, SP));
+}
+
+/// Attempt to cast \p V to type \p Ty.
+template <typename IRBTy>
+Value *tryToCast(IRBTy &IRB, Value *V, Type *Ty, const DataLayout &DL,
+                 bool AllowTruncate = false) {
+  if (!V)
+    return Constant::getAllOnesValue(Ty);
+  Type *VTy = V->getType();
+  if (VTy == Ty)
+    return V;
+  if (VTy->isAggregateType())
+    return V;
+  TypeSize RequestedSize = DL.getTypeSizeInBits(Ty);
+  TypeSize ValueSize = DL.getTypeSizeInBits(VTy);
+  bool ShouldTruncate = RequestedSize < ValueSize;
+  if (ShouldTruncate && !AllowTruncate)
+    return V;
+  if (ShouldTruncate && AllowTruncate)
+    return tryToCast(IRB,
+                     IRB.CreateIntCast(V, IRB.getIntNTy(RequestedSize),
+                                       /*IsSigned=*/false),
+                     Ty, DL, AllowTruncate);
+  if (VTy->isPointerTy() && Ty->isPointerTy())
+    return IRB.CreatePointerBitCastOrAddrSpaceCast(V, Ty);
+  if (VTy->isIntegerTy() && Ty->isIntegerTy())
+    return IRB.CreateIntCast(V, Ty, /*IsSigned=*/false);
+  if (VTy->isFloatingPointTy() && Ty->isIntOrPtrTy()) {
+    return tryToCast(IRB, IRB.CreateBitCast(V, IRB.getIntNTy(ValueSize)), Ty,
+                     DL, AllowTruncate);
+  }
+  return IRB.CreateBitOrPointerCast(V, Ty);
+}
+
+/// Get a constant integer/boolean of type \p IT and value \p Val.
+template <typename Ty> Constant *getCI(Type *IT, Ty Val) {
+  return ConstantInt::get(IT, Val);
+}
+
+/// The core of the instrumentor pass, which instruments the module as the
+/// instrumentation configuration mandates.
+class InstrumentorImpl final {
+public:
+  /// Construct an instrumentor implementation using the configuration \p IConf.
+  InstrumentorImpl(InstrumentationConfig &IConf, InstrumentorIRBuilderTy &IIRB,
+                   Module &M)
+      : IConf(IConf), M(M), IIRB(IIRB) {
+    IConf.populate(IIRB);
+  }
+
+  /// Instrument the module, public entry point.
+  bool instrument();
+
+private:
+  /// Indicate if the module should be instrumented based on the target.
+  bool shouldInstrumentTarget();
+
+  /// Indicate if the function \p Fn should be instrumented.
+  bool shouldInstrumentFunction(Function &Fn);
+
+  /// Instrument instruction \p I if needed, and use the argument caches in \p
+  /// ICaches.
+  bool instrumentInstruction(Instruction &I, InstrumentationCaches &ICaches);
+
+  /// Instrument function \p Fn.
+  bool instrumentFunction(Function &Fn);
+
+  /// The instrumentation opportunities for instructions indexed by
+  /// their opcode.
+  DenseMap<unsigned, InstrumentationOpportunity *> InstChoicesPRE,
+      InstChoicesPOST;
+
+  /// The instrumentor configuration.
+  InstrumentationConfig &IConf;
+
+  /// The underlying module.
+  Module &M;
+
+protected:
+  /// A special IR builder that keeps track of the inserted instructions.
+  InstrumentorIRBuilderTy &IIRB;
+};
+
+} // end anonymous namespace
+
+bool InstrumentorImpl::shouldInstrumentTarget() {
+  const Triple &T = M.getTargetTriple();
+  const bool IsGPU = T.isAMDGPU() || T.isNVPTX();
+
+  bool RegexMatches = true;
+  const auto TargetRegexStr = IConf.TargetRegex->getString();
+  if (!TargetRegexStr.empty()) {
+    llvm::Regex TargetRegex(TargetRegexStr);
+    std::string ErrMsg;
+    if (!TargetRegex.isValid(ErrMsg)) {
+      IIRB.Ctx.diagnose(DiagnosticInfoInstrumentation(
+          Twine("failed to parse target regex: ") + ErrMsg, DS_Warning));
+      return false;
+    }
+    RegexMatches = TargetRegex.match(T.str());
+  }
+
+  // Only instrument the module if the target has to be instrumented.
+  return ((IsGPU && IConf.GPUEnabled->getBool()) ||
+          (!IsGPU && IConf.HostEnabled->getBool())) &&
+         RegexMatches;
+}
+
+bool InstrumentorImpl::shouldInstrumentFunction(Function &Fn) {
+  if (Fn.isDeclaration())
+    return false;
+  return !Fn.getName().starts_with(IConf.getRTName()) ||
+         Fn.hasFnAttribute("instrument");
+}
+
+bool InstrumentorImpl::instrumentInstruction(Instruction &I,
+                                             InstrumentationCaches &ICaches) {
+  bool Changed = false;
+
+  // Skip instrumentation instructions.
+  if (IIRB.NewInsts.contains(&I))
+    return Changed;
+
+  // Count epochs eagerly.
+  ++IIRB.Epoch;
+
+  Value *IPtr = &I;
+  if (auto *IO = InstChoicesPRE.lookup(I.getOpcode())) {
+    IIRB.IRB.SetInsertPoint(&I);
+    ensureDbgLoc(IIRB.IRB);
+    Changed |= bool(IO->instrument(IPtr, IConf, IIRB, ICaches));
+  }
+
+  if (auto *IO = InstChoicesPOST.lookup(I.getOpcode())) {
+    IIRB.IRB.SetInsertPoint(I.getNextNode());
+    ensureDbgLoc(IIRB.IRB);
+    Changed |= bool(IO->instrument(IPtr, IConf, IIRB, ICaches));
+  }
+  IIRB.returnAllocas();
+
+  return Changed;
+}
+
+bool InstrumentorImpl::instrumentFunction(Function &Fn) {
+  bool Changed = false;
+  if (!shouldInstrumentFunction(Fn))
+    return Changed;
+
+  InstrumentationCaches ICaches;
+  ReversePostOrderTraversal<Function *> RPOT(&Fn);
+  for (auto &It : RPOT)
+    for (auto &I : *It)
+      Changed |= instrumentInstruction(I, ICaches);
+
+  return Changed;
+}
+
+bool InstrumentorImpl::instrument() {
+  bool Changed = false;
+  if (!shouldInstrumentTarget())
+    return Changed;
+
+  for (auto &It : IConf.IChoices[InstrumentationLocation::INSTRUCTION_PRE])
+    if (It.second->Enabled)
+      InstChoicesPRE[It.second->getOpcode()] = It.second;
+  for (auto &It : IConf.IChoices[InstrumentationLocation::INSTRUCTION_POST])
+    if (It.second->Enabled)
+      InstChoicesPOST[It.second->getOpcode()] = It.second;
+
+  for (Function &Fn : M)
+    Changed |= instrumentFunction(Fn);
+
+  return Changed;
+}
+
+PreservedAnalyses InstrumentorPass::run(Module &M, InstrumentationConfig &IConf,
+                                        InstrumentorIRBuilderTy &IIRB,
+                                        bool ReadConfig) {
+  InstrumentorImpl Impl(IConf, IIRB, M);
+  if (ReadConfig && !readConfigFromJSON(IConf, ReadConfigFile, IIRB.Ctx))
+    return PreservedAnalyses::all();
+
+  writeConfigToJSON(IConf, WriteConfigFile, IIRB.Ctx);
+
+  bool Changed = Impl.instrument();
+  if (!Changed)
+    return PreservedAnalyses::all();
+  return PreservedAnalyses::none();
+}
+
+PreservedAnalyses InstrumentorPass::run(Module &M, ModuleAnalysisManager &MAM) {
+  // Only create them if the user did not provide them.
+  std::unique_ptr<InstrumentationConfig> IConfInt(
+      !UserIConf ? new InstrumentationConfig() : nullptr);
+  std::unique_ptr<InstrumentorIRBuilderTy> IIRBInt(
+      !UserIIRB ? new InstrumentorIRBuilderTy(M) : nullptr);
+
+  auto *IConf = IConfInt ? IConfInt.get() : UserIConf;
+  auto *IIRB = IIRBInt ? IIRBInt.get() : UserIIRB;
+
+  auto PA = run(M, *IConf, *IIRB, !UserIConf);
+
+  assert(!verifyModule(M, &errs()));
+  return PA;
+}
+
+BaseConfigurationOption *
+BaseConfigurationOption::getBoolOption(InstrumentationConfig &IConf,
+                                       StringRef Name, StringRef Description,
+                                       bool DefaultValue) {
+  auto *BCO = new BaseConfigurationOption();
+  BCO->Name = Name;
+  BCO->Description = Description;
+  BCO->Kind = BOOLEAN;
+  BCO->Value.Bool = DefaultValue;
----------------
arsenm wrote:

Move to constructor arguments?

https://github.com/llvm/llvm-project/pull/138958