[MCA] Parameterize variant scheduling classes by hash (#92849)

This patch looks up variant scheduling classes using a hash of the
instruction. Keying by the pointer 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.).

Author: boomanaiden154

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"
@@ -71,9 +72,9 @@ class InstrBuilder {
            std::unique_ptr<const InstrDesc>>
       Descriptors;
 
-  // Key is the MCIInst and SchedClassID the describe the value InstrDesc
-  DenseMap<std::pair<const MCInst *, unsigned>,
-           std::unique_ptr<const 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;
 
   bool FirstCallInst;
@@ -83,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 &>

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)
@@ -521,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) {
@@ -539,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.
@@ -605,7 +638,10 @@ InstrBuilder::createInstrDescImpl(const MCInst &MCI,
     return *Descriptors[DKey];
   }
 
-  auto VDKey = std::make_pair(&MCI, SchedClassID);
+  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];
 }
@@ -620,9 +656,15 @@ 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);
-  auto VDKey = std::make_pair(&MCI, SchedClassID);
+  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];
 

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

@@ -1,13 +1,16 @@
+#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"
 #include "llvm/MCA/Pipeline.h"
 #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,49 @@ 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.
+TEST_F(X86TestBase, TestVariantInstructionsSameAddress) {
+  mca::Context MCA(*MRI, *STI);
+
+  mca::IncrementalSourceMgr ISM;
+  // Empty CustomBehaviour.
+  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, 100);
+
+  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);
+  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);
+  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));
+}