add memspace filter functions

include/umf/base.h

@@ -1,6 +1,6 @@
 /*
  *
- * Copyright (C) 2023 Intel Corporation
+ * Copyright (C) 2023-2024 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -43,7 +43,8 @@ typedef enum umf_result_t {
     UMF_RESULT_ERROR_INVALID_ALIGNMENT =
         4,                              ///< Invalid alignment of an argument
     UMF_RESULT_ERROR_NOT_SUPPORTED = 5, ///< Operation not supported
-
+    UMF_RESULT_ERROR_USER_SPECIFIC =
+        6, ///< Failure in user provider code (i.e in user provided callback)
     UMF_RESULT_ERROR_UNKNOWN = 0x7ffffffe ///< Unknown or internal error
 } umf_result_t;
 

include/umf/memspace.h

@@ -131,6 +131,58 @@ umf_result_t
 umfMemspaceMemtargetRemove(umf_memspace_handle_t hMemspace,
                            umf_const_memtarget_handle_t hMemtarget);
 
+/// \brief Clones memspace.
+///
+/// \param hMemspace handle to memspace
+/// \param hNewMemspace [out] handle to the newly created memspace
+/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+///
+umf_result_t umfMemspaceClone(umf_const_memspace_handle_t hMemspace,
+                              umf_memspace_handle_t *hNewMemspace);
+
+/// \brief Custom filter function for umfMemspaceUserFilter
+///
+/// \param hMemspace handle to memspace
+/// \param hMemtarget handle to memory target
+/// \param args user provided arguments
+/// \return zero if hMemtarget should be removed from memspace, positive otherwise, and negative on error
+///
+typedef int (*umf_memspace_filter_func_t)(
+    umf_const_memspace_handle_t hMemspace,
+    umf_const_memtarget_handle_t hMemtarget, void *args);
+
+/// \brief Removes all memory targets with not matching numa node ids.
+///
+/// \param hMemspace handle to memspace
+/// \param ids array of numa node ids
+/// \param size size of the array
+/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+/// If the error code is UMF_RESULT_UNKNOWN the memspace is corrupted, otherwise the memspace is not modified.
+///
+umf_result_t umfMemspaceFilterById(umf_memspace_handle_t hMemspace,
+                                   unsigned *ids, size_t size);
+
+/// \brief Filters out memory targets that capacity is less than specified size.
+///
+/// \param hMemspace handle to memspace
+/// \param size minimum capacity of memory target
+/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+/// If the error code is UMF_RESULT_UNKNOWN the memspace is corrupted, otherwise the memspace is not modified.
+///
+umf_result_t umfMemspaceFilterByCapacity(umf_memspace_handle_t hMemspace,
+                                         size_t size);
+
+/// \brief Filters out memory targets based on user provided function
+///
+/// \param hMemspace handle to memspace
+/// \param filter user provided function
+/// \param args user provided arguments
+/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+/// If the error code is UMF_RESULT_UNKNOWN the memspace is corrupted, otherwise the memspace is not modified.
+///
+umf_result_t umfMemspaceUserFilter(umf_memspace_handle_t hMemspace,
+                                   umf_memspace_filter_func_t filter,
+                                   void *args);
 #ifdef __cplusplus
 }
 #endif

include/umf/memtarget.h

@@ -39,6 +39,13 @@ umf_result_t umfMemtargetGetType(umf_const_memtarget_handle_t hMemtarget,
 umf_result_t umfMemtargetGetCapacity(umf_const_memtarget_handle_t hMemtarget,
                                      size_t *capacity);
 
+/// \brief Get physical ID of the memory target.
+/// \param hMemtarget handle to the memory target
+/// \param id [out] id of the memory target
+/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+umf_result_t umfMemtargetGetId(umf_const_memtarget_handle_t hMemtarget,
+                               unsigned *id);
+
 #ifdef __cplusplus
 }
 #endif

src/libumf.def.in

@@ -40,13 +40,18 @@ EXPORTS
     umfMempolicyDestroy
     umfMempolicySetCustomSplitPartitions
     umfMempolicySetInterleavePartSize
-    umfMemspaceNew
+    umfMemspaceClone
     umfMemspaceDestroy
-    umfMemspaceMemtargetNum
-    umfMemspaceMemtargetGet
-    umfMemtargetGetCapacity
+    umfMemspaceFilterByCapacity
+    umfMemspaceFilterById
     umfMemspaceMemtargetAdd
+    umfMemspaceMemtargetGet
+    umfMemspaceMemtargetNum
     umfMemspaceMemtargetRemove
+    umfMemspaceNew
+    umfMemspaceUserFilter
+    umfMemtargetGetCapacity
+    umfMemtargetGetId
     umfMemtargetGetType
     umfOpenIPCHandle    
     umfPoolAlignedMalloc

src/libumf.map.in

@@ -34,14 +34,19 @@ UMF_1.0 {
         umfMempolicyDestroy;
         umfMempolicySetCustomSplitPartitions;
         umfMempolicySetInterleavePartSize;
-        umfMemspaceNew;
+        umfMemspaceClone;
         umfMemspaceCreateFromNumaArray;
         umfMemspaceDestroy;
+        umfMemspaceFilterByCapacity;
+        umfMemspaceFilterById;
         umfMemspaceMemtargetAdd;
-        umfMemspaceMemtargetRemove;
-        umfMemspaceMemtargetNum;
         umfMemspaceMemtargetGet;
+        umfMemspaceMemtargetNum;
+        umfMemspaceMemtargetRemove;
+        umfMemspaceNew;
+        umfMemspaceUserFilter;
         umfMemtargetGetCapacity;
+        umfMemtargetGetId;
         umfMemtargetGetType;
         umfOpenIPCHandle;
         umfPoolAlignedMalloc;

src/memspace.c

@@ -426,3 +426,146 @@ umfMemspaceMemtargetRemove(umf_memspace_handle_t hMemspace,
     hMemspace->size--;
     return UMF_RESULT_SUCCESS;
 }
+
+// Helper function - returns zero on success, negative in case of error in filter function
+// and positive error code, in case of other errors.
+static int umfMemspaceFilterHelper(umf_memspace_handle_t memspace,
+                                   umf_memspace_filter_func_t filter,
+                                   void *args) {
+
+    if (!memspace || !filter) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    size_t idx = 0;
+    int ret;
+    umf_memtarget_handle_t *nodesToRemove =
+        umf_ba_global_alloc(sizeof(*nodesToRemove) * memspace->size);
+    if (!nodesToRemove) {
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    }
+
+    for (size_t i = 0; i < memspace->size; i++) {
+        ret = filter(memspace, memspace->nodes[i], args);
+        if (ret < 0) {
+            LOG_ERR("filter function failed");
+            goto free_mem;
+        } else if (ret == 0) {
+            nodesToRemove[idx++] = memspace->nodes[i];
+        }
+    }
+
+    size_t i = 0;
+    for (; i < idx; i++) {
+        ret = umfMemspaceMemtargetRemove(memspace, nodesToRemove[i]);
+        if (ret != UMF_RESULT_SUCCESS) {
+            goto re_add;
+        }
+    }
+
+    umf_ba_global_free(nodesToRemove);
+    return UMF_RESULT_SUCCESS;
+
+re_add:
+    // If target removal failed, add back previously removed targets.
+    for (size_t j = 0; j < i; j++) {
+        umf_result_t ret2 = umfMemspaceMemtargetAdd(memspace, nodesToRemove[j]);
+        if (ret2 != UMF_RESULT_SUCCESS) {
+            ret =
+                UMF_RESULT_ERROR_UNKNOWN; // indicate that memspace is corrupted
+            break;
+        }
+    }
+free_mem:
+    umf_ba_global_free(nodesToRemove);
+    return ret;
+}
+
+umf_result_t umfMemspaceUserFilter(umf_memspace_handle_t memspace,
+                                   umf_memspace_filter_func_t filter,
+                                   void *args) {
+
+    if (!memspace || !filter) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    int ret = umfMemspaceFilterHelper(memspace, filter, args);
+    if (ret < 0) {
+        return UMF_RESULT_ERROR_USER_SPECIFIC;
+    }
+
+    return ret;
+}
+
+typedef struct filter_by_id_args {
+    unsigned *ids; // array of numa nodes ids
+    size_t size;   // size of the array
+} filter_by_id_args_t;
+
+/*
+ * The following predefined filter callbacks returns umf_result_t codes as negative value
+ * as negative only value are treated as an error. umfMemspaceFilter will passthru this error code
+ * so and calling filter function has to translate this code to positive umf_result_t code.
+ */
+
+static int filterById(umf_const_memspace_handle_t memspace,
+                      umf_const_memtarget_handle_t target, void *args) {
+    if (!memspace || !target || !args) {
+        return -UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    filter_by_id_args_t *filterArgs = args;
+    for (size_t i = 0; i < filterArgs->size; i++) {
+        unsigned id;
+        umf_result_t ret = umfMemtargetGetId(target, &id);
+        if (ret != UMF_RESULT_SUCCESS) {
+            return -ret;
+        }
+
+        if (id == filterArgs->ids[i]) {
+            return 1;
+        }
+    }
+    return 0;
+}
+
+static int filterByCapacity(umf_const_memspace_handle_t memspace,
+                            umf_const_memtarget_handle_t target, void *args) {
+    if (!memspace || !target || !args) {
+        return -UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    size_t capacity;
+    umf_result_t ret = umfMemtargetGetCapacity(target, &capacity);
+    if (ret != UMF_RESULT_SUCCESS) {
+        return -ret;
+    }
+
+    size_t *targetCapacity = args;
+    return (capacity >= *targetCapacity) ? 1 : 0;
+}
+
+umf_result_t umfMemspaceFilterById(umf_memspace_handle_t memspace,
+                                   unsigned *ids, size_t size) {
+    if (!memspace || !ids || size == 0) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    filter_by_id_args_t args = {ids, size};
+    int ret = umfMemspaceFilterHelper(memspace, &filterById, &args);
+
+    // if umfMemspaceFilter() returned negative umf_result_t change it to positive
+    return ret < 0 ? -ret : ret;
+}
+
+umf_result_t umfMemspaceFilterByCapacity(umf_memspace_handle_t memspace,
+                                         size_t capacity) {
+    if (!memspace) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    int ret = umfMemspaceFilterHelper(memspace, &filterByCapacity, &capacity);
+
+    // if umfMemspaceFilter() returned negative umf_result_t change it to positive
+    return ret < 0 ? -ret : ret;
+}

src/memspace_internal.h

@@ -24,12 +24,6 @@ struct umf_memspace_t {
     umf_memtarget_handle_t *nodes;
 };
 
-///
-/// \brief Clones memspace
-///
-umf_result_t umfMemspaceClone(umf_const_memspace_handle_t hMemspace,
-                              umf_memspace_handle_t *outHandle);
-
 typedef umf_result_t (*umfGetPropertyFn)(umf_const_memtarget_handle_t,
                                          uint64_t *);
 

src/memtarget.c

@@ -107,6 +107,14 @@ umf_result_t umfMemtargetGetLatency(umf_memtarget_handle_t srcMemoryTarget,
                                              dstMemoryTarget->priv, latency);
 }
 
+umf_result_t umfMemtargetGetId(umf_const_memtarget_handle_t hMemtarget,
+                               unsigned *id) {
+    if (!hMemtarget || !id) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    return hMemtarget->ops->get_id(hMemtarget->priv, id);
+}
 umf_result_t umfMemtargetGetType(umf_const_memtarget_handle_t memoryTarget,
                                  umf_memtarget_type_t *type) {
     if (!memoryTarget || !type) {

src/memtarget_ops.h

@@ -45,6 +45,7 @@ typedef struct umf_memtarget_ops_t {
                                 size_t *latency);
 
     umf_result_t (*get_type)(void *memoryTarget, umf_memtarget_type_t *type);
+    umf_result_t (*get_id)(void *memoryTarget, unsigned *type);
     umf_result_t (*compare)(void *memTarget, void *otherMemTarget, int *result);
 
 } umf_memtarget_ops_t;

src/memtargets/memtarget_numa.c

@@ -326,6 +326,15 @@ static umf_result_t numa_get_type(void *memTarget, umf_memtarget_type_t *type) {
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t numa_get_id(void *memTarget, unsigned *id) {
+    if (!memTarget || !id) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    *id = ((struct numa_memtarget_t *)memTarget)->physical_id;
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_result_t numa_compare(void *memTarget, void *otherMemTarget,
                                  int *result) {
     if (!memTarget || !otherMemTarget || !result) {
@@ -350,6 +359,7 @@ struct umf_memtarget_ops_t UMF_MEMTARGET_NUMA_OPS = {
     .get_bandwidth = numa_get_bandwidth,
     .get_latency = numa_get_latency,
     .get_type = numa_get_type,
+    .get_id = numa_get_id,
     .compare = numa_compare,
     .memory_provider_create_from_memspace =
         numa_memory_provider_create_from_memspace};

src/provider/provider_level_zero.c

@@ -188,12 +188,7 @@ static void ze_memory_provider_finalize(void *provider) {
         return;
     }
 
-    utils_init_once(&ze_is_initialized, init_ze_global_state);
     umf_ba_global_free(provider);
-
-    // portable version of "ze_is_initialized = UTIL_ONCE_FLAG_INIT;"
-    static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
-    memcpy(&ze_is_initialized, &is_initialized, sizeof(ze_is_initialized));
 }
 
 static bool use_relaxed_allocation(ze_memory_provider_t *ze_provider,

test/memspaces/memspace_fixtures.hpp

@@ -102,6 +102,18 @@ struct memspaceProviderTest : ::memspaceGetTest {
     umf_memory_provider_handle_t hProvider = nullptr;
 };
 
+struct numaNodesCapacityTest : numaNodesTest {
+    void SetUp() override {
+        numaNodesTest::SetUp();
+
+        for (auto nodeId : nodeIds) {
+            capacities.push_back(numa_node_size64(nodeId, nullptr));
+        }
+    }
+
+    std::vector<size_t> capacities;
+};
+
 GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(memspaceGetTest);
 GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(memspaceProviderTest);