Skip to content

Commit 3a4f9a8

Browse files
ldorauldorau
authored
Merge pull request #1064 from wlemkows/tbb-symbols
Init TBB symbols only once
2 parents c80d2b2 + 59eee62 commit 3a4f9a8

File tree

3 files changed

+83
-47
lines changed

3 files changed

+83
-47
lines changed

‎src/libumf.c‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
#include "base_alloc_global.h"
1313
#include "ipc_cache.h"
1414
#include "memspace_internal.h"
15+
#include "pool/pool_scalable_internal.h"
1516
#include "provider_cuda_internal.h"
1617
#include "provider_level_zero_internal.h"
1718
#include "provider_tracking.h"
@@ -83,6 +84,7 @@ void umfTearDown(void) {
8384
fini_umfTearDown:
8485
fini_ze_global_state();
8586
fini_cu_global_state();
87+
fini_tbb_global_state();
8688
LOG_DEBUG("UMF library finalized");
8789
}
8890
}

‎src/pool/pool_scalable.c‎

Lines changed: 71 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
#include "base_alloc_global.h"
2222
#include "libumf.h"
23+
#include "pool_scalable_internal.h"
2324
#include "utils_common.h"
2425
#include "utils_concurrency.h"
2526
#include "utils_load_library.h"
@@ -33,6 +34,7 @@ static __TLS umf_result_t TLS_last_allocation_error;
3334
static __TLS umf_result_t TLS_last_free_error;
3435

3536
static const size_t DEFAULT_GRANULARITY = 2 * 1024 * 1024; // 2MB
37+
3638
typedef struct tbb_mem_pool_policy_t {
3739
raw_alloc_tbb_type pAlloc;
3840
raw_free_tbb_type pFree;
@@ -66,7 +68,6 @@ typedef struct tbb_callbacks_t {
6668
typedef struct tbb_memory_pool_t {
6769
umf_memory_provider_handle_t mem_provider;
6870
void *tbb_pool;
69-
tbb_callbacks_t tbb_callbacks;
7071
} tbb_memory_pool_t;
7172

7273
typedef enum tbb_enums_t {
@@ -82,6 +83,10 @@ typedef enum tbb_enums_t {
8283
TBB_POOL_SYMBOLS_MAX // it has to be the last one
8384
} tbb_enums_t;
8485

86+
static UTIL_ONCE_FLAG tbb_initialized = UTIL_ONCE_FLAG_INIT;
87+
static int tbb_init_result = 0;
88+
static tbb_callbacks_t tbb_callbacks = {0};
89+
8590
static const char *tbb_symbol[TBB_POOL_SYMBOLS_MAX] = {
8691
#ifdef _WIN32
8792
// symbols copied from oneTBB/src/tbbmalloc/def/win64-tbbmalloc.def
@@ -109,46 +114,60 @@ static const char *tbb_symbol[TBB_POOL_SYMBOLS_MAX] = {
109114
#endif
110115
};
111116

112-
static int init_tbb_callbacks(tbb_callbacks_t *tbb_callbacks) {
113-
assert(tbb_callbacks);
114-
117+
static void init_tbb_callbacks_once(void) {
115118
const char *lib_name = tbb_symbol[TBB_LIB_NAME];
116-
tbb_callbacks->lib_handle = utils_open_library(lib_name, 0);
117-
if (!tbb_callbacks->lib_handle) {
119+
tbb_callbacks.lib_handle = utils_open_library(lib_name, 0);
120+
if (!tbb_callbacks.lib_handle) {
118121
LOG_ERR("%s required by Scalable Pool not found - install TBB malloc "
119122
"or make sure it is in the default search paths.",
120123
lib_name);
121-
return -1;
124+
tbb_init_result = -1;
125+
return;
122126
}
123-
124-
*(void **)&tbb_callbacks->pool_malloc = utils_get_symbol_addr(
125-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_MALLOC], lib_name);
126-
*(void **)&tbb_callbacks->pool_realloc = utils_get_symbol_addr(
127-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_REALLOC], lib_name);
128-
*(void **)&tbb_callbacks->pool_aligned_malloc =
129-
utils_get_symbol_addr(tbb_callbacks->lib_handle,
127+
*(void **)&tbb_callbacks.pool_malloc = utils_get_symbol_addr(
128+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_MALLOC], lib_name);
129+
*(void **)&tbb_callbacks.pool_realloc = utils_get_symbol_addr(
130+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_REALLOC], lib_name);
131+
*(void **)&tbb_callbacks.pool_aligned_malloc =
132+
utils_get_symbol_addr(tbb_callbacks.lib_handle,
130133
tbb_symbol[TBB_POOL_ALIGNED_MALLOC], lib_name);
131-
*(void **)&tbb_callbacks->pool_free = utils_get_symbol_addr(
132-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_FREE], lib_name);
133-
*(void **)&tbb_callbacks->pool_create_v1 = utils_get_symbol_addr(
134-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_CREATE_V1], lib_name);
135-
*(void **)&tbb_callbacks->pool_destroy = utils_get_symbol_addr(
136-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_DESTROY], lib_name);
137-
*(void **)&tbb_callbacks->pool_identify = utils_get_symbol_addr(
138-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_IDENTIFY], lib_name);
139-
*(void **)&tbb_callbacks->pool_msize = utils_get_symbol_addr(
140-
tbb_callbacks->lib_handle, tbb_symbol[TBB_POOL_MSIZE], lib_name);
141-
142-
if (!tbb_callbacks->pool_malloc || !tbb_callbacks->pool_realloc ||
143-
!tbb_callbacks->pool_aligned_malloc || !tbb_callbacks->pool_free ||
144-
!tbb_callbacks->pool_create_v1 || !tbb_callbacks->pool_destroy ||
145-
!tbb_callbacks->pool_identify) {
134+
*(void **)&tbb_callbacks.pool_free = utils_get_symbol_addr(
135+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_FREE], lib_name);
136+
*(void **)&tbb_callbacks.pool_create_v1 = utils_get_symbol_addr(
137+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_CREATE_V1], lib_name);
138+
*(void **)&tbb_callbacks.pool_destroy = utils_get_symbol_addr(
139+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_DESTROY], lib_name);
140+
*(void **)&tbb_callbacks.pool_identify = utils_get_symbol_addr(
141+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_IDENTIFY], lib_name);
142+
*(void **)&tbb_callbacks.pool_msize = utils_get_symbol_addr(
143+
tbb_callbacks.lib_handle, tbb_symbol[TBB_POOL_MSIZE], lib_name);
144+
145+
if (!tbb_callbacks.pool_malloc || !tbb_callbacks.pool_realloc ||
146+
!tbb_callbacks.pool_aligned_malloc || !tbb_callbacks.pool_free ||
147+
!tbb_callbacks.pool_create_v1 || !tbb_callbacks.pool_destroy ||
148+
!tbb_callbacks.pool_identify) {
146149
LOG_FATAL("Could not find all TBB symbols in %s", lib_name);
147-
utils_close_library(tbb_callbacks->lib_handle);
148-
return -1;
150+
if (utils_close_library(tbb_callbacks.lib_handle)) {
151+
LOG_ERR("Could not close %s library", lib_name);
152+
}
153+
tbb_init_result = -1;
149154
}
155+
}
150156

151-
return 0;
157+
static int init_tbb_callbacks(void) {
158+
utils_init_once(&tbb_initialized, init_tbb_callbacks_once);
159+
return tbb_init_result;
160+
}
161+
162+
void fini_tbb_global_state(void) {
163+
if (tbb_callbacks.lib_handle) {
164+
if (!utils_close_library(tbb_callbacks.lib_handle)) {
165+
tbb_callbacks.lib_handle = NULL;
166+
LOG_DEBUG("TBB library closed");
167+
} else {
168+
LOG_ERR("TBB library cannot be unloaded");
169+
}
170+
}
152171
}
153172

154173
static void *tbb_raw_alloc_wrapper(intptr_t pool_id, size_t *raw_bytes) {
@@ -264,35 +283,41 @@ static umf_result_t tbb_pool_initialize(umf_memory_provider_handle_t provider,
264283
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
265284
}
266285

267-
int ret = init_tbb_callbacks(&pool_data->tbb_callbacks);
286+
umf_result_t res = UMF_RESULT_SUCCESS;
287+
int ret = init_tbb_callbacks();
268288
if (ret != 0) {
269289
LOG_FATAL("loading TBB symbols failed");
270-
return UMF_RESULT_ERROR_UNKNOWN;
290+
res = UMF_RESULT_ERROR_UNKNOWN;
291+
goto err_tbb_init;
271292
}
272293

273294
pool_data->mem_provider = provider;
274-
ret = pool_data->tbb_callbacks.pool_create_v1((intptr_t)pool_data, &policy,
275-
&(pool_data->tbb_pool));
295+
ret = tbb_callbacks.pool_create_v1((intptr_t)pool_data, &policy,
296+
&(pool_data->tbb_pool));
276297
if (ret != 0 /* TBBMALLOC_OK */) {
277-
return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
298+
res = UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
299+
goto err_tbb_init;
278300
}
279301

280302
*pool = (void *)pool_data;
281303

282-
return UMF_RESULT_SUCCESS;
304+
return res;
305+
306+
err_tbb_init:
307+
umf_ba_global_free(pool_data);
308+
return res;
283309
}
284310

285311
static void tbb_pool_finalize(void *pool) {
286312
tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
287-
pool_data->tbb_callbacks.pool_destroy(pool_data->tbb_pool);
288-
utils_close_library(pool_data->tbb_callbacks.lib_handle);
313+
tbb_callbacks.pool_destroy(pool_data->tbb_pool);
289314
umf_ba_global_free(pool_data);
290315
}
291316

292317
static void *tbb_malloc(void *pool, size_t size) {
293318
tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
294319
TLS_last_allocation_error = UMF_RESULT_SUCCESS;
295-
void *ptr = pool_data->tbb_callbacks.pool_malloc(pool_data->tbb_pool, size);
320+
void *ptr = tbb_callbacks.pool_malloc(pool_data->tbb_pool, size);
296321
if (ptr == NULL) {
297322
if (TLS_last_allocation_error == UMF_RESULT_SUCCESS) {
298323
TLS_last_allocation_error = UMF_RESULT_ERROR_UNKNOWN;
@@ -319,8 +344,7 @@ static void *tbb_calloc(void *pool, size_t num, size_t size) {
319344
static void *tbb_realloc(void *pool, void *ptr, size_t size) {
320345
tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
321346
TLS_last_allocation_error = UMF_RESULT_SUCCESS;
322-
void *new_ptr =
323-
pool_data->tbb_callbacks.pool_realloc(pool_data->tbb_pool, ptr, size);
347+
void *new_ptr = tbb_callbacks.pool_realloc(pool_data->tbb_pool, ptr, size);
324348
if (new_ptr == NULL) {
325349
if (TLS_last_allocation_error == UMF_RESULT_SUCCESS) {
326350
TLS_last_allocation_error = UMF_RESULT_ERROR_UNKNOWN;
@@ -334,8 +358,8 @@ static void *tbb_realloc(void *pool, void *ptr, size_t size) {
334358
static void *tbb_aligned_malloc(void *pool, size_t size, size_t alignment) {
335359
tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
336360
TLS_last_allocation_error = UMF_RESULT_SUCCESS;
337-
void *ptr = pool_data->tbb_callbacks.pool_aligned_malloc(
338-
pool_data->tbb_pool, size, alignment);
361+
void *ptr =
362+
tbb_callbacks.pool_aligned_malloc(pool_data->tbb_pool, size, alignment);
339363
if (ptr == NULL) {
340364
if (TLS_last_allocation_error == UMF_RESULT_SUCCESS) {
341365
TLS_last_allocation_error = UMF_RESULT_ERROR_UNKNOWN;
@@ -360,7 +384,7 @@ static umf_result_t tbb_free(void *pool, void *ptr) {
360384
utils_annotate_release(pool);
361385

362386
tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
363-
if (pool_data->tbb_callbacks.pool_free(pool_data->tbb_pool, ptr)) {
387+
if (tbb_callbacks.pool_free(pool_data->tbb_pool, ptr)) {
364388
return UMF_RESULT_SUCCESS;
365389
}
366390

@@ -373,7 +397,7 @@ static umf_result_t tbb_free(void *pool, void *ptr) {
373397

374398
static size_t tbb_malloc_usable_size(void *pool, void *ptr) {
375399
tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
376-
return pool_data->tbb_callbacks.pool_msize(pool_data->tbb_pool, ptr);
400+
return tbb_callbacks.pool_msize(pool_data->tbb_pool, ptr);
377401
}
378402

379403
static umf_result_t tbb_get_last_allocation_error(void *pool) {

‎src/pool/pool_scalable_internal.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
/*
2+
*
3+
* Copyright (C) 2025 Intel Corporation
4+
*
5+
* Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
6+
* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7+
*
8+
*/
9+
10+
void fini_tbb_global_state(void);

0 commit comments

Comments
 (0)