2 files changed, 82 insertions, 73 deletions

diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index 621f99f29..78603be4f 100644
--- a/src/cpu/simple/timing.cc
+++ b/src/cpu/simple/timing.cc
@@ -94,11 +94,10 @@ TimingSimpleCPU::TimingCPUPort::TickEvent::schedule(PacketPtr _pkt, Tick t)
 TimingSimpleCPU::TimingSimpleCPU(TimingSimpleCPUParams *p)
     : BaseSimpleCPU(p), fetchTranslation(this), icachePort(this),
       dcachePort(this), ifetch_pkt(NULL), dcache_pkt(NULL), previousCycle(0),
-      fetchEvent(this)
+      fetchEvent(this), drainManager(NULL)
 {
     _status = Idle;
 
-    setDrainState(Drainable::Running);
     system->totalNumInsts = 0;
 }
 
@@ -107,36 +106,27 @@ TimingSimpleCPU::~TimingSimpleCPU()
 {
 }
 
-void
-TimingSimpleCPU::serialize(ostream &os)
-{
-    Drainable::State so_state(getDrainState());
-    SERIALIZE_ENUM(so_state);
-    BaseSimpleCPU::serialize(os);
-}
-
-void
-TimingSimpleCPU::unserialize(Checkpoint *cp, const string &section)
-{
-    Drainable::State so_state;
-    UNSERIALIZE_ENUM(so_state);
-    BaseSimpleCPU::unserialize(cp, section);
-}
-
 unsigned int
 TimingSimpleCPU::drain(DrainManager *drain_manager)
 {
-    // TimingSimpleCPU is ready to drain if it's not waiting for
-    // an access to complete.
     if (_status == Idle ||
-        _status == BaseSimpleCPU::Running ||
+        (_status == BaseSimpleCPU::Running && isDrained()) ||
         _status == SwitchedOut) {
-        setDrainState(Drainable::Drained);
+        assert(!fetchEvent.scheduled());
+        DPRINTF(Drain, "No need to drain.\n");
         return 0;
     } else {
-        setDrainState(Drainable::Draining);
         drainManager = drain_manager;
-        DPRINTF(Drain, "CPU not drained\n");
+        DPRINTF(Drain, "Requesting drain: %s\n", pcState());
+
+        // The fetch event can become descheduled if a drain didn't
+        // succeed on the first attempt. We need to reschedule it if
+        // the CPU is waiting for a microcode routine to complete.
+        if (_status == BaseSimpleCPU::Running && !isDrained() &&
+            !fetchEvent.scheduled()) {
+            schedule(fetchEvent, nextCycle());
+        }
+
         return 1;
     }
 }
@@ -144,6 +134,8 @@ TimingSimpleCPU::drain(DrainManager *drain_manager)
 void
 TimingSimpleCPU::drainResume()
 {
+    assert(!fetchEvent.scheduled());
+
     DPRINTF(SimpleCPU, "Resume\n");
     if (_status != SwitchedOut && _status != Idle) {
         if (system->getMemoryMode() != Enums::timing) {
@@ -151,13 +143,25 @@ TimingSimpleCPU::drainResume()
                   "'timing' mode.\n");
         }
 
-        if (fetchEvent.scheduled())
-           deschedule(fetchEvent);
-
         schedule(fetchEvent, nextCycle());
     }
+}
+
+bool
+TimingSimpleCPU::tryCompleteDrain()
+{
+    if (!drainManager)
+        return false;
+
+    DPRINTF(Drain, "tryCompleteDrain: %s\n", pcState());
+    if (!isDrained())
+        return false;
+
+    DPRINTF(Drain, "CPU done draining, processing drain event\n");
+    drainManager->signalDrainDone();
+    drainManager = NULL;
 
-    setDrainState(Drainable::Running);
+    return true;
 }
 
 void
@@ -165,14 +169,13 @@ TimingSimpleCPU::switchOut()
 {
     BaseSimpleCPU::switchOut();
 
+    assert(!fetchEvent.scheduled());
     assert(_status == BaseSimpleCPU::Running || _status == Idle);
+    assert(!stayAtPC);
+    assert(microPC() == 0);
+
     _status = SwitchedOut;
     numCycles += curCycle() - previousCycle;
-
-    // If we've been scheduled to resume but are then told to switch out,
-    // we'll need to cancel it.
-    if (fetchEvent.scheduled())
-        deschedule(fetchEvent);
 }
 
 
@@ -344,12 +347,7 @@ TimingSimpleCPU::translationFault(Fault fault)
 
     postExecute();
 
-    if (getDrainState() == Drainable::Draining) {
-        advancePC(fault);
-        completeDrain();
-    } else {
-        advanceInst(fault);
-    }
+    advanceInst(fault);
 }
 
 void
@@ -618,7 +616,6 @@ TimingSimpleCPU::sendFetch(Fault fault, RequestPtr req, ThreadContext *tc)
 void
 TimingSimpleCPU::advanceInst(Fault fault)
 {
-
     if (_status == Faulting)
         return;
 
@@ -634,6 +631,9 @@ TimingSimpleCPU::advanceInst(Fault fault)
     if (!stayAtPC)
         advancePC(fault);
 
+    if (tryCompleteDrain())
+            return;
+
     if (_status == BaseSimpleCPU::Running) {
         // kick off fetch of next instruction... callback from icache
         // response will cause that instruction to be executed,
@@ -660,16 +660,6 @@ TimingSimpleCPU::completeIfetch(PacketPtr pkt)
     numCycles += curCycle() - previousCycle;
     previousCycle = curCycle();
 
-    if (getDrainState() == Drainable::Draining) {
-        if (pkt) {
-            delete pkt->req;
-            delete pkt;
-        }
-
-        completeDrain();
-        return;
-    }
-
     preExecute();
     if (curStaticInst && curStaticInst->isMemRef()) {
         // load or store: just send to dcache
@@ -816,25 +806,9 @@ TimingSimpleCPU::completeDataAccess(PacketPtr pkt)
 
     postExecute();
 
-    if (getDrainState() == Drainable::Draining) {
-        advancePC(fault);
-        completeDrain();
-
-        return;
-    }
-
     advanceInst(fault);
 }
 
-
-void
-TimingSimpleCPU::completeDrain()
-{
-    DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    setDrainState(Drainable::Drained);
-    drainManager->signalDrainDone();
-}
-
 bool
 TimingSimpleCPU::DcachePort::recvTimingResp(PacketPtr pkt)
 {
diff --git a/src/cpu/simple/timing.hh b/src/cpu/simple/timing.hh
index e7f5122d0..af780265f 100644
--- a/src/cpu/simple/timing.hh
+++ b/src/cpu/simple/timing.hh
@@ -1,4 +1,16 @@
 /*
+ * Copyright (c) 2012 ARM Limited
+ * All rights reserved
+ *
+ * The license below extends only to copyright in the software and shall
+ * not be construed as granting a license to any other intellectual
+ * property including but not limited to intellectual property relating
+ * to a hardware implementation of the functionality of the software
+ * licensed hereunder.  You may use the software subject to the license
+ * terms below provided that you ensure that this notice is replicated
+ * unmodified and in its entirety in all distributions of the software,
+ * modified or unmodified, in source code or in binary form.
+ *
  * Copyright (c) 2002-2005 The Regents of The University of Michigan
  * All rights reserved.
  *
@@ -44,9 +56,6 @@ class TimingSimpleCPU : public BaseSimpleCPU
 
     virtual void init();
 
-  public:
-    DrainManager *drainManager;
-
   private:
 
     /*
@@ -246,9 +255,6 @@ class TimingSimpleCPU : public BaseSimpleCPU
 
   public:
 
-    virtual void serialize(std::ostream &os);
-    virtual void unserialize(Checkpoint *cp, const std::string &section);
-
     unsigned int drain(DrainManager *drain_manager);
     void drainResume();
 
@@ -302,7 +308,36 @@ class TimingSimpleCPU : public BaseSimpleCPU
         virtual const char *description() const;
     };
 
-    void completeDrain();
+    /**
+     * Check if a system is in a drained state.
+     *
+     * We need to drain if:
+     * <ul>
+     * <li>We are in the middle of a microcode sequence as some CPUs
+     *     (e.g., HW accelerated CPUs) can't be started in the middle
+     *     of a gem5 microcode sequence.
+     *
+     * <li>Stay at PC is true.
+     * </ul>
+     */
+    bool isDrained() {
+        return microPC() == 0 &&
+            !stayAtPC;
+    }
+
+    /**
+     * Try to complete a drain request.
+     *
+     * @returns true if the CPU is drained, false otherwise.
+     */
+    bool tryCompleteDrain();
+
+    /**
+     * Drain manager to use when signaling drain completion
+     *
+     * This pointer is non-NULL when draining and NULL otherwise.
+     */
+    DrainManager *drainManager;
 };
 
 #endif // __CPU_SIMPLE_TIMING_HH__