use the provider's GetCachelineSize() to calculate bucket sizes in DP

Author: bratpiorka

include/umf/pools/pool_disjoint.h

@@ -1,6 +1,6 @@
 /*
  *
- * Copyright (C) 2023-2025 Intel Corporation
+ * Copyright (C) 2023-2026 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -82,7 +82,7 @@ umfDisjointPoolParamsSetCapacity(umf_disjoint_pool_params_handle_t hParams,
                                  size_t maxCapacity);
 
 /// @brief Set minimum bucket allocation size.
-/// @details Default value for minimum bucket size is 8.
+/// @details If not set, the pool uses the provider cache line size.
 /// @param hParams handle to the parameters of the disjoint pool.
 /// @param minBucketSize minimum bucket size. Must be power of 2.
 /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.

src/pool/pool_disjoint.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022-2025 Intel Corporation
+ * Copyright (C) 2022-2026 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -981,24 +981,22 @@ static bool bucket_can_pool(bucket_t *bucket) {
 static size_t size_to_idx(disjoint_pool_t *pool, size_t size) {
     assert(size <= CutOff && "Unexpected size");
     assert(size > 0 && "Unexpected size");
+    assert(pool->buckets_num > 0 && "Unexpected buckets count");
 
-    size_t min_bucket_size = (size_t)1 << pool->min_bucket_size_exp;
-    if (size < min_bucket_size) {
-        return 0;
-    }
+    size_t left = 0;
+    size_t right = pool->buckets_num - 1;
 
-    // get the position of the leftmost set bit
-    size_t position = utils_msb64(size);
+    while (left < right) {
+        size_t mid = left + (right - left) / 2;
 
-    bool is_power_of_2 = IS_POWER_OF_2(size);
-    bool larger_than_halfway_between_powers_of_2 =
-        !is_power_of_2 &&
-        (bool)((size - 1) & ((uint64_t)(1) << (position - 1)));
-    size_t index = (position - pool->min_bucket_size_exp) * 2 +
-                   (int)(!is_power_of_2) +
-                   (int)larger_than_halfway_between_powers_of_2;
+        if (pool->buckets[mid]->size < size) {
+            left = mid + 1;
+        } else {
+            right = mid;
+        }
+    }
 
-    return index;
+    return left;
 }
 
 static umf_disjoint_pool_shared_limits_t *
@@ -1120,6 +1118,7 @@ static const umf_disjoint_pool_params_t default_params = {
     .max_poolable_size = 2 * 1024 * 1024, // 2MB default
     .capacity = 4,                        // default
     .min_bucket_size = 8,                 // default
+    .min_bucket_size_set = false,
     .cur_pool_size = 0,
     .pool_trace = 0,
     .shared_limits = NULL,
@@ -1161,6 +1160,10 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
 umf_result_t disjoint_pool_post_initialize(void *ppPool) {
     disjoint_pool_t *disjoint_pool = (disjoint_pool_t *)ppPool;
 
+    const size_t max_bucket_num = (size_t)utils_msb64(CutOff) * 2 + 1;
+    size_t bucket_sizes[max_bucket_num];
+    size_t unique_bucket_sizes[max_bucket_num];
+
     disjoint_pool->post_initialized = true;
 
     if (disjoint_pool->params_overridden) {
@@ -1191,53 +1194,67 @@ umf_result_t disjoint_pool_post_initialize(void *ppPool) {
         goto err_free_disjoint_pool;
     }
 
+    if (!disjoint_pool->params.min_bucket_size_set) {
+        size_t cache_line_size;
+        umf_result_t ret = umfMemoryProviderGetCacheLineSize(
+            disjoint_pool->provider, &cache_line_size);
+        if (ret == UMF_RESULT_SUCCESS) {
+            disjoint_pool->params.min_bucket_size = cache_line_size;
+        }
+    }
+
     // Generate buckets sized such as: 64, 96, 128, 192, ..., CutOff.
     // Powers of 2 and the value halfway between the powers of 2.
-    size_t Size1 = disjoint_pool->params.min_bucket_size;
+    size_t size1 = disjoint_pool->params.min_bucket_size;
 
     // min_bucket_size cannot be larger than CutOff.
-    Size1 = utils_min(Size1, CutOff);
+    size1 = utils_min(size1, CutOff);
+
+    // Buckets sized smaller than the bucket default size - 8 aren't needed.
+    size1 = utils_max(size1, UMF_DISJOINT_POOL_MIN_BUCKET_DEFAULT_SIZE);
 
-    // Buckets sized smaller than the bucket default size- 8 aren't needed.
-    Size1 = utils_max(Size1, UMF_DISJOINT_POOL_MIN_BUCKET_DEFAULT_SIZE);
+    // Calculate the size of the "halfway" bucket
+    size_t size2 = size1 + size1 / 2;
 
     // Calculate the exponent for min_bucket_size used for finding buckets.
-    disjoint_pool->min_bucket_size_exp = (size_t)utils_msb64(Size1);
+    disjoint_pool->min_bucket_size_exp = (size_t)utils_msb64(size1);
     if (umfDisjointPoolSharedLimitsCreate(
             SIZE_MAX, &disjoint_pool->default_shared_limits) !=
         UMF_RESULT_SUCCESS) {
         goto err_free_known_slabs;
     }
 
-    // count number of buckets, start from 1
-    disjoint_pool->buckets_num = 1;
-    size_t Size2 = Size1 + Size1 / 2;
-    size_t ts2 = Size2, ts1 = Size1;
-    while (Size2 < CutOff) {
-        disjoint_pool->buckets_num += 2;
-        Size2 *= 2;
+    // Calclate sizes of buckets and store them in the array.
+    size_t i = 0;
+    for (; size2 < CutOff; size1 *= 2, size2 *= 2, i += 2) {
+        bucket_sizes[i] =
+            ALIGN_UP_SAFE(size1, disjoint_pool->params.min_bucket_size);
+        bucket_sizes[i + 1] =
+            ALIGN_UP_SAFE(size2, disjoint_pool->params.min_bucket_size);
+    }
+    bucket_sizes[i] = CutOff;
+
+    // After rounding up some sizes may be equal, we need to remove duplicates
+    // and count the number of unique sizes.
+    size_t unique_count = 0;
+    for (size_t j = 0; j <= i; j++) {
+        if (j == 0 || bucket_sizes[j] != bucket_sizes[j - 1]) {
+            unique_bucket_sizes[unique_count++] = bucket_sizes[j];
+        }
     }
+    disjoint_pool->buckets_num = unique_count;
 
     disjoint_pool->buckets = umf_ba_global_alloc(
         sizeof(*disjoint_pool->buckets) * disjoint_pool->buckets_num);
     if (disjoint_pool->buckets == NULL) {
         goto err_free_shared_limits;
     }
 
-    size_t i = 0;
-    Size1 = ts1;
-    Size2 = ts2;
-    for (; Size2 < CutOff; Size1 *= 2, Size2 *= 2, i += 2) {
-        disjoint_pool->buckets[i] = create_bucket(
-            Size1, disjoint_pool, disjoint_pool_get_limits(disjoint_pool));
-        disjoint_pool->buckets[i + 1] = create_bucket(
-            Size2, disjoint_pool, disjoint_pool_get_limits(disjoint_pool));
-    }
-    disjoint_pool->buckets[i] = create_bucket(
-        CutOff, disjoint_pool, disjoint_pool_get_limits(disjoint_pool));
-
-    // check if all buckets were created successfully
+    // Create buckets for each unique size.
     for (i = 0; i < disjoint_pool->buckets_num; i++) {
+        disjoint_pool->buckets[i] =
+            create_bucket(unique_bucket_sizes[i], disjoint_pool,
+                          disjoint_pool_get_limits(disjoint_pool));
         if (disjoint_pool->buckets[i] == NULL) {
             goto err_free_buckets;
         }
@@ -1733,6 +1750,7 @@ umfDisjointPoolParamsSetMinBucketSize(umf_disjoint_pool_params_handle_t hParams,
     }
 
     hParams->min_bucket_size = minBucketSize;
+    hParams->min_bucket_size_set = true;
     return UMF_RESULT_SUCCESS;
 }
 

src/pool/pool_disjoint_internal.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2025 Intel Corporation
+ * Copyright (C) 2025-2026 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -120,7 +120,8 @@ typedef struct umf_disjoint_pool_params_t {
 
     // Holds the minimum bucket size valid for allocation of a memory type.
     // This value must be a power of 2.
-    size_t min_bucket_size; // Default: 8
+    size_t min_bucket_size; // Default: provider cache line size
+    bool min_bucket_size_set;
 
     // Holds size of the pool managed by the allocator.
     size_t cur_pool_size; // Default: 0

test/pools/disjoint_pool.cpp

@@ -1,4 +1,4 @@
-// Copyright (C) 2023-2025 Intel Corporation
+// Copyright (C) 2023-2026 Intel Corporation
 // Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
@@ -44,6 +44,11 @@ TEST_F(test, internals) {
             *pageSize = 1024;
             return UMF_RESULT_SUCCESS;
         }
+
+        umf_result_t get_cache_line_size(size_t *size) noexcept {
+            *size = 128;
+            return UMF_RESULT_SUCCESS;
+        }
     };
     umf_memory_provider_ops_t provider_ops =
         umf_test::providerMakeCOps<memory_provider, void>();
@@ -78,22 +83,15 @@ TEST_F(test, internals) {
     EXPECT_EQ(pool->provider_min_page_size, (size_t)1024);
 
     // check buckets sizes
-    size_t expected_size = DEFAULT_DISJOINT_MIN_BUCKET_SIZE;
-    EXPECT_EQ(pool->buckets[0]->size, expected_size);
+    EXPECT_EQ(pool->buckets[0]->size, params->min_bucket_size);
     EXPECT_EQ(pool->buckets[pool->buckets_num - 1]->size,
               (size_t)1 << 31); // 2GB
+    size_t prev_size = params->min_bucket_size;
     for (size_t i = 0; i < pool->buckets_num; i++) {
         bucket_t *bucket = pool->buckets[i];
         EXPECT_NE(bucket, nullptr);
-        EXPECT_EQ(bucket->size, expected_size);
-
-        // assuming DEFAULT_DISJOINT_MIN_BUCKET_SIZE = 64, expected bucket
-        // sizes are: 64, 96, 128, 192, 256, ..., 2GB
-        if (i % 2 == 0) {
-            expected_size += expected_size / 2;
-        } else {
-            expected_size = DEFAULT_DISJOINT_MIN_BUCKET_SIZE << ((i + 1) / 2);
-        }
+        EXPECT_GE(bucket->size, prev_size);
+        prev_size = bucket->size;
     }
 
     // test small allocations

test/pools/disjoint_pool_ctl.cpp

@@ -1,4 +1,4 @@
-// Copyright (C) 2025 Intel Corporation
+// Copyright (C) 2025-2026 Intel Corporation
 // Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exceptiongi
 
@@ -614,19 +614,28 @@ TEST_F(test, disjointCtlBucketStats) {
                              sizeof(count), poolWrapper.get()));
     EXPECT_GE(count, 0ull);
 
-    auto expected_bucket_size = [](size_t i) -> size_t {
-        // Even indexes: 8 << (i/2)  => 8,16,32,64,...
-        // Odd  indexes: 12 << (i/2) => 12,24,48,96,...
-        return (i % 2 == 0) ? (size_t(8) << (i / 2)) : (size_t(12) << (i / 2));
-    };
+    size_t min_bucket_size = 0;
+    ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.{}.params.min_bucket_size",
+                             &min_bucket_size, sizeof(min_bucket_size),
+                             poolWrapper.get()));
 
+    // Check bucket sizes
+    size_t prev_size = 0;
     for (size_t i = 0; i < count; i++) {
         ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.{}.buckets.{}.size", &arg,
                                  sizeof(arg), poolWrapper.get(), i));
-        EXPECT_EQ(arg, expected_bucket_size(i)) << "Failed for bucket: " << i;
+
+        // Each bucket size should be:
+        // * greater than the previous one for >= minimum bucket size
+        // * size aligned to the minimum bucket size
+        EXPECT_GE(arg, prev_size);
+        EXPECT_EQ(arg % min_bucket_size, 0ull);
+
         if (arg >= alloc_size && used_bucket == SIZE_MAX) {
             used_bucket = i; // Find the bucket that matches alloc_size
         }
+
+        prev_size = arg;
     }
 
     std::unordered_map<std::string, size_t> stats = {