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"
@@ -73,6 +74,8 @@ typedef struct ze_memory_provider_t {
7374 size_t min_page_size ;
7475
7576 uint32_t device_ordinal ;
77+
78+ ctl_stats_t stats ;
7679} ze_memory_provider_t ;
7780
7881typedef struct ze_ops_t {
@@ -115,6 +118,16 @@ static void store_last_native_error(int32_t native_error) {
115118 TLS_last_native_error = native_error ;
116119}
117120
121+ #define CTL_PROVIDER_TYPE ze_memory_provider_t
122+ #include "provider_ctl_stats_impl.h"
123+
124+ struct ctl ze_memory_ctl_root ;
125+ static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT ;
126+
127+ static void initialize_ze_ctl (void ) {
128+ CTL_REGISTER_MODULE (& ze_memory_ctl_root , stats );
129+ }
130+
118131static umf_result_t ze2umf_result (ze_result_t result ) {
119132 switch (result ) {
120133 case ZE_RESULT_SUCCESS :
@@ -429,13 +442,13 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
429442 }
430443 }
431444
432- return ze2umf_result (ze_result );
445+ provider_ctl_stats_alloc (ze_provider , size );
446+
447+ return UMF_RESULT_SUCCESS ;
433448}
434449
435450static umf_result_t ze_memory_provider_free (void * provider , void * ptr ,
436451 size_t bytes ) {
437- (void )bytes ;
438-
439452 if (ptr == NULL ) {
440453 return UMF_RESULT_SUCCESS ;
441454 }
@@ -451,8 +464,15 @@ static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
451464 .pNext = NULL ,
452465 .freePolicy = ze_provider -> freePolicyFlags };
453466
454- return ze2umf_result (
455- g_ze_ops .zeMemFreeExt (ze_provider -> context , & desc , ptr ));
467+ umf_result_t ret =
468+ ze2umf_result (g_ze_ops .zeMemFreeExt (ze_provider -> context , & desc , ptr ));
469+ if (ret != UMF_RESULT_SUCCESS ) {
470+ return ret ;
471+ }
472+
473+ provider_ctl_stats_free (ze_provider , bytes );
474+
475+ return UMF_RESULT_SUCCESS ;
456476}
457477
458478static umf_result_t query_min_page_size (ze_memory_provider_t * ze_provider ,
@@ -790,6 +810,14 @@ ze_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
790810 return UMF_RESULT_SUCCESS ;
791811}
792812
813+ static umf_result_t ze_ctl (void * provider , int operationType , const char * name ,
814+ void * arg , size_t size ,
815+ umf_ctl_query_type_t query_type ) {
816+ utils_init_once (& ctl_initialized , initialize_ze_ctl );
817+ return ctl_query (& ze_memory_ctl_root , provider , operationType , name ,
818+ query_type , arg , size );
819+ }
820+
793821static umf_memory_provider_ops_t UMF_LEVEL_ZERO_MEMORY_PROVIDER_OPS = {
794822 .version = UMF_PROVIDER_OPS_VERSION_CURRENT ,
795823 .initialize = ze_memory_provider_initialize ,
@@ -809,6 +837,7 @@ static umf_memory_provider_ops_t UMF_LEVEL_ZERO_MEMORY_PROVIDER_OPS = {
809837 .ext_put_ipc_handle = ze_memory_provider_put_ipc_handle ,
810838 .ext_open_ipc_handle = ze_memory_provider_open_ipc_handle ,
811839 .ext_close_ipc_handle = ze_memory_provider_close_ipc_handle ,
840+ .ext_ctl = ze_ctl ,
812841};
813842
814843const umf_memory_provider_ops_t * umfLevelZeroMemoryProviderOps (void ) {
0 commit comments