1 files changed, 52 insertions, 8 deletions

diff --git a/llvm/tools/llvm-mca/llvm-mca.cpp b/llvm/tools/llvm-mca/llvm-mca.cpp
index 6f7b74f..2a27fea 100644
--- a/llvm/tools/llvm-mca/llvm-mca.cpp
+++ b/llvm/tools/llvm-mca/llvm-mca.cpp
@@ -231,6 +231,12 @@ static cl::opt<bool> DisableCustomBehaviour(
         "Disable custom behaviour (use the default class which does nothing)."),
     cl::cat(ViewOptions), cl::init(false));
 
+static cl::opt<bool> DisableInstrumentManager(
+    "disable-im",
+    cl::desc("Disable instrumentation manager (use the default class which "
+             "ignores instruments.)."),
+    cl::cat(ViewOptions), cl::init(false));
+
 namespace {
 
 const Target *getTarget(const char *ProgName) {
@@ -407,7 +413,7 @@ int main(int argc, char **argv) {
 
   // Need to initialize an MCInstPrinter as it is
   // required for initializing the MCTargetStreamer
-  // which needs to happen within the CRG.parseCodeRegions() call below.
+  // which needs to happen within the CRG.parseAnalysisRegions() call below.
   // Without an MCTargetStreamer, certain assembly directives can trigger a
   // segfault. (For example, the .cv_fpo_proc directive on x86 will segfault if
   // we don't initialize the MCTargetStreamer.)
@@ -424,9 +430,10 @@ int main(int argc, char **argv) {
   }
 
   // Parse the input and create CodeRegions that llvm-mca can analyze.
-  mca::AsmCodeRegionGenerator CRG(*TheTarget, SrcMgr, Ctx, *MAI, *STI, *MCII);
-  Expected<const mca::CodeRegions &> RegionsOrErr =
-      CRG.parseCodeRegions(std::move(IPtemp));
+  mca::AsmAnalysisRegionGenerator CRG(*TheTarget, SrcMgr, Ctx, *MAI, *STI,
+                                      *MCII);
+  Expected<const mca::AnalysisRegions &> RegionsOrErr =
+      CRG.parseAnalysisRegions(std::move(IPtemp));
   if (!RegionsOrErr) {
     if (auto Err =
             handleErrors(RegionsOrErr.takeError(), [](const StringError &E) {
@@ -437,7 +444,7 @@ int main(int argc, char **argv) {
     }
     return 1;
   }
-  const mca::CodeRegions &Regions = *RegionsOrErr;
+  const mca::AnalysisRegions &Regions = *RegionsOrErr;
 
   // Early exit if errors were found by the code region parsing logic.
   if (!Regions.isValid())
@@ -448,6 +455,39 @@ int main(int argc, char **argv) {
     return 1;
   }
 
+  std::unique_ptr<mca::InstrumentManager> IM;
+  if (!DisableInstrumentManager) {
+    IM = std::unique_ptr<mca::InstrumentManager>(
+        TheTarget->createInstrumentManager(*STI, *MCII));
+  }
+  if (!IM) {
+    // If the target doesn't have its own IM implemented (or the -disable-cb
+    // flag is set) then we use the base class (which does nothing).
+    IM = std::make_unique<mca::InstrumentManager>(*STI, *MCII);
+  }
+
+  // Parse the input and create InstrumentRegion that llvm-mca
+  // can use to improve analysis.
+  mca::AsmInstrumentRegionGenerator IRG(*TheTarget, SrcMgr, Ctx, *MAI, *STI,
+                                        *MCII, *IM);
+  Expected<const mca::InstrumentRegions &> InstrumentRegionsOrErr =
+      IRG.parseInstrumentRegions(std::move(IPtemp));
+  if (!InstrumentRegionsOrErr) {
+    if (auto Err = handleErrors(InstrumentRegionsOrErr.takeError(),
+                                [](const StringError &E) {
+                                  WithColor::error() << E.getMessage() << '\n';
+                                })) {
+      // Default case.
+      WithColor::error() << toString(std::move(Err)) << '\n';
+    }
+    return 1;
+  }
+  const mca::InstrumentRegions &InstrumentRegions = *InstrumentRegionsOrErr;
+
+  // Early exit if errors were found by the instrumentation parsing logic.
+  if (!InstrumentRegions.isValid())
+    return 1;
+
   // Now initialize the output file.
   auto OF = getOutputStream();
   if (std::error_code EC = OF.getError()) {
@@ -491,7 +531,7 @@ int main(int argc, char **argv) {
   }
 
   // Create an instruction builder.
-  mca::InstrBuilder IB(*STI, *MCII, *MRI, MCIA.get());
+  mca::InstrBuilder IB(*STI, *MCII, *MRI, MCIA.get(), *IM);
 
   // Create a context to control ownership of the pipeline hardware.
   mca::Context MCA(*MRI, *STI);
@@ -512,7 +552,7 @@ int main(int argc, char **argv) {
   assert(MAB && "Unable to create asm backend!");
 
   json::Object JSONOutput;
-  for (const std::unique_ptr<mca::CodeRegion> &Region : Regions) {
+  for (const std::unique_ptr<mca::AnalysisRegion> &Region : Regions) {
     // Skip empty code regions.
     if (Region->empty())
       continue;
@@ -527,8 +567,12 @@ int main(int argc, char **argv) {
 
     SmallVector<std::unique_ptr<mca::Instruction>> LoweredSequence;
     for (const MCInst &MCI : Insts) {
+      SMLoc Loc = MCI.getLoc();
+      const SmallVector<mca::SharedInstrument> Instruments =
+          InstrumentRegions.getActiveInstruments(Loc);
+
       Expected<std::unique_ptr<mca::Instruction>> Inst =
-          IB.createInstruction(MCI);
+          IB.createInstruction(MCI, Instruments);
       if (!Inst) {
         if (auto NewE = handleErrors(
                 Inst.takeError(),