Skip to content

Commit eacda78

Browse files
mfbalinlijialin03
authored andcommitted
[GraphBolt] Per-thread RandomEngine initialization fix. (dmlc#7557)
1 parent 66f38d0 commit eacda78

File tree

3 files changed

+5
-7
lines changed

3 files changed

+5
-7
lines changed

graphbolt/CMakeLists.txt

-1
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,6 @@ endif()
7373
add_library(${LIB_GRAPHBOLT_NAME} SHARED ${BOLT_SRC} ${BOLT_HEADERS})
7474
include_directories(BEFORE ${BOLT_DIR}
7575
${BOLT_HEADERS}
76-
"../third_party/dmlc-core/include"
7776
"../third_party/pcg/include"
7877
"../third_party/phmap")
7978
target_link_libraries(${LIB_GRAPHBOLT_NAME} "${TORCH_LIBRARIES}")

graphbolt/src/random.cc

+5-4
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,12 @@ std::optional<uint64_t> RandomEngine::manual_seed;
3535
RandomEngine::RandomEngine() {
3636
std::random_device rd;
3737
std::lock_guard lock(manual_seed_mutex);
38-
uint64_t seed = manual_seed.value_or(rd());
39-
SetSeed(seed);
38+
if (!manual_seed.has_value()) manual_seed = rd();
39+
SetSeed(manual_seed.value());
4040
}
4141

4242
/** @brief Constructor with given seed. */
43-
RandomEngine::RandomEngine(uint64_t seed) { RandomEngine(seed, GetThreadId()); }
43+
RandomEngine::RandomEngine(uint64_t seed) : RandomEngine(seed, GetThreadId()) {}
4444

4545
/** @brief Constructor with given seed. */
4646
RandomEngine::RandomEngine(uint64_t seed, uint64_t stream) {
@@ -49,7 +49,8 @@ RandomEngine::RandomEngine(uint64_t seed, uint64_t stream) {
4949

5050
/** @brief Get the thread-local random number generator instance. */
5151
RandomEngine* RandomEngine::ThreadLocal() {
52-
return dmlc::ThreadLocalStore<RandomEngine>::Get();
52+
static thread_local RandomEngine engine;
53+
return &engine;
5354
}
5455

5556
/** @brief Set the seed. */

graphbolt/src/random.h

-2
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@
88
#ifndef GRAPHBOLT_RANDOM_H_
99
#define GRAPHBOLT_RANDOM_H_
1010

11-
#include <dmlc/thread_local.h>
12-
1311
#include <mutex>
1412
#include <optional>
1513
#include <pcg_random.hpp>

0 commit comments

Comments
 (0)