[umf] optimize bucket selection for allocation sizes

source/common/umf_pools/CMakeLists.txt

@@ -6,6 +6,7 @@
 add_ur_library(disjoint_pool STATIC
     disjoint_pool.cpp
     disjoint_pool_config_parser.cpp
+    ../unified_malloc_framework/src/utils/utils_math.cpp
 )
 
 add_library(${PROJECT_NAME}::disjoint_pool ALIAS disjoint_pool)

source/common/umf_pools/disjoint_pool.cpp

@@ -11,6 +11,7 @@
 #include <bitset>
 #include <cassert>
 #include <cctype>
+#include <cmath>
 #include <iomanip>
 #include <limits>
 #include <list>
@@ -25,6 +26,7 @@
 // TODO: replace with logger?
 #include <iostream>
 
+#include "../unified_malloc_framework/src/utils/utils_math.h"
 #include "disjoint_pool.hpp"
 
 namespace usm {
@@ -283,6 +285,9 @@ class DisjointPool::AllocImpl {
     // Configuration for this instance
     DisjointPoolConfig params;
 
+    // Used in alghoritm for finding buckets
+    std::size_t MinBucketSizeExp;
+
     // Coarse-grain allocation min alignment
     size_t ProviderMinPageSize;
 
@@ -293,8 +298,12 @@ class DisjointPool::AllocImpl {
         // Generate buckets sized such as: 64, 96, 128, 192, ..., CutOff.
         // Powers of 2 and the value halfway between the powers of 2.
         auto Size1 = params.MinBucketSize;
+        // MinBucketSize cannot be larger than CutOff.
+        Size1 = std::min(Size1, CutOff);
         // Buckets sized smaller than the bucket default size- 8 aren't needed.
         Size1 = std::max(Size1, MIN_BUCKET_DEFAULT_SIZE);
+        // Calculate the exponent for MinBucketSize used for finding buckets.
+        MinBucketSizeExp = (size_t)log2(Size1);
         auto Size2 = Size1 + Size1 / 2;
         for (; Size2 < CutOff; Size1 *= 2, Size2 *= 2) {
             Buckets.push_back(std::make_unique<Bucket>(Size1, *this));
@@ -331,6 +340,7 @@ class DisjointPool::AllocImpl {
 
   private:
     Bucket &findBucket(size_t Size);
+    std::size_t sizeToIdx(size_t Size);
 };
 
 static void *memoryProviderAlloc(umf_memory_provider_handle_t hProvider,
@@ -818,16 +828,39 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
     return nullptr;
 }
 
-Bucket &DisjointPool::AllocImpl::findBucket(size_t Size) {
+std::size_t DisjointPool::AllocImpl::sizeToIdx(size_t Size) {
     assert(Size <= CutOff && "Unexpected size");
+    assert(Size > 0 && "Unexpected size");
+
+    size_t MinBucketSize = (size_t)1 << MinBucketSizeExp;
+    if (Size < MinBucketSize) {
+        return 0;
+    }
+
+    // Get the position of the leftmost set bit.
+    size_t position = getLeftmostSetBitPos(Size);
 
-    auto It = std::find_if(
-        Buckets.begin(), Buckets.end(),
-        [Size](const auto &BucketPtr) { return BucketPtr->getSize() >= Size; });
+    size_t lower_bound = (size_t)1 << position;
+    size_t diff = (position - MinBucketSizeExp) << 1;
 
-    assert((It != Buckets.end()) && "Bucket should always exist");
+    if (Size == lower_bound) {
+        return diff;
+    } else if (Size <= (lower_bound | (lower_bound >> 1))) {
+        return diff + 1;
+    } else {
+        return diff + 2;
+    }
+}
+
+Bucket &DisjointPool::AllocImpl::findBucket(size_t Size) {
+    auto calculatedIdx = sizeToIdx(Size);
+    assert(calculatedIdx >= 0);
+    assert((*(Buckets[calculatedIdx])).getSize() >= Size);
+    if (calculatedIdx > 0) {
+        assert((*(Buckets[calculatedIdx - 1])).getSize() < Size);
+    }
 
-    return *(*It);
+    return *(Buckets[calculatedIdx]);
 }
 
 void DisjointPool::AllocImpl::deallocate(void *Ptr, bool &ToPool) {

source/common/unified_malloc_framework/CMakeLists.txt

@@ -12,9 +12,9 @@ set(UMF_SOURCES
 )
 
 if(MSVC)
-    set(UMF_SOURCES ${UMF_SOURCES} src/utils/utils_windows.cpp src/memory_tracker_windows.cpp)
+    set(UMF_SOURCES ${UMF_SOURCES} src/utils/utils_windows_concurrency.cpp src/memory_tracker_windows.cpp)
 else()
-    set(UMF_SOURCES ${UMF_SOURCES} src/utils/utils_posix.c)
+    set(UMF_SOURCES ${UMF_SOURCES} src/utils/utils_posix_concurrency.c)
 endif()
 
 if(UMF_BUILD_SHARED_LIBRARY)

source/common/unified_malloc_framework/src/critnib/critnib.c

@@ -58,7 +58,7 @@
 #include <stdbool.h>
 #include <stddef.h>
 
-#include "../utils/utils.h"
+#include "../utils/utils_concurrency.h"
 #include "critnib.h"
 
 /*

source/common/unified_malloc_framework/src/utils/utils_math.cpp

@@ -0,0 +1,25 @@
+/*
+ *
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+ * See LICENSE.TXT
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include "utils_math.h"
+
+// Retrieves the position of the leftmost set bit.
+// The position of the bit is counted from 0
+// e.g. for 01000011110 the position equals 9.
+size_t getLeftmostSetBitPos(size_t num) {
+    // From C++20 countl_zero could be used for that.
+    size_t position = 0;
+    while (num > 0) {
+        num >>= 1;
+        position++;
+    }
+    position--;
+    return position;
+}

source/common/unified_malloc_framework/src/utils/utils_math.h

@@ -0,0 +1,13 @@
+/*
+ *
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+ * See LICENSE.TXT
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include <cstddef>
+
+size_t getLeftmostSetBitPos(size_t num);

source/common/unified_malloc_framework/src/utils/utils_posix.c → source/common/unified_malloc_framework/src/utils/utils_posix_concurrency.c

@@ -11,7 +11,7 @@
 #include <pthread.h>
 #include <stdlib.h>
 
-#include "utils.h"
+#include "utils_concurrency.h"
 
 struct os_mutex_t *util_mutex_create(void) {
     pthread_mutex_t *mutex = (pthread_mutex_t *)malloc(sizeof(pthread_mutex_t));

source/common/unified_malloc_framework/src/utils/utils_windows.cpp → source/common/unified_malloc_framework/src/utils/utils_windows_concurrency.cpp

@@ -10,7 +10,7 @@
 
 #include <mutex>
 
-#include "utils.h"
+#include "utils_concurrency.h"
 
 struct os_mutex_t *util_mutex_create(void) {
     return reinterpret_cast<struct os_mutex_t *>(new std::mutex);

test/unified_malloc_framework/memoryPool.hpp

@@ -36,6 +36,8 @@ struct umfPoolTest : umf_test::test,
 
     umf::pool_unique_handle_t pool;
     static constexpr int NTHREADS = 5;
+    static constexpr std::array<int, 7> nonAlignedAllocSizes = {5,  7,   23, 55,
+                                                                80, 119, 247};
 };
 
 struct umfMultiPoolTest : umfPoolTest {
@@ -86,6 +88,15 @@ TEST_P(umfPoolTest, allocFree) {
     umfPoolFree(pool.get(), ptr);
 }
 
+TEST_P(umfPoolTest, allocFreeNonAlignedSizes) {
+    for (const auto &allocSize : nonAlignedAllocSizes) {
+        auto *ptr = umfPoolMalloc(pool.get(), allocSize);
+        ASSERT_NE(ptr, nullptr);
+        std::memset(ptr, 0, allocSize);
+        umfPoolFree(pool.get(), ptr);
+    }
+}
+
 TEST_P(umfPoolTest, reallocFree) {
     if (!umf_test::isReallocSupported(pool.get())) {
         GTEST_SKIP();