From 981ee4b9d99658b3fce4422be73e7de31097f2bc Mon Sep 17 00:00:00 2001
From: Damian Duy <damian.duy@intel.com>
Date: Tue, 23 May 2023 12:25:15 +0200
Subject: [PATCH] [UMA] add more tests for allocations

---
 test/unified_memory_allocation/memoryPool.hpp | 94 +++++++++++++++++++
 1 file changed, 94 insertions(+)

diff --git a/test/unified_memory_allocation/memoryPool.hpp b/test/unified_memory_allocation/memoryPool.hpp
index 068e83dde6..0777436178 100644
--- a/test/unified_memory_allocation/memoryPool.hpp
+++ b/test/unified_memory_allocation/memoryPool.hpp
@@ -6,6 +6,7 @@
 #include <cstring>
 #include <functional>
 #include <random>
+#include <thread>
 
 #ifndef UMA_TEST_MEMORY_POOL_OPS_HPP
 #define UMA_TEST_MEMORY_POOL_OPS_HPP
@@ -27,6 +28,7 @@ struct umaPoolTest : uma_test::test,
         return pool;
     }
 
+    static constexpr int NTHREADS = 5;
     uma::pool_unique_handle_t pool;
 };
 
@@ -73,6 +75,98 @@ TEST_P(umaPoolTest, pow2AlignedAlloc) {
     }
 }
 
+TEST_P(umaPoolTest, allocZeroSize) {
+    static constexpr size_t allocSize = 0;
+    auto *ptr = umaPoolMalloc(pool.get(), allocSize);
+    ASSERT_NE(ptr, nullptr);
+    umaPoolFree(pool.get(), ptr);
+}
+
+TEST_P(umaPoolTest, freeNullptr) {
+    void *ptr = nullptr;
+    umaPoolFree(pool.get(), ptr);
+}
+
+TEST_P(umaPoolTest, allocOutOfMem) {
+    // test whether memory is kept in a pool accordingly to MaxSize
+    static constexpr size_t allocSize = 16;
+    // MaxSize equals 16 * 1024 * 1024;
+    static constexpr size_t maxAllocs = 1024 * 1024;
+
+    // allocate until oom
+    void *ptr = nullptr;
+    std::vector<void*> allocations;
+
+    for (size_t i = 0; i <= maxAllocs; ++i) {
+        allocations.emplace_back(umaPoolMalloc(pool.get(), allocSize));
+        ASSERT_NE(allocations.back(), nullptr);
+    }
+
+    // check if the last allocation belongs to the pool
+    // it returns different values each time
+    // ASSERT_EQ(umaPoolByPtr(allocations.back()), pool.get());
+
+    // free some memory
+    umaPoolFree(pool.get(), allocations.back());
+
+    allocations.pop_back();
+
+    ptr = umaPoolMalloc(pool.get(), allocSize);
+    ASSERT_NE(ptr, nullptr);
+
+    umaPoolFree(pool.get(), ptr);
+
+    for (auto allocation : allocations) {
+        umaPoolFree(pool.get(), allocation);
+    }
+}
+
+TEST_P(umaPoolTest, multiThreadedMallocFree) {
+    auto poolMalloc = [](size_t allocSize, uma_memory_pool_handle_t pool) {
+        std::vector<void*> allocations;
+        for (size_t i = 0; i <= 10; ++i) {
+            allocations.emplace_back(umaPoolMalloc(pool, allocSize));
+            ASSERT_NE(allocations.back(), nullptr);
+        }
+
+        for (auto allocation : allocations) {
+        umaPoolFree(pool, allocation);
+        }
+    };
+
+    std::vector<std::thread> threads;
+    for (int i = 0; i < NTHREADS; i++) {
+        threads.push_back(std::thread(poolMalloc, 64, pool.get()));
+    }
+
+    for (auto &thread : threads) {
+        thread.join();
+    }
+}
+
+TEST_P(umaPoolTest, multiThreadedMallocFreeRandomSizes) {
+    auto poolMalloc = [](size_t allocSize, uma_memory_pool_handle_t pool) {
+        std::vector<void*> allocations;
+        for (size_t i = 0; i <= 10; ++i) {
+            allocations.emplace_back(umaPoolMalloc(pool, allocSize));
+            ASSERT_NE(allocations.back(), nullptr);
+        }
+
+        for (auto allocation : allocations) {
+        umaPoolFree(pool, allocation);
+        }
+    };
+
+    std::vector<std::thread> threads;
+    for (int i = 0; i < NTHREADS; i++) {
+        threads.push_back(std::thread(poolMalloc, rand()%64 + 1, pool.get()));
+    }
+
+    for (auto &thread : threads) {
+        thread.join();
+    }
+}
+
 // TODO: add similar tests for realloc/aligned_alloc, etc.
 // TODO: add multithreaded tests
 TEST_P(umaMultiPoolTest, memoryTracking) {