11 files changed, 41 insertions, 16 deletions

diff --git a/src/cpu/o3/commit.hh b/src/cpu/o3/commit.hh
index 49ff5cdad..c39bc10f9 100644
--- a/src/cpu/o3/commit.hh
+++ b/src/cpu/o3/commit.hh
@@ -188,7 +188,7 @@ class DefaultCommit
     void initStage();
 
     /** Initializes the draining of commit. */
-    void drain();
+    bool drain();
 
     /** Resumes execution after draining. */
     void resume();
diff --git a/src/cpu/o3/commit_impl.hh b/src/cpu/o3/commit_impl.hh
index 2eb05afac..b50c9a898 100644
--- a/src/cpu/o3/commit_impl.hh
+++ b/src/cpu/o3/commit_impl.hh
@@ -350,10 +350,18 @@ DefaultCommit<Impl>::initStage()
 }
 
 template <class Impl>
-void
+bool
 DefaultCommit<Impl>::drain()
 {
     drainPending = true;
+
+    // If it's already drained, return true.
+    if (rob->isEmpty() && !iewStage->hasStoresToWB()) {
+        cpu->signalDrained();
+        return true;
+    }
+
+    return false;
 }
 
 template <class Impl>
diff --git a/src/cpu/o3/cpu.cc b/src/cpu/o3/cpu.cc
index b182d5ca7..3a52fe4c2 100644
--- a/src/cpu/o3/cpu.cc
+++ b/src/cpu/o3/cpu.cc
@@ -712,19 +712,27 @@ bool
 FullO3CPU<Impl>::drain(Event *drain_event)
 {
     drainCount = 0;
-    drainEvent = drain_event;
     fetch.drain();
     decode.drain();
     rename.drain();
     iew.drain();
     commit.drain();
+    // A bit of a hack...set the drainEvent after all the drain()
+    // calls have been made, that way if all of the stages drain
+    // immediately, the signalDrained() function knows not to call
+    // process on the drain event.
+    drainEvent = drain_event;
 
     // Wake the CPU and record activity so everything can drain out if
-    // the CPU is currently idle.
-    wakeCPU();
-    activityRec.activity();
+    // the CPU was not able to immediately drain.
+    if (_status != Drained) {
+        wakeCPU();
+        activityRec.activity();
 
-    return false;
+        return false;
+    } else {
+        return true;
+    }
 }
 
 template <class Impl>
@@ -751,8 +759,13 @@ FullO3CPU<Impl>::signalDrained()
     if (++drainCount == NumStages) {
         if (tickEvent.scheduled())
             tickEvent.squash();
+
         _status = Drained;
-        drainEvent->process();
+
+        if (drainEvent) {
+            drainEvent->process();
+            drainEvent = NULL;
+        }
     }
     assert(drainCount <= 5);
 }
diff --git a/src/cpu/o3/decode.hh b/src/cpu/o3/decode.hh
index 1e96f1884..7f5ecbc26 100644
--- a/src/cpu/o3/decode.hh
+++ b/src/cpu/o3/decode.hh
@@ -110,7 +110,7 @@ class DefaultDecode
     void setActiveThreads(std::list<unsigned> *at_ptr);
 
     /** Drains the decode stage. */
-    void drain();
+    bool drain();
 
     /** Resumes execution after a drain. */
     void resume() { }
diff --git a/src/cpu/o3/decode_impl.hh b/src/cpu/o3/decode_impl.hh
index 71637883b..8b851c032 100644
--- a/src/cpu/o3/decode_impl.hh
+++ b/src/cpu/o3/decode_impl.hh
@@ -165,11 +165,12 @@ DefaultDecode<Impl>::setActiveThreads(list<unsigned> *at_ptr)
 }
 
 template <class Impl>
-void
+bool
 DefaultDecode<Impl>::drain()
 {
     // Decode is done draining at any time.
     cpu->signalDrained();
+    return true;
 }
 
 template <class Impl>
diff --git a/src/cpu/o3/fetch.hh b/src/cpu/o3/fetch.hh
index 9611f0455..a793c7361 100644
--- a/src/cpu/o3/fetch.hh
+++ b/src/cpu/o3/fetch.hh
@@ -181,7 +181,7 @@ class DefaultFetch
     void processCacheCompletion(PacketPtr pkt);
 
     /** Begins the drain of the fetch stage. */
-    void drain();
+    bool drain();
 
     /** Resumes execution after a drain. */
     void resume();
diff --git a/src/cpu/o3/fetch_impl.hh b/src/cpu/o3/fetch_impl.hh
index 500b5304e..c0cc189f2 100644
--- a/src/cpu/o3/fetch_impl.hh
+++ b/src/cpu/o3/fetch_impl.hh
@@ -385,12 +385,13 @@ DefaultFetch<Impl>::processCacheCompletion(PacketPtr pkt)
 }
 
 template <class Impl>
-void
+bool
 DefaultFetch<Impl>::drain()
 {
     // Fetch is ready to drain at any time.
     cpu->signalDrained();
     drainPending = true;
+    return true;
 }
 
 template <class Impl>
diff --git a/src/cpu/o3/iew.hh b/src/cpu/o3/iew.hh
index 774b6dcbd..4908a6649 100644
--- a/src/cpu/o3/iew.hh
+++ b/src/cpu/o3/iew.hh
@@ -144,7 +144,7 @@ class DefaultIEW
     void setScoreboard(Scoreboard *sb_ptr);
 
     /** Drains IEW stage. */
-    void drain();
+    bool drain();
 
     /** Resumes execution after a drain. */
     void resume();
diff --git a/src/cpu/o3/iew_impl.hh b/src/cpu/o3/iew_impl.hh
index c3aa748ae..0d82645e3 100644
--- a/src/cpu/o3/iew_impl.hh
+++ b/src/cpu/o3/iew_impl.hh
@@ -354,11 +354,12 @@ DefaultIEW<Impl>::setScoreboard(Scoreboard *sb_ptr)
 }
 
 template <class Impl>
-void
+bool
 DefaultIEW<Impl>::drain()
 {
     // IEW is ready to drain at any time.
     cpu->signalDrained();
+    return true;
 }
 
 template <class Impl>
diff --git a/src/cpu/o3/rename.hh b/src/cpu/o3/rename.hh
index 538dd9bb4..034087feb 100644
--- a/src/cpu/o3/rename.hh
+++ b/src/cpu/o3/rename.hh
@@ -158,7 +158,7 @@ class DefaultRename
     void setScoreboard(Scoreboard *_scoreboard);
 
     /** Drains the rename stage. */
-    void drain();
+    bool drain();
 
     /** Resumes execution after a drain. */
     void resume() { }
diff --git a/src/cpu/o3/rename_impl.hh b/src/cpu/o3/rename_impl.hh
index fddbae3db..805a72808 100644
--- a/src/cpu/o3/rename_impl.hh
+++ b/src/cpu/o3/rename_impl.hh
@@ -257,11 +257,12 @@ DefaultRename<Impl>::setScoreboard(Scoreboard *_scoreboard)
 }
 
 template <class Impl>
-void
+bool
 DefaultRename<Impl>::drain()
 {
     // Rename is ready to switch out at any time.
     cpu->signalDrained();
+    return true;
 }
 
 template <class Impl>