Add test for multiple NUMA nodes

Author: Damian Duy

test/CMakeLists.txt

@@ -102,6 +102,9 @@ if(UMF_BUILD_OS_MEMORY_PROVIDER AND LINUX) # OS-specific functions are implement
     add_umf_test(NAME provider_os_memory
                  SRCS provider_os_memory.cpp
                  LIBS umf_utils)
+    add_umf_test(NAME provider_os_memory_multiple_NUMA_nodes
+                 SRCS provider_os_memory_multiple_NUMA_nodes.cpp
+                 LIBS umf_utils)
     add_umf_test(NAME memspace_numa
                  SRCS memspace_numa.cpp
                  LIBS numa)

test/provider_os_memory.cpp

@@ -7,6 +7,8 @@
 #include "cpp_helpers.hpp"
 #include "provider_os_memory_internal.h"
 #include "umf/providers/provider_os_memory.h"
+#include <numa.h>
+#include <numaif.h>
 #include <umf/memory_provider.h>
 
 using umf_test::test;

test/provider_os_memory_multiple_NUMA_nodes.cpp

@@ -0,0 +1,102 @@
+// Copyright (C) 2024 Intel Corporation
+// Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "base.hpp"
+
+#include "cpp_helpers.hpp"
+#include "provider_os_memory_internal.h"
+#include "umf/providers/provider_os_memory.h"
+#include <numa.h>
+#include <numaif.h>
+#include <umf/memory_provider.h>
+
+using umf_test::test;
+
+static umf_os_memory_provider_params_t UMF_OS_MEMORY_PROVIDER_PARAMS_TEST = {
+    /* .protection = */ UMF_PROTECTION_READ | UMF_PROTECTION_WRITE,
+    /* .visibility = */ UMF_VISIBILITY_PRIVATE,
+
+    // NUMA config
+    /* .nodemask = */ NULL,
+    /* .maxnode = */ 0,
+    /* .numa_mode = */ UMF_NUMA_MODE_BIND,
+    /* .numa_flags = */ 0,
+
+    // others
+    /* .traces = */ 1,
+};
+
+std::vector<std::tuple<int, int>> get_available_NUMA_nodes_numbers() { 
+    int max_node = numa_max_node() + 1;
+    std::vector<std::tuple<int, int>> available_NUMA_nodes(max_node);
+
+    // Get all available NUMA nodes numbers
+    for (int i = 0; i <= max_node; ++i) {
+        available_NUMA_nodes.emplace_back(i, numa_bitmask_isbitset(numa_all_nodes_ptr, i));
+    }
+
+    return available_NUMA_nodes;
+}
+
+std::tuple<int, int> get_numa_node(void* addr) {
+    int max_node = numa_max_node() + 1;
+    int numa_node;
+    int ret = get_mempolicy(&numa_node, NULL, max_node, addr, MPOL_F_NODE | MPOL_F_ADDR);
+
+    return {ret, numa_node};
+}
+
+class testNUMAnodes:
+    public testing::TestWithParam<std::tuple<int, int>> {};
+
+INSTANTIATE_TEST_SUITE_P(testNUMAnodesAllocations, testNUMAnodes,
+                         ::testing::ValuesIn(get_available_NUMA_nodes_numbers()));
+
+// Tests for allocations on numa nodes
+TEST_P(testNUMAnodes, checkNUMAnodesAllocations) {
+    if (numa_available() == -1 || numa_num_task_nodes() <= 1) {
+        GTEST_SKIP();
+    }
+
+    size_t alloc_size = 1024;
+    auto [i, is_bit_set]= GetParam();
+    int max_node = numa_max_node() + 1;
+
+    umf_result_t umf_result;
+    umf_memory_provider_handle_t os_memory_provider = nullptr;
+    umf_os_memory_provider_params_t os_memory_provider_params =
+        UMF_OS_MEMORY_PROVIDER_PARAMS_TEST;
+    os_memory_provider_params.maxnode = max_node;
+
+    bitmask * nodemask = numa_allocate_nodemask();
+    ASSERT_NE(nodemask, nullptr);
+
+    if (is_bit_set == 1) {
+            numa_bitmask_setbit(nodemask, i);
+            os_memory_provider_params.nodemask = nodemask->maskp;
+            umf_result = umfMemoryProviderCreate(&UMF_OS_MEMORY_PROVIDER_OPS,
+                                            &os_memory_provider_params,
+                                            &os_memory_provider);
+            ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+            ASSERT_NE(os_memory_provider, nullptr);
+
+            void *ptr = nullptr;
+            umf_result = umfMemoryProviderAlloc(os_memory_provider, alloc_size, 0, &ptr);
+            ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+            ASSERT_NE(ptr, nullptr);
+
+            auto [ret, numa_node] = get_numa_node(ptr);
+            ASSERT_EQ(ret, 0);
+            ASSERT_EQ(numa_node, i);
+
+            memset(ptr, 0xFF, alloc_size);
+            umf_result = umfMemoryProviderFree(os_memory_provider, ptr, alloc_size);
+            ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
+
+            umfMemoryProviderDestroy(os_memory_provider);
+
+            numa_bitmask_clearall(nodemask);
+            numa_free_nodemask(nodemask);
+        }
+}