SWDEV-497841 - Add VmHeapArray support (#76)

Add VmHeapArray class to reduce the pressure on VA reservation, since multiple memory pools can be active at the same time. [ROCm/clr commit: e974f7fde1]
2025-04-03 11:34:18 -04:00
@@ -378,6 +378,7 @@ hipError_t MemoryPool::SetAttribute(hipMemPoolAttr attr, void* value) {
        return hipErrorInvalidValue;
      }
      max_total_size_ = reset;
+      ResetMaxMappedSize();
      break;
    case hipMemPoolAttrUsedMemCurrent:
      // Should be GetAttribute only
@@ -424,7 +425,8 @@ hipError_t MemoryPool::GetAttribute(hipMemPoolAttr attr, void* value) {
      break;
    case hipMemPoolAttrReservedMemHigh:
      // High watermark of all allocated memory in OS, since the last reset
-      *reinterpret_cast<uint64_t*>(value) = max_total_size_;
+      *reinterpret_cast<uint64_t*>(value) = (state_.use_vm_heap_)
+          ? MaxMappedSize() : max_total_size_;
      break;
    case hipMemPoolAttrUsedMemCurrent:
      // Total currently used memory by the pool
@@ -102,7 +102,7 @@ class Heap : public amd::EmbeddedObject {
 public:
  typedef std::map<std::pair<size_t, amd::Memory*>, MemoryTimestamp> SortedMap;

-  Heap(hip::Device* device, amd::VmHeap& vm_heap)
+  Heap(hip::Device* device, amd::VmHeapArray& vm_heap)
    : total_size_(0)
    , max_total_size_(0)
    , release_threshold_(0)
@@ -190,7 +190,7 @@ private:
  uint64_t release_threshold_;  //!< Threshold size in bytes for memory release from heap, default 0

  hip::Device*  device_;        //!< Hip device the allocations will reside
-  amd::VmHeap&  vm_heap_;       //!< Managed heap for memory allocaitons
+  amd::VmHeapArray&  vm_heap_;  //!< Managed heap for memory allocaitons
  bool use_vm_heap_ = false;    //!< Use virtual heap or direct allocations
 };

@@ -198,7 +198,7 @@ private:
 /// @note: the logic also will look in free_heap for possible reuse.
 /// hipMemPoolReuseAllowOpportunistic option will validate if HIP event,
 /// associated with memory is done, then reuse can be performed.
-class MemoryPool : public amd::ReferenceCountedObject, amd::VmHeap {
+class MemoryPool : public amd::ReferenceCountedObject, amd::VmHeapArray {
 public:
  struct SharedAccess {
    int device_id_;             //!< Device ID for access with a specified shared resource
@@ -214,7 +214,8 @@ class MemoryPool : public amd::ReferenceCountedObject, amd::VmHeap {
  };

  MemoryPool(hip::Device* device, const hipMemPoolProps* props = nullptr, bool phys_mem = false)
-      : VmHeap(device->devices()[0]),
+      : VmHeapArray(device->devices()[0],
+                    [this]()->amd::HostQueue&{ return *device_->NullStream(); }),
        busy_heap_(device, *this),
        free_heap_(device, *this),
        lock_pool_ops_(true),
@@ -259,9 +260,6 @@ class MemoryPool : public amd::ReferenceCountedObject, amd::VmHeap {
    }
  }

-  /// Returns a queue for virtual memory map/unmap operations
-  virtual amd::HostQueue& GetVmQueue() final { return *device_->NullStream(); }
-
  /// The same stream can reuse memory without HIP event validation
  void* AllocateMemory(size_t size, Stream* stream, void* dptr = nullptr);

@@ -2231,7 +2231,7 @@ bool ResourceCache::addGpuMemory(Resource::Descriptor* desc, GpuMemoryReference*
  // Check if runtime can free suballocation
  if (desc->type_ == Resource::VaRange) {
    // We do no sub allocate VA Range.
-    result = true;
+    result = false;
  } else if ((desc->type_ == Resource::Local) && !desc->SVMRes_) {
    result = mem_sub_alloc_local_.Free(&lockCacheOps_, ref, offset);
  } else if ((desc->type_ == Resource::Local) && desc->SVMRes_) {
@@ -97,41 +97,45 @@ bool VmHeap::UncommitMemory(void* addr, size_t size) {
 }

 // ================================================================================================
-VmHeap::VmHeap(Device* device, size_t va_size, size_t chunk_size)
+VmHeap::VmHeap(Device* device, size_t va_size, size_t chunk_size, GetQueueFunc get_queue)
  : block_alignment_(kMinBlockAlignment)
  , chunk_size_(chunk_size)
  , lock_(true)
-  , device_(device) {
+  , device_(device)
+  , get_vm_queue_(get_queue) {
  va_size_ = alignUp(va_size, chunk_size);
+  free_size_ = va_size_;
 }

 // ================================================================================================
 VmHeap::~VmHeap() {
-  ScopedLock k(lock_);
+  if (created_) {
+    ScopedLock k(lock_);

-  // Release all heap blocks
-  HeapBlock* walk, * next;
-  walk = busy_list_;
-  while (walk) {
-    next = walk->next_;
-    FreeBlock(walk);
-    walk = next;
-  }
+    // Release all heap blocks
+    HeapBlock* walk, * next;
+    walk = busy_list_;
+    while (walk) {
+      next = walk->next_;
+      FreeBlock(walk);
+      walk = next;
+    }

-  walk = free_list_;
-  while (walk) {
-    next = walk->next_;
-    delete walk;
-    walk = next;
-  }
+    walk = free_list_;
+    while (walk) {
+      next = walk->next_;
+      delete walk;
+      walk = next;
+    }

-  if (mapped_mem_.size() > 0) {
-    // Unmap the entire memory range
-    UnmapPhysMemory(0, va_size_);
-  }
-  // Destroy virtual address space
-  if (base_address_ != nullptr) {
-    ReleaseAddressRange(base_address_);
+    if (mapped_mem_.size() > 0) {
+      // Unmap the entire memory range
+      UnmapPhysMemory(0, va_size_);
+    }
+    // Destroy virtual address space
+    if (base_address_ != nullptr) {
+      ReleaseAddressRange(base_address_);
+    }
  }
 }

@@ -211,6 +215,7 @@ void VmHeap::UnmapPhysMemory(size_t offset, size_t size) {

 // ================================================================================================
 void VmHeap::TrimPhysMemory(size_t unmap_threshold) {
+  ScopedLock k(lock_);
  auto current = free_list_;
  auto unmap_org = unmap_threshold_;
  unmap_threshold_ = unmap_threshold;
@@ -416,4 +421,106 @@ void VmHeap::MergeBlock(HeapBlock** head, HeapBlock* blk) {
  }
 }

+// ================================================================================================
+address VmHeapArray::Alloc(size_t size) {
+  address addr = nullptr;
+  for (uint32_t i = 0; i < kMaxArraySize; ++i) {
+    if (vm_heaps_[i]->free_size_ > (size + VmHeap::kChunkSize)) {
+      addr = vm_heaps_[i]->Alloc(size);
+      if (addr != nullptr) {
+        break;
+      }
+    }
+  }
+  return addr;
+}
+
+// ================================================================================================
+void VmHeapArray::Free(amd::Memory* memory) {
+  const device::Memory* dev_mem = memory->getDeviceMemory(*device_);
+  void* addr = reinterpret_cast<void*>(dev_mem->virtualAddress());
+  if (addr == nullptr) {
+    addr = memory->getSvmPtr();
+  }
+  for (uint32_t i = 0; i < kMaxArraySize; ++i) {
+    if (vm_heaps_[i]->created_ && vm_heaps_[i]->InRange(addr)) {
+      vm_heaps_[i]->Free(memory);
+      break;
+    }
+  }
+  uint64_t freed = 0;
+  for (uint32_t i = 0; i < kMaxArraySize; ++i) {
+    freed += vm_heaps_[i]->FreeMappedSize();
+  }
+  if (freed > unmap_threshold_) {
+    uint64_t extra = freed - unmap_threshold_;
+    uint64_t trim = (extra < unmap_threshold_) ? (unmap_threshold_ - extra) : 0;
+    TrimPhysMemory(trim);
+  }
+}
+
+// ================================================================================================
+void VmHeapArray::TrimPhysMemory(size_t unmap_threshold) {
+  for (uint i = 0; i < kMaxArraySize; ++i) {
+    // Check the threshold against the accumulated sizes in all heaps
+    if (vm_heaps_[i]->created_ && [this]() {
+      uint64_t size = 0;
+      for (uint i = 0; i < kMaxArraySize; ++i) {
+        size += vm_heaps_[i]->FreeMappedSize();
+      }
+      return size;
+     }() > unmap_threshold) {
+      vm_heaps_[i]->TrimPhysMemory(unmap_threshold);
+    } else {
+      break;
+    }
+  }
+}
+
+// ================================================================================================
+void VmHeapArray::SetUnmapThreshold(uint64_t threshold) {
+  for (uint i = 0; i < kMaxArraySize; ++i) {
+    // Note: it's not precisely correct to use the same threshold in all heaps,
+    // but the logic will trim heaps in Free()
+    if (vm_heaps_[i]->created_) {
+      vm_heaps_[i]->SetUnmapThreshold(threshold);
+    }
+  }
+  unmap_threshold_ = threshold;
+}
+
+// ================================================================================================
+uint64_t VmHeapArray::MappedSize() const {
+  uint64_t size = 0;
+  for (uint i = 0; i < kMaxArraySize; ++i) {
+    size += vm_heaps_[i]->MappedSize();
+  }
+  return size;
+}
+
+// ================================================================================================
+uint64_t VmHeapArray::FreeMappedSize() const {
+  uint64_t size = 0;
+  for (uint i = 0; i < kMaxArraySize; ++i) {
+    size += vm_heaps_[i]->FreeMappedSize();
+  }
+  return size;
+}
+
+// ================================================================================================
+uint64_t VmHeapArray::MaxMappedSize() const {
+  uint64_t size = 0;
+  for (uint i = 0; i < kMaxArraySize; ++i) {
+    size += vm_heaps_[i]->max_mapped_size_;
+  }
+  return size;
+}
+
+// ================================================================================================
+void VmHeapArray::ResetMaxMappedSize() {
+  for (uint i = 0; i < kMaxArraySize; ++i) {
+    vm_heaps_[i]->max_mapped_size_ = 0;
+  }
+}
+
 } // namespace amd
@@ -20,14 +20,17 @@

 #pragma once

+#include <functional>
 #include "top.hpp"
 #include "device/device.hpp"
 #include "object.hpp"
 #include "commandqueue.hpp"

 namespace amd {
-class VmHeap;
+
 class HeapBlock;
+class VmHeap;
+class VmHeapArray;

 class HeapBlock : public amd::HeapObject {
 public:
@@ -66,27 +69,25 @@ private:

 class VmHeap {
 public:
+  friend VmHeapArray;
  static const size_t kChunkSize = 32 * Mi; //!< Chunk size, must be power of 2
  static const size_t kMinBlockAlignment = 256;
+  typedef std::function<amd::HostQueue&()> GetQueueFunc;

-  VmHeap(Device* device      //!< GPU device object
+  VmHeap(Device* device,        //!< GPU device object
+         GetQueueFunc get_queue //!< Function to retrieve a map queue
         )
-      : VmHeap(device, device->info().globalMemSize_ / 8, kChunkSize) {}
+      : VmHeap(device, device->info().globalMemSize_ / 8, kChunkSize, get_queue) {}

  VmHeap(Device* device,        //!< GPU device object
         size_t  va_size,       //!< The size of the allocated heap (bytes).Virtual address space
-         size_t  chunk_size     //!< The size of single chunk for physical memory growth
+         size_t  chunk_size,    //!< The size of single chunk for physical memory growth
+         GetQueueFunc get_queue //!< Function to retrieve a map/unmap queue
         );

-  //! Ceates heap object. Reserves virtual address range for the heap operation
-  bool Create();
-
  //! Heap destructor
  virtual ~VmHeap();

-  //! Returns a queue for VM map/unmap operations
-  virtual amd::HostQueue& GetVmQueue() = 0;
-
  //! Returns a pointer to the allocated device memory from a heap
  address Alloc(
      size_t size     //! The allocation size
@@ -107,11 +108,19 @@ public:
  //! Returns mapped memory size (allocated physical memory) without actual allocations
  uint64_t FreeMappedSize() const { return mapped_size_ - (va_size_ - free_size_); }

+  //! Returns true if the address is in the range of this heap
+  bool InRange(void* addr) {
+    return ((addr >= base_address_) && (addr <= (base_address_ + va_size_))) ? true : false;
+  }
+
 private:
  VmHeap() = delete;
  VmHeap(const VmHeap&) = delete;
  VmHeap& operator=(const VmHeap&) = delete;

+  //! Ceates heap object. Reserves virtual address range for the heap operation
+  bool Create();
+
  //! Reseves address range for memory allocations
  address ReserveAddressRange(address start, size_t size, size_t alignment);

@@ -151,6 +160,9 @@ private:
  //! Join two blocks, transferring the size of the second into the first and deleting the second
  void Join2Blocks(HeapBlock* first, HeapBlock* second) const;

+  //! Returns a queue for VM map/unmap operations
+  amd::HostQueue& GetVmQueue() const { return get_vm_queue_(); }
+
  address       base_address_ = nullptr;  //!< GPU virtual address base of the heap
  amd::Memory*  base_memory_ = nullptr;   //!< VA space base object, used in the view creation
  HeapBlock*    free_list_ = nullptr;     //!< Head block for free list
@@ -165,8 +177,64 @@ private:
  bool          created_ = false;         //!< Used for deferred VM heap allocation
  amd::Monitor  lock_;                    //!< Lock to serialise heap accesses
  Device*       device_;                  //!< Device that owns this heap
+  GetQueueFunc  get_vm_queue_;            //!< Queue for VM operations

  std::vector<bool> mapped_mem_;  //!< A map of mapped memory, the size is total_size/chunk_size
 };

+//! Implements an array of vm heaps of different sizes for more efficient management
+class VmHeapArray {
+public:
+  VmHeapArray(Device* device,    //!< GPU device object
+              VmHeap::GetQueueFunc get_queue  //!< Function to retrieve a map queue
+  ) : heap0_(device, device->info().globalMemSize_ / 8, VmHeap::kChunkSize, get_queue)
+    , heap1_(device, device->info().globalMemSize_ / 4, VmHeap::kChunkSize, get_queue)
+    , heap2_(device, device->info().globalMemSize_ * 5 / 8, VmHeap::kChunkSize, get_queue)
+    , heap3_(device, device->info().globalMemSize_, VmHeap::kChunkSize, get_queue)
+    , device_(device) {}
+
+  //! Returns a pointer to the allocated device memory from a heap
+  address Alloc(
+    size_t size     //! The allocation size
+    );
+
+  //! Release memory back to the VM heap
+  void Free(amd::Memory* memory);
+
+  //! Unmaps freed memory based on the threshold
+  void TrimPhysMemory(size_t unmap_threshold);
+
+  //! Enable memory unmap threashold (default 0 unmap always)
+  void SetUnmapThreshold(uint64_t threshold);
+
+  //! Returns mapped memory size (total allocated physical memory)
+  uint64_t MappedSize() const;
+
+  //! Returns mapped memory size (allocated physical memory) without actual allocations
+  uint64_t FreeMappedSize() const;
+
+  //! Returns the maximum mapped memory size
+  uint64_t MaxMappedSize() const;
+
+  //! Returns the maximum mapped memory size
+  void ResetMaxMappedSize();
+
+private:
+  VmHeapArray() = delete;
+  VmHeapArray(const VmHeapArray&) = delete;
+  VmHeapArray& operator=(const VmHeapArray&) = delete;
+
+  static const uint32_t kMaxArraySize = 4;  //!< The number of vm heap in the array
+  // @note: gcc10.2 or lower wrongly uses copy constructor in the initialization
+  // of VmHeap array of objects. Hence, use an array of VmHeap pointers for now
+  VmHeap* vm_heaps_[kMaxArraySize] = {&heap0_, &heap1_, &heap2_, &heap3_};  //!< The array of heaps
+  VmHeap  heap0_;
+  VmHeap  heap1_;
+  VmHeap  heap2_;
+  VmHeap  heap3_;
+
+  uint64_t unmap_threshold_ = 0;  //!< Unmap threshold in bytes,used to release phys memory
+  Device* device_;                //!< Device that owns this heap
+};
+
 } // namespace amd