Skip to content

Commit d188efc

Browse files
committed
make umfInit threadsafe
1 parent a2a0919 commit d188efc

File tree

1 file changed

+33
-21
lines changed

1 file changed

+33
-21
lines changed

src/libumf.c

Lines changed: 33 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "provider_level_zero_internal.h"
2020
#include "provider_tracking.h"
2121
#include "utils_common.h"
22+
#include "utils_concurrency.h"
2223
#include "utils_log.h"
2324
#if !defined(UMF_NO_HWLOC)
2425
#include "topology.h"
@@ -34,36 +35,46 @@ static umf_ctl_node_t CTL_NODE(umf)[] = {CTL_CHILD(provider), CTL_CHILD(pool),
3435
void initialize_global_ctl(void) { CTL_REGISTER_MODULE(NULL, umf); }
3536

3637
umf_result_t umfInit(void) {
37-
if (utils_fetch_and_add_u64(&umfRefCount, 1) == 0) {
38-
utils_log_init();
39-
umf_result_t umf_result = umfMemoryTrackerCreate(&TRACKER);
40-
if (umf_result != UMF_RESULT_SUCCESS) {
41-
LOG_ERR("Failed to create memory tracker");
42-
return umf_result;
38+
uint64_t refCount;
39+
do {
40+
if (utils_fetch_and_add_u64(&umfRefCount, 1) == 0) {
41+
utils_log_init();
42+
umf_result_t umf_result = umfMemoryTrackerCreate(&TRACKER);
43+
if (!TRACKER) {
44+
LOG_ERR("Failed to create memory tracker");
45+
utils_atomic_decrement_size_t(&umfRefCount);
46+
return umf_result;
47+
}
48+
49+
LOG_DEBUG("UMF tracker created");
50+
51+
umf_result = umfIpcCacheGlobalInit();
52+
if (umf_result != UMF_RESULT_SUCCESS) {
53+
LOG_ERR("Failed to initialize IPC cache");
54+
umfMemoryTrackerDestroy(TRACKER);
55+
utils_atomic_decrement_size_t(&umfRefCount);
56+
return umf_result;
57+
}
58+
59+
LOG_DEBUG("UMF IPC cache initialized");
60+
initialize_global_ctl();
61+
62+
utils_atomic_increment_u64(&umfRefCount);
4363
}
4464

45-
LOG_DEBUG("UMF tracker created");
65+
do {
66+
utils_atomic_load_acquire_u64(&umfRefCount, &refCount);
67+
} while (refCount == 1);
4668

47-
umf_result = umfIpcCacheGlobalInit();
48-
if (umf_result != UMF_RESULT_SUCCESS) {
49-
LOG_ERR("Failed to initialize IPC cache");
50-
umfMemoryTrackerDestroy(TRACKER);
51-
return umf_result;
52-
}
53-
54-
LOG_DEBUG("UMF IPC cache initialized");
55-
initialize_global_ctl();
56-
}
69+
} while (refCount < 2);
5770

58-
if (TRACKER) {
59-
LOG_DEBUG("UMF library initialized");
60-
}
71+
LOG_DEBUG("UMF library initialized");
6172

6273
return UMF_RESULT_SUCCESS;
6374
}
6475

6576
umf_result_t umfTearDown(void) {
66-
if (utils_fetch_and_sub_u64(&umfRefCount, 1) == 1) {
77+
if (utils_fetch_and_sub_u64(&umfRefCount, 1) == 2) {
6778
#if !defined(_WIN32) && !defined(UMF_NO_HWLOC)
6879
umfMemspaceHostAllDestroy();
6980
umfMemspaceHighestCapacityDestroy();
@@ -95,6 +106,7 @@ umf_result_t umfTearDown(void) {
95106
fini_cu_global_state();
96107
fini_tbb_global_state();
97108
LOG_DEBUG("UMF library finalized");
109+
utils_atomic_decrement_u64(&umfRefCount);
98110
}
99111
return UMF_RESULT_SUCCESS;
100112
}

0 commit comments

Comments
 (0)