memory mode information now contained in system object

States are now running, draining, or drained. memory state information moved into system object
system parameter is not fs only for cpus
Implement drain() support in devices
Update for drain() call that returns number of times drain_event->process() will be called

Break O3 CPU! No sense in putting in a hack change that kevin is going to remove in a few minutes i imagine

src/cpu/simple/atomic.cc:
src/cpu/simple/atomic.hh:
    Since se mode has a system, allow access to it
    Verify that the atomic cpu is connected to an atomic system on resume
src/cpu/simple/base.cc:
    Since se mode has a system, allow access to it
src/cpu/simple/timing.cc:
src/cpu/simple/timing.hh:
    Update for new drain() call that returns number of times drain_event->process() will be called and memory state being moved into the system
    Since se mode has a system, allow access to it
    Verify that the timing cpu is connected to an timing system on resume
src/dev/ide_disk.cc:
src/dev/io_device.cc:
src/dev/io_device.hh:
src/dev/ns_gige.cc:
src/dev/ns_gige.hh:
src/dev/pcidev.cc:
src/dev/pcidev.hh:
src/dev/sinic.cc:
src/dev/sinic.hh:
    Implement drain() support in devices
src/python/m5/config.py:
    Allow drain to return number of times drain_event->process() will be called. Normally 0 or 1 but things like O3 cpu or devices with multiple ports may want to call it many times
src/python/m5/objects/BaseCPU.py:
    move system parameter out of fs to everyone
src/sim/sim_object.cc:
src/sim/sim_object.hh:
    States are now running, draining, or drained. memory state information moved into system object
src/sim/system.cc:
src/sim/system.hh:
    memory mode information now contained in system object

--HG--
extra : convert_revision : 1389c77e66ee6d9710bf77b4306fb47e107b21cf

20 files changed, 272 insertions, 120 deletions

diff --git a/src/cpu/simple/atomic.cc b/src/cpu/simple/atomic.cc
index 12bfdeb9b..1752b2b5b 100644
--- a/src/cpu/simple/atomic.cc
+++ b/src/cpu/simple/atomic.cc
@@ -33,6 +33,7 @@
 #include "cpu/simple/atomic.hh"
 #include "mem/packet_impl.hh"
 #include "sim/builder.hh"
+#include "sim/system.hh"
 
 using namespace std;
 using namespace TheISA;
@@ -173,6 +174,13 @@ AtomicSimpleCPU::unserialize(Checkpoint *cp, const string &section)
 }
 
 void
+AtomicSimpleCPU::resume()
+{
+    assert(system->getMemoryMode() == System::Atomic);
+    changeState(SimObject::Running);
+}
+
+void
 AtomicSimpleCPU::switchOut()
 {
     assert(status() == Running || status() == Idle);
@@ -451,11 +459,11 @@ BEGIN_DECLARE_SIM_OBJECT_PARAMS(AtomicSimpleCPU)
     Param<Counter> max_loads_any_thread;
     Param<Counter> max_loads_all_threads;
     SimObjectParam<MemObject *> mem;
+    SimObjectParam<System *> system;
 
 #if FULL_SYSTEM
     SimObjectParam<AlphaITB *> itb;
     SimObjectParam<AlphaDTB *> dtb;
-    SimObjectParam<System *> system;
     Param<int> cpu_id;
     Param<Tick> profile;
 #else
@@ -483,11 +491,11 @@ BEGIN_INIT_SIM_OBJECT_PARAMS(AtomicSimpleCPU)
     INIT_PARAM(max_loads_all_threads,
                "terminate when all threads have reached this load count"),
     INIT_PARAM(mem, "memory"),
+    INIT_PARAM(system, "system object"),
 
 #if FULL_SYSTEM
     INIT_PARAM(itb, "Instruction TLB"),
     INIT_PARAM(dtb, "Data TLB"),
-    INIT_PARAM(system, "system object"),
     INIT_PARAM(cpu_id, "processor ID"),
     INIT_PARAM(profile, ""),
 #else
@@ -520,11 +528,11 @@ CREATE_SIM_OBJECT(AtomicSimpleCPU)
     params->width = width;
     params->simulate_stalls = simulate_stalls;
     params->mem = mem;
+    params->system = system;
 
 #if FULL_SYSTEM
     params->itb = itb;
     params->dtb = dtb;
-    params->system = system;
     params->cpu_id = cpu_id;
     params->profile = profile;
 #else
diff --git a/src/cpu/simple/atomic.hh b/src/cpu/simple/atomic.hh
index 179b4a721..d59ca01aa 100644
--- a/src/cpu/simple/atomic.hh
+++ b/src/cpu/simple/atomic.hh
@@ -127,6 +127,7 @@ class AtomicSimpleCPU : public BaseSimpleCPU
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
+    virtual void resume();
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
 
diff --git a/src/cpu/simple/base.cc b/src/cpu/simple/base.cc
index a50541189..af10e64d7 100644
--- a/src/cpu/simple/base.cc
+++ b/src/cpu/simple/base.cc
@@ -55,10 +55,10 @@
 #include "sim/sim_events.hh"
 #include "sim/sim_object.hh"
 #include "sim/stats.hh"
+#include "sim/system.hh"
 
 #if FULL_SYSTEM
 #include "base/remote_gdb.hh"
-#include "sim/system.hh"
 #include "arch/tlb.hh"
 #include "arch/stacktrace.hh"
 #include "arch/vtophys.hh"
diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index e55301c6b..d2c2c7c47 100644
--- a/src/cpu/simple/timing.cc
+++ b/src/cpu/simple/timing.cc
@@ -33,6 +33,7 @@
 #include "cpu/simple/timing.hh"
 #include "mem/packet_impl.hh"
 #include "sim/builder.hh"
+#include "sim/system.hh"
 
 using namespace std;
 using namespace TheISA;
@@ -91,7 +92,7 @@ TimingSimpleCPU::TimingSimpleCPU(Params *p)
     ifetch_pkt = dcache_pkt = NULL;
     drainEvent = NULL;
     fetchEvent = NULL;
-    state = SimObject::Timing;
+    changeState(SimObject::Running);
 }
 
 
@@ -113,18 +114,18 @@ TimingSimpleCPU::unserialize(Checkpoint *cp, const string &section)
     BaseSimpleCPU::unserialize(cp, section);
 }
 
-bool
+unsigned int
 TimingSimpleCPU::drain(Event *drain_event)
 {
     // TimingSimpleCPU is ready to drain if it's not waiting for
     // an access to complete.
     if (status() == Idle || status() == Running || status() == SwitchedOut) {
-        changeState(SimObject::DrainedTiming);
-        return true;
+        changeState(SimObject::Drained);
+        return 0;
     } else {
         changeState(SimObject::Draining);
         drainEvent = drain_event;
-        return false;
+        return 1;
     }
 }
 
@@ -142,12 +143,9 @@ TimingSimpleCPU::resume()
             new EventWrapper<TimingSimpleCPU, &TimingSimpleCPU::fetch>(this, false);
         fetchEvent->schedule(curTick);
     }
-}
 
-void
-TimingSimpleCPU::setMemoryMode(State new_mode)
-{
-    assert(new_mode == SimObject::Timing);
+    assert(system->getMemoryMode() == System::Timing);
+    changeState(SimObject::Running);
 }
 
 void
@@ -514,7 +512,7 @@ void
 TimingSimpleCPU::completeDrain()
 {
     DPRINTF(Config, "Done draining\n");
-    changeState(SimObject::DrainedTiming);
+    changeState(SimObject::Drained);
     drainEvent->process();
 }
 
@@ -551,11 +549,11 @@ BEGIN_DECLARE_SIM_OBJECT_PARAMS(TimingSimpleCPU)
     Param<Counter> max_loads_any_thread;
     Param<Counter> max_loads_all_threads;
     SimObjectParam<MemObject *> mem;
+    SimObjectParam<System *> system;
 
 #if FULL_SYSTEM
     SimObjectParam<AlphaITB *> itb;
     SimObjectParam<AlphaDTB *> dtb;
-    SimObjectParam<System *> system;
     Param<int> cpu_id;
     Param<Tick> profile;
 #else
@@ -583,11 +581,11 @@ BEGIN_INIT_SIM_OBJECT_PARAMS(TimingSimpleCPU)
     INIT_PARAM(max_loads_all_threads,
                "terminate when all threads have reached this load count"),
     INIT_PARAM(mem, "memory"),
+    INIT_PARAM(system, "system object"),
 
 #if FULL_SYSTEM
     INIT_PARAM(itb, "Instruction TLB"),
     INIT_PARAM(dtb, "Data TLB"),
-    INIT_PARAM(system, "system object"),
     INIT_PARAM(cpu_id, "processor ID"),
     INIT_PARAM(profile, ""),
 #else
@@ -618,11 +616,11 @@ CREATE_SIM_OBJECT(TimingSimpleCPU)
     params->functionTrace = function_trace;
     params->functionTraceStart = function_trace_start;
     params->mem = mem;
+    params->system = system;
 
 #if FULL_SYSTEM
     params->itb = itb;
     params->dtb = dtb;
-    params->system = system;
     params->cpu_id = cpu_id;
     params->profile = profile;
 #else
diff --git a/src/cpu/simple/timing.hh b/src/cpu/simple/timing.hh
index 0a3f91e6c..ac36e5c99 100644
--- a/src/cpu/simple/timing.hh
+++ b/src/cpu/simple/timing.hh
@@ -137,9 +137,8 @@ class TimingSimpleCPU : public BaseSimpleCPU
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
-    virtual bool drain(Event *drain_event);
+    virtual unsigned int drain(Event *drain_event);
     virtual void resume();
-    virtual void setMemoryMode(State new_mode);
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
diff --git a/src/dev/ide_disk.cc b/src/dev/ide_disk.cc
index dc78021f8..12564ddd0 100644
--- a/src/dev/ide_disk.cc
+++ b/src/dev/ide_disk.cc
@@ -318,7 +318,7 @@ IdeDisk::doDmaTransfer()
         panic("Inconsistent DMA transfer state: dmaState = %d devState = %d\n",
               dmaState, devState);
 
-    if (ctrl->dmaPending()) {
+    if (ctrl->dmaPending() || ctrl->getState() != SimObject::Running) {
         dmaTransferEvent.schedule(curTick + DMA_BACKOFF_PERIOD);
         return;
     } else
@@ -398,8 +398,7 @@ IdeDisk::doDmaRead()
                 curPrd.getByteCount(), TheISA::PageBytes);
 
     }
-    if (ctrl->dmaPending()) {
-        panic("shouldn't be reentant??");
+    if (ctrl->dmaPending() || ctrl->getState() != SimObject::Running) {
         dmaReadWaitEvent.schedule(curTick + DMA_BACKOFF_PERIOD);
         return;
     } else if (!dmaReadCG->done()) {
@@ -474,8 +473,7 @@ IdeDisk::doDmaWrite()
         dmaWriteCG = new ChunkGenerator(curPrd.getBaseAddr(),
                 curPrd.getByteCount(), TheISA::PageBytes);
     }
-    if (ctrl->dmaPending()) {
-        panic("shouldn't be reentant??");
+    if (ctrl->dmaPending() || ctrl->getState() != SimObject::Running) {
         dmaWriteWaitEvent.schedule(curTick + DMA_BACKOFF_PERIOD);
         return;
     } else if (!dmaWriteCG->done()) {
diff --git a/src/dev/io_device.cc b/src/dev/io_device.cc
index cb4850108..ed2862065 100644
--- a/src/dev/io_device.cc
+++ b/src/dev/io_device.cc
@@ -32,10 +32,12 @@
 #include "base/trace.hh"
 #include "dev/io_device.hh"
 #include "sim/builder.hh"
+#include "sim/system.hh"
 
 
-PioPort::PioPort(PioDevice *dev, Platform *p, std::string pname)
-    : Port(dev->name() + pname), device(dev), platform(p)
+PioPort::PioPort(PioDevice *dev, System *s, std::string pname)
+    : Port(dev->name() + pname), device(dev), sys(s),
+      outTiming(0), drainEvent(NULL)
 { }
 
 
@@ -68,13 +70,23 @@ PioPort::recvRetry()
         if (result)
             transmitList.pop_front();
     }
+   if (transmitList.size() == 0 && drainEvent) {
+       drainEvent->process();
+       drainEvent = NULL;
+   }
 }
 
 void
 PioPort::SendEvent::process()
 {
+    port->outTiming--;
+    assert(port->outTiming >= 0);
     if (port->Port::sendTiming(packet))
-        return;
+       if (port->transmitList.size() == 0 && port->drainEvent) {
+           port->drainEvent->process();
+           port->drainEvent = NULL;
+       }
+       return;
 
     port->transmitList.push_back(packet);
 }
@@ -105,6 +117,15 @@ PioPort::recvTiming(Packet *pkt)
     return true;
 }
 
+unsigned int
+PioPort::drain(Event *de)
+{
+    if (outTiming == 0 && transmitList.size() == 0)
+        return 0;
+    drainEvent = de;
+    return 1;
+}
+
 PioDevice::~PioDevice()
 {
     if (pioPort)
@@ -119,6 +140,19 @@ PioDevice::init()
     pioPort->sendStatusChange(Port::RangeChange);
 }
 
+
+unsigned int
+PioDevice::drain(Event *de)
+{
+    unsigned int count;
+    count = pioPort->drain(de);
+    if (count)
+        changeState(Draining);
+    else
+        changeState(Drained);
+    return count;
+}
+
 void
 BasicPioDevice::addressRanges(AddrRangeList &range_list)
 {
@@ -128,8 +162,9 @@ BasicPioDevice::addressRanges(AddrRangeList &range_list)
 }
 
 
-DmaPort::DmaPort(DmaDevice *dev, Platform *p)
-    : Port(dev->name() + "-dmaport"), device(dev), platform(p), pendingCount(0)
+DmaPort::DmaPort(DmaDevice *dev, System *s)
+    : Port(dev->name() + "-dmaport"), device(dev), sys(s), pendingCount(0),
+      actionInProgress(0), drainEvent(NULL)
 { }
 
 bool
@@ -159,6 +194,11 @@ DmaPort::recvTiming(Packet *pkt)
         }
         delete pkt->req;
         delete pkt;
+
+        if (pendingCount == 0 && drainEvent) {
+            drainEvent->process();
+            drainEvent = NULL;
+        }
     }  else {
         panic("Got packet without sender state... huh?\n");
     }
@@ -170,6 +210,29 @@ DmaDevice::DmaDevice(Params *p)
     : PioDevice(p), dmaPort(NULL)
 { }
 
+
+unsigned int
+DmaDevice::drain(Event *de)
+{
+    unsigned int count;
+    count = pioPort->drain(de) + dmaPort->drain(de);
+    if (count)
+        changeState(Draining);
+    else
+        changeState(Drained);
+    return count;
+}
+
+unsigned int
+DmaPort::drain(Event *de)
+{
+    if (pendingCount == 0)
+        return 0;
+    drainEvent = de;
+    return 1;
+}
+
+
 void
 DmaPort::recvRetry()
 {
@@ -195,6 +258,8 @@ DmaPort::dmaAction(Packet::Command cmd, Addr addr, int size, Event *event,
 {
     assert(event);
 
+    assert(device->getState() == SimObject::Running);
+
     DmaReqState *reqState = new DmaReqState(event, this, size);
 
     for (ChunkGenerator gen(addr, size, peerBlockSize());
@@ -212,51 +277,53 @@ DmaPort::dmaAction(Packet::Command cmd, Addr addr, int size, Event *event,
             pendingCount++;
             sendDma(pkt);
     }
+
 }
 
 
 void
 DmaPort::sendDma(Packet *pkt, bool front)
 {
-   // some kind of selction between access methods
-   // more work is going to have to be done to make
-   // switching actually work
-  /* MemState state = device->platform->system->memState;
-
-   if (state == Timing) {  */
-       DPRINTF(DMA, "Attempting to send Packet %#x with addr: %#x\n",
-               pkt, pkt->getAddr());
-       if (transmitList.size() || !sendTiming(pkt)) {
-           if (front)
-               transmitList.push_front(pkt);
-           else
-               transmitList.push_back(pkt);
-           DPRINTF(DMA, "-- Failed: queued\n");
-       } else {
-           DPRINTF(DMA, "-- Done\n");
-       }
-  /*  } else if (state == Atomic) {
-       sendAtomic(pkt);
-       if (pkt->senderState) {
-           DmaReqState *state = dynamic_cast<DmaReqState*>(pkt->senderState);
-           assert(state);
-           state->completionEvent->schedule(curTick + (pkt->time -
-           pkt->req->getTime()) +1);
-           delete state;
-       }
-       pendingCount--;
-       assert(pendingCount >= 0);
-       delete pkt->req;
-       delete pkt;
-
-   } else if (state == Functional) {
-       sendFunctional(pkt);
-       // Is this correct???
-       completionEvent->schedule(pkt->req->responseTime - pkt->req->requestTime);
-       completionEvent == NULL;
+    // some kind of selction between access methods
+    // more work is going to have to be done to make
+    // switching actually work
+    System::MemoryMode state = sys->getMemoryMode();
+    if (state == System::Timing) {
+        DPRINTF(DMA, "Attempting to send Packet %#x with addr: %#x\n",
+                pkt, pkt->getAddr());
+        if (transmitList.size() || !sendTiming(pkt)) {
+            if (front)
+                transmitList.push_front(pkt);
+            else
+                transmitList.push_back(pkt);
+            DPRINTF(DMA, "-- Failed: queued\n");
+        } else {
+            DPRINTF(DMA, "-- Done\n");
+        }
+    } else if (state == System::Atomic) {
+        sendAtomic(pkt);
+        assert(pkt->senderState);
+        DmaReqState *state = dynamic_cast<DmaReqState*>(pkt->senderState);
+        assert(state);
+
+        state->numBytes += pkt->req->getSize();
+        if (state->totBytes == state->numBytes) {
+            state->completionEvent->schedule(curTick +
+                    (pkt->time - pkt->req->getTime()) +1);
+            delete state;
+            delete pkt->req;
+        }
+        pendingCount--;
+        assert(pendingCount >= 0);
+        delete pkt;
+
+        if (pendingCount == 0 && drainEvent) {
+            drainEvent->process();
+            drainEvent = NULL;
+        }
+
    } else
        panic("Unknown memory command state.");
-  */
 }
 
 DmaDevice::~DmaDevice()
diff --git a/src/dev/io_device.hh b/src/dev/io_device.hh
index 40edf6875..fa3f98247 100644
--- a/src/dev/io_device.hh
+++ b/src/dev/io_device.hh
@@ -60,9 +60,9 @@ class PioPort : public Port
     /** The device that this port serves. */
     PioDevice *device;
 
-    /** The platform that device/port are in. This is used to select which mode
+    /** The system that device/port are in. This is used to select which mode
      * we are currently operating in. */
-    Platform *platform;
+    System *sys;
 
     /** A list of outgoing timing response packets that haven't been serviced
      * yet. */
@@ -106,16 +106,27 @@ class PioPort : public Port
         friend class PioPort;
     };
 
+    /** Number of timing requests that are emulating the device timing before
+     * attempting to end up on the bus.
+     */
+    int outTiming;
+
+    /** If we need to drain, keep the drain event around until we're done
+     * here.*/
+    Event *drainEvent;
+
     /** Schedule a sendTiming() event to be called in the future. */
     void sendTiming(Packet *pkt, Tick time)
-    { new PioPort::SendEvent(this, pkt, time); }
+    { outTiming++; new PioPort::SendEvent(this, pkt, time); }
 
     /** This function is notification that the device should attempt to send a
      * packet again. */
     virtual void recvRetry();
 
   public:
-    PioPort(PioDevice *dev, Platform *p, std::string pname = "-pioport");
+    PioPort(PioDevice *dev, System *s, std::string pname = "-pioport");
+
+    unsigned int drain(Event *de);
 
   friend class PioPort::SendEvent;
 };
@@ -147,13 +158,20 @@ class DmaPort : public Port
     DmaDevice *device;
     std::list<Packet*> transmitList;
 
-    /** The platform that device/port are in. This is used to select which mode
+    /** The system that device/port are in. This is used to select which mode
      * we are currently operating in. */
-    Platform *platform;
+    System *sys;
 
     /** Number of outstanding packets the dma port has. */
     int pendingCount;
 
+    /** If a dmaAction is in progress. */
+    int actionInProgress;
+
+    /** If we need to drain, keep the drain event around until we're done
+     * here.*/
+    Event *drainEvent;
+
     virtual bool recvTiming(Packet *pkt);
     virtual Tick recvAtomic(Packet *pkt)
     { panic("dma port shouldn't be used for pio access."); }
@@ -171,13 +189,14 @@ class DmaPort : public Port
     void sendDma(Packet *pkt, bool front = false);
 
   public:
-    DmaPort(DmaDevice *dev, Platform *p);
+    DmaPort(DmaDevice *dev, System *s);
 
     void dmaAction(Packet::Command cmd, Addr addr, int size, Event *event,
                    uint8_t *data = NULL);
 
     bool dmaPending() { return pendingCount > 0; }
 
+    unsigned int drain(Event *de);
 };
 
 /**
@@ -196,6 +215,8 @@ class PioDevice : public MemObject
      * transaction we should perform. */
     Platform *platform;
 
+    System *sys;
+
     /** The pioPort that handles the requests for us and provides us requests
      * that it sees. */
     PioPort *pioPort;
@@ -240,20 +261,22 @@ class PioDevice : public MemObject
     const Params *params() const { return _params; }
 
     PioDevice(Params *p)
-              : MemObject(p->name),  platform(p->platform), pioPort(NULL),
-                _params(p)
+              : MemObject(p->name),  platform(p->platform), sys(p->system),
+              pioPort(NULL), _params(p)
               {}
 
     virtual ~PioDevice();
 
     virtual void init();
 
+    virtual unsigned int drain(Event *de);
+
     virtual Port *getPort(const std::string &if_name, int idx = -1)
     {
         if (if_name == "pio") {
             if (pioPort != NULL)
                 panic("pio port already connected to.");
-            pioPort = new PioPort(this, params()->platform);
+            pioPort = new PioPort(this, sys);
             return pioPort;
         } else
             return NULL;
@@ -310,17 +333,19 @@ class DmaDevice : public PioDevice
 
     bool dmaPending() { return dmaPort->dmaPending(); }
 
+    virtual unsigned int drain(Event *de);
+
     virtual Port *getPort(const std::string &if_name, int idx = -1)
     {
         if (if_name == "pio") {
             if (pioPort != NULL)
                 panic("pio port already connected to.");
-            pioPort = new PioPort(this, params()->platform);
+            pioPort = new PioPort(this, sys);
             return pioPort;
         } else if (if_name == "dma") {
             if (dmaPort != NULL)
                 panic("dma port already connected to.");
-            dmaPort = new DmaPort(this, params()->platform);
+            dmaPort = new DmaPort(this, sys);
             return dmaPort;
         } else
             return NULL;
diff --git a/src/dev/ns_gige.cc b/src/dev/ns_gige.cc
index 179a2c62d..bf2279d93 100644
--- a/src/dev/ns_gige.cc
+++ b/src/dev/ns_gige.cc
@@ -1377,7 +1377,7 @@ NSGigE::doRxDmaRead()
     assert(rxDmaState == dmaIdle || rxDmaState == dmaReadWaiting);
     rxDmaState = dmaReading;
 
-    if (dmaPending())
+    if (dmaPending() || getState() != Running)
         rxDmaState = dmaReadWaiting;
     else
         dmaRead(rxDmaAddr, rxDmaLen, &rxDmaReadEvent, (uint8_t*)rxDmaData);
@@ -1408,7 +1408,7 @@ NSGigE::doRxDmaWrite()
     assert(rxDmaState == dmaIdle || rxDmaState == dmaWriteWaiting);
     rxDmaState = dmaWriting;
 
-    if (dmaPending())
+    if (dmaPending() || getState() != Running)
         rxDmaState = dmaWriteWaiting;
     else
         dmaWrite(rxDmaAddr, rxDmaLen, &rxDmaWriteEvent, (uint8_t*)rxDmaData);
@@ -1826,7 +1826,7 @@ NSGigE::doTxDmaRead()
     assert(txDmaState == dmaIdle || txDmaState == dmaReadWaiting);
     txDmaState = dmaReading;
 
-    if (dmaPending())
+    if (dmaPending() || getState() != Running)
         txDmaState = dmaReadWaiting;
     else
         dmaRead(txDmaAddr, txDmaLen, &txDmaReadEvent, (uint8_t*)txDmaData);
@@ -1857,7 +1857,7 @@ NSGigE::doTxDmaWrite()
     assert(txDmaState == dmaIdle || txDmaState == dmaWriteWaiting);
     txDmaState = dmaWriting;
 
-    if (dmaPending())
+    if (dmaPending() || getState() != Running)
         txDmaState = dmaWriteWaiting;
     else
         dmaWrite(txDmaAddr, txDmaLen, &txDmaWriteEvent, (uint8_t*)txDmaData);
@@ -2406,6 +2406,20 @@ NSGigE::recvPacket(EthPacketPtr packet)
     return true;
 }
 
+
+void
+NSGigE::resume()
+{
+    SimObject::resume();
+
+    // During drain we could have left the state machines in a waiting state and
+    // they wouldn't get out until some other event occured to kick them.
+    // This way they'll get out immediately
+    txKick();
+    rxKick();
+}
+
+
 //=====================================================================
 //
 //
diff --git a/src/dev/ns_gige.hh b/src/dev/ns_gige.hh
index ea7243777..080c0b1f3 100644
--- a/src/dev/ns_gige.hh
+++ b/src/dev/ns_gige.hh
@@ -391,6 +391,8 @@ class NSGigE : public PciDev
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
+    virtual void resume();
+
   public:
     void regStats();
 
diff --git a/src/dev/pcidev.cc b/src/dev/pcidev.cc
index 62a7324ad..e81e0d1ee 100644
--- a/src/dev/pcidev.cc
+++ b/src/dev/pcidev.cc
@@ -56,8 +56,8 @@ using namespace std;
 
 PciDev::PciConfigPort::PciConfigPort(PciDev *dev, int busid, int devid,
         int funcid, Platform *p)
-        : PioPort(dev,p,"-pciconf"), device(dev), busId(busid), deviceId(devid),
-        functionId(funcid)
+        : PioPort(dev,p->system,"-pciconf"), device(dev), platform(p),
+          busId(busid), deviceId(devid), functionId(funcid)
 {
     configAddr = platform->calcConfigAddr(busId, deviceId, functionId);
 }
@@ -132,6 +132,18 @@ PciDev::init()
    PioDevice::init();
 }
 
+unsigned int
+PciDev::drain(Event *de)
+{
+    unsigned int count;
+    count = pioPort->drain(de) + dmaPort->drain(de) + configPort->drain(de);
+    if (count)
+        changeState(Draining);
+    else
+        changeState(Drained);
+    return count;
+}
+
 Tick
 PciDev::readConfig(Packet *pkt)
 {
diff --git a/src/dev/pcidev.hh b/src/dev/pcidev.hh
index 20ab9364a..847fb07d0 100644
--- a/src/dev/pcidev.hh
+++ b/src/dev/pcidev.hh
@@ -95,6 +95,8 @@ class PciDev : public DmaDevice
 
         virtual void getDeviceAddressRanges(AddrRangeList &resp, AddrRangeList &snoop);
 
+        Platform *platform;
+
         int busId;
         int deviceId;
         int functionId;
@@ -249,6 +251,9 @@ class PciDev : public DmaDevice
      */
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
+
+    virtual unsigned int drain(Event *de);
+
     virtual Port *getPort(const std::string &if_name, int idx = -1)
     {
         if (if_name == "config") {
diff --git a/src/dev/sinic.cc b/src/dev/sinic.cc
index dddda1f1c..815cecca5 100644
--- a/src/dev/sinic.cc
+++ b/src/dev/sinic.cc
@@ -921,7 +921,7 @@ Device::rxKick()
         break;
 
       case rxBeginCopy:
-        if (dmaPending())
+        if (dmaPending() || getState() != Running)
             goto exit;
 
         rxDmaAddr = params()->platform->pciToDma(
@@ -1109,7 +1109,7 @@ Device::txKick()
         break;
 
       case txBeginCopy:
-        if (dmaPending())
+        if (dmaPending() || getState() != Running)
             goto exit;
 
         txDmaAddr = params()->platform->pciToDma(
@@ -1287,6 +1287,18 @@ Device::recvPacket(EthPacketPtr packet)
     return true;
 }
 
+void
+Device::resume()
+{
+    SimObject::resume();
+
+    // During drain we could have left the state machines in a waiting state and
+    // they wouldn't get out until some other event occured to kick them.
+    // This way they'll get out immediately
+    txKick();
+    rxKick();
+}
+
 //=====================================================================
 //
 //
diff --git a/src/dev/sinic.hh b/src/dev/sinic.hh
index f6c229039..eece4ba6b 100644
--- a/src/dev/sinic.hh
+++ b/src/dev/sinic.hh
@@ -266,6 +266,7 @@ class Device : public Base
   public:
     virtual Tick read(Packet *pkt);
     virtual Tick write(Packet *pkt);
+    virtual void resume();
 
     void prepareIO(int cpu, int index);
     void prepareRead(int cpu, int index);
diff --git a/src/python/m5/config.py b/src/python/m5/config.py
index 8291e1e1b..8eed28dcc 100644
--- a/src/python/m5/config.py
+++ b/src/python/m5/config.py
@@ -547,8 +547,7 @@ class SimObject(object):
         count = 0
         # ParamContexts don't serialize
         if isinstance(self, SimObject) and not isinstance(self, ParamContext):
-            if not self._ccObject.drain(drain_event):
-                count = 1
+            count += self._ccObject.drain(drain_event)
         if recursive:
             for child in self._children.itervalues():
                 count += child.startDrain(drain_event, True)
@@ -561,7 +560,7 @@ class SimObject(object):
             child.resume()
 
     def changeTiming(self, mode):
-        if isinstance(self, SimObject) and not isinstance(self, ParamContext):
+        if isinstance(self, System):
             self._ccObject.setMemoryMode(mode)
         for child in self._children.itervalues():
             child.changeTiming(mode)
diff --git a/src/python/m5/objects/BaseCPU.py b/src/python/m5/objects/BaseCPU.py
index 2e78578df..5bf98be9c 100644
--- a/src/python/m5/objects/BaseCPU.py
+++ b/src/python/m5/objects/BaseCPU.py
@@ -6,10 +6,10 @@ class BaseCPU(SimObject):
     abstract = True
     mem = Param.MemObject("memory")
 
+    system = Param.System(Parent.any, "system object")
     if build_env['FULL_SYSTEM']:
         dtb = Param.AlphaDTB("Data TLB")
         itb = Param.AlphaITB("Instruction TLB")
-        system = Param.System(Parent.any, "system object")
         cpu_id = Param.Int(-1, "CPU identifier")
     else:
         workload = VectorParam.Process("processes to run")
diff --git a/src/sim/sim_object.cc b/src/sim/sim_object.cc
index a0278dba0..d12b06b7a 100644
--- a/src/sim/sim_object.cc
+++ b/src/sim/sim_object.cc
@@ -72,7 +72,7 @@ SimObject::SimObject(Params *p)
 
     doRecordEvent = !Stats::event_ignore.match(name());
     simObjectList.push_back(this);
-    state = Atomic;
+    state = Running;
 }
 
 //
@@ -88,7 +88,7 @@ SimObject::SimObject(const string &_name)
 
     doRecordEvent = !Stats::event_ignore.match(name());
     simObjectList.push_back(this);
-    state = Atomic;
+    state = Running;
 }
 
 void
@@ -269,38 +269,23 @@ SimObject::recordEvent(const std::string &stat)
         Stats::recordEvent(stat);
 }
 
-bool
+unsigned int
 SimObject::drain(Event *drain_event)
 {
-    if (state != DrainedAtomic && state != Atomic) {
-        panic("Must implement your own drain function if it is to be used "
-              "in timing mode!");
-    }
-    state = DrainedAtomic;
-    return true;
+    state = Drained;
+    return 0;
 }
 
 void
 SimObject::resume()
 {
-    if (state == DrainedAtomic) {
-        state = Atomic;
-    } else if (state == DrainedTiming) {
-        state = Timing;
-    }
+    state = Running;
 }
 
 void
 SimObject::setMemoryMode(State new_mode)
 {
-    assert(new_mode == Timing || new_mode == Atomic);
-    if (state == DrainedAtomic && new_mode == Timing) {
-        state = DrainedTiming;
-    } else if (state == DrainedTiming && new_mode == Atomic) {
-        state = DrainedAtomic;
-    } else {
-        state = new_mode;
-    }
+    panic("setMemoryMode() should only be called on systems");
 }
 
 void
diff --git a/src/sim/sim_object.hh b/src/sim/sim_object.hh
index 7ecc00958..38f2bdd23 100644
--- a/src/sim/sim_object.hh
+++ b/src/sim/sim_object.hh
@@ -60,16 +60,15 @@ class SimObject : public Serializable, protected StartupCallback
     };
 
     enum State {
-        Atomic,
-        Timing,
+        Running,
         Draining,
-        DrainedAtomic,
-        DrainedTiming
+        Drained
     };
+  private:
+    State state;
 
   protected:
     Params *_params;
-    State state;
 
     void changeState(State new_state) { state = new_state; }
 
@@ -116,8 +115,10 @@ class SimObject : public Serializable, protected StartupCallback
 
     // Methods to drain objects in order to take checkpoints
     // Or switch from timing -> atomic memory model
-    // Drain returns false if the SimObject cannot drain immediately.
-    virtual bool drain(Event *drain_event);
+    // Drain returns 0 if the simobject can drain immediately or
+    // the number of times the drain_event's process function will be called
+    // before the object will be done draining. Normally this should be 1
+    virtual unsigned int drain(Event *drain_event);
     virtual void resume();
     virtual void setMemoryMode(State new_mode);
     virtual void switchOut();
diff --git a/src/sim/system.cc b/src/sim/system.cc
index 89e7b8542..2780d0fda 100644
--- a/src/sim/system.cc
+++ b/src/sim/system.cc
@@ -143,6 +143,14 @@ int rgdb_wait = -1;
 
 #endif // FULL_SYSTEM
 
+
+void
+System::setMemoryMode(MemoryMode mode)
+{
+    assert(getState() == Drained);
+    memoryMode = mode;
+}
+
 int
 System::registerThreadContext(ThreadContext *tc, int id)
 {
diff --git a/src/sim/system.hh b/src/sim/system.hh
index 059dc92dc..77777aed4 100644
--- a/src/sim/system.hh
+++ b/src/sim/system.hh
@@ -61,6 +61,21 @@ class RemoteGDB;
 class System : public SimObject
 {
   public:
+    enum MemoryMode {
+        Invalid=0,
+        Atomic,
+        Timing
+    };
+
+
+    MemoryMode getMemoryMode() { assert(memoryMode); return memoryMode; }
+
+    /** Change the memory mode of the system. This should only be called by the
+     * python!!
+     * @param mode Mode to change to (atomic/timing)
+     */
+    void setMemoryMode(MemoryMode mode);
+
     PhysicalMemory *physmem;
     PCEventQueue pcEventQueue;
 
@@ -108,6 +123,8 @@ class System : public SimObject
 
   protected:
 
+    MemoryMode memoryMode;
+
 #if FULL_SYSTEM
     /**
      * Fix up an address used to match PCs for hooking simulator