2 files changed, 14 insertions, 11 deletions

diff --git a/src/cpu/simple/atomic.cc b/src/cpu/simple/atomic.cc
index 490be20ae..fe421ae6c 100644
--- a/src/cpu/simple/atomic.cc
+++ b/src/cpu/simple/atomic.cc
@@ -182,9 +182,9 @@ AtomicSimpleCPU::unserialize(Checkpoint *cp, const string &section)
 void
 AtomicSimpleCPU::resume()
 {
-    assert(system->getMemoryMode() == System::Atomic);
     changeState(SimObject::Running);
     if (thread->status() == ThreadContext::Active) {
+        assert(system->getMemoryMode() == System::Atomic);
         if (!tickEvent.scheduled())
             tickEvent.schedule(curTick);
     }
diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index 48362c42a..88aa882e3 100644
--- a/src/cpu/simple/timing.cc
+++ b/src/cpu/simple/timing.cc
@@ -146,6 +146,8 @@ void
 TimingSimpleCPU::resume()
 {
     if (_status != SwitchedOut && _status != Idle) {
+        assert(system->getMemoryMode() == System::Timing);
+
         // Delete the old event if it existed.
         if (fetchEvent) {
             if (fetchEvent->scheduled())
@@ -159,7 +161,6 @@ TimingSimpleCPU::resume()
         fetchEvent->schedule(curTick);
     }
 
-    assert(system->getMemoryMode() == System::Timing);
     changeState(SimObject::Running);
 }
 
@@ -190,6 +191,10 @@ TimingSimpleCPU::takeOverFrom(BaseCPU *oldCPU)
             break;
         }
     }
+
+    if (_status != Running) {
+        _status = Idle;
+    }
 }
 
 
@@ -533,15 +538,6 @@ TimingSimpleCPU::completeDataAccess(Packet *pkt)
     assert(_status == DcacheWaitResponse);
     _status = Running;
 
-    if (getState() == SimObject::Draining) {
-        completeDrain();
-
-        delete pkt->req;
-        delete pkt;
-
-        return;
-    }
-
     Fault fault = curStaticInst->completeAcc(pkt, this, traceData);
 
     if (pkt->isRead() && pkt->req->isLocked()) {
@@ -551,6 +547,13 @@ TimingSimpleCPU::completeDataAccess(Packet *pkt)
     delete pkt->req;
     delete pkt;
 
+    if (getState() == SimObject::Draining) {
+        advancePC(fault);
+        completeDrain();
+
+        return;
+    }
+
     postExecute();
     advanceInst(fault);
 }