From ed90e9bf1dbd0f0f42f3b6656f336a15582b1c5b Mon Sep 17 00:00:00 2001
From: andreroesti <an.roesti@gmail.com>
Date: Sun, 1 Dec 2024 12:58:47 -0800
Subject: [PATCH] add capacity limit and LRU policy for naive branch predictor

---
 CustomHWUnits/NaiveBranchPredictorUnit.cpp | 23 ++++++++++++++++++++--
 CustomHWUnits/NaiveBranchPredictorUnit.h   | 14 +++++++++++--
 MCAWorker.cpp                              |  9 +++++++--
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/CustomHWUnits/NaiveBranchPredictorUnit.cpp b/CustomHWUnits/NaiveBranchPredictorUnit.cpp
index 4d344fc..6885317 100644
--- a/CustomHWUnits/NaiveBranchPredictorUnit.cpp
+++ b/CustomHWUnits/NaiveBranchPredictorUnit.cpp
@@ -5,12 +5,31 @@ namespace llvm {
 namespace mcad {
 
 void NaiveBranchPredictorUnit::recordTakenBranch(MDInstrAddr instrAddr, BranchDirection nextInstrDirection) {
-    branchHistory[instrAddr] = nextInstrDirection;
+    // If no entry exists yet, add one
+    if(branchHistory.find(instrAddr) == branchHistory.end()) {
+        branchHistory[instrAddr] = {};
+    }
+    // Update the entry
+    branchHistory[instrAddr].lastDirection = nextInstrDirection;
+    branchHistory[instrAddr].lastUse = nAccesses;
+
+    // Evict the least recently used entry if we are at capacity
+    if(branchHistory.size() >= historyCapacity) {
+        auto min_it = branchHistory.begin();
+        for(auto it = branchHistory.begin(); it != branchHistory.end(); it++) {
+            if(it->second.lastUse < min_it->second.lastUse) {
+                min_it = it;
+            }
+        }
+        branchHistory.erase(min_it);
+    }
 }
 
 AbstractBranchPredictorUnit::BranchDirection NaiveBranchPredictorUnit::predictBranch(MDInstrAddr instrAddr) {
+    ++nAccesses;
     if(branchHistory.find(instrAddr) != branchHistory.end()) {
-        return branchHistory[instrAddr];
+        branchHistory[instrAddr].lastUse = nAccesses;
+        return branchHistory[instrAddr].lastDirection;
     }
     // We have no history on this; predict a fall-through branch
     return NOT_TAKEN;
diff --git a/CustomHWUnits/NaiveBranchPredictorUnit.h b/CustomHWUnits/NaiveBranchPredictorUnit.h
index 323d55e..a4016dc 100644
--- a/CustomHWUnits/NaiveBranchPredictorUnit.h
+++ b/CustomHWUnits/NaiveBranchPredictorUnit.h
@@ -8,11 +8,21 @@ namespace llvm {
 namespace mcad {
 
 class NaiveBranchPredictorUnit : public AbstractBranchPredictorUnit {
+
+    struct HistoryEntry {
+        AbstractBranchPredictorUnit::BranchDirection lastDirection;
+        unsigned lastUse;
+    };
+
     unsigned mispredictionPenalty;
-    std::map<MDInstrAddr, AbstractBranchPredictorUnit::BranchDirection> branchHistory = {};
+    unsigned historyCapacity;
+    unsigned nAccesses;
+    std::map<MDInstrAddr, struct HistoryEntry> branchHistory = {};
 
 public:
-    NaiveBranchPredictorUnit(unsigned mispredictionPenalty = 20) : mispredictionPenalty(mispredictionPenalty) {};
+    NaiveBranchPredictorUnit(unsigned mispredictionPenalty, unsigned branchHistoryTableSize)
+     : mispredictionPenalty(mispredictionPenalty),
+       historyCapacity(branchHistoryTableSize) { };
 
     void recordTakenBranch(MDInstrAddr instrAddr, BranchDirection nextInstrDirection) override;
 
diff --git a/MCAWorker.cpp b/MCAWorker.cpp
index 2387232..f50e405 100644
--- a/MCAWorker.cpp
+++ b/MCAWorker.cpp
@@ -109,6 +109,11 @@ static cl::opt<unsigned>
                            cl::desc("Delay (# cycles) added to the fetch stage of the next instruction after a branch misprediction"),
                            cl::init(20U));
 
+static cl::opt<unsigned>
+  BranchHistoryTableSize("bht-size",
+                         cl::desc("Size of the simulated branch history table for branch prediction"),
+                         cl::init(10U));
+
 void BrokerFacade::setBroker(std::unique_ptr<Broker> &&B) {
   Worker.TheBroker = std::move(B);
 }
@@ -187,7 +192,7 @@ std::unique_ptr<mca::Pipeline> MCAWorker::createDefaultPipeline() {
                                           MCAPO.StoreQueueSize,
                                           MCAPO.AssumeNoAlias, &MDRegistry);
   auto HWS = std::make_unique<Scheduler>(SM, *LSU);
-  auto BPU = std::make_unique<NaiveBranchPredictorUnit>(BranchMispredictionDelay);
+  auto BPU = std::make_unique<NaiveBranchPredictorUnit>(BranchMispredictionDelay, BranchHistoryTableSize);
 
   // Create the pipeline stages.
   auto Fetch = std::make_unique<EntryStage>(SrcMgr);
@@ -232,7 +237,7 @@ std::unique_ptr<mca::Pipeline> MCAWorker::createInOrderPipeline() {
   auto LSU = std::make_unique<MCADLSUnit>(SM, MCAPO.LoadQueueSize,
                                           MCAPO.StoreQueueSize,
                                           MCAPO.AssumeNoAlias, &MDRegistry);
-  auto BPU = std::make_unique<NaiveBranchPredictorUnit>(20);
+  auto BPU = std::make_unique<NaiveBranchPredictorUnit>(BranchMispredictionDelay, BranchHistoryTableSize);
 
   // Create the pipeline stages.
   auto Entry = std::make_unique<EntryStage>(SrcMgr);