Add resident device change call

include/umf/memory_pool_ops.h

@@ -167,7 +167,7 @@ typedef struct umf_memory_pool_ops_t {
     ///
     /// @return umf_result_t result of the control operation.
     ///
-    umf_result_t (*ext_ctl)(void *hPool, umf_ctl_query_source_t source,
+    umf_result_t (*ext_ctl)(void *pool, umf_ctl_query_source_t source,
                             const char *name, void *arg, size_t size,
                             umf_ctl_query_type_t queryType, va_list args);
 
@@ -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 device device handle
+    /// @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, void *device,
+                                               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 device device handle
+/// @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,
+                                      void *device, 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,16 @@ 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 device handle
+    /// @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, void *device,
+                                               bool is_adding);
+
 } umf_memory_provider_ops_t;
 
 #ifdef __cplusplus

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

@@ -19,6 +19,7 @@ EXPORTS
     umfCUDAMemoryProviderParamsSetContext
     umfCUDAMemoryProviderParamsSetDevice
     umfCUDAMemoryProviderParamsSetMemoryType
+    umfLevelZeroMemoryProviderParamsSetResidentDevices
     umfDevDaxMemoryProviderOps
     umfDevDaxMemoryProviderParamsCreate
     umfDevDaxMemoryProviderParamsDestroy
@@ -41,7 +42,6 @@ EXPORTS
     umfLevelZeroMemoryProviderParamsSetContext
     umfLevelZeroMemoryProviderParamsSetDevice
     umfLevelZeroMemoryProviderParamsSetMemoryType
-    umfLevelZeroMemoryProviderParamsSetResidentDevices
     umfMemoryProviderAlloc
     umfMemoryProviderAllocationMerge
     umfMemoryProviderAllocationSplit
@@ -60,6 +60,7 @@ EXPORTS
     umfMemoryProviderPurgeForce
     umfMemoryProviderPurgeLazy
     umfMemoryProviderPutIPCHandle
+    umfMemoryProviderResidentDeviceChange
     umfMempolicyCreate
     umfMempolicyDestroy
     umfMempolicySetCustomSplitPartitions

src/libumf.map

@@ -151,6 +151,7 @@ UMF_1.1 {
     umfGetMemoryProperty;
     umfGetMemoryPropertySize;
     umfJemallocPoolParamsSetName;
+    umfMemoryProviderResidentDeviceChange;
     umfLevelZeroMemoryProviderParamsSetName;
     umfOsMemoryProviderParamsSetName;
     umfPoolTrimMemory;

src/memory_provider.c

@@ -154,6 +154,14 @@ static umf_result_t umfDefaultGetAllocationPropertiesSize(
     return UMF_RESULT_ERROR_NOT_SUPPORTED;
 }
 
+static umf_result_t umfDefaultResidentDeviceChange(void *provider, void *device,
+                                                   bool is_adding) {
+    (void)provider;
+    (void)device;
+    (void)is_adding;
+    return UMF_RESULT_ERROR_NOT_SUPPORTED;
+}
+
 void assignOpsExtDefaults(umf_memory_provider_ops_t *ops) {
     if (!ops->ext_purge_lazy) {
         ops->ext_purge_lazy = umfDefaultPurgeLazy;
@@ -183,6 +191,10 @@ void assignOpsExtDefaults(umf_memory_provider_ops_t *ops) {
         ops->ext_get_allocation_properties_size =
             umfDefaultGetAllocationPropertiesSize;
     }
+
+    if (!ops->ext_resident_device_change) {
+        ops->ext_resident_device_change = umfDefaultResidentDeviceChange;
+    }
 }
 
 void assignOpsIpcDefaults(umf_memory_provider_ops_t *ops) {
@@ -606,3 +618,16 @@ umf_result_t umfMemoryProviderGetAllocationPropertiesSize(
     checkErrorAndSetLastProvider(res, hProvider);
     return res;
 }
+
+umf_result_t
+umfMemoryProviderResidentDeviceChange(umf_memory_provider_handle_t hProvider,
+                                      void *device, bool isAdding) {
+
+    UMF_CHECK((hProvider != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT);
+
+    umf_result_t res = hProvider->ops.ext_resident_device_change(
+        hProvider->provider_priv, device, isAdding);
+
+    checkErrorAndSetLastProvider(res, hProvider);
+    return res;
+}

src/provider/provider_cuda.c

@@ -829,6 +829,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 = NULL,
 };
 
 const umf_memory_provider_ops_t *umfCUDAMemoryProviderOps(void) {