[umf] optimize bucket selection for allocation sizes

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.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);
+ 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,
@@ -404,15 +414,15 @@ size_t Slab::FindFirstAvailableChunkIdx() const {
  return It - Chunks.begin();
  }
 
- return std::numeric_limits<size_t>::max();
+ return (std::numeric_limits<size_t>::max)();
 }
 
 void *Slab::getChunk() {
  // assert(NumAllocated != Chunks.size());
 
  const size_t ChunkIdx = FindFirstAvailableChunkIdx();
  // Free chunk must exist, otherwise we would have allocated another slab
- assert(ChunkIdx != (std::numeric_limits<size_t>::max()));
+ assert(ChunkIdx != ((std::numeric_limits<size_t>::max)()));
 
  void *const FreeChunk =
  (static_cast<uint8_t *>(getPtr())) + ChunkIdx * getChunkSize();
@@ -680,7 +690,7 @@ umf_memory_provider_handle_t Bucket::getMemHandle() {
 
 size_t Bucket::SlabMinSize() { return OwnAllocCtx.getParams().SlabMinSize; }
 
-size_t Bucket::SlabAllocSize() { return std::max(getSize(), SlabMinSize()); }
+size_t Bucket::SlabAllocSize() { return (std::max)(getSize(), SlabMinSize()); }
 
 size_t Bucket::Capacity() {
  // For buckets used in chunked mode, just one slab in pool is sufficient.
@@ -712,9 +722,9 @@ void Bucket::updateStats(int InUse, int InPool) {
  return;
  }
  currSlabsInUse += InUse;
- maxSlabsInUse = std::max(currSlabsInUse, maxSlabsInUse);
+ maxSlabsInUse = (std::max)(currSlabsInUse, maxSlabsInUse);
  currSlabsInPool += InPool;
- maxSlabsInPool = std::max(currSlabsInPool, maxSlabsInPool);
+ maxSlabsInPool = (std::max)(currSlabsInPool, maxSlabsInPool);
  // Increment or decrement current pool sizes based on whether
  // slab was added to or removed from pool.
  OwnAllocCtx.getParams().CurPoolSize += InPool * SlabAllocSize();
@@ -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) {
@@ -883,9 +916,9 @@ void DisjointPool::AllocImpl::printStats(bool &TitlePrinted,
  HighPeakSlabsInUse = 0;
  for (auto &B : Buckets) {
  (*B).printStats(TitlePrinted, MTName);
- HighPeakSlabsInUse = std::max((*B).maxSlabsInUse, HighPeakSlabsInUse);
+ HighPeakSlabsInUse = (std::max)((*B).maxSlabsInUse, HighPeakSlabsInUse);
  if ((*B).allocCount) {
- HighBucketSize = std::max((*B).SlabAllocSize(), HighBucketSize);
+ HighBucketSize = (std::max)((*B).SlabAllocSize(), HighBucketSize);
  }
  }
 }

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

@@ -14,8 +14,10 @@
 #if defined(_WIN32)
 #include <windows.h>
 #else
+#ifndef __cplusplus
 #include <stdatomic.h>
 #endif
+#endif
 
 #ifdef __cplusplus
 extern "C" {
@@ -72,6 +74,20 @@ static inline void *Zalloc(size_t s) {
  return m;
 }
 
+// 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.
+static inline 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;
+}
+
 #define NOFUNCTION \
  do { \
  } while (0)

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();