[OpenMP] Introduced a bump-like allocator into the target memory management

[tianshilei1992]

Summary

In this patch, a bump-like allocator is introduced to improve the efficiency of
target memory allocation and deallocation. The reason it is called a bump-like
allocator is it supports deallocation. Everytime a memory allocation request is
from the memory manager, it first checks whether there is a fitted slab. If not,
create a new slab. Bump the pointer in the slab and return to the memory manager,
like a regular bump allocator. We set a reference count for each slab to track
the number of memory requests from the slab.

Of course it is possible that we may run out of target memory. Then the memory
manager will try to release all its buffered memory in its FreeLists. For each
deallocation, the allocator first finds the slab it is from. If it cannot find
a slab, then it must be allocated directly on device so it will be deallocated
directly from the device. Otherwise, just decrease the reference count of the
slab. If the reference count is zero, which means all memory allocated are
returned. So we will call the device routine to deallocate the slab.

It is worth noting that the deallocation might only be called in three cases:

1. Destructor of the memory manager;
2. We are out of memory. The memory manager wants to free all its hold free memory;
3. A large block of memory which is not allocated from a slab. For this case,

we might optimize it such that if the size is not too large, we may take it as
another slab when it is gonna be released by the memory manager.

Tests, comments, and debug output have not been added yet, but will soon.

[JonChesterfield]

The reference counting to free is interesting. I don't think I've seen that on a bump allocator before. It doesn't allow memory reuse within a slab. I wonder whether there is usually some outstanding reference into the slab for most of the execution.

A common allocator pattern is a variant on N allocators, each for a fixed size. That makes alloc/free a stack operation, where you free the whole block if the stack becomes empty. Alloc as cheap as here, incremental free possible. Cost is the unused parts of the stack.

Go much beyond that, into allocating from the next size up etc, and one ends up with a heap. It's possible that's the end point here - do we actually just want malloc?

[grokos]

I just revisited this patch. It seems it's based on a very early implementation of the base memory manager patch (D81054). Can you rebase the patch so that we review it? Thanks!

[tianshilei1992]

In D85274#2311569, @grokos wrote:

I just revisited this patch. It seems it's based on a very early implementation of the base memory manager patch (D81054). Can you rebase the patch so that we review it? Thanks!

Sure, no problem. I'm currently working on another patches, and will rebase it later.

openmp/libomptarget/src/CMakeLists.txt

 ##===----------------------------------------------------------------------===##
 #
 # Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 #
 ##===----------------------------------------------------------------------===##
 #
 # Build offloading library libomptarget.so.
 #
 ##===----------------------------------------------------------------------===##
 
 libomptarget_say("Building offloading runtime library libomptarget.")
 
 set(src_files
+  allocator.cpp
   api.cpp
   device.cpp
   interface.cpp
   memory.cpp
   rtl.cpp
   omptarget.cpp
 )
 

openmp/libomptarget/src/allocator.h

-//===----------- memory.h - Target independent memory manager -------------===//
+//===----------- allocator.h - Target memory allocators -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// Declarations for target independent memory manager.
+// Declarations for target memory allocators.
 //
 //===----------------------------------------------------------------------===//
 
 #pragma once
 
 #include <cstddef>
 #include <memory>
 
 // Forward declaration
 struct DeviceTy;
 
 namespace memory {
 namespace impl {
-class MemoryManagerImplTy;
+class BumpAllocatorImplTy;
 } // namespace impl
 
-class MemoryManagerTy {
-  std::shared_ptr<impl::MemoryManagerImplTy> Impl;
+class BumpAllocatorTy {
+  std::shared_ptr<impl::BumpAllocatorImplTy> Impl;
 
 public:
-  /// Constructor
-  MemoryManagerTy(DeviceTy &D, size_t Threshold = 0);
+  BumpAllocatorTy(DeviceTy &Dev);
 
-  /// Allocate memory of size \p Size from target device. \p HstPtr is used to
-  /// assist the allocation.
   void *allocate(size_t Size, void *HstPtr);
 
-  /// Deallocate memory pointed by \p TgtPtr
-  int free(void *TgtPtr);
+  int32_t deallocate(void *Ptr);
 };
 } // namespace memory

openmp/libomptarget/src/allocator.cpp

+//===----------- allocator.cpp - Target memory allocators -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Functionality for target memory allocators.
+//
+//===----------------------------------------------------------------------===//
+
+#include "allocator.h"
+#include "device.h"
+#include "private.h"
+#include "rtl.h"
+
+#include <atomic>
+#include <map>
+#include <mutex>
+#include <set>
+
+namespace {
+struct SlabTy {
+  void *PtrBegin;
+  void *PtrEnd;
+  void *Cur;
+  const size_t Size;
+  size_t RefCount;
+
+  SlabTy(void *TgtPtr, size_t Size)
+      : PtrBegin(TgtPtr), PtrEnd(reinterpret_cast<char *>(PtrBegin) + Size),
+        Cur(PtrEnd), Size(Size), RefCount(0) {}
+};
+} // namespace
+
+namespace memory {
+/// Slab size. By default it is 8MB.
+size_t SlabSize = 1U << 23;
+
+constexpr const size_t Alignment = 8;
+
+namespace impl {
+class BumpAllocatorImplTy {
+  using SlabPtrTy = std::shared_ptr<SlabTy>;
+  using SlabTableTy = std::map<void *, SlabPtrTy>;
+
+  /// A reference to the device object
+  DeviceTy &Device;
+  /// A table mapping from the target pointer to its slab. The target pointer
+  /// must be the \p PtrBegin of each slab.
+  SlabTableTy PtrToSlabTable;
+  /// The mutex for the map table
+  std::mutex MapTableMtx;
+
+  /// Request memory from target device
+  void *allocateFromDevice(size_t Size, void *HstPtr) const {
+    return Device.RTL->data_alloc(Device.RTLDeviceID, Size, HstPtr);
+  }
+
+  /// Deallocate data from device
+  int deleteFromDevice(void *Ptr) const {
+    return Device.RTL->data_delete(Device.RTLDeviceID, Ptr);
+  }
+
+  /// Create a new slab of size \p Size
+  SlabPtrTy createNewSlab(size_t Size) {
+    void *TgtPtr = allocateFromDevice(Size, nullptr);
+    if (TgtPtr == nullptr)
+      return nullptr;
+    return std::make_shared<SlabTy>(TgtPtr, Size);
+  }
+
+public:
+  /// Constructor
+  BumpAllocatorImplTy(DeviceTy &Dev) : Device(Dev) {}
+
+  /// Destructor
+  ~BumpAllocatorImplTy() {
+    for (SlabTableTy::iterator Itr = PtrToSlabTable.begin();
+         Itr != PtrToSlabTable.end(); ++Itr)
+      deleteFromDevice(Itr->first);
+  }
+
+  /// Allocate memory
+  void *allocate(size_t Size, void *HstPtr) {
+    // If the size is larger than SlabSize, we will allocate it from device
+    // directly.
+    // TODO: We might want something similar to the CustomSizedSlabs in
+    // BumpPtrAllocatorImpl in LLVM. Or, we could use it as another slab when it
+    // is freed.
+    if (Size > SlabSize)
+      return allocateFromDevice(Size, HstPtr);
+
+    // Add padding if needed
+    Size += Size % Alignment;
+
+    // Find an available slab
+    {
+      std::lock_guard<std::mutex> G(MapTableMtx);
+      for (SlabTableTy::iterator Itr = PtrToSlabTable.begin();
+           Itr != PtrToSlabTable.end(); ++Itr) {
+        SlabPtrTy Slab = Itr->second;
+
+        // If the size left in a slab is less than Size, we move to the next one
+        if (reinterpret_cast<char *>(Slab->Cur) -
+                reinterpret_cast<char *>(Slab->PtrBegin) <
+            Size)
+          continue;
+
+        // Now we find one. Allocate it from the slab.
+        Slab->Cur = reinterpret_cast<char *>(Slab->Cur) - Size;
+        Slab->RefCount += 1;
+        return Slab->Cur;
+      }
+    }
+
+    // We cannot find a fitted slab. Create a new one
+    SlabPtrTy Slab = createNewSlab(SlabSize);
+
+    // Failed to create a new slab. Return nullptr directly.
+    if (Slab == nullptr)
+      return nullptr;
+
+    // We could manipulate it w/o any lock because it has not been inserted to
+    // the table yet
+    Slab->Cur = reinterpret_cast<char *>(Slab->Cur) - Size;
+    Slab->RefCount += 1;
+
+    void *TgtPtr = Slab->Cur;
+
+    // Add the slab into map table
+    {
+      std::lock_guard<std::mutex> G(MapTableMtx);
+      PtrToSlabTable[Slab->PtrBegin] = Slab;
+    }
+
+    return TgtPtr;
+  }
+
+  /// Deallocate memory
+  int32_t deallocate(void *Ptr) {
+    // Find which slab Ptr is from
+    {
+      std::lock_guard<std::mutex> G(MapTableMtx);
+      SlabTableTy::iterator Itr = PtrToSlabTable.upper_bound(Ptr);
+      if (Itr != PtrToSlabTable.end() || !PtrToSlabTable.empty()) {
+        assert(Itr != PtrToSlabTable.begin() && "Itr should NOT be begin!");
+
+        Itr = std::prev(Itr);
+        SlabPtrTy Slab = Itr->second;
+
+        // If Ptr is not in the range, it is not allocated from the slab
+        if (Ptr >= Slab->PtrBegin && Ptr < Slab->PtrEnd) {
+          size_t &RC = Slab->RefCount;
+          RC -= 1;
+
+          // This is not the last piece of the slab. Return directly.
+          if (RC != 0)
+            return OFFLOAD_SUCCESS;
+
+          // The last piece of the slab. Remove the slab from the map table.
+          PtrToSlabTable.erase(Itr);
+          // Set Ptr to the beginning of the slab.
+          Ptr = Slab->PtrBegin;
+        }
+      }
+    }
+
+    // There are only two cases that we can reach this point:
+    // 1. Ptr is not allocated from a slab;
+    // 2. Ptr has been "free" in the slab, and it is the last piece in the slab,
+    // so we need to free the slab. In this case, Ptr has already been set to
+    // the beginning of the slab in the code block above.
+    return deleteFromDevice(Ptr);
+  }
+};
+} // namespace impl
+
+BumpAllocatorTy::BumpAllocatorTy(DeviceTy &Dev)
+    : Impl(std::make_shared<impl::BumpAllocatorImplTy>(Dev)) {}
+
+void *BumpAllocatorTy::allocate(size_t Size, void *HstPtr) {
+  return Impl->allocate(Size, HstPtr);
+}
+
+int32_t BumpAllocatorTy::deallocate(void *Ptr) { return Impl->deallocate(Ptr); }
+} // namespace memory

openmp/libomptarget/src/memory.h

 
 #include <cstddef>
 #include <memory>
 
 // Forward declaration
 struct DeviceTy;
 
 namespace memory {
+class BumpAllocatorTy;
 namespace impl {
-class MemoryManagerImplTy;
+template <typename Allocator> class MemoryManagerImplTy;
 } // namespace impl
 
 class MemoryManagerTy {
-  std::shared_ptr<impl::MemoryManagerImplTy> Impl;
+  std::shared_ptr<impl::MemoryManagerImplTy<BumpAllocatorTy>> Impl;
 
 public:
   /// Constructor
   MemoryManagerTy(DeviceTy &D, size_t Threshold = 0);
 
   /// Allocate memory of size \p Size from target device. \p HstPtr is used to
   /// assist the allocation.
   void *allocate(size_t Size, void *HstPtr);
 
   /// Deallocate memory pointed by \p TgtPtr
   int free(void *TgtPtr);
 };
 } // namespace memory

openmp/libomptarget/src/memory.cpp

 #include <algorithm>
 #include <cassert>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <unordered_map>
 #include <vector>
 
+#include "allocator.h"
 #include "device.h"
 #include "memory.h"
 #include "rtl.h"
 
 #ifdef OMPTARGET_DEBUG
 #include "private.h"
 #endif
 
struct NodeTy {
 
   /// Constructor
   NodeTy(size_t Size, void *Ptr) : Size(Size), Ptr(Ptr) {}
 
   /// To make Nodes ordered when they're put into \p std::multiset.
   bool operator<(const NodeTy &RHS) { return Size < RHS.Size; }
 };
 
-class MemoryManagerImplTy {
+template <typename AllocatorTy> class MemoryManagerImplTy {
   /// Nodes are used in a format of \p std::shared_ptr<Node>
   using NodePtrTy = std::shared_ptr<NodeTy>;
   /// A \p FreeList is a set of Nodes. We're using \p std::multiset here to make
   /// the \p find procedure more efficient.
   using FreeListTy = std::multiset<NodePtrTy>;
 
   /// A list of \p FreeListTy entries, each of which is a multiset of Nodes
   /// whose size is less or equal to a specific bucket size.
   std::vector<FreeListTy> FreeLists;
   /// A table to map from a target pointer to its node
   std::unordered_map<void *, NodePtrTy> PtrToNodeTable;
   /// A list of mutex for each \p FreeListTy entry
   std::vector<std::mutex> FreeListLocks;
   /// The mutex for the table \p PtrToNodeTable
   std::mutex MapTableLock;
-  /// A reference to its corresponding \p DeviceTy object
-  DeviceTy &Device;
+  /// Allocator
+  AllocatorTy Allocator;
 
-  /// Request memory from target device
-  void *allocateFromDevice(size_t Size, void *HstPtr) const {
-    return Device.RTL->data_alloc(Device.RTLDeviceID, Size, HstPtr);
+  /// Allocate data via \p Allocator
+  void *allocateOnDevice(size_t Size, void *HstPtr) {
+    return Allocator.allocate(Size, HstPtr);
   }
 
-  /// Deallocate data from device
-  int deleteFromDevice(void *Ptr) const {
-    return Device.RTL->data_delete(Device.RTLDeviceID, Ptr);
-  }
+  /// Deallocate data via \p Allocator
+  int deallocateOnDevice(void *Ptr) { return Allocator.deallocate(Ptr); }
 
   /// This function is called when it tries to allocate memory on device but the
   /// device returns out of memory. It will first free all memory in the
   /// FreeList and try to allocate again.
   void *freeAndAllocate(size_t Size, void *HstPtr) {
     // Deallocate all memory in FreeList
     for (int I = 0; I < NumBuckets; ++I) {
       FreeListTy &List = FreeLists[I];
       if (List.empty())
         continue;
       std::lock_guard<std::mutex> Lock(FreeListLocks[I]);
       for (const NodePtrTy &N : List)
-        deleteFromDevice(N->Ptr);
+        deallocateOnDevice(N->Ptr);
       FreeLists[I].clear();
     }
 
     // Try allocate memory again
-    return allocateFromDevice(Size, HstPtr);
+    return allocateOnDevice(Size, HstPtr);
   }
 
 public:
   /// Constructor
   MemoryManagerImplTy(DeviceTy &Dev)
-      : FreeLists(NumBuckets), FreeListLocks(NumBuckets), Device(Dev) {}
+      : FreeLists(NumBuckets), FreeListLocks(NumBuckets), Allocator(Dev) {}
 
   /// Destructor
   ~MemoryManagerImplTy() {
     // TODO: There is a little issue that target plugin is destroyed before this
     // object, therefore the memory free will not succeed.
     // Deallocate all memory in FreeList
     for (int I = 0; I < NumBuckets; ++I) {
       // We don't need lock here because only one thread can execute it
       FreeListTy &List = FreeLists[I];
       for (const NodePtrTy &N : List)
-        deleteFromDevice(N->Ptr);
+        deallocateOnDevice(N->Ptr);
     }
 
     // Deallocate all memory in map
     for (std::pair<void *, NodePtrTy> P : PtrToNodeTable) {
       assert(P.second->Ptr && "nullptr in map table");
-      deleteFromDevice(P.second->Ptr);
+      deallocateOnDevice(P.second->Ptr);
     }
   }
 
   void *allocate(size_t Size, void *HstPtr) {
     // If the size is zero, we will not bother the target device. Just return
     // nullptr directly.
     if (Size == 0)
       return nullptr;
 
     DP("MemoryManagerImplTy::allocate: size %zu with host pointer " DPxMOD
        ".\n",
        Size, DPxPTR(HstPtr));
 
     // If the size is greater than the threshold, allocate it directly from
     // device.
     if (Size > SizeThreshold) {
       DP("%zu is greater than the threshold %zu. Allocate it directly from "
          "device\n",
          Size, SizeThreshold);
-      void *TgtPtr = allocateFromDevice(Size, HstPtr);
+      void *TgtPtr = allocateOnDevice(Size, HstPtr);
       // We cannot get memory from the device. It might be due to OOM. Let's
       // free all memory in FreeLists and try again.
       if (TgtPtr == nullptr) {
         DP("Failed to get memory from device. Free all memory in FreeLists and "
            "try again.\n");
         TgtPtr = freeAndAllocate(Size, HstPtr);
       }
 
// Try to get a node from FreeList
       }
     }
 
     // We cannot find a valid node in FreeLists. Let's allocate from device and
     // create a node for it.
     if (NodePtr == nullptr) {
       DP("Cannot find a node in the FreeLists. Allocate from device.\n");
       // Allocate one from device
-      void *TgtPtr = allocateFromDevice(Size, HstPtr);
+      void *TgtPtr = allocateOnDevice(Size, HstPtr);
 
       // If TgtPtr is nullptr, it might be due to OOM. Call freeAndAllocate to
       // free some memory in FreeList and then allocate again
       if (TgtPtr == nullptr) {
         DP("Failed to get memory from device. Free all memory in FreeLists and "
            "try again.\n");
         TgtPtr = freeAndAllocate(Size, HstPtr);
       }
// Look it up into the table
         P = Itr->second;
         PtrToNodeTable.erase(Itr);
       }
     }
 
     // The memory is not managed by the manager
     if (P == nullptr) {
       DP("Cannot find its node. Delete it from device directly.\n");
-      return deleteFromDevice(TgtPtr);
+      return deallocateOnDevice(TgtPtr);
     }
 
     // Insert the node to the free list
     const int B = findBucket(P->Size);
     FreeListTy &List = FreeLists[B];
 
     DP("Found its node " DPxMOD ". Insert it to bucket %d.\n", DPxPTR(P.get()),
        B);
 
     std::lock_guard<std::mutex> G(FreeListLocks[B]);
     List.insert(P);
 
     return OFFLOAD_SUCCESS;
   }
 };
 } // namespace impl
 
 void *MemoryManagerTy::allocate(size_t Size, void *HstPtr) {
   return Impl->allocate(Size, HstPtr);
 }
 
 int MemoryManagerTy::free(void *TgtPtr) { return Impl->free(TgtPtr); }
 
 MemoryManagerTy::MemoryManagerTy(DeviceTy &D, size_t Threshold)
-    : Impl(new impl::MemoryManagerImplTy(D)) {
+    : Impl(new impl::MemoryManagerImplTy<BumpAllocatorTy>(D)) {
   if (Threshold)
     impl::SizeThreshold = Threshold;
 }
 } // namespace memory