1414#include <umf/memory_provider_ops.h>
1515#include <umf/providers/provider_level_zero.h>
1616
17+ #include "provider_ctl_stats_type.h"
1718#include "provider_level_zero_internal.h"
1819#include "utils_load_library.h"
1920#include "utils_log.h"
@@ -27,6 +28,7 @@ void fini_ze_global_state(void) {
2728 }
2829}
2930
31+ #define UMF_BUILD_LEVEL_ZERO_PROVIDER 1
3032#if UMF_BUILD_LEVEL_ZERO_PROVIDER
3133
3234#include "base_alloc_global.h"
@@ -73,6 +75,8 @@ typedef struct ze_memory_provider_t {
7375 size_t min_page_size ;
7476
7577 uint32_t device_ordinal ;
78+
79+ ctl_stats_t stats ;
7680} ze_memory_provider_t ;
7781
7882typedef struct ze_ops_t {
@@ -115,6 +119,16 @@ static void store_last_native_error(int32_t native_error) {
115119 TLS_last_native_error = native_error ;
116120}
117121
122+ #define CTL_PROVIDER_TYPE ze_memory_provider_t
123+ #include "provider_ctl_stats_impl.h"
124+
125+ struct ctl ze_memory_ctl_root ;
126+ static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT ;
127+
128+ static void initialize_ze_ctl (void ) {
129+ CTL_REGISTER_MODULE (& ze_memory_ctl_root , stats );
130+ }
131+
118132static umf_result_t ze2umf_result (ze_result_t result ) {
119133 switch (result ) {
120134 case ZE_RESULT_SUCCESS :
@@ -429,13 +443,13 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
429443 }
430444 }
431445
432- return ze2umf_result (ze_result );
446+ provider_ctl_stats_alloc (ze_provider , size );
447+
448+ return UMF_RESULT_SUCCESS ;
433449}
434450
435451static umf_result_t ze_memory_provider_free (void * provider , void * ptr ,
436452 size_t bytes ) {
437- (void )bytes ;
438-
439453 if (ptr == NULL ) {
440454 return UMF_RESULT_SUCCESS ;
441455 }
@@ -451,8 +465,15 @@ static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
451465 .pNext = NULL ,
452466 .freePolicy = ze_provider -> freePolicyFlags };
453467
454- return ze2umf_result (
455- g_ze_ops .zeMemFreeExt (ze_provider -> context , & desc , ptr ));
468+ umf_result_t ret =
469+ ze2umf_result (g_ze_ops .zeMemFreeExt (ze_provider -> context , & desc , ptr ));
470+ if (ret != UMF_RESULT_SUCCESS ) {
471+ return ret ;
472+ }
473+
474+ provider_ctl_stats_free (ze_provider , bytes );
475+
476+ return UMF_RESULT_SUCCESS ;
456477}
457478
458479static umf_result_t query_min_page_size (ze_memory_provider_t * ze_provider ,
@@ -790,6 +811,14 @@ ze_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
790811 return UMF_RESULT_SUCCESS ;
791812}
792813
814+ static umf_result_t ze_ctl (void * provider , int operationType , const char * name ,
815+ void * arg , size_t size ,
816+ umf_ctl_query_type_t query_type ) {
817+ utils_init_once (& ctl_initialized , initialize_ze_ctl );
818+ return ctl_query (& ze_memory_ctl_root , provider , operationType , name ,
819+ query_type , arg , size );
820+ }
821+
793822static umf_memory_provider_ops_t UMF_LEVEL_ZERO_MEMORY_PROVIDER_OPS = {
794823 .version = UMF_PROVIDER_OPS_VERSION_CURRENT ,
795824 .initialize = ze_memory_provider_initialize ,
@@ -809,6 +838,7 @@ static umf_memory_provider_ops_t UMF_LEVEL_ZERO_MEMORY_PROVIDER_OPS = {
809838 .ext_put_ipc_handle = ze_memory_provider_put_ipc_handle ,
810839 .ext_open_ipc_handle = ze_memory_provider_open_ipc_handle ,
811840 .ext_close_ipc_handle = ze_memory_provider_close_ipc_handle ,
841+ .ext_ctl = ze_ctl ,
812842};
813843
814844const umf_memory_provider_ops_t * umfLevelZeroMemoryProviderOps (void ) {
0 commit comments