Skip to content

Commit 992e0da

Browse files
lplewalplewa
committed
add ctl_stats to remaining providers
1 parent fe3aea6 commit 992e0da

10 files changed

+480
-26
lines changed

‎src/provider/provider_cuda.c‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include "provider_cuda_internal.h"
1616
#include "utils_load_library.h"
1717
#include "utils_log.h"
18+
#include "provider_ctl_stats_type.h"
1819

1920
static void *cu_lib_handle = NULL;
2021

@@ -53,8 +54,12 @@ typedef struct cu_memory_provider_t {
5354
umf_usm_memory_type_t memory_type;
5455
size_t min_alignment;
5556
unsigned int alloc_flags;
57+
ctl_stats_t stats;
5658
} cu_memory_provider_t;
5759

60+
#define CTL_PROVIDER_TYPE cu_memory_provider_t
61+
#include "provider_ctl_stats_impl.h"
62+
5863
// CUDA Memory Provider settings struct
5964
typedef struct umf_cuda_memory_provider_params_t {
6065
// Handle to the CUDA context
@@ -98,6 +103,9 @@ static cu_ops_t g_cu_ops;
98103
static UTIL_ONCE_FLAG cu_is_initialized = UTIL_ONCE_FLAG_INIT;
99104
static bool Init_cu_global_state_failed;
100105

106+
struct ctl cu_memory_ctl_root;
107+
static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
108+
101109
// forward decl needed for alloc
102110
static umf_result_t cu_memory_provider_free(void *provider, void *ptr,
103111
size_t bytes);
@@ -133,6 +141,10 @@ static umf_result_t cu2umf_result(CUresult result) {
133141
}
134142
}
135143

144+
static void initialize_cu_ctl(void) {
145+
CTL_REGISTER_MODULE(&cu_memory_ctl_root, stats);
146+
}
147+
136148
static void init_cu_global_state(void) {
137149
#ifdef _WIN32
138150
const char *lib_name = "nvcuda.dll";
@@ -333,6 +345,7 @@ static umf_result_t cu_memory_provider_initialize(const void *params,
333345
if (!cu_provider) {
334346
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
335347
}
348+
memset(cu_provider, 0, sizeof(cu_memory_provider_t));
336349

337350
// CUDA alloc functions doesn't allow to provide user alignment - get the
338351
// minimum one from the driver
@@ -460,6 +473,8 @@ static umf_result_t cu_memory_provider_alloc(void *provider, size_t size,
460473
LOG_ERR("unsupported alignment size");
461474
return UMF_RESULT_ERROR_INVALID_ALIGNMENT;
462475
}
476+
477+
provider_ctl_stats_alloc(cu_provider, size);
463478
return umf_result;
464479
}
465480

@@ -504,7 +519,11 @@ static umf_result_t cu_memory_provider_free(void *provider, void *ptr,
504519
LOG_ERR("Failed to restore CUDA context, ret = %d", umf_result);
505520
}
506521

507-
return cu2umf_result(cu_result);
522+
umf_result_t ret = cu2umf_result(cu_result);
523+
if (ret == UMF_RESULT_SUCCESS) {
524+
provider_ctl_stats_free(cu_provider, bytes);
525+
}
526+
return ret;
508527
}
509528

510529
static umf_result_t
@@ -685,6 +704,14 @@ cu_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
685704
return UMF_RESULT_SUCCESS;
686705
}
687706

707+
static umf_result_t cu_ctl(void *provider, umf_ctl_query_source_t operationType,
708+
const char *name, void *arg, size_t size,
709+
umf_ctl_query_type_t query_type, va_list args) {
710+
utils_init_once(&ctl_initialized, initialize_cu_ctl);
711+
return ctl_query(&cu_memory_ctl_root, provider, operationType, name,
712+
query_type, arg, size, args);
713+
}
714+
688715
static umf_memory_provider_ops_t UMF_CUDA_MEMORY_PROVIDER_OPS = {
689716
.version = UMF_PROVIDER_OPS_VERSION_CURRENT,
690717
.initialize = cu_memory_provider_initialize,
@@ -707,6 +734,7 @@ static umf_memory_provider_ops_t UMF_CUDA_MEMORY_PROVIDER_OPS = {
707734
.ext_put_ipc_handle = cu_memory_provider_put_ipc_handle,
708735
.ext_open_ipc_handle = cu_memory_provider_open_ipc_handle,
709736
.ext_close_ipc_handle = cu_memory_provider_close_ipc_handle,
737+
.ext_ctl = cu_ctl,
710738
};
711739

712740
const umf_memory_provider_ops_t *umfCUDAMemoryProviderOps(void) {

‎src/provider/provider_devdax_memory.c‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ umf_result_t umfDevDaxMemoryProviderParamsSetProtection(
7070
#include "utils_common.h"
7171
#include "utils_concurrency.h"
7272
#include "utils_log.h"
73+
#include "provider_ctl_stats_type.h"
7374

7475
#define DEVDAX_PAGE_SIZE_2MB ((size_t)(2 * 1024 * 1024)) // == 2 MB
7576

@@ -83,8 +84,12 @@ typedef struct devdax_memory_provider_t {
8384
utils_mutex_t lock; // lock of ptr and offset
8485
unsigned protection; // combination of OS-specific protection flags
8586
coarse_t *coarse; // coarse library handle
87+
ctl_stats_t stats;
8688
} devdax_memory_provider_t;
8789

90+
#define CTL_PROVIDER_TYPE devdax_memory_provider_t
91+
#include "provider_ctl_stats_impl.h"
92+
8893
// DevDax Memory provider settings struct
8994
typedef struct umf_devdax_memory_provider_params_t {
9095
char *path;
@@ -112,6 +117,9 @@ static __TLS devdax_last_native_error_t TLS_last_native_error;
112117
#define _UMF_DEVDAX_RESULT_ERROR_PURGE_FORCE_FAILED \
113118
(UMF_DEVDAX_RESULT_ERROR_PURGE_FORCE_FAILED - UMF_DEVDAX_RESULT_SUCCESS)
114119

120+
struct ctl devdax_memory_ctl_root;
121+
static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
122+
115123
static const char *Native_error_str[] = {
116124
[_UMF_DEVDAX_RESULT_SUCCESS] = "success",
117125
[_UMF_DEVDAX_RESULT_ERROR_ALLOC_FAILED] = "memory allocation failed",
@@ -127,6 +135,10 @@ static void devdax_store_last_native_error(int32_t native_error,
127135
TLS_last_native_error.errno_value = errno_value;
128136
}
129137

138+
static void initialize_devdax_ctl(void) {
139+
CTL_REGISTER_MODULE(&devdax_memory_ctl_root, stats);
140+
}
141+
130142
static umf_result_t
131143
devdax_translate_params(const umf_devdax_memory_provider_params_t *in_params,
132144
devdax_memory_provider_t *provider) {
@@ -284,7 +296,12 @@ static umf_result_t devdax_alloc(void *provider, size_t size, size_t alignment,
284296
void **resultPtr) {
285297
devdax_memory_provider_t *devdax_provider =
286298
(devdax_memory_provider_t *)provider;
287-
return coarse_alloc(devdax_provider->coarse, size, alignment, resultPtr);
299+
umf_result_t ret =
300+
coarse_alloc(devdax_provider->coarse, size, alignment, resultPtr);
301+
if (ret == UMF_RESULT_SUCCESS) {
302+
provider_ctl_stats_alloc(devdax_provider, size);
303+
}
304+
return ret;
288305
}
289306

290307
static umf_result_t devdax_get_last_native_error(void *provider,
@@ -531,10 +548,23 @@ static umf_result_t devdax_close_ipc_handle(void *provider, void *ptr,
531548
return UMF_RESULT_SUCCESS;
532549
}
533550

551+
static umf_result_t devdax_ctl(void *provider,
552+
umf_ctl_query_source_t operationType,
553+
const char *name, void *arg, size_t size,
554+
umf_ctl_query_type_t query_type, va_list args) {
555+
utils_init_once(&ctl_initialized, initialize_devdax_ctl);
556+
return ctl_query(&devdax_memory_ctl_root, provider, operationType, name,
557+
query_type, arg, size, args);
558+
}
559+
534560
static umf_result_t devdax_free(void *provider, void *ptr, size_t size) {
535561
devdax_memory_provider_t *devdax_provider =
536562
(devdax_memory_provider_t *)provider;
537-
return coarse_free(devdax_provider->coarse, ptr, size);
563+
umf_result_t ret = coarse_free(devdax_provider->coarse, ptr, size);
564+
if (ret == UMF_RESULT_SUCCESS) {
565+
provider_ctl_stats_free(devdax_provider, size);
566+
}
567+
return ret;
538568
}
539569

540570
static umf_memory_provider_ops_t UMF_DEVDAX_MEMORY_PROVIDER_OPS = {
@@ -555,7 +585,8 @@ static umf_memory_provider_ops_t UMF_DEVDAX_MEMORY_PROVIDER_OPS = {
555585
.ext_get_ipc_handle = devdax_get_ipc_handle,
556586
.ext_put_ipc_handle = devdax_put_ipc_handle,
557587
.ext_open_ipc_handle = devdax_open_ipc_handle,
558-
.ext_close_ipc_handle = devdax_close_ipc_handle};
588+
.ext_close_ipc_handle = devdax_close_ipc_handle,
589+
.ext_ctl = devdax_ctl};
559590

560591
const umf_memory_provider_ops_t *umfDevDaxMemoryProviderOps(void) {
561592
return &UMF_DEVDAX_MEMORY_PROVIDER_OPS;

‎src/provider/provider_file_memory.c‎

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include <umf/memory_provider_ops.h>
1919
#include <umf/providers/provider_file_memory.h>
2020

21+
#include "provider_ctl_stats_type.h"
2122
#include "utils_log.h"
2223

2324
#if defined(_WIN32) || defined(UMF_NO_HWLOC)
@@ -77,6 +78,7 @@ umf_result_t umfFileMemoryProviderParamsSetVisibility(
7778
#include "utils_common.h"
7879
#include "utils_concurrency.h"
7980
#include "utils_log.h"
81+
#include "provider_ctl_stats_type.h"
8082

8183
#define FSDAX_PAGE_SIZE_2MB ((size_t)(2 * 1024 * 1024)) // == 2 MB
8284

@@ -112,8 +114,12 @@ typedef struct file_memory_provider_t {
112114
critnib *fd_offset_map;
113115

114116
coarse_t *coarse; // coarse library handle
117+
ctl_stats_t stats;
115118
} file_memory_provider_t;
116119

120+
#define CTL_PROVIDER_TYPE file_memory_provider_t
121+
#include "provider_ctl_stats_impl.h"
122+
117123
// File Memory Provider settings struct
118124
typedef struct umf_file_memory_provider_params_t {
119125
char *path;
@@ -139,6 +145,9 @@ static __TLS file_last_native_error_t TLS_last_native_error;
139145
#define _UMF_FILE_RESULT_ERROR_PURGE_FORCE_FAILED \
140146
(UMF_FILE_RESULT_ERROR_PURGE_FORCE_FAILED - UMF_FILE_RESULT_SUCCESS)
141147

148+
struct ctl file_memory_ctl_root;
149+
static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
150+
142151
static const char *Native_error_str[] = {
143152
[_UMF_FILE_RESULT_SUCCESS] = "success",
144153
[_UMF_FILE_RESULT_ERROR_ALLOC_FAILED] = "memory allocation failed",
@@ -152,6 +161,10 @@ static void file_store_last_native_error(int32_t native_error,
152161
TLS_last_native_error.errno_value = errno_value;
153162
}
154163

164+
static void initialize_file_ctl(void) {
165+
CTL_REGISTER_MODULE(&file_memory_ctl_root, stats);
166+
}
167+
155168
static umf_result_t
156169
file_translate_params(const umf_file_memory_provider_params_t *in_params,
157170
file_memory_provider_t *provider) {
@@ -496,7 +509,12 @@ static umf_result_t file_alloc_aligned(file_memory_provider_t *file_provider,
496509
static umf_result_t file_alloc(void *provider, size_t size, size_t alignment,
497510
void **resultPtr) {
498511
file_memory_provider_t *file_provider = (file_memory_provider_t *)provider;
499-
return coarse_alloc(file_provider->coarse, size, alignment, resultPtr);
512+
umf_result_t ret =
513+
coarse_alloc(file_provider->coarse, size, alignment, resultPtr);
514+
if (ret == UMF_RESULT_SUCCESS) {
515+
provider_ctl_stats_alloc(file_provider, size);
516+
}
517+
return ret;
500518
}
501519

502520
static umf_result_t file_alloc_cb(void *provider, size_t size, size_t alignment,
@@ -859,9 +877,22 @@ static umf_result_t file_close_ipc_handle(void *provider, void *ptr,
859877
return UMF_RESULT_SUCCESS;
860878
}
861879

880+
static umf_result_t file_ctl(void *provider,
881+
umf_ctl_query_source_t operationType,
882+
const char *name, void *arg, size_t size,
883+
umf_ctl_query_type_t query_type, va_list args) {
884+
utils_init_once(&ctl_initialized, initialize_file_ctl);
885+
return ctl_query(&file_memory_ctl_root, provider, operationType, name,
886+
query_type, arg, size, args);
887+
}
888+
862889
static umf_result_t file_free(void *provider, void *ptr, size_t size) {
863890
file_memory_provider_t *file_provider = (file_memory_provider_t *)provider;
864-
return coarse_free(file_provider->coarse, ptr, size);
891+
umf_result_t ret = coarse_free(file_provider->coarse, ptr, size);
892+
if (ret == UMF_RESULT_SUCCESS) {
893+
provider_ctl_stats_free(file_provider, size);
894+
}
895+
return ret;
865896
}
866897

867898
static umf_memory_provider_ops_t UMF_FILE_MEMORY_PROVIDER_OPS = {
@@ -882,7 +913,8 @@ static umf_memory_provider_ops_t UMF_FILE_MEMORY_PROVIDER_OPS = {
882913
.ext_get_ipc_handle = file_get_ipc_handle,
883914
.ext_put_ipc_handle = file_put_ipc_handle,
884915
.ext_open_ipc_handle = file_open_ipc_handle,
885-
.ext_close_ipc_handle = file_close_ipc_handle};
916+
.ext_close_ipc_handle = file_close_ipc_handle,
917+
.ext_ctl = file_ctl};
886918

887919
const umf_memory_provider_ops_t *umfFileMemoryProviderOps(void) {
888920
return &UMF_FILE_MEMORY_PROVIDER_OPS;

‎src/provider/provider_level_zero.c‎

Lines changed: 38 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -376,9 +376,10 @@ static ze_relaxed_allocation_limits_exp_desc_t relaxed_device_allocation_desc =
376376
.pNext = NULL,
377377
.flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE};
378378

379-
static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
380-
size_t alignment,
381-
void **resultPtr) {
379+
static umf_result_t ze_memory_provider_alloc_helper(void *provider, size_t size,
380+
size_t alignment,
381+
int update_stats,
382+
void **resultPtr) {
382383
ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
383384

384385
ze_result_t ze_result = ZE_RESULT_SUCCESS;
@@ -442,47 +443,65 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
442443
}
443444
}
444445

445-
provider_ctl_stats_alloc(ze_provider, size);
446-
446+
if (update_stats) {
447+
provider_ctl_stats_alloc(ze_provider, size);
448+
}
447449
return UMF_RESULT_SUCCESS;
448450
}
449451

450-
static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
451-
size_t bytes) {
452+
static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
453+
size_t alignment,
454+
void **resultPtr) {
455+
return ze_memory_provider_alloc_helper(provider, size, alignment, 1,
456+
resultPtr);
457+
}
458+
459+
static umf_result_t ze_memory_provider_free_helper(void *provider, void *ptr,
460+
size_t bytes,
461+
int update_stats) {
452462
if (ptr == NULL) {
453463
return UMF_RESULT_SUCCESS;
454464
}
455465

456466
ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
457-
467+
umf_result_t ret;
458468
if (ze_provider->freePolicyFlags == 0) {
459-
return ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
460-
}
469+
ret = ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
470+
} else {
471+
ze_memory_free_ext_desc_t desc = {
472+
.stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
473+
.pNext = NULL,
474+
.freePolicy = ze_provider->freePolicyFlags};
461475

462-
ze_memory_free_ext_desc_t desc = {
463-
.stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
464-
.pNext = NULL,
465-
.freePolicy = ze_provider->freePolicyFlags};
476+
ret = ze2umf_result(
477+
g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
478+
}
466479

467-
umf_result_t ret =
468-
ze2umf_result(g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
469480
if (ret != UMF_RESULT_SUCCESS) {
470481
return ret;
471482
}
472483

473-
provider_ctl_stats_free(ze_provider, bytes);
484+
if (update_stats) {
485+
provider_ctl_stats_free(ze_provider, bytes);
486+
}
474487

475488
return UMF_RESULT_SUCCESS;
476489
}
477490

491+
static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
492+
size_t bytes) {
493+
return ze_memory_provider_free_helper(provider, ptr, bytes, 1);
494+
}
495+
478496
static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
479497
size_t *min_page_size) {
480498
assert(min_page_size);
481499

482500
LOG_DEBUG("Querying minimum page size");
483501

484502
void *ptr;
485-
umf_result_t result = ze_memory_provider_alloc(ze_provider, 1, 0, &ptr);
503+
umf_result_t result =
504+
ze_memory_provider_alloc_helper(ze_provider, 1, 0, 0, &ptr);
486505
if (result != UMF_RESULT_SUCCESS) {
487506
return result;
488507
}
@@ -494,7 +513,7 @@ static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
494513

495514
*min_page_size = properties.pageSize;
496515

497-
ze_memory_provider_free(ze_provider, ptr, 1);
516+
ze_memory_provider_free_helper(ze_provider, ptr, 1, 0);
498517

499518
return ze2umf_result(ze_result);
500519
}

0 commit comments

Comments
 (0)