Add test for TBB pool manager

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

README.md

@@ -25,7 +25,7 @@ A memory provider that provides memory from an operating system.
 Required packages:
 - C++ compiler with C++17 support
 - [CMake](https://cmake.org/) >= 3.14.0
-- Linux only: libnuma-dev
+- Linux only: libnuma, libdl
 
 For development and contributions:
 - clang-format-15.0 (can be installed with `python -m pip install clang-format==15.0.7`)

test/CMakeLists.txt

@@ -19,7 +19,9 @@ enable_testing()
 set(UMF_TEST_DIR ${CMAKE_CURRENT_SOURCE_DIR})
 
 if(LINUX)
-    set(LIBS_OS_SPECIFIC numa)
+    set(LIBS_OS_SPECIFIC
+        dl
+        numa)
 else()
     set(LIBS_OS_SPECIFIC "")
 endif()
@@ -57,4 +59,5 @@ add_umf_test(disjointPool disjoint_pool.cpp)
 
 if(LINUX) # OS-specific functions are implemented only for Linux now
     add_umf_test(provider_os_memory provider_os_memory.cpp)
+    add_umf_test(pool_tbb pool_tbb.cpp)
 endif()

test/pool_tbb.cpp

@@ -0,0 +1,298 @@
+// Copyright (C) 2023 Intel Corporation
+// Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// This file contains tests for UMF provider API
+
+#include <random>
+#include <set>
+#include <string>
+#include <unordered_map>
+
+#include "base.hpp"
+
+#include <umf/memory_pool.h>
+#include <umf/memory_provider.h>
+#include <umf/pools/pool_tbb.h>
+#include <umf/providers/provider_os_memory.h>
+// #include "provider_os_memory_internal.h"
+
+using umf_test::test;
+
+constexpr int PAGE_SIZE = 4 * 1024;
+
+static umf_os_memory_provider_params_t UMF_OS_MEMORY_PROVIDER_PARAMS = {
+    /* .protection = */ UMF_PROTECTION_READ | UMF_PROTECTION_WRITE,
+    /* .visibility = */ UMF_VISIBILITY_PRIVATE,
+
+    // NUMA config
+    /* .nodemask = */ NULL,
+    /* .maxnode = */ 0,
+    /* .numa_mode = */ UMF_NUMA_MODE_DEFAULT,
+    /* .numa_flags = */ 0,
+
+    // others
+    /* .traces = */ 1,
+};
+
+const size_t KB = 1024;
+const size_t MB = 1024 * KB;
+const size_t init_buffer_size = 20 * MB;
+
+TEST_F(test, pool_tbb_MallocPool_basic) {
+    umf_memory_provider_handle_t os_memory_provider;
+    umfMemoryProviderCreate(&UMF_OS_MEMORY_PROVIDER_OPS,
+                            &UMF_OS_MEMORY_PROVIDER_PARAMS,
+                            &os_memory_provider);
+    ASSERT_NE(os_memory_provider, nullptr);
+
+    umf_memory_pool_handle_t pool;
+    umfPoolCreate(&UMF_TBB_POOL_OPS, os_memory_provider, NULL, &pool);
+    ASSERT_NE(pool, nullptr);
+
+    // test
+
+    umf_memory_provider_handle_t prov = NULL;
+    umfPoolGetMemoryProvider(pool, &prov);
+    ASSERT_NE(prov, nullptr);
+
+    // alloc 2x 2MB
+    void *p1 = umfPoolMalloc(pool, 2 * MB);
+    void *p2 = umfPoolMalloc(pool, 2 * MB);
+    ASSERT_NE(p1, p2);
+
+    // swap pointers to get p1 < p2
+    if (p1 > p2) {
+        std::swap(p1, p2);
+    }
+
+    // free + alloc first block
+    // the block should be reused
+    // currently there is no purging, so the alloc size shouldn't change
+    // there should be no block merging between used and not-used blocks
+    umf_result_t res = umfPoolFree(pool, p1);
+    ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+    p1 = umfPoolMalloc(pool, 2 * MB);
+
+    // free all allocs
+    // overall alloc size shouldn't change
+    // block p2 should merge with the prev free block p1
+    // and the remaining init block
+    res = umfPoolFree(pool, p1);
+    ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+    res = umfPoolFree(pool, p2);
+    ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+
+    // alloc whole buffer
+    // after this, there should be one single block
+    p1 = umfPoolMalloc(pool, init_buffer_size);
+
+    // free all memory
+    // alloc 2 MB block - the init block should be split
+    res = umfPoolFree(pool, p1);
+    p1 = umfPoolMalloc(pool, 2 * MB);
+
+    // alloc additional 2 MB
+    // the non-used block should be used
+    p2 = umfPoolMalloc(pool, 2 * MB);
+    ASSERT_NE(p1, p2);
+
+    // make sure that p1 < p2
+    if (p1 > p2) {
+        std::swap(p1, p2);
+    }
+
+    // free blocks in order: p2, p1
+    // block p1 should merge with the next block p2
+    // swap pointers to get p1 < p2
+    umfPoolFree(pool, p2);
+    umfPoolFree(pool, p1);
+
+    // alloc 10x 2 MB - this should occupy all allocated memory
+    constexpr int allocs_size = 10;
+    void *allocs[allocs_size] = {0};
+    for (int i = 0; i < allocs_size; i++) {
+        allocs[i] = umfPoolMalloc(pool, 2 * MB);
+        ASSERT_NE(allocs[i], nullptr);
+    }
+    // there should be no block with the free memory
+
+    // free all memory
+    for (int i = 0; i < allocs_size; i++) {
+        res = umfPoolFree(pool, allocs[i]);
+        ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+    }
+
+    umfPoolDestroy(pool);
+    umfMemoryProviderDestroy(os_memory_provider);
+}
+
+TEST_F(test, pool_tbb_MallocPool_simple1) {
+    umf_memory_provider_handle_t os_memory_provider;
+    umfMemoryProviderCreate(&UMF_OS_MEMORY_PROVIDER_OPS,
+                            &UMF_OS_MEMORY_PROVIDER_PARAMS,
+                            &os_memory_provider);
+    ASSERT_NE(os_memory_provider, nullptr);
+
+    umf_memory_pool_handle_t pool;
+    umfPoolCreate(&UMF_TBB_POOL_OPS, os_memory_provider, NULL, &pool);
+    ASSERT_NE(pool, nullptr);
+
+    umf_memory_provider_handle_t prov = NULL;
+    umfPoolGetMemoryProvider(pool, &prov);
+    ASSERT_NE(prov, nullptr);
+
+    // test 1
+
+    size_t s1 = 74659 * KB;
+    size_t s2 = 8206 * KB;
+
+    size_t max_alloc_size = 0;
+
+    // s1
+    for (int j = 0; j < 2; j++) {
+        void *t[6] = {0};
+        for (int i = 0; i < 6; i++) {
+            t[i] = umfPoolMalloc(pool, s1);
+            ASSERT_NE(t[i], nullptr);
+        }
+
+        if (max_alloc_size == 0) {
+        }
+
+        for (int i = 0; i < 6; i++) {
+            umf_result_t res = umfPoolFree(pool, t[i]);
+            ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+        }
+    }
+
+    // s2
+    for (int j = 0; j < 2; j++) {
+        void *t[6] = {0};
+        for (int i = 0; i < 6; i++) {
+            t[i] = umfPoolMalloc(pool, s2);
+            ASSERT_NE(t[i], nullptr);
+        }
+
+        for (int i = 0; i < 6; i++) {
+            umf_result_t res = umfPoolFree(pool, t[i]);
+            ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+        }
+    }
+
+    umfPoolDestroy(pool);
+    umfMemoryProviderDestroy(os_memory_provider);
+}
+
+TEST_F(test, pool_tbb_MallocPool_simple2) {
+    umf_memory_provider_handle_t os_memory_provider;
+    umfMemoryProviderCreate(&UMF_OS_MEMORY_PROVIDER_OPS,
+                            &UMF_OS_MEMORY_PROVIDER_PARAMS,
+                            &os_memory_provider);
+    ASSERT_NE(os_memory_provider, nullptr);
+
+    umf_memory_pool_handle_t pool;
+    umfPoolCreate(&UMF_TBB_POOL_OPS, os_memory_provider, NULL, &pool);
+    ASSERT_NE(pool, nullptr);
+
+    // test
+    double sizes[] = {2, 4, 0.5, 1, 8, 0.25};
+    for (int i = 0; i < 6; i++) {
+        size_t s = (size_t)(sizes[i] * MB);
+        void *t[8] = {0};
+        for (int j = 0; j < 8; j++) {
+            t[j] = umfPoolMalloc(pool, s);
+            ASSERT_NE(t[j], nullptr);
+        }
+
+        for (int j = 0; j < 8; j++) {
+            umf_result_t res = umfPoolFree(pool, t[j]);
+            ASSERT_EQ(res, UMF_RESULT_SUCCESS);
+        }
+    }
+
+    umfPoolDestroy(pool);
+    umfMemoryProviderDestroy(os_memory_provider);
+}
+
+TEST_F(test, pool_tbb_MallocPool_random) {
+    umf_memory_provider_handle_t os_memory_provider;
+    umfMemoryProviderCreate(&UMF_OS_MEMORY_PROVIDER_OPS,
+                            &UMF_OS_MEMORY_PROVIDER_PARAMS,
+                            &os_memory_provider);
+    ASSERT_NE(os_memory_provider, nullptr);
+
+    umf_memory_pool_handle_t pool;
+    umfPoolCreate(&UMF_TBB_POOL_OPS, os_memory_provider, NULL, &pool);
+    ASSERT_NE(pool, nullptr);
+
+    // set constant seed so each test run will have the same scenario
+    uint32_t seed = 1234;
+    std::mt19937 mt(seed);
+
+    // different sizes to alloc
+    std::vector<size_t> sizes = {
+        15,       49,       588,       1025,     2 * KB,  5 * KB,
+        160 * KB, 511 * KB, 1000 * KB, MB,       3 * MB,  7 * MB,
+        19 * MB,  26 * MB,  69 * MB,   109 * MB, 111 * MB};
+    std::uniform_int_distribution<int> sizes_dist(0, (int)(sizes.size() - 1));
+
+    // each alloc would be done few times
+    std::vector<size_t> counts = {1, 3, 4, 8, 9, 11};
+    std::uniform_int_distribution<int> counts_dist(0, (int)(counts.size() - 1));
+
+    // action to take will be random
+    // alloc = <0, .5), free = <.5, 1)
+    std::uniform_real_distribution<float> actions_dist(0, 1);
+
+    std::set<void *> allocs;
+    for (size_t i = 0; i < 100; i++) {
+        size_t size = sizes[sizes_dist(mt)];
+        size_t count = counts[counts_dist(mt)];
+        float action = actions_dist(mt);
+
+        std::cout << "size: " << size << " count: " << count
+                  << " action: " << ((action < 0.5) ? "alloc" : "free")
+                  << std::endl;
+
+        if (action < 0.5) {
+            // alloc
+            for (size_t j = 0; j < count; j++) {
+                void *ptr = umfPoolMalloc(pool, size);
+                ASSERT_NE(ptr, nullptr);
+
+                allocs.insert(ptr);
+            }
+        } else {
+            // free random allocs
+            for (size_t j = 0; j < count; j++) {
+                if (allocs.size() == 0) {
+                    continue;
+                }
+
+                std::uniform_int_distribution<int> free_dist(
+                    0, (int)(allocs.size() - 1));
+                size_t free_id = free_dist(mt);
+                auto it = allocs.begin();
+                std::advance(it, free_id);
+                void *ptr = (*it);
+                ASSERT_NE(ptr, nullptr);
+
+                umf_result_t ret = umfPoolFree(pool, ptr);
+                ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+                allocs.erase(ptr);
+            }
+        }
+    }
+
+    std::cout << "cleanup" << std::endl;
+
+    while (allocs.size()) {
+        umf_result_t ret = umfPoolFree(pool, *allocs.begin());
+        ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+        allocs.erase(allocs.begin());
+    }
+
+    umfPoolDestroy(pool);
+    umfMemoryProviderDestroy(os_memory_provider);
+}