3 files changed, 32 insertions, 0 deletions
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.