optimize bf allocator

Author: zhaoguochun1995

dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBFCachingAllocator.cpp

@@ -1,5 +1,6 @@
 // Copyright (c) 2023, DeepLink.
 
+#include <cstddef>
 #include <functional>
 #include <memory>
 #include <stack>
@@ -14,10 +15,12 @@
 
 namespace dipu {
 
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-const size_t kMaxExtendSize = get_env_or_default("DIPU_MAX_EXTEND_SIZE", 1024)
-                              << 20U;
-
+inline size_t round_up_to_alignment(size_t nbytes, size_t alignment_size) {
+  if (nbytes <= 0) {
+    return alignment_size;
+  }
+  return ((nbytes - 1) | (alignment_size - 1)) + 1;
+}
 class BFCachingAllocatorImpl {
  public:
   using allocate_fn_t = std::function<void*(size_t)>;
@@ -32,9 +35,11 @@ class BFCachingAllocatorImpl {
   static constexpr int kNumSubBins = 4;
   static constexpr int kLogNumSubBins = 2;
   // Allocation parameters
-  static constexpr size_t kMinAllocationSize = 512;
-  static constexpr size_t kMaxInternalFragmentation = 8U << 20U;  // 8MB
-  static constexpr size_t kMinExtendSize = 8U << 20U;             // 8MB
+  static constexpr int kMinAllocationSize = 512;
+  static constexpr int kSmallBlockSize = 2 << 20;
+  static constexpr int kMiddleBlockSize = 20 << 20;
+  static constexpr int kLargeBlockSize = 200 << 20;
+  static constexpr int kLargeAlignSize = 1024 << 20;
 
   size_t cachedBytes = 0;
   size_t allocatedBytes = 0;
@@ -67,8 +72,6 @@ class BFCachingAllocatorImpl {
     __uint128_t bits = 0;
     // Virtual chunks which are the heads of the bins
     std::array<int, static_cast<size_t>(kNumBigBins* kNumSubBins)> binHeads_{};
-    // The extending size next time
-    size_t currExtendSize_ = kMinExtendSize;
 
     explicit StreamSet(size_t id) : id(id) {}
 
@@ -140,7 +143,10 @@ class BFCachingAllocatorImpl {
   mutable mutex_t mut_;
 
   static size_t roundBytes(size_t nbytes) {
-    return ((nbytes - 1) | (kMinAllocationSize - 1)) + 1;
+    if (nbytes < kLargeBlockSize) {
+      return round_up_to_alignment(nbytes, kMinAllocationSize);
+    }
+    return round_up_to_alignment(nbytes, kSmallBlockSize);
   }
 
   int newChunk(void* ptr, size_t size, size_t stream) {
@@ -293,20 +299,21 @@ class BFCachingAllocatorImpl {
 
   int extend(size_t nbytes, StreamSetHandle& set) {
     emptyCacheWithoutLock();
-    auto& extSize = set->currExtendSize_;
     bool increased = false;
-    while (extSize < nbytes && extSize < kMaxExtendSize) {
-      extSize *= 2;
-      increased = true;
+    size_t allocateSize = nbytes;
+    if (nbytes < kSmallBlockSize) {
+      allocateSize = kSmallBlockSize;
+    } else if (nbytes < kMiddleBlockSize) {
+      allocateSize = kMiddleBlockSize;
+    } else if (nbytes < kLargeBlockSize) {
+      allocateSize = round_up_to_alignment(nbytes, kMiddleBlockSize);
+    } else {
+      allocateSize = round_up_to_alignment(nbytes, kLargeAlignSize);
     }
 
-    size_t currBytes = std::max(nbytes, extSize);
+    size_t currBytes = std::max(nbytes, allocateSize);
     void* ptr = allocateOnDevice(currBytes);
-    if (ptr) {
-      if (!increased && extSize < kMaxExtendSize) {
-        extSize *= 2;
-      }
-    } else {
+    if (!ptr) {
       if (currBytes > nbytes) {
         currBytes = nbytes;
         ptr = allocateOnDevice(currBytes);
@@ -371,8 +378,17 @@ class BFCachingAllocatorImpl {
     }
 
     if (id) {
-      if (chunks_[id].size >= nbytes * 2 ||
-          chunks_[id].size >= nbytes + kMaxInternalFragmentation) {
+      int internlalMaxFragnmentSize = 0;
+      const int chunk_size = static_cast<int>(chunks_[id].size);
+      if (chunk_size < kSmallBlockSize) {
+        internlalMaxFragnmentSize = kMinAllocationSize;
+      } else if (chunk_size < kLargeAlignSize) {
+        internlalMaxFragnmentSize = kSmallBlockSize;
+      } else {
+        internlalMaxFragnmentSize = kLargeAlignSize;
+      }
+      if ((chunk_size >= (nbytes << 1)) ||
+          (chunk_size > (nbytes + internlalMaxFragnmentSize))) {
         id = split(id, nbytes);
       }
       chunks_[id].allocated = true;

dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPURawAllocator.cpp

@@ -50,7 +50,8 @@ c10::DataPtr DIPURawDeviceAllocator::allocate(
   if (nbytes > 0) {
     devproxy::mallocDevice(&data, nbytes);
     DIPU_DEBUG_ALLOCATOR(1, "devproxy::mallocDevice: malloc "
-                                << nbytes << " nbytes, ptr:" << data);
+                                << nbytes << "(" << (nbytes >> 20) << "MB)"
+                                << " nbytes, ptr:" << data);
   }
   return {data, data, &DIPURawDeviceAllocatorDeleter,
           c10::Device(dipu::DIPU_DEVICE_TYPE, device_index)};

dipu/torch_dipu/csrc_dipu/vendor/ascend/AscendCustomAllocatorStrategy.cpp

@@ -1,10 +1,15 @@
+#include <cstdint>
+
 #include "csrc_dipu/runtime/core/allocator/DIPUCachingAllocator.h"
+#include "csrc_dipu/utils/env.hpp"
 
 const static int ascend_init = []() {
   // – 内存大小向上对齐成32整数倍+32字节(m=ALIGN_UP[len,32]+32字节);
   // – 内存起始地址需满足64字节对齐(ALIGN_UP[m,64])。
   //  nbytes = align_64(1 * nbytes + 32);
-  static dipu::MemoryAlignmentStrategy memoryAlignStrategy(64, 1, 32);
+
+  int64_t align = dipu::get_env_or_default("ASCEND_ALLOCATE_ALIGN", 512);
+  static dipu::MemoryAlignmentStrategy memoryAlignStrategy(align, 1, 32);
   dipu::setMemoryAlignmentStrategy(&memoryAlignStrategy);
   return 0;
 }();