Add naive branch predictor unit (#29)

Author: andrej

CMakeLists.txt

@@ -97,6 +97,7 @@ set(_MCAVIEWS_SOURCE_FILES
 set(_CUSTOMHW_SOURCE_FILES
     CustomHWUnits/CustomSourceMgr.cpp
     CustomHWUnits/MCADLSUnit.cpp
+    CustomHWUnits/NaiveBranchPredictorUnit.cpp
    )
 
 set(_CUSTOM_STAGES_SOURCE_FILES
@@ -118,6 +119,7 @@ set(_SOURCE_FILES
     ${_CUSTOM_STAGES_SOURCE_FILES}
     ${_BROKERS_SOURCE_FILES}
     MCAWorker.cpp
+    MetadataCategories.cpp
     PipelinePrinter.cpp
     )
 

CustomHWUnits/AbstractBranchPredictorUnit.h

@@ -0,0 +1,26 @@
+#ifndef LLVM_MCAD_ABSTRACT_BRANCH_PREDICTOR_UNIT_H
+#define LLVM_MCAD_ABSTRACT_BRANCH_PREDICTOR_UNIT_H
+
+#include <optional>
+#include "llvm/MCA/Instruction.h"
+#include "llvm/MCA/HardwareUnits/HardwareUnit.h"
+#include "MetadataRegistry.h"
+#include "MetadataCategories.h"
+
+namespace llvm {
+namespace mcad {
+
+class AbstractBranchPredictorUnit : public llvm::mca::HardwareUnit {
+
+public:
+    ~AbstractBranchPredictorUnit() {}
+    virtual void recordTakenBranch(MDInstrAddr IA, MDInstrAddr destAddr) = 0;
+    virtual MDInstrAddr predictBranch(MDInstrAddr IA) = 0;
+    virtual unsigned getMispredictionPenalty() = 0;
+
+};
+
+}
+}
+
+#endif

CustomHWUnits/NaiveBranchPredictorUnit.cpp

@@ -0,0 +1,22 @@
+#include <map>
+#include "CustomHWUnits/NaiveBranchPredictorUnit.h"
+
+namespace llvm {
+namespace mcad {
+
+void NaiveBranchPredictorUnit::recordTakenBranch(MDInstrAddr IA, MDInstrAddr destAddr) {
+    branchHistory[IA] = destAddr;
+}
+
+MDInstrAddr NaiveBranchPredictorUnit::predictBranch(MDInstrAddr IA) {
+    if(branchHistory.find(IA) != branchHistory.end()) {
+        return branchHistory[IA];
+    }
+    // We have no history on this; predict a fall-through branch
+    // FIXME: fix this to use actual branch instruction size, which is likely
+    // larger than one byte.
+    return MDInstrAddr { IA.addr + 1 };
+}
+
+}
+}

CustomHWUnits/NaiveBranchPredictorUnit.h

@@ -0,0 +1,28 @@
+#ifndef LLVM_MCAD_NAIVE_BRANCH_PREDICTOR_UNIT_H
+#define LLVM_MCAD_NAIVE_BRANCH_PREDICTOR_UNIT_H
+
+#include <map>
+#include "CustomHWUnits/AbstractBranchPredictorUnit.h"
+
+namespace llvm {
+namespace mcad {
+
+class NaiveBranchPredictorUnit : public AbstractBranchPredictorUnit {
+    unsigned mispredictionPenalty;
+    std::map<MDInstrAddr, MDInstrAddr> branchHistory = {};
+
+public:
+    NaiveBranchPredictorUnit(unsigned mispredictionPenalty = 20) : mispredictionPenalty(mispredictionPenalty) {};
+
+    void recordTakenBranch(MDInstrAddr IA, MDInstrAddr destAddr) override;
+    MDInstrAddr predictBranch(MDInstrAddr IA) override;
+    unsigned getMispredictionPenalty() override {
+        return mispredictionPenalty;
+    }
+
+};
+
+}
+}
+
+#endif

CustomStages/MCADFetchDelayStage.cpp

@@ -1,5 +1,11 @@
-#include <iostream>
 #include "CustomStages/MCADFetchDelayStage.h"
+#include "MetadataCategories.h"
+#include "llvm/Support/Debug.h"
+#define DEBUG_TYPE "llvm-mca"
+
+#include <iostream>
+#include <iomanip>
+#include <optional>
 
 namespace llvm {
 namespace mcad {
@@ -34,11 +40,35 @@ llvm::Error MCADFetchDelayStage::execute(llvm::mca::InstRef &IR) {
     const llvm::MCInstrDesc &MCID = MCII.get(I->getOpcode());
     bool immediatelyExecute = true;
     unsigned delayCyclesLeft = 0;
-    if(MCID.isBranch()) {
-        // delayed, will have to wait
-        delayCyclesLeft = 100;
+    std::optional<MDInstrAddr> instrAddr = getMDInstrAddrForInstr(MD, IR);
+    // Check if previous instruction was a branch, and if so if the predicted
+    // branch target matched what we ended up executing
+    if(predictedNextInstrAddr.has_value() && instrAddr.has_value()) {
+        if(previousInstrAddr.has_value()) {
+            BPU.recordTakenBranch(*previousInstrAddr, *instrAddr);
+        }
+        if(*predictedNextInstrAddr != *instrAddr) {
+            // Previous prediction was wrong; this instruction will have extra
+            // latency due to misprediction.
+            delayCyclesLeft += BPU.getMispredictionPenalty();
+            LLVM_DEBUG(dbgs() << "[MCAD FetchDelayStage] Previous branch at "); 
+            LLVM_DEBUG(dbgs().write_hex(instrAddr->addr));
+            LLVM_DEBUG(dbgs() << " mispredicted, delaying next instruction by " 
+                       << delayCyclesLeft << " cycle(s).\n");
+        } else {
+            LLVM_DEBUG(dbgs() << "[MCAD FetchDelayStage] Previous branch at ");
+            LLVM_DEBUG(dbgs().write_hex(instrAddr->addr));
+            LLVM_DEBUG(dbgs() << " predicted correctly.\n" );
+        }
+    }
+    // Update branch prediction state
+    if(MCID.isBranch() && instrAddr.has_value()) {
+        predictedNextInstrAddr = BPU.predictBranch(*instrAddr);
+    } else {
+        predictedNextInstrAddr = std::nullopt;
     }
     instrQueue.emplace_back(DelayedInstr { delayCyclesLeft, IR });
+    previousInstrAddr = instrAddr;
     // if the instruction is not delayed, execute it immediately (it will
     // have a delayCyclesLeft of 0 and be at the top of the queue)
     return forwardDueInstrs();

CustomStages/MCADFetchDelayStage.h

@@ -8,9 +8,12 @@
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MCA/SourceMgr.h"
 #include "llvm/MCA/Stages/Stage.h"
+#include "CustomHWUnits/AbstractBranchPredictorUnit.h"
+#include "MetadataRegistry.h"
 
 #include <vector>
 #include <queue>
+#include <optional>
 
 namespace llvm {
 namespace mcad {
@@ -25,23 +28,30 @@ class MCADFetchDelayStage : public llvm::mca::Stage {
     const llvm::MCInstrInfo &MCII;
     std::deque<DelayedInstr> instrQueue = {};
 
+    AbstractBranchPredictorUnit &BPU;
+    MetadataRegistry &MD;
+
+    // Whenever a branch instruction is executed, we run the branch predictor 
+    // and store the predicted instruction address here.
+    // At the next instruction, we compare the predicted address to the actual
+    // address and add a penalty if there is a mismatch.
+    // Non-branch instructions set this member to nullopt.
+    std::optional<MDInstrAddr> predictedNextInstrAddr = std::nullopt;
+    
+    // Stores the address of the last executed instruction.
+    std::optional<MDInstrAddr> previousInstrAddr = std::nullopt;
+
 public:
-    MCADFetchDelayStage(const llvm::MCInstrInfo &MCII) : MCII(MCII) {}
+    MCADFetchDelayStage(const llvm::MCInstrInfo &MCII, MetadataRegistry &MD, AbstractBranchPredictorUnit &BPU) : MCII(MCII), MD(MD), BPU(BPU) {}
 
     bool hasWorkToComplete() const override;
     bool isAvailable(const llvm::mca::InstRef &IR) const override;
     llvm::Error execute(llvm::mca::InstRef &IR) override;
 
-    //llvm::Error cycleStart() override;
     llvm::Error cycleStart() override;
 
     llvm::Error forwardDueInstrs();
 
-    ///// Called after the pipeline is resumed from pausing state.
-    //virtual Error cycleResume() { return ErrorSuccess(); }
-
-    ///// Called once at the end of each cycle.
-
 };
 
 }

MCAWorker.cpp

@@ -37,6 +37,7 @@
 #include <unistd.h>
 
 #include "CustomHWUnits/MCADLSUnit.h"
+#include "CustomHWUnits/NaiveBranchPredictorUnit.h"
 #include "CustomStages/MCADFetchDelayStage.h"
 #include "MCAViews/SummaryView.h"
 #include "MCAViews/TimelineView.h"
@@ -181,10 +182,11 @@ 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>(20);
 
   // Create the pipeline stages.
   auto Fetch = std::make_unique<EntryStage>(SrcMgr);
-  auto FetchDelay = std::make_unique<MCADFetchDelayStage>(MCII);
+  auto FetchDelay = std::make_unique<MCADFetchDelayStage>(MCII, MDRegistry, *BPU);
   auto Dispatch = std::make_unique<DispatchStage>(STI, MRI, MCAPO.DispatchWidth,
                                                   *RCU, *PRF);
   auto Execute =
@@ -196,6 +198,7 @@ std::unique_ptr<mca::Pipeline> MCAWorker::createDefaultPipeline() {
   TheMCA.addHardwareUnit(std::move(PRF));
   TheMCA.addHardwareUnit(std::move(LSU));
   TheMCA.addHardwareUnit(std::move(HWS));
+  TheMCA.addHardwareUnit(std::move(BPU));
 
   // Build the pipeline.
   auto StagePipeline = std::make_unique<Pipeline>();
@@ -224,16 +227,18 @@ 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);
 
   // Create the pipeline stages.
   auto Entry = std::make_unique<EntryStage>(SrcMgr);
-  auto FetchDelay = std::make_unique<MCADFetchDelayStage>(MCII);
+  auto FetchDelay = std::make_unique<MCADFetchDelayStage>(MCII, MDRegistry, *BPU);
   auto InOrderIssue = std::make_unique<InOrderIssueStage>(STI, *PRF, *CB, *LSU);
   auto StagePipeline = std::make_unique<Pipeline>();
 
   // Pass the ownership of all the hardware units to this Context.
   TheMCA.addHardwareUnit(std::move(PRF));
   TheMCA.addHardwareUnit(std::move(LSU));
+  TheMCA.addHardwareUnit(std::move(BPU));
 
   // Build the pipeline.
   StagePipeline->appendStage(std::move(Entry));

MetadataCategories.cpp

@@ -0,0 +1,21 @@
+#include "llvm/MCA/Instruction.h"
+#include <optional>
+#include "MetadataRegistry.h"
+#include "MetadataCategories.h"
+
+namespace llvm {
+namespace mcad {
+
+std::optional<MDInstrAddr> getMDInstrAddrForInstr(MetadataRegistry &MD, const llvm::mca::InstRef &IR) {
+    const llvm::mca::Instruction *I = IR.getInstruction();
+    auto instrId = I->getIdentifier();
+    if (instrId.has_value()) {
+        auto &Registry = MD[llvm::mcad::MD_InstrAddr];
+        auto instrAddr = Registry.get<MDInstrAddr>(*instrId);
+        return instrAddr;
+    }
+    return std::nullopt;
+}
+
+}
+}

MetadataCategories.h

@@ -22,6 +22,21 @@ MD_BinaryRegionMarkers
 
 };
 
+struct MDInstrAddr {
+    unsigned long long addr;
+    const bool operator<(const MDInstrAddr &b) const {
+        return addr < b.addr;
+    }
+    const bool operator==(const MDInstrAddr &b) const {
+        return addr == b.addr;
+    }
+    const bool operator!=(const MDInstrAddr &b) const {
+        return addr != b.addr;
+    }
+};
+
+std::optional<MDInstrAddr> getMDInstrAddrForInstr(MetadataRegistry &MD, const llvm::mca::InstRef &IR);
+
 } // end namespace mcad
 } // end namespace llvm
 #endif

plugins/vivisect-broker/Broker.cpp

@@ -40,6 +40,7 @@ using namespace mcad;
 // Needed so the TypeID of the shared library and main executable refer to the
 // same type.
 extern template class Any::TypeId<MDMemoryAccess>;
+extern template class Any::TypeId<MDInstrAddr>;
 
 class EmulatorService final : public Emulator::Service {
   grpc::Status RecordEmulatorActions(grpc::ServerContext *ctxt,
@@ -143,7 +144,7 @@ class VivisectBroker : public Broker {
               IndexMap[i] = TotalNumTraces;
 
               auto &InstrAddrCat = Registry[MD_InstrAddr];
-              InstrAddrCat[TotalNumTraces] = insn.addr();
+              InstrAddrCat[TotalNumTraces] = MDInstrAddr { insn.addr() };
 
               if (insn.has_memory_access()) {
                   auto MemAccess = insn.memory_access();