Add initial Mercury TLS implementation

Currently this builds but it untested.
jmlapre · Jan 13, 2025 · 4d4e04e · 4d4e04e
1 parent 659211a
commit 4d4e04e
Show file tree

Hide file tree

Showing 4 changed files with 132 additions and 48 deletions.
diff --git a/src/sst/elements/mercury/components/operating_system.cc b/src/sst/elements/mercury/components/operating_system.cc
@@ -183,7 +183,7 @@ OperatingSystem::startThread(Thread* t)
             stack,
             StackAlloc::stacksize(),
             parent->globalsStorage(),
-            nullptr);
+            parent->newTlsStorage());
     }
   running_threads_[t->tid()] = t;
 }

diff --git a/src/sst/elements/mercury/operating_system/process/app.cc b/src/sst/elements/mercury/operating_system/process/app.cc
@@ -75,28 +75,38 @@ App::allocateTlsKey(destructor_fxn fxn)
   return next;
 }
 
-//static char* get_data_segment(SST::Params& params,
-//                             const char* param_name, GlobalVariableContext& ctx)
-//{
-//  int allocSize = ctx.allocSize();
-//  if (params.contains(param_name)){
-//    allocSize = params.find<int>(param_name);
-//    if (ctx.allocSize() != allocSize){
-//      ctx.setAllocSize(allocSize);
-//    }
-//  }
-//  if (allocSize != 0){
-//    char* segment = new char[allocSize];
-//    ::memcpy(segment, ctx.globalInit(), ctx.globalsSize());
-//    return segment;
-//  } else {
-//    return nullptr;
-//  }
-//}
+static char* get_data_segment(SST::Params& params,
+                            const char* param_name, GlobalVariableContext& ctx)
+{
+ int allocSize = ctx.allocSize();
+ if (params.contains(param_name)){
+   allocSize = params.find<int>(param_name);
+   if (ctx.allocSize() != allocSize){
+     ctx.setAllocSize(allocSize);
+   }
+ }
+ if (allocSize != 0){
+   char* segment = new char[allocSize];
+   ::memcpy(segment, ctx.globalInit(), ctx.globalsSize());
+   return segment;
+ } else {
+   return nullptr;
+ }
+}
 
 
 static thread_lock dlopen_lock;
 
+void
+GlobalVariableContext::callInitFxns(void *globals)
+{
+  for (auto& pair : initFxns){
+    int offset = pair.first;
+    char* ptr = ((char*)globals) + offset;
+    (pair.second)(ptr);
+  }
+}
+
 void
 App::lockDlopen(int aid)
 {
@@ -226,15 +236,15 @@ App::dlcloseCheck_Library(std::string api_name)
   dlopen_lock.unlock();
 }
 
-//char*
-//App::allocateDataSegment(bool tls)
-//{
-//  if (tls){
-//    return get_data_segment(params_, "tls_size", GlobalVariable::tlsCtx);
-//  } else {
-//    return get_data_segment(params_, "globals_size", GlobalVariable::glblCtx);
-//  }
-//}
+char*
+App::allocateDataSegment(bool tls)
+{
+ if (tls){
+   return get_data_segment(params_, "tls_size", GlobalVariable::tlsCtx);
+ } else {
+   return get_data_segment(params_, "globals_size", GlobalVariable::glblCtx);
+ }
+}
 
 App::App(SST::Params& params, SoftwareId sid,
          OperatingSystem* os) :

diff --git a/src/sst/elements/mercury/operating_system/process/app.h b/src/sst/elements/mercury/operating_system/process/app.h
@@ -36,6 +36,71 @@
 namespace SST {
 namespace Hg {
 
+class GlobalVariableContext {
+ public:
+  void init();
+
+  ~GlobalVariableContext();
+
+  int append(const int size, const char* name);
+
+  int globalsSize() {
+    return stackOffset;
+  }
+
+  int allocSize() {
+    return allocSize_;
+  }
+
+  void setAllocSize(int sz){
+    allocSize_ = sz;
+  }
+
+  void* globalInit() {
+    return globalInits;
+  }
+
+  void addActiveSegment(void* globals){
+    activeGlobalMaps_.insert(globals);
+  }
+
+  void removeActiveSegment(void* globals){
+    activeGlobalMaps_.erase(globals);
+  }
+
+  void initGlobalSpace(void* ptr, int size, int offset);
+
+  void callInitFxns(void* globals);
+
+  void unregisterInitFxn(int offset){
+    initFxns.erase(offset);
+  }
+
+  void registerInitFxn(int offset, std::function<void(void*)>&& fxn);
+
+ private:
+  int stackOffset;
+  char* globalInits;
+  int allocSize_;
+  //these should be ordered by the offset in the data segment
+  //this ensures as much as possible that global variables
+  //are initialized in the same order in SST/macro as they would be in the real app
+  std::map<int, std::function<void(void*)>> initFxns;
+
+ private:
+  std::unordered_set<void*> activeGlobalMaps_;
+
+};
+
+class GlobalVariable {
+ public:
+  static int init(const int size, const char* name, bool tls = false);
+
+  static GlobalVariableContext glblCtx;
+  static GlobalVariableContext tlsCtx;
+  static bool inited;
+};
+
 /**
  * The app derived class adds to the thread base class by providing
  * facilities to allow applications to simulate computation.
@@ -160,9 +225,9 @@ class App : public Thread
     return globals_storage_;
   }
 
-//  void* newTlsStorage() {
-//    return allocateDataSegment(true);
-//  }
+ void* newTlsStorage() {
+   return allocateDataSegment(true);
+ }
 
   const std::string& uniqueName() const {
     return unique_name_;

diff --git a/src/sst/elements/mercury/operating_system/process/thread_info.cc b/src/sst/elements/mercury/operating_system/process/thread_info.cc
@@ -16,6 +16,7 @@
 #include <mercury/common/errors.h>
 #include <mercury/components/operating_system.h>
 #include <mercury/operating_system/process/thread_info.h>
+#include <mercury/operating_system/process/app.h>
 #include <mercury/operating_system/process/tls.h>
 #include <mercury/operating_system/threading/thread_lock.h>
 #include <mercury/operating_system/threading/stack_alloc.h>
@@ -103,23 +104,31 @@ ThreadInfo::registerUserSpaceVirtualThread(int phys_thread_id, void *stack,
 //  void** currentGlobalsPtr = (void**) &fake_globals[0];
 //  void** currentTlsPtr = (void**) &fake_tls[0];
 
-//  globals_lock.lock();
-//  if (globalsMap && isAppStartup){
-//    void* currentGlobals = *currentGlobalsPtr;
-//    *currentGlobalsPtr = globalsMap;
-//    GlobalVariable::glblCtx.addActiveSegment(globalsMap);
-//    GlobalVariable::glblCtx.callInitFxns(globalsMap);
-//    *currentGlobalsPtr = currentGlobals;
-//  }
-
-//  if (tlsMap && isThreadStartup){
-//    void* currentTls = *currentTlsPtr;
-//    *currentTlsPtr = tlsMap;
-//    GlobalVariable::tlsCtx.addActiveSegment(tlsMap);
-//    GlobalVariable::tlsCtx.callInitFxns(tlsMap);
-//    *currentTlsPtr = currentTls;
-//  }
-//  globals_lock.unlock();
+  int activeStack; int* activeStackPtr = &activeStack;
+  intptr_t stackTopInt = sst_hg_global_stacksize //avoid errors if this is zero
+      ? ((intptr_t)activeStackPtr/sst_hg_global_stacksize)*sst_hg_global_stacksize
+      : 0;
+
+  void** currentGlobalsPtr = (void**)(stackTopInt + SST_HG_TLS_GLOBAL_MAP);
+  void** currentTlsPtr = (void**)(stackTopInt + SST_HG_TLS_GLOBAL_MAP);
+
+ globals_lock.lock();
+ if (globalsMap && isAppStartup){
+   void* currentGlobals = *currentGlobalsPtr;
+   *currentGlobalsPtr = globalsMap;
+   GlobalVariable::glblCtx.addActiveSegment(globalsMap);
+   GlobalVariable::glblCtx.callInitFxns(globalsMap);
+   *currentGlobalsPtr = currentGlobals;
+ }
+
+ if (tlsMap && isThreadStartup){
+   void* currentTls = *currentTlsPtr;
+   *currentTlsPtr = tlsMap;
+   GlobalVariable::tlsCtx.addActiveSegment(tlsMap);
+   GlobalVariable::tlsCtx.callInitFxns(tlsMap);
+   *currentTlsPtr = currentTls;
+ }
+ globals_lock.unlock();
 }
 
 void