[llvm] [MCA] Parameterize variant scheduling classes by explicit variable (PR #92849)

Aiden Grossman via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 28 10:10:53 PDT 2024


https://github.com/boomanaiden154 updated https://github.com/llvm/llvm-project/pull/92849

>From 37235b9927c3b60810b05100aa9381ec8b5febf9 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Tue, 21 May 2024 03:36:53 +0000
Subject: [PATCH 1/8] [MCA] Parameterize variant scheduling classes by explicit
 variable

This patch looks up variant scheduling classes using an explicit
variable (InstructionAddress passed to createInstruction) rather than
simply using the address of the MCInst that gets passed into
createInstruction. This breaks certain use cases that might occur out of
tree, like decoding an execution trace instruction by instruction and
creating MCA instructions as one goes along, like in the MCAD case. In
this case, the MCInst will always have the same address and thus all
instructions with the same variant scheduling class will end up with the
same instruction description, leading to undesired behavior (assertions,
uses after free, invalid results, etc.).
---
 llvm/include/llvm/MCA/InstrBuilder.h           | 18 +++++++++++++-----
 llvm/lib/MCA/InstrBuilder.cpp                  | 18 +++++++++++-------
 llvm/tools/llvm-mca/llvm-mca.cpp               |  4 ++--
 llvm/unittests/tools/llvm-mca/MCATestBase.cpp  |  4 ++--
 .../tools/llvm-mca/X86/TestIncrementalMCA.cpp  |  4 ++--
 5 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index 3594372489148..67f9c736054e1 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -72,7 +72,7 @@ class InstrBuilder {
       Descriptors;
 
   // Key is the MCIInst and SchedClassID the describe the value InstrDesc
-  DenseMap<std::pair<const MCInst *, unsigned>,
+  DenseMap<std::pair<uint64_t, unsigned>,
            std::unique_ptr<const InstrDesc>>
       VariantDescriptors;
 
@@ -83,10 +83,11 @@ class InstrBuilder {
   InstRecycleCallback InstRecycleCB;
 
   Expected<const InstrDesc &>
-  createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
+  createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
+                      uint64_t InstructionAddress);
   Expected<const InstrDesc &>
-  getOrCreateInstrDesc(const MCInst &MCI,
-                       const SmallVector<Instrument *> &IVec);
+  getOrCreateInstrDesc(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
+                       uint64_t InstructionAddress);
 
   InstrBuilder(const InstrBuilder &) = delete;
   InstrBuilder &operator=(const InstrBuilder &) = delete;
@@ -111,8 +112,15 @@ class InstrBuilder {
   /// or null if there isn't any.
   void setInstRecycleCallback(InstRecycleCallback CB) { InstRecycleCB = CB; }
 
+  /// Create an MCA Instruction from a MC Instruction that contains all the
+  /// relevant state MCA needs for modeling. Variant instructions (e.g.,
+  /// register zeroing idioms) each need their own instruction description
+  /// which is uniqued based on InstructionAddress. This can be an actual
+  /// address or something unique per instruction like a loop iteration
+  /// variable, but it must uniquely identify the instruction being passed in.
   Expected<std::unique_ptr<Instruction>>
-  createInstruction(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
+  createInstruction(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
+                    uint64_t InstructionAddress);
 };
 } // namespace mca
 } // namespace llvm
diff --git a/llvm/lib/MCA/InstrBuilder.cpp b/llvm/lib/MCA/InstrBuilder.cpp
index bcf065c566918..378b88812513f 100644
--- a/llvm/lib/MCA/InstrBuilder.cpp
+++ b/llvm/lib/MCA/InstrBuilder.cpp
@@ -521,7 +521,8 @@ Error InstrBuilder::verifyInstrDesc(const InstrDesc &ID,
 
 Expected<const InstrDesc &>
 InstrBuilder::createInstrDescImpl(const MCInst &MCI,
-                                  const SmallVector<Instrument *> &IVec) {
+                                  const SmallVector<Instrument *> &IVec,
+                                  uint64_t InstructionAddress) {
   assert(STI.getSchedModel().hasInstrSchedModel() &&
          "Itineraries are not yet supported!");
 
@@ -603,14 +604,15 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
     return *Descriptors[DKey];
   }
 
-  auto VDKey = std::make_pair(&MCI, SchedClassID);
+  auto VDKey = std::make_pair(InstructionAddress, SchedClassID);
   VariantDescriptors[VDKey] = std::move(ID);
   return *VariantDescriptors[VDKey];
 }
 
 Expected<const InstrDesc &>
 InstrBuilder::getOrCreateInstrDesc(const MCInst &MCI,
-                                   const SmallVector<Instrument *> &IVec) {
+                                   const SmallVector<Instrument *> &IVec,
+                                   uint64_t InstructionAddress) {
   // Cache lookup using SchedClassID from Instrumentation
   unsigned SchedClassID = IM.getSchedClassID(MCII, MCI, IVec);
 
@@ -620,19 +622,21 @@ InstrBuilder::getOrCreateInstrDesc(const MCInst &MCI,
 
   unsigned CPUID = STI.getSchedModel().getProcessorID();
   SchedClassID = STI.resolveVariantSchedClass(SchedClassID, &MCI, &MCII, CPUID);
-  auto VDKey = std::make_pair(&MCI, SchedClassID);
+  auto VDKey = std::make_pair(InstructionAddress, SchedClassID);
   if (VariantDescriptors.contains(VDKey))
     return *VariantDescriptors[VDKey];
 
-  return createInstrDescImpl(MCI, IVec);
+  return createInstrDescImpl(MCI, IVec, InstructionAddress);
 }
 
 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,
+                                uint64_t InstructionAddress) {
+  Expected<const InstrDesc &> DescOrErr =
+      getOrCreateInstrDesc(MCI, IVec, InstructionAddress);
   if (!DescOrErr)
     return DescOrErr.takeError();
   const InstrDesc &D = *DescOrErr;
diff --git a/llvm/tools/llvm-mca/llvm-mca.cpp b/llvm/tools/llvm-mca/llvm-mca.cpp
index 03d7d7944b9cd..d497d67efbd9b 100644
--- a/llvm/tools/llvm-mca/llvm-mca.cpp
+++ b/llvm/tools/llvm-mca/llvm-mca.cpp
@@ -611,8 +611,8 @@ int main(int argc, char **argv) {
       const SmallVector<mca::Instrument *> Instruments =
           InstrumentRegions.getActiveInstruments(Loc);
 
-      Expected<std::unique_ptr<mca::Instruction>> Inst =
-          IB.createInstruction(MCI, Instruments);
+      Expected<std::unique_ptr<mca::Instruction>> Inst = IB.createInstruction(
+          MCI, Instruments, reinterpret_cast<uint64_t>(&MCI));
       if (!Inst) {
         if (auto NewE = handleErrors(
                 Inst.takeError(),
diff --git a/llvm/unittests/tools/llvm-mca/MCATestBase.cpp b/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
index 4f444fae3d4c3..4032d5bdf4bd2 100644
--- a/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
+++ b/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
@@ -71,8 +71,8 @@ Error MCATestBase::runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
   const SmallVector<mca::Instrument *> Instruments;
   SmallVector<std::unique_ptr<mca::Instruction>> LoweredInsts;
   for (const auto &MCI : Insts) {
-    Expected<std::unique_ptr<mca::Instruction>> Inst =
-        IB.createInstruction(MCI, Instruments);
+    Expected<std::unique_ptr<mca::Instruction>> Inst = IB.createInstruction(
+        MCI, Instruments, reinterpret_cast<uint64_t>(&MCI));
     if (!Inst) {
       if (auto NewE =
               handleErrors(Inst.takeError(),
diff --git a/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp b/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
index 00a44dc1bab15..e3cb922462e90 100644
--- a/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
+++ b/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
@@ -40,7 +40,7 @@ TEST_F(X86TestBase, TestResumablePipeline) {
   for (unsigned i = 0U, E = MCIs.size(); i < E;) {
     for (unsigned TE = i + 7; i < TE && i < E; ++i) {
       Expected<std::unique_ptr<mca::Instruction>> InstOrErr =
-          IB.createInstruction(MCIs[i], Instruments);
+          IB.createInstruction(MCIs[i], Instruments, i);
       ASSERT_TRUE(bool(InstOrErr));
       ISM.addInst(std::move(InstOrErr.get()));
     }
@@ -132,7 +132,7 @@ TEST_F(X86TestBase, TestInstructionRecycling) {
   for (unsigned i = 0U, E = MCIs.size(); i < E;) {
     for (unsigned TE = i + 7; i < TE && i < E; ++i) {
       Expected<std::unique_ptr<mca::Instruction>> InstOrErr =
-          IB.createInstruction(MCIs[i], Instruments);
+          IB.createInstruction(MCIs[i], Instruments, i);
 
       if (!InstOrErr) {
         mca::Instruction *RecycledInst = nullptr;

>From 51a8a907fe1a809d808be67c0adf64624ff65b1c Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Tue, 21 May 2024 03:48:34 +0000
Subject: [PATCH 2/8] Fix formatting

---
 llvm/include/llvm/MCA/InstrBuilder.h | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index 67f9c736054e1..b625dfda76780 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -72,8 +72,7 @@ class InstrBuilder {
       Descriptors;
 
   // Key is the MCIInst and SchedClassID the describe the value InstrDesc
-  DenseMap<std::pair<uint64_t, unsigned>,
-           std::unique_ptr<const InstrDesc>>
+  DenseMap<std::pair<uint64_t, unsigned>, std::unique_ptr<const InstrDesc>>
       VariantDescriptors;
 
   bool FirstCallInst;

>From a0e3c4a22393fad74a0e3df522eb76baa279f5bf Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Thu, 30 May 2024 09:56:53 +0000
Subject: [PATCH 3/8] Add unittest

---
 .../tools/llvm-mca/X86/TestIncrementalMCA.cpp | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp b/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
index e3cb922462e90..48ee19f211c21 100644
--- a/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
+++ b/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
@@ -1,6 +1,8 @@
+#include "MCTargetDesc/X86MCTargetDesc.h"
 #include "Views/SummaryView.h"
 #include "X86TestBase.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/MC/MCInstBuilder.h"
 #include "llvm/MCA/CustomBehaviour.h"
 #include "llvm/MCA/IncrementalSourceMgr.h"
 #include "llvm/MCA/InstrBuilder.h"
@@ -8,6 +10,7 @@
 #include "llvm/Support/Format.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
+#include <memory>
 #include <unordered_map>
 
 using namespace llvm;
@@ -185,3 +188,51 @@ TEST_F(X86TestBase, TestInstructionRecycling) {
     ASSERT_EQ(*BV, *V) << "Value of '" << F << "' does not match";
   }
 }
+
+// Test that we do not depend upon the MCInst address for variant description
+// construction. This test creates two instructions that will use variant
+// description as they are both zeroing idioms, but write to different
+// registers. If the key used to access the variant instruction description is
+// the same between the descriptions (like the MCInst pointer), we will run into
+// an assertion failure due to the different writes. Within this test, setting
+// both instruction addresses to the same builder within the calls to the MCA
+// instruction builder will cause the problem.
+TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
+  mca::Context MCA(*MRI, *STI);
+
+  mca::IncrementalSourceMgr ISM;
+  // Empty CustomBehavioour.
+  auto CB = std::make_unique<mca::CustomBehaviour>(*STI, ISM, *MCII);
+
+  auto PO = getDefaultPipelineOptions();
+  auto P = MCA.createDefaultPipeline(PO, ISM, *CB);
+  ASSERT_TRUE(P);
+
+  auto IM = std::make_unique<mca::InstrumentManager>(*STI, *MCII);
+  mca::InstrBuilder IB(*STI, *MCII, *MRI, MCIA.get(), *IM);
+
+  const SmallVector<mca::Instrument *> Instruments;
+
+  MCInst InstructionToAdd;
+  InstructionToAdd = MCInstBuilder(X86::XOR64rr)
+                         .addReg(X86::RAX)
+                         .addReg(X86::RAX)
+                         .addReg(X86::RAX);
+  Expected<std::unique_ptr<mca::Instruction>> Instruction1OrErr =
+      IB.createInstruction(InstructionToAdd, Instruments, 0);
+  ASSERT_TRUE(static_cast<bool>(Instruction1OrErr));
+  ISM.addInst(std::move(Instruction1OrErr.get()));
+
+  InstructionToAdd = MCInstBuilder(X86::XORPSrr)
+                         .addReg(X86::XMM0)
+                         .addReg(X86::XMM0)
+                         .addReg(X86::XMM0);
+  Expected<std::unique_ptr<mca::Instruction>> Instruction2OrErr =
+      IB.createInstruction(InstructionToAdd, Instruments, 1);
+  ASSERT_TRUE(static_cast<bool>(Instruction2OrErr));
+  ISM.addInst(std::move(Instruction2OrErr.get()));
+
+  ISM.endOfStream();
+  Expected<unsigned> Cycles = P->run();
+  ASSERT_TRUE(static_cast<bool>(Cycles));
+}

>From db8a40cbd4f8cb7f4f25afb0314ad917bc08c690 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Thu, 30 May 2024 09:58:28 +0000
Subject: [PATCH 4/8] Address reviewer feedback

---
 llvm/include/llvm/MCA/InstrBuilder.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index b625dfda76780..6f227431ac2df 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -71,7 +71,8 @@ class InstrBuilder {
            std::unique_ptr<const InstrDesc>>
       Descriptors;
 
-  // Key is the MCIInst and SchedClassID the describe the value InstrDesc
+  // Key is the instruction address and SchedClassID the describe the value
+  // InstrDesc
   DenseMap<std::pair<uint64_t, unsigned>, std::unique_ptr<const InstrDesc>>
       VariantDescriptors;
 

>From c7fcab8c25f958ffd72cbbd773390af1028e03ac Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Fri, 21 Jun 2024 03:54:37 +0000
Subject: [PATCH 5/8] Address reviewer feedback

---
 llvm/include/llvm/MCA/InstrBuilder.h          | 18 +++-------
 llvm/lib/MCA/InstrBuilder.cpp                 | 36 +++++++++++++------
 llvm/tools/llvm-mca/llvm-mca.cpp              |  2 +-
 llvm/unittests/tools/llvm-mca/MCATestBase.cpp |  2 +-
 .../tools/llvm-mca/X86/TestIncrementalMCA.cpp | 16 ++++-----
 5 files changed, 40 insertions(+), 34 deletions(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index ee8d6c32e3c5f..1d94f71e36e27 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -14,6 +14,7 @@
 #ifndef LLVM_MCA_INSTRBUILDER_H
 #define LLVM_MCA_INSTRBUILDER_H
 
+#include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/MC/MCInstrAnalysis.h"
 #include "llvm/MC/MCInstrInfo.h"
@@ -73,7 +74,7 @@ class InstrBuilder {
 
   // Key is the instruction address and SchedClassID the describe the value
   // InstrDesc
-  DenseMap<std::pair<uint64_t, unsigned>, std::unique_ptr<const InstrDesc>>
+  DenseMap<std::pair<hash_code, unsigned>, std::unique_ptr<const InstrDesc>>
       VariantDescriptors;
 
   bool FirstCallInst;
@@ -84,11 +85,9 @@ class InstrBuilder {
   InstRecycleCallback InstRecycleCB;
 
   Expected<const InstrDesc &>
-  createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
-                      uint64_t InstructionAddress);
+  createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
   Expected<const InstrDesc &>
-  getOrCreateInstrDesc(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
-                       uint64_t InstructionAddress);
+  getOrCreateInstrDesc(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
 
   InstrBuilder(const InstrBuilder &) = delete;
   InstrBuilder &operator=(const InstrBuilder &) = delete;
@@ -113,15 +112,8 @@ class InstrBuilder {
   /// or null if there isn't any.
   void setInstRecycleCallback(InstRecycleCallback CB) { InstRecycleCB = CB; }
 
-  /// Create an MCA Instruction from a MC Instruction that contains all the
-  /// relevant state MCA needs for modeling. Variant instructions (e.g.,
-  /// register zeroing idioms) each need their own instruction description
-  /// which is uniqued based on InstructionAddress. This can be an actual
-  /// address or something unique per instruction like a loop iteration
-  /// variable, but it must uniquely identify the instruction being passed in.
   Expected<std::unique_ptr<Instruction>>
-  createInstruction(const MCInst &MCI, const SmallVector<Instrument *> &IVec,
-                    uint64_t InstructionAddress);
+  createInstruction(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
 };
 } // namespace mca
 } // namespace llvm
diff --git a/llvm/lib/MCA/InstrBuilder.cpp b/llvm/lib/MCA/InstrBuilder.cpp
index 623fb45e0cb0f..0fea555806856 100644
--- a/llvm/lib/MCA/InstrBuilder.cpp
+++ b/llvm/lib/MCA/InstrBuilder.cpp
@@ -14,6 +14,7 @@
 #include "llvm/MCA/InstrBuilder.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/Support/Debug.h"
@@ -504,6 +505,24 @@ void InstrBuilder::populateReads(InstrDesc &ID, const MCInst &MCI,
   ID.Reads.resize(CurrentUse);
 }
 
+hash_code hashMCOperand(const MCOperand &MCO) {
+  hash_code TypeHash = hash_combine(MCO.isReg(), MCO.isImm(), MCO.isSFPImm(),
+                                    MCO.isDFPImm(), MCO.isExpr(), MCO.isInst());
+  if (MCO.isReg())
+    return hash_combine(TypeHash, MCO.getReg());
+
+  return TypeHash;
+}
+
+hash_code hashMCInst(const MCInst &MCI) {
+  hash_code InstructionHash = hash_combine(MCI.getOpcode(), MCI.getFlags());
+  for (unsigned I = 0; I < MCI.getNumOperands(); ++I) {
+    InstructionHash =
+        hash_combine(InstructionHash, hashMCOperand(MCI.getOperand(I)));
+  }
+  return InstructionHash;
+}
+
 Error InstrBuilder::verifyInstrDesc(const InstrDesc &ID,
                                     const MCInst &MCI) const {
   if (ID.NumMicroOps != 0)
@@ -523,8 +542,7 @@ Error InstrBuilder::verifyInstrDesc(const InstrDesc &ID,
 
 Expected<const InstrDesc &>
 InstrBuilder::createInstrDescImpl(const MCInst &MCI,
-                                  const SmallVector<Instrument *> &IVec,
-                                  uint64_t InstructionAddress) {
+                                  const SmallVector<Instrument *> &IVec) {
   assert(STI.getSchedModel().hasInstrSchedModel() &&
          "Itineraries are not yet supported!");
 
@@ -606,15 +624,14 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
     return *Descriptors[DKey];
   }
 
-  auto VDKey = std::make_pair(InstructionAddress, SchedClassID);
+  auto VDKey = std::make_pair(hashMCInst(MCI), SchedClassID);
   VariantDescriptors[VDKey] = std::move(ID);
   return *VariantDescriptors[VDKey];
 }
 
 Expected<const InstrDesc &>
 InstrBuilder::getOrCreateInstrDesc(const MCInst &MCI,
-                                   const SmallVector<Instrument *> &IVec,
-                                   uint64_t InstructionAddress) {
+                                   const SmallVector<Instrument *> &IVec) {
   // Cache lookup using SchedClassID from Instrumentation
   unsigned SchedClassID = IM.getSchedClassID(MCII, MCI, IVec);
 
@@ -624,21 +641,20 @@ InstrBuilder::getOrCreateInstrDesc(const MCInst &MCI,
 
   unsigned CPUID = STI.getSchedModel().getProcessorID();
   SchedClassID = STI.resolveVariantSchedClass(SchedClassID, &MCI, &MCII, CPUID);
-  auto VDKey = std::make_pair(InstructionAddress, SchedClassID);
+  auto VDKey = std::make_pair(hashMCInst(MCI), SchedClassID);
   if (VariantDescriptors.contains(VDKey))
     return *VariantDescriptors[VDKey];
 
-  return createInstrDescImpl(MCI, IVec, InstructionAddress);
+  return createInstrDescImpl(MCI, IVec);
 }
 
 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,
-                                uint64_t InstructionAddress) {
+                                const SmallVector<Instrument *> &IVec) {
   Expected<const InstrDesc &> DescOrErr =
-      getOrCreateInstrDesc(MCI, IVec, InstructionAddress);
+      getOrCreateInstrDesc(MCI, IVec);
   if (!DescOrErr)
     return DescOrErr.takeError();
   const InstrDesc &D = *DescOrErr;
diff --git a/llvm/tools/llvm-mca/llvm-mca.cpp b/llvm/tools/llvm-mca/llvm-mca.cpp
index 9190cd7f828ff..ca961771f0fb4 100644
--- a/llvm/tools/llvm-mca/llvm-mca.cpp
+++ b/llvm/tools/llvm-mca/llvm-mca.cpp
@@ -617,7 +617,7 @@ int main(int argc, char **argv) {
           InstrumentRegions.getActiveInstruments(Loc);
 
       Expected<std::unique_ptr<mca::Instruction>> Inst = IB.createInstruction(
-          MCI, Instruments, reinterpret_cast<uint64_t>(&MCI));
+          MCI, Instruments);
       if (!Inst) {
         if (auto NewE = handleErrors(
                 Inst.takeError(),
diff --git a/llvm/unittests/tools/llvm-mca/MCATestBase.cpp b/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
index 66192436dca1d..6e03ff4edd03b 100644
--- a/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
+++ b/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
@@ -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, reinterpret_cast<uint64_t>(&MCI));
+        MCI, Instruments);
     if (!Inst) {
       if (auto NewE =
               handleErrors(Inst.takeError(),
diff --git a/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp b/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
index 94998144f9193..1a14c687295ca 100644
--- a/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
+++ b/llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp
@@ -43,7 +43,7 @@ TEST_F(X86TestBase, TestResumablePipeline) {
   for (unsigned i = 0U, E = MCIs.size(); i < E;) {
     for (unsigned TE = i + 7; i < TE && i < E; ++i) {
       Expected<std::unique_ptr<mca::Instruction>> InstOrErr =
-          IB.createInstruction(MCIs[i], Instruments, i);
+          IB.createInstruction(MCIs[i], Instruments);
       ASSERT_TRUE(bool(InstOrErr));
       ISM.addInst(std::move(InstOrErr.get()));
     }
@@ -135,7 +135,7 @@ TEST_F(X86TestBase, TestInstructionRecycling) {
   for (unsigned i = 0U, E = MCIs.size(); i < E;) {
     for (unsigned TE = i + 7; i < TE && i < E; ++i) {
       Expected<std::unique_ptr<mca::Instruction>> InstOrErr =
-          IB.createInstruction(MCIs[i], Instruments, i);
+          IB.createInstruction(MCIs[i], Instruments);
 
       if (!InstOrErr) {
         mca::Instruction *RecycledInst = nullptr;
@@ -194,14 +194,12 @@ TEST_F(X86TestBase, TestInstructionRecycling) {
 // description as they are both zeroing idioms, but write to different
 // registers. If the key used to access the variant instruction description is
 // the same between the descriptions (like the MCInst pointer), we will run into
-// an assertion failure due to the different writes. Within this test, setting
-// both instruction addresses to the same builder within the calls to the MCA
-// instruction builder will cause the problem.
+// an assertion failure due to the different writes.
 TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
   mca::Context MCA(*MRI, *STI);
 
   mca::IncrementalSourceMgr ISM;
-  // Empty CustomBehavioour.
+  // Empty CustomBehaviour.
   auto CB = std::make_unique<mca::CustomBehaviour>(*STI, ISM, *MCII);
 
   auto PO = getDefaultPipelineOptions();
@@ -209,7 +207,7 @@ TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
   ASSERT_TRUE(P);
 
   auto IM = std::make_unique<mca::InstrumentManager>(*STI, *MCII);
-  mca::InstrBuilder IB(*STI, *MCII, *MRI, MCIA.get(), *IM);
+  mca::InstrBuilder IB(*STI, *MCII, *MRI, MCIA.get(), *IM, 100);
 
   const SmallVector<mca::Instrument *> Instruments;
 
@@ -219,7 +217,7 @@ TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
                          .addReg(X86::RAX)
                          .addReg(X86::RAX);
   Expected<std::unique_ptr<mca::Instruction>> Instruction1OrErr =
-      IB.createInstruction(InstructionToAdd, Instruments, 0);
+      IB.createInstruction(InstructionToAdd, Instruments);
   ASSERT_TRUE(static_cast<bool>(Instruction1OrErr));
   ISM.addInst(std::move(Instruction1OrErr.get()));
 
@@ -228,7 +226,7 @@ TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
                          .addReg(X86::XMM0)
                          .addReg(X86::XMM0);
   Expected<std::unique_ptr<mca::Instruction>> Instruction2OrErr =
-      IB.createInstruction(InstructionToAdd, Instruments, 1);
+      IB.createInstruction(InstructionToAdd, Instruments);
   ASSERT_TRUE(static_cast<bool>(Instruction2OrErr));
   ISM.addInst(std::move(Instruction2OrErr.get()));
 

>From 94f5e2012a707de7ffe7e9300c12b8a12cd0dc64 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Fri, 21 Jun 2024 03:58:08 +0000
Subject: [PATCH 6/8] Fix formatting changes

---
 llvm/include/llvm/MCA/InstrBuilder.h          | 3 ++-
 llvm/lib/MCA/InstrBuilder.cpp                 | 3 +--
 llvm/tools/llvm-mca/llvm-mca.cpp              | 4 ++--
 llvm/unittests/tools/llvm-mca/MCATestBase.cpp | 4 ++--
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index 1d94f71e36e27..d15e2557339c0 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -87,7 +87,8 @@ class InstrBuilder {
   Expected<const InstrDesc &>
   createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
   Expected<const InstrDesc &>
-  getOrCreateInstrDesc(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
+  getOrCreateInstrDesc(const MCInst &MCI,
+                       const SmallVector<Instrument *> &IVec);
 
   InstrBuilder(const InstrBuilder &) = delete;
   InstrBuilder &operator=(const InstrBuilder &) = delete;
diff --git a/llvm/lib/MCA/InstrBuilder.cpp b/llvm/lib/MCA/InstrBuilder.cpp
index 0fea555806856..8222d4a2d8e5a 100644
--- a/llvm/lib/MCA/InstrBuilder.cpp
+++ b/llvm/lib/MCA/InstrBuilder.cpp
@@ -653,8 +653,7 @@ 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 = getOrCreateInstrDesc(MCI, IVec);
   if (!DescOrErr)
     return DescOrErr.takeError();
   const InstrDesc &D = *DescOrErr;
diff --git a/llvm/tools/llvm-mca/llvm-mca.cpp b/llvm/tools/llvm-mca/llvm-mca.cpp
index ca961771f0fb4..cc5d4f5fa05de 100644
--- a/llvm/tools/llvm-mca/llvm-mca.cpp
+++ b/llvm/tools/llvm-mca/llvm-mca.cpp
@@ -616,8 +616,8 @@ int main(int argc, char **argv) {
       const SmallVector<mca::Instrument *> Instruments =
           InstrumentRegions.getActiveInstruments(Loc);
 
-      Expected<std::unique_ptr<mca::Instruction>> Inst = IB.createInstruction(
-          MCI, Instruments);
+      Expected<std::unique_ptr<mca::Instruction>> Inst =
+          IB.createInstruction(MCI, Instruments);
       if (!Inst) {
         if (auto NewE = handleErrors(
                 Inst.takeError(),
diff --git a/llvm/unittests/tools/llvm-mca/MCATestBase.cpp b/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
index 6e03ff4edd03b..4a39f5e663f23 100644
--- a/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
+++ b/llvm/unittests/tools/llvm-mca/MCATestBase.cpp
@@ -71,8 +71,8 @@ Error MCATestBase::runBaselineMCA(json::Object &Result, ArrayRef<MCInst> Insts,
   const SmallVector<mca::Instrument *> Instruments;
   SmallVector<std::unique_ptr<mca::Instruction>> LoweredInsts;
   for (const auto &MCI : Insts) {
-    Expected<std::unique_ptr<mca::Instruction>> Inst = IB.createInstruction(
-        MCI, Instruments);
+    Expected<std::unique_ptr<mca::Instruction>> Inst =
+        IB.createInstruction(MCI, Instruments);
     if (!Inst) {
       if (auto NewE =
               handleErrors(Inst.takeError(),

>From 508b09059e18fc580071f56a7f16eb7a7f2c72eb Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Fri, 21 Jun 2024 03:59:18 +0000
Subject: [PATCH 7/8] Cleanup comment

---
 llvm/include/llvm/MCA/InstrBuilder.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index d15e2557339c0..9bfb8c0d2fca7 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -72,8 +72,8 @@ class InstrBuilder {
            std::unique_ptr<const InstrDesc>>
       Descriptors;
 
-  // Key is the instruction address and SchedClassID the describe the value
-  // InstrDesc
+  // Key is a hash of the MCInstruction and a SchedClassID that describe the
+  // value InstrDesc
   DenseMap<std::pair<hash_code, unsigned>, std::unique_ptr<const InstrDesc>>
       VariantDescriptors;
 

>From 1df4c0b5251018cd5eecccfd414186f302a1c7f1 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Fri, 28 Jun 2024 17:10:39 +0000
Subject: [PATCH 8/8] Fix bug

---
 llvm/include/llvm/MCA/InstrBuilder.h |  1 +
 llvm/lib/MCA/InstrBuilder.cpp        | 43 +++++++++++++++++++++-------
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/MCA/InstrBuilder.h b/llvm/include/llvm/MCA/InstrBuilder.h
index 9bfb8c0d2fca7..00b19a65dc166 100644
--- a/llvm/include/llvm/MCA/InstrBuilder.h
+++ b/llvm/include/llvm/MCA/InstrBuilder.h
@@ -84,6 +84,7 @@ class InstrBuilder {
   using InstRecycleCallback = std::function<Instruction *(const InstrDesc &)>;
   InstRecycleCallback InstRecycleCB;
 
+  Expected<unsigned> getVariantSchedClassID(const MCInst &MCI, unsigned SchedClassID);
   Expected<const InstrDesc &>
   createInstrDescImpl(const MCInst &MCI, const SmallVector<Instrument *> &IVec);
   Expected<const InstrDesc &>
diff --git a/llvm/lib/MCA/InstrBuilder.cpp b/llvm/lib/MCA/InstrBuilder.cpp
index 8222d4a2d8e5a..32b20d758ee70 100644
--- a/llvm/lib/MCA/InstrBuilder.cpp
+++ b/llvm/lib/MCA/InstrBuilder.cpp
@@ -540,6 +540,22 @@ Error InstrBuilder::verifyInstrDesc(const InstrDesc &ID,
   return make_error<InstructionError<MCInst>>(std::string(Message), MCI);
 }
 
+Expected<unsigned> InstrBuilder::getVariantSchedClassID(const MCInst &MCI,
+                                                        unsigned SchedClassID) {
+  const MCSchedModel &SM = STI.getSchedModel();
+  unsigned CPUID = SM.getProcessorID();
+  while (SchedClassID && SM.getSchedClassDesc(SchedClassID)->isVariant())
+    SchedClassID =
+        STI.resolveVariantSchedClass(SchedClassID, &MCI, &MCII, CPUID);
+
+  if (!SchedClassID) {
+    return make_error<InstructionError<MCInst>>(
+        "unable to resolve scheduling class for write variant.", MCI);
+  }
+
+  return SchedClassID;
+}
+
 Expected<const InstrDesc &>
 InstrBuilder::createInstrDescImpl(const MCInst &MCI,
                                   const SmallVector<Instrument *> &IVec) {
@@ -558,15 +574,13 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
 
   // Try to solve variant scheduling classes.
   if (IsVariant) {
-    unsigned CPUID = SM.getProcessorID();
-    while (SchedClassID && SM.getSchedClassDesc(SchedClassID)->isVariant())
-      SchedClassID =
-          STI.resolveVariantSchedClass(SchedClassID, &MCI, &MCII, CPUID);
-
-    if (!SchedClassID) {
-      return make_error<InstructionError<MCInst>>(
-          "unable to resolve scheduling class for write variant.", MCI);
+    Expected<unsigned> VariantSchedClassIDOrErr =
+        getVariantSchedClassID(MCI, SchedClassID);
+    if (!VariantSchedClassIDOrErr) {
+      return VariantSchedClassIDOrErr.takeError();
     }
+
+    SchedClassID = *VariantSchedClassIDOrErr;
   }
 
   // Check if this instruction is supported. Otherwise, report an error.
@@ -625,6 +639,9 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
   }
 
   auto VDKey = std::make_pair(hashMCInst(MCI), SchedClassID);
+  assert(
+      !VariantDescriptors.contains(VDKey) &&
+      "Expected VariantDescriptors to not already have a value for this key.");
   VariantDescriptors[VDKey] = std::move(ID);
   return *VariantDescriptors[VDKey];
 }
@@ -639,8 +656,14 @@ InstrBuilder::getOrCreateInstrDesc(const MCInst &MCI,
   if (Descriptors.find_as(DKey) != Descriptors.end())
     return *Descriptors[DKey];
 
-  unsigned CPUID = STI.getSchedModel().getProcessorID();
-  SchedClassID = STI.resolveVariantSchedClass(SchedClassID, &MCI, &MCII, CPUID);
+  Expected<unsigned> VariantSchedClassIDOrErr =
+      getVariantSchedClassID(MCI, SchedClassID);
+  if (!VariantSchedClassIDOrErr) {
+    return VariantSchedClassIDOrErr.takeError();
+  }
+
+  SchedClassID = *VariantSchedClassIDOrErr;
+
   auto VDKey = std::make_pair(hashMCInst(MCI), SchedClassID);
   if (VariantDescriptors.contains(VDKey))
     return *VariantDescriptors[VDKey];



More information about the llvm-commits mailing list