3 files changed, 15 insertions, 11 deletions

diff --git a/llvm/tools/llvm-mca/include/Stages/EntryStage.h b/llvm/tools/llvm-mca/include/Stages/EntryStage.h
index ed9f32f5d53..21b8331e95b 100644
--- a/llvm/tools/llvm-mca/include/Stages/EntryStage.h
+++ b/llvm/tools/llvm-mca/include/Stages/EntryStage.h
@@ -19,16 +19,16 @@
 
 #include "SourceMgr.h"
 #include "Stages/Stage.h"
-#include <map>
+#include "llvm/ADT/SmallVector.h"
 
 namespace llvm {
 namespace mca {
 
 class EntryStage final : public Stage {
   InstRef CurrentInstruction;
-  using InstMap = std::map<unsigned, std::unique_ptr<Instruction>>;
-  InstMap Instructions;
+  SmallVector<std::unique_ptr<Instruction>, 16> Instructions;
   SourceMgr &SM;
+  unsigned NumRetired;
 
   // Updates the program counter, and sets 'CurrentInstruction'.
   void getNextInstruction();
@@ -37,7 +37,7 @@ class EntryStage final : public Stage {
   EntryStage &operator=(const EntryStage &Other) = delete;
 
 public:
-  EntryStage(SourceMgr &SM) : CurrentInstruction(), SM(SM) {}
+  EntryStage(SourceMgr &SM) : CurrentInstruction(), SM(SM), NumRetired(0) { }
 
   bool isAvailable(const InstRef &IR) const override;
   bool hasWorkToComplete() const override;
diff --git a/llvm/tools/llvm-mca/lib/HardwareUnits/RetireControlUnit.cpp b/llvm/tools/llvm-mca/lib/HardwareUnits/RetireControlUnit.cpp
index 0456e1d7a5b..bd7b411af11 100644
--- a/llvm/tools/llvm-mca/lib/HardwareUnits/RetireControlUnit.cpp
+++ b/llvm/tools/llvm-mca/lib/HardwareUnits/RetireControlUnit.cpp
@@ -60,9 +60,10 @@ const RetireControlUnit::RUToken &RetireControlUnit::peekCurrentToken() const {
 }
 
 void RetireControlUnit::consumeCurrentToken() {
-  const RetireControlUnit::RUToken &Current = peekCurrentToken();
+  RetireControlUnit::RUToken &Current = Queue[CurrentInstructionSlotIdx];
   assert(Current.NumSlots && "Reserved zero slots?");
   assert(Current.IR && "Invalid RUToken in the RCU queue.");
+  Current.IR.getInstruction()->retire();
 
   // Update the slot index to be the next item in the circular queue.
   CurrentInstructionSlotIdx += Current.NumSlots;
diff --git a/llvm/tools/llvm-mca/lib/Stages/EntryStage.cpp b/llvm/tools/llvm-mca/lib/Stages/EntryStage.cpp
index 929d51e824c..f552132cac6 100644
--- a/llvm/tools/llvm-mca/lib/Stages/EntryStage.cpp
+++ b/llvm/tools/llvm-mca/lib/Stages/EntryStage.cpp
@@ -34,7 +34,7 @@ void EntryStage::getNextInstruction() {
   SourceRef SR = SM.peekNext();
   std::unique_ptr<Instruction> Inst = llvm::make_unique<Instruction>(SR.second);
   CurrentInstruction = InstRef(SR.first, Inst.get());
-  Instructions[SR.first] = std::move(Inst);
+  Instructions.emplace_back(std::move(Inst));
   SM.updateNext();
 }
 
@@ -57,14 +57,17 @@ llvm::Error EntryStage::cycleStart() {
 
 llvm::Error EntryStage::cycleEnd() {
   // Find the first instruction which hasn't been retired.
-  const InstMap::iterator It =
-      llvm::find_if(Instructions, [](const InstMap::value_type &KeyValuePair) {
-        return !KeyValuePair.second->isRetired();
-      });
+  auto Range = make_range(&Instructions[NumRetired], Instructions.end());
+  auto It = find_if(Range, [](const std::unique_ptr<Instruction> &I) {
+    return !I->isRetired();
+  });
 
+  NumRetired = std::distance(Instructions.begin(), It);
   // Erase instructions up to the first that hasn't been retired.
-  if (It != Instructions.begin())
+  if ((NumRetired * 2) >= Instructions.size()) {
     Instructions.erase(Instructions.begin(), It);
+    NumRetired = 0;
+  }
 
   return llvm::ErrorSuccess();
 }