sim: Track ThreadContext-s in the Workload object.

Change-Id: I00bf9fa36d3993f55d41e50196ad8a89a3d506c4
Reviewed-on: https://gem5-review.googlesource.com/c/public/gem5/+/44616
Maintainer: Gabe Black <gabe.black@gmail.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Boris Shingarov <shingarov@gmail.com>

src/sim/SConscript

@@ -83,6 +83,7 @@ Source('mathexpr.cc')
 Source('power_state.cc')
 Source('power_domain.cc')
 Source('stats.cc')
+Source('workload.cc')
 
 GTest('byteswap.test', 'byteswap.test.cc', '../base/types.cc')
 GTest('guest_abi.test', 'guest_abi.test.cc')

src/sim/system.cc

@@ -301,6 +301,9 @@ System::registerThreadContext(ThreadContext *tc, ContextID assigned)
 {
     threads.insert(tc, assigned);
 
+    if (workload)
+        workload->registerThreadContext(tc);
+
     for (auto *e: liveEvents)
         tc->schedule(e);
 }
@@ -331,6 +334,9 @@ System::replaceThreadContext(ThreadContext *tc, ContextID context_id)
     auto *otc = threads[context_id];
     threads.replace(tc, context_id);
 
+    if (workload)
+        workload->replaceThreadContext(tc);
+
     for (auto *e: liveEvents) {
         otc->remove(e);
         tc->schedule(e);

src/sim/workload.cc

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2021 Google Inc.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#include "sim/workload.hh"
+
+#include "cpu/thread_context.hh"
+
+void
+Workload::registerThreadContext(ThreadContext *tc)
+{
+    std::set<ThreadContext *>::iterator it;
+    bool success;
+    std::tie(it, success) = threads.insert(tc);
+    panic_if(!success, "Failed to add thread context %d.",
+            tc->contextId());
+}
+
+void
+Workload::replaceThreadContext(ThreadContext *tc)
+{
+    ContextID id = tc->contextId();
+
+    for (auto *old: threads) {
+        if (old->contextId() != id)
+            continue;
+        threads.erase(old);
+
+        std::set<ThreadContext *>::iterator it;
+        bool success;
+        std::tie(it, success) = threads.insert(tc);
+        panic_if(!success,
+                "Failed to insert replacement thread context %d.", id);
+        return;
+    }
+    panic("Replacement thread context %d doesn't match any known id.", id);
+}

src/sim/workload.hh

@@ -28,6 +28,9 @@
 #ifndef __SIM_WORKLOAD_HH__
 #define __SIM_WORKLOAD_HH__
 
+#include <set>
+#include <tuple>
+
 #include "base/loader/object_file.hh"
 #include "base/loader/symtab.hh"
 #include "params/Workload.hh"
@@ -63,6 +66,8 @@ class Workload : public SimObject
         {}
     } stats;
 
+    std::set<ThreadContext *> threads;
+
   public:
     Workload(const WorkloadParams &params) : SimObject(params), stats(this)
     {}
@@ -72,6 +77,9 @@ class Workload : public SimObject
 
     System *system = nullptr;
 
+    virtual void registerThreadContext(ThreadContext *tc);
+    virtual void replaceThreadContext(ThreadContext *tc);
+
     virtual Addr getEntry() const = 0;
     virtual Loader::Arch getArch() const = 0;