4 files changed, 40 insertions, 35 deletions

diff --git a/src/cpu/simple/atomic.cc b/src/cpu/simple/atomic.cc
index df7e780e6..ca4627bbf 100644
--- a/src/cpu/simple/atomic.cc
+++ b/src/cpu/simple/atomic.cc
@@ -126,6 +126,17 @@ AtomicSimpleCPU::CpuPort::recvRetry()
     panic("AtomicSimpleCPU doesn't expect recvRetry callback!");
 }
 
+void
+AtomicSimpleCPU::DcachePort::setPeer(Port *port)
+{
+    Port::setPeer(port);
+
+#if FULL_SYSTEM
+    // Update the ThreadContext's memory ports (Functional/Virtual
+    // Ports)
+    cpu->tcBase()->connectMemPorts();
+#endif
+}
 
 AtomicSimpleCPU::AtomicSimpleCPU(Params *p)
     : BaseSimpleCPU(p), tickEvent(this),
@@ -211,7 +222,7 @@ AtomicSimpleCPU::switchOut()
 void
 AtomicSimpleCPU::takeOverFrom(BaseCPU *oldCPU)
 {
-    BaseCPU::takeOverFrom(oldCPU);
+    BaseCPU::takeOverFrom(oldCPU, &icachePort, &dcachePort);
 
     assert(!tickEvent.scheduled());
 
@@ -242,12 +253,6 @@ AtomicSimpleCPU::activateContext(int thread_num, int delay)
 
     notIdleFraction++;
 
-#if FULL_SYSTEM
-    // Connect the ThreadContext's memory ports (Functional/Virtual
-    // Ports)
-    tc->connectMemPorts();
-#endif
-
     //Make sure ticks are still on multiples of cycles
     tickEvent.schedule(nextCycle(curTick + cycles(delay)));
     _status = Running;
diff --git a/src/cpu/simple/atomic.hh b/src/cpu/simple/atomic.hh
index 5bffb7666..ad4aa4708 100644
--- a/src/cpu/simple/atomic.hh
+++ b/src/cpu/simple/atomic.hh
@@ -81,9 +81,6 @@ class AtomicSimpleCPU : public BaseSimpleCPU
 
     class CpuPort : public Port
     {
-
-        AtomicSimpleCPU *cpu;
-
       public:
 
         CpuPort(const std::string &_name, AtomicSimpleCPU *_cpu)
@@ -94,6 +91,8 @@ class AtomicSimpleCPU : public BaseSimpleCPU
 
       protected:
 
+        AtomicSimpleCPU *cpu;
+
         virtual bool recvTiming(PacketPtr pkt);
 
         virtual Tick recvAtomic(PacketPtr pkt);
@@ -110,7 +109,17 @@ class AtomicSimpleCPU : public BaseSimpleCPU
 
     };
     CpuPort icachePort;
-    CpuPort dcachePort;
+
+    class DcachePort : public CpuPort
+    {
+      public:
+        DcachePort(const std::string &_name, AtomicSimpleCPU *_cpu)
+            : CpuPort(_name, _cpu)
+        { }
+
+        virtual void setPeer(Port *port);
+    };
+    DcachePort dcachePort;
 
     Request  *ifetch_req;
     PacketPtr ifetch_pkt;
diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index 7f857c68d..6f3604678 100644
--- a/src/cpu/simple/timing.cc
+++ b/src/cpu/simple/timing.cc
@@ -194,7 +194,7 @@ TimingSimpleCPU::switchOut()
 void
 TimingSimpleCPU::takeOverFrom(BaseCPU *oldCPU)
 {
-    BaseCPU::takeOverFrom(oldCPU);
+    BaseCPU::takeOverFrom(oldCPU, &icachePort, &dcachePort);
 
     // if any of this CPU's ThreadContexts are active, mark the CPU as
     // running and schedule its tick event.
@@ -209,23 +209,6 @@ TimingSimpleCPU::takeOverFrom(BaseCPU *oldCPU)
     if (_status != Running) {
         _status = Idle;
     }
-
-    Port *peer;
-    if (icachePort.getPeer() == NULL) {
-        peer = oldCPU->getPort("icache_port")->getPeer();
-        icachePort.setPeer(peer);
-    } else {
-        peer = icachePort.getPeer();
-    }
-    peer->setPeer(&icachePort);
-
-    if (dcachePort.getPeer() == NULL) {
-        peer = oldCPU->getPort("dcache_port")->getPeer();
-        dcachePort.setPeer(peer);
-    } else {
-        peer = dcachePort.getPeer();
-    }
-    peer->setPeer(&dcachePort);
 }
 
 
@@ -240,12 +223,6 @@ TimingSimpleCPU::activateContext(int thread_num, int delay)
     notIdleFraction++;
     _status = Running;
 
-#if FULL_SYSTEM
-    // Connect the ThreadContext's memory ports (Functional/Virtual
-    // Ports)
-    tc->connectMemPorts();
-#endif
-
     // kick things off by initiating the fetch of the next instruction
     fetchEvent =
         new EventWrapper<TimingSimpleCPU, &TimingSimpleCPU::fetch>(this, false);
@@ -649,6 +626,18 @@ TimingSimpleCPU::completeDrain()
     drainEvent->process();
 }
 
+void
+TimingSimpleCPU::DcachePort::setPeer(Port *port)
+{
+    Port::setPeer(port);
+
+#if FULL_SYSTEM
+    // Update the ThreadContext's memory ports (Functional/Virtual
+    // Ports)
+    cpu->tcBase()->connectMemPorts();
+#endif
+}
+
 bool
 TimingSimpleCPU::DcachePort::recvTiming(PacketPtr pkt)
 {
diff --git a/src/cpu/simple/timing.hh b/src/cpu/simple/timing.hh
index abcb224bf..ef062d24a 100644
--- a/src/cpu/simple/timing.hh
+++ b/src/cpu/simple/timing.hh
@@ -144,6 +144,8 @@ class TimingSimpleCPU : public BaseSimpleCPU
             : CpuPort(_cpu->name() + "-dport", _cpu, _lat), tickEvent(_cpu)
         { }
 
+        virtual void setPeer(Port *port);
+
       protected:
 
         virtual bool recvTiming(PacketPtr pkt);