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);
 

include/umf/providers/provider_level_zero.h

@@ -8,6 +8,8 @@
 #ifndef UMF_LEVEL_ZERO_PROVIDER_H
 #define UMF_LEVEL_ZERO_PROVIDER_H
 
+#include <stdbool.h>
+
 #include <umf/memory_provider_gpu.h>
 
 #ifdef __cplusplus
@@ -101,6 +103,17 @@ umf_result_t umfLevelZeroMemoryProviderParamsSetDeviceOrdinal(
 umf_result_t umfLevelZeroMemoryProviderParamsSetName(
     umf_level_zero_memory_provider_params_handle_t hParams, const char *name);
 
+/// @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 umfLevelZeroMemoryProviderResidentDeviceChange(
+    umf_memory_provider_handle_t provider, ze_device_handle_t device,
+    bool is_adding);
+
 const umf_memory_provider_ops_t *umfLevelZeroMemoryProviderOps(void);
 
 #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

@@ -42,6 +42,7 @@ EXPORTS
     umfLevelZeroMemoryProviderParamsSetDevice
     umfLevelZeroMemoryProviderParamsSetMemoryType
     umfLevelZeroMemoryProviderParamsSetResidentDevices
+    umfLevelZeroMemoryProviderResidentDeviceChange
     umfMemoryProviderAlloc
     umfMemoryProviderAllocationMerge
     umfMemoryProviderAllocationSplit

src/libumf.map

@@ -152,6 +152,7 @@ UMF_1.1 {
     umfGetMemoryPropertySize;
     umfJemallocPoolParamsSetName;
     umfLevelZeroMemoryProviderParamsSetName;
+    umfLevelZeroMemoryProviderResidentDeviceChange;
     umfOsMemoryProviderParamsSetName;
     umfPoolTrimMemory;
     umfScalablePoolParamsSetName;