[MCA] Enable customization of individual instructions

llvm/include/llvm/MCA/InstrBuilder.h

@@ -25,6 +25,7 @@
 #include "llvm/MCA/Support.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Error.h"
+#include <functional>
 
 namespace llvm {
 namespace mca {
@@ -78,6 +79,10 @@ class InstrBuilder {
   DenseMap<std::pair<hash_code, unsigned>, std::unique_ptr<const InstrDesc>>
       VariantDescriptors;
 
+  // These descriptors are customized for particular instructions and cannot
+  // be reused
+  SmallVector<std::unique_ptr<const InstrDesc>> CustomDescriptors;
+
   bool FirstCallInst;
   bool FirstReturnInst;
   unsigned CallLatency;
@@ -87,7 +92,8 @@ class InstrBuilder {
 
   Expected<unsigned> getVariantSchedClassID(const MCInst &MCI, unsigned SchedClassID);
   Expected<const InstrDesc &>
-  createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
+  createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
+                      std::function<void(InstrDesc &)> Customizer = {});
   Expected<const InstrDesc &>
   getOrCreateInstrDesc(const MCInst &MCI,
                        const SmallVector<Instrument *> &IVec);
@@ -116,7 +122,8 @@ class InstrBuilder {
   void setInstRecycleCallback(InstRecycleCallback CB) { InstRecycleCB = CB; }
 
   LLVM_ABI Expected<std::unique_ptr<Instruction>>
-  createInstruction(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
+  createInstruction(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
+                    std::function<void(InstrDesc &)> Customizer = {});
 };
 } // namespace mca
 } // namespace llvm

llvm/lib/MCA/InstrBuilder.cpp

@@ -558,7 +558,8 @@ Expected<unsigned> InstrBuilder::getVariantSchedClassID(const MCInst &MCI,
 
 Expected<const InstrDesc &>
 InstrBuilder::createInstrDescImpl(const MCInst &MCI,
-                                  const SmallVector<Instrument *> &IVec) {
+                                  const SmallVector<Instrument *> &IVec,
+                                  std::function<void(InstrDesc &)> Customizer) {
   assert(STI.getSchedModel().hasInstrSchedModel() &&
          "Itineraries are not yet supported!");
 
@@ -632,6 +633,12 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
     return std::move(Err);
 
   // Now add the new descriptor.
+
+  if (Customizer) {
+    Customizer(*ID);
+    return *CustomDescriptors.emplace_back(std::move(ID));
+  }
+
   bool IsVariadic = MCDesc.isVariadic();
   if ((ID->IsRecyclable = !IsVariadic && !IsVariant)) {
     auto DKey = std::make_pair(MCI.getOpcode(), SchedClassID);
@@ -675,8 +682,11 @@ STATISTIC(NumVariantInst, "Number of MCInsts that doesn't have static Desc");
 
 Expected<std::unique_ptr<Instruction>>
 InstrBuilder::createInstruction(const MCInst &MCI,
-                                const SmallVector<Instrument *> &IVec) {
-  Expected<const InstrDesc &> DescOrErr = getOrCreateInstrDesc(MCI, IVec);
+                                const SmallVector<Instrument *> &IVec,
+                                std::function<void(InstrDesc &)> Customizer) {
+  Expected<const InstrDesc &> DescOrErr =
+      Customizer ? createInstrDescImpl(MCI, IVec, Customizer)
+                 : getOrCreateInstrDesc(MCI, IVec);
   if (!DescOrErr)
     return DescOrErr.takeError();
   const InstrDesc &D = *DescOrErr;

llvm/test/tools/llvm-mca/X86/llvm-mca-markers-13.s

@@ -0,0 +1,10 @@
+# NOTE: Assertions have been autogenerated by utils/update_mca_test_checks.py
+# RUN: llvm-mca -mtriple=x86_64-unknown-unknown -mcpu=btver2 -iterations=10 %s 2>&1 | FileCheck %s
+
+add (%eax), %eax // LLVM-MCA-LATENCY : 100
+mov %eax, (%ebx)
+
+# CHECK:      Iterations:        10
+# CHECK-NEXT: Instructions:      20
+# CHECK-NEXT: Total Cycles:      1004
+# CHECK-NEXT: Total uOps:        20

llvm/test/tools/llvm-mca/X86/llvm-mca-markers-14.s

@@ -0,0 +1,10 @@
+# NOTE: Assertions have been autogenerated by utils/update_mca_test_checks.py
+# RUN: llvm-mca -mtriple=x86_64-unknown-unknown -mcpu=btver2 -iterations=10 %s 2>&1 | FileCheck %s
+
+add (%eax), %eax // LLVM-MCA-LATENCY : 100
+mov %eax, (%ebx) // LLVM-MCA-LATENCY : 100
+
+# CHECK:      Iterations:        10
+# CHECK-NEXT: Instructions:      20
+# CHECK-NEXT: Total Cycles:      1103
+# CHECK-NEXT: Total uOps:        20

llvm/tools/llvm-mca/CodeRegion.h

@@ -68,11 +68,16 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/SourceMgr.h"
+#include <utility>
 #include <vector>
 
 namespace llvm {
 namespace mca {
 
+struct InstAnnotation {
+  std::optional<unsigned> Latency;
+};
+
 /// A region of assembly code.
 ///
 /// It identifies a sequence of machine instructions.
@@ -155,6 +160,9 @@ class CodeRegions {
   llvm::StringMap<unsigned> ActiveRegions;
   bool FoundErrors;
 
+  // Annotations specified in comments, indexed by SMLoc value
+  llvm::DenseMap<const char *, InstAnnotation> Annotations;
+
 public:
   CodeRegions(llvm::SourceMgr &S) : SM(S), FoundErrors(false) {}
   virtual ~CodeRegions() = default;
@@ -170,6 +178,18 @@ class CodeRegions {
   void addInstruction(const llvm::MCInst &Instruction);
   llvm::SourceMgr &getSourceMgr() const { return SM; }
 
+  void Annotate(llvm::SMLoc Loc, const InstAnnotation &A) {
+    Annotations[Loc.getPointer()] = A;
+  }
+  std::optional<unsigned> getExplicitLatency(llvm::SMLoc Loc) const {
+    const auto It = Annotations.find(Loc.getPointer());
+    if (It != Annotations.end()) {
+      return It->second.Latency;
+    } else {
+      return {};
+    }
+  }
+
   llvm::ArrayRef<llvm::MCInst> getInstructionSequence(unsigned Idx) const {
     return Regions[Idx]->getInstructions();
   }

llvm/tools/llvm-mca/CodeRegionGenerator.cpp

@@ -95,6 +95,18 @@ void AnalysisRegionCommentConsumer::HandleComment(SMLoc Loc,
     return;
 
   Comment = Comment.drop_front(Position);
+  if (Comment.starts_with("LLVM-MCA-LATENCY")) {
+    auto Parts = Comment.split(':');
+    Position = Parts.second.find_first_not_of(" \t");
+    if (Position >= Parts.second.size())
+      return;
+    auto LatStr = Parts.second.drop_front(Position);
+    unsigned Latency = 0;
+    if (!LatStr.getAsInteger(10, Latency))
+      Streamer.AddLatencyAnnotation(Latency);
+    return;
+  }
+
   if (Comment.consume_front("LLVM-MCA-END")) {
     // Skip spaces and tabs.
     Position = Comment.find_first_not_of(" \t");

llvm/tools/llvm-mca/CodeRegionGenerator.h

@@ -31,6 +31,51 @@
 namespace llvm {
 namespace mca {
 
+// This class provides the callbacks that occur when parsing input assembly.
+class MCStreamerWrapper : public MCStreamer {
+protected:
+  CodeRegions &Regions;
+  std::optional<InstAnnotation> CurrentAnnotation;
+
+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);
+    if (CurrentAnnotation) {
+      Regions.Annotate(Inst.getLoc(), *CurrentAnnotation);
+      CurrentAnnotation = {};
+    }
+  }
+
+  void AddLatencyAnnotation(unsigned Lat) {
+    if (!CurrentAnnotation)
+      CurrentAnnotation = InstAnnotation();
+    CurrentAnnotation->Latency = Lat;
+  }
+
+  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 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 MCACommentConsumer : public AsmCommentConsumer {
 protected:
   bool FoundError = false;
@@ -44,9 +89,11 @@ class MCACommentConsumer : public AsmCommentConsumer {
 /// A comment consumer that parses strings.  The only valid tokens are strings.
 class AnalysisRegionCommentConsumer : public MCACommentConsumer {
   AnalysisRegions &Regions;
+  MCStreamerWrapper &Streamer;
 
 public:
-  AnalysisRegionCommentConsumer(AnalysisRegions &R) : Regions(R) {}
+  AnalysisRegionCommentConsumer(AnalysisRegions &R, MCStreamerWrapper &S)
+      : Regions(R), Streamer(S) {}
 
   /// 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.
@@ -82,40 +129,6 @@ class InstrumentRegionCommentConsumer : public MCACommentConsumer {
   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 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;
 
@@ -210,15 +223,15 @@ class AsmCodeRegionGenerator : public virtual CodeRegionGenerator {
 
 class AsmAnalysisRegionGenerator final : public AnalysisRegionGenerator,
                                          public AsmCodeRegionGenerator {
-  AnalysisRegionCommentConsumer CC;
   MCStreamerWrapper Streamer;
+  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), Streamer(Ctx, Regions) {}
+        Streamer(Ctx, Regions), CC(Regions, Streamer) {}
 
   MCACommentConsumer *getCommentConsumer() override { return &CC; };
   CodeRegions &getRegions() override { return Regions; };

llvm/tools/llvm-mca/llvm-mca.cpp

@@ -633,8 +633,15 @@ int main(int argc, char **argv) {
       const SmallVector<mca::Instrument *> Instruments =
           InstrumentRegions.getActiveInstruments(Loc);
 
+      auto Latency = Regions.getExplicitLatency(Loc);
       Expected<std::unique_ptr<mca::Instruction>> Inst =
-          IB.createInstruction(MCI, Instruments);
+          Latency ? IB.createInstruction(MCI, Instruments,
+                                         [=](llvm::mca::InstrDesc &ID) {
+                                           for (auto &W : ID.Writes)
+                                             W.Latency = *Latency;
+                                           ID.MaxLatency = *Latency;
+                                         })
+                  : IB.createInstruction(MCI, Instruments);
       if (!Inst) {
         if (auto NewE = handleErrors(
                 Inst.takeError(),

llvm/unittests/tools/llvm-mca/MCATestBase.cpp

@@ -61,7 +61,7 @@ void MCATestBase::SetUp() {
 
 Error MCATestBase::runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
                                   ArrayRef<mca::View *> Views,
-                                  const mca::PipelineOptions *PO) {
+                                  const mca::PipelineOptions *PO, Builder B) {
   mca::Context MCA(*MRI, *STI);
 
   // Default InstrumentManager
@@ -72,7 +72,7 @@ Error MCATestBase::runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
   SmallVector<std::unique_ptr<mca::Instruction>> LoweredInsts;
   for (const auto &MCI : Insts) {
     Expected<std::unique_ptr<mca::Instruction>> Inst =
-        IB.createInstruction(MCI, Instruments);
+        B ? B(IB, MCI, Instruments) : IB.createInstruction(MCI, Instruments);
     if (!Inst) {
       if (auto NewE =
               handleErrors(Inst.takeError(),

llvm/unittests/tools/llvm-mca/MCATestBase.h

@@ -24,6 +24,7 @@
 #include "llvm/MC/MCTargetOptions.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/MCA/Context.h"
+#include "llvm/MCA/InstrBuilder.h"
 #include "llvm/TargetParser/SubtargetFeature.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -70,12 +71,17 @@ class MCATestBase : public ::testing::Test {
 
   void SetUp() override;
 
+  using Builder = std::function<Expected<std::unique_ptr<mca::Instruction>>(
+      mca::InstrBuilder &, const MCInst &,
+      const SmallVector<mca::Instrument *> &)>;
+
   /// Utility function to run MCA with (nearly) the same configuration as the
   /// `llvm-mca` tool to verify result correctness.
   /// This function only displays on SummaryView by default.
   virtual Error runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
                                ArrayRef<mca::View *> Views = {},
-                               const mca::PipelineOptions *PO = nullptr);
+                               const mca::PipelineOptions *PO = nullptr,
+                               Builder B = {});
 };
 
 } // end namespace mca

llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp

@@ -234,3 +234,31 @@ TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
   Expected<unsigned> Cycles = P->run();
   ASSERT_TRUE(static_cast<bool>(Cycles));
 }
+
+TEST_F(X86TestBase, TestInstructionCustomization) {
+  const unsigned ExplicitLatency = 100;
+  SmallVector<MCInst> MCIs;
+  MCInst InstructionToAdd = MCInstBuilder(X86::XOR64rr)
+                                .addReg(X86::RAX)
+                                .addReg(X86::RAX)
+                                .addReg(X86::RAX);
+  MCIs.push_back(InstructionToAdd);
+
+  // Run the baseline.
+  json::Object BaselineResult;
+  auto E = runBaselineMCA(BaselineResult, MCIs, {}, nullptr,
+                          [=](InstrBuilder &IB, const MCInst &MCI,
+                              const SmallVector<Instrument *> &Instruments) {
+                            return IB.createInstruction(
+                                MCI, Instruments, [=](InstrDesc &ID) {
+                                  for (auto &W : ID.Writes)
+                                    W.Latency = ExplicitLatency;
+                                  ID.MaxLatency = ExplicitLatency;
+                                });
+                          });
+  auto *BaselineObj = BaselineResult.getObject("SummaryView");
+  auto V = BaselineObj->getInteger("TotalCycles");
+  ASSERT_TRUE(V);
+  // Additional 3 cycles for Dispatch, Executed and Retired states
+  ASSERT_EQ(unsigned(*V), ExplicitLatency + 3) << "Total cycles do not match";
+}