From 968f3c67c770ecf0fd00e4b4679d81df75125b52 Mon Sep 17 00:00:00 2001
From: Aditya Badole <supershinx@gmail.com>
Date: Thu, 12 Oct 2023 15:44:54 +0800
Subject: [PATCH] Update serialization helper to properly store tagged regs

Change-Id: I5ce2a79eb028d861249b03edaf424a62dd8f00f4
---
 src/arch/riscvcapstone/o3/thread_context.cc | 18 ++++++++
 src/arch/riscvcapstone/o3/thread_context.hh |  4 ++
 src/cpu/regfile.hh                          |  3 +-
 src/cpu/thread_context.cc                   | 48 +++++++++++++++++++++
 src/cpu/thread_context.hh                   | 15 +++++++
 5 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/src/arch/riscvcapstone/o3/thread_context.cc b/src/arch/riscvcapstone/o3/thread_context.cc
index 5c99242f85..592ec27274 100644
--- a/src/arch/riscvcapstone/o3/thread_context.cc
+++ b/src/arch/riscvcapstone/o3/thread_context.cc
@@ -156,6 +156,24 @@ ThreadContext::getRegFlat(const RegId &reg) const
     return cpu->getArchReg(reg, thread->threadId());
 }
 
+ConstTaggedRegVal
+ThreadContext::getTaggedRegFlat(const RegId &reg) const
+{
+    return cpu->getTaggedArchReg(reg, thread->threadId());
+}
+
+void
+ThreadContext::setTaggedRegFlat(int i, ConstTaggedRegVal v)
+{
+    cpu->setTaggedArchReg(RegId(IntRegClass, i), v, thread->threadId());
+    conditionalSquash();
+}
+
+void
+ThreadContext::printRegs() {
+    cpu->printRegs();
+}
+
 void *
 ThreadContext::getWritableRegFlat(const RegId &reg)
 {
diff --git a/src/arch/riscvcapstone/o3/thread_context.hh b/src/arch/riscvcapstone/o3/thread_context.hh
index 08c08c20b2..5ae9bafa58 100644
--- a/src/arch/riscvcapstone/o3/thread_context.hh
+++ b/src/arch/riscvcapstone/o3/thread_context.hh
@@ -249,6 +249,10 @@ class ThreadContext : public gem5::ThreadContext
     void getRegFlat(const RegId &reg, void *val) const override;
     void *getWritableRegFlat(const RegId &reg) override;
 
+    ConstTaggedRegVal getTaggedRegFlat(const RegId &reg) const override;
+    void setTaggedRegFlat(int i, ConstTaggedRegVal v) override;
+    void printRegs() override;
+
     void setRegFlat(const RegId &reg, RegVal val) override;
     void setRegFlat(const RegId &reg, const void *val) override;
 
diff --git a/src/cpu/regfile.hh b/src/cpu/regfile.hh
index 63c4bec380..7b5e0f82aa 100644
--- a/src/cpu/regfile.hh
+++ b/src/cpu/regfile.hh
@@ -46,9 +46,8 @@ class RegFile
     const size_t _regShift;
     const size_t _regBytes;
 
-
-    const bool hasTags;
 #ifdef TARGET_RISCVCapstone
+    const bool hasTags;
     std::vector<bool> tags;
 #endif
 
diff --git a/src/cpu/thread_context.cc b/src/cpu/thread_context.cc
index 400bc165a7..a6b7e2a286 100644
--- a/src/cpu/thread_context.cc
+++ b/src/cpu/thread_context.cc
@@ -242,10 +242,34 @@ serialize(const ThreadContext &tc, CheckpointOut &cp)
     SERIALIZE_CONTAINER(vecPredRegs);
 
     const size_t numInts = regClasses.at(IntRegClass).numRegs();
+
+#ifdef TARGET_RISCVCapstone
+    uint64_t intRegValLo[numInts];
+    uint64_t intRegValHi[numInts];
+    bool intRegsTag[numInts];
+    ConstTaggedRegVal temp;
+    uint128_t cap;
+    for (int i = 0; i < numInts; ++i)
+    {
+        temp = tc.getTaggedRegFlat(RegId(IntRegClass, i));
+        intRegsTag[i] = temp.getTag();
+        if(intRegsTag[i]) {
+            cap = temp.getRegVal().rawCapVal();
+            memcpy(&intRegValLo[i], &cap.lo, sizeof(cap.lo));
+            memcpy(&intRegValHi[i], &cap.hi, sizeof(cap.hi));
+        } else {
+            intRegValLo[i] = temp.getRegVal().intVal();
+        }
+    }
+    SERIALIZE_ARRAY(intRegValLo, numInts);
+    SERIALIZE_ARRAY(intRegValHi, numInts);
+    SERIALIZE_ARRAY(intRegsTag, numInts);
+#else
     RegVal intRegs[numInts];
     for (int i = 0; i < numInts; ++i)
         intRegs[i] = tc.readIntRegFlat(i);
     SERIALIZE_ARRAY(intRegs, numInts);
+#endif
 
     const size_t numCcs = regClasses.at(CCRegClass).numRegs();
     if (numCcs) {
@@ -288,10 +312,34 @@ unserialize(ThreadContext &tc, CheckpointIn &cp)
     }
 
     const size_t numInts = regClasses.at(IntRegClass).numRegs();
+
+#ifdef TARGET_RISCVCapstone
+    uint64_t intRegValLo[numInts];
+    uint64_t intRegValHi[numInts];
+    bool intRegsTag[numInts];
+    ConstTaggedRegVal temp;
+    uint128_t cap;
+    UNSERIALIZE_ARRAY(intRegValLo, numInts);
+    UNSERIALIZE_ARRAY(intRegValHi, numInts);
+    UNSERIALIZE_ARRAY(intRegsTag, numInts);
+    for (int i = 0; i < numInts; ++i) {
+        temp.setTag(intRegsTag[i]);
+        if(intRegsTag[i]) {
+            memcpy(&cap.lo, &intRegValLo[i], sizeof(intRegValLo[i]));
+            memcpy(&cap.hi, &intRegValHi[i], sizeof(intRegValHi[i]));
+            memcpy(&temp.getRegVal().rawCapVal(), &cap, sizeof(cap));
+        } else {
+            temp.getRegVal().intVal() = intRegValLo[i];
+        }
+        tc.setTaggedRegFlat(i, temp);
+    }
+    // tc.printRegs();
+#else
     RegVal intRegs[numInts];
     UNSERIALIZE_ARRAY(intRegs, numInts);
     for (int i = 0; i < numInts; ++i)
         tc.setIntRegFlat(i, intRegs[i]);
+#endif
 
     const size_t numCcs = regClasses.at(CCRegClass).numRegs();
     if (numCcs) {
diff --git a/src/cpu/thread_context.hh b/src/cpu/thread_context.hh
index 9875d00dbd..41f9583f67 100644
--- a/src/cpu/thread_context.hh
+++ b/src/cpu/thread_context.hh
@@ -323,6 +323,21 @@ class ThreadContext : public PCEventScope
     virtual void getRegFlat(const RegId &reg, void *val) const = 0;
     virtual void *getWritableRegFlat(const RegId &reg) = 0;
 
+#ifdef TARGET_RISCVCapstone
+    virtual ConstTaggedRegVal getTaggedRegFlat(const RegId &reg) const {
+        ConstTaggedRegVal v;
+        return v;
+    };
+
+    virtual void setTaggedRegFlat(int i, ConstTaggedRegVal v) {
+        return;
+    }
+
+    virtual void printRegs() {
+        return;
+    }
+#endif
+
     virtual void setRegFlat(const RegId &reg, RegVal val);
     virtual void setRegFlat(const RegId &reg, const void *val) = 0;