3 files changed, 25 insertions, 9 deletions

diff --git a/src/cpu/base.cc b/src/cpu/base.cc
index 1fa7add65..167606135 100644
--- a/src/cpu/base.cc
+++ b/src/cpu/base.cc
@@ -285,7 +285,17 @@ BaseCPU::registerThreadContexts()
     for (int i = 0; i < threadContexts.size(); ++i) {
         ThreadContext *tc = threadContexts[i];
 
-        tc->setContextId(system->registerThreadContext(tc));
+        /** This is so that contextId and cpuId match where there is a
+         * 1cpu:1context relationship.  Otherwise, the order of registration
+         * could affect the assignment and cpu 1 could have context id 3, for
+         * example.  We may even want to do something like this for SMT so that
+         * cpu 0 has the lowest thread contexts and cpu N has the highest, but
+         * I'll just do this for now
+         */
+        if (number_of_threads == 1)
+            tc->setContextId(system->registerThreadContext(tc, _cpuId));
+        else
+            tc->setContextId(system->registerThreadContext(tc));
 #if !FULL_SYSTEM
         tc->getProcessPtr()->assignThreadContext(tc->contextId());
 #endif
diff --git a/src/sim/system.cc b/src/sim/system.cc
index 9704c83f0..29b1d1f61 100644
--- a/src/sim/system.cc
+++ b/src/sim/system.cc
@@ -166,16 +166,22 @@ bool System::breakpoint()
 }
 
 int
-System::registerThreadContext(ThreadContext *tc)
+System::registerThreadContext(ThreadContext *tc, int assigned)
 {
     int id;
-    for (id = 0; id < threadContexts.size(); id++) {
-        if (!threadContexts[id])
-            break;
-    }
+    if (assigned == -1) {
+        for (id = 0; id < threadContexts.size(); id++) {
+            if (!threadContexts[id])
+                break;
+        }
 
-    if (threadContexts.size() <= id)
-        threadContexts.resize(id + 1);
+        if (threadContexts.size() <= id)
+            threadContexts.resize(id + 1);
+    } else {
+        if (threadContexts.size() <= assigned)
+            threadContexts.resize(assigned + 1);
+        id = assigned;
+    }
 
     if (threadContexts[id])
         panic("Cannot have two CPUs with the same id (%d)\n", id);
diff --git a/src/sim/system.hh b/src/sim/system.hh
index e993a7a50..9715767b1 100644
--- a/src/sim/system.hh
+++ b/src/sim/system.hh
@@ -224,7 +224,7 @@ class System : public SimObject
 
 #endif // FULL_SYSTEM
 
-    int registerThreadContext(ThreadContext *tc);
+    int registerThreadContext(ThreadContext *tc, int assigned=-1);
     void replaceThreadContext(ThreadContext *tc, int context_id);
 
     void serialize(std::ostream &os);