SWDEV-497841 - Add virtual memory heap

Add initial implementation of virtual memory heap with dynamic virtual memory mapping support for memory pools. DEBUG_HIP_MEM_POOL_VMHEAP controls the new method. Change-Id: I8dc5be2e0f34ab472f1800f43bb6243639a5e500 [ROCm/clr commit: 296dce5570]
2025-01-24 19:26:39 -05:00
parent 6f2a603277
commit f9d9b2c441
11 changed files with 670 additions and 39 deletions
@@ -315,7 +315,7 @@ bool Device::existsActiveStreamForDevice() {

 // ================================================================================================
 Device::~Device() {
-  if (default_mem_pool_ != nullptr) {
+  if ((IS_LINUX || !DEBUG_HIP_MEM_POOL_VMHEAP) && (default_mem_pool_ != nullptr)) {
    default_mem_pool_->release();
  }

@@ -1,4 +1,4 @@
-/* Copyright (c) 2022-2023 Advanced Micro Devices, Inc.
+/* Copyright (c) 2022-2025 Advanced Micro Devices, Inc.

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -90,16 +90,19 @@ bool Heap::RemoveMemory(amd::Memory* memory, MemoryTimestamp* ts) {
 }

 // ================================================================================================
-Heap::SortedMap::iterator Heap::EraseAllocaton(Heap::SortedMap::iterator& it) {
+Heap::SortedMap::iterator Heap::EraseAllocation(Heap::SortedMap::iterator& it) {
  auto memory = it->first.second;
  const device::Memory* dev_mem = memory->getDeviceMemory(*device_->devices()[0]);
  void* dev_mem_vaddr = reinterpret_cast<void*>(dev_mem->virtualAddress());
  total_size_ -= it->first.first;
+  if (dev_mem_vaddr == nullptr) {
+    dev_mem_vaddr = memory->getSvmPtr();
+  }

-  if (dev_mem_vaddr != nullptr) {
-    amd::SvmBuffer::free(memory->getContext(), dev_mem_vaddr);
+  if (use_vm_heap_) {
+    vm_heap_.Free(memory);
  } else {
-    amd::SvmBuffer::free(memory->getContext(), memory->getSvmPtr());
+    amd::SvmBuffer::free(memory->getContext(), dev_mem_vaddr);
  }
  // Clear HIP event
  it->second.SetEvent(nullptr);
@@ -119,11 +122,15 @@ bool Heap::ReleaseAllMemory(size_t min_bytes_to_hold, bool safe_release) {
      it->second.Wait();
    }
    if (it->second.IsSafeRelease()) {
-      it = EraseAllocaton(it);
+      it = EraseAllocation(it);
    } else {
      ++it;
    }
  }
+  // Handle managed pool with trim
+  if (vm_heap_.FreeMappedSize() > min_bytes_to_hold) {
+    vm_heap_.TrimPhysMemory(min_bytes_to_hold);
+  }
  return true;
 }

@@ -131,11 +138,12 @@ bool Heap::ReleaseAllMemory(size_t min_bytes_to_hold, bool safe_release) {
 bool Heap::ReleaseAllMemory() {
  for (auto it = allocations_.begin(); it != allocations_.end();) {
    // Make sure the heap holds the minimum number of bytes
-    if (total_size_ <= release_threshold_) {
+    // @note: Managed memory controls the threshold on its own
+    if (!use_vm_heap_ && (total_size_ <= release_threshold_)) {
      return true;
    }
    if (it->second.IsSafeRelease()) {
-      it = EraseAllocaton(it);
+      it = EraseAllocation(it);
    } else {
      ++it;
    }
@@ -187,13 +195,17 @@ void* MemoryPool::AllocateMemory(size_t size, Stream* stream, void* dptr) {
    }
    cl_svm_mem_flags flags = (state_.interprocess_) ? ROCCLR_MEM_INTERPROCESS : 0;
    flags |= (state_.phys_mem_) ? ROCCLR_MEM_PHYMEM : 0;
-    dev_ptr = amd::SvmBuffer::malloc(*context, flags, size, dev_info.memBaseAddrAlign_, nullptr);
+    if (state_.use_vm_heap_) {
+      dev_ptr = Alloc(size);
+    } else {
+      dev_ptr = amd::SvmBuffer::malloc(*context, flags, size, dev_info.memBaseAddrAlign_, nullptr);
+    }
    if (dev_ptr == nullptr) {
      size_t free = 0, total =0;
      hipError_t err = hipMemGetInfo(&free, &total);
      if (err == hipSuccess) {
-        LogPrintfError("Allocation failed : Device memory : required :%zu | free :%zu | total :%zu",
-          size, free, total);
+        LogPrintfError("Allocation failed : Device memory : required :\
+          %zu | free :%zu | total :%zu", size, free, total);
      }
      return nullptr;
    }
@@ -236,7 +248,7 @@ bool MemoryPool::FreeMemory(amd::Memory* memory, Stream* stream, Event* event) {
  {
    amd::ScopedLock lock(lock_pool_ops_);

-    if (memory->getUserData().phys_mem_obj != nullptr) {
+    if (!state_.use_vm_heap_ && memory->getUserData().phys_mem_obj != nullptr) {
      memory = memory->getUserData().phys_mem_obj;
    }

@@ -408,8 +420,9 @@ hipError_t MemoryPool::GetAttribute(hipMemPoolAttr attr, void* value) {
      *reinterpret_cast<uint64_t*>(value) = free_heap_.GetReleaseThreshold();
      break;
    case hipMemPoolAttrReservedMemCurrent:
-      // All allocate memory by the pool in OS
-      *reinterpret_cast<uint64_t*>(value) = busy_heap_.GetTotalSize() + free_heap_.GetTotalSize();
+      // All allocated memory by the pool in OS
+      *reinterpret_cast<uint64_t*>(value) = (state_.use_vm_heap_) ? MappedSize() :
+        (busy_heap_.GetTotalSize() + free_heap_.GetTotalSize());
      break;
    case hipMemPoolAttrReservedMemHigh:
      // High watermark of all allocated memory in OS, since the last reset
@@ -1,4 +1,4 @@
-/* Copyright (c) 2022 Advanced Micro Devices, Inc.
+/* Copyright (c) 2022-2025 Advanced Micro Devices, Inc.

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -23,6 +23,7 @@
 #include <hip/hip_runtime.h>
 #include "hip_event.hpp"
 #include "hip_internal.hpp"
+#include "platform/vmheap.hpp"
 #include <unordered_map>
 #include <unordered_set>

@@ -101,8 +102,12 @@ class Heap : public amd::EmbeddedObject {
 public:
  typedef std::map<std::pair<size_t, amd::Memory*>, MemoryTimestamp> SortedMap;

-  Heap(hip::Device* device):
-    total_size_(0), max_total_size_(0), release_threshold_(0), device_(device) {}
+  Heap(hip::Device* device, amd::VmHeap& vm_heap)
+    : total_size_(0)
+    , max_total_size_(0)
+    , release_threshold_(0)
+    , device_(device)
+    , vm_heap_(vm_heap) {}
  ~Heap() {}

  /// Adds allocation into the heap on a specific stream
@@ -134,7 +139,10 @@ public:
  bool IsEmpty() const { return (allocations_.size() == 0) ? true : false; }

  /// Set the memory release threshold
-  void SetReleaseThreshold(uint64_t value) { release_threshold_ = value; }
+  void SetReleaseThreshold(uint64_t value) {
+    release_threshold_ = value;
+    vm_heap_.SetUnmapThreshold(value);
+  }

  /// Set the memory release threshold
  uint64_t GetReleaseThreshold() const { return release_threshold_; }
@@ -149,7 +157,7 @@ public:
  void SetMaxTotalSize(uint64_t value) { max_total_size_ = value; }

  /// Erases single allocation form the heap's map
-  SortedMap::iterator EraseAllocaton(SortedMap::iterator& it);
+  SortedMap::iterator EraseAllocation(SortedMap::iterator& it);

  /// Add a safe stream for  quick looks-ups in all allocations
  void AddSafeStream(Stream* event_stream, Stream* wait_stream) {
@@ -162,6 +170,13 @@ public:
  bool IsActiveMemory(amd::Memory* memory) const {
    return (allocations_.find({memory->getSize(), memory}) != allocations_.end());
  }
+
+  /// Enabled VM heap for memory, instead of direct allocations
+  void EnableVmHeap() { use_vm_heap_ = true; }
+
+  /// Returns true if heap uses virtual memory
+  bool UseVmHeap() const { return use_vm_heap_; }
+
  const auto& Allocations() { return allocations_; }

 private:
@@ -174,14 +189,16 @@ private:
  uint64_t max_total_size_;     //!< Maximum heap allocation size
  uint64_t release_threshold_;  //!< Threshold size in bytes for memory release from heap, default 0

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

 /// Allocates memory in the pool on the specified stream and places the allocation into busy_heap_
 /// @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 {
+class MemoryPool : public amd::ReferenceCountedObject, amd::VmHeap {
 public:
  struct SharedAccess {
    int device_id_;             //!< Device ID for access with a specified shared resource
@@ -197,9 +214,10 @@ class MemoryPool : public amd::ReferenceCountedObject {
  };

  MemoryPool(hip::Device* device, const hipMemPoolProps* props = nullptr, bool phys_mem = false)
-      : busy_heap_(device),
-        free_heap_(device),
-        lock_pool_ops_(true), /* Pool operations */
+      : VmHeap(device->asContext()->devices()[0], *device->NullStream()),
+        busy_heap_(device, *this),
+        free_heap_(device, *this),
+        lock_pool_ops_(true),
        device_(device),
        shared_(nullptr),
        max_total_size_(0) {
@@ -220,9 +238,15 @@ class MemoryPool : public amd::ReferenceCountedObject {
                     .reserved = {}};
    }
    state_.interprocess_ = properties_.handleTypes != hipMemHandleTypeNone;
+    // Check if VM heap can be enabled
+    if (DEBUG_HIP_MEM_POOL_VMHEAP && !state_.phys_mem_ && !state_.interprocess_) {
+      state_.use_vm_heap_ = true;
+      busy_heap_.EnableVmHeap();
+      free_heap_.EnableVmHeap();
+    }
  }

-  virtual ~MemoryPool() {
+  virtual ~MemoryPool() override {
    if (!busy_heap_.IsEmpty()) {
      LogError("Shouldn't destroy pool with busy allocations!");
    }
@@ -320,6 +344,7 @@ class MemoryPool : public amd::ReferenceCountedObject {
      uint32_t interprocess_ : 1;   //!< Memory pool can be used in interprocess communications
      uint32_t graph_in_use_ : 1;   //!< Memory pool was used in a graph execution
      uint32_t phys_mem_ : 1;       //!< Mempool is used for graphs and will have physical allocations
+      uint32_t use_vm_heap_ : 1;    //!< Use VM heap or direct allocations
    };
    uint32_t value_;
  } state_;
@@ -327,7 +352,8 @@ class MemoryPool : public amd::ReferenceCountedObject {
  hipMemPoolProps properties_;  //!< Properties of the memory pool
  amd::Monitor lock_pool_ops_;  //!< Access to the pool must be lock protected
  std::map<hip::Device*, hipMemAccessFlags> access_map_;  //!< Map of access to the pool from devices
-  hip::Device*  device_;    //!< Hip device the heap will reside
+
+  hip::Device* device_;     //!< Hip device the heap will reside
  SharedMemPool* shared_;   //!< Pointer to shared memory for IPC
  uint64_t max_total_size_; //!< Max of total reserved memory in the pool since last reset
 };
@@ -86,6 +86,7 @@ target_sources(rocclr PRIVATE
  ${ROCCLR_SRC_DIR}/platform/commandqueue.cpp
  ${ROCCLR_SRC_DIR}/platform/context.cpp
  ${ROCCLR_SRC_DIR}/platform/kernel.cpp
+  ${ROCCLR_SRC_DIR}/platform/vmheap.cpp
  ${ROCCLR_SRC_DIR}/platform/memory.cpp
  ${ROCCLR_SRC_DIR}/platform/ndrange.cpp
  ${ROCCLR_SRC_DIR}/platform/program.cpp
@@ -82,7 +82,8 @@ GpuMemoryReference* GpuMemoryReference::Create(const Device& dev,
      return nullptr;
    }
  }
-  if (!createInfo.flags.sdiExternal) {
+  memRef->va_range_ = createInfo.flags.virtualAlloc;
+  if (!createInfo.flags.sdiExternal && !createInfo.flags.virtualAlloc) {
    // Update free memory size counters
    dev.updateAllocedMemory(memRef->gpuMem_->Desc().heaps[0], memRef->gpuMem_->Desc().size, false);
  }
@@ -280,7 +281,7 @@ GpuMemoryReference::~GpuMemoryReference() {
    iMem()->Unmap();
  }
  if (!(iMem()->Desc().flags.isShared || iMem()->Desc().flags.isExternal ||
-        iMem()->Desc().flags.isExternPhys)) {
+        iMem()->Desc().flags.isExternPhys || va_range_)) {
    // Update free memory size counters
    device_.updateAllocedMemory(iMem()->Desc().heaps[0], iMem()->Desc().size, true);
  }
@@ -70,6 +70,7 @@ class GpuMemoryReference : public amd::ReferenceCountedObject {
  const Device& device_;     //!< GPU device
  //! @note: This field is necessary for the thread safe release only
  VirtualGPU* gpu_;  //!< Resource will be used only on this queue
+  bool va_range_ = false; //!< Resource is a VA range(@note: PAL doesn't provide this info)

 protected:
  //! Default destructor
@@ -1,4 +1,4 @@
-/* Copyright (c) 2010 - 2022 Advanced Micro Devices, Inc.
+/* Copyright (c) 2010 - 2025 Advanced Micro Devices, Inc.

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -192,12 +192,12 @@ void Memory::operator delete(void* p, const Context& context) { Memory::operator

 void Memory::addSubBuffer(Memory* view) {
  amd::ScopedLock lock(lockMemoryOps());
-  subBuffers_.push_back(view);
+  subBuffers_.emplace(view);
 }

 void Memory::removeSubBuffer(Memory* view) {
  amd::ScopedLock lock(lockMemoryOps());
-  subBuffers_.remove(view);
+  subBuffers_.erase(view);
 }

 bool Memory::allocHostMemory(void* initFrom, bool allocHostMem, bool forceCopy) {
@@ -1,4 +1,4 @@
-/* Copyright (c) 2010 - 2023 Advanced Micro Devices, Inc.
+/* Copyright (c) 2010 - 2025 Advanced Micro Devices, Inc.

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -18,8 +18,7 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE. */

-#ifndef MEMORY_H_
-#define MEMORY_H_
+#pragma once

 #include "top.hpp"
 #include "utils/flags.hpp"
@@ -226,7 +225,7 @@ class Memory : public amd::RuntimeObject {
  Memory(const Memory&);

  Monitor lockMemoryOps_;          //!< Lock to serialize memory operations
-  std::list<Memory*> subBuffers_;  //!< List of all subbuffers for this memory object
+  std::set<Memory*> subBuffers_;   //!< List of all subbuffers for this memory object
  device::Memory* svmBase_;        //!< svmBase allocation for MGPU case

 protected:
@@ -286,7 +285,7 @@ class Memory : public amd::RuntimeObject {
  void removeSubBuffer(Memory* item);

  //! Returns the list of all subbuffers
-  std::list<Memory*>& subBuffers() { return subBuffers_; }
+  std::set<Memory*>& subBuffers() { return subBuffers_; }

  //! Returns the number of devices
  size_t numDevices() const { return numDevices_; }
@@ -711,7 +710,5 @@ class IpcBuffer : public Buffer {
  const void* handle_;  //!< Ipc handle, associated with this memory object
 };

-
 }  // namespace amd

-#endif  // MEMORY_H_
@@ -0,0 +1,419 @@
+/* Copyright (c) 2025 Advanced Micro Devices, Inc.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE. */
+
+#include <string>
+#include <fstream>
+#include <sstream>
+#include <iostream>
+#include "vmheap.hpp"
+#include "command.hpp"
+
+namespace amd {
+
+// ================================================================================================
+address VmHeap::ReserveAddressRange(address start, size_t size, size_t alignment) {
+  // Reserve a virtual address range on the device
+  void* ptr = device_->virtualAlloc(start, size, alignment);
+  // Save base memory object to accelerate access in the future
+  base_memory_ = MemObjMap::FindVirtualMemObj(ptr);
+  return reinterpret_cast<address>(ptr);
+}
+
+// ================================================================================================
+bool VmHeap::ReleaseAddressRange(void* addr) {
+  Memory* memObj = MemObjMap::FindVirtualMemObj(addr);
+  assert(memObj != nullptr && "Cannot find the Virtual MemObj entry");
+
+  // Frees address range on the device
+  device_->virtualFree(addr);
+  memObj->release();
+
+  return true;
+}
+
+// ================================================================================================
+bool VmHeap::CommitMemory(void* addr, size_t size) {
+  const auto& dev_info = device_->info();
+  size_t granularity = dev_info.virtualMemAllocGranularity_;
+  auto padded_size = alignUp(size, granularity);
+
+  // Allocate physical memory
+  void* ptr = SvmBuffer::malloc(device_->context(), ROCCLR_MEM_PHYMEM, padded_size,
+                                dev_info.memBaseAddrAlign_, nullptr);
+  if (ptr == nullptr) {
+    LogPrintfError("Failed to allocate physical memory %zd", padded_size);
+    return false;
+  }
+
+  size_t offset = 0;  // this is ignored
+  // Find physical memory in the map of all objects
+  Memory* phys_mem_obj = MemObjMap::FindMemObj(ptr, &offset);
+
+  // Map the physical memory to a virtual address
+  Command* cmd = new VirtualMapCommand(
+      map_queue_, Command::EventWaitList{}, addr, padded_size, phys_mem_obj);
+  cmd->enqueue();
+  cmd->awaitCompletion();
+  cmd->release();
+
+  // Enable memory access
+  if (!device_->SetMemAccess(addr, padded_size, Device::VmmAccess::kReadWrite)) {
+    LogError("SetAccess failed for the commited memory in VmHeap!");
+  }
+  return true;
+}
+
+// ================================================================================================
+bool VmHeap::UncommitMemory(void* addr, size_t size) {
+  Memory* vaddr_sub_obj = MemObjMap::FindMemObj(addr);
+  Memory* phys_mem_obj = vaddr_sub_obj->getUserData().phys_mem_obj;
+
+  // Unmap the physical memory from a virtual address
+  Command* cmd = new VirtualMapCommand(map_queue_, Command::EventWaitList{}, addr, size, nullptr);
+  cmd->enqueue();
+  cmd->awaitCompletion();
+  cmd->release();
+  vaddr_sub_obj->release();
+  SvmBuffer::free(device_->context(), phys_mem_obj->getSvmPtr());
+  return true;
+}
+
+// ================================================================================================
+VmHeap::VmHeap(Device* device, HostQueue& queue, size_t va_size, size_t chunk_size)
+  : block_alignment_(kMinBlockAlignment)
+  , chunk_size_(chunk_size)
+  , lock_(true)
+  , device_(device)
+  , map_queue_(queue) {
+  va_size_ = alignUp(va_size, chunk_size);
+}
+
+// ================================================================================================
+VmHeap::~VmHeap() {
+  ScopedLock k(lock_);
+
+  // 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;
+  }
+
+  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_);
+  }
+}
+
+// ================================================================================================
+bool VmHeap::Create() {
+  // Create a new GPU resource
+  base_address_ = ReserveAddressRange(0, va_size_, kChunkSize);
+  if (base_address_ == nullptr) {
+    return false;
+  }
+  free_size_ = va_size_;
+  // Set up initial free list
+  free_list_ = new HeapBlock(this, va_size_, 0);
+  if (free_list_ == nullptr) {
+    return false;
+  }
+  mapped_mem_.resize(va_size_ / chunk_size_);
+  return true;
+}
+
+// ================================================================================================
+bool VmHeap::MapPhysMemory(size_t offset, size_t size) {
+  auto start_chunk = offset / chunk_size_;
+  auto end_chunk = alignUp(offset + size, chunk_size_) / chunk_size_;
+
+  for (auto i = start_chunk; i < end_chunk; ++i) {
+    if (!mapped_mem_[i]) {
+      auto address = base_address_ + i * chunk_size_;
+      if (CommitMemory(address, chunk_size_)) {
+        mapped_size_ += chunk_size_;
+        if (mapped_size_ > max_mapped_size_) {
+          ClPrint(LOG_INFO, LOG_MEM_POOL, "VM heap grows in physical alloc to %d GB\n",
+            static_cast<int>(mapped_size_ / Gi));
+        }
+        max_mapped_size_ = std::max(max_mapped_size_, mapped_size_);
+        mapped_mem_[i] = true;
+      } else {
+        assert(false);
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+// ================================================================================================
+void VmHeap::UnmapPhysMemory(size_t offset, size_t size) {
+  auto busy_size = va_size_ - free_size_;
+  uint64_t free_mapped = alignDown(mapped_size_ - busy_size, kChunkSize);
+
+  int start_chunk = alignUp(offset, chunk_size_) / chunk_size_;
+  int end_chunk = alignDown(offset + size, chunk_size_) / chunk_size_;
+
+  for (int i = end_chunk - 1; i >= start_chunk; i--) {
+    // If free mapped memory lower than the threshold, then stop unmapping
+    if (free_mapped <= unmap_threshold_) {
+      return;
+    }
+    if (i >= mapped_mem_.size()) {
+      assert(false);
+      LogError("VM heap allocation is beyond the range!");
+      return;
+    }
+    if (mapped_mem_[i]) {
+      auto address = base_address_ + i * chunk_size_;
+      if (UncommitMemory(address, chunk_size_)) {
+        mapped_size_ -= chunk_size_;
+        free_mapped -= chunk_size_;
+        mapped_mem_[i] = false;
+      }
+      else {
+        assert(false);
+      }
+    }
+  }
+}
+
+// ================================================================================================
+void VmHeap::TrimPhysMemory(size_t unmap_threshold) {
+  auto current = free_list_;
+  auto unmap_org = unmap_threshold_;
+  unmap_threshold_ = unmap_threshold;
+  while (current != nullptr) {
+    UnmapPhysMemory(current->offset_, current->size_);
+    current = current->next_;
+  }
+  unmap_threshold_ = unmap_org;
+}
+
+// ================================================================================================
+address VmHeap::Alloc(size_t size) {
+  ScopedLock k(lock_);
+
+  if (!created_) {
+    // Create VM heap if it's not created
+    created_ = Create();
+    if (!created_) {
+      return nullptr;
+    }
+  }
+  address ptr = nullptr;
+  size_t offset = 0;
+  auto hb = AllocBlock(size + block_alignment_);
+  if (hb != nullptr) {
+    offset = ((hb->Offset() & ~kChunkSize) == 0) ? hb->Offset() + block_alignment_ : hb->Offset();
+    ptr = base_address_ + offset;
+  } else {
+    return nullptr;
+  }
+  auto memory = new (base_memory_->getContext()) Buffer(*base_memory_, 0, offset, size);
+  if (nullptr == memory || !memory->create(nullptr)) {
+    FreeBlock(hb);
+    return nullptr;
+  }
+  MemObjMap::AddMemObj(ptr, memory);
+  if (memory->getUserData().data == nullptr) {
+    memory->getUserData().data = hb;
+  }
+  ClPrint(LOG_INFO, LOG_MEM_POOL,
+    "VmHeap Alloc: %p offset(%zx + %zx) hb(%p)", ptr, hb->Offset(), memory->getSize(), hb);
+  return ptr;
+}
+
+// ================================================================================================
+void VmHeap::Free(Memory* memory) {
+  const device::Memory* dev_mem = memory->getDeviceMemory(*device_);
+  void* addr = reinterpret_cast<void*>(dev_mem->virtualAddress());
+  if (addr == nullptr) {
+    addr = memory->getSvmPtr();
+  }
+
+  if (!created_ || (addr < base_address_)) {
+    return;
+  }
+  ScopedLock k(lock_);
+  if (memory->getUserData().data != nullptr) {
+    auto hb = reinterpret_cast<HeapBlock*>(memory->getUserData().data);
+    ClPrint(LOG_INFO, LOG_MEM_POOL, "VmHeap Free: %p offset(%zx + %zx) hb(%p)",
+      addr, hb->Offset(), memory->getSize(), hb);
+    FreeBlock(hb);
+  }
+  MemObjMap::RemoveMemObj(addr);
+  memory->release();
+}
+
+// ================================================================================================
+HeapBlock* VmHeap::AllocBlock(size_t un_size) {
+  assert(un_size != 0);
+  ScopedLock k(lock_);
+  HeapBlock* walk = free_list_;
+  HeapBlock* best = nullptr;
+
+  // Round size
+  auto size = alignUp(un_size, block_alignment_);
+
+  // Walk the free list looking for a suitable block (currently best-fit)
+  while (walk) {
+    if ((walk->size_ > size) && (best == nullptr || walk->size_ < best->size_)) {
+      best = walk;
+    } else if (walk->size_ == size) {
+      // No need to split, just move to busy list
+      DetachBlock(&free_list_, walk);
+      walk->busy_ = true;
+      InsertBlock(&busy_list_, walk);
+      free_size_ -= size;
+      if (!MapPhysMemory(walk->Offset(), size)) {
+        free(walk);
+        return nullptr;
+      }
+      return walk;
+    }
+    walk = walk->next_;
+  }
+
+  if (best != nullptr) {
+    // Got one, but need to split it. Keep first part in free list,
+    // put second part into busy list
+    HeapBlock *newblock = SplitBlock(best, size);
+    newblock->busy_ = true;
+    InsertBlock(&busy_list_, newblock);
+    free_size_ -= size;
+    if (!MapPhysMemory(newblock->Offset(), size)) {
+      free(newblock);
+      return nullptr;
+    }
+    return newblock;
+  }
+
+  return nullptr;
+}
+
+// ================================================================================================
+void VmHeap::FreeBlock(HeapBlock* blk) {
+  DetachBlock(&busy_list_, blk);
+  blk->busy_ = false;
+  free_size_ += blk->size_;
+  UnmapPhysMemory(blk->offset_, blk->size_);
+  MergeBlock(&free_list_, blk);
+}
+
+// ================================================================================================
+void VmHeap::DetachBlock(HeapBlock** list, HeapBlock* blk) {
+  if (*list == blk) {
+    *list = blk->next_;
+  }
+  if (blk->prev_) {
+    blk->prev_->next_ = blk->next_;
+  }
+  if (blk->next_) {
+    blk->next_->prev_ = blk->prev_;
+  }
+}
+
+// ================================================================================================
+void VmHeap::InsertBlock(HeapBlock** head, HeapBlock* blk) {
+  if (nullptr == *head) {
+    *head = blk;
+    blk->prev_ = nullptr;
+    blk->next_ = nullptr;
+    return;
+  }
+
+  // Find the place to insert it at
+  HeapBlock* walk = *head;
+  while (walk->next_ && walk->next_->offset_ < blk->offset_) {
+    walk = walk->next_;
+  }
+
+  // Insert it
+  if (walk == *head) {
+    if (walk->offset_ >= blk->offset_) {
+      *head = blk;
+      blk->prev_ = nullptr;
+      blk->next_ = walk;
+      walk->prev_ = *head;
+      return;
+    }
+  }
+
+  blk->next_ = walk->next_;
+  blk->prev_ = walk;
+  if (walk->next_) {
+      walk->next_->prev_ = blk;
+  }
+  walk->next_ = blk;
+}
+
+// ================================================================================================
+HeapBlock* VmHeap::SplitBlock(HeapBlock* blk, size_t tailsize) {
+  // Create a new block from the beginning of the current
+  HeapBlock* nb = new HeapBlock(blk->owner_, tailsize, blk->offset_);
+
+  // Resize the old block
+  blk->offset_ += tailsize;
+  blk->size_ -= tailsize;
+  return nb;
+}
+
+// ================================================================================================
+void VmHeap::Join2Blocks(HeapBlock* first, HeapBlock* second) const {
+  // Do the join
+  first->size_ = first->size_ + second->size_;
+  first->next_ = second->next_;
+  if (second->next_) {
+      second->next_->prev_ = first;
+  }
+  delete second;
+}
+
+// ================================================================================================
+void VmHeap::MergeBlock(HeapBlock** head, HeapBlock* blk) {
+  InsertBlock(head, blk);
+
+  // Merge with successor if possible
+  if ((blk->next_ != nullptr) && (blk->offset_ + blk->size_ == blk->next_->offset_)) {
+      Join2Blocks(blk, blk->next_);
+  }
+
+  // Merge with predecessor if possible
+  if ((blk->prev_ != nullptr) && (blk->prev_->offset_ + blk->prev_->size_ == blk->offset_)) {
+      Join2Blocks(blk->prev_, blk);
+  }
+}
+
+} // namespace amd
@@ -0,0 +1,171 @@
+/* Copyright (c) 2025 Advanced Micro Devices, Inc.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE. */
+
+#pragma once
+
+#include "top.hpp"
+#include "device/device.hpp"
+#include "object.hpp"
+#include "commandqueue.hpp"
+
+namespace amd {
+class VmHeap;
+class HeapBlock;
+
+class HeapBlock : public amd::HeapObject {
+public:
+  friend VmHeap;
+  //! Constructor
+  HeapBlock(
+      VmHeap* owner = nullptr,  //!< VmHeap object that owns this heap block
+      size_t size = 0,          //!< Heap block size for allocation
+      size_t offset = 0)        //!< Heap block offset
+      : owner_(owner)
+      , size_(size)
+      , offset_(offset)
+      , next_(nullptr)
+      , prev_(nullptr)
+      , busy_(false)
+      {}
+
+  //! Destructor does some sanity checks
+  ~HeapBlock() { assert(!busy_ && "The blocked must be destroyed explicitly!"); }
+
+  //! Gets the offset
+  size_t Offset() const { return offset_; }
+
+private:
+  HeapBlock() = delete;
+  HeapBlock(const HeapBlock&) = delete;
+  HeapBlock& operator=(const HeapBlock&) = delete;
+
+  VmHeap*     owner_;   //!< Heap that owns this block
+  size_t      size_;    //!< Size of the block in bytes
+  size_t      offset_;  //!< Offset of this block in the heap
+  HeapBlock*  next_;    //!< Next block on the list, or nullptr
+  HeapBlock*  prev_;    //!< Previous block on the list, or nullptr
+  bool        busy_;    //!< True if the block is in use
+};
+
+class VmHeap {
+public:
+  static const size_t kChunkSize = 32 * Mi; //!< Chunk size, must be power of 2
+  static const size_t kMinBlockAlignment = 256;
+  VmHeap(Device* device,    //!< GPU device object
+         HostQueue& queue   //!< Queue, used for map/unmap of physical memory
+         )
+      : VmHeap(device, queue, device->info().globalMemSize_ / 8, kChunkSize) {}
+
+  VmHeap(Device* device,    //!< GPU device object
+         HostQueue& queue,  //!< Queue, usde for map/unmap of physical memory
+         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
+         );
+
+  //! Ceates heap object. Reserves virtual address range for the heap operation
+  bool Create();
+
+  //! Heap destructor
+  virtual ~VmHeap();
+
+  //! 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) { unmap_threshold_ = threshold; }
+
+  //! Returns mapped memory size (total allocated physical memory)
+  uint64_t MappedSize() const { return mapped_size_; }
+
+  //! Returns mapped memory size (allocated physical memory) without actual allocations
+  uint64_t FreeMappedSize() const { return mapped_size_ - (va_size_ - free_size_); }
+
+private:
+  VmHeap() = delete;
+  VmHeap(const VmHeap&) = delete;
+  VmHeap& operator=(const VmHeap&) = delete;
+
+  //! Reseves address range for memory allocations
+  address ReserveAddressRange(address start, size_t size, size_t alignment);
+
+  //! Releases address range specified by the address
+  bool ReleaseAddressRange(void* addr);
+
+  //! Commits actual physical memory on the specified address
+  bool CommitMemory(void* addr, size_t size);
+
+  //! Uncommits physical memory from the spcified address
+  bool UncommitMemory(void* addr, size_t size);
+
+  HeapBlock* AllocBlock(size_t size  //! The allocation size
+                        );
+
+  //! Release memory back to a heap
+  void FreeBlock(HeapBlock* blk);
+
+  //! Insert a block into a list
+  void InsertBlock(HeapBlock** list, HeapBlock* node);
+
+  //! Merge a block into a list
+  void MergeBlock(HeapBlock** list, HeapBlock* node);
+
+  //! Remove a block from a list
+  void DetachBlock(HeapBlock** list, HeapBlock* node);
+
+  //! Splits a block into two pieces
+  HeapBlock* SplitBlock(HeapBlock* node, size_t size);
+
+  //! Maps physical memory into specified address space
+  bool MapPhysMemory(size_t offset, size_t size);
+
+  //! Unmaps physical memory from the specified address
+  void UnmapPhysMemory(size_t offset, size_t size);
+
+  //! Join two blocks, transferring the size of the second into the first and deleting the second
+  void Join2Blocks(HeapBlock* first, HeapBlock* second) const;
+
+  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
+  HeapBlock*    busy_list_ = nullptr;     //!< Head block for busy list
+  size_t        free_size_ = 0;           //!< Total free size of the heap (both mapped and unmapped)
+  size_t        va_size_ = 0;             //!< Heap virtual address space size
+  size_t        block_alignment_ = 1;     //!< Size of an allocation page
+  size_t        chunk_size_ = 0;          //!< Chunk size (min physical allocation for the growth)
+  uint64_t      unmap_threshold_ = 0;     //!< Unmap threshold in bytes,used to release phys memory
+  uint64_t      mapped_size_ = 0;         //!< Size of mapped memory
+  uint64_t      max_mapped_size_ = 0;     //!< Max size of mapped memory in this heap
+  bool          created_ = false;         //!< Used for deferred VM heap allocation
+  amd::Monitor  lock_;                    //!< Lock to serialise heap accesses
+  Device*       device_;                  //!< Device that owns this heap
+  HostQueue&    map_queue_;               //!< Queue, used to map/unmap
+
+  std::vector<bool> mapped_mem_;  //!< A map of mapped memory, the size is total_size/chunk_size
+};
+
+} // namespace amd
@@ -202,6 +202,8 @@ release(bool, HIP_MEM_POOL_SUPPORT, true,                                     \
        "Enables memory pool support in HIP")                                 \
 release(bool, HIP_MEM_POOL_USE_VM, true,                                      \
        "Enables memory pool support in HIP")                                 \
+release(bool, DEBUG_HIP_MEM_POOL_VMHEAP, false,                               \
+        "Enables virtual memory for memory pools")                            \
 release(bool, PAL_HIP_IPC_FLAG, true,                                         \
        "Enable interprocess flag for device allocation in PAL HIP")          \
 release(uint, PAL_FORCE_ASIC_REVISION, 0,                                     \