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

[Andrea Di Biagio via llvm-commits]

Author: Andrea Di Biagio
Date: 2021-04-20T13:30:45+01:00
New Revision: 2226d21896d6d30d51e13385361ea0706ee9d9fb

URL: https://github.com/llvm/llvm-project/commit/2226d21896d6d30d51e13385361ea0706ee9d9fb
DIFF: https://github.com/llvm/llvm-project/commit/2226d21896d6d30d51e13385361ea0706ee9d9fb.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
    llvm/lib/MCA/HardwareUnits/LSUnit.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h b/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
index 0f1fac55af4f1..7eddd067aa0cc 100644
--- a/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
+++ b/llvm/include/llvm/MCA/HardwareUnits/LSUnit.h
@@ -160,11 +160,16 @@ class MemoryGroup {
       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 4594368fc0e96..07be7b077bc9d 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);
 }