[llvm-mca][RISCV] vsetivli and vsetvli act as instruments

[michaelmaitland]

Summary

Since the LMUL data that is needed to create an instrument is
avaliable statically from vsetivli and vsetvli instructions,
LMUL instruments can be automatically generated so that clients
of the tool do no need to manually insert instrument comments.

Instrument comments may be placed after a vset{i}vli instruction,
which will override instrument that was automatically inserted.
As a result, clients of llvm-mca instruments do not need to update
their existing instrument comments. However, if the instrument
has the same LMUL as the vset{i}vli, then it is reccomended to
remove the instrument comment as it becomes redundant.

[andreadb]

llvm/tools/llvm-mca/CodeRegionGenerator.h

121–127

Have you considered different approaches? Regions are already available from the base class. We don't need two references to the same data.

131

You can remove this comment.

135

Please remove this.

[myhsu]

llvm/include/llvm/MCA/CustomBehaviour.h

166

I feel like the comment here (and maybe that of createInstrument(StringRef,StringRef)) can hint that there are two ways to create Instruments -- by string description + data or from MCInst

llvm/tools/llvm-mca/CodeRegionGenerator.h

135

please remove this or put into LLVM_DEBUG

136

we can use Instruments returned earlier rather than create a new one

[michaelmaitland]

• Avoid duplicate store of CodeRegions and InstrumentRegions in InstrumentMCStreamer
• Remove errs()
• Remove duplicate call to createInstruments

[michaelmaitland]

llvm/include/llvm/MCA/CustomBehaviour.h

166

I've added this hint on both functions.

llvm/tools/llvm-mca/CodeRegionGenerator.h

121–127

I've updated this to avoid duplicate reference.

[michaelmaitland]

Use virtual functions to avoid downcast

[michaelmaitland]

• Revert Regions pointer back to Regions reference

[andreadb]

Much better thanks!

Patch looks good to me. If @myhsu is also happy with it, then feel free to commit it.

[myhsu]

LGTM thanks!

llvm/include/llvm/MCA/CustomBehaviour.h

public:
 
   /// Returns true if llvm-mca should ignore instruments.
   virtual bool shouldIgnoreInstruments() const { return true; }
 
   // Returns true if this supports processing Instrument with
   // Instrument.Desc equal to Type
   virtual bool supportsInstrumentType(StringRef Type) const { return false; }
 
-  /// Allocate an Instrument, and return a unique pointer to it.
+  /// Allocate an Instrument, and return a unique pointer to it. This function
+  /// may be useful to create instruments coming from comments in the assembly.
+  /// See createInstruments to create Instruments from MCInst
   virtual UniqueInstrument createInstrument(StringRef Desc, StringRef Data);
 
+  /// Return a list of unique pointers to Instruments, where each Instrument
+  /// is allocated by this function. See createInstrument to create Instrument
+  /// from a description and data.
+  virtual SmallVector<UniqueInstrument> createInstruments(const MCInst &Inst);
+
   /// Given an MCInst and a vector of Instrument, a target can
   /// return a SchedClassID. This can be used by a subtarget to return a
   /// PseudoInstruction SchedClassID instead of the one that belongs to the
   /// BaseInstruction This can be useful when a BaseInstruction does not convey
   /// the correct scheduling information without additional data. By default,
   /// it returns the SchedClassID that belongs to MCI.
   virtual unsigned getSchedClassID(const MCInstrInfo &MCII, const MCInst &MCI,
                                    const SmallVector<Instrument *> &IVec) const;
 };
 
 } // namespace mca
 } // namespace llvm
 
 #endif /* LLVM_MCA_CUSTOMBEHAVIOUR_H */

llvm/lib/MCA/CustomBehaviour.cpp

CustomBehaviour::getEndViews(llvm::MCInstPrinter &IP,
   return std::vector<std::unique_ptr<View>>();
 }
 
 UniqueInstrument InstrumentManager::createInstrument(llvm::StringRef Desc,
                                                      llvm::StringRef Data) {
   return std::make_unique<Instrument>(Desc, Data);
 }
 
+SmallVector<UniqueInstrument>
+InstrumentManager::createInstruments(const MCInst &Inst) {
+  return SmallVector<UniqueInstrument>();
+}
+
 unsigned InstrumentManager::getSchedClassID(
     const MCInstrInfo &MCII, const MCInst &MCI,
     const llvm::SmallVector<Instrument *> &IVec) const {
   return MCII.get(MCI.getOpcode()).getSchedClass();
 }
 
 } // namespace mca
 } // namespace llvm

llvm/lib/Target/RISCV/MCA/RISCVCustomBehaviour.h

RISCVInstrumentManager(const MCSubtargetInfo &STI, const MCInstrInfo &MCII)
       : InstrumentManager(STI, MCII) {}
 
   bool shouldIgnoreInstruments() const override { return false; }
   bool supportsInstrumentType(StringRef Type) const override;
 
   /// Create a Instrument for RISC-V target
   UniqueInstrument createInstrument(StringRef Desc, StringRef Data) override;
 
+  SmallVector<UniqueInstrument> createInstruments(const MCInst &Inst) override;
+
   /// Using the Instrument, returns a SchedClassID to use instead of
   /// the SchedClassID that belongs to the MCI or the original SchedClassID.
   unsigned
   getSchedClassID(const MCInstrInfo &MCII, const MCInst &MCI,
                   const SmallVector<Instrument *> &IVec) const override;
 };
 
 } // namespace mca
 } // namespace llvm
 
 #endif

llvm/lib/Target/RISCV/MCA/RISCVCustomBehaviour.cpp

 //===------------------- RISCVCustomBehaviour.cpp ---------------*-C++ -* -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 /// \file
 ///
 /// This file implements methods from the RISCVCustomBehaviour class.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "RISCVCustomBehaviour.h"
 #include "MCTargetDesc/RISCVMCTargetDesc.h"
+#include "RISCV.h"
 #include "RISCVInstrInfo.h"
 #include "TargetInfo/RISCVTargetInfo.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "llvm-mca-riscv-custombehaviour"
 
 // This brings in a table with primary key of
RISCVInstrumentManager::createInstrument(llvm::StringRef Desc,
   if (!RISCVLMULInstrument::isDataValid(Data)) {
     LLVM_DEBUG(dbgs() << "RVCB: Bad data for instrument kind " << Desc << ": "
                       << Data << '\n');
     return nullptr;
   }
   return std::make_unique<RISCVLMULInstrument>(Data);
 }
 
+SmallVector<UniqueInstrument>
+RISCVInstrumentManager::createInstruments(const MCInst &Inst) {
+  if (Inst.getOpcode() == RISCV::VSETVLI ||
+      Inst.getOpcode() == RISCV::VSETIVLI) {
+    LLVM_DEBUG(dbgs() << "RVCB: Found VSETVLI and creating instrument for it: "
+                      << Inst << "\n");
+    unsigned VTypeI = Inst.getOperand(2).getImm();
+    RISCVII::VLMUL VLMUL = RISCVVType::getVLMUL(VTypeI);
+
+    StringRef LMUL;
+    switch (VLMUL) {
+    case RISCVII::LMUL_1:
+      LMUL = "M1";
+      break;
+    case RISCVII::LMUL_2:
+      LMUL = "M2";
+      break;
+    case RISCVII::LMUL_4:
+      LMUL = "M4";
+      break;
+    case RISCVII::LMUL_8:
+      LMUL = "M8";
+      break;
+    case RISCVII::LMUL_F2:
+      LMUL = "MF2";
+      break;
+    case RISCVII::LMUL_F4:
+      LMUL = "MF4";
+      break;
+    case RISCVII::LMUL_F8:
+      LMUL = "MF8";
+      break;
+    case RISCVII::LMUL_RESERVED:
+      llvm_unreachable("Cannot create instrument for LMUL_RESERVED");
+    }
+    SmallVector<UniqueInstrument> Instruments;
+    Instruments.emplace_back(
+        createInstrument(RISCVLMULInstrument::DESC_NAME, LMUL));
+    return Instruments;
+  }
+  return SmallVector<UniqueInstrument>();
+}
+
 unsigned RISCVInstrumentManager::getSchedClassID(
     const MCInstrInfo &MCII, const MCInst &MCI,
     const llvm::SmallVector<Instrument *> &IVec) const {
   unsigned short Opcode = MCI.getOpcode();
   unsigned SchedClassID = MCII.get(Opcode).getSchedClass();
 
   for (const auto &I : IVec) {
     // Unknown Instrument kind

llvm/test/tools/llvm-mca/RISCV/no-vsetvli-to-start.s

+# NOTE: Assertions have been autogenerated by utils/update_mca_test_checks.py
+# RUN: llvm-mca -mtriple=riscv64 -mcpu=sifive-x280 -timeline -iterations=1 < %s | FileCheck %s
+
+vadd.vv v12, v12, v12
+vsetvli zero, a0, e8, m1, tu, mu
+vadd.vv v12, v12, v12
+
+# CHECK:      Iterations:        1
+# CHECK-NEXT: Instructions:      3
+# CHECK-NEXT: Total Cycles:      21
+# CHECK-NEXT: Total uOps:        3
+
+# CHECK:      Dispatch Width:    2
+# CHECK-NEXT: uOps Per Cycle:    0.14
+# CHECK-NEXT: IPC:               0.14
+# CHECK-NEXT: Block RThroughput: 18.0
+
+# CHECK:      Instruction Info:
+# CHECK-NEXT: [1]: #uOps
+# CHECK-NEXT: [2]: Latency
+# CHECK-NEXT: [3]: RThroughput
+# CHECK-NEXT: [4]: MayLoad
+# CHECK-NEXT: [5]: MayStore
+# CHECK-NEXT: [6]: HasSideEffects (U)
+
+# CHECK:      [1]    [2]    [3]    [4]    [5]    [6]    Instructions:
+# CHECK-NEXT:  1      4     16.00                       vadd.vv v12, v12, v12
+# CHECK-NEXT:  1      3     1.00                  U     vsetvli zero, a0, e8, m1, tu, mu
+# CHECK-NEXT:  1      4     2.00                        vadd.vv v12, v12, v12
+
+# CHECK:      Resources:
+# CHECK-NEXT: [0]   - SiFive7FDiv
+# CHECK-NEXT: [1]   - SiFive7IDiv
+# CHECK-NEXT: [2]   - SiFive7PipeA
+# CHECK-NEXT: [3]   - SiFive7PipeB
+# CHECK-NEXT: [4]   - SiFive7PipeV
+# CHECK-NEXT: [5]   - SiFive7VA
+# CHECK-NEXT: [6]   - SiFive7VL
+# CHECK-NEXT: [7]   - SiFive7VS
+
+# CHECK:      Resource pressure per iteration:
+# CHECK-NEXT: [0]    [1]    [2]    [3]    [4]    [5]    [6]    [7]
+# CHECK-NEXT:  -      -     1.00    -     18.00  18.00   -      -
+
+# CHECK:      Resource pressure by instruction:
+# CHECK-NEXT: [0]    [1]    [2]    [3]    [4]    [5]    [6]    [7]    Instructions:
+# CHECK-NEXT:  -      -      -      -     16.00  16.00   -      -     vadd.vv v12, v12, v12
+# CHECK-NEXT:  -      -     1.00    -      -      -      -      -     vsetvli zero, a0, e8, m1, tu, mu
+# CHECK-NEXT:  -      -      -      -     2.00   2.00    -      -     vadd.vv v12, v12, v12
+
+# CHECK:      Timeline view:
+# CHECK-NEXT:                     0123456789
+# CHECK-NEXT: Index     0123456789          0
+
+# CHECK:      [0,0]     DeeeE.    .    .    .   vadd.vv v12, v12, v12
+# CHECK-NEXT: [0,1]     .DeeE.    .    .    .   vsetvli zero, a0, e8, m1, tu, mu
+# CHECK-NEXT: [0,2]     .    .    .    .DeeeE   vadd.vv v12, v12, v12
+
+# CHECK:      Average Wait times (based on the timeline view):
+# CHECK-NEXT: [0]: Executions
+# CHECK-NEXT: [1]: Average time spent waiting in a scheduler's queue
+# CHECK-NEXT: [2]: Average time spent waiting in a scheduler's queue while ready
+# CHECK-NEXT: [3]: Average time elapsed from WB until retire stage
+
+# CHECK:            [0]    [1]    [2]    [3]
+# CHECK-NEXT: 0.     1     0.0    0.0    0.0       vadd.vv  v12, v12, v12
+# CHECK-NEXT: 1.     1     0.0    0.0    0.0       vsetvli  zero, a0, e8, m1, tu, mu
+# CHECK-NEXT: 2.     1     0.0    0.0    0.0       vadd.vv  v12, v12, v12
+# CHECK-NEXT:        1     0.0    0.0    0.0       <total>

llvm/test/tools/llvm-mca/RISCV/vsetivli-lmul-instrument.s

+# NOTE: Assertions have been autogenerated by utils/update_mca_test_checks.py
+# RUN: llvm-mca -mtriple=riscv64 -mcpu=sifive-x280 -timeline -iterations=1 < %s | FileCheck %s
+
+vsetivli zero, 8, e8, m1, tu, mu
+vadd.vv v12, v12, v12
+vsetivli zero, 8, e8, m8, tu, mu
+vadd.vv v12, v12, v12
+
+# CHECK:      Iterations:        1
+# CHECK-NEXT: Instructions:      4
+# CHECK-NEXT: Total Cycles:      12
+# CHECK-NEXT: Total uOps:        4
+
+# CHECK:      Dispatch Width:    2
+# CHECK-NEXT: uOps Per Cycle:    0.33
+# CHECK-NEXT: IPC:               0.33
+# CHECK-NEXT: Block RThroughput: 18.0
+
+# CHECK:      Instruction Info:
+# CHECK-NEXT: [1]: #uOps
+# CHECK-NEXT: [2]: Latency
+# CHECK-NEXT: [3]: RThroughput
+# CHECK-NEXT: [4]: MayLoad
+# CHECK-NEXT: [5]: MayStore
+# CHECK-NEXT: [6]: HasSideEffects (U)
+
+# CHECK:      [1]    [2]    [3]    [4]    [5]    [6]    Instructions:
+# CHECK-NEXT:  1      3     1.00                  U     vsetivli  zero, 8, e8, m1, tu, mu
+# CHECK-NEXT:  1      4     2.00                        vadd.vv v12, v12, v12
+# CHECK-NEXT:  1      3     1.00                  U     vsetivli  zero, 8, e8, m8, tu, mu
+# CHECK-NEXT:  1      4     16.00                       vadd.vv v12, v12, v12
+
+# CHECK:      Resources:
+# CHECK-NEXT: [0]   - SiFive7FDiv
+# CHECK-NEXT: [1]   - SiFive7IDiv
+# CHECK-NEXT: [2]   - SiFive7PipeA
+# CHECK-NEXT: [3]   - SiFive7PipeB
+# CHECK-NEXT: [4]   - SiFive7PipeV
+# CHECK-NEXT: [5]   - SiFive7VA
+# CHECK-NEXT: [6]   - SiFive7VL
+# CHECK-NEXT: [7]   - SiFive7VS
+
+# CHECK:      Resource pressure per iteration:
+# CHECK-NEXT: [0]    [1]    [2]    [3]    [4]    [5]    [6]    [7]
+# CHECK-NEXT:  -      -     2.00    -     18.00  18.00   -      -
+
+# CHECK:      Resource pressure by instruction:
+# CHECK-NEXT: [0]    [1]    [2]    [3]    [4]    [5]    [6]    [7]    Instructions:
+# CHECK-NEXT:  -      -     1.00    -      -      -      -      -     vsetivli  zero, 8, e8, m1, tu, mu
+# CHECK-NEXT:  -      -      -      -     2.00   2.00    -      -     vadd.vv v12, v12, v12
+# CHECK-NEXT:  -      -     1.00    -      -      -      -      -     vsetivli  zero, 8, e8, m8, tu, mu
+# CHECK-NEXT:  -      -      -      -     16.00  16.00   -      -     vadd.vv v12, v12, v12
+
+# CHECK:      Timeline view:
+# CHECK-NEXT:                     01
+# CHECK-NEXT: Index     0123456789
+
+# CHECK:      [0,0]     DeeE .    ..   vsetivli zero, 8, e8, m1, tu, mu
+# CHECK-NEXT: [0,1]     .  DeeeE  ..   vadd.vv  v12, v12, v12
+# CHECK-NEXT: [0,2]     .   DeeE  ..   vsetivli zero, 8, e8, m8, tu, mu
+# CHECK-NEXT: [0,3]     .    . DeeeE   vadd.vv  v12, v12, v12
+
+# CHECK:      Average Wait times (based on the timeline view):
+# CHECK-NEXT: [0]: Executions
+# CHECK-NEXT: [1]: Average time spent waiting in a scheduler's queue
+# CHECK-NEXT: [2]: Average time spent waiting in a scheduler's queue while ready
+# CHECK-NEXT: [3]: Average time elapsed from WB until retire stage
+
+# CHECK:            [0]    [1]    [2]    [3]
+# CHECK-NEXT: 0.     1     0.0    0.0    0.0       vsetivli zero, 8, e8, m1, tu, mu
+# CHECK-NEXT: 1.     1     0.0    0.0    0.0       vadd.vv  v12, v12, v12
+# CHECK-NEXT: 2.     1     0.0    0.0    0.0       vsetivli zero, 8, e8, m8, tu, mu
+# CHECK-NEXT: 3.     1     0.0    0.0    0.0       vadd.vv  v12, v12, v12
+# CHECK-NEXT:        1     0.0    0.0    0.0       <total>

llvm/test/tools/llvm-mca/RISCV/vsetvli-lmul-instrument.s

+# NOTE: Assertions have been autogenerated by utils/update_mca_test_checks.py
+# RUN: llvm-mca -mtriple=riscv64 -mcpu=sifive-x280 -timeline -iterations=1 < %s | FileCheck %s
+
+vsetvli zero, a0, e8, m1, tu, mu
+vadd.vv v12, v12, v12
+vsetvli zero, a0, e8, m8, tu, mu
+vadd.vv v12, v12, v12
+
+# CHECK:      Iterations:        1
+# CHECK-NEXT: Instructions:      4
+# CHECK-NEXT: Total Cycles:      12
+# CHECK-NEXT: Total uOps:        4
+
+# CHECK:      Dispatch Width:    2
+# CHECK-NEXT: uOps Per Cycle:    0.33
+# CHECK-NEXT: IPC:               0.33
+# CHECK-NEXT: Block RThroughput: 18.0
+
+# CHECK:      Instruction Info:
+# CHECK-NEXT: [1]: #uOps
+# CHECK-NEXT: [2]: Latency
+# CHECK-NEXT: [3]: RThroughput
+# CHECK-NEXT: [4]: MayLoad
+# CHECK-NEXT: [5]: MayStore
+# CHECK-NEXT: [6]: HasSideEffects (U)
+
+# CHECK:      [1]    [2]    [3]    [4]    [5]    [6]    Instructions:
+# CHECK-NEXT:  1      3     1.00                  U     vsetvli zero, a0, e8, m1, tu, mu
+# CHECK-NEXT:  1      4     2.00                        vadd.vv v12, v12, v12
+# CHECK-NEXT:  1      3     1.00                  U     vsetvli zero, a0, e8, m8, tu, mu
+# CHECK-NEXT:  1      4     16.00                       vadd.vv v12, v12, v12
+
+# CHECK:      Resources:
+# CHECK-NEXT: [0]   - SiFive7FDiv
+# CHECK-NEXT: [1]   - SiFive7IDiv
+# CHECK-NEXT: [2]   - SiFive7PipeA
+# CHECK-NEXT: [3]   - SiFive7PipeB
+# CHECK-NEXT: [4]   - SiFive7PipeV
+# CHECK-NEXT: [5]   - SiFive7VA
+# CHECK-NEXT: [6]   - SiFive7VL
+# CHECK-NEXT: [7]   - SiFive7VS
+
+# CHECK:      Resource pressure per iteration:
+# CHECK-NEXT: [0]    [1]    [2]    [3]    [4]    [5]    [6]    [7]
+# CHECK-NEXT:  -      -     2.00    -     18.00  18.00   -      -
+
+# CHECK:      Resource pressure by instruction:
+# CHECK-NEXT: [0]    [1]    [2]    [3]    [4]    [5]    [6]    [7]    Instructions:
+# CHECK-NEXT:  -      -     1.00    -      -      -      -      -     vsetvli zero, a0, e8, m1, tu, mu
+# CHECK-NEXT:  -      -      -      -     2.00   2.00    -      -     vadd.vv v12, v12, v12
+# CHECK-NEXT:  -      -     1.00    -      -      -      -      -     vsetvli zero, a0, e8, m8, tu, mu
+# CHECK-NEXT:  -      -      -      -     16.00  16.00   -      -     vadd.vv v12, v12, v12
+
+# CHECK:      Timeline view:
+# CHECK-NEXT:                     01
+# CHECK-NEXT: Index     0123456789
+
+# CHECK:      [0,0]     DeeE .    ..   vsetvli  zero, a0, e8, m1, tu, mu
+# CHECK-NEXT: [0,1]     .  DeeeE  ..   vadd.vv  v12, v12, v12
+# CHECK-NEXT: [0,2]     .   DeeE  ..   vsetvli  zero, a0, e8, m8, tu, mu
+# CHECK-NEXT: [0,3]     .    . DeeeE   vadd.vv  v12, v12, v12
+
+# CHECK:      Average Wait times (based on the timeline view):
+# CHECK-NEXT: [0]: Executions
+# CHECK-NEXT: [1]: Average time spent waiting in a scheduler's queue
+# CHECK-NEXT: [2]: Average time spent waiting in a scheduler's queue while ready
+# CHECK-NEXT: [3]: Average time elapsed from WB until retire stage
+
+# CHECK:            [0]    [1]    [2]    [3]
+# CHECK-NEXT: 0.     1     0.0    0.0    0.0       vsetvli  zero, a0, e8, m1, tu, mu
+# CHECK-NEXT: 1.     1     0.0    0.0    0.0       vadd.vv  v12, v12, v12
+# CHECK-NEXT: 2.     1     0.0    0.0    0.0       vsetvli  zero, a0, e8, m8, tu, mu
+# CHECK-NEXT: 3.     1     0.0    0.0    0.0       vadd.vv  v12, v12, v12
+# CHECK-NEXT:        1     0.0    0.0    0.0       <total>

llvm/tools/llvm-mca/CodeRegion.h

bool empty() const {
     });
   }
 
   bool isValid() const { return !FoundErrors; }
 
   bool isRegionActive(llvm::StringRef Description) const {
     return ActiveRegions.contains(Description);
   }
+
+  virtual void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc) = 0;
+  virtual void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc,
+                           UniqueInstrument Instrument) = 0;
+  virtual void endRegion(llvm::StringRef Description, llvm::SMLoc Loc) = 0;
 };
 
 struct AnalysisRegions : public CodeRegions {
   AnalysisRegions(llvm::SourceMgr &S);
 
-  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc);
-  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc);
+  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc) override;
+  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc,
+                   UniqueInstrument Instrument) override {}
+  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc) override;
 };
 
 struct InstrumentRegions : public CodeRegions {
 
   InstrumentRegions(llvm::SourceMgr &S);
 
+  void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc) override{};
   void beginRegion(llvm::StringRef Description, llvm::SMLoc Loc,
-                   UniqueInstrument Instrument);
-  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc);
+                   UniqueInstrument Instrument) override;
+  void endRegion(llvm::StringRef Description, llvm::SMLoc Loc) override;
 
   const SmallVector<Instrument *> getActiveInstruments(llvm::SMLoc Loc) const;
 };
 
 } // namespace mca
 } // namespace llvm
 
 #endif

llvm/tools/llvm-mca/CodeRegionGenerator.h

 
 #ifndef LLVM_TOOLS_LLVM_MCA_CODEREGION_GENERATOR_H
 #define LLVM_TOOLS_LLVM_MCA_CODEREGION_GENERATOR_H
 
 #include "CodeRegion.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCParser/MCAsmLexer.h"
+#include "llvm/MC/MCStreamer.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>
 
 namespace llvm {
InstrumentRegionCommentConsumer(llvm::SourceMgr &SM, InstrumentRegions &R,
       : 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;
+
+  InstrumentManager &getInstrumentManager() { return IM; }
+};
+
+// This class provides the callbacks that occur when parsing input assembly.
+class MCStreamerWrapper : public MCStreamer {
+protected:
+  CodeRegions &Regions;
+
+public:
+  MCStreamerWrapper(MCContext &Context, mca::CodeRegions &R)
+      : MCStreamer(Context), Regions(R) {}
+
+  // We only want to intercept the emission of new instructions.
+  void emitInstruction(const MCInst &Inst,
+                       const MCSubtargetInfo & /* unused */) override {
+    Regions.addInstruction(Inst);
+  }
+
+  bool emitSymbolAttribute(MCSymbol *Symbol, MCSymbolAttr Attribute) override {
+    return true;
+  }
+
+  void emitCommonSymbol(MCSymbol *Symbol, uint64_t Size,
+                        Align ByteAlignment) override {}
+  void emitZerofill(MCSection *Section, MCSymbol *Symbol = nullptr,
+                    uint64_t Size = 0, Align ByteAlignment = Align(1),
+                    SMLoc Loc = SMLoc()) override {}
+  void emitGPRel32Value(const MCExpr *Value) override {}
+  void beginCOFFSymbolDef(const MCSymbol *Symbol) override {}
+  void emitCOFFSymbolStorageClass(int StorageClass) override {}
+  void emitCOFFSymbolType(int Type) override {}
+  void endCOFFSymbolDef() override {}
+
+  ArrayRef<MCInst> GetInstructionSequence(unsigned Index) const {
+    return Regions.getInstructionSequence(Index);
+  }
+};
+
+class InstrumentMCStreamer : public MCStreamerWrapper {
+  InstrumentManager &IM;
+
+public:
+  InstrumentMCStreamer(MCContext &Context, mca::InstrumentRegions &R,
+                       InstrumentManager &IM)
+      : MCStreamerWrapper(Context, R), IM(IM) {}
+
+  void emitInstruction(const MCInst &Inst,
+                       const MCSubtargetInfo &MCSI) override {
+    MCStreamerWrapper::emitInstruction(Inst, MCSI);
+
+    // We know that Regions is an InstrumentRegions by the constructor.
+    for (UniqueInstrument &I : IM.createInstruments(Inst)) {
+      StringRef InstrumentKind = I.get()->getDesc();
+      // End InstrumentType region if one is open
+      if (Regions.isRegionActive(InstrumentKind))
+        Regions.endRegion(InstrumentKind, Inst.getLoc());
+      // Start new instrumentation region
+      Regions.beginRegion(InstrumentKind, Inst.getLoc(), std::move(I));
+    }
+  }
 };
 
 /// 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:
   CodeRegionGenerator(const CodeRegionGenerator &) = delete;
   CodeRegionGenerator &operator=(const CodeRegionGenerator &) = delete;
public:
   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;
   const MCSubtargetInfo &STI;
   const MCInstrInfo &MCII;
   unsigned AssemblerDialect; // This is set during parsing.
 
+protected:
+  MCContext &Ctx;
+
 public:
   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) {}
+      : TheTarget(T), MAI(A), STI(S), MCII(I), AssemblerDialect(0), Ctx(C) {}
 
   virtual MCACommentConsumer *getCommentConsumer() = 0;
   virtual CodeRegions &getRegions() = 0;
+  virtual MCStreamerWrapper *getMCStreamer() = 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;
+  MCStreamerWrapper Streamer;
 
 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) {}
+        CC(Regions), Streamer(Ctx, Regions) {}
 
   MCACommentConsumer *getCommentConsumer() override { return &CC; };
   CodeRegions &getRegions() override { return Regions; };
+  MCStreamerWrapper *getMCStreamer() override { return &Streamer; }
 
   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);
   }
 
   Expected<const CodeRegions &>
   parseCodeRegions(const std::unique_ptr<MCInstPrinter> &IP) override {
     return AsmCodeRegionGenerator::parseCodeRegions(IP);
   }
 };
 
 class AsmInstrumentRegionGenerator final : public InstrumentRegionGenerator,
                                            public AsmCodeRegionGenerator {
   InstrumentRegionCommentConsumer CC;
+  InstrumentMCStreamer Streamer;
 
 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) {}
+        CC(SM, Regions, IM), Streamer(Ctx, Regions, IM) {}
 
   MCACommentConsumer *getCommentConsumer() override { return &CC; };
   CodeRegions &getRegions() override { return Regions; };
+  MCStreamerWrapper *getMCStreamer() override { return &Streamer; }
 
   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);

llvm/tools/llvm-mca/CodeRegionGenerator.cpp

 /// so the classes here provide the input-to-CodeRegions translation.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodeRegionGenerator.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/MC/MCParser/MCTargetAsmParser.h"
-#include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCTargetOptions.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/SMLoc.h"
 #include <memory>
 
 namespace llvm {
 namespace mca {
 
 // This virtual dtor serves as the anchor for the CodeRegionGenerator class.
 CodeRegionGenerator::~CodeRegionGenerator() {}
 
-// This class provides the callbacks that occur when parsing input assembly.
-class MCStreamerWrapper final : public MCStreamer {
-  CodeRegions &Regions;
-
-public:
-  MCStreamerWrapper(MCContext &Context, mca::CodeRegions &R)
-      : MCStreamer(Context), Regions(R) {}
-
-  // We only want to intercept the emission of new instructions.
-  void emitInstruction(const MCInst &Inst,
-                       const MCSubtargetInfo & /* unused */) override {
-    Regions.addInstruction(Inst);
-  }
-
-  bool emitSymbolAttribute(MCSymbol *Symbol, MCSymbolAttr Attribute) override {
-    return true;
-  }
-
-  void emitCommonSymbol(MCSymbol *Symbol, uint64_t Size,
-                        Align ByteAlignment) override {}
-  void emitZerofill(MCSection *Section, MCSymbol *Symbol = nullptr,
-                    uint64_t Size = 0, Align ByteAlignment = Align(1),
-                    SMLoc Loc = SMLoc()) override {}
-  void emitGPRel32Value(const MCExpr *Value) override {}
-  void beginCOFFSymbolDef(const MCSymbol *Symbol) override {}
-  void emitCOFFSymbolStorageClass(int StorageClass) override {}
-  void emitCOFFSymbolType(int Type) override {}
-  void endCOFFSymbolDef() override {}
-
-  ArrayRef<MCInst> GetInstructionSequence(unsigned Index) const {
-    return Regions.getInstructionSequence(Index);
-  }
-};
-
 Expected<const CodeRegions &> AsmCodeRegionGenerator::parseCodeRegions(
     const std::unique_ptr<MCInstPrinter> &IP) {
   MCTargetOptions Opts;
   Opts.PreserveAsmComments = false;
   CodeRegions &Regions = getRegions();
-  MCStreamerWrapper Str(Ctx, Regions);
+  MCStreamerWrapper *Str = getMCStreamer();
 
   // 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(),
+  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));
+      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));