[GraphBolt] Per-thread RandomEngine initialization fix. (#7557)

dmlc · Jul 22, 2024 · c83acdd · c83acdd
1 parent 5b4635a
commit c83acdd
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 7 deletions.
diff --git a/graphbolt/CMakeLists.txt b/graphbolt/CMakeLists.txt
@@ -73,7 +73,6 @@ endif()
 add_library(${LIB_GRAPHBOLT_NAME} SHARED ${BOLT_SRC} ${BOLT_HEADERS})
 include_directories(BEFORE ${BOLT_DIR}
                            ${BOLT_HEADERS}
-                           "../third_party/dmlc-core/include"
                            "../third_party/pcg/include"
                            "../third_party/phmap")
 target_link_libraries(${LIB_GRAPHBOLT_NAME} "${TORCH_LIBRARIES}")

diff --git a/graphbolt/src/random.cc b/graphbolt/src/random.cc
@@ -35,12 +35,12 @@ std::optional<uint64_t> RandomEngine::manual_seed;
 RandomEngine::RandomEngine() {
   std::random_device rd;
   std::lock_guard lock(manual_seed_mutex);
-  uint64_t seed = manual_seed.value_or(rd());
-  SetSeed(seed);
+  if (!manual_seed.has_value()) manual_seed = rd();
+  SetSeed(manual_seed.value());
 }
 
 /** @brief Constructor with given seed. */
-RandomEngine::RandomEngine(uint64_t seed) { RandomEngine(seed, GetThreadId()); }
+RandomEngine::RandomEngine(uint64_t seed) : RandomEngine(seed, GetThreadId()) {}
 
 /** @brief Constructor with given seed. */
 RandomEngine::RandomEngine(uint64_t seed, uint64_t stream) {
@@ -49,7 +49,8 @@ RandomEngine::RandomEngine(uint64_t seed, uint64_t stream) {
 
 /** @brief Get the thread-local random number generator instance. */
 RandomEngine* RandomEngine::ThreadLocal() {
-  return dmlc::ThreadLocalStore<RandomEngine>::Get();
+  static thread_local RandomEngine engine;
+  return &engine;
 }
 
 /** @brief Set the seed. */

diff --git a/graphbolt/src/random.h b/graphbolt/src/random.h
@@ -8,8 +8,6 @@
 #ifndef GRAPHBOLT_RANDOM_H_
 #define GRAPHBOLT_RANDOM_H_
 
-#include <dmlc/thread_local.h>
-
 #include <mutex>
 #include <optional>
 #include <pcg_random.hpp>