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

via llvm-commits llvm-commits at lists.llvm.org
Mon May 20 20:41:20 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-tools-llvm-mca

Author: Aiden Grossman (boomanaiden154)

<details>
<summary>Changes</summary>

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.).

CC: @<!-- -->mtrofin 

---
Full diff: https://github.com/llvm/llvm-project/pull/92849.diff


5 Files Affected:

- (modified) llvm/include/llvm/MCA/InstrBuilder.h (+13-5) 
- (modified) llvm/lib/MCA/InstrBuilder.cpp (+11-7) 
- (modified) llvm/tools/llvm-mca/llvm-mca.cpp (+2-2) 
- (modified) llvm/unittests/tools/llvm-mca/MCATestBase.cpp (+2-2) 
- (modified) llvm/unittests/tools/llvm-mca/X86/TestIncrementalMCA.cpp (+2-2) 


``````````diff
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;

``````````

</details>


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


More information about the llvm-commits mailing list