[MCA][LSUnit] Fix a potential use after free in the logic that updates memory groups.

Make sure that the `CriticalMemoryInstruction` of a memory group is invalidated
if it references an already executed instruction.  This avoids a potential
use-after-free if the critical memory info becomes stale, and the value is
read after the instruction has executed.

2 files changed, 7 insertions, 2 deletions

diff --git a/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h b/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
index 0f1fac55af4f..7eddd067aa0c 100644
--- a/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
+++ b/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
@@ -160,11 +160,16 @@ public:
       MG->onGroupIssued(CriticalMemoryInstruction, true);
   }
 
-  void onInstructionExecuted() {
+  void onInstructionExecuted(const InstRef &IR) {
     assert(isReady() && !isExecuted() && "Invalid internal state!");
     --NumExecuting;
     ++NumExecuted;
 
+    if (CriticalMemoryInstruction &&
+        CriticalMemoryInstruction.getSourceIndex() == IR.getSourceIndex()) {
+      CriticalMemoryInstruction.invalidate();
+    }
+
     if (!isExecuted())
       return;
 
diff --git a/llvm/lib/MCA/HardwareUnits/LSUnit.cpp b/llvm/lib/MCA/HardwareUnits/LSUnit.cpp
index 4594368fc0e9..07be7b077bc9 100644
--- a/llvm/lib/MCA/HardwareUnits/LSUnit.cpp
+++ b/llvm/lib/MCA/HardwareUnits/LSUnit.cpp
@@ -205,7 +205,7 @@ void LSUnitBase::onInstructionExecuted(const InstRef &IR) {
   unsigned GroupID = IR.getInstruction()->getLSUTokenID();
   auto It = Groups.find(GroupID);
   assert(It != Groups.end() && "Instruction not dispatched to the LS unit");
-  It->second->onInstructionExecuted();
+  It->second->onInstructionExecuted(IR);
   if (It->second->isExecuted())
     Groups.erase(It);
 }