4 files changed, 19 insertions, 4 deletions
diff --git a/src/python/m5/simulate.py b/src/python/m5/simulate.py
index 0cf0a254e..cd2f8bb64 100644
--- a/src/python/m5/simulate.py
+++ b/src/python/m5/simulate.py
@@ -88,8 +88,12 @@ def instantiate(ckpt_dir=None):
 
     # Restore checkpoint (if any)
     if ckpt_dir:
-        internal.core.unserializeAll(ckpt_dir)
+        ckpt = internal.core.getCheckpoint(ckpt_dir)
+        internal.core.unserializeGlobals(ckpt);
+        for obj in root.descendants(): obj.loadState(ckpt)
         need_resume.append(root)
+    else:
+        for obj in root.descendants(): obj.initState()
 
     # Reset to put the stats in a consistent state.
     stats.reset()
diff --git a/src/python/swig/core.i b/src/python/swig/core.i
index 81085dd06..f48fe9590 100644
--- a/src/python/swig/core.i
+++ b/src/python/swig/core.i
@@ -75,8 +75,11 @@ void setClockFrequency(Tick ticksPerSecond);
 %immutable curTick;
 Tick curTick;
 
+class Checkpoint;
+
 void serializeAll(const std::string &cpt_dir);
-void unserializeAll(const std::string &cpt_dir);
+Checkpoint *getCheckpoint(const std::string &cpt_dir);
+void unserializeGlobals(Checkpoint *cp);
 
 bool want_warn, warn_verbose;
 bool want_info, info_verbose;
diff --git a/src/python/swig/pyobject.hh b/src/python/swig/pyobject.hh
index a27080d08..b18a2a76c 100644
--- a/src/python/swig/pyobject.hh
+++ b/src/python/swig/pyobject.hh
@@ -52,8 +52,14 @@ serializeAll(const std::string &cpt_dir)
     Serializable::serializeAll(cpt_dir);
 }
 
+inline Checkpoint *
+getCheckpoint(const std::string &cpt_dir)
+{
+    return new Checkpoint(cpt_dir);
+}
+
 inline void
-unserializeAll(const std::string &cpt_dir)
+unserializeGlobals(Checkpoint *cp)
 {
-    Serializable::unserializeAll(cpt_dir);
+    Serializable::unserializeGlobals(cp);
 }
diff --git a/src/python/swig/sim_object.i b/src/python/swig/sim_object.i
index 8cd8e8beb..af9afd057 100644
--- a/src/python/swig/sim_object.i
+++ b/src/python/swig/sim_object.i
@@ -51,6 +51,8 @@ class SimObject {
     };
 
     void init();
+    void loadState(Checkpoint *cp);
+    void initState();
     void regStats();
     void regFormulas();
     void resetStats();