[MCA] Enable customization of individual instructions

llvm/include/llvm/MCA/CustomBehaviour.h

@@ -130,6 +130,9 @@ class Instrument {
 
   virtual ~Instrument() = default;
 
+  virtual bool canCustomize() const { return false; }
+  virtual void customize(InstrDesc &) const {}
+
   StringRef getDesc() const { return Desc; }
   StringRef getData() const { return Data; }
 };
@@ -143,19 +146,27 @@ class LLVM_ABI InstrumentManager {
 protected:
   const MCSubtargetInfo &STI;
   const MCInstrInfo &MCII;
+  bool EnableDefaults;
+  std::unique_ptr<InstrumentManager> TargetIM;
 
 public:
-  InstrumentManager(const MCSubtargetInfo &STI, const MCInstrInfo &MCII)
-      : STI(STI), MCII(MCII) {}
+  InstrumentManager(const MCSubtargetInfo &STI, const MCInstrInfo &MCII,
+                    bool EnableDefaults = false,
+                    std::unique_ptr<InstrumentManager> TargetIM = {})
+      : STI(STI), MCII(MCII), EnableDefaults(EnableDefaults),
+        TargetIM(std::move(TargetIM)) {}
 
   virtual ~InstrumentManager() = default;
 
   /// Returns true if llvm-mca should ignore instruments.
-  virtual bool shouldIgnoreInstruments() const { return true; }
+  virtual bool shouldIgnoreInstruments() const {
+    return !EnableDefaults &&
+           (!TargetIM || TargetIM->shouldIgnoreInstruments());
+  }
 
   // Returns true if this supports processing Instrument with
   // Instrument.Desc equal to Type
-  virtual bool supportsInstrumentType(StringRef Type) const { return false; }
+  virtual bool supportsInstrumentType(StringRef Type) const;
 
   /// Allocate an Instrument, and return a unique pointer to it. This function
   /// may be useful to create instruments coming from comments in the assembly.
@@ -175,6 +186,13 @@ class LLVM_ABI InstrumentManager {
   /// it returns the SchedClassID that belongs to MCI.
   virtual unsigned getSchedClassID(const MCInstrInfo &MCII, const MCInst &MCI,
                                    const SmallVector<Instrument *> &IVec) const;
+
+  // Return true if instruments can modify instruction description
+  virtual bool canCustomize(const SmallVector<Instrument *> &IVec) const;
+
+  // Customize instruction description
+  virtual void customize(const SmallVector<Instrument *> &IVec,
+                         llvm::mca::InstrDesc &Desc) const;
 };
 
 } // namespace mca

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;

llvm/lib/MCA/CustomBehaviour.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MCA/CustomBehaviour.h"
+#include "llvm/MCA/Instruction.h"
 
 namespace llvm {
 namespace mca {
@@ -42,19 +43,91 @@ CustomBehaviour::getEndViews(llvm::MCInstPrinter &IP,
   return std::vector<std::unique_ptr<View>>();
 }
 
+class CustomInstrument : public Instrument {
+  std::optional<unsigned> Latency;
+
+public:
+  static const llvm::StringRef DESC_NAME;
+  explicit CustomInstrument(StringRef Data) : Instrument(DESC_NAME, Data) {
+    // Skip spaces and tabs.
+    unsigned Position = Data.find_first_not_of(" \t");
+    if (Position >= Data.size())
+      // We reached the end of the comment. Bail out.
+      return;
+    Data = Data.drop_front(Position);
+    auto [Name, Value] = Data.split(':');
+    if (Name.upper() == "LATENCY") {
+      Position = Value.find_first_not_of(" \t");
+      if (Position >= Value.size())
+        return;
+      auto Stripped = Value.drop_front(Position);
+      unsigned L = 0;
+      if (!Stripped.getAsInteger(10, L))
+        Latency = L;
+    }
+  }
+
+  bool canCustomize() const override { return bool(Latency); }
+
+  void customize(InstrDesc &ID) const override {
+    if (Latency) {
+      for (auto &W : ID.Writes)
+        W.Latency = *Latency;
+      ID.MaxLatency = *Latency;
+    }
+  }
+};
+
+const llvm::StringRef CustomInstrument::DESC_NAME = "CUSTOMIZE";
+
+bool InstrumentManager::supportsInstrumentType(StringRef Type) const {
+  if (EnableDefaults && Type == CustomInstrument::DESC_NAME)
+    return true;
+  if (TargetIM)
+    return TargetIM->supportsInstrumentType(Type);
+  return false;
+}
+
+bool InstrumentManager::canCustomize(
+    const llvm::SmallVector<Instrument *> &IVec) const {
+  for (const auto I : IVec) {
+    if (I->canCustomize())
+      return true;
+  }
+  return false;
+}
+
+void InstrumentManager::customize(const llvm::SmallVector<Instrument *> &IVec,
+                                  InstrDesc &ID) const {
+  for (const auto I : IVec) {
+    if (I->canCustomize())
+      I->customize(ID);
+  }
+}
+
 UniqueInstrument InstrumentManager::createInstrument(llvm::StringRef Desc,
                                                      llvm::StringRef Data) {
+  if (!EnableDefaults)
+    return std::make_unique<Instrument>(Desc, Data);
+  if (Desc == CustomInstrument::DESC_NAME)
+    return std::make_unique<CustomInstrument>(Data);
+  if (TargetIM && TargetIM->supportsInstrumentType(Desc))
+    return TargetIM->createInstrument(Desc, Data);
   return std::make_unique<Instrument>(Desc, Data);
 }
 
 SmallVector<UniqueInstrument>
 InstrumentManager::createInstruments(const MCInst &Inst) {
+  if (TargetIM)
+    return TargetIM->createInstruments(Inst);
   return SmallVector<UniqueInstrument>();
 }
 
 unsigned InstrumentManager::getSchedClassID(
     const MCInstrInfo &MCII, const MCInst &MCI,
     const llvm::SmallVector<Instrument *> &IVec) const {
+  if (TargetIM)
+    return TargetIM->getSchedClassID(MCII, MCI, IVec);
   return MCII.get(MCI.getOpcode()).getSchedClass();
 }
 

llvm/lib/MCA/InstrBuilder.cpp

@@ -632,6 +632,12 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
     return std::move(Err);
 
   // Now add the new descriptor.
+
+  if (IM.canCustomize(IVec)) {
+    IM.customize(IVec, *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);
@@ -676,7 +682,9 @@ 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);
+  Expected<const InstrDesc &> DescOrErr = IM.canCustomize(IVec)
+                                              ? createInstrDescImpl(MCI, IVec)
+                                              : 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,12 @@
+# 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
+
+# LLVM-MCA-CUSTOMIZE Latency:100
+add (%eax), %eax
+# LLVM-MCA-CUSTOMIZE
+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,11 @@
+# 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
+
+# LLVM-MCA-CUSTOMIZE Latency:100
+add (%eax), %eax
+mov %eax, (%ebx)
+
+# CHECK:      Iterations:        10
+# CHECK-NEXT: Instructions:      20
+# CHECK-NEXT: Total Cycles:      1103
+# CHECK-NEXT: Total uOps:        20

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

@@ -511,12 +511,15 @@ int main(int argc, char **argv) {
 
   std::unique_ptr<mca::InstrumentManager> IM;
   if (!DisableInstrumentManager) {
-    IM = std::unique_ptr<mca::InstrumentManager>(
-        TheTarget->createInstrumentManager(*STI, *MCII));
+    std::unique_ptr<mca::InstrumentManager> TargetIM =
+        std::unique_ptr<mca::InstrumentManager>(
+            TheTarget->createInstrumentManager(*STI, *MCII));
+    IM = std::make_unique<mca::InstrumentManager>(*STI, *MCII, true,
+                                                  std::move(TargetIM));
   }
   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).
+    // If -disable-cb flag is set then we use the base class with default
+    // behavior (which does nothing).
     IM = std::make_unique<mca::InstrumentManager>(*STI, *MCII);
   }
 

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

@@ -59,16 +59,25 @@ void MCATestBase::SetUp() {
   ASSERT_TRUE(IP);
 }
 
-Error MCATestBase::runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
-                                  ArrayRef<mca::View *> Views,
-                                  const mca::PipelineOptions *PO) {
+Error MCATestBase::runBaselineMCA(
+    json::Object &Result, ArrayRef<MCInst> Insts,
+    ArrayRef<mca::View *> Views,
+    const mca::PipelineOptions *PO,
+    SmallVector<std::pair<StringRef, StringRef>> Descs) {
   mca::Context MCA(*MRI, *STI);
 
-  // Default InstrumentManager
-  auto IM = std::make_unique<mca::InstrumentManager>(*STI, *MCII);
+  // Enable instruments when descriptions are provided
+  auto IM =
+      std::make_unique<mca::InstrumentManager>(*STI, *MCII, !Descs.empty());
   mca::InstrBuilder IB(*STI, *MCII, *MRI, MCIA.get(), *IM, /*CallLatency=*/100);
 
-  const SmallVector<mca::Instrument *> Instruments;
+  SmallVector<mca::Instrument *> Instruments;
+  SmallVector<mca::UniqueInstrument> InstrumentsOwner;
+  for (const auto &Desc : Descs) {
+    auto I = IM->createInstrument(Desc.first, Desc.second);
+    Instruments.push_back(I.get());
+    InstrumentsOwner.push_back(std::move(I));
+  }
   SmallVector<std::unique_ptr<mca::Instruction>> LoweredInsts;
   for (const auto &MCI : Insts) {
     Expected<std::unique_ptr<mca::Instruction>> Inst =

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"
 
@@ -73,9 +74,11 @@ class MCATestBase : public ::testing::Test {
   /// 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);
+  virtual Error
+  runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
+                 ArrayRef<mca::View *> Views = {},
+                 const mca::PipelineOptions *PO = nullptr,
+                 SmallVector<std::pair<StringRef, StringRef>> Descs = {});
 };
 
 } // end namespace mca

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

@@ -234,3 +234,26 @@ 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);
+  SmallVector<std::pair<StringRef, StringRef>> InstrDescs;
+  InstrDescs.push_back(
+      std::make_pair(StringRef("CUSTOMIZE"), StringRef("Latency:100")));
+
+  // Run the baseline.
+  json::Object BaselineResult;
+  auto E = runBaselineMCA(BaselineResult, MCIs, {}, nullptr, InstrDescs);
+  ASSERT_FALSE(bool(E)) << "Failed to run baseline";
+  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";
+}