4 files changed, 36 insertions, 3 deletions
diff --git a/src/arch/sparc/tlb.hh b/src/arch/sparc/tlb.hh
index 8ed10ff0e..7246cd4f6 100644
--- a/src/arch/sparc/tlb.hh
+++ b/src/arch/sparc/tlb.hh
@@ -114,6 +114,10 @@ class TLB : public BaseTLB
      */
     TlbEntry *lookup(Addr va, int partition_id, bool real, int context_id = 0,
             bool update_used = true);
+
+    /** Remove all entries from the TLB */
+    void flushAll();
+
   protected:
     /** Insert a PTE into the TLB. */
     void insert(Addr vpn, int partition_id, int context_id, bool real,
@@ -122,9 +126,6 @@ class TLB : public BaseTLB
     /** Given an entry id, read that tlb entries' tag. */
     uint64_t TagRead(int entry);
 
-    /** Remove all entries from the TLB */
-    void flushAll();
-
     /** Remove all non-locked entries from the tlb that match partition id. */
     void demapAll(int partition_id);
 
diff --git a/src/cpu/BaseCPU.py b/src/cpu/BaseCPU.py
index 6673b9d41..900a23991 100644
--- a/src/cpu/BaseCPU.py
+++ b/src/cpu/BaseCPU.py
@@ -96,6 +96,7 @@ class BaseCPU(MemObject):
     void switchOut();
     void takeOverFrom(BaseCPU *cpu);
     bool switchedOut();
+    void flushTLBs();
 ''')
 
     def takeOverFrom(self, old_cpu):
diff --git a/src/cpu/base.cc b/src/cpu/base.cc
index 202dc476a..14b5586c8 100644
--- a/src/cpu/base.cc
+++ b/src/cpu/base.cc
@@ -361,6 +361,10 @@ BaseCPU::switchOut()
     _switchedOut = true;
     if (profileEvent && profileEvent->scheduled())
         deschedule(profileEvent);
+
+    // Flush all TLBs in the CPU to avoid having stale translations if
+    // it gets switched in later.
+    flushTLBs();
 }
 
 void
@@ -482,6 +486,22 @@ BaseCPU::takeOverFrom(BaseCPU *oldCPU)
     getDataPort().bind(data_peer_port);
 }
 
+void
+BaseCPU::flushTLBs()
+{
+    for (ThreadID i = 0; i < threadContexts.size(); ++i) {
+        ThreadContext &tc(*threadContexts[i]);
+        CheckerCPU *checker(tc.getCheckerCpuPtr());
+
+        tc.getITBPtr()->flushAll();
+        tc.getDTBPtr()->flushAll();
+        if (checker) {
+            checker->getITBPtr()->flushAll();
+            checker->getDTBPtr()->flushAll();
+        }
+    }
+}
+
 
 BaseCPU::ProfileEvent::ProfileEvent(BaseCPU *_cpu, Tick _interval)
     : cpu(_cpu), interval(_interval)
diff --git a/src/cpu/base.hh b/src/cpu/base.hh
index 633b7f2a7..cd30d29bc 100644
--- a/src/cpu/base.hh
+++ b/src/cpu/base.hh
@@ -324,6 +324,17 @@ class BaseCPU : public MemObject
     virtual void takeOverFrom(BaseCPU *cpu);
 
     /**
+     * Flush all TLBs in the CPU.
+     *
+     * This method is mainly used to flush stale translations when
+     * switching CPUs. It is also exported to the Python world to
+     * allow it to request a TLB flush after draining the CPU to make
+     * it easier to compare traces when debugging
+     * handover/checkpointing.
+     */
+    void flushTLBs();
+
+    /**
      * Determine if the CPU is switched out.
      *
      * @return True if the CPU is switched out, false otherwise.