[RISCV][llvm-mca] Use LMUL Instruments to provide more accurate reports on RISCV

On x86 and AArch, SIMD instructions encode all of the scheduling information in the instruction
itself. For example, VADD.I16 q0, q1, q2 is a neon instruction that operates on 16-bit integer
elements stored in 128-bit Q registers, which leads to eight 16-bit lanes in parallel. This kind
of information impacts how the instruction takes to execute and what dependencies this may cause.

On RISCV however, the data that impacts scheduling is encoded in CSR registers such as vtype or
vl, in addition with the instruction itself. But MCA does not track or use the data in these
registers. This patch fixes this problem by introducing Instruments into MCA.

* Replace `CodeRegions` with `AnalysisRegions`
* Add `Instrument` and `InstrumentManager`
* Add `InstrumentRegions`
* Add RISCV Instrument and `InstrumentManager`
* Parse `Instruments` in driver
* Use instruments to override schedule class
* RISCV use lmul instrument to override schedule class
* Fix unit tests to pass empty instruments
* Add -ignore-im clopt to disable this change

A prior version of this patch was commited in 5e82ee537321. 2323a4ee610f reverted
that change because the unit test files caused build errors. The change with fixes
were committed in b88b8307bf9e but reverted once again e8e92c8313a0 due to more
build errors.

This commit adds the prior changes and fixes the build error.

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

5 files changed, 444 insertions, 95 deletions

diff --git a/llvm/tools/llvm-mca/CodeRegion.cpp b/llvm/tools/llvm-mca/CodeRegion.cpp
index 7662538..c91ed75 100644
--- a/llvm/tools/llvm-mca/CodeRegion.cpp
+++ b/llvm/tools/llvm-mca/CodeRegion.cpp
@@ -16,11 +16,6 @@
 namespace llvm {
 namespace mca {
 
-CodeRegions::CodeRegions(llvm::SourceMgr &S) : SM(S), FoundErrors(false) {
-  // Create a default region for the input code sequence.
-  Regions.emplace_back(std::make_unique<CodeRegion>("", SMLoc()));
-}
-
 bool CodeRegion::isLocInRange(SMLoc Loc) const {
   if (RangeEnd.isValid() && Loc.getPointer() > RangeEnd.getPointer())
     return false;
@@ -29,7 +24,19 @@ bool CodeRegion::isLocInRange(SMLoc Loc) const {
   return true;
 }
 
-void CodeRegions::beginRegion(StringRef Description, SMLoc Loc) {
+void CodeRegions::addInstruction(const MCInst &Instruction) {
+  SMLoc Loc = Instruction.getLoc();
+  for (UniqueCodeRegion &Region : Regions)
+    if (Region->isLocInRange(Loc))
+      Region->addInstruction(Instruction);
+}
+
+AnalysisRegions::AnalysisRegions(llvm::SourceMgr &S) : CodeRegions(S) {
+  // Create a default region for the input code sequence.
+  Regions.emplace_back(std::make_unique<CodeRegion>("", SMLoc()));
+}
+
+void AnalysisRegions::beginRegion(StringRef Description, SMLoc Loc) {
   if (ActiveRegions.empty()) {
     // Remove the default region if there is at least one user defined region.
     // By construction, only the default region has an invalid start location.
@@ -44,17 +51,17 @@ void CodeRegions::beginRegion(StringRef Description, SMLoc Loc) {
     if (It != ActiveRegions.end()) {
       const CodeRegion &R = *Regions[It->second];
       if (Description.empty()) {
-        SM.PrintMessage(Loc, SourceMgr::DK_Error,
+        SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
                         "found multiple overlapping anonymous regions");
-        SM.PrintMessage(R.startLoc(), SourceMgr::DK_Note,
+        SM.PrintMessage(R.startLoc(), llvm::SourceMgr::DK_Note,
                         "Previous anonymous region was defined here");
         FoundErrors = true;
         return;
       }
 
-      SM.PrintMessage(Loc, SourceMgr::DK_Error,
+      SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
                       "overlapping regions cannot have the same name");
-      SM.PrintMessage(R.startLoc(), SourceMgr::DK_Note,
+      SM.PrintMessage(R.startLoc(), llvm::SourceMgr::DK_Note,
                       "region " + Description + " was previously defined here");
       FoundErrors = true;
       return;
@@ -65,7 +72,7 @@ void CodeRegions::beginRegion(StringRef Description, SMLoc Loc) {
   Regions.emplace_back(std::make_unique<CodeRegion>(Description, Loc));
 }
 
-void CodeRegions::endRegion(StringRef Description, SMLoc Loc) {
+void AnalysisRegions::endRegion(StringRef Description, SMLoc Loc) {
   if (Description.empty()) {
     // Special case where there is only one user defined region,
     // and this LLVM-MCA-END directive doesn't provide a region name.
@@ -94,22 +101,73 @@ void CodeRegions::endRegion(StringRef Description, SMLoc Loc) {
   }
 
   FoundErrors = true;
-  SM.PrintMessage(Loc, SourceMgr::DK_Error,
+  SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
                   "found an invalid region end directive");
   if (!Description.empty()) {
-    SM.PrintMessage(Loc, SourceMgr::DK_Note,
+    SM.PrintMessage(Loc, llvm::SourceMgr::DK_Note,
                     "unable to find an active region named " + Description);
   } else {
-    SM.PrintMessage(Loc, SourceMgr::DK_Note,
+    SM.PrintMessage(Loc, llvm::SourceMgr::DK_Note,
                     "unable to find an active anonymous region");
   }
 }
 
-void CodeRegions::addInstruction(const MCInst &Instruction) {
-  SMLoc Loc = Instruction.getLoc();
-  for (UniqueCodeRegion &Region : Regions)
-    if (Region->isLocInRange(Loc))
-      Region->addInstruction(Instruction);
+InstrumentRegions::InstrumentRegions(llvm::SourceMgr &S) : CodeRegions(S) {}
+
+void InstrumentRegions::beginRegion(StringRef Description, SMLoc Loc,
+                                    SharedInstrument I) {
+  if (Description.empty()) {
+    SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
+                    "anonymous instrumentation regions are not permitted");
+    FoundErrors = true;
+    return;
+  }
+
+  auto It = ActiveRegions.find(Description);
+  if (It != ActiveRegions.end()) {
+    const CodeRegion &R = *Regions[It->second];
+    SM.PrintMessage(
+        Loc, llvm::SourceMgr::DK_Error,
+        "overlapping instrumentation regions cannot be of the same kind");
+    SM.PrintMessage(R.startLoc(), llvm::SourceMgr::DK_Note,
+                    "instrumentation region " + Description +
+                        " was previously defined here");
+    FoundErrors = true;
+    return;
+  }
+
+  ActiveRegions[Description] = Regions.size();
+  Regions.emplace_back(std::make_unique<InstrumentRegion>(Description, Loc, I));
+}
+
+void InstrumentRegions::endRegion(StringRef Description, SMLoc Loc) {
+  auto It = ActiveRegions.find(Description);
+  if (It != ActiveRegions.end()) {
+    Regions[It->second]->setEndLocation(Loc);
+    ActiveRegions.erase(It);
+    return;
+  }
+
+  FoundErrors = true;
+  SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
+                  "found an invalid instrumentation region end directive");
+  if (!Description.empty()) {
+    SM.PrintMessage(Loc, llvm::SourceMgr::DK_Note,
+                    "unable to find an active instrumentation region named " +
+                        Description);
+  }
+}
+
+const SmallVector<SharedInstrument>
+InstrumentRegions::getActiveInstruments(SMLoc Loc) const {
+  SmallVector<SharedInstrument> AI;
+  for (auto &R : Regions) {
+    if (R->isLocInRange(Loc)) {
+      InstrumentRegion *IR = static_cast<InstrumentRegion *>(R.get());
+      AI.emplace_back(IR->getInstrument());
+    }
+  }
+  return AI;
 }
 
 } // namespace mca
diff --git a/llvm/tools/llvm-mca/CodeRegion.h b/llvm/tools/llvm-mca/CodeRegion.h
index 0e1e02a..b5b2f3a 100644
--- a/llvm/tools/llvm-mca/CodeRegion.h
+++ b/llvm/tools/llvm-mca/CodeRegion.h
@@ -7,7 +7,8 @@
 //===----------------------------------------------------------------------===//
 /// \file
 ///
-/// This file implements class CodeRegion and CodeRegions.
+/// This file implements class CodeRegion and CodeRegions, InstrumentRegion,
+/// AnalysisRegions, and InstrumentRegions.
 ///
 /// A CodeRegion describes a region of assembly code guarded by special LLVM-MCA
 /// comment directives.
@@ -25,8 +26,32 @@
 /// description; internally, regions are described by a range of source
 /// locations (SMLoc objects).
 ///
-/// An instruction (a MCInst) is added to a region R only if its location is in
-/// range [R.RangeStart, R.RangeEnd].
+/// An instruction (a MCInst) is added to a CodeRegion R only if its
+/// location is in range [R.RangeStart, R.RangeEnd].
+///
+/// A InstrumentRegion describes a region of assembly code guarded by
+/// special LLVM-MCA comment directives.
+///
+///   # LLVM-MCA-<INSTRUMENTATION_TYPE> <data>
+///     ...  ## asm
+///
+/// where INSTRUMENTATION_TYPE is a type defined in llvm and expects to use
+/// data.
+///
+/// A comment starting with substring LLVM-MCA-<INSTRUMENTATION_TYPE>
+/// brings data into scope for llvm-mca to use in its analysis for
+/// all following instructions.
+///
+/// If the same INSTRUMENTATION_TYPE is found later in the instruction list,
+/// then the original InstrumentRegion will be automatically ended,
+/// and a new InstrumentRegion will begin.
+///
+/// If there are comments containing the different INSTRUMENTATION_TYPEs,
+/// then both data sets remain available. In contrast with a CodeRegion,
+/// an InstrumentRegion does not need a comment to end the region.
+//
+// An instruction (a MCInst) is added to an InstrumentRegion R only
+// if its location is in range [R.RangeStart, R.RangeEnd].
 //
 //===----------------------------------------------------------------------===//
 
@@ -38,6 +63,7 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/MC/MCInst.h"
+#include "llvm/MCA/CustomBehaviour.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/SourceMgr.h"
@@ -81,9 +107,31 @@ public:
   llvm::StringRef getDescription() const { return Description; }
 };
 
+/// Alias AnalysisRegion with CodeRegion since CodeRegionGenerator
+/// is absract and AnalysisRegionGenerator operates on AnalysisRegions
+using AnalysisRegion = CodeRegion;
+
+/// A CodeRegion that contains instrumentation that can be used
+/// in analysis of the region.
+class InstrumentRegion : public CodeRegion {
+  /// Instrument for this region.
+  SharedInstrument Instrument;
+
+public:
+  InstrumentRegion(llvm::StringRef Desc, llvm::SMLoc Start, SharedInstrument I)
+      : CodeRegion(Desc, Start), Instrument(I) {}
+
+public:
+  SharedInstrument getInstrument() const { return Instrument; }
+};
+
 class CodeRegionParseError final : public Error {};
 
 class CodeRegions {
+  CodeRegions(const CodeRegions &) = delete;
+  CodeRegions &operator=(const CodeRegions &) = delete;
+
+protected:
   // A source manager. Used by the tool to generate meaningful warnings.
   llvm::SourceMgr &SM;
 
@@ -92,11 +140,8 @@ class CodeRegions {
   llvm::StringMap<unsigned> ActiveRegions;
   bool FoundErrors;
 
-  CodeRegions(const CodeRegions &) = delete;
-  CodeRegions &operator=(const CodeRegions &) = delete;
-
 public:
-  CodeRegions(llvm::SourceMgr &S);
+  CodeRegions(llvm::SourceMgr &S) : SM(S), FoundErrors(false) {}
 
   typedef std::vector<UniqueCodeRegion>::iterator iterator;
   typedef std::vector<UniqueCodeRegion>::const_iterator const_iterator;
@@ -106,8 +151,6 @@ public:
   const_iterator begin() const { return Regions.cbegin(); }
   const_iterator end() const { return Regions.cend(); }
 
-  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc);
-  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc);
   void addInstruction(const llvm::MCInst &Instruction);
   llvm::SourceMgr &getSourceMgr() const { return SM; }
 
@@ -122,6 +165,28 @@ public:
   }
 
   bool isValid() const { return !FoundErrors; }
+
+  bool isRegionActive(llvm::StringRef Description) const {
+    return ActiveRegions.find(Description) != ActiveRegions.end();
+  }
+};
+
+struct AnalysisRegions : public CodeRegions {
+  AnalysisRegions(llvm::SourceMgr &S);
+
+  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc);
+  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc);
+};
+
+struct InstrumentRegions : public CodeRegions {
+  InstrumentRegions(llvm::SourceMgr &S);
+
+  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc,
+                   SharedInstrument Instrument);
+  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc);
+
+  const SmallVector<SharedInstrument>
+  getActiveInstruments(llvm::SMLoc Loc) const;
 };
 
 } // namespace mca
diff --git a/llvm/tools/llvm-mca/CodeRegionGenerator.cpp b/llvm/tools/llvm-mca/CodeRegionGenerator.cpp
index cdecfba..d643234 100644
--- a/llvm/tools/llvm-mca/CodeRegionGenerator.cpp
+++ b/llvm/tools/llvm-mca/CodeRegionGenerator.cpp
@@ -16,7 +16,6 @@
 #include "CodeRegionGenerator.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
-#include "llvm/MC/MCParser/MCAsmLexer.h"
 #include "llvm/MC/MCParser/MCTargetAsmParser.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCTargetOptions.h"
@@ -30,15 +29,6 @@ namespace mca {
 // This virtual dtor serves as the anchor for the CodeRegionGenerator class.
 CodeRegionGenerator::~CodeRegionGenerator() {}
 
-// A comment consumer that parses strings.  The only valid tokens are strings.
-class MCACommentConsumer : public AsmCommentConsumer {
-public:
-  CodeRegions &Regions;
-
-  MCACommentConsumer(CodeRegions &R) : Regions(R) {}
-  void HandleComment(SMLoc Loc, StringRef CommentText) override;
-};
-
 // This class provides the callbacks that occur when parsing input assembly.
 class MCStreamerWrapper final : public MCStreamer {
   CodeRegions &Regions;
@@ -73,7 +63,53 @@ public:
   }
 };
 
-void MCACommentConsumer::HandleComment(SMLoc Loc, StringRef CommentText) {
+Expected<const CodeRegions &> AsmCodeRegionGenerator::parseCodeRegions(
+    const std::unique_ptr<MCInstPrinter> &IP) {
+  MCTargetOptions Opts;
+  Opts.PreserveAsmComments = false;
+  CodeRegions &Regions = getRegions();
+  MCStreamerWrapper Str(Ctx, Regions);
+
+  // Need to initialize an MCTargetStreamer otherwise
+  // certain asm directives will cause a segfault.
+  // Using nulls() so that anything emitted by the MCTargetStreamer
+  // doesn't show up in the llvm-mca output.
+  raw_ostream &OSRef = nulls();
+  formatted_raw_ostream FOSRef(OSRef);
+  TheTarget.createAsmTargetStreamer(Str, FOSRef, IP.get(),
+                                    /*IsVerboseAsm=*/true);
+
+  // Create a MCAsmParser and setup the lexer to recognize llvm-mca ASM
+  // comments.
+  std::unique_ptr<MCAsmParser> Parser(
+      createMCAsmParser(Regions.getSourceMgr(), Ctx, Str, MAI));
+  MCAsmLexer &Lexer = Parser->getLexer();
+  MCACommentConsumer *CCP = getCommentConsumer();
+  Lexer.setCommentConsumer(CCP);
+  // Enable support for MASM literal numbers (example: 05h, 101b).
+  Lexer.setLexMasmIntegers(true);
+
+  std::unique_ptr<MCTargetAsmParser> TAP(
+      TheTarget.createMCAsmParser(STI, *Parser, MCII, Opts));
+  if (!TAP)
+    return make_error<StringError>(
+        "This target does not support assembly parsing.",
+        inconvertibleErrorCode());
+  Parser->setTargetParser(*TAP);
+  Parser->Run(false);
+
+  if (CCP->hadErr())
+    return make_error<StringError>("There was an error parsing comments.",
+                                   inconvertibleErrorCode());
+
+  // Set the assembler dialect from the input. llvm-mca will use this as the
+  // default dialect when printing reports.
+  AssemblerDialect = Parser->getAssemblerDialect();
+  return Regions;
+}
+
+void AnalysisRegionCommentConsumer::HandleComment(SMLoc Loc,
+                                                  StringRef CommentText) {
   // Skip empty comments.
   StringRef Comment(CommentText);
   if (Comment.empty())
@@ -107,44 +143,66 @@ void MCACommentConsumer::HandleComment(SMLoc Loc, StringRef CommentText) {
   Regions.beginRegion(Comment, Loc);
 }
 
-Expected<const CodeRegions &> AsmCodeRegionGenerator::parseCodeRegions(
-    const std::unique_ptr<MCInstPrinter> &IP) {
-  MCTargetOptions Opts;
-  Opts.PreserveAsmComments = false;
-  MCStreamerWrapper Str(Ctx, Regions);
+void InstrumentRegionCommentConsumer::HandleComment(SMLoc Loc,
+                                                    StringRef CommentText) {
+  // Skip empty comments.
+  StringRef Comment(CommentText);
+  if (Comment.empty())
+    return;
 
-  // Need to initialize an MCTargetStreamer otherwise
-  // certain asm directives will cause a segfault.
-  // Using nulls() so that anything emitted by the MCTargetStreamer
-  // doesn't show up in the llvm-mca output.
-  raw_ostream &OSRef = nulls();
-  formatted_raw_ostream FOSRef(OSRef);
-  TheTarget.createAsmTargetStreamer(Str, FOSRef, IP.get(),
-                                    /*IsVerboseAsm=*/true);
+  // Skip spaces and tabs.
+  unsigned Position = Comment.find_first_not_of(" \t");
+  if (Position >= Comment.size())
+    // We reached the end of the comment. Bail out.
+    return;
+  Comment = Comment.drop_front(Position);
 
-  // Create a MCAsmParser and setup the lexer to recognize llvm-mca ASM
-  // comments.
-  std::unique_ptr<MCAsmParser> Parser(
-      createMCAsmParser(Regions.getSourceMgr(), Ctx, Str, MAI));
-  MCAsmLexer &Lexer = Parser->getLexer();
-  MCACommentConsumer CC(Regions);
-  Lexer.setCommentConsumer(&CC);
-  // Enable support for MASM literal numbers (example: 05h, 101b).
-  Lexer.setLexMasmIntegers(true);
+  // Bail out if not an MCA style comment
+  if (!Comment.consume_front("LLVM-MCA-"))
+    return;
 
-  std::unique_ptr<MCTargetAsmParser> TAP(
-      TheTarget.createMCAsmParser(STI, *Parser, MCII, Opts));
-  if (!TAP)
-    return make_error<StringError>(
-        "This target does not support assembly parsing.",
-        inconvertibleErrorCode());
-  Parser->setTargetParser(*TAP);
-  Parser->Run(false);
+  // Skip AnalysisRegion comments
+  if (Comment.consume_front("BEGIN") || Comment.consume_front("END"))
+    return;
 
-  // Set the assembler dialect from the input. llvm-mca will use this as the
-  // default dialect when printing reports.
-  AssemblerDialect = Parser->getAssemblerDialect();
-  return Regions;
+  if (IM.shouldIgnoreInstruments())
+    return;
+
+  auto [InstrumentKind, Data] = Comment.split(" ");
+
+  // An error if not of the form LLVM-MCA-TARGET-KIND
+  if (!IM.supportsInstrumentType(InstrumentKind)) {
+    if (InstrumentKind.empty())
+      SM.PrintMessage(
+          Loc, llvm::SourceMgr::DK_Error,
+          "No instrumentation kind was provided in LLVM-MCA comment");
+    else
+      SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
+                      "Unknown instrumentation type in LLVM-MCA comment: " +
+                          InstrumentKind);
+    FoundError = true;
+    return;
+  }
+
+  SharedInstrument I = IM.createInstrument(InstrumentKind, Data);
+  if (!I) {
+    if (Data.empty())
+      SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
+                      "Failed to create " + InstrumentKind +
+                          " instrument with no data");
+    else
+      SM.PrintMessage(Loc, llvm::SourceMgr::DK_Error,
+                      "Failed to create " + InstrumentKind +
+                          " instrument with data: " + Data);
+    FoundError = true;
+    return;
+  }
+
+  // End InstrumentType region if one is open
+  if (Regions.isRegionActive(InstrumentKind))
+    Regions.endRegion(InstrumentKind, Loc);
+  // Start new instrumentation region
+  Regions.beginRegion(InstrumentKind, Loc, I);
 }
 
 } // namespace mca
diff --git a/llvm/tools/llvm-mca/CodeRegionGenerator.h b/llvm/tools/llvm-mca/CodeRegionGenerator.h
index ac02131..d9e9be2 100644
--- a/llvm/tools/llvm-mca/CodeRegionGenerator.h
+++ b/llvm/tools/llvm-mca/CodeRegionGenerator.h
@@ -19,8 +19,10 @@
 #include "CodeRegion.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCParser/MCAsmLexer.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/TargetRegistry.h"
+#include "llvm/MCA/CustomBehaviour.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/SourceMgr.h"
 #include <memory>
@@ -28,24 +30,96 @@
 namespace llvm {
 namespace mca {
 
-/// This class is responsible for parsing the input given to the llvm-mca
-/// driver, and converting that into a CodeRegions instance.
+class MCACommentConsumer : public AsmCommentConsumer {
+protected:
+  bool FoundError;
+
+public:
+  MCACommentConsumer() : FoundError(false) {}
+
+  bool hadErr() const { return FoundError; }
+};
+
+/// A comment consumer that parses strings.  The only valid tokens are strings.
+class AnalysisRegionCommentConsumer : public MCACommentConsumer {
+  AnalysisRegions &Regions;
+
+public:
+  AnalysisRegionCommentConsumer(AnalysisRegions &R) : Regions(R) {}
+
+  /// Parses a comment. It begins a new region if it is of the form
+  /// LLVM-MCA-BEGIN. It ends a region if it is of the form LLVM-MCA-END.
+  /// Regions can be optionally named if they are of the form
+  /// LLVM-MCA-BEGIN <name> or LLVM-MCA-END <name>. Subregions are
+  /// permitted, but a region that begins while another region is active
+  /// must be ended before the outer region is ended. If thre is only one
+  /// active region, LLVM-MCA-END does not need to provide a name.
+  void HandleComment(SMLoc Loc, StringRef CommentText) override;
+};
+
+/// A comment consumer that parses strings to create InstrumentRegions.
+/// The only valid tokens are strings.
+class InstrumentRegionCommentConsumer : public MCACommentConsumer {
+  llvm::SourceMgr &SM;
+
+  InstrumentRegions &Regions;
+
+  InstrumentManager &IM;
+
+public:
+  InstrumentRegionCommentConsumer(llvm::SourceMgr &SM, InstrumentRegions &R,
+                                  InstrumentManager &IM)
+      : SM(SM), Regions(R), IM(IM) {}
+
+  /// Parses a comment. It begins a new region if it is of the form
+  /// LLVM-MCA-<INSTRUMENTATION_TYPE> <data> where INSTRUMENTATION_TYPE
+  /// is a valid InstrumentKind. If there is already an active
+  /// region of type INSTRUMENATION_TYPE, then it will end the active
+  /// one and begin a new one using the new data.
+  void HandleComment(SMLoc Loc, StringRef CommentText) override;
+};
+
+/// This abstract class is responsible for parsing the input given to
+/// the llvm-mca driver, and converting that into a CodeRegions instance.
 class CodeRegionGenerator {
 protected:
-  CodeRegions Regions;
   CodeRegionGenerator(const CodeRegionGenerator &) = delete;
   CodeRegionGenerator &operator=(const CodeRegionGenerator &) = delete;
+  virtual Expected<const CodeRegions &>
+  parseCodeRegions(const std::unique_ptr<MCInstPrinter> &IP) = 0;
 
 public:
-  CodeRegionGenerator(llvm::SourceMgr &SM) : Regions(SM) {}
+  CodeRegionGenerator() {}
   virtual ~CodeRegionGenerator();
-  virtual Expected<const CodeRegions &>
-  parseCodeRegions(const std::unique_ptr<MCInstPrinter> &IP) = 0;
 };
 
-/// This class is responsible for parsing input ASM and generating
-/// a CodeRegions instance.
-class AsmCodeRegionGenerator final : public CodeRegionGenerator {
+/// Abastract CodeRegionGenerator with AnalysisRegions member
+class AnalysisRegionGenerator : public virtual CodeRegionGenerator {
+protected:
+  AnalysisRegions Regions;
+
+public:
+  AnalysisRegionGenerator(llvm::SourceMgr &SM) : Regions(SM) {}
+
+  virtual Expected<const AnalysisRegions &>
+  parseAnalysisRegions(const std::unique_ptr<MCInstPrinter> &IP) = 0;
+};
+
+/// Abstract CodeRegionGenerator with InstrumentRegionsRegions member
+class InstrumentRegionGenerator : public virtual CodeRegionGenerator {
+protected:
+  InstrumentRegions Regions;
+
+public:
+  InstrumentRegionGenerator(llvm::SourceMgr &SM) : Regions(SM) {}
+
+  virtual Expected<const InstrumentRegions &>
+  parseInstrumentRegions(const std::unique_ptr<MCInstPrinter> &IP) = 0;
+};
+
+/// This abstract class is responsible for parsing input ASM and
+/// generating a CodeRegions instance.
+class AsmCodeRegionGenerator : public virtual CodeRegionGenerator {
   const Target &TheTarget;
   MCContext &Ctx;
   const MCAsmInfo &MAI;
@@ -54,17 +128,67 @@ class AsmCodeRegionGenerator final : public CodeRegionGenerator {
   unsigned AssemblerDialect; // This is set during parsing.
 
 public:
-  AsmCodeRegionGenerator(const Target &T, llvm::SourceMgr &SM, MCContext &C,
-                         const MCAsmInfo &A, const MCSubtargetInfo &S,
-                         const MCInstrInfo &I)
-      : CodeRegionGenerator(SM), TheTarget(T), Ctx(C), MAI(A), STI(S), MCII(I),
-        AssemblerDialect(0) {}
+  AsmCodeRegionGenerator(const Target &T, MCContext &C, const MCAsmInfo &A,
+                         const MCSubtargetInfo &S, const MCInstrInfo &I)
+      : TheTarget(T), Ctx(C), MAI(A), STI(S), MCII(I), AssemblerDialect(0) {}
+
+  virtual MCACommentConsumer *getCommentConsumer() = 0;
+  virtual CodeRegions &getRegions() = 0;
 
   unsigned getAssemblerDialect() const { return AssemblerDialect; }
   Expected<const CodeRegions &>
   parseCodeRegions(const std::unique_ptr<MCInstPrinter> &IP) override;
 };
 
+class AsmAnalysisRegionGenerator final : public AnalysisRegionGenerator,
+                                         public AsmCodeRegionGenerator {
+  AnalysisRegionCommentConsumer CC;
+
+public:
+  AsmAnalysisRegionGenerator(const Target &T, llvm::SourceMgr &SM, MCContext &C,
+                             const MCAsmInfo &A, const MCSubtargetInfo &S,
+                             const MCInstrInfo &I)
+      : AnalysisRegionGenerator(SM), AsmCodeRegionGenerator(T, C, A, S, I),
+        CC(Regions) {}
+
+  MCACommentConsumer *getCommentConsumer() override { return &CC; };
+  CodeRegions &getRegions() override { return Regions; };
+
+  Expected<const AnalysisRegions &>
+  parseAnalysisRegions(const std::unique_ptr<MCInstPrinter> &IP) override {
+    Expected<const CodeRegions &> RegionsOrErr = parseCodeRegions(IP);
+    if (!RegionsOrErr)
+      return RegionsOrErr.takeError();
+    else
+      return static_cast<const AnalysisRegions &>(*RegionsOrErr);
+  }
+};
+
+class AsmInstrumentRegionGenerator final : public InstrumentRegionGenerator,
+                                           public AsmCodeRegionGenerator {
+  InstrumentRegionCommentConsumer CC;
+
+public:
+  AsmInstrumentRegionGenerator(const Target &T, llvm::SourceMgr &SM,
+                               MCContext &C, const MCAsmInfo &A,
+                               const MCSubtargetInfo &S, const MCInstrInfo &I,
+                               InstrumentManager &IM)
+      : InstrumentRegionGenerator(SM), AsmCodeRegionGenerator(T, C, A, S, I),
+        CC(SM, Regions, IM) {}
+
+  MCACommentConsumer *getCommentConsumer() override { return &CC; };
+  CodeRegions &getRegions() override { return Regions; };
+
+  Expected<const InstrumentRegions &>
+  parseInstrumentRegions(const std::unique_ptr<MCInstPrinter> &IP) override {
+    Expected<const CodeRegions &> RegionsOrErr = parseCodeRegions(IP);
+    if (!RegionsOrErr)
+      return RegionsOrErr.takeError();
+    else
+      return static_cast<const InstrumentRegions &>(*RegionsOrErr);
+  }
+};
+
 } // namespace mca
 } // namespace llvm
 
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(),