1 files changed, 34 insertions, 13 deletions

diff --git a/src/arch/arm/table_walker.cc b/src/arch/arm/table_walker.cc
index dfb406543..7dbe92d9b 100644
--- a/src/arch/arm/table_walker.cc
+++ b/src/arch/arm/table_walker.cc
@@ -43,6 +43,7 @@
 #include "cpu/base.hh"
 #include "cpu/thread_context.hh"
 #include "debug/Checkpoint.hh"
+#include "debug/Drain.hh"
 #include "debug/TLB.hh"
 #include "debug/TLBVerbose.hh"
 #include "sim/system.hh"
@@ -51,7 +52,7 @@ using namespace ArmISA;
 
 TableWalker::TableWalker(const Params *p)
     : MemObject(p), port(this, params()->sys, params()->min_backoff,
-                         params()->max_backoff),
+                         params()->max_backoff), drainEvent(NULL),
       tlb(NULL), currState(NULL), pending(false),
       masterId(p->sys->getMasterId(name())),
       doL1DescEvent(this), doL2DescEvent(this), doProcessEvent(this)
@@ -64,20 +65,38 @@ TableWalker::~TableWalker()
     ;
 }
 
+void
+TableWalker::completeDrain()
+{
+    if (drainEvent && stateQueueL1.empty() && stateQueueL2.empty() &&
+        pendingQueue.empty()) {
+        changeState(Drained);
+        DPRINTF(Drain, "TableWalker done draining, processing drain event\n");
+        drainEvent->process();
+        drainEvent = NULL;
+    }
+}
+
 unsigned int
 TableWalker::drain(Event *de)
 {
-    if (stateQueueL1.size() || stateQueueL2.size() || pendingQueue.size())
-    {
-        changeState(Draining);
-        DPRINTF(Checkpoint, "TableWalker busy, wait to drain\n");
-        return 1;
-    }
-    else
-    {
+    unsigned int count = port.drain(de);
+
+    if (stateQueueL1.empty() && stateQueueL2.empty() &&
+        pendingQueue.empty()) {
         changeState(Drained);
-        DPRINTF(Checkpoint, "TableWalker free, no need to drain\n");
-        return 0;
+        DPRINTF(Drain, "TableWalker free, no need to drain\n");
+
+        // table walker is drained, but its ports may still need to be drained
+        return count;
+    } else {
+        drainEvent = de;
+        changeState(Draining);
+        DPRINTF(Drain, "TableWalker not drained\n");
+
+        // return port drain count plus the table walker itself needs to drain
+        return count + 1;
+
     }
 }
 
@@ -86,8 +105,8 @@ TableWalker::resume()
 {
     MemObject::resume();
     if ((params()->sys->getMemoryMode() == Enums::timing) && currState) {
-            delete currState;
-            currState = NULL;
+        delete currState;
+        currState = NULL;
     }
 }
 
@@ -667,6 +686,7 @@ TableWalker::doL1DescriptorWrapper()
     doL1Descriptor();
 
     stateQueueL1.pop_front();
+    completeDrain();
     // Check if fault was generated
     if (currState->fault != NoFault) {
         currState->transState->finish(currState->fault, currState->req,
@@ -723,6 +743,7 @@ TableWalker::doL2DescriptorWrapper()
 
 
     stateQueueL2.pop_front();
+    completeDrain();
     pending = false;
     nextWalk(currState->tc);