Merge pull request #1444 from lplewa/ctl_stats_cnt

add ctl_stats to remaining providers

Author: KFilipek

src/provider/provider_cuda.c

@@ -12,6 +12,7 @@
 #include <umf.h>
 #include <umf/providers/provider_cuda.h>
 
+#include "provider_ctl_stats_type.h"
 #include "provider_cuda_internal.h"
 #include "utils_load_library.h"
 #include "utils_log.h"
@@ -53,8 +54,12 @@ typedef struct cu_memory_provider_t {
     umf_usm_memory_type_t memory_type;
     size_t min_alignment;
     unsigned int alloc_flags;
+    ctl_stats_t stats;
 } cu_memory_provider_t;
 
+#define CTL_PROVIDER_TYPE cu_memory_provider_t
+#include "provider_ctl_stats_impl.h"
+
 // CUDA Memory Provider settings struct
 typedef struct umf_cuda_memory_provider_params_t {
     // Handle to the CUDA context
@@ -98,6 +103,9 @@ static cu_ops_t g_cu_ops;
 static UTIL_ONCE_FLAG cu_is_initialized = UTIL_ONCE_FLAG_INIT;
 static bool Init_cu_global_state_failed;
 
+struct ctl cu_memory_ctl_root;
+static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
+
 // forward decl needed for alloc
 static umf_result_t cu_memory_provider_free(void *provider, void *ptr,
                                             size_t bytes);
@@ -133,6 +141,10 @@ static umf_result_t cu2umf_result(CUresult result) {
     }
 }
 
+static void initialize_cu_ctl(void) {
+    CTL_REGISTER_MODULE(&cu_memory_ctl_root, stats);
+}
+
 static void init_cu_global_state(void) {
 #ifdef _WIN32
     const char *lib_name = "nvcuda.dll";
@@ -333,6 +345,7 @@ static umf_result_t cu_memory_provider_initialize(const void *params,
     if (!cu_provider) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
+    memset(cu_provider, 0, sizeof(cu_memory_provider_t));
 
     // CUDA alloc functions doesn't allow to provide user alignment - get the
     // minimum one from the driver
@@ -460,6 +473,8 @@ static umf_result_t cu_memory_provider_alloc(void *provider, size_t size,
         LOG_ERR("unsupported alignment size");
         return UMF_RESULT_ERROR_INVALID_ALIGNMENT;
     }
+
+    provider_ctl_stats_alloc(cu_provider, size);
     return umf_result;
 }
 
@@ -504,7 +519,11 @@ static umf_result_t cu_memory_provider_free(void *provider, void *ptr,
         LOG_ERR("Failed to restore CUDA context, ret = %d", umf_result);
     }
 
-    return cu2umf_result(cu_result);
+    umf_result_t ret = cu2umf_result(cu_result);
+    if (ret == UMF_RESULT_SUCCESS) {
+        provider_ctl_stats_free(cu_provider, bytes);
+    }
+    return ret;
 }
 
 static umf_result_t
@@ -685,6 +704,14 @@ cu_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t cu_ctl(void *provider, umf_ctl_query_source_t operationType,
+                           const char *name, void *arg, size_t size,
+                           umf_ctl_query_type_t query_type, va_list args) {
+    utils_init_once(&ctl_initialized, initialize_cu_ctl);
+    return ctl_query(&cu_memory_ctl_root, provider, operationType, name,
+                     query_type, arg, size, args);
+}
+
 static umf_memory_provider_ops_t UMF_CUDA_MEMORY_PROVIDER_OPS = {
     .version = UMF_PROVIDER_OPS_VERSION_CURRENT,
     .initialize = cu_memory_provider_initialize,
@@ -707,6 +734,7 @@ static umf_memory_provider_ops_t UMF_CUDA_MEMORY_PROVIDER_OPS = {
     .ext_put_ipc_handle = cu_memory_provider_put_ipc_handle,
     .ext_open_ipc_handle = cu_memory_provider_open_ipc_handle,
     .ext_close_ipc_handle = cu_memory_provider_close_ipc_handle,
+    .ext_ctl = cu_ctl,
 };
 
 const umf_memory_provider_ops_t *umfCUDAMemoryProviderOps(void) {

src/provider/provider_devdax_memory.c

@@ -67,6 +67,7 @@ umf_result_t umfDevDaxMemoryProviderParamsSetProtection(
 #include "base_alloc_global.h"
 #include "coarse.h"
 #include "libumf.h"
+#include "provider_ctl_stats_type.h"
 #include "utils_common.h"
 #include "utils_concurrency.h"
 #include "utils_log.h"
@@ -83,8 +84,12 @@ typedef struct devdax_memory_provider_t {
     utils_mutex_t lock;  // lock of ptr and offset
     unsigned protection; // combination of OS-specific protection flags
     coarse_t *coarse;    // coarse library handle
+    ctl_stats_t stats;
 } devdax_memory_provider_t;
 
+#define CTL_PROVIDER_TYPE devdax_memory_provider_t
+#include "provider_ctl_stats_impl.h"
+
 // DevDax Memory provider settings struct
 typedef struct umf_devdax_memory_provider_params_t {
     char *path;
@@ -112,6 +117,9 @@ static __TLS devdax_last_native_error_t TLS_last_native_error;
 #define _UMF_DEVDAX_RESULT_ERROR_PURGE_FORCE_FAILED                            \
     (UMF_DEVDAX_RESULT_ERROR_PURGE_FORCE_FAILED - UMF_DEVDAX_RESULT_SUCCESS)
 
+struct ctl devdax_memory_ctl_root;
+static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
+
 static const char *Native_error_str[] = {
     [_UMF_DEVDAX_RESULT_SUCCESS] = "success",
     [_UMF_DEVDAX_RESULT_ERROR_ALLOC_FAILED] = "memory allocation failed",
@@ -127,6 +135,10 @@ static void devdax_store_last_native_error(int32_t native_error,
     TLS_last_native_error.errno_value = errno_value;
 }
 
+static void initialize_devdax_ctl(void) {
+    CTL_REGISTER_MODULE(&devdax_memory_ctl_root, stats);
+}
+
 static umf_result_t
 devdax_translate_params(const umf_devdax_memory_provider_params_t *in_params,
                         devdax_memory_provider_t *provider) {
@@ -284,7 +296,12 @@ static umf_result_t devdax_alloc(void *provider, size_t size, size_t alignment,
                                  void **resultPtr) {
     devdax_memory_provider_t *devdax_provider =
         (devdax_memory_provider_t *)provider;
-    return coarse_alloc(devdax_provider->coarse, size, alignment, resultPtr);
+    umf_result_t ret =
+        coarse_alloc(devdax_provider->coarse, size, alignment, resultPtr);
+    if (ret == UMF_RESULT_SUCCESS) {
+        provider_ctl_stats_alloc(devdax_provider, size);
+    }
+    return ret;
 }
 
 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,
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t devdax_ctl(void *provider,
+                               umf_ctl_query_source_t operationType,
+                               const char *name, void *arg, size_t size,
+                               umf_ctl_query_type_t query_type, va_list args) {
+    utils_init_once(&ctl_initialized, initialize_devdax_ctl);
+    return ctl_query(&devdax_memory_ctl_root, provider, operationType, name,
+                     query_type, arg, size, args);
+}
+
 static umf_result_t devdax_free(void *provider, void *ptr, size_t size) {
     devdax_memory_provider_t *devdax_provider =
         (devdax_memory_provider_t *)provider;
-    return coarse_free(devdax_provider->coarse, ptr, size);
+    umf_result_t ret = coarse_free(devdax_provider->coarse, ptr, size);
+    if (ret == UMF_RESULT_SUCCESS) {
+        provider_ctl_stats_free(devdax_provider, size);
+    }
+    return ret;
 }
 
 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 = {
     .ext_get_ipc_handle = devdax_get_ipc_handle,
     .ext_put_ipc_handle = devdax_put_ipc_handle,
     .ext_open_ipc_handle = devdax_open_ipc_handle,
-    .ext_close_ipc_handle = devdax_close_ipc_handle};
+    .ext_close_ipc_handle = devdax_close_ipc_handle,
+    .ext_ctl = devdax_ctl};
 
 const umf_memory_provider_ops_t *umfDevDaxMemoryProviderOps(void) {
     return &UMF_DEVDAX_MEMORY_PROVIDER_OPS;

src/provider/provider_file_memory.c

@@ -74,6 +74,7 @@ umf_result_t umfFileMemoryProviderParamsSetVisibility(
 #include "coarse.h"
 #include "critnib.h"
 #include "libumf.h"
+#include "provider_ctl_stats_type.h"
 #include "utils_common.h"
 #include "utils_concurrency.h"
 #include "utils_log.h"
@@ -112,8 +113,12 @@ typedef struct file_memory_provider_t {
     critnib *fd_offset_map;
 
     coarse_t *coarse; // coarse library handle
+    ctl_stats_t stats;
 } file_memory_provider_t;
 
+#define CTL_PROVIDER_TYPE file_memory_provider_t
+#include "provider_ctl_stats_impl.h"
+
 // File Memory Provider settings struct
 typedef struct umf_file_memory_provider_params_t {
     char *path;
@@ -139,6 +144,9 @@ static __TLS file_last_native_error_t TLS_last_native_error;
 #define _UMF_FILE_RESULT_ERROR_PURGE_FORCE_FAILED                              \
     (UMF_FILE_RESULT_ERROR_PURGE_FORCE_FAILED - UMF_FILE_RESULT_SUCCESS)
 
+struct ctl file_memory_ctl_root;
+static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
+
 static const char *Native_error_str[] = {
     [_UMF_FILE_RESULT_SUCCESS] = "success",
     [_UMF_FILE_RESULT_ERROR_ALLOC_FAILED] = "memory allocation failed",
@@ -152,6 +160,10 @@ static void file_store_last_native_error(int32_t native_error,
     TLS_last_native_error.errno_value = errno_value;
 }
 
+static void initialize_file_ctl(void) {
+    CTL_REGISTER_MODULE(&file_memory_ctl_root, stats);
+}
+
 static umf_result_t
 file_translate_params(const umf_file_memory_provider_params_t *in_params,
                       file_memory_provider_t *provider) {
@@ -496,7 +508,12 @@ static umf_result_t file_alloc_aligned(file_memory_provider_t *file_provider,
 static umf_result_t file_alloc(void *provider, size_t size, size_t alignment,
                                void **resultPtr) {
     file_memory_provider_t *file_provider = (file_memory_provider_t *)provider;
-    return coarse_alloc(file_provider->coarse, size, alignment, resultPtr);
+    umf_result_t ret =
+        coarse_alloc(file_provider->coarse, size, alignment, resultPtr);
+    if (ret == UMF_RESULT_SUCCESS) {
+        provider_ctl_stats_alloc(file_provider, size);
+    }
+    return ret;
 }
 
 static umf_result_t file_alloc_cb(void *provider, size_t size, size_t alignment,
@@ -859,9 +876,22 @@ static umf_result_t file_close_ipc_handle(void *provider, void *ptr,
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t file_ctl(void *provider,
+                             umf_ctl_query_source_t operationType,
+                             const char *name, void *arg, size_t size,
+                             umf_ctl_query_type_t query_type, va_list args) {
+    utils_init_once(&ctl_initialized, initialize_file_ctl);
+    return ctl_query(&file_memory_ctl_root, provider, operationType, name,
+                     query_type, arg, size, args);
+}
+
 static umf_result_t file_free(void *provider, void *ptr, size_t size) {
     file_memory_provider_t *file_provider = (file_memory_provider_t *)provider;
-    return coarse_free(file_provider->coarse, ptr, size);
+    umf_result_t ret = coarse_free(file_provider->coarse, ptr, size);
+    if (ret == UMF_RESULT_SUCCESS) {
+        provider_ctl_stats_free(file_provider, size);
+    }
+    return ret;
 }
 
 static umf_memory_provider_ops_t UMF_FILE_MEMORY_PROVIDER_OPS = {
@@ -882,7 +912,8 @@ static umf_memory_provider_ops_t UMF_FILE_MEMORY_PROVIDER_OPS = {
     .ext_get_ipc_handle = file_get_ipc_handle,
     .ext_put_ipc_handle = file_put_ipc_handle,
     .ext_open_ipc_handle = file_open_ipc_handle,
-    .ext_close_ipc_handle = file_close_ipc_handle};
+    .ext_close_ipc_handle = file_close_ipc_handle,
+    .ext_ctl = file_ctl};
 
 const umf_memory_provider_ops_t *umfFileMemoryProviderOps(void) {
     return &UMF_FILE_MEMORY_PROVIDER_OPS;

src/provider/provider_level_zero.c

@@ -376,9 +376,10 @@ static ze_relaxed_allocation_limits_exp_desc_t relaxed_device_allocation_desc =
      .pNext = NULL,
      .flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE};
 
-static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
-                                             size_t alignment,
-                                             void **resultPtr) {
+static umf_result_t ze_memory_provider_alloc_helper(void *provider, size_t size,
+                                                    size_t alignment,
+                                                    int update_stats,
+                                                    void **resultPtr) {
     ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
 
     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,
         }
     }
 
-    provider_ctl_stats_alloc(ze_provider, size);
-
+    if (update_stats) {
+        provider_ctl_stats_alloc(ze_provider, size);
+    }
     return UMF_RESULT_SUCCESS;
 }
 
-static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
-                                            size_t bytes) {
+static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
+                                             size_t alignment,
+                                             void **resultPtr) {
+    return ze_memory_provider_alloc_helper(provider, size, alignment, 1,
+                                           resultPtr);
+}
+
+static umf_result_t ze_memory_provider_free_helper(void *provider, void *ptr,
+                                                   size_t bytes,
+                                                   int update_stats) {
     if (ptr == NULL) {
         return UMF_RESULT_SUCCESS;
     }
 
     ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
-
+    umf_result_t ret;
     if (ze_provider->freePolicyFlags == 0) {
-        return ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
-    }
+        ret = ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
+    } else {
+        ze_memory_free_ext_desc_t desc = {
+            .stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
+            .pNext = NULL,
+            .freePolicy = ze_provider->freePolicyFlags};
 
-    ze_memory_free_ext_desc_t desc = {
-        .stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
-        .pNext = NULL,
-        .freePolicy = ze_provider->freePolicyFlags};
+        ret = ze2umf_result(
+            g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
+    }
 
-    umf_result_t ret =
-        ze2umf_result(g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
     if (ret != UMF_RESULT_SUCCESS) {
         return ret;
     }
 
-    provider_ctl_stats_free(ze_provider, bytes);
+    if (update_stats) {
+        provider_ctl_stats_free(ze_provider, bytes);
+    }
 
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
+                                            size_t bytes) {
+    return ze_memory_provider_free_helper(provider, ptr, bytes, 1);
+}
+
 static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
                                         size_t *min_page_size) {
     assert(min_page_size);
 
     LOG_DEBUG("Querying minimum page size");
 
     void *ptr;
-    umf_result_t result = ze_memory_provider_alloc(ze_provider, 1, 0, &ptr);
+    umf_result_t result =
+        ze_memory_provider_alloc_helper(ze_provider, 1, 0, 0, &ptr);
     if (result != UMF_RESULT_SUCCESS) {
         return result;
     }
@@ -494,7 +513,7 @@ static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
 
     *min_page_size = properties.pageSize;
 
-    ze_memory_provider_free(ze_provider, ptr, 1);
+    ze_memory_provider_free_helper(ze_provider, ptr, 1, 0);
 
     return ze2umf_result(ze_result);
 }