Add resident device change call

Author: lslusarczyk

include/umf/memory_pool_ops.h

@@ -191,6 +191,18 @@ typedef struct umf_memory_pool_ops_t {
     ///         failure.
     ///
     umf_result_t (*ext_trim_memory)(void *pool, size_t minBytesToKeep);
+
+    ///
+    /// @brief Adds or removes devices on which allocations should be made
+    ///        resident.
+    /// @param pool pointer to the memory pool
+    /// @param peerIdx identifier of device
+    /// @param isAdding boolean indicating if peer is to be removed or added
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on
+    ///         failure.
+    ///
+    umf_result_t (*ext_resident_device_change)(void *pool, uint32_t peerIdx,
+                                               _Bool isAdding);
 } umf_memory_pool_ops_t;
 
 #ifdef __cplusplus

include/umf/memory_provider.h

@@ -265,6 +265,17 @@ umf_result_t
 umfMemoryProviderAllocationMerge(umf_memory_provider_handle_t hProvider,
                                  void *lowPtr, void *highPtr, size_t totalSize);
 
+/// @brief Adds or removes devices on which allocations should be made
+///        resident.
+/// @param hProvider handle to the memory provider
+/// @param deviceIndex identifier of device
+/// @param isAdding boolean indicating if peer is to be removed or added
+/// @return UMF_RESULT_SUCCESS on success or appropriate error code on
+///         failure.
+umf_result_t
+umfMemoryProviderResidentDeviceChange(umf_memory_provider_handle_t hProvider,
+                                      uint32_t deviceIndex, bool isAdding);
+
 #ifdef __cplusplus
 }
 #endif

include/umf/memory_provider_ops.h

@@ -11,6 +11,7 @@
 #define UMF_MEMORY_PROVIDER_OPS_H 1
 
 #include <stdarg.h>
+#include <stdbool.h>
 
 #include <umf/base.h>
 
@@ -322,6 +323,17 @@ typedef struct umf_memory_provider_ops_t {
         void *provider, umf_memory_property_id_t memory_property_id,
         size_t *size);
 
+    /// @brief Adds or removes devices on which allocations should be made
+    ///        resident.
+    /// @param provider handle to the memory provider
+    /// @param device_index identifier of device
+    /// @param is_adding boolean indicating if peer is to be removed or added
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on
+    ///         failure.
+    umf_result_t (*ext_resident_device_change)(void *provider,
+                                               uint32_t device_index,
+                                               bool is_adding);
+
 } umf_memory_provider_ops_t;
 
 #ifdef __cplusplus

include/umf/providers/provider_level_zero.h

@@ -66,7 +66,8 @@ umf_result_t umfLevelZeroMemoryProviderParamsSetMemoryType(
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
 umf_result_t umfLevelZeroMemoryProviderParamsSetResidentDevices(
     umf_level_zero_memory_provider_params_handle_t hParams,
-    ze_device_handle_t *hDevices, uint32_t deviceCount);
+    ze_device_handle_t *hDevices, uint32_t deviceCount,
+    uint32_t *residentDevicesIndices, uint32_t residentDevicesCount);
 
 typedef enum umf_level_zero_memory_provider_free_policy_t {
     UMF_LEVEL_ZERO_MEMORY_PROVIDER_FREE_POLICY_DEFAULT =

src/critnib/critnib.c

@@ -1094,7 +1094,8 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
  *
  * If func() returns non-zero, the search is aborted.
  */
-static int iter(struct critnib_node *__restrict n, word min, word max,
+static int iter(struct critnib_node *__restrict n, const word min,
+                const word max,
                 int (*func)(word key, void *value, void *privdata),
                 void *privdata) {
     if (is_leaf(n)) {
@@ -1129,9 +1130,21 @@ static int iter(struct critnib_node *__restrict n, word min, word max,
 void critnib_iter(critnib *c, uintptr_t min, uintptr_t max,
                   int (*func)(uintptr_t key, void *value, void *privdata),
                   void *privdata) {
+    bool wasIterating = false;
     utils_mutex_lock(&c->mutex);
     if (c->root) {
         iter(c->root, min, max, func, privdata);
+        wasIterating = true;
     }
     utils_mutex_unlock(&c->mutex);
+    if (!wasIterating) {
+        LOG_DEBUG("there was no root, iterating critnib:%p was skipped",
+                  (void *)c);
+    }
+}
+
+void critnib_iter_all(critnib *c,
+                      int (*func)(uintptr_t key, void *value, void *privdata),
+                      void *privdata) {
+    critnib_iter(c, 0, (uintptr_t)-1, func, privdata);
 }

src/critnib/critnib.h

@@ -35,6 +35,10 @@ int critnib_insert(critnib *c, uintptr_t key, void *value, int update);
 void critnib_iter(critnib *c, uintptr_t min, uintptr_t max,
                   int (*func)(uintptr_t key, void *value, void *privdata),
                   void *privdata);
+void critnib_iter_all(critnib *c,
+                      int (*func)(uintptr_t key, void *value, void *privdata),
+                      void *privdata);
+
 int critnib_remove_release(critnib *c, uintptr_t key);
 
 /*

src/libumf.def

@@ -60,6 +60,7 @@ EXPORTS
     umfMemoryProviderPurgeForce
     umfMemoryProviderPurgeLazy
     umfMemoryProviderPutIPCHandle
+    umfMemoryProviderResidentDeviceChange
     umfMempolicyCreate
     umfMempolicyDestroy
     umfMempolicySetCustomSplitPartitions

src/libumf.map

@@ -54,6 +54,7 @@ UMF_1.0 {
         umfMemoryProviderPurgeForce;
         umfMemoryProviderPurgeLazy;
         umfMemoryProviderPutIPCHandle;
+	umfMemoryProviderResidentDeviceChange;
         umfMempolicyCreate;
         umfMempolicyDestroy;
         umfMempolicySetCustomSplitPartitions;

src/memory_provider.c

@@ -606,3 +606,13 @@ umf_result_t umfMemoryProviderGetAllocationPropertiesSize(
     checkErrorAndSetLastProvider(res, hProvider);
     return res;
 }
+
+umf_result_t
+umfMemoryProviderResidentDeviceChange(umf_memory_provider_handle_t hProvider,
+                                      uint32_t deviceIndex, bool isAdding) {
+    UMF_CHECK((hProvider != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT);
+    umf_result_t res = hProvider->ops.ext_resident_device_change(
+        hProvider->provider_priv, deviceIndex, isAdding);
+    checkErrorAndSetLastProvider(res, hProvider);
+    return res;
+}

src/provider/provider_cuda.c

@@ -802,6 +802,15 @@ static umf_result_t cu_memory_provider_get_allocation_properties_size(
     return UMF_RESULT_ERROR_INVALID_ARGUMENT;
 }
 
+static umf_result_t
+cu_memory_provider_resident_device_change(void *provider, uint32_t device_index,
+                                          bool is_adding) {
+    (void)provider;
+    (void)device_index;
+    (void)is_adding;
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_memory_provider_ops_t UMF_CUDA_MEMORY_PROVIDER_OPS = {
     .version = UMF_PROVIDER_OPS_VERSION_CURRENT,
     .initialize = cu_memory_provider_initialize,
@@ -829,6 +838,7 @@ static umf_memory_provider_ops_t UMF_CUDA_MEMORY_PROVIDER_OPS = {
         cu_memory_provider_get_allocation_properties,
     .ext_get_allocation_properties_size =
         cu_memory_provider_get_allocation_properties_size,
+    .ext_resident_device_change = cu_memory_provider_resident_device_change,
 };
 
 const umf_memory_provider_ops_t *umfCUDAMemoryProviderOps(void) {